feat(qwen4): add the QSA sparse-attention kernel and its golden by sysuTobo · Pull Request #1151 · pegainfer-project/pegainfer

sysuTobo · 2026-10-12T13:27:25Z

Summary

Adds the QSA sparse-attention forward for the Qwen4 line, together with the component golden and gate that pin it. The kernel consumes the indexer/selection mask (it never re-derives it) and produces the gated attention output that o_proj reads. Prefill only, over a contiguous key array, with an 8192-token cap; paged KV, decode, and the streaming softmax are later work.

What changed

pegainfer-kernels/csrc/qwen4/qwen4_qsa_attn_forward.cu — one key pass that (1+w)-norms and partially rotates every token's two key heads, then one CTA per (query row, one/two/four query heads) for the bf16 score, fp32 softmax, bf16-rounded weights, context, and sigmoid output gate. The CTA's heads share one key/value scan; the head count follows the score-row shared-memory budget.
pegainfer-kernels/csrc/qwen4/qwen4_qsa_norm_rope.cuh — shares the (1+w) RMSNorm and partial-RoPE rounding points with the indexer. Both kernels now read the same shared cos/sin table, so neither computes trig on the device.
qwen4_qsa_attn_forward_cuda / qsa_attn_forward_into — the FFI entry and Rust wrapper. The wrapper validates QsaAttnGeometry against the compiled constants before sizing any buffer, refuses a short cos/sin table, traps rows/positions outside 0..tokens, and treats an empty launch as a no-op like the sibling wrappers.
tests/qwen4_qsa_attn_parity.rs — 8 tests: the four-case replay, three input-level controls (causal mask, KV-head swap, saturated gate), the exact MAX_TOKENS launch, an empty-launch test, and the refusal contract. The first three cases' weights and the first two cases' outputs are asserted bit-for-bit; the 2055-token case keeps the calibrated tolerance.
tests/qwen4_qsa_attn_rows_{,negative_}trap.rs — the two halves of the out-of-range rows guard, one trap per process.
tools/accuracy/dump_qwen38fn_qsa_attn_golden.py + test_data/qwen38fn-qsa-attn-golden.safetensors — the four-case component golden (S=8/10/64, and S=2055 with rows 2049..=2054), 38 tensors, one JSON metadata key. The dump is a pure recorder of the reference's own forward hooks; the indexer has its own oracle.

Semantics

The kernel reproduces the reference's q|gate split of the fused q_proj row, (1+w) RMSNorm on q and k at rms_norm_eps, 64-of-256 partial RoPE from the shared table, GQA group 12, q·k/√256, the stored mask's polarity (non-zero = visible), fp32 softmax over the rounded bf16 scores, weights rounded to bf16 before the context reduction, and the sigmoid gate after the head-major reshape. A fully hidden row is outside the contract: the reference's two masks sum to -inf and its softmax is NaN, while this kernel's single minimum produces a uniform row.

Validation

PEGAINFER_REQUIRE_GPU=1 cargo test --release -p pegainfer-kernels --features qwen4 \
  --test qwen4_qsa_attn_parity -- --nocapture

CUDA 12.9.1, one RTX 5060 (sm_120): test result: ok. 8 passed; 0 failed, 398160 weights and 540672 output elements.

case	shape	`attn_weights`	`attn_out`
0	S=8, R=8	1536/1536 bit-identical	49152/49152 bit-identical
1	S=10, R=10	2400/2400 bit-identical	61440/61440 bit-identical
2	S=64, R=64	98304/98304 bit-identical	393207/393216 (9 one-ulp), max abs 1.5e-5
3	S=2055, R=6	295911/295920 (9 one-quantum), max abs 7.6e-6	36793/36864, max abs 1.2e-4, max rel 1.1e-2, 13 over one ulp, 4 over two

Worst bound fraction across the four cases: 0.17 (weights), 0.17 (output). The 2055-token residual is the float32 reduction order; moving the trig source to the shared table is what took it down from the earlier 0.42.

The three controls keep the gate from being vacuous: replacing the stored mask with plain causality leaves the causal cases at 0.000 of the bound and fails case 3 at 4.1e4 (weights) / 2.2e1 (output); exchanging the two KV heads fails every case (weights 9.1e3..5.5e4, output 6.1e2..7.1e3); saturating the gate leaves the weights at 0.000..0.17 and fails the output at 1.3e2..2.5e2.

Mutation calibration: dropping the bf16 rounding of the sigmoid now fails the case 0 exact assertion (36077/49152 outputs bit-identical). Before the short cases were asserted bit-for-bit, that mutation passed the whole gate.

The indexer score/selection gate was re-run after the shared-helper refactor: test result: ok. 4 passed; 0 failed, 3674 block scores and 133 selection rows. cargo fmt and the CI-equivalent PEGAINFER_CUDA_SM=80 cargo clippy --release --locked -p pegainfer-kernels --features qwen4 --all-targets -- -D warnings row are clean. The two trap binaries compile but were not run here; a trap poisons the context, so they must run last on a disposable GPU.

