fix(ascend): repair Qwen3.6 accuracy on Ascend (attention/moe/gdn)

修复 Qwen3.6 在 Ascend 上的精度问题（注意力/MoE/GDN 四项）

1. attention: drop the mathematically invalid head-split fallback for
   head_size=256 prefill and call npu_fused_infer_attention_score (TND)
   directly, matching vllm-ascend. Splitting a 256-dim head into two
   independent 128-dim softmax attentions is not an identity transform
   (measured max err 0.0076 direct vs 4.30 head-split on NPU).
   注意力：删除 head_dim=256 时数学上不成立的拆头回退路径，prefill 直连
   FIA（TND），与 vllm-ascend 一致。将 256 维头拆成两个独立 128 维
   softmax 注意力并非恒等变换（NPU 实测直连误差 0.0076，旧路径 4.30）。

2. attention: derive attn_state from num_computed_tokens; use
   PrefillNoCache only when every request starts at token 0. Pure-prefill
   batches containing prefix-cache hits or chunked-prefill continuations
   previously dropped all cached KV (few-shot context lost).
   注意力：attn_state 依据 num_computed_tokens 判定，仅当所有请求从第 0
   个 token 开始才走 PrefillNoCache。此前含 prefix 命中或续写 chunk 的
   纯 prefill batch 会丢弃全部已缓存 KV（few-shot 上下文丢失）。

3. moe: cast router probs to hidden dtype (bf16) before
   npu_moe_token_unpermute, matching vllm-ascend; the op expects bf16
   probs.
   MoE：npu_moe_token_unpermute 前将路由 probs 转为 hidden dtype
   （bf16），与 vllm-ascend 一致；该算子 probs 期望 bf16。

4. gdn: default VLLM_FL_DISABLE_PTO_GDN to 1. The PTO megakernel carries
   A/A_inv/recurrent state in fp16, while the Triton chunk path keeps fp32
   state; the fp16 pipeline measurably degrades accuracy. Re-enable after
   reworking the kernel state path to fp32.
   GDN：VLLM_FL_DISABLE_PTO_GDN 默认改为 1。PTO megakernel 的
   A/A_inv/循环状态均为 fp16，而 Triton chunk 路径保持 fp32 状态，fp16
   链路实测显著拉低精度。待 kernel 状态链路 fp32 化后可重新开启。

Verified on 910B: op-level tests for FIA head256 / chunked-prefill FIA /
recurrent GDN decode / causal-conv1d decode all pass; gsm8k Qwen3.6-27B
(128 samples) recovers to 0.9453 (previously in the 0.47-0.88 range).
验证：FIA head256 / chunked prefill FIA / GDN recurrent decode /
causal-conv1d decode 算子级测试全部通过；gsm8k Qwen3.6-27B（128 样本）
恢复至 0.9453（此前在 0.47-0.88 区间）。

Author: OrangePure

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

@@ -53,10 +53,6 @@
 
 logger = logging.getLogger(__name__)
 
-# Ascend npu_fused_infer_attention_score with input_layout="TND" only supports
-# head_dim in {64, 128, 192} (plus the special qD=kD=192, vD=128 case).
-_PFA_TND_SUPPORTED_HEAD_DIMS = frozenset({64, 128, 192})
-
 # Check torch_npu availability and setup NPU compatibility
 _TORCH_NPU_AVAILABLE = False
 try:
