[wip] Ascend 020

vllm_fl/dispatch/backends/vendor/ascend/ascend.py

@@ -146,3 +146,83 @@ def attention_backend(self, use_mla: bool = False, use_sparse: bool = False) ->
                 raise NotImplementedError("MLA with sparse attention is not implemented for Ascend yet.")
             return "vllm_fl.dispatch.backends.vendor.ascend.impl.attention.AscendMLABackend"
         return "vllm_fl.dispatch.backends.vendor.ascend.impl.attention.AscendAttentionBackend"
+
+    def invoke_fused_moe_triton_kernel(
+        self,
+        A,
+        B,
+        C,
+        A_scale,
+        B_scale,
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        mul_routed_weight,
+        top_k,
+        config,
+        compute_type=None,
+        use_fp8_w8a8=False,
+        use_int8_w8a8=False,
+        use_int8_w8a16=False,
+        use_int4_w4a16=False,
+        per_channel_quant=False,
+        block_shape=None,
+        B_bias=None,
+    ):
+        """Ascend NPU fused MoE kernel using torch.mm.
+
+        Replaces the FlagGems Triton kernel which overflows the NPU's
+        unified buffer on certain model shapes.
+        """
+        from .impl.fused_moe_kernel import invoke_fused_moe_torch
+        invoke_fused_moe_torch(
+            A, B, C, A_scale, B_scale, topk_weights,
+            sorted_token_ids, expert_ids, num_tokens_post_padded,
+            mul_routed_weight, top_k, config,
+            use_fp8_w8a8=use_fp8_w8a8,
+            use_int8_w8a8=use_int8_w8a8,
+            B_bias=B_bias,
+        )
+
+    def moe_align_block_size(
+        self,
+        topk_ids,
+        block_size,
+        num_experts,
+        expert_map=None,
+        pad_sorted_ids=False,
+        ignore_invalid_experts=False,
+    ):
+        """Pure-torch moe_align_block_size for Ascend NPU.
+
+        Replaces the FlagGems Triton kernel which causes DDR address OOB
+        errors on Ascend NPU hardware.
+        """
+        from .impl.fused_moe_kernel import moe_align_block_size_torch
+        return moe_align_block_size_torch(
+            topk_ids, block_size, num_experts, expert_map,
+            pad_sorted_ids, ignore_invalid_experts,
+        )
+
+    def moe_sum(self, inp, out):
+        """Pure-torch moe_sum: sum over top_k dimension."""
+        # inp is (M, top_k, N), out is (M, N)
+        # Avoid out= parameter which can cause NPU issues
+        result = inp.sum(dim=1)
+        out.copy_(result)
+
+    def topk_softmax(
+        self, topk_weights, topk_indices, token_expert_indices, gating_output,
+        renormalize=False,
+    ):
+        """Pure-torch topk_softmax for Ascend NPU."""
+        scores = torch.softmax(gating_output.float(), dim=-1)
+        topk = topk_weights.shape[1]
+        tk_weights, tk_indices = torch.topk(scores, k=topk, dim=-1)
+        topk_weights.copy_(tk_weights.to(topk_weights.dtype))
+        topk_indices.copy_(tk_indices.to(topk_indices.dtype))
+        if renormalize:
+            s = topk_weights.sum(dim=-1, keepdim=True)
+            topk_weights.div_(s.clamp(min=1e-8))
+        return topk_weights, topk_indices

vllm_fl/dispatch/backends/vendor/ascend/impl/attention.py

@@ -37,6 +37,7 @@
 from vllm.config import VllmConfig, get_current_vllm_config
 from vllm.utils.math_utils import cdiv
 from vllm.v1.attention.backend import AttentionCGSupport