Performance

The kernel was tuned with a local harness that is not committed (the A/B numbers are the evidence; a scratch benchmark stays out of the tree). Same-session ABBA medians (10 iterations x 5 rounds), CUDA 12.9.1, one RTX 5060 (sm_120), 24 heads, head_dim 256, an all-visible mask (the kernel scores every token, so mask density does not change the work), including the attn_weights write:

rows x tokens	before	after	speedup	TFLOP/s
512 x 512	12.18 ms	6.93 ms	1.76x	0.53 -> 0.93
1024 x 1024	48.12 ms	26.67 ms	1.80x	0.54 -> 0.97
2048 x 2048	192.84 ms	111.55 ms	1.73x	0.54 -> 0.93

Two changes, measured one at a time: the 128-bit key load alone took 2048 from 192.8 to 133.3 ms (1.44x); two-head CTAs took it to 118.2 ms (1.63x); four-head CTAs reached 111.5 ms (1.73x).

Key load. Each lane's eight dims are 16 bytes and the row stride keeps them aligned, so one uint4 load replaces eight scalar bf16 loads. The baseline ncu profile reported 11.27B excessive sectors (77% of total sectors); the vectorized kernel has no uncoalesced-global warning.
Head grouping. One, two, or four query heads per CTA share the key and value scans, since all heads in a CTA map to one KV head. Four heads up to 2048 tokens, two up to 4096, one above, so the score row stays inside the 48 KiB default shared-memory cap.
Remaining bound. The four-head 2048 launch profiles at 93.6% memory/L1 and compute, 33.3% occupancy (shared-memory-limited: 32.77 KiB dynamic + 4.22 KiB static), 64 registers/thread, and a 12,288-CTA grid. The memory pipe is the limit; further head grouping is capped by the score-row shared memory, and a tiled/TMA or streaming design is a different change.

Tuning did not change the numerical contract: the attention gate is still 8/8 and case 3's parity counts are unchanged (295911/295920 weights; 36793/36864 output). compute-sanitizer memcheck reports 0 errors at 1024 and 2048 tokens (four-head path), and racecheck reports 0 hazards at 1024.

Limits

Prefill only, contiguous key array, 8192-token cap; paged KV, decode, and the streaming softmax are not covered. Nothing on the serving path calls this yet.
The four plain projections are outside the fixture by design; the indexer that produces the mask has its own oracle.
The fixture's positions are arange, so the key side cannot distinguish position from index. The query side is pinned by case 3 (tokens 2049..=2054).
Performance was measured on sm_120 only; the sm_89 path was not re-run after the tuning. The 2048-token four-head launch is benchmarked and memchecked, but the reference gate's four-head cases are the 64-token case; the 2055 case uses the two-head path. The full 8192-token cap uses the one-head path.
The measurement includes the full attn_weights write (192 MiB at 2048). A serving caller that only needs attn_out could skip that observable, but the component contract still writes both.
The Qwen4 sm_80 CUDA Clippy row now builds the feature in CI. The GPU correctness gate is still manual: PEGAINFER_REQUIRE_GPU=1 is required so a missing device cannot pass silently, and CI wiring with a GPU runner is a follow-up.

Refs #1103 #1105

sysuTobo · 2026-10-12T14:39:48Z

Code review

范围：2 文件、无运行时数据流（dump 脚本 + fixture）。逐文件认领：tools/accuracy/dump_qwen38fn_qsa_attn_golden.py（227 行）、test_data/qwen38fn-qsa-attn-golden.safetensors（33 张量 = 3 provenance + 9 权重 + 3 case × 7，header 与脚本写入逐项核对一致）。

正向路径（本机 torch 2.14.1+cpu / transformers 5.16.1 实跑复核）：

三条断言均非空转：零质量断言（tools/accuracy/dump_qwen38fn_qsa_attn_golden.py:194）是精确相等，实测 case2 有 4 个因果有效 (q,k) 对被稀疏丢弃、被 block 位置 softmax 质量恒为 0.0——attention 若无视稀疏 mask，任何 ≥ bf16 最小正数的质量泄漏即红；rowsum 断言（:197）实测最大偏差 2.1e-3，任何 >3% 的质量丢失/重复即红；expect_sparse 断言（:202-205）实测 dropped = 0/0/4，直接守卫 body「只有缩 budget case 丢因果有效 key」的声称。
结构陷阱逐项核对均有覆盖：Q+gate 融合（fixture q_proj [1024,512] = 4·128·2，融合维完整存下）；partial RoPE（rotary 64 作用于 head_dim 128 的前段，与参考 apply_rotary_pos_emb 的 rotary_dim 切片语义一致）；GQA group 2（参考 repeat_kv 被执行）；sigmoid 输出门（参考在 o_proj 前乘 sigmoid(gate)）；budget 截断（case2 budget 16 → block_topk 4，5 个完整块 > 4）。
fixture 权重覆盖完整：attention 模块全部 9 个参数（q/k/v/o_proj + q/k_norm + indexer 3 个）都在，消费者可离线重建。
重跑复现：31/33 张量与 committed fixture 逐字节相同，仅 case1_attn_out（maxabs 2.4e-4）/case2_attn_out（9.5e-7）有 bf16 ULP 级环境漂移（原 fixture 在 torch 2.5.1+cu121 dump，归约顺序随版本/设备变化）。