@@ -324,7 +320,9 @@ def build(
 
         # Determine attention state
         attn_state = self._determine_attn_state(
-            num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens
+            num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens,
+            num_computed_tokens_cpu=getattr(
+                common_attn_metadata, 'num_computed_tokens_cpu', None),
         )
 
         # Create attention mask based on state
@@ -368,13 +366,21 @@ def _determine_attn_state(
         num_prefills: int,
         num_decode_tokens: int,
         num_prefill_tokens: int,
+        num_computed_tokens_cpu: Optional[torch.Tensor] = None,
     ) -> AscendAttentionState:
         """Determine attention state based on batch composition."""
         if num_prefills == 0:
             return AscendAttentionState.DecodeOnly
+        # PrefillNoCache is only valid when every request starts from token 0.
+        # A pure-prefill batch may still contain prefix-cache hits or chunked
+        # prefill continuations; those must attend to the paged KV cache via
+        # the ChunkedPrefill path (same handling as vllm-ascend).
+        if (num_computed_tokens_cpu is not None
+                and len(num_computed_tokens_cpu) > 0
+                and int(num_computed_tokens_cpu.max()) > 0):
+            return AscendAttentionState.ChunkedPrefill
         elif num_decodes == 0 and num_prefill_tokens > 0:
-            # Pure prefill - check if cache hit or no cache
-            # For simplicity, use ChunkedPrefill as default
+            # Fresh prefill with no cached context
             return AscendAttentionState.PrefillNoCache
         else:
             # Mixed decode and prefill
@@ -563,13 +569,6 @@ def __init__(
         self.num_queries_per_kv = self.num_heads // self.num_kv_heads
         self.key_cache = None
         self.value_cache = None
-        self._use_fusion_fallback = self.head_size not in _PFA_TND_SUPPORTED_HEAD_DIMS
-        if self._use_fusion_fallback:
-            logger.info(
-                "AscendAttentionBackendImpl: head_size=%d is not supported by "
-                "npu_fused_infer_attention_score TND layout, falling back to "
-                "head-split attention (2x%d).", self.head_size,
-                self.head_size // 2)
 
     @classmethod
     def update_graph_params(
@@ -707,130 +706,6 @@ def reshape_and_cache(
             )
         return key, value
 
-    def _gather_kv_cache(
-        self,
-        attn_metadata: AscendMetadata,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """Gather paged KV cache into dense (T_kv, num_kv_heads, head_size).
-
-        Only the blocks actually referenced by each request's sequence length
-        are gathered.  The block table is padded to ``max_num_blocks_per_req``,
-        so gathering the whole table would allocate ``batch_size`` times more
-        memory than necessary and can OOM on long-context configs.
-        """
-        if self.key_cache is None or self.value_cache is None:
-            raise RuntimeError(
-                "key_cache/value_cache is not initialized for KV gather")
-
-        seq_lens = attn_metadata.seq_lens_list
-        batch_size = len(seq_lens)
-        block_size = self.key_cache.shape[1]
-        block_table = attn_metadata.block_tables[:batch_size]
-
-        k_list: List[torch.Tensor] = []
-        v_list: List[torch.Tensor] = []
-        for i, sl in enumerate(seq_lens):
-            num_blocks = (sl + block_size - 1) // block_size
-            block_ids = block_table[i, :num_blocks].long()
-            gathered_k = self.key_cache[block_ids].view(
-                -1, self.num_kv_heads, self.head_size)[:sl]
-            gathered_v = self.value_cache[block_ids].view(
-                -1, self.num_kv_heads, self.head_size)[:sl]
-            k_list.append(gathered_k)
-            v_list.append(gathered_v)
-
-        return torch.cat(k_list, dim=0), torch.cat(v_list, dim=0)
-
-    def _split_heads_for_pfa(
-        self,
-        x: torch.Tensor,
-        num_heads: int,
-    ) -> torch.Tensor:
-        """Split head dimension to a supported PFA TND head size.
-
-        Ascend's npu_fused_infer_attention_score TND layout only supports
-        head_dim in {64, 128, 192}.  For unsupported head sizes such as 256,
-        we split each head into two heads of half the dimension and double the
-        head count, which is mathematically equivalent for self-attention.
-
-        Accepts either 2D input [T, num_heads * head_size] or 3D input
-        [T, num_heads, head_size].  Output shape: [T, num_heads * 2,
-        head_size // 2].
-        """
-        head_size_128 = self.head_size // 2
-        if x.dim() == 2:
-            x = x.view(-1, num_heads, self.head_size)
-        return x.view(-1, num_heads, 2, head_size_128).reshape(
-            -1, num_heads * 2, head_size_128)
-
-    def _merge_heads_from_pfa(
-        self,
-        x: torch.Tensor,
-    ) -> torch.Tensor:
-        """Merge split heads back to original head dimension.
-
-        Input shape:  [T, num_heads * 2, head_size // 2]
-        Output shape: [T, num_heads * head_size]
-        """
-        head_size_128 = self.head_size // 2
-        return x.view(-1, self.num_heads, 2, head_size_128).reshape(
-            -1, self.num_heads * self.head_size)
-
-    def _forward_fusion_attention(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: AscendMetadata,
-        output: torch.Tensor,
-    ) -> torch.Tensor:
-        """Forward pass for unsupported head sizes using head-split PFA.
-
-        Splits head_dim into two supported heads and runs
-        npu_fused_infer_attention_score (PFA) with the standard TND path.
-        This keeps the implementation compatible with ACL graph capture while
-        avoiding the per-batch dense causal mask limitation of
-        npu_fusion_attention.
-        """
-        num_tokens = attn_metadata.actual_seq_lengths_q[-1]
-        query = query[:num_tokens]
-
-        if attn_metadata.attn_state == AscendAttentionState.PrefillNoCache:
-            key = key[:num_tokens]
-            value = value[:num_tokens]
-            actual_seq_kvlen = attn_metadata.actual_seq_lengths_q
-        else:
-            key, value = self._gather_kv_cache(attn_metadata)
-            actual_seq_kvlen = [
-                sum(attn_metadata.seq_lens_list[:i + 1])
-                for i in range(len(attn_metadata.seq_lens_list))
-            ]
-
-        query = self._split_heads_for_pfa(query, self.num_heads)
-        key = self._split_heads_for_pfa(key, self.num_kv_heads)
-        value = self._split_heads_for_pfa(value, self.num_kv_heads)
-
-        attn_output, _ = torch_npu.npu_fused_infer_attention_score(
-            query=query,
-            key=key,
-            value=value,
-            atten_mask=attn_metadata.attn_mask,
-            block_table=None,
-            input_layout="TND",
-            block_size=AscendAttentionBackend.get_supported_block_size()[0],
-            actual_seq_lengths=attn_metadata.actual_seq_lengths_q,
-            actual_seq_lengths_kv=actual_seq_kvlen,
-            num_key_value_heads=self.num_kv_heads * 2,
-            num_heads=self.num_heads * 2,
-            scale=self.scale,
-            sparse_mode=3,
-        )
-
-        attn_output = self._merge_heads_from_pfa(attn_output)
-        attn_output = attn_output.view(num_tokens, self.num_heads, self.head_size)
-        output[:num_tokens] = attn_output[:num_tokens]
-        return output
-
     def forward_fused_infer_attention(
         self,
         query: torch.Tensor,
@@ -840,11 +715,6 @@ def forward_fused_infer_attention(
         output: torch.Tensor,
     ) -> torch.Tensor:
         """Forward pass using fused_infer_attention_score."""
-        if (self._use_fusion_fallback
-                and attn_metadata.attn_state != AscendAttentionState.DecodeOnly):
-            return self._forward_fusion_attention(
-                query, key, value, attn_metadata, output)
-
         key, value, block_size, block_table, actual_seq_lengths_kv = \
             self._get_fia_params(key, value, attn_metadata)
 

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

@@ -228,10 +228,11 @@ def _ascendc_fused_experts_impl(
     )[0]
 
     # Scatter/sum the expert outputs back to the token dimension.
+    # Match vllm-ascend: npu_moe_token_unpermute expects bf16 probs.
     out = torch_npu.npu_moe_token_unpermute(
         permuted_tokens=down,
         sorted_indices=expanded_row_idx.abs(),
-        probs=probs,
+        probs=probs.to(down.dtype) if probs is not None else None,
     )
 
     if inplace:
@@ -357,10 +358,11 @@ def _torch_fused_experts_impl(
     )[0]
 
     # Scatter/sum the expert outputs back to the token dimension.
+    # Match vllm-ascend: npu_moe_token_unpermute expects bf16 probs.
     out = torch_npu.npu_moe_token_unpermute(
         permuted_tokens=down,
         sorted_indices=sorted_indices,
-        probs=probs,
+        probs=probs.to(down.dtype) if probs is not None else None,
     )
 
     if inplace:

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

@@ -274,7 +274,11 @@ def _pto_available() -> bool:
     global _PTO_AVAILABLE
     if _PTO_AVAILABLE is not None:
         return _PTO_AVAILABLE
-    if os.environ.get("VLLM_FL_DISABLE_PTO_GDN", "0") == "1":
+    # Default OFF: the PTO megakernel carries A/A_inv/recurrent state in fp16
+    # (see vllm_fl/ops/pto_chunk_gdn/mega_kernel.py), which measurably degrades
+    # accuracy vs the fp32-state Triton chunk path. Re-enable with
+    # VLLM_FL_DISABLE_PTO_GDN=0 once the kernel is reworked to fp32 state.
+    if os.environ.get("VLLM_FL_DISABLE_PTO_GDN", "1") == "1":
         logger.info("VLLM_FL_DISABLE_PTO_GDN=1, keep Triton chunk path")
         _PTO_AVAILABLE = False
         return False