[WIP]feat(attention): add FP8 HISA prefill indexer for Ascend 950-ops-tilelang-AtomGit

MOAT.xml+4-4

	@@ -21,12 +21,12 @@
21	<filteritem type="filename" name="*.yaml" desc="YAML configuration files do not require license headers"/>	21	<filteritem type="filename" name="*.yaml" desc="YAML configuration files do not require license headers"/>
22	<filteritem type="filename" name="*.yml" desc="YAML configuration files do not require license headers"/>	22	<filteritem type="filename" name="*.yml" desc="YAML configuration files do not require license headers"/>
23	<filteritem type="filename" name="pyproject.toml" desc="Python package and tool configuration; project license is declared in its metadata"/>	23	<filteritem type="filename" name="pyproject.toml" desc="Python package and tool configuration; project license is declared in its metadata"/>
24	- <filteritem type="filepath" name="tests/[^/]+/benchmark/.*\.jsonl" desc="Machine-readable benchmark baselines cannot contain comment headers"/>	24	+ <filteritem type="filepath" name="tests/[^/]+/benchmark/.*[.]jsonl" desc="Machine-readable benchmark baselines cannot contain comment headers"/>
25	- <filteritem type="filepath" name="tests/[^/]+/npu_mem_profiles/.*\.json" desc="Machine-readable NPU memory profiles cannot contain comment headers"/>	25	+ <filteritem type="filepath" name="tests/[^/]+/npu_mem_profiles/.*[.]json" desc="Machine-readable NPU memory profiles cannot contain comment headers"/>
26	</filefilter>	26	</filefilter>
27	<filefilter name="copyrightPolicyFilter" desc="Exceptions for copyright header policies">	27	<filefilter name="copyrightPolicyFilter" desc="Exceptions for copyright header policies">
28	- <filteritem type="filepath" name="tests/[^/]+/benchmark/.*\.jsonl" desc="JSONL benchmark records cannot contain comment headers"/>	28	+ <filteritem type="filepath" name="tests/[^/]+/benchmark/.*[.]jsonl" desc="JSONL benchmark records cannot contain comment headers"/>
29	- <filteritem type="filepath" name="tests/[^/]+/npu_mem_profiles/.*\.json" desc="JSON memory profiles cannot contain comment headers"/>	29	+ <filteritem type="filepath" name="tests/[^/]+/npu_mem_profiles/.*[.]json" desc="JSON memory profiles cannot contain comment headers"/>
30	</filefilter>	30	</filefilter>
31	<filefilter name="licenseFileNamePolicyFilter" desc="Exceptions for LICENSE file policies">	31	<filefilter name="licenseFileNamePolicyFilter" desc="Exceptions for LICENSE file policies">
32	</filefilter>	32	</filefilter>

Adocs/dsa_hisa.md+87-0

