Repository navigation
feat(qwen4): score QSA token-indexer blocks in native CUDA - #1160
Merged
FeathBow merged 12 commits intoOct 20, 2026
Merged
Conversation
QSA's token indexer is the highest-lead-time piece of this line's bring-up, and it gates on a golden that has been on `main` since pegainfer-project#1150 with nothing to replay it. This lands the middle of the three units pegainfer-project#1105 splits QSA into: the scoring half, in `pegainfer-kernels/csrc/qwen38fn/` 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 pools every complete micro-block of `indexer_compress_ratio` tokens into a single key — the block's raw keys averaged in float32, cast back to bf16 as the reference casts them, then `(1 + w)` RMSNorm and partial RoPE at the block's first token — and the other norms each of a query row's four heads at the row's own position and reduces `sum_h relu(q_h . k_b) / sqrt(indexer_head_dim)`. The pooled half is a pass of its own because it does not depend on the query row: the reference rebuilds it per row from the same keys through the same bf16 rounding points, so one pass over the keys is the same numbers for every row that can see them. It follows pegainfer-project#1105's D3: A/C native CUDA in a model-local directory rather than a vendored FlashInfer patch, and B the sm_89-first portable subset — no TMA, no cluster, no wgmma. Both kernels are SIMT shuffle reductions rather than `mma.sync` bodies, because one block's score is a 128-wide dot over 4 query heads: a shape m16n8k16 cannot tile without padding 4 rows to 16. The gate replays `test_data/qwen38fn-qsa-indexer-golden.safetensors` and holds its recorded `format_version` / `transformers_version` / `config_sha256`, plus the frozen config's own indexer geometry, against this build — a golden from another checkpoint fails loudly instead of skipping. It reads each row's explicit validity prefix rather than comparing through the reference's NaN padding, and proves the kernel writes exactly that prefix by checking the untouched tail still holds a sentinel. Six of seven cases land within 9.6e-7 of the reference; the 2055-token case, the only one with positions past 63, within 2.4e-3 of its 5e-3 bound. Selection (top-k over blocks plus the unconditional tail) is pegainfer-project#1105's next unit, which is why this kernel takes no budget: the reference scores every complete block and truncates only at selection. Refs pegainfer-project#1105 Signed-off-by: bohaotu <bohaotu@pony.ai>
The entry point of the QSA indexer's scoring kernel validates before it forms a launch, so its half of the refusal contract is pinnable with no device — and nothing pinned it. It is the only place the compiled row stride is checked, which is the reason `fused_width` is an argument at all: the body reads `fused_qk` with the width it was compiled for, so a caller whose rows are strided differently is read out of bounds rather than computed wrongly, the one mismatch the pointer arithmetic cannot survive. `rotary_dim` and `rope_theta` are compiled constants with no other check anywhere. The new test calls the entry point with the compiled geometry and a zero extent first, so the guards are the whole of the call and no launch follows: that positive control is what keeps every refusal below from being vacuous. Then each refusal: a row stride either side of 640, a rotary split of 32 and of 128, a rope_theta of 1e6 and one f32 ulp off 1e7, a negative and a NaN rms_eps, three negative extents, a score stride under the block grid, and one null per kind of argument. The pointers dangle and are never dereferenced, because a launch is the only thing that would touch them. Also prints the golden's recorded torch and platform in the provenance half. The parity bound's justification leans on the dumping host's reduction order, which was otherwise invisible to whoever reads this gate's output. Refs pegainfer-project#1105 Signed-off-by: bohaotu <bohaotu@pony.ai>
github-actions
Bot
added
qwen4
Contributor
Author
Code review覆盖 10 个文件（2 个 CUDA/FFI/ops 新文件 + 门禁 + build/Cargo 接线 + bring-up 文档），无遗漏；三条数据流的正向路径全部符合预期：
修改措施
|
sysuTobo
marked this pull request as ready for review
October 17, 2026 07:40
The scoring kernels move from csrc/qwen38fn/ to csrc/qwen4/ and take the line's name, like qwen35 and glm52: the directory, the .cu, both kernel entry points, the FFI symbol, the Rust modules and the parity test all drop the checkpoint qualifier. The golden fixture keeps its qwen38fn name — it is PR pegainfer-project#1150's committed data, recorded in its own provenance manifest, and renaming it would churn every branch that replays it. No semantic change: the same golden gate replays with the same numbers (7 cases, worst bound fraction 0.164 of the 5e-3 bound). Signed-off-by: bohaotu <bohaotu@pony.ai>
sysuTobo
force-pushed
the
feat/qwen4-qsa-indexer-score
branch
from
October 17, 2026 07:40
513d01b to
de9c111
Compare
validate()'s first ensure pins rotary_dim == 64 and compress_ratio == 4 bit-for-bit, so the range checks behind it (rotary in 1..=head_dim, compress_ratio > 0) can never fire — dead guards that read as extra safety but guard nothing. The parity test's mutations are still caught by the exact-equality pin, and it passes unchanged. The docs now also say what the untouched score tail implies for a caller: telling an unwritten entry from a real one means poisoning the buffer yourself, which is what the gate's sentinel does. Found in a data-flow audit of the PR; no behavior change. Signed-off-by: bohaotu <bohaotu@pony.ai>
Condense the kernel, wrapper and gate comment blocks to what a reader
needs — the frozen math, the grid contract, the rounding points and the
bound's justification — and drop the internal planning labels
("Stage 1, plan P2", "decision D3-B") from the code. No semantic
change; the parity gate replays with the same numbers.
Signed-off-by: bohaotu <bohaotu@pony.ai>
FeathBow
requested changes
Oct 17, 2026
FeathBow
left a comment
FeathBow
left a comment
Collaborator
There was a problem hiding this comment.
Thanks for this. On one sm_89 GPU the head commit builds with the qwen4 feature, clippy is clean and the gate prints exactly the output in the description. The fixture firmly pins the relu placement and the pooled key's RoPE anchor. Requesting changes because several statements in the description and the test comments do not hold when measured, and a reader would take them to mean the gate catches faults it lets through.
We changed one thing in the kernel at a time and ran the unchanged gate on the head commit.
| Kernel change | The description says | Worst score vs the bound | Gate |
|---|---|---|---|
| Relu after the head sum | "17x to 100x the bound" | 166x | fails |
| Pooled key RoPE at the block end | not listed | 536x | fails |
| Pool by sum instead of mean | "17x to 100x the bound" | 0.66x | passes |
rms_eps removed |
"pin every scalar constant" | 0.66x | passes |
rms_eps times 100 |
"pin every scalar constant" | 2.5x | fails narrowly |
| Mean kept in f32 | "17x to 100x the bound" | 3.1x | fails narrowly |
| cos and sin kept in f32 | "Every bf16 rounding point of the reference is reproduced" | 2.3x | fails narrowly |
The description also says the bound cannot be tightened because rsqrtf, cosf and sinf vary across architectures. Your sm_120 run and our sm_89 run printed identical output, and swapping the device trig for a host table built like pegainfer_core::rope leaves all 3,674 scores unchanged. The six short cases sit at 9.5e-7 while every passing or narrow row above moves them by 2.1e-3 or more, so holding cases 0 to 5 at that level and keeping 5e-3 only for the long case would turn all five into clear failures.
Before merge we would ask for four things.
- Correct the coverage and tolerance statements in the description, the test comment and the bring-up note to what was measured (inline).
- Tighten the short cases as above. The long case can keep its bound until its residual is traced.
- Make out of range
rowsorvisible_lenfail observably, and keep a NaN visible instead of scoring it zero (inline, both reproduced). - Print the worst element's own coordinates next to
max_abs(inline).
Smaller points.
- The geometry plumbing and the budget guard callers that do not exist yet (inline).
- The
static_assertonkDimPerLanerepeats its initializer, and therow_blocks > blocksclamp is implied by thevisible <= tokensclamp above it. - The
full_layers == 12assertion is outside this kernel's contract, the finalworst_scaled <= 1.0assertion cannot fail becausecheck_casepanics first, and nothing reads the manifestassert_provenancereturns. - The description still names the old
qwen38fnpaths and test target, so its commands do not run on this head. - Optionally, the kernel could read the cos and sin table that
pegainfer_core::ropebuilds, since the reference indexer shares the position embeddings of attention. That dropskRopeTheta,rotary_dimand their checks, and on this fixture it left every score identical.
Comment on lines
+44
to
+63
| /// Score parity bound, `abs(got - want) <= ATOL + RTOL * abs(want)`. | ||
| /// | ||
| /// Six cases land within 9.6e-7 — every bf16 rounding decision the reference | ||
| /// makes is reproduced bit-for-bit. The residual sits in the seventh (the only | ||
| /// case with positions past 63): 152 of 3076 scores above 1e-4, worst 2.4e-3, | ||
| /// sign-balanced with a -5.2e-7 bias — the shape of a flipped bf16 rounding | ||
| /// decision, not a semantic difference. The reference rounds its float32 | ||
| /// reductions in a torch-defined order; where a value sits within one f32 ulp | ||
| /// of a bf16 midpoint the kernel's butterfly tree rounds apart from it, and | ||
| /// one flipped element moves every score it touches. The reference does the | ||
| /// same to itself: re-deriving case 2 on x86 (torch 2.14.1 vs the golden's | ||
| /// 2.14.0/arm64) moves one of its 496 scores by 2.2e-3, three orders more than | ||
| /// this kernel's 9.5e-7 there. | ||
| /// | ||
| /// So the bound holds ~6 flips at a typical score of ~2, against a measured | ||
| /// worst of 0.16 of it, while any real fault — a block off by one, a missing | ||
| /// relu, an unrotated half — moves a score by ~1, two orders out. Selection | ||
| /// must treat a block tied at the cut as an equally valid pick, for the same | ||
| /// reason. | ||
| const ATOL: f32 = 5.0e-3; |
Collaborator
There was a problem hiding this comment.
This comment and the matching note in the bring-up doc claim more than the gate shows. Pooling by sum and removing rms_eps both pass, dropping the mean's bf16 round trip fails at only 3.1x, and the reduction order account has no intermediate comparison behind it, so please replace the block shift estimate and that account with what was measured.
Comment on lines
+334
to
+337
| /// One case's parity: how much was compared, the worst element and where, and | ||
| /// the shape of the residual — the sign split and the bias catch a | ||
| /// *systematic* error (a lean one way, deviations piling up in one row) | ||
| /// that a max alone cannot. |
Collaborator
There was a problem hiding this comment.
Nothing asserts on the sign split, the bias or the widest row, so these fields cannot catch a systematic error as this comment says. Every score is nonnegative and RTOL is 5e-3, so a uniform scale error up to 0.5% passes every comparison with no sign to flip, and the sum pooling mutation shows the same balanced signs as the baseline. Please assert something these fields would catch, or drop them with the comment.
Comment on lines
+441
to
+450
| max_abs = max_abs.max(abs); | ||
| signed_sum += f64::from(delta); | ||
| over_1e4 += usize::from(abs > 1.0e-4); | ||
| over_1e5 += usize::from(abs > 1.0e-5); | ||
| pos_1e4 += usize::from(abs > 1.0e-4 && delta > 0.0); | ||
| row_over_1e4 += usize::from(abs > 1.0e-4); | ||
| if scaled > max_scaled { | ||
| max_scaled = scaled; | ||
| case_worst = Some((block, *x, *want, abs)); | ||
| worst = (*row, block, *x, *want); |
Collaborator
There was a problem hiding this comment.
max_abs takes the largest absolute error while worst follows the largest scaled error, and the print presents them as one element. In the long case it prints max abs 2.401e-3 beside an element whose own difference is 2.255e-3. Please keep each metric's coordinates, or print only the metric they belong to.
Comment on lines
+169
to
+173
| // `rows[row]` is the row's own position and the fused_qk row its query | ||
| // heads are read from; clamped, like `visible_len` below, so a broken | ||
| // caller gets a wrong position rather than an out-of-bounds read. | ||
| const int raw_pos = rows[row]; | ||
| const int pos = raw_pos < 0 ? 0 : (raw_pos >= tokens ? tokens - 1 : raw_pos); |
Collaborator
There was a problem hiding this comment.
On our run rows = [8] and rows = [1000] against an 8 token buffer both return Ok with the scores of row 7, and rows = [-3] with visible_len = [-5] returns Ok and writes nothing. The repository already traps out of range device indices without a host copy (csr_page_row_checked and resident_row_checked in common.cuh, exercised by hd512_qk_rope_trap_page.rs), so the D2H cost the description cites does not apply. Please trap here as well and keep the valid empty prefix case successful.
Comment on lines
+251
to
+254
| // relu per head, then the sum over heads, then the 1/sqrt(head_dim) | ||
| // scale: all three are the reference's own order. | ||
| const float relu_sum = fmaxf(acc[0], 0.0f) + fmaxf(acc[1], 0.0f) + | ||
| fmaxf(acc[2], 0.0f) + fmaxf(acc[3], 0.0f); |
Collaborator
There was a problem hiding this comment.
fmaxf returns the other operand when one is NaN, so a poisoned block scores exactly zero. With one raw key element of case 2 set to NaN, all 57 rows that see that block scored 0.0 where the clean reference has 2.62, 0.68 and 1.61, and the launch reported success. torch.relu keeps the NaN, so please let it through or fail instead of handing selection a valid looking lowest score.
Comment on lines
+76
to
+80
| /// Refuse a geometry the kernel was not compiled for. Called by | ||
| /// [`qsa_indexer_block_scores_into`] before it forms a pointer, and public so | ||
| /// a caller deriving geometry from a checkpoint's config can check it | ||
| /// without a launch. | ||
| pub fn validate(&self) -> Result<()> { |
Collaborator
There was a problem hiding this comment.
Only the literal from qwen38_flash_next() passes validate(), and the wrapper then hands the kernel the fused_width, rotary_dim and rope_theta it just pinned. The matching entry point checks can therefore fail only through the raw FFI, which is the path the refusal test exercises, while only_the_compiled_geometry_passes asserts the literals. Passing only rms_eps at run time, and checking the checkpoint's indexer fields where the model crate reads its config once that caller lands, removes the struct, both refusal layers and both tests.
Comment on lines
+42
to
+45
| /// `indexer_budget`. Not read by this stage: the reference scores every | ||
| /// complete block and truncates at selection, so the budget only bounds | ||
| /// what comes next. Carried here so the caller passes one geometry around. | ||
| pub budget: usize, |
Collaborator
There was a problem hiding this comment.
The budget is unused by this stage and its divisibility check serves selection, so both can land with selection. The same holds for the page size reasoning in the max_blocks doc and for the gate's truncating and boundary assertions, which depend only on the fixture.
…rows A mutation audit found the gate letting real semantic faults through: with one 5e-3 bound everywhere, pooling by sum and dropping rms_eps both passed. The pooled key's RMSNorm is scale-invariant, so those faults only move the result through eps at the pooled key's scale — 3.3e-3 on this fixture — and the long case's residual forced one bound on every case. Cases 0-5 sit at <=9.6e-7 and now carry a 2e-6 bound; the long case keeps 5e-3 until its residual is traced. The tight bound's floor is the measured mutation minimum (3.3e-3 for the eps/pooling faults, 1.15e-2 and 1.55e-2 for the missing bf16 round trips, 0.83 for a post-sum relu, 2.7 for a shifted RoPE anchor), not architecture variance: sm_89 and sm_120 print identical numbers and a host trig table leaves all 3,674 scores unchanged. Correctness: - Out-of-range `rows`/`visible_len` trap instead of clamping, following common.cuh's checked-row precedent; the valid empty prefix still succeeds. Two trap binaries, one per case, since __trap() is sticky. - relu is a comparison rather than fmaxf, so a poisoned key stays NaN like torch.relu instead of scoring a valid-looking zero. New gate: one poisoned key makes exactly the rows that see its block score NaN, every other block finite, the untouched tail sentinel. - max-abs and max-scaled each keep their own element coordinates; they were printed as one element. Scope, from the same review: - Geometry drops the unused budget and its divisibility guard (selection's business); the gate drops the fixture-only truncating/boundary asserts and the unreachable final bound assert. - Removes the tautological static_assert, the dead row_blocks clamp, the out-of-contract full_layers assert, and assert_provenance's unused return. - The bring-up note and gate comments state the measured numbers. Signed-off-by: bohaotu <bohaotu@pony.ai>
The indexer shares the attention line's position embeddings, so the two kernels now read the cos/sin table `pegainfer_core::rope` builds instead of computing trig on the device. The entry point takes the two tables and checks `rotary_dim` as their compiled row stride; `rope_theta` disappears from the kernel, the FFI and the geometry, and the wrapper checks the tables cover every position a call can read. The gate builds the same table from the frozen config's own `rope_theta`. Measured: every score is bit-identical to the device-trig version on all seven cases (same max abs, coordinates and bias), so the parity numbers are unchanged. 4 parity + 2 trap tests and 20 lib tests pass; fmt and clippy are clean. Signed-off-by: bohaotu <bohaotu@pony.ai>
…dant Only the literal from `qwen38_flash_next()` could pass `validate()`, and the wrapper then handed the kernel the same values, so the Rust validation and the entry-point geometry checks could only fail through the raw FFI — the path the device-free refusal test exercised. Following the review, the struct, both layers and both tests are gone: the wrapper takes only `rms_eps` at run time, sizes its buffer checks from the compiled constants, and the checkpoint's indexer fields are checked where the model crate reads its config once that caller lands. The gate derives its geometry from the frozen config and cross-checks it against the fixture's own tensor shapes instead of comparing it to a Rust copy of the constants, so a config/geometry drift reaches the kernel as a score mismatch rather than as a duplicated contract. The trap tests size their buffers from local compiled-geometry constants. No semantic change: parity 3/3 (provenance, comparison, NaN), traps 2/2, lib 19, fmt and clippy clean. Signed-off-by: bohaotu <bohaotu@pony.ai>
sysuTobo
marked this pull request as draft
October 17, 2026 17:19
Contributor
Author
|
Thanks @FeathBow — the mutation audit found real gaps, and the measured table is what changed this gate. All three rounds are on the head:
The four asks1. Statements corrected. The description is rewritten (including the stale 2. Short cases tightened. Cases 0–5 now carry 3. Out-of-range metadata traps. 4. Read-out coordinates. Smaller points
Optional: the shared RoPE table (done)The kernels take the two Verification (head
|
sysuTobo
marked this pull request as ready for review
October 17, 2026 17:24
The trap tests only asserted when the call returned Ok, so any wrapper-side
Err — a drifted sizing constant, a new capacity check — would leave the gate
green without ever exercising the trap. The call's Ok is now the expected
path, and an Err is accepted only when it is the launch's own ("launch
failed"); a wrapper rejection fails the test. Both binaries still pass, so
the trap surfaces with no launch-side error on this box.
Signed-off-by: bohaotu <bohaotu@pony.ai>
Collaborator
|
@codex review please |
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: aadacfd44d
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| /// `[max_pos, ROTARY_DIM]`, position-major): the indexer rotates with the same | ||
| /// position embeddings the attention uses, so this stage computes no trig. The | ||
| /// tables must cover every position this call can read, `0..tokens`. | ||
| pub fn qsa_indexer_block_scores_into( |
There was a problem hiding this comment.
In the Qwen4 path I checked, pegainfer-qwen4 has no dependency on this kernel crate and pegainfer-qwen4/src/model_line.rs:29-30 still makes every launch call bail, so no production request can invoke this new operator. The golden and trap tests therefore validate only the isolated buffers and cannot exercise the projection/layout/stream integration required by the repository's production-E2E rule for features; either land this with the text-graph call site and a production E2E or defer the unreachable feature until that path exists.
AGENTS.md reference: AGENTS.md:L124-L126
Useful? React with 👍 / 👎.
Contributor
Author
There was a problem hiding this comment.
Thanks — the facts check out, and this is the one place this PR deliberately stops short.
pegainfer-qwen4/src/model_line.rs:28-30does bail onlaunch; that is the line's shape since feat(qwen4): enforce the tensor contract before any device work #1149 — detection, config validation and the tensor contract, with no text graph yet.- Nothing production can reach this op today: it is behind
pegainfer-kernels's additiveqwen4feature, which no production build enables (the server'sqwen4feature only turns ondep:pegainfer-qwen4, a CPU-only crate —pegainfer-server/Cargo.toml:48), and its only callers are the three test files. - The production call site needs the slice-B text graph: the
index_qk_projprojection (2560 → 640), the stream and layout it feeds, and the scheduler step that consumes the scores.
So P2's deliverable is deliberately the component-level gate — golden parity, the trap contract and NaN propagation — and that boundary is written into the PR's "what is NOT verified" section rather than hidden. It is the same staging shape as this line's earlier component PRs (#1149, #1150, #1153), none of which has a production request path yet.
We will wire the op and carry a production E2E with the text-graph slice. If maintainers would rather not merge a component ahead of its caller, the alternative is to keep this PR draft and fold it into that slice — flagging so the merge-order call is explicit rather than accidental.
|
This pull request has merge conflicts with main. Please rebase it so it can be reviewed and merged. |
…ct the gate notes Signed-off-by: Feathbow <feathbow@gmail.com>
Signed-off-by: Feathbow <feathbow@gmail.com>
Commit attribution noticeThis pull request contains commits attributed to an account other than the PR author @sysuTobo without a verified signature from that account: 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. |
… stale notes Signed-off-by: Feathbow <feathbow@gmail.com>
Collaborator
|
We pushed three commits to this branch on top of your four, two fixes and a merge of main. Your four commits close the four points of the earlier review. The traps, the NaN path, the tight bound and the separate read-out coordinates all hold up. We checked each one by breaking the kernel on one sm_89 GPU with the gate unchanged. The numbers are how far the first failing score is from the reference.
What our commits change.
What our trap changes catch, checked the same way.
The description still refers to code this head no longer has. The gate has one device free test, Verification at the new head on one sm_89 GPU, with main's FlashInfer pin.
|
FeathBow
approved these changes
Oct 20, 2026
sysuTobo
added a commit
to sysuTobo/pegainfer
that referenced
this pull request
Oct 21, 2026
The Qwen3.8-Flash-Next MoE golden had a device-free consumer only, so it pinned a contract without executing any production code. Add its device half inside pegainfer-kernels, behind the line's additive `qwen4` feature, appending to the module pegainfer-project#1160 opened, so the model crate stays CUDA-free and the feature keeps the new translation unit out of default builds: csrc/qwen4/qwen4_moe_route.cu routing + shared-expert combine src/ffi/qwen4.rs the two C entry points, appended src/ops/qwen4.rs geometry/semantics guards, wrappers tests/qwen4_moe_block_parity.rs the gate, through the ops layer Routing is fp32 softmax over 512 -> top-10 in the canonical order (descending probability, ascending index inside a tie) -> optional `norm_topk_prob` renorm summing in that same order -> bf16 weights and i32 ids, one block per token over fixed shuffle trees so the picks are deterministic run to run; a non-finite logit fails its row closed by a block-uniform vote before the softmax, so no NaN reaches the arithmetic. The combine folds the learned per-token sigmoid gate into the block sum, rounding where the reference rounds. Everything else is composition: the router GEMM, the shared expert's three projections and its SwiGLU go through `gemm_into_checked` / `silu_mul_batch_into`, so the gate grades the production path. Nothing in-tree carries this routing contract, which is why the kernel is new: gemma4's router is 128 experts with `__expf`, an unconditional renorm and f32 weights; glm52 (256/top-8) and kimi (384/top-8) score with sigmoid + bias; FlashInfer's FilteredTopK does no softmax at all. The combine is the one near-duplicate -- `add_scaled_bf16_kernel`'s arithmetic, which takes its scale as one host f32 and so cannot carry a per-token device gate -- and the doc names it as such. Measured on an RTX 5060 (sm_120), CUDA 12.9.1, every bound derived before it was checked: routing bitwise 210/210 on the stored logits, the chain 0.0397 relative against the derived 0.1332, the combine bitwise 53760/53760 on stored operands and <=0.371 of its bound on the device branch, and three controls that are not vacuous thresholds (no-renorm 210/210 at 146x, bf16 softmax 56/210 with no re-ordered picks, no shared gate 76% of outputs at 125x). The device-free replay is untouched and still green. Signed-off-by: bohaotu <bohaotu@pony.ai>
This was referenced Oct 21, 2026
sysuTobo
added a commit
to sysuTobo/pegainfer
that referenced
this pull request
Oct 23, 2026
One indexer, one golden gate: scoring and selection now live in `qwen4_qsa_indexer_parity`, sharing the fixture path, provenance check and tensor readers instead of duplicating them across two files. All five tests and every assertion are preserved; the select half still replays the fixture's recorded scores. The five trap binaries keep one trap per process, because `__trap()` leaves the context poisoned. Their bodies move into the shared `qwen4_qsa_trap` launcher, leaving each binary about eighteen lines like pegainfer-project#1160's, and the device-reset warning is recorded once in the launcher. Signed-off-by: bohaotu <bohaotu@pony.ai>
sysuTobo
added a commit
to sysuTobo/pegainfer
that referenced
this pull request
Oct 23, 2026
One indexer, one golden gate: scoring and selection now live in `qwen4_qsa_indexer_parity`, sharing the fixture path, provenance check and tensor readers instead of duplicating them across two files. All five tests and every assertion are preserved; the select half still replays the fixture's recorded scores. The five trap binaries keep one trap per process, because `__trap()` leaves the context poisoned. Their bodies move into the shared `qwen4_qsa_trap` launcher, leaving each binary about eighteen lines like pegainfer-project#1160's, and the device-reset warning is recorded once in the launcher. Signed-off-by: bohaotu <bohaotu@pony.ai>
sysuTobo
added a commit
to sysuTobo/pegainfer
that referenced
this pull request
Oct 24, 2026
The kernel, its C ABI, the ops surface and six exported constants carried a gate-only mutation parameter. Measure the five controls with temporary kernel patches instead, as the pegainfer-project#1160 gate did, and keep the production launch free of the knob. The golden, tie and structural assertions stay; the fixture still covers both selection paths. Signed-off-by: bohaotu <bohaotu@pony.ai>
FeathBow
added a commit
that referenced
this pull request
Nov 1, 2026
…1178) * feat(qwen4): select QSA indexer blocks and their tail in native CUDA Signed-off-by: bohaotu <bohaotu@pony.ai> * refactor(qwen4): keep the selector's gate knob out of the production entry The only public entry took the gate's mutation variant as a required argument, so a production caller had to spell a value that must always be 0 and could silently select a negative-control path. Split it the way the PLE block does: `qsa_indexer_select_into` fixes the production variant, and `qsa_indexer_select_variant_into` carries the knob for the parity gate. The kernel is unchanged. Also record the trap binaries' device side effect: on an RTX 5060 under CUDA 12.9.1 a deliberate `__trap()` left the whole GPU uninitializable for later processes, so they must run last or on a disposable device. Signed-off-by: bohaotu <bohaotu@pony.ai> * fix(qwen4): tighten the QSA selector's guards and docs - Treat an empty launch as a no-op in the selector, matching the score stage and the other kernels' convention; the C entry still refuses one reached directly. - Count the selector's static shared state against the 48 KiB per-block ceiling so an exact-boundary launch cannot be refused by the driver. - Drop the dead int64 overflow guard: an int x int product always fits. - Correct the budget error text and the module/feature docs, which still described selection as a later stage. - Trap a negative `rows` entry in its own binary; the negative half of the row-position guard was untested. - Record in the bring-up doc that ties at the cut are not an edge case at long context, so an end-to-end accuracy gate must not require mask equality with the reference there. Signed-off-by: bohaotu <bohaotu@pony.ai> * test(qwen4): merge the indexer gates and share the trap launcher One indexer, one golden gate: scoring and selection now live in `qwen4_qsa_indexer_parity`, sharing the fixture path, provenance check and tensor readers instead of duplicating them across two files. All five tests and every assertion are preserved; the select half still replays the fixture's recorded scores. The five trap binaries keep one trap per process, because `__trap()` leaves the context poisoned. Their bodies move into the shared `qwen4_qsa_trap` launcher, leaving each binary about eighteen lines like #1160's, and the device-reset warning is recorded once in the launcher. Signed-off-by: bohaotu <bohaotu@pony.ai> * test(qwen4): run every indexer trap from one child-process gate The five trap binaries each carried a case, but the cases are one line of setup. Keep one trap per process — `__trap()` poisons the context — and let one gate re-execute itself with `QSA_TRAP_CASE`, so the indexer has two test files: the golden parity gate and this trap gate. The child sets `PEGAINFER_REQUIRE_GPU=1` so a missing device fails instead of reporting a trap it never reached. Signed-off-by: bohaotu <bohaotu@pony.ai> * style(qwen4): keep the qwen4 re-exports rustfmt-ordered Signed-off-by: bohaotu <bohaotu@pony.ai> * fix(qwen4): take the selector's score stride explicitly The selector inferred `score_cols` from `scores.len() / rows.len()`. The score stage allows a buffer with spare capacity, so handing the selector that buffer produced a larger stride, passed both host checks and read every row from the wrong offset. Take `score_cols` like the score stage does, check `scores.len() >= rows * score_cols`, and cover the spare- capacity case in the launcher tests. Signed-off-by: bohaotu <bohaotu@pony.ai> * refactor(qwen4): drop the selector's gate-only variant knob The kernel, its C ABI, the ops surface and six exported constants carried a gate-only mutation parameter. Measure the five controls with temporary kernel patches instead, as the #1160 gate did, and keep the production launch free of the knob. The golden, tie and structural assertions stay; the fixture still covers both selection paths. Signed-off-by: bohaotu <bohaotu@pony.ai> * perf(qwen4): skip the selection search when every block is kept Rows under the budget keep every complete block, so the mask is the whole visible prefix. Skip the masked argmax loop and its barriers for them; the parity gate already counts those rows and asserts the fixture covers both paths. Signed-off-by: bohaotu <bohaotu@pony.ai> * fix(qwen4): fail the trap gate when its child runs no test A filter that matches nothing (`--exact` after a rename) exits 0 with `0 passed`, so checking only the exit status let every case pass without reaching a trap. Require the child to report one passed test. Signed-off-by: bohaotu <bohaotu@pony.ai> * perf(qwen4): radix-select the indexer's top blocks The masked-argmax selection cost one full row scan and two barriers per kept block: 86 ms for a single 256k-token row, 18 ms at 64k and 0.9 ms at the first truncating row. Replace it with an exact 8-bit radix select: four histogram passes find the keep-th sortable key, then the blocks equal to the threshold are filled in ascending block-index order (the canonical tie rule). The all-kept fast path stays. Measured on one L20, one row, budget 2048, nsys kernel medians at 256 threads: 20.8 us at 2052 tokens, 34.4 us at 8k, 67.0 us at 64k, 298 us at 256k, and 1.4 us for the all-kept fast path. The parity gate passes 5/5 and compute-sanitizer memcheck/racecheck are clean. Add the manual benchmark example used for those numbers. Signed-off-by: bohaotu <bohaotu@pony.ai> * docs(qwen4): label the selector thread-sweep timing protocols Signed-off-by: bohaotu <bohaotu@pony.ai> * fix(qwen4): keep the benchmark's long literals readable Clippy's unreadable_literal gate rejects 65535 and 262143 without separators; use separators so the manual benchmark stays clean under -D warnings. Signed-off-by: bohaotu <bohaotu@pony.ai> * refactor(qwen4): keep the branch to the selector and its parity gate The score stage's test files, parity gate and trap binaries are restored byte-for-byte from main; selection tests sit beside the score test in `qwen4_qsa_indexer_score_parity.rs` without moving the score half. The select kernel header is the one place for the contract; the C entry and the Rust wrapper keep only what each side owns. The `SelectPlan` split and its five device-free tests fold back into the wrapper, the benchmark example and its Cargo entry go, and the bring-up doc keeps only what the code cannot say: the radix select's 256k result, the token-0 grid, and the tied-cut rule. Signed-off-by: bohaotu <bohaotu@pony.ai> * docs(qwen4): bump the bring-up freshness stamp Signed-off-by: bohaotu <bohaotu@pony.ai> * refactor(qwen4): trim the QSA selector's unreachable checks and repeated docs Signed-off-by: Feathbow <feathbow@gmail.com> --------- Signed-off-by: bohaotu <bohaotu@pony.ai> Signed-off-by: Feathbow <feathbow@gmail.com> Co-authored-by: Feathbow <feathbow@gmail.com>
sysuTobo
added a commit
to sysuTobo/pegainfer
that referenced
this pull request
Nov 1, 2026
The Qwen3.8-Flash-Next MoE golden had a device-free consumer only, so it pinned a contract without executing any production code. Add its device half inside pegainfer-kernels, behind the line's additive `qwen4` feature, appending to the module pegainfer-project#1160 opened, so the model crate stays CUDA-free and the feature keeps the new translation unit out of default builds: csrc/qwen4/qwen4_moe_route.cu routing + shared-expert combine src/ffi/qwen4.rs the two C entry points, appended src/ops/qwen4.rs geometry/semantics guards, wrappers tests/qwen4_moe_block_parity.rs the gate, through the ops layer Routing is fp32 softmax over 512 -> top-10 in the canonical order (descending probability, ascending index inside a tie) -> optional `norm_topk_prob` renorm summing in that same order -> bf16 weights and i32 ids, one block per token over fixed shuffle trees so the picks are deterministic run to run; a non-finite logit fails its row closed by a block-uniform vote before the softmax, so no NaN reaches the arithmetic. The combine folds the learned per-token sigmoid gate into the block sum, rounding where the reference rounds. Everything else is composition: the router GEMM, the shared expert's three projections and its SwiGLU go through `gemm_into_checked` / `silu_mul_batch_into`, so the gate grades the production path. Nothing in-tree carries this routing contract, which is why the kernel is new: gemma4's router is 128 experts with `__expf`, an unconditional renorm and f32 weights; glm52 (256/top-8) and kimi (384/top-8) score with sigmoid + bias; FlashInfer's FilteredTopK does no softmax at all. The combine is the one near-duplicate -- `add_scaled_bf16_kernel`'s arithmetic, which takes its scale as one host f32 and so cannot carry a per-token device gate -- and the doc names it as such. Measured on an RTX 5060 (sm_120), CUDA 12.9.1, every bound derived before it was checked: routing bitwise 210/210 on the stored logits, the chain 0.0397 relative against the derived 0.1332, the combine bitwise 53760/53760 on stored operands and <=0.371 of its bound on the device branch, and three controls that are not vacuous thresholds (no-renorm 210/210 at 146x, bf16 softmax 56/210 with no re-ordered picks, no shared gate 76% of outputs at 125x). The device-free replay is untouched and still green. Signed-off-by: bohaotu <bohaotu@pony.ai>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This lands Stage 1 / P2 of the QSA (sparse attention) bring-up for the Qwen3.8-Flash-Next line: the token-indexer block-scoring kernel in native CUDA, plus the parity gate that replays the committed reference golden. Selection is P3 and out of scope here.
Summary
pegainfer-kernels/csrc/qwen4/qwen4_qsa_indexer_score.cu— two kernels behind an additiveqwen4feature onpegainfer-kernels, so the model cratepegainfer-qwen4keeps having no CUDA dependency at all (that is the point of its shape, and #1149's CPU-only test set is untouched). No PyTorch, no new build-time dependency, BF16 only.The gate is
pegainfer-kernels/tests/qwen4_qsa_indexer_score_parity.rs: it replaystest_data/qwen38fn-qsa-indexer-golden.safetensors(merged in #1150) through the kernel and compares block scores against the reference. Two trap binaries pin the kernel's out-of-range behavior (qwen4_qsa_indexer_{rows,visible}_trap.rs).pegainfer-kernels/csrc/qwen4/qwen4_qsa_indexer_score.cuextern "C"entry pointpegainfer-kernels/src/ffi/qwen4.rspegainfer-kernels/src/ops/qwen4.rsqsa_indexer_block_scores_intoand the compiled-geometry constants its buffer checks usepegainfer-kernels/tests/qwen4_qsa_indexer_score_parity.rspegainfer-kernels/tests/qwen4_qsa_indexer_rows_trap.rsrowstraps (one trap per binary)pegainfer-kernels/tests/qwen4_qsa_indexer_visible_trap.rsvisible_lentrapspegainfer-kernels/{Cargo.toml,build.rs}+src/{ffi,ops}.rsqwen4feature and its directory gate (4 small edits, no existing row touched)docs/models/qwen4/bring-up.mdKernel design, and the decisions it follows
Frozen semantics are transformers 5.16.1
qwen4_exp,Qwen4ExpTextQSAIndexer.forward, per visible micro-blockbof a query row at positionp:indexer_kv_heads = 1), and the relu is per head before the sum over heads — both visible in the structure, not just the arithmetic.visible_lenseparately fromrows. Nothing in the kernel can detect a visible list anchored elsewhere; that caller must hand it a shifted key array.page_size = 16 / compress_ratio = 4divisibility is taken as given and not used by this stage, which reads the raw key array; what this stage does depend on from D2 is the grid anchor (token 0) above.csrc/qwen4/, named after the model line likeqwen35andglm52), no vendored FlashInfer change — the FlashInfer submodule is untouched, so no new submodule patch to maintain.mma.syncbodies, because one block's score is a 128-wide dot over 4 query heads, a shapem16n8k16cannot tile without padding 4 rows to 16. Nomma/cp.asyncis used at all; the exclusion list is what this variant obeys.The kernel is compiled for one geometry; the wrapper carries only
rms_epsat run time. The fused row is 640 wide, the shared tables stride 64, and four tokens make a micro-block; the wrapper sizes its buffer checks from those numbers. The checkpoint's own indexer fields are checked where the model crate reads its config, once that caller lands — this stage does not carry a second geometry object to pin, and the score comparison is what holds the compiled kernel to the config.rowsandvisible_lenare caller-provided device values: an out-of-range entry traps on the device (the checked-row pattern ofcommon.cuh, no D2H copy), while the valid empty prefix —visible_len = 0— still succeeds and writes nothing. The relu is written as a comparison rather thanfmaxf, so a poisoned key keeps its NaN the waytorch.reludoes, instead of turning into a valid-looking lowest score.Evidence
Exact commands, from this worktree, in the dev container (
./docker/dev.sh run ..., imagepegainfer-dev:cu132, one sm_120 GPU):What the gate asserts, beyond the score comparison: the fixture's
format_version/transformers_version/config/config_sha256against this build's frozen config, and that config's indexer geometry (indexer_n_heads4,indexer_kv_heads1,indexer_head_dim128,indexer_compress_ratio4,rope_theta1e7,rms_norm_eps1e-6,rotary_dim=head_dim×partial_rotary_factor= 64) againstQsaIndexerGeometry::qwen38_flash_next(); per row it compares only the fixture's explicit validity prefix(row + 1) / compress_ratio(the reference NaN-pads the rest) and asserts the NaN padding is really there; and it asserts the kernel left every entry past that prefix holding the sentinel it was pre-filled with, so a wrong block count cannot pass as a shorter comparison. One fixture property is still asserted because this stage depends on it — at least one case must carry an incomplete tail; the budget and block-count bookkeeping is selection's business and is printed, not asserted.The device-free half
fixture_records_the_frozen_referenceneeds no device and is not behinddevice_or_skip, so it cannot be skipped: it holds the fixture's recordedformat_version,transformers_versionandconfig_sha256against the frozen config. The gate then derives its geometry from that config and cross-checks the numbers against the fixture's own tensor shapes, so a checkpoint or config drift reaches the kernel as a score mismatch rather than as a Rust-side duplicate of the constants. (An earlier revision also carried a device-free refusal test for aQsaIndexerGeometrystruct; review pointed out that it only compared the literal to itself, and both are gone.)Tolerance: 2e-6 on cases 0–5, 5e-3 on the one case whose residual is untraced
abs(got - want) <= ATOL + RTOL * abs(want), with2e-6 / 2e-6on cases 0–5 and5e-3 / 5e-3on the 2055-token case.rms_eps×100 2.5x — and, under the old single 5e-3 bound, pooling by sum and droppingrms_epsboth passed at 0.66x. The last two are the ones the tight bound exists for: they move scores by 3.3e-3 (the pooled key's RMSNorm is scale-invariant, so pooling by sum differs from the mean only throughepsat the pooled key's scale), which is ~1600x the new 2e-6 bound. With the split bound every measured fault now fails clearly.torch 2.14.1+cpu, against a golden dumped undertorch 2.14.0on arm64 — the gate prints both), moves one of its 496 scores by 2.2e-3, the same single-flip size, while on that case this kernel is three orders closer (9.5e-7). It keeps 5e-3 until that residual is traced.pegainfer_core::rope, built once and passed in) instead of computing trig, so there is no device-specific surface left to vary — the table-driven run reproduces the previous scores bit for bit. The short cases' 2e-6 is calibrated to the measured mutation floor, not to a cross-host drift.Why this gate does fail if the kernel is wrong. The fixture measures the distance itself: a one-block grid shift replaces each score with its neighbour's, and over the fixture's own reference scores that is a median
|delta|of 1.49 (p10 0.26, p90 3.78) at the long case — 166x the bound at the typical score for the measured relu fault, 536x for the RoPE anchor, on nearly every element. The same order covers a missingsqrt(128), summing the heads before the relu, rotating 128 dims instead of 64, and reading the key head at the wrong fused column; the rounding faults measured in 2 are now two to three orders past the tight bound.Diagnostics, not assertions. The sign split, the bias and the widest row are printed as evidence about the residual's shape, not asserted: independent flips cancel in a mean, and a scale error cancels too (it flips the residual's sign with the score's), so such an assertion would add a brittle threshold with little power. The assertion is the bound, and each of
max absandmax bound fractionnow prints its own element's coordinates.Kernel limitations / what is NOT verified
index_qk_proj(2560 → 640) is a GEMM, not part of this kernel, and the golden starts after it (its docstring says so): this stage gates the semantics from the projection's output onward.rows/visible_lentrap on the device rather than being clamped: an in-range contract is asserted, the valid empty prefix still succeeds, and both trap paths have their own binaries. A poisoned key keeps its NaN score (torch.relusemantics), gated bya_poisoned_key_keeps_its_nan_score.tokens < 4(no complete block in the whole buffer) is partly covered: rows 0–2 of case 0 exercise the per-row empty prefix, and the whole-buffer zero extent is pinned by the entry-point test (which asserts it returns success and forms no launch), but no score comparison covers it.qwen4feature — the CUDA-clippy matrix namespegainfer-kernelsonly on the default-features row, so every guard in this PR (the two device-free halves, the ops unit test and the traps) is compiled by CI only as an empty target and runs today only under the manual commands above. Wiring a row for this line's kernels is a follow-up, and it must setPEGAINFER_REQUIRE_GPU=1so the device half cannot skip and still report green.pegainfer-qwen4has no runtime,launchrefuses, and Stage 2/4 (sparse attention forward, assembly into the layer) do not exist — no request path reaches this kernel. The golden parity is the whole of the evidence, and thenvcc does not list compute_120f; compiling CUDA kernels for raw sm_120warning is expected, not a failure.Next: P3
Selection —
topk(min(budget / compress_ratio, n_complete))over the block scores, expanding each kept block back to itscompress_ratiotokens, then appending the unconditionallen(visible) % compress_ratiotail (buffer upper boundbudget + compress_ratio - 1= 2051). It depends on the same golden'scase*_mask(already in the fixture) and not on this kernel. The reference leaves tied scores' order open, so a block tied with the last kept one at the cut is an equally valid pick — the long case's bf16 noise documented above is exactly the perturbation that can reach that tie, which is worth keeping in mind when that gate picks its own bound.Refs
#1105 (QSA), #1150 (the golden this gate replays), #1103 / #1104 (the line)