ID	级别	问题	证据
M1	M	`config_digest` 规范化串漏掉 budget 常量（`REAL_BUDGET`/`BOUNDARY_BUDGET`），与姊妹脚本（indexer golden 的 digest 含 `budget=`）约定不一致。改 `BOUNDARY_BUDGET` 会改变 case2 的稀疏语义与存储的 `case2_budget`，但 `config_sha256` 不变——按 digest 断言配置的消费者拿到假 provenance。	`tools/accuracy/dump_qwen38fn_qsa_attn_golden.py:126-132`
M2	M	body Validation 声称 "Fixture reloads and reproduces output + weights bit-for-bit from the stored weights"，但脚本没有任何 reload 校验——该声称目前只能人工复验。姊妹脚本的对应守卫是「独立重推导 == 参考 mask」交叉验证，本脚本没有等价物：权重名映射写错（如 q_norm/k_norm 互换）、漏存参数、序列化损坏、`attn_out` 出 NaN，都能通过现有全部断言并静默污染 fixture（NaN 会到未来 gate 才炸，且报错指向消费者而非 dump）。	`tools/accuracy/dump_qwen38fn_qsa_attn_golden.py:158-171,220-223`
L1	L	rowsum `atol=3e-2` 无标定依据（拍脑袋还是量过？）。实测噪声底 ~2.1e-3、最小有意义 bug（20 key 丢 1 个）偏差 ~5e-2，3e-2 恰在两者之间且余量 >10×——数值本身成立，但标定应写进注释，否则后人不敢动也不知道能不能收紧。	`tools/accuracy/dump_qwen38fn_qsa_attn_golden.py:197`
L2	L	与姊妹脚本最终态的小分歧：`positions` 存 [1,S]（姊妹为 [S]）；无 per-case `ratio` 张量（姊妹有——但 ratio 已在 digest canon 里，姊妹那份实为冗余，本脚本不存反而更干净）；`FORMAT_VERSION` 缺「gate 加载时断言」注释。均不影响消费者语义。	`tools/accuracy/dump_qwen38fn_qsa_attn_golden.py:63,211`
L3	L	家族级缺口：`RMS_EPS`（以及 SEED、CASES 表）不在任何一方的 digest canon 里——改 eps 两个 fixture 的 digest 都不红。属 #1150-#1155 系列共同约定，不宜单改一个造成分歧，记录为已知边界。	`tools/accuracy/dump_qwen38fn_qsa_attn_golden.py:50,126-132`
L4	L	「bit-for-bit 复现」只在同环境成立：跨 torch 版本重跑 `attn_out` 有 ≤2.4e-4 的 bf16 ULP 漂移（见上）。docstring 应告知 gate 消费者对 `attn_out` 用容差比较而非逐位，避免未来 parity gate 把环境噪声当回归追。	`tools/accuracy/dump_qwen38fn_qsa_attn_golden.py:17-22`

最小加固：修 M1（digest canon 补两个 budget 常量，与姊妹脚本对齐）、M2（脚本内固化 reload-and-reproduce：strict load 证明权重映射完整 + 逐位复现 out/weights/mask，顺带 isfinite(attn_out) 显式守卫）；L1/L4 补注释与 docstring；L2/L3 记录为已知边界（L3 留给家族层面统一处理）。修 M1 会改变 config_sha256 张量 → 属有意变更 fixture，将重生成替换。

sysuTobo · 2026-10-12T14:45:33Z

修改措施

Pushed 33ac1e90。

