upgrade vllm to 0.20.2 on ascend platform (#307)

<!--
 Copyright 2026 FlagOS Contributors

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 -->

### PR Category
Vendor

### Description

test cases:

Test | Model | Type | Result
-- | -- | -- | --
Text math (7×8) | 27B | Text | ✅ 56
Text knowledge (Capital of France) | 27B | Text | ✅ Paris
Text math (7×8) | 35B | Text | ✅ 56
Text knowledge (Capital of France) | 35B | Text | ✅ Paris
Image description | 27B | Image | ✅ Correct (black text, white
background)
Image OCR | 27B | Image | ✅ "Hello"
Image description | 35B | Image | ✅ Correct (black text, white
background)
Image OCR | 35B | Image | ✅ "Hello"
Concurrent text (8 requests) | 27B | Text | ✅ 8/8
Concurrent images (8 requests) | 27B | Image | ✅ 8/8
Concurrent mixed (8 requests) | 27B | Mixed | ✅ 8/8
Concurrent text (8 requests) | 35B | Text | ✅ 8/8
Concurrent images (8 requests) | 35B | Image | ✅ 8/8
Concurrent mixed (8 requests) | 35B | Mixed | ✅ 8/8
Prime numbers | 27B | Text | ✅ 2, 3, 5, 7, 11
Prime numbers | 35B | Text | ✅ 2, 3, 5, 7, 11
Code generation | 27B | Text | ✅ Valid Python
Code generation | 35B | Text | ✅ Valid Python


### Usage
```
export VLLM_PLUGINS=fl
export VLLM_FL_PLATFORM=ascend


vllm serve /models/Qwen3.6-35B-A3B \
  --host 0.0.0.0 \
  --trust-remote-code \
  --max-model-len 4096 \
  --enforce-eager \
  --served-model-name qwen \
  --tensor-parallel-size 2 \
  --gpu-memory-utilization 0.8 \
  --no-enable-chunked-prefill \
  --no-async-scheduling \
  --no-enable-prefix-caching
```

Author: cyber-pioneer

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
 
@@ -593,16 +606,16 @@ def reshape_and_cache(
                 self.key_cache, self.value_cache = kv_cache[0], kv_cache[1]
             slots = attn_metadata.slot_mapping
             # torch_npu requires int32 for slot_indices
-            # 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
+
+            num_actual = attn_metadata.num_actual_tokens
             torch_npu._npu_reshape_and_cache(
-                key=key[:attn_metadata.num_actual_tokens],
-                value=value[:attn_metadata.num_actual_tokens],
+                key=key[:num_actual],
+                value=value[:num_actual],
                 key_cache=self.key_cache,
                 value_cache=self.value_cache,
-                slot_indices=slots[:attn_metadata.num_actual_tokens]
+                slot_indices=slots[:num_actual]
             )
         return key, value
 
@@ -625,9 +638,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 +654,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 +792,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()

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

@@ -713,3 +713,97 @@ def grid(META):
     if unsqueeze:
         out = out.squeeze(-1)
     return out.to(original_x_dtype)
+
+
+def causal_conv1d_update_ref(
+    x: torch.Tensor,
+    conv_state: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    activation: bool | str | None = None,
+    conv_state_indices: torch.Tensor | None = None,
+    num_accepted_tokens: torch.Tensor | None = None,
+    query_start_loc: torch.Tensor | None = None,
+    max_query_len: int = -1,
+    pad_slot_id: int = PAD_SLOT_ID,
+    block_idx_last_scheduled_token: torch.Tensor | None = None,
+    initial_state_idx: torch.Tensor | None = None,
+    validate_data=False,
+):
+    """Pure-PyTorch causal_conv1d_update for decode (single token per seq).
+
+    Handles the simple case: x is [batch, dim] (single token), conv_state
+    is [num_cache_lines, dim, width-1], and we do a sliding window update.
+    """
+    if isinstance(activation, bool):
+        activation = "silu" if activation is True else None
+    elif activation is not None:
+        assert activation in ["silu", "swish"]
+
+    original_x_dtype = x.dtype
+    x = x.to(conv_state.dtype)
+    unsqueeze = query_start_loc is None and x.dim() == 2
+    if unsqueeze:
+        x = x.unsqueeze(-1)  # [batch, dim, 1]
+
+    if query_start_loc is not None:
+        # varlen mode
+        batch = conv_state_indices.size(0)
+        dim = x.size(1)
+    else:
+        batch, dim, seqlen = x.shape
+
+    _, width = weight.shape
+    state_len = width - 1
+
+    out = torch.empty_like(x)
+
+    if query_start_loc is None:
+        # Simple batched mode: x is [batch, dim, seqlen]
+        seqlen = x.shape[-1]
+        for b_idx in range(batch):
+            if conv_state_indices is not None:
+                s_idx = conv_state_indices[b_idx].item()
+                if pad_slot_id is not None and s_idx == pad_slot_id:
+                    out[b_idx] = 0
+                    continue
+            else:
+                s_idx = b_idx
+
+            state = conv_state[s_idx]  # [dim, state_len]
+            for t in range(seqlen):
+                # Shift state left, append new token
+                state = torch.cat([state[:, 1:], x[b_idx, :, t:t+1]], dim=-1)
+                # Dot product with weight for each channel (depthwise conv)
+                val = (state * weight).sum(dim=-1)  # [dim]
+                if bias is not None:
+                    val = val + bias
+                if activation in ["silu", "swish"]:
+                    val = val * torch.sigmoid(val)
+                out[b_idx, :, t] = val
+            conv_state[s_idx] = state
+    else:
+        # Varlen mode
+        cu_cpu = query_start_loc.cpu().tolist()
+        for b_idx in range(batch):
+            s_idx = conv_state_indices[b_idx].item()
+            if pad_slot_id is not None and s_idx == pad_slot_id:
+                continue
+
+            bos = cu_cpu[b_idx] if b_idx < len(cu_cpu) else cu_cpu[-1]
+            eos = cu_cpu[b_idx + 1] if b_idx + 1 < len(cu_cpu) else cu_cpu[-1]
+
+            state = conv_state[s_idx]  # [dim, state_len]
+            for t in range(bos, eos):
+                state = torch.cat([state[:, 1:], x[t:t+1, :].T], dim=-1)
+                val = (state * weight).sum(dim=-1)
+                if bias is not None:
+                    val = val + bias
+                if activation in ["silu", "swish"]:
+                    val = val * torch.sigmoid(val)
+                out[t, :] = val  # varlen: out is [num_tokens, dim]
+            conv_state[s_idx] = state
+
+    if unsqueeze:
+        out = out.squeeze(-1)
+    return out.to(original_x_dtype)