使用GLM-5.3大模型辅助开发支持NPU平台适配DS v32 DSA稀疏注意力算子前向反馈-ops-tilelang-AtomGit

Open

使用GLM-5.3大模型辅助开发支持NPU平台适配DS v32 DSA稀疏注意力算子前向反馈 #71

arick_369Add on 1 day agoMerged fromarick_369/ops-tilelang： master into cann/ops-tilelang： master

使用GLM-5.3大模型辅助开发支持NPU平台适配DS v32 DSA稀疏注意力算子前向反馈 #71

Open

arick_369Add on 1 day ago

Merged fromarick_369/ops-tilelang masterintocann/ops-tilelang master

Total: 6 files changed+932-0

Msrc/cann_ops_tilelang/__init__.py+2-0

	@@ -80,6 +80,7 @@ __all__ = [
80	'shuffle_image_tokens',	80	'shuffle_image_tokens',
81	'softmax',	81	'softmax',
82	'softmax_and_logsoftmax',	82	'softmax_and_logsoftmax',
		83	+ 'sparse_mla_fwd',
83	'store_index_kv_cache',	84	'store_index_kv_cache',
84	'store_kv_cache',	85	'store_kv_cache',
85	'store_kv_cache_llama',	86	'store_kv_cache_llama',
	@@ -99,6 +100,7 @@ from .attention import (
99	prepare_v33_decode_metadata,	100	prepare_v33_decode_metadata,
100	rotary_embedding,	101	rotary_embedding,
101	select_state_slice,	102	select_state_slice,
		103	+ sparse_mla_fwd,
102	store_index_kv_cache,	104	store_index_kv_cache,
103	store_kv_cache,	105	store_kv_cache,
104	store_kv_cache_llama,	106	store_kv_cache_llama,

Msrc/cann_ops_tilelang/attention/__init__.py+2-0

	@@ -13,6 +13,7 @@ __all__ = [
13	'prepare_v33_decode_metadata',	13	'prepare_v33_decode_metadata',
14	'rotary_embedding',	14	'rotary_embedding',
15	'select_state_slice',	15	'select_state_slice',
		16	+ 'sparse_mla_fwd',
16	'store_index_kv_cache',	17	'store_index_kv_cache',
17	'store_kv_cache',	18	'store_kv_cache',
18	'store_kv_cache_llama',	19	'store_kv_cache_llama',
	@@ -33,4 +34,5 @@ from .prepare_cp_mirror import prepare_for_mirror_cp_tilelang
33	from .prepare_llama_decode_metadata import prepare_llama_decode_metadata	34	from .prepare_llama_decode_metadata import prepare_llama_decode_metadata
34	from .prepare_v33_decode_metadata import prepare_v33_decode_metadata	35	from .prepare_v33_decode_metadata import prepare_v33_decode_metadata
35	from .rope import rotary_embedding	36	from .rope import rotary_embedding
		37	+from .sparse_mla_fwd import sparse_mla_fwd
36	from .state_cache import load_state_to_cache, select_state_slice	38	from .state_cache import load_state_to_cache, select_state_slice

Asrc/cann_ops_tilelang/attention/sparse_mla_fwd.py+124-0

@@ -0,0 +1,124 @@
	1	+"""Public interface for the DeepSeek sparse MLA forward operator."""
	2	+
	3	+from __future__ import annotations
	4	+
	5	+import torch
	6	+
	7	+from cann_ops_tilelang.attention.sparse_mla_fwd_kernel import D_QK, D_V, get_sparse_mla_fwd_kernel
	8	+from cann_ops_tilelang.config import require_ascend, should_print_kernel_source
	9	+
	10	+
	11	+def _validate_common(q: torch.Tensor, kv: torch.Tensor, indices: torch.Tensor) -> None:
	12	+ if not isinstance(q, torch.Tensor) or not isinstance(kv, torch.Tensor) or not isinstance(indices, torch.Tensor):
	13	+ raise TypeError('q, kv and indices must be torch.Tensor instances.')
	14	+ if q.device.type != 'npu' or kv.device.type != 'npu' or indices.device.type != 'npu':
	15	+ raise ValueError(f'q, kv and indices must be NPU tensors, got {q.device}, {kv.device} and {indices.device}.')
	16	+ if q.device != kv.device or q.device != indices.device:
	17	+ raise ValueError('q, kv and indices must be on the same NPU device.')
	18	+ if q.dtype != torch.bfloat16 or kv.dtype != torch.bfloat16:
	19	+ raise TypeError(f'q and kv must use torch.bfloat16, got {q.dtype} and {kv.dtype}.')
	20	+ if indices.dtype != torch.int32:
	21	+ raise TypeError(f'indices must use torch.int32, got {indices.dtype}.')
	22	+ if q.ndim != 4 or kv.ndim != 4 or indices.ndim != 4:
	23	+ raise ValueError(f'q, kv and indices must be four-dimensional, got {q.ndim}, {kv.ndim} and {indices.ndim}.')
	24	+ if q.shape[3] != D_QK or kv.shape[3] != D_QK:
	25	+ raise ValueError(f'q and kv must have a last dimension of {D_QK}, got {q.shape[3]} and {kv.shape[3]}.')
	26	+ if not q.is_contiguous() or not kv.is_contiguous() or not indices.is_contiguous():
	27	+ raise ValueError('q, kv and indices must be contiguous.')
	28	+
	29	+ batch, seq_len, heads = q.shape[0], q.shape[1], q.shape[2]
	30	+ kv_group = kv.shape[2]
	31	+ topk = indices.shape[3]
	32	+ if q.shape[0] != kv.shape[0] or q.shape[0] != indices.shape[0]:
	33	+ raise ValueError(f'batch dims must agree: q={q.shape[0]}, kv={kv.shape[0]}, indices={indices.shape[0]}.')
	34	+ if q.shape[1] != indices.shape[1]:
	35	+ raise ValueError(f'seq_len dims must agree: q={q.shape[1]}, indices={indices.shape[1]}.')
	36	+ if kv.shape[2] != indices.shape[2]:
	37	+ raise ValueError(f'kv_group dims must agree: kv={kv.shape[2]}, indices={indices.shape[2]}.')
	38	+ if indices.shape != (batch, seq_len, kv_group, topk):
	39	+ raise ValueError(
	40	+ f'indices must have shape ({batch}, {seq_len}, {kv_group}, {topk}), got {tuple(indices.shape)}.'
	41	+ )
	42	+ if heads % kv_group != 0:
	43	+ raise ValueError(f'heads ({heads}) must be divisible by kv_group ({kv_group}).')
	44	+ if topk % 64 != 0:
	45	+ raise ValueError(f'topk ({topk}) must be a multiple of 64.')
	46	+ head_kv = heads // kv_group
	47	+ if head_kv not in (16, 32, 64) and head_kv % 64 != 0:
	48	+ raise ValueError(
	49	+ f'unsupported heads/kv_group ratio {head_kv}: for head_kv <= 64 use one of 16/32/64; '
	50	+ 'larger values must be multiples of 64.'
	51	+ )
	52	+
	53	+
	54	+def _validate_out(out: torch.Tensor, lse: torch.Tensor, shape: tuple[int, ...], device: torch.device) -> None:
	55	+ for name, tensor, dtype, expected_shape in (
	56	+ ('out', out, torch.bfloat16, (*shape, D_V)),
	57	+ ('lse', lse, torch.float32, shape),
	58	+ ):
	59	+ if tensor.device.type != 'npu' or tensor.device != device:
	60	+ raise ValueError(f'{name} must be on the same NPU device as the inputs, got {tensor.device}.')
	61	+ if tensor.dtype != dtype:
	62	+ raise TypeError(f'{name} must use {dtype}, got {tensor.dtype}.')
	63	+ if not tensor.is_contiguous():
	64	+ raise ValueError(f'{name} must be contiguous.')
	65	+ if tensor.shape != expected_shape:
	66	+ raise ValueError(f'{name} must have shape {expected_shape}, got {tuple(tensor.shape)}.')
	67	+ if out.numel() and lse.numel() and out.untyped_storage().data_ptr() == lse.untyped_storage().data_ptr():
	68	+ raise ValueError('out and lse must not alias each other.')
	69	+
	70	+
	71	+def sparse_mla_fwd(
	72	+ q: torch.Tensor,
	73	+ kv: torch.Tensor,
	74	+ indices: torch.Tensor,
	75	+ sm_scale: float \| None = None,
	76	+ *,
	77	+ out: torch.Tensor \| None = None,
	78	+ lse: torch.Tensor \| None = None,
	79	+) -> tuple[torch.Tensor, torch.Tensor]:
	80	+ """DeepSeek v3.2 DSA sparse MLA forward on Ascend.
	81	+
	82	+ Args:
	83	+ q: [batch, seq_len, heads, 576] bfloat16, contiguous NPU tensor.
	84	+ kv: [batch, seq_len_kv, kv_group, 576] bfloat16, contiguous NPU tensor.
	85	+ indices: [batch, seq_len, kv_group, topk] int32 selected KV slots
	86	+ (topk must be a multiple of 64; out-of-range slots are clamped and
	87	+ causally masked, so padding with ``seq_len_kv`` is safe).
	88	+ sm_scale: softmax scale; defaults to ``576 ** -0.5``.
	89	+ out: optional [batch, seq_len, heads, 512] bfloat16 output buffer.
	90	+ lse: optional [batch, seq_len, heads] float32 log-sum-exp buffer.
	91	+
	92	+ Returns:
	93	+ Tuple ``(out, lse)`` with the attention output and the log-domain
	94	+ softmax statistics. Inputs are never modified.
	95	+ """
	96	+ _validate_common(q, kv, indices)
	97	+ require_ascend()
	98	+
	99	+ batch, seq_len, heads = q.shape[0], q.shape[1], q.shape[2]
	100	+ kv_group = kv.shape[2]
	101	+ topk = indices.shape[3]
	102	+ if out is None:
	103	+ out = torch.empty((batch, seq_len, heads, D_V), dtype=torch.bfloat16, device=q.device)
	104	+ if lse is None:
	105	+ lse = torch.empty((batch, seq_len, heads), dtype=torch.float32, device=q.device)
	106	+ _validate_out(out, lse, (batch, seq_len, heads), q.device)
	107	+ if out.numel() == 0:
	108	+ return out, lse
	109	+ if kv.shape[1] < 1:
	110	+ raise ValueError('seq_len_kv must be positive for non-empty queries.')
	111	+
	112	+ kernel = get_sparse_mla_fwd_kernel(
	113	+ batch,
	114	+ seq_len,
	115	+ kv.shape[1],
	116	+ heads,
	117	+ topk,
	118	+ kv_group=kv_group,
	119	+ sm_scale=sm_scale,
	120	+ )
	121	+ if should_print_kernel_source():
	122	+ print(kernel.get_kernel_source())
	123	+ kernel(q, kv, indices, out.view(-1, D_V), lse.view(-1))
	124	+ return out, lse

Asrc/cann_ops_tilelang/attention/sparse_mla_fwd_kernel.py+498-0

@@ -0,0 +1,498 @@
	1	+"""TileLang Ascend implementation of the DeepSeek sparse MLA forward kernel.
	2	+
	3	+Computes, per (batch, query position, kv group, head block):
	4	+
	5	+ S = Q[:, :576] @ gather(KV, Indices)^T (topk selected KV rows, causal-masked)
	6	+ P = online_softmax(S * sm_scale) (natural exp; no log2(e) factor)
	7	+ O = P @ gather(KV, Indices)[:, :512]
	8	+ Lse = m + ln(l) (log-domain softmax statistics)
	9	+
	10	+Ported from ``dsa_sparse_attention_ascend950/sparse_mla_fwd.py`` (itself a port of
	11	+``tilelang/examples/deepseek_v32/sparse_mla_fwd.py``). The AIC/AIV + UB/L1/L0C
	12	+structure follows ``examples/ascend/flash_attention/core.py`` (GQA): persistent
	13	+1-D core grid, explicit L0A/L0B operand GEMMs into L0C, ``T.dual_copy``
	14	+L0C->UB M-split across the two AIVs, SIMD-intrinsic online softmax.
	15	+
	16	+Key Ascend-specific constraints preserved from the DT-verified port:
	17	+
	18	+ * ALL GEMMs run on EXPLICIT L0A/L0B operands fed by plain L1->L0 copies. The
	19	+ L1-input GEMM path stages operands at L0 address 0 with zero scheduler
	20	+ visibility, so explicit L0 buffers collide with it once a Pipelined loop
	21	+ overlaps iterations (silent corruption for NI >= 3; DT-verified 2026-10-26).
	22	+ * The K axis (576) is tiled into four whole ``[., 128]`` L1 chunk buffers plus
	23	+ a ``[., 64]`` tail buffer: the L0A/L0B budget is 64 KB and column-sliced
	24	+ L1->L0 copies load garbage in this build (NaN; DT-verified 2026-10-26). The
	25	+ kv chunks are gathered ONCE per iteration and shared by the S phase (plain
	26	+ copies) and the PV phase (transposing copies).
	27	+ * The sparse gather is one whole-chunk ``T.copy`` per destination row with the
	28	+ GM source row as a runtime scalar (clamped for OOB safety). Fusing multiple
	29	+ same-row copies into one loop body corrupts the auto-scheduled MTE2 stream
	30	+ (DT-verified 2026-10-26), and rank-1 index staging is rejected by layout
	31	+ inference (rank >= 2 requirement).
	32	+ * The softmax P tile is written DIRECTLY in compact-NZ layout from UB with a
	33	+ ``vdintlv(x, x)`` butterfly after ``vcvt(part=0)`` — without the densify
	34	+ step every other P column is zero (DT-verified 2026-10-26). This avoids the
	35	+ ``nd2nz_copy.h`` ``ROWS % 16 == 0`` constraint for HPB=16.
	36	+ * Persistent grids where every core gets exactly ONE tile trip an
	37	+ auto-schedule defect that corrupts PV output while leaving Lse exact
	38	+ (DT-verified 2026-10-26); :func:`pick_num_blocks` caps the grid at
	39	+ ``num_tiles // 2``.
	40	+ * ``num_stages=1`` ships by default: stages 2/3 are correct but measured no
	41	+ faster (the small-burst gather dominates, not overlap).
	42	+ * exp2/log2 (CUDA) become natural exp/ln SIMD intrinsics; ``sm_scale``
	43	+ therefore does NOT carry the ``log2(e)`` factor.
	44	+ * The PTO backend is rejected at build time: its L0C->UB MTE lowering
	45	+ requires ``n_size <= ceil(M/16)*16``, which this kernel's [HPB, 64/128]
	46	+ readout tiles violate for every supported head-block shape.
	47	+"""
	48	+
	49	+import tilelang
	50	+from tilelang.ascend import language as T
	51	+from tilelang.ascend.language import simd as S
	52	+from tilelang.layout import make_ascend_compact_nz_layout
	53	+
	54	+from cann_ops_tilelang.attention._simd import is_pto_target
	55	+from cann_ops_tilelang.config import get_num_ai_cores
	56	+
	57	+BLOCK_I = 64 # KV rows per sparse block; must divide topk (CUDA ref constraint)
	58	+D_V = 512 # output / value head dim
	59	+D_TAIL = 64 # absorbing-brownian tail of the qk head dim
	60	+D_QK = D_V + D_TAIL # 576
	61	+LANES = 64 # fp32 lanes of one 2048-bit SIMD vector register
	62	+
	63	+# Scaled-domain online-softmax constants (S is pre-multiplied by sm_scale): m is
	64	+# initialized to M_INIT, which doubles as the running-max floor, so masked
	65	+# entries at MASK_S = M_INIT - 1e4 give exp(MASK_S - m) <= exp(-1e4) -> exact 0
	66	+# in fp32 — no sumexp pollution on all-masked blocks. Real scaled scores are
	67	+# \|S*scale\| << 1e4, far above both constants.
	68	+M_INIT = -1.0e6
	69	+MASK_S = -1.01e6
	70	+
	71	+
	72	+def _next_power_of_2(x: int) -> int:
	73	+ return 1 << (x - 1).bit_length()
	74	+
	75	+
	76	+def _check_backend_support() -> None:
	77	+ """Reject the PTO backend at build time with an actionable error.
	78	+
	79	+ The PTO lowering requires ``n_size <= ceil(M/16)*16`` on every L0C->UB MTE
	80	+ copy. This kernel's S accumulator is [HPB, 64] and its O accumulators are
	81	+ [HPB, 128] with HPB in {16, 32, 64}: S violates the constraint for HPB
	82	+ 16/32 and O violates it for every supported HPB. The AscendC backend lowers
	83	+ the same descriptor without the check and is the verified path.
	84	+ """
	85	+ if is_pto_target():
	86	+ raise RuntimeError(
	87	+ 'sparse_mla_fwd requires the AscendC backend (set TILELANG_DEFAULT_TARGET=ascend, '
	88	+ 'or leave it unset for auto selection on NPU hosts). The PTO backend rejects this '
	89	+ "kernel's L0C->UB tiles: its MTE lowering requires n_size <= ceil(M/16)*16, but "
	90	+ 'the S tile is [HPB, 64] and the O tile is [HPB, 128] with HPB in {16, 32, 64}.'
	91	+ )
	92	+
	93	+
	94	+def pick_num_blocks(num_tiles: int, preferred: int) -> int:
	95	+ """Persistent grid extent: ``min(preferred, num_tiles // 2)``.
	96	+
	97	+ Non-divisor grids are supported via a per-core tail: the first
	98	+ ``num_tiles % num_blocks`` cores process one extra tile behind a runtime
	99	+ ``T.If`` guard. ``preferred`` is capped at ``num_tiles // 2`` because a grid
	100	+ where every core gets exactly ONE tile trips an auto-schedule defect that
	101	+ corrupts the PV output while leaving Lse exact (DT-verified 2026-10-26).
	102	+ """
	103	+ cap = num_tiles // 2 if num_tiles >= 2 else num_tiles
	104	+ return max(1, min(preferred, num_tiles, cap))
	105	+
	106	+
	107	+@tilelang.jit(pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True})
	108	+def get_sparse_mla_fwd_kernel(
	109	+ batch: int,
	110	+ seq_len: int,
	111	+ seq_len_kv: int,
	112	+ heads: int,
	113	+ topk: int,
	114	+ kv_group: int = 1,
	115	+ dim: int = D_V,
	116	+ tail_dim: int = D_TAIL,
	117	+ sm_scale: float \| None = None,
	118	+ num_blocks: int \| None = None,
	119	+ num_stages: int = 1,
	120	+):
	121	+ """Build (and lazily compile/cache) the sparse MLA forward Ascend kernel.
	122	+
	123	+ Args match the CUDA reference ``sparse_mla_fwd``; ``num_blocks`` defaults to
	124	+ the device cube-core count via :func:`pick_num_blocks`. The kernel takes
	125	+ (Q, KV, Indices, Output, Lse) with flat GM outputs — Output is
	126	+ ``[batchseq_lenheads, dim]`` and Lse is ``[batchseq_lenheads]``; the
	127	+ public interface restores the logical views with zero-copy reshapes.
	128	+ """
	129	+ _check_backend_support()
	130	+ assert dim == _next_power_of_2(dim), f'dim must be a power of 2, got {dim}'
	131	+ assert tail_dim == _next_power_of_2(tail_dim), f'tail_dim must be a power of 2, got {tail_dim}'
	132	+ assert heads % kv_group == 0, 'heads must be divisible by kv_group'
	133	+ assert topk % BLOCK_I == 0, f'topk ({topk}) must be a multiple of BLOCK_I ({BLOCK_I})'
	134	+ assert dim == 512, (
	135	+ f'this build hardcodes {512 // 128} K-chunk buffers (dim must be 512, got {dim}); '
	136	+ 'adjust the version-annotation list when changing the chunk count'
	137	+ )
	138	+ if sm_scale is None:
	139	+ # Natural exp on NPU: unlike the CUDA ref, do NOT fold in log2(e).
	140	+ sm_scale = (1.0 / (dim + tail_dim)) ** 0.5
	141	+
	142	+ BI = BLOCK_I
	143	+ D, D_tail, D_qk = dim, tail_dim, dim + tail_dim
	144	+ dtype, accum_dtype = 'bfloat16', 'float32'
	145	+ scale = float(sm_scale)
	146	+
	147	+ head_kv = heads // kv_group
	148	+ padded_H = max(_next_power_of_2(head_kv), 16)
	149	+ # The CUDA reference tolerates padded_H != head_kv via GM accesses past the
	150	+ # heads axis (L2-forgiven on GPUs). NPU DMA has no such tolerance, so this
	151	+ # kernel requires exact head blocks instead: head_kv <= 64 must be one of
	152	+ # 16/32/64; larger values must be multiples of 64 (then 64-row head blocks
	153	+ # tile head_kv exactly, no padding).
	154	+ assert padded_H == head_kv or head_kv % 64 == 0, (
	155	+ f'unsupported head_kv={head_kv} (heads={heads}, kv_group={kv_group}): '
	156	+ 'for head_kv<=64 use one of 16/32/64; larger values must be multiples of 64. '
	157	+ 'The H-padding path of the CUDA reference accesses GM out of bounds and '
	158	+ 'is intentionally not ported.'
	159	+ )
	160	+ replicate_h = head_kv // 64 if head_kv > 64 else 1
	161	+ H_per_block = padded_H if replicate_h == 1 else 64
	162	+
	163	+ ROWS = H_per_block // 2 # per-AIV rows after the L0C->UB M-split
	164	+ NI = topk // BI
	165	+ num_tiles = batch * seq_len * kv_group * replicate_h
	166	+ if num_blocks is None:
	167	+ num_blocks = pick_num_blocks(num_tiles, get_num_ai_cores())
	168	+ # Tail grid: non-divisor num_blocks is allowed; the first ``num_tiles %
	169	+ # num_blocks`` cores process one extra tile behind a runtime T.If guard.
	170	+ # Core bx handles a contiguous, non-overlapping partition of [0, num_tiles).
	171	+ tpc = num_tiles // num_blocks
	172	+ rem = num_tiles % num_blocks
	173	+ assert num_blocks <= num_tiles, f'num_blocks ({num_blocks}) exceeds num_tiles ({num_tiles})'
	174	+ assert num_tiles == 1 or tpc >= 2, (
	175	+ f'tiles_per_core ({tpc}) < 2 trips an auto-schedule defect (PV output corruption, '
	176	+ 'Lse unaffected; DT-verified 2026-10-26). Use fewer blocks (pick_num_blocks '
	177	+ 'already caps at num_tiles // 2).'
	178	+ )
	179	+ serial_len = tpc + (1 if rem else 0)
	180	+
	181	+ RH, HPB = replicate_h, H_per_block
	182	+ SKV = seq_len_kv
	183	+ # NOTE: do NOT bind a closure variable named ``S`` here — it would shadow
	184	+ # the module-level ``simd as S`` import inside the traced prim_func and fail
	185	+ # at trace time. ``SEQ`` is the sequence-length symbol.
	186	+ B, SEQ, G, H = batch, seq_len, kv_group, heads
	187	+
	188	+ @T.prim_func
	189	+ def sparse_mla_fwd_kernel(
	190	+ Q: T.Tensor((B, SEQ, H, D_qk), dtype),
	191	+ KV: T.Tensor((B, SKV, G, D_qk), dtype),
	192	+ Indices: T.Tensor((B, SEQ, G, topk), 'int32'),
	193	+ # Output/Lse are declared flat: T.dual_copy derives split direction from
	194	+ # allocation shapes when a sliced destination materializes with its
	195	+ # allocation rank, so GM destinations must keep the rank the verified
	196	+ # GQA example uses (2-D O, 1-D LSE). The host interface restores the
	197	+ # logical [B, S, H, ...] views with zero-copy reshape.
	198	+ Output: T.Tensor((B * SEQ * H, D), dtype),
	199	+ Lse: T.Tensor((B * SEQ * H,), accum_dtype),
	200	+ ):
	201	+ with T.Kernel(num_blocks) as bx:
	202	+ # ---- on-chip storage ------------------------------------------
	203	+ # L1 (CBuf, AIC-side, common to the core's AIVs). Q and KV are
	204	+ # staged as SEPARATE whole buffers per K-segment (512 + 64), never
	205	+ # column slices of one buffer: the L1-input GEMM path drops K-axis
	206	+ # offsets of sliced operands, and column-sliced L1->L0 copies load
	207	+ # garbage in this build. The kv chunks are gathered ONCE per
	208	+ # iteration and shared by the S (plain copy) and PV (transposing
	209	+ # copy) phases — copy-only consumers share a single MajorK anchor.
	210	+ KC = 128
	211	+ NKC = D // KC
	212	+ q_l1 = [T.alloc_l1((HPB, KC), dtype) for _ in range(NKC)]
	213	+ qt_l1 = T.alloc_l1((HPB, D_tail), dtype)
	214	+ kv_l1 = [T.alloc_l1((BI, KC), dtype) for _ in range(NKC)]
	215	+ kvt_l1 = T.alloc_l1((BI, D_tail), dtype)
	216	+ # P tile: NZ layout staged in L1 for the PV GEMM A-operand.
	217	+ P_l1 = T.alloc_l1((HPB, BI), dtype)
	218	+ # S-phase L0 operands: ONE [HPB, KC]/[BI, KC] pair serially reused
	219	+ # by the NKC chunk GEMMs (plus the 64-wide tail pair).
	220	+ s_a_l0 = T.alloc_l0a((HPB, KC), dtype)
	221	+ s_b_l0 = T.alloc_l0b((BI, KC), dtype)
	222	+ st_a_l0 = T.alloc_l0a((HPB, D_tail), dtype)
	223	+ st_b_l0 = T.alloc_l0b((BI, D_tail), dtype)
	224	+ # PV-phase L0 operands: A = P [HPB, BI]; B = V-chunk^T [KC, BI],
	225	+ # serially reused by the NKC N-split GEMMs (transposing copies).
	226	+ pv_a_l0 = T.alloc_l0a((HPB, BI), dtype)
	227	+ pv_b_l0 = T.alloc_l0b((KC, BI), dtype)
	228	+ # L0C accumulators: S [HPB, BI]; O N-split into NKC pieces
	229	+ # (zero-origin L0 regions forbid column-sliced C operands).
	230	+ S_l0c = T.alloc_l0c((HPB, BI), accum_dtype)
	231	+ o_l0c = [T.alloc_l0c((HPB, KC), accum_dtype) for _ in range(NKC)]
	232	+ # UB (per-AIV); ROWS rows each after the M-split. The row-state
	233	+ # vectors are padded to 64 lanes so whole-vector SIMD idioms
	234	+ # (vld/vsts at [0]) stay in-bounds for any ROWS <= 64.
	235	+ S_ub = T.alloc_shared((ROWS, BI), accum_dtype)
	236	+ # P tile written DIRECTLY in NZ layout from UB (GQA core.py softmax
	237	+ # path): the softmax VF stores 64-bf16 column blocks with vsstb into
	238	+ # compact-NZ order, so the UB->L1 handoff is a plain dual_copy with
	239	+ # NO ND->NZ scatter. This sidesteps the nd2nz_copy.h
	240	+ # static_assert(ROWS % 16 == 0) which HPB=16 (ROWS=8) cannot
	241	+ # satisfy. Shape is (ROWS + 1) rows so the NZ column-stride
	242	+ # (ROWS+1) blocks has a full 32B-aligned pad row.
	243	+ P_nz_ub = T.alloc_shared((ROWS + 1, BI), dtype)
	244	+ T.annotate_layout({P_nz_ub: make_ascend_compact_nz_layout(P_nz_ub)})
	245	+ O_ub = T.alloc_shared((ROWS, D), accum_dtype)
	246	+ O_tmp_ub = T.alloc_shared((ROWS, D), accum_dtype)
	247	+ # O_cast_ub is declared in the epilogue as a byte-alias T.view of O_ub
	248	+ idx_ub = T.alloc_shared((BI,), 'int32')
	249	+ # lse: computed into the padded vector, then masked-stored into the
	250	+ # exact-size buffer so the final dual_copy moves the WHOLE buffer
	251	+ # (GQA symmetry; sub-region 1-D dual copies are untested).
	252	+ lse_ub = T.alloc_shared((LANES,), accum_dtype)
	253	+ lse_out_ub = T.alloc_shared((ROWS,), accum_dtype)
	254	+ m_ub = T.alloc_shared((LANES,), accum_dtype)
	255	+ m_prev_ub = T.alloc_shared((LANES,), accum_dtype)
	256	+ alpha_ub = T.alloc_shared((LANES,), accum_dtype)
	257	+ l_ub = T.alloc_shared((LANES,), accum_dtype)
	258	+ rs_ub = T.alloc_shared((LANES,), accum_dtype)
	259	+ T.annotate_buffer_versions(
	260	+ {
	261	+ # kv chunks: double-buffered across the NI loop. Their
	262	+ # L1->L0 copy consumers are not RAW-tracked across Pipelined
	263	+ # iterations, so single buffering lets iteration i+1's
	264	+ # gather clobber iteration i's tile before the PV
	265	+ # transposing copy reads it (DT-verified 2026-10-26).
	266	+ kv_l1[0]: 2,
	267	+ kv_l1[1]: 2,
	268	+ kv_l1[2]: 2,
	269	+ kv_l1[3]: 2,
	270	+ kvt_l1: 2,
	271	+ S_ub: 2,
	272	+ P_nz_ub: 2,
	273	+ O_tmp_ub: 2,
	274	+ }
	275	+ )
	276	+
	277	+ for local_t in T.serial(serial_len):
	278	+ # Tail grid: cores [0, rem) process tpc+1 tiles, the rest tpc;
	279	+ # the trailing over-iteration of the non-tail cores is disabled
	280	+ # by the runtime guard (tile ownership stays a partition).
	281	+ t = bx * tpc + T.if_then_else(bx < rem, bx, rem) + local_t
	282	+ with T.If(local_t < T.if_then_else(bx < rem, tpc + 1, tpc)): # noqa: SIM117
	283	+ with T.Then():
	284	+ # tile decomposition: fastest axis = head replica
	285	+ # (consecutive tiles share (b, s, g) -> Q tiles hit L1 hot)
	286	+ h_blk = t % RH
	287	+ tg = t // RH
	288	+ g_i = tg % G
	289	+ tb = tg // G
	290	+ s_i = tb % SEQ
	291	+ b_i = tb // SEQ
	292	+ # g-stride is head_kv (exact blocks, see builder asserts)
	293	+ H0 = g_i * head_kv + h_blk * HPB
	294	+
	295	+ # Q chunks are tile-invariant: load once per tile.
	296	+ T.copy(Q[b_i, s_i, H0 : H0 + HPB, 0:KC], q_l1[0][0:HPB, 0:KC])
	297	+ T.copy(Q[b_i, s_i, H0 : H0 + HPB, KC : 2 * KC], q_l1[1][0:HPB, 0:KC])
	298	+ T.copy(Q[b_i, s_i, H0 : H0 + HPB, 2 * KC : 3 * KC], q_l1[2][0:HPB, 0:KC])
	299	+ T.copy(Q[b_i, s_i, H0 : H0 + HPB, 3 * KC : D], q_l1[3][0:HPB, 0:KC])
	300	+ T.copy(Q[b_i, s_i, H0 : H0 + HPB, D:D_qk], qt_l1[0:HPB, 0:D_tail])
	301	+
	302	+ # ---- init softmax state + O ----------------------
	303	+ with T.SimdVF():
	304	+ init_mask = S.pset(32, 'PAT_ALL')
	305	+ zero = S.vdup(T.float32(0.0), accum_dtype, init_mask)
	306	+ init_m = S.vdup(T.float32(M_INIT), accum_dtype, init_mask)
	307	+ S.vsts(l_ub[0], zero, init_mask, 'NORM_B32')
	308	+ S.vsts(m_ub[0], init_m, init_mask, 'NORM_B32')
	309	+ S.vsts(rs_ub[0], zero, init_mask, 'NORM_B32')
	310	+ for r in range(ROWS):
	311	+ for c in range(0, D, LANES):
	312	+ S.vsts(O_ub[r, c], zero, init_mask, 'NORM_B32')
	313	+
	314	+ for i_i in T.Pipelined(NI, num_stages=num_stages, annotations={'enable_offset': True}):
	315	+ # ---- sparse gather (AIC side) -----------------
	316	+ # NOTE: no rank-1 L1 staging for the indices — a
	317	+ # standalone GM->L1 copy whose destination has no
	318	+ # GEMM consumer gets a fractal MajorK layout forced
	319	+ # by AscendLayoutInference, which requires rank >= 2
	320	+ # (ascend_layouts.cc ICHECK). Read the index scalars
	321	+ # straight from GM instead, and stage the block to
	322	+ # UB only for the AIV-side causal mask.
	323	+ T.copy(Indices[b_i, s_i, g_i, i_i * BI : i_i * BI + BI], idx_ub[0:BI])
	324	+ # ---- sparse gather: one whole-chunk copy per row
	325	+ # (index scalars read straight from GM; clamped for
	326	+ # OOB safety)
	327	+ for bi in T.serial(BI):
	328	+ row = T.min(T.max(Indices[b_i, s_i, g_i, i_i * BI + bi], 0), SKV - 1)
	329	+ T.copy(KV[b_i, row, g_i, 0 * KC : 1 * KC], kv_l1[0][bi, 0:KC])
	330	+ for bi in T.serial(BI):
	331	+ row = T.min(T.max(Indices[b_i, s_i, g_i, i_i * BI + bi], 0), SKV - 1)
	332	+ T.copy(KV[b_i, row, g_i, 1 * KC : 2 * KC], kv_l1[1][bi, 0:KC])
	333	+ for bi in T.serial(BI):
	334	+ row = T.min(T.max(Indices[b_i, s_i, g_i, i_i * BI + bi], 0), SKV - 1)
	335	+ T.copy(KV[b_i, row, g_i, 2 * KC : 3 * KC], kv_l1[2][bi, 0:KC])
	336	+ for bi in T.serial(BI):
	337	+ row = T.min(T.max(Indices[b_i, s_i, g_i, i_i * BI + bi], 0), SKV - 1)
	338	+ T.copy(KV[b_i, row, g_i, 3 * KC : 4 * KC], kv_l1[3][bi, 0:KC])
	339	+ for bi in T.serial(BI):
	340	+ row = T.min(T.max(Indices[b_i, s_i, g_i, i_i * BI + bi], 0), SKV - 1)
	341	+ T.copy(KV[b_i, row, g_i, D:D_qk], kvt_l1[bi, 0:D_tail])
	342	+
	343	+ # ---- S = Q @ K^T (4 x K=128 chunks + K=64 tail)
	344	+ # Explicit L0A/L0B operands, serially reused pairs;
	345	+ # each chunk GEMM accumulates into S_l0c.
	346	+ T.copy(q_l1[0][0:HPB, 0:KC], s_a_l0[0:HPB, 0:KC])
	347	+ T.copy(kv_l1[0][0:BI, 0:KC], s_b_l0[0:BI, 0:KC])
	348	+ T.gemm(s_a_l0, s_b_l0, S_l0c, transpose_B=True, clear_accum=True)
	349	+ T.copy(q_l1[1][0:HPB, 0:KC], s_a_l0[0:HPB, 0:KC])
	350	+ T.copy(kv_l1[1][0:BI, 0:KC], s_b_l0[0:BI, 0:KC])
	351	+ T.gemm(s_a_l0, s_b_l0, S_l0c, transpose_B=True, clear_accum=False)
	352	+ T.copy(q_l1[2][0:HPB, 0:KC], s_a_l0[0:HPB, 0:KC])
	353	+ T.copy(kv_l1[2][0:BI, 0:KC], s_b_l0[0:BI, 0:KC])
	354	+ T.gemm(s_a_l0, s_b_l0, S_l0c, transpose_B=True, clear_accum=False)
	355	+ T.copy(q_l1[3][0:HPB, 0:KC], s_a_l0[0:HPB, 0:KC])
	356	+ T.copy(kv_l1[3][0:BI, 0:KC], s_b_l0[0:BI, 0:KC])
	357	+ T.gemm(s_a_l0, s_b_l0, S_l0c, transpose_B=True, clear_accum=False)
	358	+ T.copy(qt_l1[0:HPB, 0:D_tail], st_a_l0[0:HPB, 0:D_tail])
	359	+ T.copy(kvt_l1[0:BI, 0:D_tail], st_b_l0[0:BI, 0:D_tail])
	360	+ T.gemm(st_a_l0, st_b_l0, S_l0c, transpose_B=True, clear_accum=False)
	361	+
	362	+ # ---- online softmax (vector side, per-AIV halves)
	363	+ # GQA core.py softmax structure, adapted to
	364	+ # [ROWS, 64] tiles: one fp32 vector register per row.
	365	+ T.dual_copy(S_l0c, S_ub)
	366	+ with T.SimdVF():
	367	+ full = S.pset(32, 'PAT_ALL')
	368	+ one_pt = S.pset(32, 'PAT_VL1')
	369	+ # causal validity: raw index <= query position
	370	+ # (int32 vector-vs-scalar compare, NPU-tested)
	371	+ idxv = S.vld(idx_ub[0], 'NORM')
	372	+ valid = S.vcmps(idxv, s_i, full, 'le')
	373	+ one = S.vdup(T.float32(1.0), accum_dtype, full)
	374	+ zero = S.vdup(T.float32(0.0), accum_dtype, full)
	375	+ w = S.vsel(one, zero, valid)
	376	+ ws = S.vmuls(w, T.float32(scale), full)
	377	+ mask_term = S.vmuls(S.vsub(one, w, full), T.float32(MASK_S), full)
	378	+ # save the previous running max (whole-row vector)
	379	+ S.vsts(m_prev_ub[0], S.vld(m_ub[0], 'NORM'), full, 'NORM_B32')
	380	+ # mask + prescale scores in place, then per-row max
	381	+ for r in range(ROWS):
	382	+ sv = S.vld(S_ub[r, 0], 'NORM')
	383	+ sv_scaled = S.vadd(S.vmul(sv, ws, full), mask_term, full)
	384	+ S.vsts(S_ub[r, 0], sv_scaled, full, 'NORM_B32')
	385	+ S.vsts(m_ub[r], S.vcmax(sv_scaled, full), one_pt, 'ONEPT_B32')
	386	+ S.mem_bar('VST_VLD')
	387	+ merged = S.vmax(S.vld(m_ub[0], 'NORM'), S.vld(m_prev_ub[0], 'NORM'), full)
	388	+ S.vsts(m_ub[0], merged, full, 'NORM_B32')
	389	+ S.mem_bar('VST_VLD')
	390	+ # p = exp(S_scaled - m), row sums, bf16 P tile
	391	+ # written DIRECTLY in compact-NZ order. Each row
	392	+ # is one fp32 VL (64 lanes) -> 64 bf16 = 4 32B
	393	+ # column blocks; vsstb scatters them to the NZ
	394	+ # column slots with blockStride=(ROWS+1) blocks
	395	+ # and POST_UPDATE advancing 1 block per row.
	396	+ p16 = S.pset(16, 'PAT_VL64')
	397	+ p_ptr = S.make_ubuf_ptr(
	398	+ T.access_ptr(P_nz_ub[0, 0], 'w', ROWS + 1, BI),
	399	+ 'bfloat16',
	400	+ )
	401	+ for r in range(ROWS):
	402	+ m_brc = S.vld(m_ub[r], 'BRC_B32')
	403	+ e = S.vexpdif(S.vld(S_ub[r, 0], 'NORM'), m_brc, full)
	404	+ S.vsts(rs_ub[r], S.vcadd(e, full), one_pt, 'ONEPT_B32')
	405	+ # Densify the bf16 conversion: vcvt(part=0)
	406	+ # places element i at EVEN 16-bit lanes only
	407	+ # (DT-verified 2026-10-26). With one 64-lane
	408	+ # row the fix is a vdintlv(x, x) butterfly
	409	+ # that compresses the even lanes into a
	410	+ # dense low half.
	411	+ c0 = S.vcvt(e, dtype, full, part=0)
	412	+ d0, d1 = S.vdintlv(c0, c0)
	413	+ p_ptr = S.vsstb(
	414	+ d0,
	415	+ p_ptr,
	416	+ T.int32(((ROWS + 1) << 16) \| 1),
	417	+ p16,
	418	+ update=True,
	419	+ )
	420	+ S.mem_bar('VST_VLD')
	421	+ # alpha + l (whole-row vectors; padding lanes stay 0)
	422	+ alpha_reg = S.vexpdif(S.vld(m_prev_ub[0], 'NORM'), S.vld(m_ub[0], 'NORM'), full)
	423	+ S.vsts(alpha_ub[0], alpha_reg, full, 'NORM_B32')
	424	+ l_updated = S.vadd(
	425	+ S.vmul(S.vld(l_ub[0], 'NORM'), alpha_reg, full),
	426	+ S.vld(rs_ub[0], 'NORM'),
	427	+ full,
	428	+ )
	429	+ S.vsts(l_ub[0], l_updated, full, 'NORM_B32')
	430	+
	431	+ # ---- O += P @ V --------------------------------
	432	+ # P is already in NZ order in UB (vsstb wrote
	433	+ # compact-NZ directly), so the UB->L1 handoff is a
	434	+ # plain dual_copy with NO ND->NZ scatter. Each AIV
	435	+ # writes its ROWS rows; the dual_copy M-split merges
	436	+ # the two AIV halves into HPB rows.
	437	+ T.dual_copy(P_nz_ub[0:ROWS, 0:BI], P_l1[0:HPB, 0:BI])
	438	+ # NT PV GEMM, N-split into NKC pieces: B = V-chunk^T
	439	+ # [KC, BI] fed by a transposing L1->L0B copy from
	440	+ # the SAME kv chunk buffers the S phase gathered
	441	+ # (copy-only consumers, one shared MajorK anchor).
	442	+ T.copy(P_l1[0:HPB, 0:BI], pv_a_l0)
	443	+ T.copy(kv_l1[0][0:BI, 0:KC], pv_b_l0[0:KC, 0:BI], transpose=True)
	444	+ T.gemm(pv_a_l0, pv_b_l0, o_l0c[0], transpose_B=True, clear_accum=True)
	445	+ T.dual_copy(o_l0c[0], O_tmp_ub[0:ROWS, 0 * KC : 1 * KC])
	446	+ T.copy(kv_l1[1][0:BI, 0:KC], pv_b_l0[0:KC, 0:BI], transpose=True)
	447	+ T.gemm(pv_a_l0, pv_b_l0, o_l0c[1], transpose_B=True, clear_accum=True)
	448	+ T.dual_copy(o_l0c[1], O_tmp_ub[0:ROWS, 1 * KC : 2 * KC])
	449	+ T.copy(kv_l1[2][0:BI, 0:KC], pv_b_l0[0:KC, 0:BI], transpose=True)
	450	+ T.gemm(pv_a_l0, pv_b_l0, o_l0c[2], transpose_B=True, clear_accum=True)
	451	+ T.dual_copy(o_l0c[2], O_tmp_ub[0:ROWS, 2 * KC : 3 * KC])
	452	+ T.copy(kv_l1[3][0:BI, 0:KC], pv_b_l0[0:KC, 0:BI], transpose=True)
	453	+ T.gemm(pv_a_l0, pv_b_l0, o_l0c[3], transpose_B=True, clear_accum=True)
	454	+ T.dual_copy(o_l0c[3], O_tmp_ub[0:ROWS, 3 * KC : 4 * KC])
	455	+ # GQA accumulate_output idiom: O = O*alpha + P@V
	456	+ with T.SimdVF():
	457	+ acc_mask = S.pset(32, 'PAT_ALL')
	458	+ for r in range(ROWS):
	459	+ alpha = S.vld(alpha_ub[r], 'BRC_B32')
	460	+ for c in range(0, D, LANES):
	461	+ prev = S.vld(O_ub[r, c], 'NORM')
	462	+ cur = S.vld(O_tmp_ub[r, c], 'NORM')
	463	+ updated = S.vadd(S.vmul(alpha, prev, acc_mask), cur, acc_mask)
	464	+ S.vsts(O_ub[r, c], updated, acc_mask, 'NORM_B32')
	465	+
	466	+ # ---- epilogue (GQA finalize / store_lse idioms) ----
	467	+ with T.SimdVF():
	468	+ norm_mask = S.pset(32, 'PAT_ALL')
	469	+ for r in range(ROWS):
	470	+ denom = S.vld(l_ub[r], 'BRC_B32')
	471	+ for c in range(0, D, LANES):
	472	+ normalized = S.vdiv(S.vld(O_ub[r, c], 'NORM'), denom, norm_mask)
	473	+ S.vsts(O_ub[r, c], normalized, norm_mask, 'NORM_B32')
	474	+ lse = S.vadd(
	475	+ S.vln(S.vld(l_ub[0], 'NORM'), norm_mask),
	476	+ S.vld(m_ub[0], 'NORM'),
	477	+ norm_mask,
	478	+ )
	479	+ S.vsts(lse_ub[0], lse, norm_mask, 'NORM_B32')
	480	+ with T.SimdVF():
	481	+ # exact-size lse staging: masked store of the first
	482	+ # ROWS lanes, then a whole-buffer dual_copy
	483	+ row_mask = S.pset(32, f'PAT_VL{ROWS}')
	484	+ S.vsts(lse_out_ub[0], S.vld(lse_ub[0], 'NORM'), row_mask, 'NORM_B32', extent=ROWS)
	485	+ # O_cast_ub aliases O_ub bytes (GQA finalize_cast idiom):
	486	+ # bf16 view of the same UB capacity, no extra allocation.
	487	+ O_cast_ub = T.view(O_ub, (ROWS * 2, D), dtype=dtype)
	488	+ with T.SimdVF():
	489	+ cast_full = S.pset(32, 'PAT_ALL')
	490	+ for r in range(ROWS):
	491	+ for c in range(0, D, LANES):
	492	+ ob = S.vcvt(S.vld(O_ub[r, c], 'NORM'), dtype, cast_full)
	493	+ S.vsts(O_cast_ub[r, c], ob, cast_full, 'PK_B32')
	494	+ row0 = (b_i * SEQ + s_i) * H + H0
	495	+ T.dual_copy(O_cast_ub[0:ROWS, 0:D], Output[row0 : row0 + HPB, 0:D])
	496	+ T.dual_copy(lse_out_ub, Lse[row0 : row0 + HPB])
	497	+
	498	+ return sparse_mla_fwd_kernel

Atests/attention/sparse_mla_fwd_ref.py+106-0

@@ -0,0 +1,106 @@
	1	+"""Pure PyTorch references for the sparse MLA forward operator.
	2	+
	3	+Serves only as a correctness oracle for tests; must not be imported by runtime
	4	+code. Adapted from ``dsa_sparse_attention_ascend950/test_sparse_mla_fwd.py``
	5	+(itself adapted from the CUDA example ``tilelang/examples/deepseek_v32``).
	6	+"""
	7	+
	8	+import torch
	9	+
	10	+
	11	+def make_sparse_mla_indices(
	12	+ batch: int,
	13	+ seq_len: int,
	14	+ kv_group: int,
	15	+ topk: int,
	16	+ seq_len_kv: int,
	17	+ device: torch.device,
	18	+ seed: int = 0,
	19	+) -> torch.Tensor:
	20	+ """CUDA-reference index generator: causal-valid prefix, SKV padding tail."""
	21	+ torch.random.manual_seed(seed)
	22	+ indices = torch.full((batch, seq_len, kv_group, topk), seq_len_kv, dtype=torch.int32, device=device)
	23	+ for b in range(batch):
	24	+ for t in range(seq_len):
	25	+ for g in range(kv_group):
	26	+ sel = torch.randperm(max(1, t))[:topk]
	27	+ indices[b, t, g, : len(sel)] = sel
	28	+ return indices
	29	+
	30	+
	31	+def sparse_mla_fwd_ref(
	32	+ q: torch.Tensor,
	33	+ kv: torch.Tensor,
	34	+ indices: torch.Tensor,
	35	+ sm_scale: float \| None = None,
	36	+) -> tuple[torch.Tensor, torch.Tensor]:
	37	+ """Float32 reference (scatter-mask semantics) returning (O, Lse)."""
	38	+ q = q.float()
	39	+ kv = kv.float()
	40	+ indices = indices.transpose(1, 2)
	41	+ b, sq, h, dim_qk = q.shape
	42	+ _, sk, g, _ = kv.shape
	43	+
	44	+ dim = 512
	45	+ k = kv
	46	+ v = kv[..., :dim]
	47	+
	48	+ device = q.device
	49	+ causal = torch.arange(0, sq, dtype=torch.int32, device=device).view(-1, 1) >= torch.arange(
	50	+ 0, sk, dtype=torch.int32, device=device
	51	+ ).view(1, -1)
	52	+ mask = q.new_zeros(b, g, sq, sk + 1, dtype=torch.bool).scatter(3, indices.long(), 1)
	53	+ mask = mask[..., :-1]
	54	+ mask = mask & causal.view(1, 1, sq, sk)
	55	+ mask[:, :, 0, 0] = True
	56	+ mask = mask.view(b, g, 1, sq, sk)
	57	+
	58	+ scale = dim_qk**-0.5 if sm_scale is None else sm_scale
	59	+ q4 = q.view(b, sq, g, -1, dim_qk)
	60	+ score = torch.einsum('bmghd,bngd->bghmn', q4, k).masked_fill(~mask, float('-inf'))
	61	+ lse = torch.logsumexp(score * scale, dim=-1) # [b, g, h_index, sq]
	62	+ p = torch.softmax(score * scale, dim=-1)
	63	+ o = torch.einsum('bghmn,bngd->bmghd', p.type(v.dtype), v).reshape(b, sq, h, dim)
	64	+ return o.to(torch.bfloat16), lse.permute(0, 3, 1, 2).reshape(b, sq, h)
	65	+
	66	+
	67	+def sparse_mla_fwd_ref_per_slot(
	68	+ q: torch.Tensor,
	69	+ kv: torch.Tensor,
	70	+ indices: torch.Tensor,
	71	+ sm_scale: float \| None = None,
	72	+) -> tuple[torch.Tensor, torch.Tensor]:
	73	+ """Per-slot reference: softmax over the topk gathered slots themselves.
	74	+
	75	+ The kernel counts duplicate indices once per slot, while
	76	+ :func:`sparse_mla_fwd_ref` collapses duplicates through its scatter mask.
	77	+ The two agree for duplicate-free (randperm) inputs; duplicate-index inputs
	78	+ follow per-slot semantics, validated with this reference.
	79	+ """
	80	+ q = q.float()
	81	+ kv = kv.float()
	82	+ b, sq, h, dim_qk = q.shape
	83	+ _, sk, g, _ = kv.shape
	84	+ _, _, _, topk = indices.shape
	85	+ dim = 512
	86	+
	87	+ idx = indices.transpose(1, 2).long() # [b, g, sq, topk]
	88	+ kv_g = (
	89	+ kv.permute(0, 2, 1, 3)
	90	+ .gather(2, idx.reshape(b, g, sq * topk, 1).expand(b, g, sq * topk, dim_qk))
	91	+ .view(b, g, sq, topk, dim_qk)
	92	+ )
	93	+ q4 = q.view(b, sq, g, -1, dim_qk).permute(0, 2, 3, 1, 4) # [b, g, hi, sq, d]
	94	+
	95	+ scale = dim_qk**-0.5 if sm_scale is None else sm_scale
	96	+ score = torch.einsum('bghsd,bgstd->bghst', q4, kv_g) * scale # [b,g,hi,sq,topk]
	97	+ valid = (idx.unsqueeze(2) <= torch.arange(sq, device=q.device).view(1, 1, 1, sq, 1)) & (idx.unsqueeze(2) < sk)
	98	+ score = score.masked_fill(~valid, float('-inf'))
	99	+ score[:, :, :, 0, 0] = torch.where(
	100	+ score[:, :, :, 0, 0].isinf(), torch.zeros_like(score[:, :, :, 0, 0]), score[:, :, :, 0, 0]
	101	+ )
	102	+ lse = torch.logsumexp(score, dim=-1) # [b, g, hi, sq]
	103	+ p = torch.softmax(score, dim=-1)
	104	+ o = torch.einsum('bghst,bgstd->bghsd', p, kv_g[..., :dim])
	105	+ o = o.permute(0, 3, 1, 2, 4).reshape(b, sq, h, dim)
	106	+ return o.to(torch.bfloat16), lse.permute(0, 3, 1, 2).reshape(b, sq, h)

Atests/attention/test_sparse_mla_fwd.py+200-0

@@ -0,0 +1,200 @@
	1	+import os
	2	+
	3	+import pytest
	4	+import torch
	5	+
	6	+import cann_ops_tilelang
	7	+from cann_ops_tilelang.config import get_device
	8	+from tests.attention.sparse_mla_fwd_ref import (
	9	+ make_sparse_mla_indices,
	10	+ sparse_mla_fwd_ref,
	11	+ sparse_mla_fwd_ref_per_slot,
	12	+)
	13	+from tests.testing.bench import make_param_id
	14	+from tests.testing.generator import get_test_level
	15	+from tests.testing.numeric import count_bytes
	16	+
	17	+os.environ.setdefault('TILELANG_PRINT_ON_COMPILATION', '0')
	18	+pytestmark = pytest.mark.npu
	19	+
	20	+
	21	+def generate_test_data(params: dict) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
	22	+ device = get_device()
	23	+ seed = params.get('seed', 0)
	24	+ torch.manual_seed(seed)
	25	+ batch, seq_len = params['b'], params['s']
	26	+ seq_len_kv, heads, kv_group, topk = params['skv'], params['h'], params['hkv'], params['topk']
	27	+ q = torch.randn((batch, seq_len, heads, 576), dtype=torch.bfloat16, device=device)
	28	+ scale = params.get('input_scale')
	29	+ if scale is not None:
	30	+ q = q * scale
	31	+ kv = torch.randn((batch, seq_len_kv, kv_group, 576), dtype=torch.bfloat16, device=device) * scale
	32	+ else:
	33	+ kv = torch.randn((batch, seq_len_kv, kv_group, 576), dtype=torch.bfloat16, device=device)
	34	+ indices = make_sparse_mla_indices(batch, seq_len, kv_group, topk, seq_len_kv, device, seed=seed)
	35	+ return q, kv, indices
	36	+
	37	+
	38	+def generate_test_params(level: int) -> list[dict]:
	39	+ core = [
	40	+ {'b': 1, 's': 128, 'skv': 256, 'h': 16, 'hkv': 1, 'topk': 64},
	41	+ {'b': 1, 's': 64, 'skv': 128, 'h': 64, 'hkv': 1, 'topk': 64},
	42	+ {'b': 1, 's': 64, 'skv': 128, 'h': 64, 'hkv': 4, 'topk': 64},
	43	+ ]
	44	+ default = core + [
	45	+ {'b': 1, 's': 64, 'skv': 128, 'h': 128, 'hkv': 1, 'topk': 64},
	46	+ {'b': 1, 's': 128, 'skv': 256, 'h': 16, 'hkv': 1, 'topk': 128},
	47	+ {'b': 2, 's': 96, 'skv': 192, 'h': 32, 'hkv': 1, 'topk': 64},
	48	+ {'b': 1, 's': 8, 'skv': 64, 'h': 16, 'hkv': 1, 'topk': 64},
	49	+ ]
	50	+ full = default + [
	51	+ {'b': 1, 's': 1, 'skv': 64, 'h': 16, 'hkv': 1, 'topk': 64},
	52	+ {'b': 1, 's': 64, 'skv': 128, 'h': 192, 'hkv': 1, 'topk': 64},
	53	+ {'b': 1, 's': 256, 'skv': 512, 'h': 16, 'hkv': 1, 'topk': 192},
	54	+ {'b': 1, 's': 64, 'skv': 64, 'h': 16, 'hkv': 1, 'topk': 256},
	55	+ {'b': 1, 's': 64, 'skv': 128, 'h': 256, 'hkv': 2, 'topk': 64},
	56	+ {'b': 1, 's': 128, 'skv': 256, 'h': 16, 'hkv': 1, 'topk': 64, 'sm_scale': 0.05},
	57	+ {'b': 1, 's': 128, 'skv': 256, 'h': 16, 'hkv': 1, 'topk': 64, 'input_scale': 8.0},
	58	+ ]
	59	+ return {0: core, 1: default, 2: full}[level]
	60	+
	61	+
	62	+def assert_similar(out: torch.Tensor, ref_out: torch.Tensor, lse: torch.Tensor, ref_lse: torch.Tensor) -> None:
	63	+ assert not torch.isnan(out.float()).any(), 'Output contains NaN'
	64	+ assert not torch.isnan(lse).any(), 'Lse contains NaN'
	65	+ out_diff = (out.float() - ref_out.float()).abs()
	66	+ denom = max(ref_out.float().abs().max().item(), 1e-6)
	67	+ rel = out_diff.max().item() / denom
	68	+ assert rel < 1e-2, f'Output rel diff {rel:.4f} >= 1e-2'
	69	+ lse_diff = (lse - ref_lse).abs()
	70	+ lse_tol = 5e-2 * (1.0 + ref_lse.abs())
	71	+ bad = (lse_diff > lse_tol) & torch.isfinite(ref_lse)
	72	+ assert not bad.any(), f'Lse mismatch: max diff {lse_diff.max().item():.4f}'
	73	+
	74	+
	75	+@pytest.mark.parametrize('params', generate_test_params(get_test_level()), ids=make_param_id)
	76	+def test_sparse_mla_fwd(params):
	77	+ q, kv, indices = generate_test_data(params)
	78	+ q_ref, kv_ref, idx_ref = q.clone(), kv.clone(), indices.clone()
	79	+
	80	+ out, lse = cann_ops_tilelang.sparse_mla_fwd(q, kv, indices, sm_scale=params.get('sm_scale'))
	81	+ torch.npu.synchronize()
	82	+
	83	+ assert out.dtype == torch.bfloat16 and lse.dtype == torch.float32
	84	+ assert out.shape == (params['b'], params['s'], params['h'], 512)
	85	+ assert lse.shape == (params['b'], params['s'], params['h'])
	86	+ assert out.is_contiguous() and lse.is_contiguous()
	87	+ assert torch.equal(q, q_ref) and torch.equal(kv, kv_ref) and torch.equal(indices, idx_ref)
	88	+ ref_out, ref_lse = sparse_mla_fwd_ref(q, kv, indices, sm_scale=params.get('sm_scale'))
	89	+ assert_similar(out, ref_out, lse, ref_lse)
	90	+
	91	+
	92	+def test_sparse_mla_fwd_duplicate_indices():
	93	+ device = get_device()
	94	+ torch.manual_seed(0)
	95	+ batch, seq_len, seq_len_kv, heads, kv_group, topk = 1, 128, 256, 16, 1, 64
	96	+ q = torch.randn((batch, seq_len, heads, 576), dtype=torch.bfloat16, device=device)
	97	+ kv = torch.randn((batch, seq_len_kv, kv_group, 576), dtype=torch.bfloat16, device=device)
	98	+ indices = torch.zeros((batch, seq_len, kv_group, topk), dtype=torch.int32, device=device)
	99	+
	100	+ out, lse = cann_ops_tilelang.sparse_mla_fwd(q, kv, indices)
	101	+ torch.npu.synchronize()
	102	+
	103	+ ref_out, ref_lse = sparse_mla_fwd_ref_per_slot(q, kv, indices)
	104	+ assert_similar(out, ref_out, lse, ref_lse)
	105	+
	106	+
	107	+def test_sparse_mla_fwd_determinism():
	108	+ q, kv, indices = generate_test_data({'b': 1, 's': 128, 'skv': 256, 'h': 16, 'hkv': 1, 'topk': 64})
	109	+
	110	+ out1, lse1 = cann_ops_tilelang.sparse_mla_fwd(q, kv, indices)
	111	+ out2, lse2 = cann_ops_tilelang.sparse_mla_fwd(q, kv, indices)
	112	+ torch.npu.synchronize()
	113	+
	114	+ assert torch.equal(out1, out2), 'Output not deterministic across launches'
	115	+ assert torch.equal(lse1, lse2), 'Lse not deterministic across launches'
	116	+
	117	+
	118	+def test_sparse_mla_fwd_out_buffer():
	119	+ params = {'b': 1, 's': 64, 'skv': 128, 'h': 16, 'hkv': 1, 'topk': 64}
	120	+ q, kv, indices = generate_test_data(params)
	121	+ ref_out, ref_lse = sparse_mla_fwd_ref(q, kv, indices)
	122	+ device = get_device()
	123	+ out = torch.full((1, 64, 16, 512), float('nan'), dtype=torch.bfloat16, device=device)
	124	+ lse = torch.empty((1, 64, 16), dtype=torch.float32, device=device)
	125	+
	126	+ returned_out, returned_lse = cann_ops_tilelang.sparse_mla_fwd(q, kv, indices, out=out, lse=lse)
	127	+ torch.npu.synchronize()
	128	+
	129	+ assert returned_out is out and returned_lse is lse
	130	+ assert_similar(out, ref_out, lse, ref_lse)
	131	+
	132	+ out_elements = 1 * 64 * 16 * 512
	133	+ buf = torch.empty(out_elements + 2 + 2 * 1 * 64 * 16, dtype=torch.bfloat16, device=device)
	134	+ aliased_out = buf[:out_elements].view(1, 64, 16, 512)
	135	+ aliased_lse = buf[out_elements + 2 :].view(torch.float32).view(1, 64, 16)
	136	+ with pytest.raises(ValueError, match='alias'):
	137	+ cann_ops_tilelang.sparse_mla_fwd(q, kv, indices, out=aliased_out, lse=aliased_lse)
	138	+ with pytest.raises(ValueError, match='shape'):
	139	+ wrong = torch.empty((1, 64, 16, 256), dtype=torch.bfloat16, device=device)
	140	+ cann_ops_tilelang.sparse_mla_fwd(q, kv, indices, out=wrong)
	141	+ with pytest.raises(TypeError, match='bfloat16'):
	142	+ wrong = torch.empty((1, 64, 16, 512), dtype=torch.float32, device=device)
	143	+ cann_ops_tilelang.sparse_mla_fwd(q, kv, indices, out=wrong)
	144	+
	145	+
	146	+def test_sparse_mla_fwd_empty():
	147	+ device = get_device()
	148	+ q = torch.empty((0, 8, 16, 576), dtype=torch.bfloat16, device=device)
	149	+ kv = torch.empty((0, 16, 1, 576), dtype=torch.bfloat16, device=device)
	150	+ indices = torch.empty((0, 8, 1, 64), dtype=torch.int32, device=device)
	151	+
	152	+ out, lse = cann_ops_tilelang.sparse_mla_fwd(q, kv, indices)
	153	+
	154	+ assert out.shape == (0, 8, 16, 512) and out.dtype == torch.bfloat16
	155	+ assert lse.shape == (0, 8, 16) and lse.dtype == torch.float32
	156	+
	157	+
	158	+def test_sparse_mla_fwd_invalid_inputs():
	159	+ device = get_device()
	160	+ q, kv, indices = generate_test_data({'b': 1, 's': 64, 'skv': 128, 'h': 16, 'hkv': 1, 'topk': 64})
	161	+
	162	+ with pytest.raises(TypeError, match='torch.Tensor'):
	163	+ cann_ops_tilelang.sparse_mla_fwd(q.tolist(), kv, indices)
	164	+ with pytest.raises(ValueError, match='NPU'):
	165	+ cann_ops_tilelang.sparse_mla_fwd(q.to('cpu'), kv.to('cpu'), indices.to('cpu'))
	166	+ with pytest.raises(TypeError, match='bfloat16'):
	167	+ cann_ops_tilelang.sparse_mla_fwd(q.to(torch.float16), kv, indices)
	168	+ with pytest.raises(TypeError, match='int32'):
	169	+ cann_ops_tilelang.sparse_mla_fwd(q, kv, indices.to(torch.int64))
	170	+ with pytest.raises(ValueError, match='contiguous'):
	171	+ cann_ops_tilelang.sparse_mla_fwd(q.transpose(1, 2), kv, indices)
	172	+ with pytest.raises(ValueError, match='topk'):
	173	+ cann_ops_tilelang.sparse_mla_fwd(q, kv, indices[..., :32].contiguous())
	174	+ with pytest.raises(ValueError, match='576'):
	175	+ cann_ops_tilelang.sparse_mla_fwd(q[..., :128].contiguous(), kv, indices)
	176	+ bad_q = torch.randn((1, 64, 14, 576), dtype=torch.bfloat16, device=device)
	177	+ bad_kv = torch.randn((1, 128, 4, 576), dtype=torch.bfloat16, device=device)
	178	+ bad_indices = torch.zeros((1, 64, 4, 64), dtype=torch.int32, device=device)
	179	+ with pytest.raises(ValueError, match='divisible'):
	180	+ cann_ops_tilelang.sparse_mla_fwd(bad_q, bad_kv, bad_indices)
	181	+ bad_heads = torch.randn((1, 64, 24, 576), dtype=torch.bfloat16, device=device)
	182	+ with pytest.raises(ValueError, match='ratio'):
	183	+ cann_ops_tilelang.sparse_mla_fwd(bad_heads, kv, indices)
	184	+
	185	+
	186	+@pytest.mark.benchmark
	187	+@pytest.mark.parametrize('params', generate_test_params(0), ids=make_param_id)
	188	+def test_sparse_mla_fwd_benchmark(benchmark_timer, benchmark_record, params):
	189	+ q, kv, indices = generate_test_data(params)
	190	+ out, lse = cann_ops_tilelang.sparse_mla_fwd(q, kv, indices)
	191	+
	192	+ time_us = benchmark_timer(lambda: cann_ops_tilelang.sparse_mla_fwd(q, kv, indices))
	193	+ num_bytes = count_bytes(q, kv, indices, out, lse)
	194	+ benchmark_record(
	195	+ kernel='sparse_mla_fwd',
	196	+ operation='fwd',
	197	+ params=params,
	198	+ time_us=time_us,
	199	+ bandwidth_gbs=num_bytes / time_us / 1e3,
	200	+ )