+from vllm.v1.attention.backends.registry import AttentionBackendEnum, register_backend
 from vllm.v1.attention.backends.utils import CommonAttentionMetadata
 
 from vllm_fl.dispatch.backends.vendor.ascend.impl.attention_mask import (
@@ -215,6 +216,7 @@ class AscendAttentionMetadataBuilder:
     # ACL graph support - ALWAYS means full graph capture is supported
     aclgraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
     reorder_batch_threshold: ClassVar[int] = 1
+    supports_update_block_table: bool = False
 
     @staticmethod
     def get_cudagraph_support(vllm_config, kv_cache_spec) -> AttentionCGSupport:
@@ -433,7 +435,7 @@ class AscendAttentionBackend(AttentionBackend):
 
     @staticmethod
     def get_name() -> str:
-        return "ASCEND_FL"
+        return "CUSTOM"
 
     @staticmethod
     def get_impl_cls() -> Type["AscendAttentionBackendImpl"]:
@@ -443,6 +445,15 @@ def get_impl_cls() -> Type["AscendAttentionBackendImpl"]:
     def get_builder_cls() -> Type["AscendAttentionMetadataBuilder"]:
         return AscendAttentionMetadataBuilder
 
+    @staticmethod
+    def get_supported_kernel_block_sizes() -> list[int]:
+        # Ascend fused_infer_attention_score and paged_attention kernels
+        # are validated for block size 128 in vllm-ascend. Allowing the
+        # default MultipleOf(1) lets the V1 engine pick unsupported merged
+        # storage block sizes (e.g. 784 for Qwen3.5 hybrid models), which
+        # causes aclnnFusedInferAttentionScoreV3 to fail with error 561002.
+        return [128]
+
     @staticmethod
     def get_kv_cache_shape(
         num_blocks: int,
@@ -488,6 +499,12 @@ def get_supported_block_size() -> list[int]:
         return [128]
 
 
+register_backend(
+    AttentionBackendEnum.CUSTOM,
+    "vllm_fl.dispatch.backends.vendor.ascend.impl.attention.AscendAttentionBackend",
+)
+
+
 class AscendAttentionBackendImpl(AttentionImpl):
     """
     Ascend attention implementation using native torch_npu operators.
@@ -561,20 +578,16 @@ def _get_fia_params(
             value = self.value_cache.view(num_block, block_size, -1)
             actual_seq_lengths_kv = attn_metadata.seq_lens_list
         elif attn_metadata.attn_state == AscendAttentionState.DecodeOnly:
-            # num_block, block_size, _, _ = self.key_cache.shape
-            # key = self.key_cache.view(num_block, block_size, -1)
-            # value = self.value_cache.view(num_block, block_size, -1)
-            key = self.key_cache.view(-1, block_size, 256)
-            value = self.value_cache.view(-1, block_size, 256)
+            num_block, block_size, _, _ = self.key_cache.shape
+            key = self.key_cache.view(num_block, block_size, -1)
+            value = self.value_cache.view(num_block, block_size, -1)
             block_table = attn_metadata.block_tables
             actual_seq_lengths_kv = attn_metadata.seq_lens_list
         else:
             # ChunkedPrefill
-            # num_block, block_size, _, _ = self.key_cache.shape
-            # key = self.key_cache.view(num_block, block_size, -1)
-            # value = self.value_cache.view(num_block, block_size, -1)
-            key = self.key_cache.view(-1, block_size, 256)
-            value = self.value_cache.view(-1, block_size, 256)
+            num_block, block_size, _, _ = self.key_cache.shape
+            key = self.key_cache.view(num_block, block_size, -1)
+            value = self.value_cache.view(num_block, block_size, -1)
             block_table = attn_metadata.block_tables
             actual_seq_lengths_kv = attn_metadata.seq_lens_list
 
@@ -596,7 +609,6 @@ def reshape_and_cache(
             # TODO(yxa): block_table.py: CUDA uses int64, NPU uses int32.
             if slots.dtype != torch.int32:
                 slots = slots.to(torch.int32)
-            # Use torch_npu reshape_and_cache
             torch_npu._npu_reshape_and_cache(
                 key=key[:attn_metadata.num_actual_tokens],
                 value=value[:attn_metadata.num_actual_tokens],
@@ -625,9 +637,9 @@ def forward_fused_infer_attention(
             key = key[:num_tokens]
             value = value[:num_tokens]
 
-        # Determine sparse_mode based on mask availability
-        # sparse_mode=3 requires attn_mask; sparse_mode=0 does not
-        # sparse_mode = 3 if attn_metadata.attn_mask is not None else 0
+        # sparse_mode: 3 = causal with mask, 0 = no mask
+        sparse_mode = 3 if attn_metadata.attn_mask is not None else 0
+
         attn_output, _ = torch_npu.npu_fused_infer_attention_score(
             query=query,
             key=key,
@@ -641,7 +653,7 @@ def forward_fused_infer_attention(
             num_key_value_heads=self.num_kv_heads,
             num_heads=self.num_heads,
             scale=self.scale,
-            sparse_mode=3,
+            sparse_mode=sparse_mode,
         )
 
         attn_output = attn_output.view(num_tokens, self.num_heads, self.head_size)
@@ -779,8 +791,9 @@ def forward(
             return output.fill_(0)
 
         # Reshape and cache KV
-        if attn_metadata != AscendAttentionState.DecodeOnly:
-            kv_cache = [i.contiguous() for i in kv_cache]
+        # Note: kv_cache[0]/[1] may be non-contiguous views of a
+        # [2, num_blocks, ...] tensor.  _npu_reshape_and_cache handles
+        # them directly via slot_indices — no contiguous copy needed.
         if key is not None and value is not None:
             key = key.contiguous()
             value = value.contiguous()