@@ -0,0 +1,87 @@
	1	+# HISA 稀疏索引
	2	+
	3	+`cann_ops_tilelang.hisa_indexer` 为 prefill 阶段的每个 query 返回可见 K 的相对 token 位置。它先按块对 K 求均值、重新量化并选出候选块，再对块内原始 token 打分并执行 top-k。调用方可以将输出索引用于后续的稀疏注意力计算。
	4	+
	5	+\| 参数 \| 形状及类型 \| 含义 \|
	6	+\| --- \| --- \| --- \|
	7	+\| `q`、`k` \| `[M,H,D]`、`[N,D]`，FP8 E4M3FN \| 索引 query 和每个 token 的共享 key \|
	8	+\| `k_scale`、`weights` \| `[N]`、`[M,H]`，FP32 \| K 的反量化系数和 query 的各头权重 \|
	9	+\| `cu_seqlen_ks`、`cu_seqlen_ke` \| 各 `[M]`，INT32 \| 每个 query 在打包 K 中可见的左闭右开窗口，调用方保证 `0 <= ks <= ke <= N` \|
	10	+\| `k_block_size`、`block_topk`、`topk_tokens` \| 正整数，关键字参数 \| 每块 token 数、候选块数上限、最终 token 数上限 \|
	11	+
	12	+输入张量在同一 NPU 上连续存放，N 须整除为完整的 K 块。块分数采用 FP32 累加，选块前按最近偶数规则舍入到 BF16。非空窗口的两个边界块先按块号升序选取，仅选一块时保留首块。内块按分数降序排列，同分时块号小者优先。候选 token 按 FP32 分数选择，同分时按候选槽位取较小者。
	13	+
	14	+接口自行分配 INT32 输出，签名不接收 `out=`。索引相对于所在行的 `ks`，无效槽位填 `-1`，输出按索引集合使用。令 `B = N // k_block_size`、`C = min(block_topk, B) * k_block_size`，输出形状为 `[M, min(topk_tokens, C)]`。空 Q 或空 K 返回对应形状的空张量。
	15	+
	16	+非空输入的 `D` 与块长须为 32 的倍数。块数 `B` 最多 1024，小于 64 时须能整除 64。每行有效候选数 `C` 须为 64 的倍数，最多 8192。位置比较使用 FP32，要求 `N <= 2**24`。入口还会按形状检查计算 tile 能否放入片上存储，例如 `H=64, k_block_size=512` 超过 token 打分的累加器容量。这些形状检查均在首个 kernel 启动前完成。
	17	+
	18	+典型配置为 H=64、D=128、块长 128，M=N 取 1024、4096、8192，每档包含单序列和打包多序列。8192 个候选选 2048 个 token 时使用专用 radix-4 内核。接口面向 Ascend 950 单流推理，支持 eager 和固定形状的 Torch 图执行。各阶段共享的中间分数和候选块是索引计算的内部数据，调用方只取得最终相对索引。GPU 算法和数据类型参考 [TileLang HISA](https://github.com/tile-ai/tilelang/tree/main/examples/dsa_hisa)。
	19	+
	20	+```python
	21	+from cann_ops_tilelang import hisa_indexer
	22	+
	23	+indices = hisa_indexer(
	24	+ q, k, k_scale, weights, cu_seqlen_ks, cu_seqlen_ke,
	25	+ k_block_size=128, block_topk=64, topk_tokens=2048,
	26	+)
	27	+```
	28	+
	29	+图模式沿用相同的公开接口、数据类型和并列规则。算子在 Torch 中注册为 `cann_ops_tilelang::hisa_indexer`，输出为独立分配的 INT32 张量。下面的例子在同一条流上以固定输入形状运行推理。
	30	+
	31	+```python
	32	+import torch
	33	+import torch_npu
	34	+from cann_ops_tilelang import hisa_indexer
	35	+
	36	+index = torch.compile(hisa_indexer, backend='npugraph_ex', fullgraph=True, dynamic=False)
	37	+with torch.inference_mode():
	38	+ indices = index(
	39	+ q, k, k_scale, weights, cu_seqlen_ks, cu_seqlen_ke,
	40	+ k_block_size=128, block_topk=64, topk_tokens=2048,
	41	+ )
	42	+```
	43	+
	44	+从仓库根目录运行 HISA 定向测试。测试包含独立的 PyTorch 参考，从原始输入计算池化、重新量化、选块、token 分数和窗口偏移。切换到已配套的 PTO 环境后可运行相同的功能用例。六组性能基线采用默认 AscendC 后端和仓库的 msprof 计时。
	45	+
	46	+```bash
	47	+OPS_TILELANG_TEST_LEVEL=2 TILELANG_DEFAULT_TARGET=auto python -m pytest --log-cli-level=INFO tests/attention/test_dsa_hisa.py -k hisa_compile
	48	+OPS_TILELANG_TEST_LEVEL=2 TILELANG_DEFAULT_TARGET=auto python -m pytest tests/attention/test_dsa_hisa.py -m 'not benchmark'
	49	+OPS_TILELANG_TEST_LEVEL=2 TILELANG_DEFAULT_TARGET=pto python -m pytest tests/attention/test_dsa_hisa.py -m 'not benchmark'
	50	+TILELANG_DEFAULT_TARGET=auto python -m pytest tests/attention/test_dsa_hisa.py -m benchmark --run-benchmark
	51	+```
	52	+
	53	+## 8192 单序列性能剖析
	54	+
	55	+在 Ascend 950 单卡运行公开入口 `hisa_indexer` 的 eager 默认后端，沿用 benchmark 的输入和种子。M=N=8192，H=64，D=128，每块 128 个 token，粗排的块数上限为 64，精排从至多 8192 个候选 token 中选 2048 个。窗口从序列起点到当前 token，含当前 token。
	56	+
	57	+`msprof PipeUtilization` 预热 5 次后采集 16 次。从首个池化内核开始到最后一个 token 选择内核结束，设备跨度中位数为 8.109 ms。图和表均选取设备跨度最接近中位数的一次调用，以池化内核开始为零点。
	58	+
	59	+```text
	60	+ms 0 2 4 6 8
	61	+ \|---------\|---------\|---------\|---------\|
	62	+pool ^
	63	+block_topk ===
	64	+token_score =========================
	65	+token_topk =============
	66	+```
	67	+
	68	+时间轴每格约为 0.2 ms。`^` 标出耗时 9 µs 的池化内核，精确起止时间见下表。
	69	+
	70	+\| 阶段 \| 开始（µs） \| 结束（µs） \| 耗时（µs） \| 距上一内核结束（µs） \|
	71	+\| --- \| ---: \| ---: \| ---: \| ---: \|
	72	+\| K 均值池化 \| 0.000 \| 9.006 \| 9.006 \| — \|
	73	+\| 块打分和选择 \| 9.028 \| 579.660 \| 570.632 \| 0.022 \|
	74	+\| Token 打分 \| 579.682 \| 5627.941 \| 5048.259 \| 0.022 \|
	75	+\| Token 选择 \| 5627.963 \| 8108.123 \| 2480.160 \| 0.022 \|
	76	+
	77	+按每次调用中四个内核的执行时间求占比，再取 16 次调用的中位数，token 打分占 62.3%，token 选择占 30.6%。下表给出各阶段在 16 次调用中的流水线活跃率中位数。各流水线可以同时工作。
	78	+
	79	+\| 阶段 \| AIC MAC \| AIC MTE2 \| AIV vector \| AIV MTE2 \| AIV MTE3 \| AIV scalar \|
	80	+\| --- \| ---: \| ---: \| ---: \| ---: \| ---: \| ---: \|
	81	+\| K 均值池化 \| 0% \| 0% \| 12.95% \| 15.15% \| 66.35% \| 6.90% \|
	82	+\| 块打分和选择 \| 10.35% \| 45.40% \| 64.90% \| 11.40% \| 8.00% \| 5.40% \|
	83	+\| Token 打分 \| 33.20% \| 53.80% \| 36.90% \| 14.70% \| 43.90% \| 12.00% \|
	84	+\| Token 选择 \| 0% \| 0% \| 89.30% \| 3.70% \| 1.10% \| 8.45% \|
	85	+
	86	+- 实际 HBM 带宽。同一输入另跑 100 次，开启 `msprof --aic-metrics=Memory --sys-hardware-mem=on --sys-hardware-mem-freq=100`。从导出数据库的 `HBM` 表逐时刻汇总各 HBM stack，取稳态设备窗口内 81 个完整区间，按区间时长加权。实测读 46.69 GB/s、写 41.66 GB/s，合计 88.36 GB/s。同卡、同次采集中连续复制 256 MiB 张量 1024 次，16 个完整区间实测 HBM 读写合计 2522.00 GB/s。两者的实测 HBM 吞吐之比为 3.50%。硬件采样间隔约 10 ms，只据此分析整段稳态负载。结合流水线指标，后续优先分析 token 打分阶段的搬运与计算重叠、token 选择阶段的向量指令开销。
	87	+- Ascend C 对照范围。在公开 `cann/ops-transformer` 源码中核对了两个最接近的实现。正式目录的 `QuantLightningIndexerV2` 直接对给定上下文做 token top-k。实验两级实现先按块归约 token 分数的最大值，其量化 consumer 仅支持 A2/A3、INT8、块长 8。HISA 先对块内 K 求均值、重新量化并打分，保留可见首尾块，再按 FP32 分数选 token。此次检索未找到可在目标设备上以同精度、同选择规则下直接对测的公开 Ascend C 实现。[QuantLightningIndexerV2](https://gitcode.com/cann/ops-transformer/blob/bbcdd0b99dfc822112363fa1058f4641f3c717e5/attention/quant_lightning_indexer_v2/README.md) · [两级实现的块归约](https://gitcode.com/cann/ops-transformer/blob/bbcdd0b99dfc822112363fa1058f4641f3c717e5/experimental/attention/quant_lightning_indexer_v2/op_kernel/arch22/quant_lightning_indexer_v2_service_vector_arch22.h#L482) · [量化 consumer 的约束](https://gitcode.com/cann/ops-transformer/blob/bbcdd0b99dfc822112363fa1058f4641f3c717e5/experimental/attention/quant_sparse_lightning_indexer/README.md)

Msrc/cann_ops_tilelang/__init__.py+6-0

	@@ -1,3 +1,7 @@
		1	+# Copyright (c) 2026 ops-tilelang contributors
		2	+# SPDX-License-Identifier: MIT
		3	+# See LICENSE in the repository root for the full license text.
		4	+
1	"""High-performance Ascend operators implemented with TileLang."""	5	"""High-performance Ascend operators implemented with TileLang."""
2		6
3	__all__ = [	7	__all__ = [
	@@ -61,6 +65,7 @@ __all__ = [
61	'get_logits_cur_by_step',	65	'get_logits_cur_by_step',
62	'get_prompt_logprobs_tensor',	66	'get_prompt_logprobs_tensor',
63	'get_sum_smaller_probs',	67	'get_sum_smaller_probs',
		68	+ 'hisa_indexer',
64	'kv_index_cat',	69	'kv_index_cat',
65	'layer_norm_fn',	70	'layer_norm_fn',
66	'load_index_kv_cache',	71	'load_index_kv_cache',
	@@ -89,6 +94,7 @@ __all__ = [
89	from .attention import (	94	from .attention import (
90	convert_kv_index,	95	convert_kv_index,
91	get_fused_dflash_get_input,	96	get_fused_dflash_get_input,
		97	+ hisa_indexer,
92	kv_index_cat,	98	kv_index_cat,
93	load_index_kv_cache,	99	load_index_kv_cache,
94	load_kv_cache,	100	load_kv_cache,

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

	@@ -1,8 +1,13 @@
		1	+# Copyright (c) 2026 ops-tilelang contributors
		2	+# SPDX-License-Identifier: MIT
		3	+# See LICENSE in the repository root for the full license text.
		4	+
1	"""Attention and KV-cache operators for Ascend."""	5	"""Attention and KV-cache operators for Ascend."""
2		6
3	__all__ = [	7	__all__ = [
4	'convert_kv_index',	8	'convert_kv_index',
5	'get_fused_dflash_get_input',	9	'get_fused_dflash_get_input',
		10	+ 'hisa_indexer',
6	'kv_index_cat',	11	'kv_index_cat',
7	'load_index_kv_cache',	12	'load_index_kv_cache',
8	'load_kv_cache',	13	'load_kv_cache',
	@@ -26,6 +31,7 @@ from .cache import (
26	store_kv_cache,	31	store_kv_cache,
27	store_kv_cache_llama,	32	store_kv_cache_llama,
28	)	33	)
		34	+from .dsa_hisa import hisa_indexer
29	from .flash_mla_utils import convert_kv_index	35	from .flash_mla_utils import convert_kv_index
30	from .get_kv_index import get_fused_dflash_get_input	36	from .get_kv_index import get_fused_dflash_get_input
31	from .kv_index_cat import kv_index_cat	37	from .kv_index_cat import kv_index_cat

Asrc/cann_ops_tilelang/attention/dsa_hisa/__init__.py+7-0

@@ -0,0 +1,7 @@
	1	+# Copyright (c) 2026 ops-tilelang contributors
	2	+# SPDX-License-Identifier: MIT
	3	+# See LICENSE in the repository root for the full license text.
	4	+
	5	+__all__ = ['hisa_indexer']
	6	+
	7	+from .hisa import hisa_indexer

Asrc/cann_ops_tilelang/attention/dsa_hisa/block_sparse_mqa_fp8.py+578-0

@@ -0,0 +1,578 @@
	1	+# Copyright (c) 2026 ops-tilelang contributors
	2	+# SPDX-License-Identifier: MIT
	3	+# See LICENSE in the repository root for the full license text.
	4	+
	5	+"""Block-sparse FP8 MQA scoring with FP32 accumulation and causal masking."""
	6	+
	7	+from math import gcd
	8	+
	9	+import tilelang.ascend.language as T
	10	+import torch
	11	+from tilelang.ascend.language import simd as S
	12	+
	13	+from .._simd import is_pto_target
	14	+from .tilelang_utils import EAGER_WORKSPACE, cached_kernel
	15	+
	16	+NEG_INF = float('-inf')
	17	+
	18	+
	19	+PAD_H = 16
	20	+
	21	+
	22	+L0C_BUDGET = 64 * 1024
	23	+
	24	+
	25	+SIMD_H = 32
	26	+
	27	+
	28	+NTILE_CAP = 4
	29	+
	30	+
	31	+def mdim_of(num_heads: int) -> int:
	32	+ return ((num_heads + 15) // 16) * 16
	33	+
	34	+
	35	+def pick_ntile(num_heads: int, kb: int, block_topk: int) -> int:
	36	+ if kb % 2 != 0:
	37	+ return 1
	38	+ budget = (L0C_BUDGET // (mdim_of(num_heads) * 4)) // kb
	39	+ ntile = max(1, min(block_topk, NTILE_CAP, budget))
	40	+ while ntile > 1 and block_topk % ntile != 0:
	41	+ ntile //= 2
	42	+ return max(1, ntile)
	43	+
	44	+
	45	+def fused_ok(num_heads: int, head_dim: int, kb: int, block_topk: int) -> bool:
	46	+ if num_heads <= 0 or head_dim % 16 != 0 or kb % 2 != 0:
	47	+ return False
	48	+ rows = mdim_of(num_heads)
	49	+ tiles = pick_ntile(num_heads, kb, block_topk)
	50	+ if rows * kb * 4 > L0C_BUDGET or (rows + tiles * kb) * head_dim > 256 * 1024:
	51	+ return False
	52	+ if is_pto_target():
	53	+ # The PTO score transfer uses a square-padded leading dimension.
	54	+ cols = gcd(kb, 128)
	55	+ return (max(rows, cols) + cols) * head_dim <= 256 * 1024
	56	+ return True
	57	+
	58	+
	59	+def _token_group_io(ntile, kb, half_block, half_group_tokens, head_dim, tile_kb=None):
	60	+ """Build transfers between packed half-blocks and their original token slots."""
	61	+
	62	+ def block_begin(blk, offset):
	63	+ return blk * kb if tile_kb is None else blk * kb + offset
	64	+
	65	+ def block_end(blk, offset):
	66	+ return (blk + 1) * kb if tile_kb is None else blk * kb + offset + tile_kb
	67	+
	68	+ @T.macro
	69	+ def load_key_group(keys, block_indices, kp_l1, bx, first, offset):
	70	+ for hf in range(2):
	71	+ for ti in range(ntile):
	72	+ blk = block_indices[bx, first + ti]
	73	+ key_half_start = T.meta_var(block_begin(blk, offset) + hf * half_block)
	74	+ key_half_end = T.meta_var(block_begin(blk, offset) + (hf + 1) * half_block)
	75	+ packed_half_start = T.meta_var(hf * half_group_tokens + ti * half_block)
	76	+ packed_half_end = T.meta_var(hf * half_group_tokens + (ti + 1) * half_block)
	77	+ T.copy(
	78	+ keys[key_half_start:key_half_end, 0:head_dim], kp_l1[packed_half_start:packed_half_end, 0:head_dim]
	79	+ )
	80	+
	81	+ @T.macro
	82	+ def load_key_metadata(key_scales, positions, block_indices, ks_ub, pas_ub, bx, first, offset):
	83	+ for ti in range(ntile):
	84	+ blk = block_indices[bx, first + ti]
	85	+ key_block_start = T.meta_var(block_begin(blk, offset))
	86	+ key_block_end = T.meta_var(block_end(blk, offset))
	87	+ half_tile_start = T.meta_var(ti * half_block)
	88	+ half_tile_end = T.meta_var((ti + 1) * half_block)
	89	+ T.dual_copy(key_scales[key_block_start:key_block_end], ks_ub[half_tile_start:half_tile_end])
	90	+ T.dual_copy(positions[key_block_start:key_block_end], pas_ub[half_tile_start:half_tile_end])
	91	+
	92	+ @T.macro
	93	+ def store_score_group(o_ub, output, bx, first, offset):
	94	+ for ti in range(ntile):
	95	+ half_tile_start = T.meta_var(ti * half_block)
	96	+ half_tile_end = T.meta_var((ti + 1) * half_block)
	97	+ row_end = T.meta_var(bx + 1)
	98	+ output_block_start = T.meta_var(block_begin(first + ti, offset))
	99	+ output_block_end = T.meta_var(block_end(first + ti, offset))
	100	+ T.dual_copy(

Asrc/cann_ops_tilelang/attention/dsa_hisa/block_topk.py+75-0

@@ -0,0 +1,75 @@
	1	+# Copyright (c) 2026 ops-tilelang contributors
	2	+# SPDX-License-Identifier: MIT
	3	+# See LICENSE in the repository root for the full license text.
	4	+
	5	+"""Select visible BF16-ranked blocks, retaining boundaries and breaking ties by block index."""
	6	+
	7	+import tilelang.ascend.language as T
	8	+from tilelang.ascend.language import simd as S
	9	+
	10	+from .tilelang_utils import VS, argmin, bf16_block_rank, block_priority, cached_kernel, ramp, redmax
	11	+
	12	+
	13	+def block_topk_select(num_rows: int, num_blocks: int, kb: int, block_topk: int):
	14	+ if num_blocks <= 0:
	15	+ raise AssertionError('the block count must be positive')
	16	+ padded_nb = ((num_blocks + VS - 1) // VS) * VS
	17	+ if block_topk > num_blocks:
	18	+ raise AssertionError('block_topk exceeds the block count')
	19	+ nv = padded_nb // VS
	20	+
	21	+ @T.macro
	22	+ def select_blocks(l_ub, idx_ub, ks_b, ke_b):
	23	+ with T.SimdVF():
	24	+ full = S.pset(32, 'PAT_ALL')
	25	+ one = S.pset(32, 'PAT_VL1')
	26	+ removed = S.vdup(T.int32(-32769), 'int32', full)
	27	+ imax = S.vdup(T.max_value(T.int32), 'int32', full)
	28	+
	29	+ v = S.alloc_local(nv, 'int32')
	30	+ for i in T.unroll(nv, explicit=True):
	31	+ indices = ramp(i)
	32	+ rank = bf16_block_rank(S.vld(l_ub[i * VS]), full)
	33	+ priority = block_priority(rank, indices, ks_b, ke_b, full)
	34	+ v[i] = S.vsel(removed, priority, S.vcmps(indices, num_blocks, full, 'ge'))
	35	+
	36	+ with T.serial(block_topk) as k:
	37	+ mx = S.vdupv(S.vcmax(redmax(v, nv, full), full), full)
	38	+ iv = S.vdupv(S.vcmin(argmin(v, nv, mx, full, imax), full), full)
	39	+ S.vsts(idx_ub[k], iv, one, 'ONEPT_B32')
	40	+ for i in T.unroll(nv, explicit=True):
	41	+ v[i] = S.vsel(removed, v[i], S.vcmp(ramp(i), iv, full, 'eq'))
	42	+
	43	+ @T.prim_func
	44	+ def main(
	45	+ scores_in: T.Tensor((num_rows, num_blocks), 'float32'),
	46	+ window_start: T.Tensor((num_rows,), 'int32'),
	47	+ window_end: T.Tensor((num_rows,), 'int32'),
	48	+ indices_out: T.Tensor((num_rows, block_topk), 'int32'),
	49	+ ):
	50	+ with T.Kernel(num_rows) as bx:
	51	+ l_ub = T.alloc_shared((padded_nb,), 'float32')
	52	+ idx_ub = T.alloc_shared((block_topk,), 'int32')
	53	+ T.fill(l_ub, -T.infinity(T.float32))
	54	+ T.copy(scores_in[bx, 0:num_blocks], l_ub[0:num_blocks])
	55	+
	56	+ ks_b = window_start[bx] // kb
	57	+ ke_b = (window_end[bx] + (kb - 1)) // kb
	58	+
	59	+ select_blocks(l_ub, idx_ub, ks_b, ke_b)
	60	+
	61	+ for k in T.serial(block_topk):
	62	+ indices_out[bx, k] = idx_ub[k]
	63	+
	64	+ return main
	65	+
	66	+
	67	+def block_topk_select_interface(logits, cu_seqlen_ks, cu_seqlen_ke, kb, block_topk):
	68	+ num_rows, num_blocks = logits.shape
	69	+ bt = min(block_topk, num_blocks)
	70	+ kern = cached_kernel(
	71	+ ('hisa-block-topk-bf16', num_rows, num_blocks, kb, bt),
	72	+ lambda: block_topk_select(num_rows, num_blocks, kb, bt),
	73	+ [3],
	74	+ )
	75	+ return kern(logits, cu_seqlen_ks, cu_seqlen_ke)

Asrc/cann_ops_tilelang/attention/dsa_hisa/fp8_block_mean_pooling.py+95-0

@@ -0,0 +1,95 @@
	1	+# Copyright (c) 2026 ops-tilelang contributors
	2	+# SPDX-License-Identifier: MIT
	3	+# See LICENSE in the repository root for the full license text.
	4	+
	5	+"""Mean-pool scaled FP8 keys in FP32 and requantize each pooled block to FP8."""
	6	+
	7	+import tilelang.ascend.language as T
	8	+from tilelang.ascend.language import simd as S
	9	+
	10	+from .tilelang_utils import cached_kernel
	11	+
	12	+FP8_MAX = 448.0
	13	+
	14	+
	15	+EPS = 1e-10
	16	+
	17	+
	18	+def fp8_native_block_mean_pooling(num_tokens: int, head_dim: int, kb: int):
	19	+ if num_tokens % kb:
	20	+ raise AssertionError('the token count must be divisible by the block size')
	21	+ num_blocks = num_tokens // kb
	22	+ if head_dim % 32:
	23	+ raise AssertionError('the head dimension must be a multiple of 32')
	24	+ chunks = (head_dim + 63) // 64
	25	+ padded_d = chunks * 64
	26	+
	27	+ @T.macro
	28	+ def clear_tail_lanes(k8):
	29	+ with T.SimdVF():
	30	+ zero = S.vcvt(S.vdup(T.float32(0), 'float32'), T.float8_e4m3fn)
	31	+ for pad_row in T.serial(kb):
	32	+ S.vsts(k8[pad_row, padded_d - 64], zero, dist='PK4_B32')
	33	+
	34	+ @T.macro
	35	+ def pool_block(k8, s_ub, q8):
	36	+ with T.SimdVF():
	37	+ full = S.pset(32, 'PAT_ALL')
	38	+ one = S.pset(32, 'PAT_VL1')
	39	+ sums = S.alloc_local(chunks, 'float32')
	40	+ for chunk in T.unroll(chunks, explicit=True):
	41	+ sums[chunk] = S.vdup(T.float32(0), 'float32', full)
	42	+ with T.serial(kb) as row:
	43	+ scale = S.vld(s_ub[row], dist='BRC_B32')
	44	+ for chunk in T.unroll(chunks, explicit=True):
	45	+ values = S.vcvt(S.vld(k8[row, chunk * 64], dist='UNPK4_B8'), T.float32)
	46	+ sums[chunk] = S.vadd(sums[chunk], S.vmul(values, scale, full), full)
	47	+ maximum = S.alloc_local(1, 'float32')
	48	+ maximum[0] = S.vdup(T.float32(0), 'float32', full)
	49	+ for chunk in T.unroll(chunks, explicit=True):
	50	+ sums[chunk] = S.vmuls(sums[chunk], T.float32(1.0 / kb), full)
	51	+ maximum[0] = S.vmax(maximum[0], S.vabs(sums[chunk], full), full)
	52	+ max_value = S.vdupv(S.vcmax(maximum[0], full), full)
	53	+ sf = S.vmax(
	54	+ S.vmuls(max_value, T.float32(1.0 / FP8_MAX), full), S.vdup(T.float32(EPS), 'float32', full), full
	55	+ )
	56	+ inv = S.vdiv(S.vdup(T.float32(1.0), 'float32', full), sf, full)
	57	+ for chunk in T.unroll(chunks, explicit=True):
	58	+ q = S.vcvt(S.vmul(sums[chunk], inv, full), T.float8_e4m3fn)
	59	+ S.vsts(q8[chunk * 64], q, dist='PK4_B32')
	60	+ S.vsts(s_ub[0], sf, one, 'ONEPT_B32')
	61	+
	62	+ @T.prim_func
	63	+ def main(
	64	+ input_keys: T.Tensor((num_tokens, head_dim), T.float8_e4m3fn),
	65	+ key_scales: T.Tensor((num_tokens,), 'float32'),
	66	+ pooled_keys: T.Tensor((num_blocks, head_dim), T.float8_e4m3fn),
	67	+ pooled_scales: T.Tensor((num_blocks,), 'float32'),
	68	+ ):
	69	+ with T.Kernel(num_blocks) as bx:
	70	+ k8 = T.alloc_shared((kb, padded_d), T.float8_e4m3fn)
	71	+ s_ub = T.alloc_shared((kb,), 'float32')
	72	+ q8 = T.alloc_shared((padded_d,), T.float8_e4m3fn)
	73	+
	74	+ if padded_d != head_dim:
	75	+ clear_tail_lanes(k8)
	76	+ next_block = T.meta_var(bx + 1)
	77	+ block_start = T.meta_var(bx * kb)
	78	+ block_end = T.meta_var(next_block * kb)
	79	+ T.copy(input_keys[block_start:block_end, 0:head_dim], k8[0:kb, 0:head_dim])
	80	+ T.copy(key_scales[block_start:block_end], s_ub)
	81	+ pool_block(k8, s_ub, q8)
	82	+ T.copy(s_ub[0:1], pooled_scales[bx:next_block])
	83	+ T.copy(q8[0:head_dim], pooled_keys[bx, 0:head_dim])
	84	+
	85	+ return main
	86	+
	87	+
	88	+def fp8_native_block_mean_pooling_interface(k, k_scale, k_block_size):
	89	+ num_tokens, head_dim = k.shape
	90	+ kern = cached_kernel(
	91	+ ('hisa-mean-pool', num_tokens, head_dim, k_block_size),
	92	+ lambda: fp8_native_block_mean_pooling(num_tokens, head_dim, k_block_size),
	93	+ [2, 3],
	94	+ )
	95	+ return kern(k, k_scale)

Asrc/cann_ops_tilelang/attention/dsa_hisa/hisa.py+242-0

@@ -0,0 +1,242 @@
	1	+# Copyright (c) Tile-AI.
	2	+# Copyright (c) 2026 ops-tilelang contributors
	3	+# SPDX-License-Identifier: MIT
	4	+# See LICENSE in the repository root for the full license text.
	5	+
	6	+"""Hierarchical sparse-attention indexing on Ascend using TileLang kernels."""
	7	+
	8	+from numbers import Integral
	9	+
	10	+import torch
	11	+
	12	+from .block_sparse_mqa_fp8 import fp8_native_block_sparse_mqa_attn_return_logits_interface
	13	+from .block_sparse_mqa_fp8 import fused_ok as token_score_tile_fits
	14	+from .block_topk import block_topk_select_interface
	15	+from .fp8_block_mean_pooling import fp8_native_block_mean_pooling_interface
	16	+from .pool_mqa_fp8 import (
	17	+ fused_bt_ok,
	18	+ pool_mqa_attn_return_logits_fp8_interface,
	19	+ pool_mqa_logits_block_topk_interface,
	20	+)
	21	+from .pool_mqa_fp8 import route as pool_score_route
	22	+from .tilelang_utils import EAGER_WORKSPACE
	23	+from .token_topk import token_topk_with_offset_interface
	24	+
	25	+_SHARED_BYTES = 184 * 1024
	26	+
	27	+
	28	+def _validate_tile_dimensions(dim, tokens, kb, blocks, candidates):
	29	+ if dim % 32 or kb % 32:
	30	+ raise ValueError('D and k_block_size must be multiples of 32 for FP8 tiles')
	31	+ if blocks < 64 and 64 % blocks:
	32	+ raise ValueError('a nonempty K with fewer than 64 blocks must have a block count that divides 64')
	33	+ if blocks > 1024:
	34	+ raise ValueError('block top-k supports at most 1024 input blocks per row')
	35	+ if candidates % 64 or candidates > 8192:
	36	+ raise ValueError('effective candidate count must be a multiple of 64 and at most 8192')
	37	+ if tokens > 2**24:
	38	+ raise ValueError('N exceeds the exact FP32 range for token position comparisons')
	39	+ # Pooling keeps the key tile, FP32 scales and requantized result in shared memory.
	40	+ padded_d = (dim + 63) // 64 * 64
	41	+ if kb * padded_d + 4 * kb + padded_d > _SHARED_BYTES:
	42	+ raise ValueError('the FP8 block pooling tile exceeds the supported shared-memory budget')
	43	+ if blocks < 64 and 64 * (dim + 4) > _SHARED_BYTES:
	44	+ raise ValueError('padding the FP8 pooled keys and scales to 64 blocks exceeds the shared-memory budget')
	45	+
	46	+
	47	+def _validate_backend_tiles(rows, heads, dim, kb, blocks, block_topk):
	48	+ if not token_score_tile_fits(heads, dim, kb, block_topk):
	49	+ raise ValueError('the FP8 operands or FP32 accumulator exceed the token scoring tile capacity')
	50	+ padded_blocks = max(64, blocks)
	51	+ route = pool_score_route(rows, heads, dim, padded_blocks)
	52	+ if route == 'vec':
	53	+ vec_bytes = 5 * (heads + padded_blocks) * max(64, dim) + 4 * (
	54	+ heads * padded_blocks + 2 * heads + 2 * padded_blocks
	55	+ )
	56	+ if vec_bytes > _SHARED_BYTES:
	57	+ raise ValueError('the vector block scoring tile exceeds the supported shared-memory budget')
	58	+ if route == 'gemm2' and (16 + padded_blocks) * dim > 256 * 1024:
	59	+ raise ValueError('the block scoring matrices exceed the supported L1 tile capacity')
	60	+
	61	+
	62	+def _q_scale_ones(rows, heads, device, workspace):
	63	+ key = ('q_scale', rows, heads, device)
	64	+ scale = workspace.get(key)
	65	+ if scale is None:
	66	+ scale = torch.ones((rows, heads), dtype=torch.float32, device=device)
	67	+ workspace[key] = scale
	68	+ return scale
	69	+
	70	+
	71	+def _validate_hisa_tensors(q, k, k_scale, weights, cu_seqlen_ks, cu_seqlen_ke):
	72	+ tensors = (q, k, k_scale, weights, cu_seqlen_ks, cu_seqlen_ke)
	73	+ if any(not isinstance(t, torch.Tensor) for t in tensors):
	74	+ raise TypeError('HISA expects tensor inputs')
	75	+ if q.device.type != 'npu' or any(t.device != q.device or not t.is_contiguous() for t in tensors):
	76	+ raise ValueError('all inputs must be contiguous and on the same NPU')
	77	+ if q.dtype != torch.float8_e4m3fn or k.dtype != q.dtype:
	78	+ raise TypeError('q and k must be float8_e4m3fn')
	79	+ if weights.dtype != torch.float32 or k_scale.dtype != torch.float32:
	80	+ raise TypeError('weights and k_scale must be float32')
	81	+ if cu_seqlen_ks.dtype != torch.int32 or cu_seqlen_ke.dtype != torch.int32:
	82	+ raise TypeError('query window bounds must be int32')
	83	+
	84	+
	85	+def _validate_hisa_inputs(q, k, k_scale, weights, cu_seqlen_ks, cu_seqlen_ke, k_block_size, block_topk, topk_tokens):
	86	+ _validate_hisa_tensors(q, k, k_scale, weights, cu_seqlen_ks, cu_seqlen_ke)
	87	+ if q.ndim != 3 or k.ndim != 2:
	88	+ raise ValueError('expected q[M, H, D] and k[N, D]')
	89	+ rows, heads, dim = q.shape
	90	+ tokens = k.shape[0]
	91	+ if min(heads, dim) < 1 or k.shape[1] != dim:
	92	+ raise ValueError('incompatible Q, K, scale or weight shapes')
	93	+ if weights.shape != (rows, heads) or k_scale.shape != (tokens,):
	94	+ raise ValueError('incompatible Q, K, scale or weight shapes')
	95	+ if cu_seqlen_ks.shape != (rows,) or cu_seqlen_ke.shape != (rows,):
	96	+ raise ValueError('both query window bounds must have shape [M]')
	97	+ if min(k_block_size, block_topk, topk_tokens) < 1 or tokens % k_block_size:
	98	+ raise ValueError('selection sizes must be positive and N must be divisible by k_block_size')
	99	+ blocks = k.shape[0] // k_block_size
	100	+ block_topk_eff = min(block_topk, blocks)
	101	+ topk_tokens_eff = min(topk_tokens, block_topk_eff * k_block_size)
	102	+ if rows and tokens:
	103	+ _validate_tile_dimensions(dim, tokens, k_block_size, blocks, block_topk_eff * k_block_size)
	104	+ return block_topk_eff, topk_tokens_eff
	105	+
	106	+
	107	+def _select_blocks(q, q_scale, k, k_scale, weights, cu_seqlen_ks, cu_seqlen_ke, k_block_size, block_topk, workspace):
	108	+ rows, heads, dim = q.shape
	109	+ blocks = k.shape[0] // k_block_size
	110	+ # Stage 0: pool K into FP8 blocks with one FP32 scale per block.
	111	+ blocked_k_fp8, blocked_k_scale = fp8_native_block_mean_pooling_interface(k, k_scale, k_block_size)
	112	+
	113	+ # Stage 1: score pooled blocks, apply causal boundaries, then select blocks.
	114	+ # Fuse scoring and selection when the Cube/Vector tile geometry permits it.
	115	+ if fused_bt_ok(rows, heads, dim, blocks, k_block_size, block_topk):
	116	+ return pool_mqa_logits_block_topk_interface(
	117	+ q,
	118	+ q_scale,
	119	+ blocked_k_fp8,
	120	+ blocked_k_scale,
	121	+ weights,
	122	+ cu_seqlen_ks,
	123	+ cu_seqlen_ke,
	124	+ k_block_size,
	125	+ block_topk,
	126	+ workspace=workspace,
	127	+ )
	128	+ block_k_score = pool_mqa_attn_return_logits_fp8_interface(
	129	+ q, q_scale, blocked_k_fp8, blocked_k_scale, weights, workspace=workspace
	130	+ )
	131	+ return block_topk_select_interface(
	132	+ block_k_score,
	133	+ cu_seqlen_ks,
	134	+ cu_seqlen_ke,
	135	+ k_block_size,
	136	+ block_topk,
	137	+ )
	138	+
	139	+
	140	+def _hisa_indexer_impl(
	141	+ q, k, k_scale, weights, cu_seqlen_ks, cu_seqlen_ke, k_block_size, block_topk, topk_tokens, workspace
	142	+):
	143	+ block_topk_eff, topk_tokens_eff = _validate_hisa_inputs(
	144	+ q, k, k_scale, weights, cu_seqlen_ks, cu_seqlen_ke, k_block_size, block_topk, topk_tokens
	145	+ )
	146	+ rows, heads, dim = q.shape
	147	+ tokens = k.shape[0]
	148	+ if rows == 0 or tokens == 0:
	149	+ return torch.empty((rows, topk_tokens_eff), device=q.device, dtype=torch.int32)
	150	+ blocks = tokens // k_block_size
	151	+ # Tile feasibility depends on the actual backend and is checked before launch.
	152	+ _validate_backend_tiles(rows, heads, dim, k_block_size, blocks, block_topk_eff)
	153	+ q_scale = _q_scale_ones(rows, heads, q.device, workspace)
	154	+
	155	+ topk_block_indices = _select_blocks(
	156	+ q, q_scale, k, k_scale, weights, cu_seqlen_ks, cu_seqlen_ke, k_block_size, block_topk_eff, workspace
	157	+ )
	158	+
	159	+ # Stage 2: compute FP32 token scores only within the selected blocks.
	160	+ block_sparse_logits = fp8_native_block_sparse_mqa_attn_return_logits_interface(
	161	+ q,
	162	+ q_scale,
	163	+ k,
	164	+ k_scale,
	165	+ weights,
	166	+ topk_block_indices,
	167	+ (cu_seqlen_ks, cu_seqlen_ke),
	168	+ k_block_size,
	169	+ workspace=workspace,
	170	+ )
	171	+
	172	+ # Stage 3: select tokens and convert candidate slots to relative K offsets.
	173	+ return token_topk_with_offset_interface(
	174	+ block_sparse_logits,
	175	+ topk_block_indices,
	176	+ cu_seqlen_ks,
	177	+ cu_seqlen_ke,
	178	+ k_block_size,
	179	+ topk_tokens_eff,
	180	+ )
	181	+
	182	+
	183	+def hisa_indexer(
	184	+ q: torch.Tensor, # [M, H, D] FP8 E4M3
	185	+ k: torch.Tensor, # [N, D] FP8 E4M3
	186	+ k_scale: torch.Tensor, # [N] FP32
	187	+ weights: torch.Tensor, # [M, H] FP32
	188	+ cu_seqlen_ks: torch.Tensor, # [M] INT32, inclusive K start
	189	+ cu_seqlen_ke: torch.Tensor, # [M] INT32, exclusive K end
	190	+ *,
	191	+ k_block_size: int,
	192	+ block_topk: int,
	193	+ topk_tokens: int,
	194	+) -> torch.Tensor:
	195	+ """Return selected K offsets relative to each query's K start.
	196	+
	197	+ The output is INT32 with shape [M, min(topk_tokens, effective_candidates)].
	198	+ Invalid selected positions become -1. Output order is unspecified; equal
	199	+ token scores choose the smaller candidate index. Block selection rounds
	200	+ scores to BF16, while token selection uses the full FP32 scores.
	201	+ """
	202	+ sizes = (k_block_size, block_topk, topk_tokens)
	203	+ if any(not isinstance(size, Integral) or isinstance(size, bool) for size in sizes):
	204	+ raise TypeError('k_block_size, block_topk and topk_tokens must be integers')
	205	+ k_block_size, block_topk, topk_tokens = map(int, sizes)
	206	+ if torch.compiler.is_compiling():
	207	+ return _compiled_hisa_indexer(
	208	+ q,
	209	+ k,
	210	+ k_scale,
	211	+ weights,
	212	+ cu_seqlen_ks,
	213	+ cu_seqlen_ke,
	214	+ k_block_size=k_block_size,
	215	+ block_topk=block_topk,
	216	+ topk_tokens=topk_tokens,
	217	+ )
	218	+ return _hisa_indexer_impl(
	219	+ q, k, k_scale, weights, cu_seqlen_ks, cu_seqlen_ke, k_block_size, block_topk, topk_tokens, EAGER_WORKSPACE
	220	+ )
	221	+
	222	+
	223	+@torch.library.custom_op(
	224	+ 'cann_ops_tilelang::hisa_indexer',
	225	+ mutates_args=(),
	226	+ schema=torch.library.infer_schema(hisa_indexer, mutates_args=()),
	227	+)
	228	+def _compiled_hisa_indexer(
	229	+ q, k, k_scale, weights, cu_seqlen_ks, cu_seqlen_ke, *, k_block_size, block_topk, topk_tokens
	230	+):
	231	+ # Graph-owned tensor storage must not escape capture into a process-wide cache.
	232	+ return _hisa_indexer_impl(
	233	+ q, k, k_scale, weights, cu_seqlen_ks, cu_seqlen_ke, k_block_size, block_topk, topk_tokens, workspace={}
	234	+ )
	235	+
	236	+
	237	+@_compiled_hisa_indexer.register_fake
	238	+def _hisa_indexer_fake(q, k, k_scale, weights, cu_seqlen_ks, cu_seqlen_ke, *, k_block_size, block_topk, topk_tokens):
	239	+ _, width = _validate_hisa_inputs(
	240	+ q, k, k_scale, weights, cu_seqlen_ks, cu_seqlen_ke, k_block_size, block_topk, topk_tokens
	241	+ )
	242	+ return torch.empty((q.shape[0], width), dtype=torch.int32, device=q.device)

Asrc/cann_ops_tilelang/attention/dsa_hisa/pool_mqa_fp8.py+537-0

@@ -0,0 +1,537 @@
	1	+# Copyright (c) 2026 ops-tilelang contributors
	2	+# SPDX-License-Identifier: MIT
	3	+# See LICENSE in the repository root for the full license text.
	4	+
	5	+"""Pooled-key MQA scoring and optional fused BF16 block selection."""
	6	+
	7	+# Reuse score buffers on one NPU. Consume each result before the next call with the same shape.
	8	+
	9	+from math import gcd
	10	+
	11	+import tilelang
	12	+import tilelang.ascend.language as T
	13	+import torch
	14	+from tilelang.ascend.language import simd as S
	15	+
	16	+from .._simd import is_pto_target
	17	+from .tilelang_utils import EAGER_WORKSPACE, cached_kernel
	18	+
	19	+PAD_H = 16
	20	+
	21	+
	22	+def _lcm(a: int, b: int) -> int:
	23	+ return a // gcd(a, b) * b
	24	+
	25	+
	26	+def fused_block_q(num_heads: int) -> int:
	27	+ return _lcm(16 // gcd(num_heads, 16), 2)
	28	+
	29	+
	30	+def fused_ok(num_rows: int, num_heads: int, head_dim: int, num_blocks: int, bq: int \| None = None) -> bool:
	31	+ bq = fused_block_q(num_heads) if bq is None else bq
	32	+ # Fused reduction consumes complete 64-score vectors. The separate scorer also handles tails.
	33	+ if num_heads <= 0 or head_dim % 16 != 0 or num_blocks % 64 != 0:
	34	+ return False
	35	+ if (bq * num_heads) % 16 != 0 or bq % 2 != 0:
	36	+ return False
	37	+ if num_rows % bq != 0:
	38	+ return False
	39	+ if is_pto_target() and bq * num_heads < num_blocks:
	40	+ return False
	41	+
	42	+ if bq * num_heads * head_dim + num_blocks * head_dim > 256 * 1024:
	43	+ return False
	44	+
	45	+ return bq * num_heads * num_blocks * 4 <= 64 * 1024
	46	+
	47	+
	48	+def _block_score_reduction(half_queries, num_heads, num_blocks):
	49	+ @T.macro
	50	+ def accumulate_head(row, head, s_ub, qs_ub, ks_ub, w_ub, acc, zero, full):
	51	+ q_scale = S.vld(qs_ub[row, head], 'BRC_B32')
	52	+ weight = S.vld(w_ub[row, head], 'BRC_B32')
	53	+ for col in T.unroll(num_blocks // 64, explicit=True):
	54	+ k_scale = S.vld(ks_ub[col * 64], 'NORM')
	55	+ score = S.vld(s_ub[row * num_heads + head, col * 64], 'NORM')
	56	+ scaled = S.vmul(score, S.vmul(q_scale, k_scale, full), full)
	57	+ weighted = S.vmul(S.vmax(scaled, zero, full), weight, full)
	58	+ acc[col] = S.vadd(acc[col], weighted, full)
	59	+
	60	+ @T.macro
	61	+ def reduce_scores(s_ub, qs_ub, ks_ub, w_ub, o_ub):
	62	+ with T.SimdVF():
	63	+ full = S.pset(32, 'PAT_ALL')
	64	+ zero = S.vdup(T.float32(0), 'float32', full)
	65	+ acc = S.alloc_local(num_blocks // 64, 'float32')
	66	+ for row in T.serial(half_queries):
	67	+ for col in T.unroll(num_blocks // 64, explicit=True):
	68	+ acc[col] = S.vdup(T.float32(0), 'float32', full)
	69	+ for head in T.serial(num_heads):
	70	+ accumulate_head(row, head, s_ub, qs_ub, ks_ub, w_ub, acc, zero, full)
	71	+ for col in T.unroll(num_blocks // 64, explicit=True):
	72	+ S.vsts(o_ub[row, col * 64], acc[col], full, 'NORM_B32')
	73	+
	74	+ return reduce_scores
	75	+
	76	+
	77	+def _alloc_pooled_gemm_tiles(tile_rows, head_dim, num_blocks):
	78	+ return (
	79	+ T.alloc_l1((tile_rows, head_dim), T.float8_e4m3fn),
	80	+ T.alloc_l1((num_blocks, head_dim), T.float8_e4m3fn),
	81	+ T.alloc_l0c((tile_rows, num_blocks), 'float32'),
	82	+ )
	83	+
	84	+
	85	+def _pooled_input_loader(bq, num_heads, head_dim, num_blocks):
	86	+ @T.macro
	87	+ def load_inputs(m0, query, pooled_keys, pooled_scales, q_l1, kp_l1, ks_ub):
	88	+ for i in range(bq):
	89	+ head_tile_start = T.meta_var(i * num_heads)
	90	+ head_tile_end = T.meta_var((i + 1) * num_heads)
	91	+ T.copy(query[m0 + i, 0:num_heads, 0:head_dim], q_l1[head_tile_start:head_tile_end, 0:head_dim])
	92	+ T.copy(pooled_keys, kp_l1)
	93	+ T.copy(pooled_scales[0:num_blocks], ks_ub)
	94	+
	95	+ return load_inputs
	96	+
	97	+
	98	+def pool_mqa_fused(num_rows: int, num_heads: int, head_dim: int, num_blocks: int, bq: int \| None = None):
	99	+ bq = fused_block_q(num_heads) if bq is None else bq
	100	+ if not fused_ok(num_rows, num_heads, head_dim, num_blocks, bq):

Asrc/cann_ops_tilelang/attention/dsa_hisa/tilelang_utils.py+106-0

@@ -0,0 +1,106 @@
	1	+# Copyright (c) 2026 ops-tilelang contributors
	2	+# SPDX-License-Identifier: MIT
	3	+# See LICENSE in the repository root for the full license text.
	4	+
	5	+"""Kernel compilation cache and shared vector helpers for top-k selection."""
	6	+
	7	+import tilelang
	8	+import tilelang.ascend.language as T
	9	+from tilelang.ascend.language import simd as S
	10	+
	11	+
	12	+def count_selectable(logits, accum, nv, full, ones, zeros, ninf):
	13	+ """Count scores > -inf and == -inf; their sum excludes every NaN."""
	14	+ width = 8 if nv % 8 == 0 else 1
	15	+ T.evaluate(S.vsts(accum[0, 0], zeros, full, 'NORM_B32'))
	16	+ T.evaluate(S.vsts(accum[1, 0], zeros, full, 'NORM_B32'))
	17	+ T.evaluate(S.mem_bar('VST_VLD'))
	18	+ with T.serial(nv // width) as tile:
	19	+ greater = S.vld(accum[0, 0])
	20	+ equal = S.vld(accum[1, 0])
	21	+ for j in range(width):
	22	+ value = S.vld(logits[(tile * width + j) * 64])
	23	+ greater = S.vadd(greater, S.vsel(ones, zeros, S.vcmp(value, ninf, full, 'gt')), full)
	24	+ equal = S.vadd(equal, S.vsel(ones, zeros, S.vcmp(value, ninf, full, 'eq')), full)
	25	+ T.evaluate(S.vsts(accum[0, 0], greater, full, 'NORM_B32'))
	26	+ T.evaluate(S.vsts(accum[1, 0], equal, full, 'NORM_B32'))
	27	+ T.evaluate(S.mem_bar('VST_VLD'))
	28	+ return S.vcadd(S.vld(accum[0, 0]), full), S.vcadd(S.vld(accum[1, 0]), full)
	29	+
	30	+
	31	+def window_score_classifier(topk, nv):
	32	+ @T.macro
	33	+ def classify_window(logits, counts, accum, span):
	34	+ # The window is a cheap upper-bound hint; scores still decide eligibility.
	35	+ # Longer windows retain the original search without classification overhead.
	36	+ counts[0] = topk + 1
	37	+ counts[1] = T.int32(0)
	38	+ if span <= topk:
	39	+ with T.SimdVF():
	40	+ full = S.pset(32, 'PAT_ALL')
	41	+ one = S.pset(32, 'PAT_VL1')
	42	+ ones = S.vdup(T.int32(1), 'int32', full)
	43	+ zeros = S.vdup(T.int32(0), 'int32', full)
	44	+ ninf = S.vdup(-T.infinity(T.float32), 'float32', full)
	45	+ greater, equal = count_selectable(logits, accum, nv, full, ones, zeros, ninf)
	46	+ S.vsts(counts[0], greater, one, 'ONEPT_B32')
	47	+ S.vsts(counts[1], equal, one, 'ONEPT_B32')
	48	+
	49	+ return classify_window
	50	+
	51	+
	52	+VS = 64
	53	+
	54	+
	55	+def ramp(i):
	56	+ return S.vci(T.int32(i * VS), 'int32')
	57	+
	58	+
	59	+def bf16_block_rank(score, full):
	60	+ """Round FP32 to nearest-even BF16 and order its signed high word; both zeros rank 0."""
	61	+ bits = T.reinterpret(score, 'int32x64')
	62	+ high = S.vshrs(bits, T.int32(16), full)
	63	+ one = S.vdup(T.int32(1), 'int32', full)
	64	+ bias = S.vadds(S.vand(high, one, full), T.int32(0x7FFF), full)
	65	+ rounded = S.vshrs(S.vadd(bits, bias, full), T.int32(16), full)
	66	+ negative = S.vcmps(rounded, T.int32(0), full, 'lt')
	67	+ reverse = S.vsub(S.vdup(T.int32(-32768), 'int32', full), rounded, full)
	68	+ return S.vsel(reverse, rounded, negative)
	69	+
	70	+
	71	+def block_priority(rank, indices, first, end, full):
	72	+ # Real BF16 ranks fit in [-32767, 32767]. Reserve one rank on either side.
	73	+ edge = S.por(S.vcmps(indices, first, full, 'eq'), S.vcmps(indices, end - 1, full, 'eq'), full)
	74	+ outside = S.por(S.vcmps(indices, first, full, 'lt'), S.vcmps(indices, end, full, 'ge'), full)
	75	+ prioritized = S.vsel(S.vdup(T.int32(32768), 'int32', full), rank, edge)
	76	+ return S.vsel(S.vdup(T.int32(-32768), 'int32', full), prioritized, outside)
	77	+
	78	+
	79	+def redmax(v, nv, full):
	80	+ m = v[0]
	81	+ for i in range(1, nv):
	82	+ m = S.vmax(m, v[i], full)
	83	+ return m
	84	+
	85	+
	86	+def argmin(v, nv, mx, full, imax):
	87	+ m = S.vsel(ramp(0), imax, S.vcmp(v[0], mx, full, 'eq'))
	88	+ for i in range(1, nv):
	89	+ m = S.vmin(m, S.vsel(ramp(i), imax, S.vcmp(v[i], mx, full, 'eq')), full)
	90	+ return m
	91	+
	92	+
	93	+# Kernel handles can be shared across executions. Eager scratch tensors are
	94	+# reused on one stream; graph capture supplies its own temporary workspace.
	95	+EAGER_WORKSPACE = {}
	96	+_KERNEL_CACHE = {}
	97	+
	98	+
	99	+def cached_kernel(key, build_fn, out_idx):
	100	+ """Compile each implementation and shape once per configured target."""
	101	+ key = (tilelang.env.get_default_target(), key)
	102	+ kernel = _KERNEL_CACHE.get(key)
	103	+ if kernel is None:
	104	+ kernel = tilelang.compile(build_fn(), out_idx=out_idx)
	105	+ _KERNEL_CACHE[key] = kernel
	106	+ return kernel

Asrc/cann_ops_tilelang/attention/dsa_hisa/token_topk.py+376-0

@@ -0,0 +1,376 @@
	1	+# Copyright (c) 2026 ops-tilelang contributors
	2	+# SPDX-License-Identifier: MIT
	3	+# See LICENSE in the repository root for the full license text.
	4	+
	5	+"""Select exact FP32 token scores, breaking ties by candidate index, and return relative offsets."""
	6	+
	7	+import tilelang.ascend.language as T
	8	+import torch
	9	+from tilelang.ascend.language import simd as S
	10	+from tvm import tirx
	11	+
	12	+from .tilelang_utils import argmin, cached_kernel, ramp, redmax, window_score_classifier
	13	+
	14	+VS = 64
	15	+
	16	+
	17	+NTHR = 15
	18	+
	19	+
	20	+S_A, S_R, S_RS, S_JS, S_B16, S_BLO, S_R2, S_BSZ, S_W, S_T = range(10)
	21	+
	22	+
	23	+def _key32_of(l_ub, i, full, zf, zi, minv):
	24	+ x = S.vadd(S.vld(l_ub[i * VS]), zf, full)
	25	+ b = tirx.reinterpret('int32x64', x)
	26	+ neg = S.vcmp(b, zi, full, 'lt')
	27	+ return S.vsel(S.vnot(b, full), S.vor(b, minv, full), neg)
	28	+
	29	+
	30	+def _cnt_ge(k_ub, thr, nv, full, oi, zi, acc_ub):
	31	+ width = 8 if nv % 8 == 0 else 1
	32	+ T.evaluate(S.vsts(acc_ub[0], zi, full, 'NORM_B32'))
	33	+ T.evaluate(S.mem_bar('VST_VLD'))
	34	+ with T.serial(nv // width) as tile:
	35	+ acc = S.vld(acc_ub[0])
	36	+ for j in range(width):
	37	+ value = S.vld(k_ub[(tile * width + j) * VS])
	38	+ mask = S.vcmps(value, thr, full, 'ge')
	39	+ acc = S.vadd(acc, S.vsel(oi, zi, mask), full)
	40	+ T.evaluate(S.vsts(acc_ub[0], acc, full, 'NORM_B32'))
	41	+ T.evaluate(S.mem_bar('VST_VLD'))
	42	+ return S.vcadd(S.vld(acc_ub[0]), full)
	43	+
	44	+
	45	+def _cnt_ge_masked(k_ub, lo_ub, thr, b16v, nv, full, oi, zi, acc_ub):
	46	+ width = 8 if nv % 8 == 0 else 1
	47	+ T.evaluate(S.vsts(acc_ub[0], zi, full, 'NORM_B32'))
	48	+ T.evaluate(S.mem_bar('VST_VLD'))
	49	+ with T.serial(nv // width) as tile:
	50	+ acc = S.vld(acc_ub[0])
	51	+ for j in range(width):
	52	+ value = S.vld(k_ub[(tile * width + j) * VS])
	53	+ low = S.vld(lo_ub[(tile * width + j) * VS])
	54	+ same = S.vcmp(value, b16v, full, 'eq')
	55	+ ge = S.vcmps(low, thr, full, 'ge')
	56	+ mask = S.pand(same, ge, full)
	57	+ acc = S.vadd(acc, S.vsel(oi, zi, mask), full)
	58	+ T.evaluate(S.vsts(acc_ub[0], acc, full, 'NORM_B32'))
	59	+ T.evaluate(S.mem_bar('VST_VLD'))
	60	+ return S.vcadd(S.vld(acc_ub[0]), full)
	61	+
	62	+
	63	+@T.macro
	64	+def _advance_digit(st_ub, cnt_ub, threshold_word, rank_slot):
	65	+ cnt_ub[16] = st_ub[S_A]
	66	+ st_ub[S_RS] = st_ub[S_A] + st_ub[rank_slot]
	67	+ st_ub[S_JS] = 0
	68	+ for bucket in T.serial(1, 16):
	69	+ if cnt_ub[bucket] >= st_ub[S_RS]:
	70	+ st_ub[S_JS] = bucket
	71	+ st_ub[S_A] = cnt_ub[st_ub[S_JS] + 1]
	72	+ st_ub[rank_slot] = st_ub[S_RS] - cnt_ub[st_ub[S_JS] + 1]
	73	+ st_ub[threshold_word] = st_ub[threshold_word] * 16 + st_ub[S_JS]
	74	+
	75	+
	76	+def token_topk_with_offset(
	77	+ num_rows: int,
	78	+ num_candidates: int,
	79	+ topk: int,
	80	+ kb: int,
	81	+ block_topk: int,
	82	+):
	83	+ if num_candidates % VS:
	84	+ raise AssertionError(f'NC={num_candidates} 必须是 {VS} 的整数倍')
	85	+ if num_rows <= 0 or not 0 < topk <= num_candidates:
	86	+ raise AssertionError('expected a positive row count and top-k within the candidate count')
	87	+ if kb <= 0 or kb * block_topk != num_candidates:
	88	+ raise AssertionError('the candidate count must equal block_topk times a positive block size')
	89	+ nv = num_candidates // VS
	90	+ topk_pad = ((topk + 127) // 128) * 128
	91	+
	92	+ @T.macro
	93	+ def split_score_keys(l_ub, k16_ub, lo16_ub):
	94	+ with T.SimdVF():
	95	+ full = S.pset(32, 'PAT_ALL')
	96	+ zf = S.vdup(T.float32(0.0), 'float32', full)
	97	+ zi = S.vdup(T.int32(0), 'int32', full)
	98	+ minv = S.vdup(T.int32(-2147483648), 'int32', full)
	99	+ m16 = S.vdup(T.int32(0xFFFF), 'int32', full)
	100	+ for i in range(nv):
	101	+ k32 = _key32_of(l_ub, i, full, zf, zi, minv)
	102	+ S.vsts(k16_ub[i * VS], S.vand(S.vshrs(k32, T.int32(16), full), m16, full), full, 'NORM_B32')
	103	+ S.vsts(lo16_ub[i * VS], S.vand(k32, m16, full), full, 'NORM_B32')
	104	+
	105	+ @T.macro
	106	+ def search_high_word(k16_ub, st_ub, sh_ub, thr_ub, cnt_ub, acc_ub):
	107	+ for p in T.unroll(4):
	108	+ for j in T.unroll(16):
	109	+ thr_ub[j] = (st_ub[S_B16] * 16 + j) * sh_ub[p]
	110	+ with T.SimdVF():
	111	+ full = S.pset(32, 'PAT_ALL')
	112	+ one = S.pset(32, 'PAT_VL1')
	113	+ oi = S.vdup(T.int32(1), 'int32', full)
	114	+ zi = S.vdup(T.int32(0), 'int32', full)
	115	+ with T.serial(NTHR) as jj:
	116	+ S.vsts(
	117	+ cnt_ub[jj + 1],
	118	+ _cnt_ge(k16_ub, thr_ub[jj + 1], nv, full, oi, zi, acc_ub),
	119	+ one,
	120	+ 'ONEPT_B32',
	121	+ )
	122	+
	123	+ _advance_digit(st_ub, cnt_ub, S_B16, S_R)
	124	+
	125	+ @T.macro
	126	+ def search_low_word(k16_ub, lo16_ub, st_ub, sh_ub, thr_ub, cnt_ub, acc_ub):
	127	+ for p in T.unroll(4):
	128	+ for j in T.unroll(16):
	129	+ thr_ub[j] = (st_ub[S_BLO] * 16 + j) * sh_ub[p]
	130	+ with T.SimdVF():
	131	+ full = S.pset(32, 'PAT_ALL')
	132	+ one = S.pset(32, 'PAT_VL1')
	133	+ oi = S.vdup(T.int32(1), 'int32', full)
	134	+ zi = S.vdup(T.int32(0), 'int32', full)
	135	+ b16v = S.vdup(st_ub[S_B16], 'int32', full)
	136	+ with T.serial(NTHR) as jj:
	137	+ S.vsts(
	138	+ cnt_ub[jj + 1],
	139	+ _cnt_ge_masked(k16_ub, lo16_ub, thr_ub[jj + 1], b16v, nv, full, oi, zi, acc_ub),
	140	+ one,
	141	+ 'ONEPT_B32',
	142	+ )
	143	+
	144	+ _advance_digit(st_ub, cnt_ub, S_BLO, S_R2)
	145	+
	146	+ @T.macro
	147	+ def search_boundary(k16_ub, lo16_ub, st_ub, sh_ub, thr_ub, cnt_ub, acc_ub):
	148	+ st_ub[S_A] = T.int32(0)
	149	+ st_ub[S_R] = T.int32(topk)
	150	+ st_ub[S_B16] = T.int32(0)
	151	+ st_ub[S_BLO] = T.int32(-1)
	152	+ st_ub[S_R2] = T.int32(topk)
	153	+
	154	+ sh_ub[0] = T.int32(4096)
	155	+ sh_ub[1] = T.int32(256)
	156	+ sh_ub[2] = T.int32(16)
	157	+ sh_ub[3] = T.int32(1)
	158	+
	159	+ search_high_word(k16_ub, st_ub, sh_ub, thr_ub, cnt_ub, acc_ub)
	160	+
	161	+ st_ub[S_R2] = st_ub[S_R]
	162	+
	163	+ st_ub[S_BLO] = T.int32(0)
	164	+ st_ub[S_A] = T.int32(0)
	165	+ search_low_word(k16_ub, lo16_ub, st_ub, sh_ub, thr_ub, cnt_ub, acc_ub)
	166	+
	167	+ @T.macro
	168	+ def collect_selected(k16_ub, lo16_ub, st_ub, sel_ub, out_off_ub, tie_off_ub):
	169	+ with T.SimdVF():
	170	+ full = S.pset(32, 'PAT_ALL')
	171	+ zero = S.vdup(T.int32(0), 'int32', full)
	172	+ for vector in T.serial(topk_pad // VS):
	173	+ S.vsts(sel_ub[vector * VS], zero, full, 'NORM_B32')
	174	+ S.vsts(out_off_ub[0], zero, full, 'NORM_B32')
	175	+ S.vsts(tie_off_ub[0], zero, full, 'NORM_B32')
	176	+ with T.SimdVF():
	177	+ full = S.pset(32, 'PAT_ALL')
	178	+ onev = S.vdup(T.int32(1), 'int32', full)
	179	+ zerov = S.vdup(T.int32(0), 'int32', full)
	180	+ boundary = S.vdup(st_ub[S_B16], 'int32', full)
	181	+ low_boundary = S.vdup(st_ub[S_BLO], 'int32', full)
	182	+ quota = S.vdup(st_ub[S_R2], 'int32', full)
	183	+ for vi in range(nv):
	184	+ key = S.vld(k16_ub[vi * VS])
	185	+ low = S.vld(lo16_ub[vi * VS])
	186	+ off = S.vld(out_off_ub[0])
	187	+ tie_off = S.vld(tie_off_ub[0])
	188	+ same = S.vcmp(key, boundary, full, 'eq')
	189	+ greater_key = S.vcmp(key, boundary, full, 'gt')
	190	+ greater_low = S.vcmp(low, low_boundary, full, 'gt')
	191	+ equal_low = S.vcmp(low, low_boundary, full, 'eq')
	192	+ strictly_greater = S.por(greater_key, S.pand(same, greater_low, full), full)
	193	+ tied = S.pand(same, equal_low, full)
	194	+ tie_prefix = S.vusqz(tied, 'int32')
	195	+ tie_rank = S.vadd(tie_prefix, tie_off, full)
	196	+ within_quota = S.vcmp(tie_rank, quota, full, 'lt')
	197	+ take_tie = S.pand(tied, within_quota, full)
	198	+ selected = S.por(strictly_greater, take_tie, full)
	199	+ prefix = S.vusqz(selected, 'int32')
	200	+ positions = S.vadd(prefix, off, full)
	201	+ indices = tirx.reinterpret('uint32x64', positions)
	202	+ candidates = S.vci(T.int32(vi * VS), 'int32')
	203	+ S.vscatter(candidates, sel_ub[0], indices, selected)
	204	+ nsel = S.vcadd(S.vsel(onev, zerov, selected), full)
	205	+ ntie = S.vcadd(S.vsel(onev, zerov, tied), full)
	206	+ nsel_all = S.vdupv(nsel, full)
	207	+ ntie_all = S.vdupv(ntie, full)
	208	+ S.vsts(out_off_ub[0], S.vadd(off, nsel_all, full), full, 'NORM_B32')
	209	+ S.vsts(tie_off_ub[0], S.vadd(tie_off, ntie_all, full), full, 'NORM_B32')
	210	+ S.mem_bar('VST_VLD')
	211	+
	212	+ classify_window = window_score_classifier(topk, nv)
	213	+
	214	+ @T.prim_func
	215	+ def main(
	216	+ scores_in: T.Tensor((num_rows, num_candidates), 'float32'),
	217	+ block_indices: T.Tensor((num_rows, block_topk), 'int32'),
	218	+ window_start: T.Tensor((num_rows,), 'int32'),
	219	+ window_end: T.Tensor((num_rows,), 'int32'),
	220	+ indices_out: T.Tensor((num_rows, topk), 'int32'),
	221	+ ):
	222	+ with T.Kernel(num_rows) as bx:
	223	+ classify_accum = T.alloc_shared((2, 64), 'int32')
	224	+ classify_counts = T.alloc_shared((2,), 'int32')
	225	+ acc_ub = T.alloc_shared((VS,), 'int32')
	226	+ l_ub = T.alloc_shared((num_candidates,), 'float32')
	227	+ k16_ub = T.alloc_shared((num_candidates,), 'int32')
	228	+ lo16_ub = T.alloc_shared((num_candidates,), 'int32')
	229	+ cnt_ub = T.alloc_shared((17,), 'int32')
	230	+ thr_ub = T.alloc_shared((16,), 'int32')
	231	+ sh_ub = T.alloc_shared((4,), 'int32')
	232	+ st_ub = T.alloc_shared((16,), 'int32')
	233	+ bi_ub = T.alloc_shared((block_topk,), 'int32')
	234	+ out_off_ub = T.alloc_shared((VS,), 'int32')
	235	+ tie_off_ub = T.alloc_shared((VS,), 'int32')
	236	+ sel_ub = T.alloc_shared((topk_pad,), 'int32')
	237	+ o_ub = T.alloc_shared((topk_pad,), 'int32')
	238	+
	239	+ T.copy(scores_in[bx, 0:num_candidates], l_ub)
	240	+ T.copy(block_indices[bx, 0:block_topk], bi_ub)
	241	+ ks_m = window_start[bx]
	242	+ span_m = window_end[bx] - window_start[bx]
	243	+
	244	+ split_score_keys(l_ub, k16_ub, lo16_ub)
	245	+
	246	+ classify_window(l_ub, classify_counts, classify_accum, span_m)
	247	+
	248	+ # -inf is a sufficient collection boundary when all greater scores fit.
	249	+ # NaNs fall back to the original full-key search.
	250	+ if (classify_counts[0] <= topk) and (classify_counts[0] + classify_counts[1] == num_candidates):
	251	+ st_ub[S_B16] = T.int32(127)
	252	+ st_ub[S_BLO] = T.int32(65535)
	253	+ st_ub[S_R2] = topk - classify_counts[0]
	254	+ else:
	255	+ search_boundary(k16_ub, lo16_ub, st_ub, sh_ub, thr_ub, cnt_ub, acc_ub)
	256	+
	257	+ collect_selected(k16_ub, lo16_ub, st_ub, sel_ub, out_off_ub, tie_off_ub)
	258	+
	259	+ with T.SimtVF(threads=128):
	260	+ for k in T.Parallel(topk_pad):
	261	+ c = sel_ub[k]
	262	+ pos = bi_ub[c // kb] * kb + (c % kb)
	263	+ rel = pos - ks_m
	264	+ o_ub[k] = T.if_then_else((rel >= 0) and (rel < span_m), rel, T.int32(-1))
	265	+ T.copy(o_ub[0:topk], indices_out[bx, 0:topk])
	266	+
	267	+ return main
	268	+
	269	+
	270	+def token_topk_small(num_rows: int, num_candidates: int, topk: int, kb: int, block_topk: int):
	271	+ if num_candidates % VS:
	272	+ raise AssertionError(f'NC={num_candidates} 必须是 {VS} 的整数倍')
	273	+ if topk > num_candidates or topk > block_topk * kb:
	274	+ raise AssertionError('top-k exceeds the available candidates')
	275	+ nv = num_candidates // VS
	276	+
	277	+ @T.macro
	278	+ def select_maxima(l_ub, idx_ub):
	279	+ with T.SimdVF():
	280	+ full = S.pset(32, 'PAT_ALL')
	281	+ one = S.pset(32, 'PAT_VL1')
	282	+ zero_score = S.vdup(T.float32(0), 'float32', full)
	283	+ zero = S.vdup(T.int32(0), 'int32', full)
	284	+ removed = S.vdup(T.int32(-2147483648), 'int32', full)
	285	+ imax = S.vdup(T.max_value(T.int32), 'int32', full)
	286	+
	287	+ keys = S.alloc_local(nv, 'int32')
	288	+ for i in T.unroll(nv, explicit=True):
	289	+ # Flipping the unsigned radix key's sign bit gives signed score order.
	290	+ # INT_MIN is below the key for -inf, so removed lanes cannot win a tie.
	291	+ keys[i] = S.vxor(_key32_of(l_ub, i, full, zero_score, zero, removed), removed, full)
	292	+ with T.serial(topk) as k:
	293	+ mx = S.vdupv(S.vcmax(redmax(keys, nv, full), full), full)
	294	+ iv = S.vdupv(S.vcmin(argmin(keys, nv, mx, full, imax), full), full)
	295	+ S.vsts(idx_ub[k], iv, one, 'ONEPT_B32')
	296	+ for i in T.unroll(nv, explicit=True):
	297	+ keys[i] = S.vsel(removed, keys[i], S.vcmp(ramp(i), iv, full, 'eq'))
	298	+
	299	+ @T.prim_func
	300	+ def main(
	301	+ scores_in: T.Tensor((num_rows, num_candidates), 'float32'),
	302	+ block_indices: T.Tensor((num_rows, block_topk), 'int32'),
	303	+ window_start: T.Tensor((num_rows,), 'int32'),
	304	+ window_end: T.Tensor((num_rows,), 'int32'),
	305	+ indices_out: T.Tensor((num_rows, topk), 'int32'),
	306	+ ):
	307	+ with T.Kernel(num_rows) as bx:
	308	+ l_ub = T.alloc_shared((num_candidates,), 'float32')
	309	+ bi_ub = T.alloc_shared((block_topk,), 'int32')
	310	+ idx_ub = T.alloc_shared((topk,), 'int32')
	311	+ o_ub = T.alloc_shared((topk,), 'int32')
	312	+ T.copy(scores_in[bx, 0:num_candidates], l_ub)
	313	+ T.copy(block_indices[bx, 0:block_topk], bi_ub)
	314	+ ks_m = window_start[bx]
	315	+ span_m = window_end[bx] - window_start[bx]
	316	+
	317	+ select_maxima(l_ub, idx_ub)
	318	+
	319	+ for k in T.serial(topk):
	320	+ j = idx_ub[k]
	321	+ pos = bi_ub[j // kb] * kb + (j % kb)
	322	+ rel = pos - ks_m
	323	+ o_ub[k] = T.if_then_else((rel >= 0) and (rel < span_m), rel, T.int32(-1))
	324	+ T.copy(o_ub, indices_out[bx, 0:topk])
	325	+
	326	+ return main
	327	+
	328	+
	329	+def token_topk_with_offset_interface(logits, block_idx, cu_ks, cu_ke, kb, topk_tokens):
	330	+ if logits.ndim != 2 or block_idx.ndim != 2:
	331	+ raise ValueError('logits and block_idx must be rank 2')
	332	+ if kb <= 0 or topk_tokens <= 0:
	333	+ raise ValueError('kb and topk_tokens must be positive')
	334	+ num_rows, num_candidates = logits.shape
	335	+ block_topk = block_idx.shape[1]
	336	+ if num_candidates % VS or block_topk * kb != num_candidates:
	337	+ raise ValueError('NC must be divisible by 64 and equal block_topk*kb')
	338	+ if (
	339	+ tuple(block_idx.shape) != (num_rows, block_topk)
	340	+ or tuple(cu_ks.shape) != (num_rows,)
	341	+ or tuple(cu_ke.shape) != (num_rows,)
	342	+ ):
	343	+ raise ValueError('row count / ks / ke shape mismatch')
	344	+ tensors = (logits, block_idx, cu_ks, cu_ke)
	345	+ if logits.dtype != torch.float32 or any(t.dtype != torch.int32 for t in tensors[1:]):
	346	+ raise TypeError('expected float32 logits and int32 metadata')
	347	+ if any(t.device != logits.device or not t.is_contiguous() for t in tensors):
	348	+ raise ValueError('inputs must be contiguous and on the same NPU')
	349	+ if logits.device.type != 'npu':
	350	+ raise ValueError('inputs must be on NPU')
	351	+ topk = min(topk_tokens, num_candidates)
	352	+ if num_candidates <= 1024 and topk <= 32:
	353	+ kern = cached_kernel(
	354	+ ('hisa-token-small', num_rows, num_candidates, topk, kb, block_topk),
	355	+ lambda: token_topk_small(num_rows, num_candidates, topk, kb, block_topk),
	356	+ [],
	357	+ )
	358	+ elif (num_candidates, topk) == (8192, 2048):
	359	+ from .token_topk_radix4 import token_topk_radix4
	360	+
	361	+ kern = cached_kernel(
	362	+ ('hisa-token-radix4', num_rows, num_candidates, topk, kb, block_topk),
	363	+ lambda: token_topk_radix4(num_rows, kb, block_topk),
	364	+ [],
	365	+ )
	366	+ else:
	367	+ kern = cached_kernel(
	368	+ ('hisa-token-radix', num_rows, num_candidates, topk, kb, block_topk),
	369	+ lambda: token_topk_with_offset(num_rows, num_candidates, topk, kb, block_topk),
	370	+ [],
	371	+ )
	372	+ # ACL Graph records the NPU storage format of each public output.
	373	+ # Allocate it with PyTorch so it has the native torch_npu storage descriptor.
	374	+ out = torch.empty((num_rows, topk), dtype=torch.int32, device=logits.device)
	375	+ kern(logits, block_idx, cu_ks, cu_ke, out)
	376	+ return out

Asrc/cann_ops_tilelang/attention/dsa_hisa/token_topk_radix4.py+189-0

@@ -0,0 +1,189 @@
	1	+# Copyright (c) 2026 ops-tilelang contributors
	2	+# SPDX-License-Identifier: MIT
	3	+# See LICENSE in the repository root for the full license text.
	4	+
	5	+"""Radix-4 selection specialized for 8192 FP32 candidates and top-k 2048."""
	6	+
	7	+import tilelang.ascend.language as T
	8	+from tilelang.ascend.language import simd as S
	9	+from tvm import tirx
	10	+
	11	+from .tilelang_utils import window_score_classifier
	12	+
	13	+NC = 8192
	14	+
	15	+
	16	+TOPK = 2048
	17	+
	18	+
	19	+LANES = 64
	20	+
	21	+
	22	+TILE_VECTORS = 8
	23	+
	24	+
	25	+TILES = NC // LANES // TILE_VECTORS
	26	+
	27	+
	28	+def _count_tile(keys, tile, b1, b2, b3, full, ones, zeros):
	29	+ a1 = a2 = a3 = zeros
	30	+ for j in range(TILE_VECTORS):
	31	+ key = S.vld(keys[(tile * TILE_VECTORS + j) * LANES])
	32	+ p1 = S.vcmps(key, b1, full, 'ge')
	33	+ p2 = S.vcmps(key, b2, full, 'ge')
	34	+ p3 = S.vcmps(key, b3, full, 'ge')
	35	+ a1 = S.vadd(a1, S.vsel(ones, zeros, p1), full)
	36	+ a2 = S.vadd(a2, S.vsel(ones, zeros, p2), full)
	37	+ a3 = S.vadd(a3, S.vsel(ones, zeros, p3), full)
	38	+ return a1, a2, a3
	39	+
	40	+
	41	+def token_topk_radix4(num_rows: int, kb: int, block_topk: int):
	42	+ if num_rows <= 0 or kb * block_topk != NC:
	43	+ raise AssertionError(f'radix-4 requires positive rows and {NC} candidates')
	44	+ row_stride = 2053 if num_rows == 8192 else 1
	45	+
	46	+ @T.macro
	47	+ def order_score_keys(logits, keys):
	48	+ with T.SimdVF():
	49	+ full = S.pset(32, 'PAT_ALL')
	50	+ zero = S.vdup(T.int32(0), 'int32', full)
	51	+ sign = S.vdup(T.int32(-2147483648), 'int32', full)
	52	+ for i in T.serial(NC // LANES):
	53	+ bits = tirx.reinterpret('int32x64', S.vld(logits[i * LANES]))
	54	+
	55	+ magnitude = S.vand(bits, S.vnot(sign, full), full)
	56	+ is_zero = S.vcmp(magnitude, zero, full, 'eq')
	57	+ bits = S.vsel(zero, bits, is_zero)
	58	+ negative = S.vcmp(bits, zero, full, 'lt')
	59	+ negative_key = S.vor(S.vnot(bits, full), sign, full)
	60	+ key = S.vsel(negative_key, bits, negative)
	61	+ S.vsts(keys[i * LANES], key, full, 'NORM_B32')
	62	+
	63	+ @T.macro
	64	+ def count_thresholds(keys, tile_counts, counts, b1, b2, b3):
	65	+ with T.SimdVF():
	66	+ full = S.pset(32, 'PAT_ALL')
	67	+ one = S.pset(32, 'PAT_VL1')
	68	+ live_tiles = S.pset(32, 'PAT_VL16')
	69	+ zeros = S.vdup(T.int32(0), 'int32', full)
	70	+ ones = S.vdup(T.int32(1), 'int32', full)
	71	+ with T.serial(TILES) as tile:
	72	+ a1, a2, a3 = _count_tile(keys, tile, b1, b2, b3, full, ones, zeros)
	73	+ S.vsts(tile_counts[0, tile], S.vcadd(a1, full), one, 'ONEPT_B32')
	74	+ S.vsts(tile_counts[1, tile], S.vcadd(a2, full), one, 'ONEPT_B32')
	75	+ S.vsts(tile_counts[2, tile], S.vcadd(a3, full), one, 'ONEPT_B32')
	76	+ S.mem_bar('VST_VLD')
	77	+ for j in range(3):
	78	+ total = S.vcadd(S.vld(tile_counts[j, 0]), live_tiles)
	79	+ S.vsts(counts[j], total, one, 'ONEPT_B32')
	80	+
	81	+ @T.macro
	82	+ def advance_quarter(state, counts, b1, b2, b3):
	83	+ if counts[2] >= TOPK:
	84	+ state[0] = b3
	85	+ elif counts[1] >= TOPK:
	86	+ state[0] = b2
	87	+ state[2] = counts[2]
	88	+ elif counts[0] >= TOPK:
	89	+ state[0] = b1
	90	+ state[2] = counts[1]
	91	+ else:
	92	+ state[2] = counts[0]
	93	+ state[1] = state[1] // 4
	94	+
	95	+ @T.macro
	96	+ def search_threshold(keys, tile_counts, counts, state):
	97	+ state[0] = T.int32(-2147483648)
	98	+ state[1] = T.int32(1073741824)
	99	+ state[2] = T.int32(0)
	100	+ for _ in T.serial(16):
	101	+ base = T.Cast('int64', state[0])
	102	+ step = T.Cast('int64', state[1])
	103	+ b1 = T.Cast('int32', base + step)
	104	+ b2 = T.Cast('int32', base + 2 * step)
	105	+ b3 = T.Cast('int32', base + 3 * step)
	106	+ count_thresholds(keys, tile_counts, counts, b1, b2, b3)
	107	+ advance_quarter(state, counts, b1, b2, b3)
	108	+
	109	+ @T.macro
	110	+ def collect_selected(keys, state, offsets, selected_indices):
	111	+ with T.SimdVF():
	112	+ full = S.pset(32, 'PAT_ALL')
	113	+ zeros = S.vdup(T.int32(0), 'int32', full)
	114	+ ones = S.vdup(T.int32(1), 'int32', full)
	115	+ S.vsts(offsets[0, 0], zeros, full, 'NORM_B32')
	116	+ S.vsts(offsets[1, 0], zeros, full, 'NORM_B32')
	117	+ S.mem_bar('VST_VLD')
	118	+ boundary = S.vdup(state[0], 'int32', full)
	119	+ quota = S.vdup(TOPK - state[2], 'int32', full)
	120	+ with T.serial(NC // LANES) as vector:
	121	+ keys_v = S.vld(keys[vector * LANES])
	122	+ off = S.vld(offsets[0, 0])
	123	+ tie_off = S.vld(offsets[1, 0])
	124	+ tied = S.vcmp(keys_v, boundary, full, 'eq')
	125	+ greater = S.vcmp(keys_v, boundary, full, 'gt')
	126	+ tie_rank = S.vadd(S.vusqz(tied, 'int32'), tie_off, full)
	127	+ take_tie = S.pand(tied, S.vcmp(tie_rank, quota, full, 'lt'), full)
	128	+ selected = S.por(greater, take_tie, full)
	129	+ positions = S.vadd(S.vusqz(selected, 'int32'), off, full)
	130	+ candidates = S.vci(T.Cast('int32', vector * LANES), 'int32')
	131	+ S.vscatter(candidates, selected_indices[0], tirx.reinterpret('uint32x64', positions), selected)
	132	+ nsel = S.vcadd(S.vsel(ones, zeros, selected), full)
	133	+ ntie = S.vcadd(S.vsel(ones, zeros, tied), full)
	134	+ S.vsts(offsets[0, 0], S.vadd(off, S.vdupv(nsel, full), full), full, 'NORM_B32')
	135	+ S.vsts(offsets[1, 0], S.vadd(tie_off, S.vdupv(ntie, full), full), full, 'NORM_B32')
	136	+ S.mem_bar('VST_VLD')
	137	+
	138	+ classify_window = window_score_classifier(TOPK, NC // LANES)
	139	+
	140	+ @T.prim_func
	141	+ def main(
	142	+ scores_in: T.Tensor((num_rows, NC), 'float32'),
	143	+ block_indices: T.Tensor((num_rows, block_topk), 'int32'),
	144	+ window_start: T.Tensor((num_rows,), 'int32'),
	145	+ window_end: T.Tensor((num_rows,), 'int32'),
	146	+ indices_out: T.Tensor((num_rows, TOPK), 'int32'),
	147	+ ):
	148	+ with T.Kernel(num_rows) as bx:
	149	+ # Permute query execution only; preserve all row-local selection semantics.
	150	+ row = (bx * row_stride) % num_rows
	151	+ classify_accum = T.alloc_shared((2, 64), 'int32')
	152	+ classify_counts = T.alloc_shared((2,), 'int32')
	153	+ logits = T.alloc_shared((NC,), 'float32')
	154	+ keys = T.alloc_shared((NC,), 'int32')
	155	+ bi = T.alloc_shared((block_topk,), 'int32')
	156	+ tile_counts = T.alloc_shared((3, LANES), 'int32')
	157	+ counts = T.alloc_shared((3,), 'int32')
	158	+
	159	+ state = T.alloc_shared((3,), 'int32')
	160	+ selected_indices = T.alloc_shared((TOPK,), 'int32')
	161	+ offsets = T.alloc_shared((2, LANES), 'int32')
	162	+ out = T.alloc_shared((TOPK,), 'int32')
	163	+
	164	+ T.copy(scores_in[row, 0:NC], logits)
	165	+ T.copy(block_indices[row, 0:block_topk], bi)
	166	+ ks = window_start[row]
	167	+ span = window_end[row] - ks
	168	+ order_score_keys(logits, keys)
	169	+
	170	+ classify_window(logits, classify_counts, classify_accum, span)
	171	+
	172	+ # Preserve the common collection and tie ordering below.
	173	+ if (classify_counts[0] <= TOPK) and (classify_counts[0] + classify_counts[1] == NC):
	174	+ state[0] = T.int32(-2139095041) # ordered FP32 key of -inf
	175	+ state[2] = classify_counts[0]
	176	+ else:
	177	+ search_threshold(keys, tile_counts, counts, state)
	178	+
	179	+ collect_selected(keys, state, offsets, selected_indices)
	180	+
	181	+ with T.SimtVF(threads=128):
	182	+ for k in T.Parallel(TOPK):
	183	+ c = selected_indices[k]
	184	+ pos = bi[c // kb] * kb + c % kb
	185	+ rel = pos - ks
	186	+ out[k] = T.if_then_else((rel >= 0) and (rel < span), rel, T.int32(-1))
	187	+ T.copy(out, indices_out[row, 0:TOPK])
	188	+
	189	+ return main

Atests/attention/benchmark/test_dsa_hisa.jsonl+6-0

@@ -0,0 +1,6 @@
	1	+{"kernel": "hisa_indexer", "operation": "indexer", "params": {"D": 128, "H": 64, "M": 1024, "block_topk": 16, "num_seqs": 1, "topk_tokens": 256}, "param_id": "1024-16-256-1", "cluster": "a5", "time_us": 526.78}
	2	+{"kernel": "hisa_indexer", "operation": "indexer", "params": {"D": 128, "H": 64, "M": 1024, "block_topk": 16, "num_seqs": 4, "topk_tokens": 256}, "param_id": "1024-16-256-4", "cluster": "a5", "time_us": 280.92}
	3	+{"kernel": "hisa_indexer", "operation": "indexer", "params": {"D": 128, "H": 64, "M": 4096, "block_topk": 32, "num_seqs": 1, "topk_tokens": 1024}, "param_id": "4096-32-1024-1", "cluster": "a5", "time_us": 4920.37}
	4	+{"kernel": "hisa_indexer", "operation": "indexer", "params": {"D": 128, "H": 64, "M": 4096, "block_topk": 32, "num_seqs": 4, "topk_tokens": 1024}, "param_id": "4096-32-1024-4", "cluster": "a5", "time_us": 1871.36}
	5	+{"kernel": "hisa_indexer", "operation": "indexer", "params": {"D": 128, "H": 64, "M": 8192, "block_topk": 64, "num_seqs": 1, "topk_tokens": 2048}, "param_id": "8192-64-2048-1", "cluster": "a5", "time_us": 8096.59}
	6	+{"kernel": "hisa_indexer", "operation": "indexer", "params": {"D": 128, "H": 64, "M": 8192, "block_topk": 64, "num_seqs": 8, "topk_tokens": 2048}, "param_id": "8192-64-2048-8", "cluster": "a5", "time_us": 4821.62}

Atests/attention/dsa_hisa_frozen_scorer.py+137-0

@@ -0,0 +1,137 @@
	1	+# Copyright (c) 2026 ops-tilelang contributors
	2	+# SPDX-License-Identifier: MIT
	3	+# See LICENSE in the repository root for the full license text.
	4	+
	5	+"""Frozen dense-traversal scorer used to verify FP32 bits across visibility optimizations."""
	6	+
	7	+import tilelang.ascend.language as T
	8	+from tilelang.ascend.language import simd as S
	9	+
	10	+
	11	+def fused_scorer_simd(rows: int, heads: int, dim: int, num_keys: int, kb: int, block_topk: int):
	12	+ """Build the historical dense scorer for the three FP8 bitwise regression shapes."""
	13	+ if (heads, dim, kb) != (64, 128, 128) or block_topk not in (8, 32, 64):
	14	+ raise ValueError('frozen score regression requires H=64, D=128, kb=128 and 8/32/64 blocks')
	15	+ # The historical 64 KiB L0C holds two FP32 blocks: 64 heads x 256 scores.
	16	+ ntile = 2
	17	+ padded_heads = heads
	18	+ tile_tokens = ntile * kb
	19	+ half_block = kb // 2
	20	+ half_tile = tile_tokens // 2
	21	+ half_group = ntile * half_block
	22	+ candidates = block_topk * kb
	23	+ groups = block_topk // ntile
	24	+ ng = half_tile // 64
	25	+
	26	+ query_shape = (rows, heads, dim)
	27	+ head_shape = (rows, heads)
	28	+ key_shape = (num_keys, dim)
	29	+ key_axis = (num_keys,)
	30	+ row_axis = (rows,)
	31	+ index_shape = (rows, block_topk)
	32	+ score_shape = (rows, candidates)
	33	+ query_tile = (padded_heads, dim)
	34	+ key_tile = (tile_tokens, dim)
	35	+ result_tile = (padded_heads, tile_tokens)
	36	+ half_result_tile = (padded_heads, half_tile)
	37	+ head_axis = (heads,)
	38	+ half_key_axis = (half_tile,)
	39	+
	40	+ @T.macro
	41	+ def load_key_pair(keys, block_ids, kp_l1, bx, g):
	42	+ for hf in range(2):
	43	+ for ti in range(ntile):
	44	+ blk = block_ids[bx, g * ntile + ti]
	45	+ key_start = T.meta_var(blk * kb + hf * half_block)
	46	+ key_end = T.meta_var(blk * kb + (hf + 1) * half_block)
	47	+ tile_start = T.meta_var(hf * half_group + ti * half_block)
	48	+ tile_end = T.meta_var(hf * half_group + (ti + 1) * half_block)
	49	+ T.copy(keys[key_start:key_end, 0:dim], kp_l1[tile_start:tile_end, 0:dim])
	50	+
	51	+ @T.macro
	52	+ def load_scales_and_positions(key_scale, positions, block_ids, ks_ub, pas_ub, bx, g):
	53	+ for ti in range(ntile):
	54	+ blk = block_ids[bx, g * ntile + ti]
	55	+ block_start = T.meta_var(blk * kb)
	56	+ block_end = T.meta_var((blk + 1) * kb)
	57	+ lane_start = T.meta_var(ti * half_block)
	58	+ lane_end = T.meta_var((ti + 1) * half_block)
	59	+ T.dual_copy(key_scale[block_start:block_end], ks_ub[lane_start:lane_end])
	60	+ T.dual_copy(positions[block_start:block_end], pas_ub[lane_start:lane_end])
	61	+
	62	+ @T.macro
	63	+ def reduce_dense_scores(s_ub, qs_ub, ks_ub, w_ub, pas_ub, o_ub, ks_m, ke_m):
	64	+ with T.SimdVF():
	65	+ acc = S.alloc_local(ng, 'float32')
	66	+ full = S.pset(32, 'PAT_ALL')
	67	+ zero = S.vdup(T.float32(0), 'float32', full)
	68	+ ninf = S.vdup(-T.infinity(T.float32), 'float32', full)
	69	+ ksm = S.vdup(ks_m, 'float32', full)
	70	+ kem = S.vdup(ke_m, 'float32', full)
	71	+ for gi in T.unroll(ng):
	72	+ acc[gi] = S.vdup(T.float32(0), 'float32', full)
	73	+ with T.serial(heads) as h:
	74	+ qsv = S.vld(qs_ub[h], 'BRC_B32')
	75	+ wv = S.vld(w_ub[h], 'BRC_B32')
	76	+ for gj in T.unroll(ng):
	77	+ kv = S.vld(ks_ub[gj * 64], 'NORM')
	78	+ sv = S.vld(s_ub[h, gj * 64], 'NORM')
	79	+ p = S.vmul(sv, S.vmul(qsv, kv, full), full)
	80	+ r = S.vmax(p, zero, full)
	81	+ acc[gj] = S.vadd(acc[gj], S.vmul(r, wv, full), full)
	82	+ for gk in T.unroll(ng):
	83	+ pv = S.vld(pas_ub[gk * 64], 'NORM')
	84	+ bad = S.por(S.vcmp(pv, ksm, full, 'lt'), S.vcmp(pv, kem, full, 'ge'), full)
	85	+ S.vsts(o_ub[0, gk * 64], S.vsel(ninf, acc[gk], bad), full, 'NORM_B32')
	86	+
	87	+ @T.macro
	88	+ def store_score_pair(o_ub, output, bx, g):
	89	+ for ti in range(ntile):
	90	+ row_end = T.meta_var(bx + 1)
	91	+ lane_start = T.meta_var(ti * half_block)
	92	+ lane_end = T.meta_var((ti + 1) * half_block)
	93	+ output_start = T.meta_var((g * ntile + ti) * kb)
	94	+ output_end = T.meta_var((g * ntile + ti + 1) * kb)
	95	+ T.dual_copy(o_ub[0:1, lane_start:lane_end], output[bx:row_end, output_start:output_end])
	96	+
	97	+ @T.prim_func
	98	+ def token_score_simd(
	99	+ query: T.Tensor(query_shape, T.float8_e4m3fn),
	100	+ query_scale: T.Tensor(head_shape, 'float32'),
	101	+ keys: T.Tensor(key_shape, T.float8_e4m3fn),
	102	+ key_scale: T.Tensor(key_axis, 'float32'),
	103	+ weights: T.Tensor(head_shape, 'float32'),
	104	+ block_ids: T.Tensor(index_shape, 'int32'),
	105	+ window_start: T.Tensor(row_axis, 'int32'),
	106	+ window_end: T.Tensor(row_axis, 'int32'),
	107	+ positions: T.Tensor(key_axis, 'float32'),
	108	+ output: T.Tensor(score_shape, 'float32'),
	109	+ ):
	110	+ with T.Kernel(rows) as bx:
	111	+ q_l1 = T.alloc_l1(query_tile, T.float8_e4m3fn)
	112	+ kp_l1 = T.alloc_l1(key_tile, T.float8_e4m3fn)
	113	+ res = T.alloc_l0c(result_tile, 'float32')
	114	+ s_ub = T.alloc_shared(half_result_tile, 'float32')
	115	+ qs_ub = T.alloc_shared(head_axis, 'float32')
	116	+ w_ub = T.alloc_shared(head_axis, 'float32')
	117	+ ks_ub = T.alloc_shared(half_key_axis, 'float32')
	118	+ pas_ub = T.alloc_shared(half_key_axis, 'float32')
	119	+ o_ub = T.alloc_shared((1, half_tile), 'float32')
	120	+
	121	+ T.copy(query[bx, 0:heads, 0:dim], q_l1[0:heads, 0:dim])
	122	+ T.copy(query_scale[bx, 0:heads], qs_ub)
	123	+ T.copy(weights[bx, 0:heads], w_ub)
	124	+ ks_m = T.Cast('float32', window_start[bx])
	125	+ ke_m = T.Cast('float32', window_end[bx])
	126	+
	127	+ for g in range(groups):
	128	+ load_key_pair(keys, block_ids, kp_l1, bx, g)
	129	+ load_scales_and_positions(key_scale, positions, block_ids, ks_ub, pas_ub, bx, g)
	130	+
	131	+ T.gemm(q_l1, kp_l1, res, transpose_B=True, clear_accum=True)
	132	+ T.dual_copy(res, s_ub)
	133	+
	134	+ reduce_dense_scores(s_ub, qs_ub, ks_ub, w_ub, pas_ub, o_ub, ks_m, ke_m)
	135	+ store_score_pair(o_ub, output, bx, g)
	136	+
	137	+ return token_score_simd

Atests/attention/dsa_hisa_ref.py+59-0

@@ -0,0 +1,59 @@
	1	+# Copyright (c) 2026 ops-tilelang contributors
	2	+# SPDX-License-Identifier: MIT
	3	+# See LICENSE in the repository root for the full license text.
	4	+
	5	+"""CPU PyTorch reference for hierarchical FP8 block and token selection."""
	6	+
	7	+import torch
	8	+
	9	+
	10	+def _weighted_scores(q, keys, scales, weights, q_scale=None):
	11	+ dots = q.float() @ keys.float().T
	12	+ if q_scale is not None:
	13	+ scales = q_scale.unsqueeze(-1) * scales
	14	+ return ((dots * scales).relu() * weights.unsqueeze(-1)).sum(dim=1)
	15	+
	16	+
	17	+def token_scores_ref(q, q_scale, k, k_scale, weights, block_indices, ks, ke, kb):
	18	+ """FP32 token logits ordered by candidate slot, including the window mask."""
	19	+ positions = (block_indices.long().unsqueeze(-1) * kb + torch.arange(kb)).flatten(1)
	20	+ dense = _weighted_scores(q, k, k_scale, weights, q_scale)
	21	+ logits = dense.gather(1, positions)
	22	+ return logits.masked_fill((positions < ks[:, None]) \| (positions >= ke[:, None]), -torch.inf)
	23	+
	24	+
	25	+def hisa_indexer_ref(q, k, k_scale, weights, ks, ke, *, k_block_size, block_topk, topk_tokens):
	26	+ """Return relative indices from raw inputs; use stable sorting for ties."""
	27	+ rows = q.shape[0]
	28	+ tokens, dim = k.shape
	29	+ blocks = tokens // k_block_size
	30	+ num_selected = min(block_topk, blocks)
	31	+ width = min(topk_tokens, num_selected * k_block_size)
	32	+ result = torch.full((rows, width), -1, dtype=torch.int32)
	33	+ if not rows or not tokens:
	34	+ return result
	35	+
	36	+ block_means = (k.float() * k_scale[:, None]).reshape(blocks, k_block_size, dim).mean(dim=1)
	37	+ block_scales = (block_means.abs().amax(dim=1) / torch.finfo(q.dtype).max).clamp_min(1e-10)
	38	+ pooled = (block_means / block_scales[:, None]).to(q.dtype)
	39	+ block_scores = _weighted_scores(q, pooled, block_scales, weights).to(torch.bfloat16)
	40	+ token_scores = _weighted_scores(q, k, k_scale, weights)
	41	+ all_blocks = torch.arange(blocks)
	42	+
	43	+ for row, (start, stop) in enumerate(zip(ks.tolist(), ke.tolist(), strict=True)):
	44	+ if start == stop:
	45	+ continue
	46	+ first, last = start // k_block_size, (stop - 1) // k_block_size
	47	+ boundary = torch.tensor([first, last] if first != last else [first])
	48	+ interior = torch.arange(first + 1, max(first + 1, last))
	49	+ order = block_scores[row, interior].argsort(descending=True, stable=True)
	50	+ outside = all_blocks[(all_blocks < first) \| (all_blocks > last)]
	51	+ selected = torch.cat((boundary, interior[order], outside))[:num_selected]
	52	+ positions = (selected[:, None] * k_block_size + torch.arange(k_block_size)).flatten()
	53	+ valid = (positions >= start) & (positions < stop)
	54	+ # A visible score may itself be -inf. Apply the mask before the stable
	55	+ # sort so a masked slot still participates in the tie at that score.
	56	+ scores = token_scores[row, positions].masked_fill(~valid, -torch.inf)
	57	+ top = scores.argsort(descending=True, stable=True)[:width]
	58	+ result[row] = (positions[top] - start).masked_fill(~valid[top], -1).int()
	59	+ return result

Atests/attention/npu_mem_profiles/test_dsa_hisa.json+104-0

@@ -0,0 +1,104 @@
	1	+{
	2	+ "attention/test_dsa_hisa.py::test_bf16_block_selection[32,]": {
	3	+ "peak_npu_mb": 0.0
	4	+ },
	5	+ "attention/test_dsa_hisa.py::test_bf16_block_selection[64,]": {
	6	+ "peak_npu_mb": 0.0
	7	+ },
	8	+ "attention/test_dsa_hisa.py::test_bf16_block_selection[8,]": {
	9	+ "peak_npu_mb": 0.0
	10	+ },
	11	+ "attention/test_dsa_hisa.py::test_fp8_score_bits[1024-8,]": {
	12	+ "peak_npu_mb": 0.5
	13	+ },
	14	+ "attention/test_dsa_hisa.py::test_fp8_score_bits[4096-32,]": {
	15	+ "peak_npu_mb": 1.3
	16	+ },
	17	+ "attention/test_dsa_hisa.py::test_fp8_score_bits[8192-64,]": {
	18	+ "peak_npu_mb": 3.3
	19	+ },
	20	+ "attention/test_dsa_hisa.py::test_hisa_against_torch[n=1024,blocks=2,topk=191,]": {
	21	+ "peak_npu_mb": 16.2
	22	+ },
	23	+ "attention/test_dsa_hisa.py::test_hisa_against_torch[n=1024,blocks=4,topk=4,dim=96,]": {
	24	+ "peak_npu_mb": 0.2
	25	+ },
	26	+ "attention/test_dsa_hisa.py::test_hisa_against_torch[n=10240,blocks=4,topk=321,]": {
	27	+ "peak_npu_mb": 1.4
	28	+ },
	29	+ "attention/test_dsa_hisa.py::test_hisa_against_torch[n=4096,blocks=1,topk=32,]": {
	30	+ "peak_npu_mb": 0.6
	31	+ },
	32	+ "attention/test_dsa_hisa.py::test_hisa_against_torch[n=8192,blocks=4,topk=321,]": {
	33	+ "peak_npu_mb": 1.2
	34	+ },
	35	+ "attention/test_dsa_hisa.py::test_hisa_against_torch[n=8192,blocks=64,topk=2048,]": {
	36	+ "peak_npu_mb": 1.4
	37	+ },
	38	+ "attention/test_dsa_hisa.py::test_hisa_compile_composition": {
	39	+ "peak_npu_mb": 16.4
	40	+ },
	41	+ "attention/test_dsa_hisa.py::test_hisa_compile_empty[0-1024-256,]": {
	42	+ "peak_npu_mb": 0.3
	43	+ },
	44	+ "attention/test_dsa_hisa.py::test_hisa_compile_empty[8-0-0,]": {
	45	+ "peak_npu_mb": 0.1
	46	+ },
	47	+ "attention/test_dsa_hisa.py::test_hisa_compile_opcheck": {
	48	+ "peak_npu_mb": 16.6
	49	+ },
	50	+ "attention/test_dsa_hisa.py::test_hisa_compile_replay[1024-4-32,]": {
	51	+ "peak_npu_mb": 16.4
	52	+ },
	53	+ "attention/test_dsa_hisa.py::test_hisa_compile_replay[4096-32-1024,]": {
	54	+ "peak_npu_mb": 17.2
	55	+ },
	56	+ "attention/test_dsa_hisa.py::test_hisa_compile_replay[8192-64-2048,]": {
	57	+ "peak_npu_mb": 18.2
	58	+ },
	59	+ "attention/test_dsa_hisa.py::test_hisa_empty_and_invalid_metadata": {
	60	+ "peak_npu_mb": 16.1
	61	+ },
	62	+ "attention/test_dsa_hisa.py::test_hisa_prunes_blocks_before_token_selection": {
	63	+ "peak_npu_mb": 24.6
	64	+ },
	65	+ "attention/test_dsa_hisa.py::test_hisa_rejects_unsupported_geometry_before_launch[384-128-64-128,]": {
	66	+ "peak_npu_mb": 16.1
	67	+ },
	68	+ "attention/test_dsa_hisa.py::test_hisa_rejects_unsupported_geometry_before_launch[64-32-1-3008,]": {
	69	+ "peak_npu_mb": 16.2
	70	+ },
	71	+ "attention/test_dsa_hisa.py::test_hisa_rejects_unsupported_geometry_before_launch[8192-512-64-128,]": {
	72	+ "peak_npu_mb": 17.1
	73	+ },
	74	+ "attention/test_dsa_hisa.py::test_hisa_selection[0-1024-16-256-1,]": {
	75	+ "peak_npu_mb": 13.5
	76	+ },
	77	+ "attention/test_dsa_hisa.py::test_hisa_selection[1-1024-16-256-4,]": {
	78	+ "peak_npu_mb": 13.5
	79	+ },
	80	+ "attention/test_dsa_hisa.py::test_hisa_selection[2-4096-32-1024-1,]": {
	81	+ "peak_npu_mb": 117.6
	82	+ },
	83	+ "attention/test_dsa_hisa.py::test_hisa_selection[3-4096-32-1024-4,]": {
	84	+ "peak_npu_mb": 115.6
	85	+ },
	86	+ "attention/test_dsa_hisa.py::test_hisa_selection[4-8192-64-2048-1,]": {
	87	+ "peak_npu_mb": 394.1
	88	+ },
	89	+ "attention/test_dsa_hisa.py::test_hisa_selection[5-8192-64-2048-8,]": {
	90	+ "peak_npu_mb": 388.1
	91	+ },
	92	+ "attention/test_dsa_hisa.py::test_token_ties_and_boundaries[1024-256,]": {
	93	+ "peak_npu_mb": 0.1
	94	+ },
	95	+ "attention/test_dsa_hisa.py::test_token_ties_and_boundaries[1024-32,]": {
	96	+ "peak_npu_mb": 0.0
	97	+ },
	98	+ "attention/test_dsa_hisa.py::test_token_ties_and_boundaries[4096-1024,]": {
	99	+ "peak_npu_mb": 0.2
	100	+ },
	101	+ "attention/test_dsa_hisa.py::test_token_ties_and_boundaries[8192-2048,]": {
	102	+ "peak_npu_mb": 0.4
	103	+ }
	104	+}

Atests/attention/test_dsa_hisa.py+494-0