fix(dsv4/sparse_attn): per-row inverse-RoPE to avoid a5 gather lowering bug

models/deepseek/v4/decode_sparse_attn.py

@@ -1,3 +1,3 @@
 # Copyright (c) PyPTO Contributors.
 # This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 # CANN Open Software License Agreement Version 2.0 (the "License").
@@ -354,7 +354,12 @@
                     qk_exp = pl.exp(pl.row_expand_sub(qk_scores, qk_mi))
                     qk_li = pl.row_sum(qk_exp)
                     qk_exp_bf16 = pl.cast(qk_exp, target_type=pl.BF16, mode="rint")
-                    qk_oi = pl.matmul(qk_exp_bf16, qk_kv, out_dtype=pl.FP32)
+                    # Materialize qk_exp_bf16 to GM before the PV cube: the a5 codegen
+                    # mis-fuses the exp->cast->cube chain in this composed task, losing
+                    # ~15% precision. Forcing a GM round-trip breaks the fusion (fixes it).
+                    _qk_exp_gm = pl.create_tensor([QK_M_TILE, ATTN_K_TILE], dtype=pl.BF16)
+                    _qk_exp_gm[:, :] = qk_exp_bf16
+                    qk_oi = pl.matmul(_qk_exp_gm, qk_kv, out_dtype=pl.FP32)
                     for qk_sub in pl.unroll(QK_M_TILE // H_TILE):
                         qk_h_idx = qk_hb * (QK_M_TILE // H_TILE) + qk_sub
                         qk_r0 = qk_sub * H_TILE
@@ -438,27 +443,31 @@
         # head-invariant for token m_t, so col_expand them over the H_TILE head rows;
         # swap_idx (j^1) pairs the interleaved real/imag lanes. Rounded to bf16 (golden
         # also rounds inverse-RoPE to bf16) and packed into o_packed's rope columns.
-        m_col = pl.col_expand_mul(
-            pl.full([H_TILE, ROPE_DIM], dtype=pl.FP32, value=1.0),
-            pl.cast(pl.arange(0, [1, ROPE_DIM], dtype=pl.INT32), target_type=pl.FP32))
-        m_dup_f = pl.cast(pl.cast(pl.mul(m_col, 0.5), target_type=pl.INT32, mode="trunc"), target_type=pl.FP32)
-        m_lane = pl.sub(m_col, pl.mul(m_dup_f, 2.0))                                              # j%2
-        m_swap_idx = pl.cast(pl.sub(pl.add(m_col, 1.0), pl.mul(m_lane, 2.0)), target_type=pl.INT32)  # j^1
-        m_rope = n_full[0 : H_TILE, NOPE_DIM : HEAD_DIM]
-        m_cos_il = rope_cos_il[m_t : m_t + 1, 0 : ROPE_DIM]
-        m_sin_signed = rope_sin_signed[m_t : m_t + 1, 0 : ROPE_DIM]
-        m_swapped = pl.gather(m_rope, dim=-1, index=m_swap_idx)
-        m_rot = pl.add(pl.col_expand_mul(m_rope, m_cos_il), pl.col_expand_mul(m_swapped, m_sin_signed))
-        n_rope_bf16 = pl.cast(m_rot, target_type=pl.BF16, mode="rint")
-
+        # Inverse RoPE per head-row (1-row [1, ROPE_DIM] tiles). On A5 the
+        # tensor.gather flat-index lowering miscompiles when the source tile spans
+        # more than one FP32 vector box (>8 rows); the old fused [H_TILE=16, ROPE_DIM]
+        # gather corrupted the ROPE columns of all heads except the first of each
+        # 8-head group. Unrolling to 1-row tiles keeps every gather in one vector
+        # box (qkv_proj_rope passes with its [8, ROPE_DIM] tile for the same reason).
+        mr_ones = pl.full([1, ROPE_DIM], dtype=pl.FP32, value=1.0)
+        mr_col = pl.col_expand_mul(mr_ones, pl.cast(pl.arange(0, [1, ROPE_DIM], dtype=pl.INT32), target_type=pl.FP32))
+        mr_dup_f = pl.cast(pl.cast(pl.mul(mr_col, 0.5), target_type=pl.INT32, mode="trunc"), target_type=pl.FP32)
+        mr_lane = pl.sub(mr_col, pl.mul(mr_dup_f, 2.0))                                            # j%2
+        mr_swap_idx = pl.cast(pl.sub(pl.add(mr_col, 1.0), pl.mul(mr_lane, 2.0)), target_type=pl.INT32)  # j^1
+        mr_cos_il = rope_cos_il[m_t : m_t + 1, 0 : ROPE_DIM]                                       # [1, ROPE_DIM]
+        mr_sin_signed = rope_sin_signed[m_t : m_t + 1, 0 : ROPE_DIM]
         for n_hi in pl.range(H_TILE):
+            mr_row = n_full[n_hi : n_hi + 1, NOPE_DIM : HEAD_DIM]
+            mr_swapped = pl.gather(mr_row, dim=-1, index=mr_swap_idx)
+            mr_rot = pl.add(pl.mul(mr_row, mr_cos_il), pl.mul(mr_swapped, mr_sin_signed))
+            mr_bf16 = pl.cast(mr_rot, target_type=pl.BF16, mode="rint")
             n_gh = m_h0 + n_hi
             n_g = n_gh // HEADS_PER_GROUP
             n_hh = n_gh - n_g * HEADS_PER_GROUP
             n_pack_row = n_g * T + m_t
             n_col = n_hh * HEAD_DIM
             o_packed[n_pack_row : n_pack_row + 1, n_col : n_col + NOPE_DIM] = n_bf16[n_hi : n_hi + 1, 0 : NOPE_DIM]
-            o_packed[n_pack_row : n_pack_row + 1, n_col + NOPE_DIM : n_col + HEAD_DIM] = n_rope_bf16[n_hi : n_hi + 1, 0 : ROPE_DIM]
+            o_packed[n_pack_row : n_pack_row + 1, n_col + NOPE_DIM : n_col + HEAD_DIM] = mr_bf16
 
     # ========================================================================
     # Back-to-back grouped output projection (manual scope, PER-GROUP INT8 quant).

models/deepseek/v4/decode_sparse_attn_hca.py

@@ -1,3 +1,3 @@
 # Copyright (c) PyPTO Contributors.
 # This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 # CANN Open Software License Agreement Version 2.0 (the "License").
@@ -307,7 +307,12 @@
             qk_exp = pl.exp(pl.row_expand_sub(qk_scores, qk_mi))
             qk_li = pl.row_sum(qk_exp)
             qk_exp_bf16 = pl.cast(qk_exp, target_type=pl.BF16, mode="rint")
-            qk_oi = pl.matmul(qk_exp_bf16, qk_kv, out_dtype=pl.FP32)
+            # Materialize qk_exp_bf16 to GM before the PV cube: the a5 codegen
+            # mis-fuses the exp->cast->cube chain in this composed task, losing
+            # ~15% precision. Forcing a GM round-trip breaks the fusion (fixes it).
+            _qk_exp_gm = pl.create_tensor([QK_M_TILE, ATTN_K_TILE], dtype=pl.BF16)
+            _qk_exp_gm[:, :] = qk_exp_bf16
+            qk_oi = pl.matmul(_qk_exp_gm, qk_kv, out_dtype=pl.FP32)
             for qk_sub in pl.unroll(QK_M_TILE // H_TILE):
                 qk_h_idx = qk_hb * (QK_M_TILE // H_TILE) + qk_sub
                 qk_r0 = qk_sub * H_TILE
@@ -385,27 +390,28 @@
         # head-invariant for token m_t, so col_expand them over the H_TILE head rows;
         # swap_idx (j^1) pairs the interleaved real/imag lanes. Rounded to bf16 (golden
         # also rounds inverse-RoPE to bf16) and packed into o_packed's rope columns.
-        m_col = pl.col_expand_mul(
-            pl.full([H_TILE, ROPE_DIM], dtype=pl.FP32, value=1.0),
-            pl.cast(pl.arange(0, [1, ROPE_DIM], dtype=pl.INT32), target_type=pl.FP32))
-        m_dup_f = pl.cast(pl.cast(pl.mul(m_col, 0.5), target_type=pl.INT32, mode="trunc"), target_type=pl.FP32)
-        m_lane = pl.sub(m_col, pl.mul(m_dup_f, 2.0))                                              # j%2
-        m_swap_idx = pl.cast(pl.sub(pl.add(m_col, 1.0), pl.mul(m_lane, 2.0)), target_type=pl.INT32)  # j^1
-        m_rope = n_full[0 : H_TILE, NOPE_DIM : HEAD_DIM]
-        m_cos_il = rope_cos_il[m_t : m_t + 1, 0 : ROPE_DIM]
-        m_sin_signed = rope_sin_signed[m_t : m_t + 1, 0 : ROPE_DIM]
-        m_swapped = pl.gather(m_rope, dim=-1, index=m_swap_idx)
-        m_rot = pl.add(pl.col_expand_mul(m_rope, m_cos_il), pl.col_expand_mul(m_swapped, m_sin_signed))
-        n_rope_bf16 = pl.cast(m_rot, target_type=pl.BF16, mode="rint")
-
+        # Per head-row inverse RoPE (1-row tiles) -- A5 gather lowering corrupts
+        # the [H_TILE, ROPE_DIM] gather for tiles >8 rows (only first head of each
+        # 8-head group survived). Mirrors qkv_proj_rope's <=8-row rotation tile.
+        mr_ones = pl.full([1, ROPE_DIM], dtype=pl.FP32, value=1.0)
+        mr_col = pl.col_expand_mul(mr_ones, pl.cast(pl.arange(0, [1, ROPE_DIM], dtype=pl.INT32), target_type=pl.FP32))
+        mr_dup_f = pl.cast(pl.cast(pl.mul(mr_col, 0.5), target_type=pl.INT32, mode="trunc"), target_type=pl.FP32)
+        mr_lane = pl.sub(mr_col, pl.mul(mr_dup_f, 2.0))
+        mr_swap_idx = pl.cast(pl.sub(pl.add(mr_col, 1.0), pl.mul(mr_lane, 2.0)), target_type=pl.INT32)
+        mr_cos_il = rope_cos_il[m_t : m_t + 1, 0 : ROPE_DIM]
+        mr_sin_signed = rope_sin_signed[m_t : m_t + 1, 0 : ROPE_DIM]
         for n_hi in pl.range(H_TILE):
+            mr_row = n_full[n_hi : n_hi + 1, NOPE_DIM : HEAD_DIM]
+            mr_swapped = pl.gather(mr_row, dim=-1, index=mr_swap_idx)
+            mr_rot = pl.add(pl.mul(mr_row, mr_cos_il), pl.mul(mr_swapped, mr_sin_signed))
+            mr_bf16 = pl.cast(mr_rot, target_type=pl.BF16, mode="rint")
             n_gh = m_h0 + n_hi
             n_g = n_gh // HEADS_PER_GROUP
             n_hh = n_gh - n_g * HEADS_PER_GROUP
             n_pack_row = n_g * T + m_t
             n_col = n_hh * HEAD_DIM
             o_packed[n_pack_row : n_pack_row + 1, n_col : n_col + NOPE_DIM] = n_bf16[n_hi : n_hi + 1, 0 : NOPE_DIM]
-            o_packed[n_pack_row : n_pack_row + 1, n_col + NOPE_DIM : n_col + HEAD_DIM] = n_rope_bf16[n_hi : n_hi + 1, 0 : ROPE_DIM]
+            o_packed[n_pack_row : n_pack_row + 1, n_col + NOPE_DIM : n_col + HEAD_DIM] = mr_bf16
 
     # ========================================================================
     # Back-to-back grouped output projection (manual scope, PER-GROUP INT8 quant).

models/deepseek/v4/decode_sparse_attn_swa.py

@@ -1,3 +1,3 @@
 # Copyright (c) PyPTO Contributors.
 # This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 # CANN Open Software License Agreement Version 2.0 (the "License").
@@ -242,7 +242,11 @@
                 qk_exp = pl.exp(pl.row_expand_sub(qk_scores, qk_mi))
                 qk_li = pl.row_sum(qk_exp)
                 qk_exp_bf16 = pl.cast(qk_exp, target_type=pl.BF16, mode="rint")
-                qk_oi = pl.matmul(qk_exp_bf16, qk_kv, out_dtype=pl.FP32)
+                # [DIAGNOSTIC] force qk_exp_bf16 through GM before the PV cube, to break any
+                # composed-task fusion (exp->cast->cube) suspected of causing the PV drift.
+                _qk_exp_gm = pl.create_tensor([QK_M_TILE, ATTN_K_TILE], dtype=pl.BF16)
+                _qk_exp_gm[:, :] = qk_exp_bf16
+                qk_oi = pl.matmul(_qk_exp_gm, qk_kv, out_dtype=pl.FP32)
                 for qk_sub in pl.unroll(QK_M_TILE // H_TILE):
                     qk_h_idx = qk_hb * (QK_M_TILE // H_TILE) + qk_sub
                     qk_r0 = qk_sub * H_TILE
@@ -318,14 +322,24 @@
         n_full = n_normalized[0:H_TILE, 0:HEAD_DIM]
         n_bf16 = pl.cast(n_full, target_type=pl.BF16, mode="rint")
 
-        m_rope = n_full[:, NOPE_DIM:HEAD_DIM]
-        m_swapped = pl.gather(m_rope, dim=-1, index=rope_swap_idx[:, :])
-        m_cos_il = rope_cos_il[m_t : m_t + 1, 0:ROPE_DIM]
-        m_sin_signed = rope_sin_signed[m_t : m_t + 1, 0:ROPE_DIM]
-        m_rope_cos = pl.col_expand_mul(m_rope, m_cos_il)
-        m_swap_sin = pl.col_expand_mul(m_swapped, m_sin_signed)
-        m_rot = pl.add(m_rope_cos, m_swap_sin)
-        n_rope_bf16 = pl.cast(m_rot, target_type=pl.BF16, mode="rint")
+        # Per head-row inverse RoPE (1-row tiles). A5 gather lowering corrupts the
+        # [H_TILE, ROPE_DIM] gather for tiles >8 rows (only the first head of each
+        # 8-head group survived). Rotating one head-row at a time keeps every gather
+        # within a single FP32 vector box; result materialized to n_rope_bf16 so the
+        # existing pack loop is unchanged.
+        # Reuse the head-invariant swap index precomputed in the rope_cs task
+        # above: every row of rope_swap_idx is identical, so its first row is
+        # mathematically equal to the [1, ROPE_DIM] pattern rebuilt here. Slicing
+        # saves the redundant vector ops while keeping the gather tile 1-row.
+        mr_swap_idx = rope_swap_idx[0:1, :]  # [1, ROPE_DIM]
+        mr_cos_il = rope_cos_il[m_t : m_t + 1, 0 : ROPE_DIM]
+        mr_sin_signed = rope_sin_signed[m_t : m_t + 1, 0 : ROPE_DIM]
+        n_rope_bf16 = pl.create_tensor([H_TILE, ROPE_DIM], dtype=pl.BF16)
+        for n_hi in pl.range(H_TILE):
+            mr_row = n_full[n_hi : n_hi + 1, NOPE_DIM : HEAD_DIM]
+            mr_swapped = pl.gather(mr_row, dim=-1, index=mr_swap_idx)
+            mr_rot = pl.add(pl.mul(mr_row, mr_cos_il), pl.mul(mr_swapped, mr_sin_signed))
+            n_rope_bf16[n_hi : n_hi + 1, :] = pl.cast(mr_rot, target_type=pl.BF16, mode="rint")
 
         m_g0 = m_h0 // HEADS_PER_GROUP
         for m_sg in pl.unroll(H_TILE // HEADS_PER_GROUP):

models/deepseek/v4/hc_pre.py

@@ -804,16 +804,19 @@ def init_hc_base():
     hc_pre = _bind_hc_pre()
     print(f"hc_pre implementation: {HC_PRE_IMPL}")
 
-    # hc_pre is specialized to Ascend 910B. The "syncall" body sets NUM_CORES=24 == the
-    # physical AIC count and its hard full-occupancy mix-syncall hangs (AICore timeout
-    # 507018) unless the launch fills every physical core; A5 (Ascend950) has a different
-    # core count. The "separate" body has no such barrier but is still 910B-tuned (tile
-    # sizes + device-only atomic-add). Either way, reject A5 rather than hang / mis-run;
-    # supporting it needs a backend-aware participant count + re-tuning/re-validation.
-    if args.platform in ("a5", "a5sim"):
+    # hc_pre "syncall" is specialized to Ascend 910B: it sets NUM_CORES=24 == the physical
+    # AIC count and its hard full-occupancy mix-syncall hangs (AICore timeout 507018) unless
+    # the launch fills every physical core; A5 (Ascend950) has a different AIC count (36), so
+    # the syncall impl is rejected on A5. The "separate" impl has no such barrier -- it uses
+    # data-derived pl.spmd(work_count) + CORE_GROUP + atomic-add (the same structural pattern
+    # as hc_head/hc_post, which already pass on A5), so it is core-count-agnostic and runs on
+    # A5. Its tile sizes are 910B-tuned (perf, not correctness); a5 perf may be suboptimal
+    # until re-swept, but precision is unaffected.
+    if args.platform in ("a5", "a5sim") and HC_PRE_IMPL == "syncall":
         raise SystemExit(
-            f"hc_pre is specialized to Ascend 910B (NUM_CORES={NUM_CORES} == physical AIC count); "
-            f"the {HC_PRE_IMPL!r} impl would hang / mis-run on {args.platform!r}. Run with -p a2a3."
+            f"hc_pre 'syncall' impl is specialized to Ascend 910B (NUM_CORES={NUM_CORES} == physical "
+            f"AIC count); its full-occupancy mix-syncall would hang (AICore timeout 507018) on "
+            f"{args.platform!r}. Re-run with --impl separate on {args.platform!r}, or -p a2a3."
         )
 
     modes_to_run = list(MODES.keys()) if args.mode == "all" else [args.mode]

models/deepseek/v4/prefill_sparse_attn.py

@@ -215,7 +215,12 @@ def prefill_sparse_attn(
                     # li sums the FP32 exp; only the PV matmul uses the BF16 cast.
                     qk_li = pl.row_sum(qk_exp)
                     qk_exp_bf16 = pl.cast(qk_exp, target_type=pl.BF16, mode="rint")
-                    qk_oi = pl.matmul(qk_exp_bf16, qk_kv_v, out_dtype=pl.FP32)
+                    # Materialize qk_exp_bf16 to GM before the PV cube: the a5 codegen
+                    # mis-fuses the exp->cast->cube chain in this composed task, losing
+                    # ~15% precision. Forcing a GM round-trip breaks the fusion (fixes it).
+                    _qk_exp_gm = pl.create_tensor([QK_M_TILE, PREFILL_ATTN_TILE], dtype=pl.BF16)
+                    _qk_exp_gm[:, :] = qk_exp_bf16
+                    qk_oi = pl.matmul(_qk_exp_gm, qk_kv_v, out_dtype=pl.FP32)
                     for qk_sub in pl.unroll(QK_M_TILE // HEAD_TILE):
                         qk_h_idx = qk_hb * (QK_M_TILE // HEAD_TILE) + qk_sub
                         qk_r0 = qk_sub * HEAD_TILE