[sysuTobo]

Summary

QSA unit three of #1105, after the scoring unit (#1160): turn the indexer's per-(row, block) scores into the token-level visibility mask, in native CUDA. csrc/qwen4/qwen4_qsa_indexer_select.cu adds one CTA per query row that writes u8 [query_rows, tokens], row stride tokens — the shape #1105's sparse attention consumes. Wrapper qsa_indexer_select_into in src/ops/qwen4.rs, FFI in src/ffi/qwen4.rs, selection gate in the select module of tests/qwen4_qsa_indexer_score_parity.rs, and a bring-up doc update. The scoring stage's parity file and its two trap binaries are byte-for-byte as on main; only the selection module is added beside them.

Selection is an exact 8-bit radix select over the row's sortable f32 keys: four histogram passes find the keep-th score, the blocks above it are kept, and the threshold-equal blocks fill the rest in ascending block-index order (the canonical tie rule). A row under the budget keeps every block and writes the visible prefix in one pass. This replaces 512 masked argmax passes — 0.9 ms at the first truncating row, 18 ms at 64k and 86 ms at 256k — with 20.8 / 67 / 298 µs on one L20.

Exact semantics

For a query row at position p (token 0 is the first visible token — the scoring pass's grid contract):

• blocks = (p + 1) / 4 complete micro-blocks, tail = (p + 1) % 4 tokens at 4*blocks ..= p;
• keep = min(budget / 4, blocks), with budget a positive multiple of 4 keeping at most 512 blocks;
• the keep highest block scores are kept under the canonical total order descending score, ascending block index;
• the mask is 1 exactly for the kept blocks' four tokens each and for the tail. The tail is unconditional, and causality is inherent because a block always covers tokens below p.

score_cols is an explicit argument, like the score stage: scores.len() may exceed query_rows * score_cols (spare capacity is allowed) and the row stride stays score_cols.

A row under the budget keeps every complete block and one pass writes its visible prefix. Otherwise the selector maps each finite score to a monotonic u32 key (-0.0 folded to +0.0, so both zeros tie), runs four 8-bit histogram passes to find the keep-th key, keeps every key above it, and fills the remaining slots from the equal keys in ascending block-index order. A non-finite score inside a row's block prefix traps rather than being skipped: the reference score is a relu sum, so a poison means the indexer input is already broken, and silently dropping that block would hand the attention softmax a shorter visible set with no trace (an all-masked row is undefined there). The wrapper treats an empty launch (rows.len() == 0 or tokens == 0) as a no-op, matching the score stage; the C entry still refuses one reached directly. The wrapper refuses a malformed budget and a short buffer; the kernel repeats the budget and score-row-width checks and traps on out-of-range rows entries.

No gate knob

The production kernel, its C ABI and the ops surface carry no mutation variant. The five mutation controls in the table below were measured by temporarily patching the kernel, as the #1160 review did; they are not shipped code.

What the gate measures

cargo test --release -p pegainfer-kernels --features qwen4 \
  --test qwen4_qsa_indexer_score_parity -- --nocapture

Selection is comparison-only over the fixture's recorded scores, so there is no tolerance: strict-cut rows must match the reference mask bitwise. A block tied with the last kept one is an equally valid torch.topk pick, so a tie row is held to: both masks structurally valid (causal prefix, complete 4-token blocks only plus the exact tail, exactly min(K, n)*4 + tail bits set), every selected block scoring at least every excluded one, and the same number of tied-at-cut picks. Tie rows are counted and reported, never skipped.

Measured on one L20 (sm_89, CUDA 12.9 plus the flashinfer tree's CCCL, 256 threads): 5/5 tests pass; 133 rows, 129 compared bitwise, 4 compared up to a tie, 0 mismatches; the 4 tie rows are case 6 at cut score 0.0 with 31–43 tied blocks; the fixture covers 128 all-kept and 5 truncating rows.

The five mutation controls, measured by temporary patches:

control	differing bits	what it proved
drop the tail	201	exactly the fixture's tail tokens
keep one fewer	460	exactly the last kept block on every keeping row
keep one more	20	exactly one added block on the five truncating rows
ignore causality	2470	exactly the columns past each row's own token
invert the order	40	two blocks per truncating row; every dropped block recovered

Performance, one row, budget 2048 (old = idle host p50; new = nsys kernel median at 256 threads):

case	masked argmax	radix select
all-kept 2k x1	2.5 µs	1.4 µs
trunc 2k+ x1 (513 blocks)	0.89 ms	20.8 µs
trunc 8k x1 (2048 blocks)	1.37 ms	34.4 µs
trunc 64k x1 (16384 blocks)	18.1 ms	67.0 µs
trunc 256k x1 (65536 blocks)	85.9 ms	298 µs
trunc 64k x8	20.4 ms	68.2 µs

Thread sweep on the shared node (nsys kernel medians): 128 threads is faster at 2k/8k (idle host p50 17.9 / 26.9 µs against 20.8 / 34.4 µs at 256), 512 is faster at 256k (187 against 298 µs); 256 is the balanced default. These numbers were measured with a temporary benchmark binary that does not ship.

Device-free coverage: the select module keeps one test for the C entry's guards (null pointers, zero/negative extents, malformed budget, overly narrow score row, oversized bitmap), which answer CUDA_ERROR_INVALID_VALUE before any pointer is formed. The Rust-side guards are inline in qsa_indexer_select_into (bad budget, short buffers, spare capacity, overflowing products).

Verification

• Static checks: cargo fmt --all -- --check, cargo check -p pegainfer-kernels --features qwen4 --tests, and cargo clippy -p pegainfer-kernels --features qwen4 --tests -- -D warnings pass at this head. The build used CUDA 12.9 plus the flashinfer tree's CCCL and a build-only guard for the CUDA-13-only cudaErrorLossyQuery case; the production CUDA 13.2 build is the CI/dev image and was not re-run here.
• The parity gate passed 5/5 at this head (CUDA 12.9, sm_120): 133 rows, 129 bitwise, 4 up to ties, 128 all-kept, 5 truncating.
• compute-sanitizer memcheck on the 256k case: 0 errors; racecheck on the 2k+ case: 0 hazards. (Measured at the previous head; this head renames the bitmap's helpers and trims comments, and the selector logic is unchanged.)
• nsys kernel medians: 25 launches of the 256k case, median 298 µs (min 297.5, max 298.8); the shared node's other tenants are why the thread sweep uses nsys medians.
• ncu is unavailable on the shared pod: ERR_NVGPUCTRPERM (no GPU performance-counter permission).

Boundaries not gated

• No CI job enables the kernels crate's qwen4 feature, so CI neither builds nor runs any of this; the gate is manual (PEGAINFER_REQUIRE_GPU=1).
• The fixture replays recorded scores. Selection over scores the scoring kernel actually produced is covered only by the scoring gate upstream, not end to end.
• Every tie in the fixture sits at cut 0.0; a tie at a non-zero cut is unexercised.
• Case 6 evaluates positions 2049–2054 only, so that length is not exercised below 2049.
• The selector derives a row's visible list from its position (un-evicted causal text, the scoring pass's token-0 grid contract); a prefix-evicted visible list is out of contract.
• The mask's consumer — the sparse attention unit of [Model] Qwen3.8-Flash-Next: QSA sparse attention #1105 — is separate and not gated here.
• The selector's two trap paths (out-of-range rows, non-finite score) are not exercised by a test in this PR; the score stage's trap binaries are unchanged from main and were not re-run here.
• The performance table is one row per case (plus one 8-row case) on one shared L20; the production CUDA 13.2 build is not re-run.

Refs #1103 #1105.

[sysuTobo]

复核 + 一处 API 收紧

在 agent 交付后按数据流与门禁纪律复核了一遍：

• 数值复现：PEGAINFER_REQUIRE_GPU=1 下 gate 3 passed；133 rows = 129 逐位 + 4 按并列比较 + 5 truncating；五条对照分别差 201 / 460 / 20 / 2470 / 40 bits，与子 agent 报告一致。
• tie 判据：4 条并列行都在 case 6、cut=0.0，两侧都要求"同分块取同数量 + 按分数选择"，不是简单跳过，也不是放宽成恒过。
• trap 覆盖：非有限分数与越界 row 各一个独立进程的 trap 测试，均通过。

修改：把 gate 的 variant 从生产入口剥离（ed2479be）

原实现的唯一公共入口把 variant 作为必填参数，生产调用方必须写 0 且可能误选到负控路径。按 PLE 块的先例拆成两条：

• qsa_indexer_select_into —— 生产入口，固定 QSA_SELECT_PRODUCTION，无旋钮；
• qsa_indexer_select_variant_into —— gate 专用，带 variant，非法值拒绝。

parity gate 的正常路径现在走生产入口，负控走 variant 入口；核本身未改。
复核后：clippy -D warnings、fmt、默认 feature cargo check、6 条 device-free 守卫全绿；device-free 的 parity 子测试通过。

环境副作用（重要）

跑完两个 trap 测试后，本机整块 GPU 不可初始化（host 侧 cuInit 返回 3，容器内 cudaErrorInvalidContext），不是本 PR 的逻辑问题，但会阻断复跑：

• 两个 trap 测试头已加警告：最后跑，或用一次性 GPU；
• 本次 API 拆分后的 device parity 复跑要等设备 reset/reboot；拆分前同一核与同一 gate 已连续两次绿（数据见上），拆分只增加一层转发；
• 已合入的 feat(qwen4): score QSA token-indexer blocks in native CUDA #1160 里也有同类 trap 测试（qwen4_qsa_indexer_*_trap.rs），建议 CI 若将来启用 qwen4，把这类测试放在 job 末尾或独立的一次性 runner。

仍开放的边界（与 PR 正文一致）

无 CI job 启用 qwen4；只重放记录分数；并列只在 cut=0 出现；case 6 只覆盖 2049–2054；prefix-evicted 可见列表不在契约内。

[github-actions]

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

[FeathBow]

The selection reads right against the reference's own mask, and the score gate moved into the shared file with its assertions unchanged. A few things before this leaves draft, details inline.

The description no longer matches the branch: the gate is now qwen4_qsa_indexer_parity, the traps run from qwen4_qsa_indexer_trap, the wrapper treats an empty launch as a no-op, and the PR also restructures the #1160 tests. It also says the device rerun after the API split is still pending, and the kernel changed after that. Could you rerun the gate and the trap binary at this head and update the numbers and commands?

[FeathBow]

The score stage takes score_cols explicitly and accepts a buffer longer than rows * score_cols; here the stride is inferred from the length. Handing this stage the score stage's buffer with spare capacity gives a larger stride, passes both checks, and selects from the wrong rows with no error. Could this take score_cols like the score stage?

[FeathBow]

These variants stay in the production kernel and its C ABI, with six constants exported from ops, for something only the gate uses. The table they produce is good evidence, and temporary patches give the same evidence, which is how the #1160 review measured its gate. Could the knob go, with the table kept in the description?

[FeathBow]

Only the exit status is checked. If --exact trap_case stops matching (a rename, a harness flag), the child runs zero tests and exits 0, so every case passes without reaching a trap. Checking that the child reports one passed test would close that.

[FeathBow]

Each pick is a full pass over the row's blocks plus two barriers: up to 512 passes over 65536 blocks per row at the checkpoint's 262144 positions, and blocks passes for every row shorter than the budget, where every block is kept anyway. Has this been timed at long context? A fast path for rows that keep all their blocks is cheap at least.

[sysuTobo]

Selector tuning: radix select instead of 512 argmax passes

The inline long-context question was right. The truncating path was 512 masked argmax passes, each a full row scan plus two barriers; at 256k tokens that was 86 ms for a single row. I replaced it with an exact 8-bit radix select:

• map each finite score to a monotonic u32 key (-0.0 folded to +0.0, so both zeros tie);
• four 8-bit histogram passes (per-warp histograms, shared atomics) find the keep-th key;
• keep every key above it and fill the remaining slots from the equal keys in ascending block-index order (the canonical tie rule);
• the all-kept fast path stays: one pass, 1.4 us.

One row, budget 2048 (before = idle host p50; after = nsys kernel median at 256 threads):

case	before	after
all-kept 2k	2.5 us	1.4 us
trunc 2k+ (513 blocks)	0.89 ms	20.8 us
trunc 8k (2048 blocks)	1.37 ms	34.4 us
trunc 64k (16384 blocks)	18.1 ms	67.0 us
trunc 256k (65536 blocks)	85.9 ms	298 us
trunc 64k x8	20.4 ms	68.2 us

Thread sweep (nsys on a shared node): 128 threads is better at 2k/8k, 512 at 256k (187 vs 298 us); 256 is the balanced default.

Correctness: the parity gate passes 5/5 at this head; selection is 133 rows, 129 bitwise, 4 up to a tie, 128 all-kept, 5 truncating. compute-sanitizer memcheck (256k) reports 0 errors and racecheck (2k+) 0 hazards. A real bug surfaced here first: the equality pass marks a shared bitmap from several threads, and the non-atomic |= lost 9 of 34 marks; atomicOr fixes it, which the gate now covers through its zero-tie rows.

Benchmark: cargo run --release -p pegainfer-kernels --features qwen4 --example qwen4_qsa_select_bench (QSA_BENCH_CASE=256k filters a case). Commits: 3bae6b7c (kernel, example, doc), bf3bde73 (timing labels), 446af7a6 (clippy literals).

Environment caveat: the L20 image ships CUDA 12.9, below the repo's CUDA 13 floor, so this gate build used the FlashInfer tree's CCCL and a build-only guard for cudaErrorLossyQuery; the production CUDA 13.2 build was not re-run here. ncu is blocked by ERR_NVGPUCTRPERM on the shared pod, and the trap gate was not run (deliberate device faults). The PR stays draft; the description has the full table and boundaries.

[sysuTobo]

Parity + trap rerun at this head (L20)

Ran both gates on one L20 (sm_89, CUDA 12.9 plus the FlashInfer tree's CCCL and the build-only cudaErrorLossyQuery guard; the production CUDA 13.2 build was not re-run):

• parity: 5/5 passed; selection 133 rows, 129 bitwise, 4 up to a tie, 128 all-kept, 5 truncating.
• trap: 5/5 cases (score rows, score visible length, select past the mask, select negative row, select NaN), each trapped in its own child process; 2/2 tests passed in 8.6 s.
• post-trap parity rerun: 5/5, so this L20 stayed usable after the traps (unlike the earlier RTX 5060 / CUDA 12.9.1 observation).

The gates ran on GPU 5 of a shared node with other tenants active. The PR stays draft; the description's Verification and Boundaries now reflect the trap run.

[FeathBow]

All four points from last round are in: score_cols is explicit, the knob is gone from the kernel and the ABI, the trap parent checks for one passed test, and the radix select brings the 256k row down to 298 µs. The radix select reads right.

What is left is what the branch carries around the kernel, and we want that cleared out: comments and doc text that describe the argmax version or repeat each other, test scaffolding, a benchmark example, and a rebuild of the score stage's tests that this feature does not need. Details inline.

[FeathBow]

This describes the argmax version. The bitmap now only marks which threshold equal blocks get taken, so "exclusion state" and "picked" no longer fit. The same leftover sits in the shared memory ceiling comment ("exclusion bitmap") and the description.

[FeathBow]

The selection contract is spread over four docs (this header, the C entry, the FFI declaration and qsa_indexer_select_into), and this header and the wrapper doc each give the full version: grid contract, tie rule, the torch.topk note and the trap. The NaN rationale alone appears here, again in the kernel at the finite check, and again in the wrapper doc. Much of it argues why the design is right (why trap instead of skip, why the kernel repeats the host checks, why the host mirrors MAX_KEPT_BLOCKS), and the header also carries derived numbers (319392 blocks, 1.28M tokens). Could this header be the one place for the semantics, tie rule and trap, in a few lines, with the wrapper doc keeping only the Rust side (stride, spare capacity, empty launch) and the rationale moving to the description?

[FeathBow]

This says an empty launch is refused because asking for nothing is a scheduling bug, but the only caller, qsa_indexer_select_into, returns early on exactly that case and treats it as nothing to do. The check can stay; the justification contradicts the wrapper and can go.

[FeathBow]

select_plan is split out so its guards can be unit tested, and the five tests check one line ensure!s. select_plan_refuses_empty_extents checks a branch the public entry never reaches, since it returns early on an empty launch. The score stage above has no such split. Could select_plan fold back into qsa_indexer_select_into and these tests go?

[FeathBow]

This still says masked argmax, and the module doc above carries this fixture's measured counts (133 rows, 129 bitwise and so on), which go stale with any fixture change and belong in the description.

More broadly, the score half here is the merged score parity file moved into mod score and reindented, with its assertions unchanged, so 750 lines are deleted and readded for code this PR doesn't change. Could the selection test sit beside the existing score test in its file, without moving the score half?

[FeathBow]

This replaces the score stage's merged trap binaries with a harness that reruns itself per case, and adds three selection traps. The header still says the parent only checks the exit status, and it carries a device model and an incident report. The selection traps cover a range check and a finite check that are a few lines each, they have to fault the device, and nothing runs them outside a disposable card. Could the selection traps go and the score stage's trap files stay as they are on main?

[FeathBow]

This reproduces the description's timing table and checks only the mask's set bit count. The numbers already live in the description, so the example and its [[example]] entry in Cargo.toml can go, along with the bench command in bring-up.md.

[FeathBow]

This paragraph is a fifth copy of the kernel contract: mask layout, tie rule, the radix algorithm, the shared memory bound, the grid contract, the wrapper's refusals and the NaN rationale. The timing paragraph below repeats the description: the before and after table, the thread sweep, the build environment and the mutation table. The doc could keep what the code can't say: the radix select result at 256k, the token 0 grid limitation, and the note that an accuracy gate must not require mask equality at a tied cut.

[sysuTobo]

Review round applied

All eight inline points are in at e477e576; the branch now adds 7 files (+1044/−14) against main.

• Kernel header / comments — the header is now the single place for the contract (semantics, tie rule, trap, token-0 grid, radix path) and no longer carries derived totals. The bitmap helpers no longer describe the argmax version: block_picked/mark_picked/picked became block_taken/mark_taken/taken, and the shared-memory comment no longer says "exclusion".
• C entry — its doc no longer argues that an empty launch is a scheduling bug; it states that the entry refuses one while the Rust wrapper treats it as a no-op.
• SelectPlan — folded back into qsa_indexer_select_into; the five unit tests are gone. The device-free guard test for the C entry stays in the selection module, and the Rust-side guards are inline.
• Parity file — qwen4_qsa_indexer_parity.rs is deleted. qwen4_qsa_indexer_score_parity.rs is restored from main and the selection test is appended as a select module beside the score tests; the score half is untouched apart from one doc line. The module doc no longer says "masked argmax" or carries this fixture's measured counts (133/129/4/31–43); those live in the description.
• Trap files — qwen4_qsa_indexer_trap.rs is deleted and qwen4_qsa_indexer_rows_trap.rs, qwen4_qsa_indexer_visible_trap.rs and qwen4_qsa_trap/mod.rs are restored byte-for-byte from main. The three selection traps are gone; so is the harness and its incident report.
• Benchmark — examples/qwen4_qsa_select_bench.rs, its [[example]] entry and the bench command in bring-up.md are removed. The timing table, thread sweep, mutation table and build caveats stay in the description.
• bring-up.md — the selection section now keeps only what the code cannot say: the radix result at 256k, the token-0 grid limitation, and the tied-cut rule. Last touched is bumped to 2026-11.

Verification at this head (e477e576): cargo fmt --check, cargo check -p pegainfer-kernels --features qwen4 --tests and cargo clippy … -- -D warnings pass; the parity gate passed 5/5 (133 rows, 129 bitwise, 4 up to ties, 128 all-kept, 5 truncating). The build used CUDA 12.9 with the flashinfer tree's CCCL and the build-only cudaErrorLossyQuery guard; the production CUDA 13.2 build was not re-run. The before/after timings, memcheck and racecheck numbers in the description were measured at the previous head; this head only renames the bitmap helpers and trims comments, so the selector logic is unchanged. The score stage's trap binaries are unchanged from main and were not re-run.

[github-actions]

Commit attribution notice

This pull request contains commits attributed to an account other than the PR author @sysuTobo without a verified signature from that account:

• d2f075e3 attributed to @FeathBow

If you are mentioned above and did not take part in this pull request, please say so here. Attributing a commit to someone else normally requires that person to be involved; maintainers should treat unacknowledged attribution as a review blocker.

This check never fails the build.

[FeathBow]

We pushed one commit on top that trims what the last review pointed at. Selection behaviour is unchanged.

• The kernel's budget and score width traps are gone. Only the C entry launches the kernel, it already refuses both, and pos < tokens is checked first, so neither could fire. The comments arguing why the kernel repeats the host checks and why a NaN traps went with them, since the header states the trap as contract.
• qsa_indexer_select_into takes budget as usize like tokens and score_cols, which drops the conversion and its toolchain comment. The buffer checks use saturating_mul inside ensure! as the score stage does.
• The select module doc keeps what the gate compares, how it treats a tied cut and the ungated tie away from zero. Provenance, the tie rule and the token 0 grid note already live in the file header, the kernel header and the bring-up doc.
• The three summary assertions that held by construction are gone. The three coverage checks stay.
• The C entry guard test is gone. The wrapper refuses each of those inputs before it reaches the entry, and a CudaSlice never passes a null pointer. The last "exclusion bitmap" wording went with it.
• The bring-up doc drops the comparison with the argmax version, which never reached main. That comparison stays in the description.

Checked on sm_89 with the branch merged into current main: fmt, clippy with -D warnings over all qwen4 kernel targets, and the parity file with 4 passed (133 selection rows, 129 bitwise and 4 up to ties, score max abs 2.401e-3 as before). A synthetic check of the radix path also matched a host sort on every row, covering deep truncation, ties away from zero, signed zeros and keys that differ only in the low byte, up to 262147 tokens. The trap binaries were not run, since they fault the device and this commit does not touch them.

The description still mentions the device-free guard test.