M1 已修：config_digest canon 补入两个 budget 常量（budget={REAL_BUDGET};bbudget={BOUNDARY_BUDGET}），并加 docstring 说明 budget 属于消费者要断言的配置面（它决定每 case 的 block_topk 与 mask 保留范围），与姊妹脚本（indexer golden 的 digest 含 budget=）对齐。tools/accuracy/dump_qwen38fn_qsa_attn_golden.py:134-144。config_sha256 由 997cf3da… 变为 f9d2bc69…（有意变更）。
M2 已修：脚本内固化 reload-and-reproduce（body Validation 声称从此可复验）：保存后重载 fixture，① provenance 三张量解码断言（format_version / config_sha256 / ref_transformers）；② 用重载的 9 个权重张量 load_state_dict strict 重建模块——证明存储的权重集完整且命名映射正确（漏存/写错名会当场红）；③ 逐 case 断言存储的 seq_len/budget 与 CASES 表一致，并从存储的 hidden/positions 重放，attn_out / attn_weights / mask 三者与文件内逐位相等。case 执行路径抽成 run_case 供 dump 与重放共用，保证重放走的是同一条参考路径。另加显式 isfinite(attn_out) 守卫（NaN 不再静默入库）。tools/accuracy/dump_qwen38fn_qsa_attn_golden.py:151-164,221,247-282
L1 已修：rowsum atol=3e-2 处补标定注释——实测噪声底 max |sum−1| ≈ 2.1e-3（bf16 行和量化），最小有意义 bug（20 key 丢 1 个）移动 ~5e-2，3e-2 在两者之间、余量 >10×。数值本身不改。tools/accuracy/dump_qwen38fn_qsa_attn_golden.py:215-220
L2 部分回执，其余接受为已知边界：FORMAT_VERSION 补「gate 加载时断言」注释（与姊妹脚本对齐，:70-71）；positions [1,S] 形状与「不存 per-case ratio」维持现状——ratio 已在 digest canon 中，per-case 冗余存储反而引入两个真相源，改形状/加张量都要动 fixture，收益不抵。
L3 不修：RMS_EPS（及 SEED、CASES 表）不进 digest 是 feat(qwen4): add Qwen3.8-Flash-Next QSA indexer golden #1150-feat(qwen4): add the PLE n-gram golden and its host/device readers #1155 系列的共同约定缺口（姊妹脚本同样不含），单改一个 fixture 会造成系列内 digest 语义分歧；记录为家族级已知边界，留待系列层面统一处理。
L4 已修：docstring 增补逐位契约段——bit-for-bit 仅同环境成立；跨 torch 版本/设备 attn_out 有 bf16 ULP 级归约漂移（实测 ≤2.5e-4），gate 消费者应以容差而非精确相等回放 attn_out。tools/accuracy/dump_qwen38fn_qsa_attn_golden.py:24-29

复验：fixture 以 torch 2.14.1+cpu / transformers 5.16.1 / CPU 重生成（33 tensor，布局与张量集不变），参考输出与原 dump 存在 bf16 ULP 级环境漂移（≤2.5e-4）——与原 committed fixture 相比 31/33 张量逐字节相同，差异仅 config_sha256（M1 有意变更）与 case1_attn_out（maxabs 2.4e-4）/case2_attn_out（9.5e-7）的已知环境漂移。provenance = 5.16.1 / config_sha256=f9d2bc69…；3 case dump 断言全过（blocked_frac 0.438/0.450/0.485，因果有效丢弃对 0/0/4，被 block 位置 softmax 质量恒 0.0，rowsum 最大偏差 2.1e-3，attn_out 全有限），reload 重放 3/3 逐位自洽；修复后脚本同环境二次重跑与 committed fixture 逐字节相同。

sysuTobo · 2026-10-14T13:59:06Z

+        # The captured o_proj input must reproduce the module's returned output.
+        assert torch.equal(module.o_proj(captured["o_in"]), out), f"case {i}: o_proj replay mismatch"
+
+        # Effective sparse+causal mask = the reference's own indexer mask AND causal.
+        selected_bool = captured["indexer_mask"] == 0
+        eff = causal_bool & selected_bool  # [1,1,S,S]
+
+        # Self-check: blocked (sparse-dropped or future) positions carry zero softmax mass.
+        blocked_mass = (attn_weights.float() * (~eff).float()).abs().max().item()
+        assert blocked_mass == 0.0, f"case {i}: blocked positions leaked softmax mass ({blocked_mass})"
+
+        # Drop accounting: causally-valid keys the sparse budget discarded, per row.
+        causal_count = causal_bool[0, 0][row_idx].sum(-1)
+        kept_count = eff[0, 0][row_idx].sum(-1)
+        drops = (causal_count - kept_count).tolist()
+        assert all(d == 0 for d in drops[:2]) if seq_len == 2055 else all(d == 0 for d in drops), (
+            f"case {i}: unexpected sparse drop on a short case: {drops}"
+        )
+        if seq_len == 2055:
+            assert all(d == 0 for d in drops[:2]) and all(d == 4 for d in drops[2:]), (
+                f"case {i}: real boundary drops unexpected: {drops}"
+            )


多此一举

sysuTobo · 2026-10-16T16:31:33Z

Code review

审计对象：0a552e22（本 PR 的 consumer 提交），基线 21dbcee8。
数据流：fixture 张量 → host Case（形状/范围断言）→ device buffer → norm_rope_keys + attn 两个 kernel → 与 attn_weights / attn_out 逐元素比较（含三条反向对照）。
覆盖清单：csrc/qwen38fn/qwen38fn_qsa_attn_forward.cu(新, 412) · src/ffi/qwen38fn_attn.rs(新, 33) · src/ops/qwen38fn_attn.rs(新, 198) · tests/qwen38fn_qsa_attn_parity.rs(新, 1205) · src/ffi.rs(+4) · src/ops.rs(+6) · Cargo.toml(+7) · build.rs(+21) · Cargo.lock(+3)，无遗漏文件。
验证（RTX 5060 / sm_120，pegainfer-dev:cu132，CUDA 12.9.1）：cargo test --release --locked -p pegainfer-kernels --features qwen4 --test qwen38fn_qsa_attn_parity → 7 passed; 0 failed；cargo fmt --all -- --check 干净；cargo clippy --release --locked -p pegainfer-kernels --all-targets 在带 qwen4 与默认 feature 两种配置下都无告警；cargo check --locked -p pegainfer-kernels（默认 feature）通过。

ID	级别	一句话	证据(文件:行)
F1	M	容差注释里的 ulp 计数用粗口径 `2^-8·	want
F2	M	token cap 只钉了拒绝侧（`MAX_TOKENS + 1`），cap 本身那次 launch（32 KiB 动态 shared，本 kernel 的最大请求）从未真正发射过	pegainfer-kernels/tests/qwen38fn_qsa_attn_parity.rs:813
F3	M	同一个 8192 上限有两个真相源（Rust `MAX_TOKENS` / CUDA `kMaxTokens`），漂移只会 fail-closed，但原先没有任何东西拦住漂移	pegainfer-kernels/src/ops/qwen38fn_attn.rs:48 · pegainfer-kernels/csrc/qwen38fn/qwen38fn_qsa_attn_forward.cu:96
F4	M	频率表与兄弟 indexer kernel 不同源：本 kernel 用 double `pow` 后一次性取整，对方用 `powf`（规范 8 ulp，在 position 2054 上等于 1e-3 rad 的角度误差）	pegainfer-kernels/csrc/qwen38fn/qwen38fn_qsa_attn_forward.cu:146
F5	L	`.cu` 头注释指向「beside this file」的 indexer kernel，而该文件在本分支并不存在（来自兄弟 PR）——假指针	pegainfer-kernels/csrc/qwen38fn/qwen38fn_qsa_attn_forward.cu:54
F6	L	覆盖断言弱于文档声明：文档说 case 3 是唯一的 row-subset、唯一非因果的 case，断言却只写 `>= 1`	pegainfer-kernels/tests/qwen38fn_qsa_attn_parity.rs:1059
F7	L	`assert_provenance` 返回 manifest，四个调用点全部丢弃	pegainfer-kernels/tests/qwen38fn_qsa_attn_parity.rs:155
F8	L	NaN 断言把成因唯一地写成 sentinel，而「整行无可见 key」的 softmax 同样给 NaN	pegainfer-kernels/tests/qwen38fn_qsa_attn_parity.rs:598
F9	L	key 侧 position 查表在本 fixture 上不可判（`positions` 是 `arange`，`positions[s] == s`），gate 原先没有写出这个覆盖空洞	pegainfer-kernels/tests/qwen38fn_qsa_attn_parity.rs:33
F10	L	`rows[row]` 越界 clamp 是文档化契约但从未被执行（`load_case` 已断言 rows 全在界内）	pegainfer-kernels/csrc/qwen38fn/qwen38fn_qsa_attn_forward.cu:229
F11	L	ops wrapper 的 8 条 `ensure!` 长度检查没有测试（需要 device + 欠尺寸 buffer 才能触发）	pegainfer-kernels/src/ops/qwen38fn_attn.rs:117

修改措施

F1 已修：统一成精确 quantum 口径。最坏权重改写为「1.68e-2 = 该量级上 8 个 quantum」（0.0018157959 的 quantum 是 2^-18，偏差 3.052e-5 正好 8 个），最坏 output 元素改写为「2.19e-4 对 2.67e-4，50 个 quantum / 18%，但绝对值只有 4.8e-5」，并写明 REL_FLOOR 之下相对读数描述的是抵消而不是 kernel（tests:73, 86, 100-113）。
F2 已修：新增 the_token_cap_is_a_launch_it_can_actually_form（tests:901-996）——tokens 恰为 8192、一行 query、全可见 mask，断言调用返回 Ok、196608 个权重全部有限、24 个 head 各自的 8192 个权重之和在 2e-3 内等于 1（权重求和是「最大 shared 配置的 launch 不但合法而且真的跑完整行」的判据，而不只是「没报错」）。实测通过。
F3 已修：两侧现在被同一个常量双向钉住——拒绝侧用 MAX_TOKENS + 1 探 kernel 的编译常量（tests:806-818），接受侧就在 MAX_TOKENS 上跑一次（F2 的测试）。单独抬高任一侧都会让其中一个测试红。这层关系写进了 MAX_TOKENS 的 doc（ops:41-48）。
F4 就地记录 + 明确留给后落地方：取舍是刻意的，不改成 powf。理由写在 .cu:146-152——参考在 host 上用 glibc powf（正确舍入）建表，device 的 powf 规范误差 8 ulp，在 position 2054 上是 1e-3 rad 的角度误差 ≈ cos 的 1/8 个 bf16 quantum，足以把「4 个 case 里 3 个逐位相同」推走。抽公共 header 的那一方应取 double 版本，这条建议写在 .cu:54-59。
F5 已修：改成按文件名指 qwen38fn_qsa_indexer_score.cu，并写明它由兄弟改动落进同一目录、helper 的重复是刻意的（抽公共 header 只是搬动重复而不是消除），lifting 留给后落地的一方（.cu:54-59）。
F6 已修：subset / sparse 改为 assert_eq!(…, 1)，与既有 boundary == 1 同一口径，fixture 若新增 case 会立刻要求作者复核覆盖声明（tests:1059-1073）。
F7 已修：assert_provenance 不再返回值（tests:149-155）。
F8 已修：断言措辞点名两种成因——sentinel 未被覆盖，或整行无可见 key 的 softmax（tests:596-601）。后者另有 load_case 的前置断言挡着（tests:344-353）。
F9 就地更正（记录，不改 fixture）：任务禁止改 dump 脚本 / 重生成 fixture，所以这条不可判性写进 gate 头注释（tests:33-38）与 PR body 的 “Still ungated”，并写明 query 侧是可判的（case 3 的 6 行是 token 2049..2054，按 slot 取 position 不可能对）。明确留给 Stage 3 的 paged KV / decode fixture 去覆盖。
F10 明确保留：与兄弟 indexer kernel 同一契约（clamp 而非拒绝），语义写在 .cu:227-233（「a broken caller reads a wrong row rather than memory it does not own」）。分支未被执行是事实，记录为「文档化防御契约，无 gate」，不为它加测试。
F11 明确留给调用方接入时：这 8 条 ensure! 是同一契约在 Rust 侧的更早失败点、同样 fail-closed；触发它需要 device + 刻意欠尺寸的 buffer，价值低于 CUDA 侧那份 device-free 的拒绝测试（tests:711-886，14 条拒绝 = 9 条几何/范围 + 5 条空指针，外加 2 条零长度正对照）。留给 qwen4 model crate 接入时的调用方测试。

github-actions · 2026-10-20T10:06:33Z

This pull request has merge conflicts with main. Please rebase it so it can be reviewed and merged.

github-actions · 2026-10-23T15:10:06Z

This pull request has merge conflicts with main. Please rebase it so it can be reviewed and merged.

Signed-off-by: bohaotu <bohaotu@pony.ai>

… pass - config_sha256 canon now covers REAL_BUDGET/BOUNDARY_BUDGET so a budget change redds the provenance digest (aligns with the indexer golden). - Reload-and-reproduce after save: strict state-dict load proves the stored weights are complete and correctly named; every case replays attn_out / attn_weights / mask bit-for-bit in-process. Docstring records that cross-environment replay is tolerance-bound (bf16 ULP drift <= 2.5e-4). - Explicit isfinite guard on attn_out and a calibration note for the 3e-2 row-sum atol (measured noise floor ~2e-3). - Fixture regenerated for the new digest: 31/33 tensors byte-identical to the previous dump; case1/case2 attn_out carry the known <= 2.4e-4 bf16 ULP environment drift (torch 2.14.1+cpu vs the original 2.5.1+cu121). Signed-off-by: bohaotu <bohaotu@pony.ai>

Rewrite the QSA sparse-attention oracle to use the real frozen checkpoint geometry (hidden 2560, 24 Q / 2 KV heads, head_dim 256, partial rotary 0.25 -> 64 dims, GQA group 12, real indexer budget 2048 -> 512 blocks) instead of the reduced proxy, and hit the real boundary at S=2055. The plain q/k/v/o GEMM weights (98 MB) are not stored; their outputs are captured with forward hooks, and only the small math-carrying weights (q_norm/k_norm, index_qk_proj, q/k_layernorm) are kept. Each case records only the evaluated query rows; metadata moves to a single JSON key (format_version 2). Signed-off-by: bohaotu <bohaotu@pony.ai>

…producer case*_q_proj/k_proj/v_proj held the captured projection outputs, not weights, so the names read like the GEMM weights the golden deliberately does not store. Rename them queries/keys/values (matching the merged indexer golden). Signed-off-by: bohaotu <bohaotu@pony.ai>

…oad pass The fixture stored the indexer's `index_qk_proj`/`q_layernorm`/`k_layernorm` (3.3 MB) that nothing could exercise: their input, the layer's hidden state, is not part of the fixture, and the sparse mask they produce is stored as data. The indexer has its own oracle. The 38 surviving tensors are bit-identical to the previous fixture and the metadata JSON is unchanged. Also drops `replay()`/`verify()` and the `--verify-only` flag. That pass re-implemented the reference forward to show the fixture reproduces itself, a property the captured values do not depend on -- and the only re-implementation this oracle carried. The in-run invariants stay: the `o_proj` replay, zero blocked softmax mass, and the per-row drop accounting. Signed-off-by: bohaotu <bohaotu@pony.ai>

…d asserts The case loop asserted two properties the architecture already guarantees -- zero softmax mass on blocked positions, and the budget's drop pattern -- which this fixture has no business re-checking: it stores the reference's own output, so its job is numerics. What remains is the capture check that `o_proj`(captured input) reproduces the module's returned output. The docstring is compressed and no longer claims the dump asserts those drops. Fixture bytes are unchanged. Signed-off-by: bohaotu <bohaotu@pony.ai>

The `o_proj`(captured input) == module output check guarded nothing a consumer needs: the pre-hook captures `o_proj`'s input by construction, so the stored `attn_out` is already the right tensor. The script is now assert-free, and the docstring no longer claims an in-run check. Fixture bytes are unchanged. Signed-off-by: bohaotu <bohaotu@pony.ai>

The QSA sparse-attention golden this branch adds has had nothing to read it, and reader: the attention half of the three units pegainfer-project#1105 splits QSA into, in `pegainfer-kernels/csrc/qwen4/` behind the kernels crate's own `qwen4` feature, so the model crate stays free of CUDA. Two kernels against the frozen transformers-5.16.1 semantics. One `(1 + w)` RMSNorms and partially rotates every token's two key heads once, because all twelve query heads of a KV head read the same prepared key; the other is one CTA per (query row, head) and does that head's own norm and rotation, the 256-wide score against every key, the float32 softmax, the weights, and the gated context. It consumes the visibility mask instead of deriving it: the mask is the indexer's selection AND causality, which is the previous unit's business. Rounding is reproduced at every point the reference rounds — the norms cast once, cos/sin are cast to bf16 and rotate in bf16, the score is the bf16 GEMM result rounded before the exact 1/16 scale, the weight is rounded before the context GEMM reads it, and the gate's sigmoid rounds before the final product. The float32 reduction *order* is this kernel's own, and that is the whole of the residual: three of the four cases are bit-identical on both observables, and in the 2055-token case 532 of 295920 weights sit one bf16 quantum off and 56 more than one. It follows pegainfer-project#1105's D3-B: the sm_89-first portable subset — no TMA, no cluster, no wgmma — and SIMT shuffle reductions rather than `mma.sync` bodies, because one CTA's score is a 256-wide dot and its context a 256-wide column reduction. The gate replays `test_data/qwen38fn-qsa-attn-golden.safetensors` and holds its recorded `format_version` / `transformers_version` / `config_sha256`, plus the frozen config's own attention geometry, against this build — a golden from another checkpoint fails loudly instead of skipping. Three negative controls keep it from being vacuous: the stored mask replaced by plain causality, where the three causal cases must still match and the boundary case must not; the two KV heads exchanged, a wrong GQA mapping; and the gate saturated, which must leave the weights alone and break only the output. Two more tests pin the entry point's refusals device-free and run one call at exactly the token cap, the widest shared memory this kernel asks for. Nothing on the serving path calls this yet; the qwen4 model crate will when the text graph lands. Refs pegainfer-project#1103 pegainfer-project#1105 Signed-off-by: bohaotu <bohaotu@pony.ai>

- Validate QsaAttnGeometry against the compiled constants before sizing buffers, and make an empty launch a no-op like the sibling wrappers. - Trap rows/positions outside the key array instead of clamping, with one trap binary per half of the guard. - Share the QSA norm/partial-RoPE rounding helpers and the attention cos/sin table with the indexer; the attention residual drops from 0.42 to 0.17 of the bound because the trig source is now the shared table. - Assert bit-identical short cases, add an empty-launch test, cover every null pointer, and correct the fully-hidden-row claims. Dropping the sigmoid bf16 rounding now fails the case0 exact assertion. - Update the qwen4 feature comment and the public bring-up doc. Signed-off-by: bohaotu <bohaotu@pony.ai>

- Load each lane's eight key dims with one 128-bit load; the old scalar pattern requested 77% excessive sectors on a 2048-token prefill. - Group one/two/four query heads per CTA so they share the key and value scans; four heads up to 2048 tokens, two up to 4096, one above, keeping the score row inside the 48 KiB default shared-memory cap. - Add a manual benchmark example. Same-session ABBA at 2048 tokens: 192.1 -> 111.5 ms median (1.73x); 1024: 48.1 -> 26.7 ms (1.80x); 512: 12.2 -> 6.9 ms (1.76x). - Attention gate 8/8 still passes; case3 parity counts unchanged. compute-sanitizer memcheck 0 errors at 1024/2048, racecheck 0 hazards at 1024. The four-head launch is shared-memory-limited at 33% occupancy. Signed-off-by: bohaotu <bohaotu@pony.ai>

The A/B numbers stay in the PR description; the harness was scratch and should not ship. Signed-off-by: bohaotu <bohaotu@pony.ai>

The merged `shared/ffi_guard.cuh` exposes `consume_last_cuda_error` in the global namespace. The QSA translation units each defined an anonymous-namespace helper with the same name, which nvcc 13 reports as an ambiguous overload (the Qwen4 sm_80 CI row). Give each file's helper a distinct name; behavior is unchanged. Verified with PEGAINFER_CUDA_SM=80 cargo clippy --release -p pegainfer-kernels --features qwen4 --all-targets -- -D warnings. Signed-off-by: bohaotu <bohaotu@pony.ai>

github-actions · 2026-11-03T04:31:55Z

This pull request has merge conflicts with main. Please rebase it so it can be reviewed and merged.

Use the shared `consume_last_cuda_error` from `shared/ffi_guard.cuh` in the Qwen4 translation units; the merge's new global helper made the local anonymous-namespace copies ambiguous under nvcc 13.

github-actions Bot added the enhancement New feature or request label Oct 12, 2026

This was referenced Oct 12, 2026

feat(qwen4): add the PLE n-gram golden and its host/device readers #1155

Open

feat(qwen4): add Qwen3.8-Flash-Next QSA indexer golden #1150

Merged

sysuTobo closed this Oct 14, 2026

sysuTobo reopened this Oct 14, 2026

sysuTobo commented Oct 14, 2026

View reviewed changes

sysuTobo marked this pull request as ready for review October 14, 2026 14:07

FeathBow mentioned this pull request Oct 15, 2026

[Model] Qwen3.8-Flash-Next roadmap #1103

Open

sysuTobo marked this pull request as draft October 16, 2026 13:45

github-actions Bot added the needs-rebase Conflicts with main, needs a rebase label Oct 20, 2026

sysuTobo force-pushed the feat/qwen38fn-qsa-attn-golden branch from 0a552e2 to 19bc006 Compare October 21, 2026 02:57

github-actions Bot added qwen4 Qwen4-architecture model crate (pegainfer-qwen4) needs-rebase Conflicts with main, needs a rebase and removed needs-rebase Conflicts with main, needs a rebase labels Oct 21, 2026

sysuTobo added 8 commits November 1, 2026 22:24

feat(accuracy): add Qwen3.8-Flash-Next QSA sparse attention golden

fe438de

Signed-off-by: bohaotu <bohaotu@pony.ai>

sysuTobo force-pushed the feat/qwen38fn-qsa-attn-golden branch from 19bc006 to 0742cae Compare November 1, 2026 14:38

github-actions Bot removed the needs-rebase Conflicts with main, needs a rebase label Nov 1, 2026

sysuTobo changed the title ~~feat(qwen4): add Qwen3.8-Flash-Next QSA sparse-attention golden~~ feat(qwen4): add the QSA sparse-attention kernel and its golden Nov 2, 2026

sysuTobo marked this pull request as ready for review November 2, 2026 14:21

sysuTobo added 2 commits November 3, 2026 12:21

chore(qwen4): drop the committed attention benchmark example

2a76885

The A/B numbers stay in the PR description; the harness was scratch and should not ship. Signed-off-by: bohaotu <bohaotu@pony.ai>

github-actions Bot added the needs-rebase Conflicts with main, needs a rebase label Nov 3, 2026

Merge origin/main into feat/qwen38fn-qsa-attn-golden

bedc7ff

Use the shared `consume_last_cuda_error` from `shared/ffi_guard.cuh` in the Qwen4 translation units; the merge's new global helper made the local anonymous-namespace copies ambiguous under nvcc 13.

github-actions Bot removed the needs-rebase Conflicts with main, needs a rebase label Nov 3, 2026

Navigation Menu

Repository navigation

feat(qwen4): add the QSA sparse-attention kernel and its golden - #1151

feat(qwen4): add the QSA sparse-attention kernel and its golden#1151
sysuTobo wants to merge 13 commits into
pegainfer-project:mainfrom
sysuTobo:feat/qwen38fn-qsa-attn-golden

sysuTobo commented Oct 12, 2026 •

edited

Loading

Uh oh!

sysuTobo commented Oct 12, 2026

Uh oh!

sysuTobo commented Oct 12, 2026

Uh oh!

sysuTobo Oct 14, 2026

Uh oh!

sysuTobo commented Oct 16, 2026

Uh oh!

github-actions Bot commented Oct 20, 2026

Uh oh!

github-actions Bot commented Oct 23, 2026

Uh oh!

github-actions Bot commented Nov 3, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

sysuTobo commented Oct 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

What changed

Semantics

Validation

Performance

Limits

Uh oh!

sysuTobo commented Oct 12, 2026

Code review

Uh oh!

sysuTobo commented Oct 12, 2026

修改措施

Uh oh!

sysuTobo Oct 14, 2026

Choose a reason for hiding this comment

Uh oh!

sysuTobo commented Oct 16, 2026

Code review

修改措施

Uh oh!

github-actions Bot commented Oct 20, 2026

Uh oh!

github-actions Bot commented Oct 23, 2026

Uh oh!

github-actions Bot commented Nov 3, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

sysuTobo commented Oct 12, 2026 •

edited

Loading