Repository navigation
feat(qwen4): add Qwen3.8-Flash-Next QSA indexer golden - #1150
Conversation
Dump the QSA token indexer reference (transformers 5.16.1 `qwen4_exp`) into a fixture the Stage 1 kernel gate replays: per-block scores + selected token indices for 6 cases, covering short/medium inputs, the budget-boundary truncation (block_topk reached with a shrunk budget) and the always-kept tail. The re-derivation is validated against the reference module's own mask in-dump. Signed-off-by: bohaotu <bohaotu@pony.ai>
Code review范围：2 文件、无运行时数据流（dump 脚本 + fixture）。逐文件认领。
最小加固：修 M1（加 provenance 张量 + 载入断言）、M2（显式有效性掩码）；L1 记录为已知边界（加真实 budget 长样例体积过大，留待集成用例）。 |
…alidity Address PR review: add `format_version` / `config_sha256` / `ref_transformers` so a consumer can assert the oracle's origin, and an explicit `scores_valid` mask so the NaN-padded block scores cannot be mis-consumed (no silent nan_to_num pass). L1 (real-budget truncation) kept as a known bound. Signed-off-by: bohaotu <bohaotu@pony.ai>
修改措施Pushed
复验：fixture 重生成（66 tensor），provenance = 5.16.1 / |
sysuTobo
changed the title
Drop references to local-only planning docs and tighten the docstring/comments. No behavior change; fixture is unaffected. Signed-off-by: bohaotu <bohaotu@pony.ai>
Address the family-level finding from the pegainfer-project#1153/pegainfer-project#1155 reviews: RMS_EPS feeds Config.rms_norm_eps but was missing from the config_digest canon, so editing eps changed the golden while the stored provenance stayed green. Canon now carries eps like the sibling scripts. Fixture regenerated (66 tensors): config_sha256 b9fd359a -> 20cb29e2; 59 tensors bitwise stable; the six case*_scores carry bf16 ULP-level drift (<=4.8e-4) from re-dumping under torch 2.14.1+cpu / transformers 5.16.1 instead of the original torch 2.5.1+cu121 environment. Second same-environment run is byte-identical; all six cases pass mask_ok and the reload assertions. Signed-off-by: bohaotu <bohaotu@pony.ai>
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. |
|
We pushed two commits on top of yours. The fixture now stores what the reference itself computes, covers the real 2048 token budget, and takes its geometry from the frozen config. The dump now records the scores and the selection directly from the reference. A scoped wrapper around Each case now stores the indexer's inputs after A new case of 2055 tokens checks its last six rows, which see 2050 to 2055 tokens. At 2050 and 2051 tokens every one of the 512 blocks is kept. At 2052 the indexer starts dropping one block, and from 2053 to 2055 it drops one block and keeps a tail of 1 to 3 tokens. In these rows 31 to 43 blocks score exactly zero, so the dropped block is one of several tied blocks. Case 5 covers a row whose dropped block scores lower than every kept block. The geometry and the fingerprint come from We removed the code these changes made unnecessary. That covers the scatter patch, which torch 2.14 does not need, the hand written We generated the fixture with transformers 5.16.1 and torch 2.14.0 on arm64. |
|
Post-merge note: the dumper needs a newer torch than the dev/pod env Verified against the merged revision on the pod env ( The reference builds Nothing in the repo pins torch ( Not a blocker — flagging it so a future re-dump is not a surprise. |
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).
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>
* feat(qwen4): score QSA token-indexer blocks in native CUDA 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 #1150 with nothing to replay it. This lands the middle of the three units #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 #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 #1105's next unit, which is why this kernel takes no budget: the reference scores every complete block and truncates only at selection. Refs #1105 Signed-off-by: bohaotu <bohaotu@pony.ai> * test(qwen4): pin the QSA indexer's refusal contract device-free 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 #1105 Signed-off-by: bohaotu <bohaotu@pony.ai> * refactor(qwen4): name the QSA indexer's kernels after the model line 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 #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> * refactor(qwen4): drop the two always-true geometry checks 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> * docs(qwen4): trim the QSA indexer's comments 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> * fix(qwen4): make the indexer gate catch what it claims, and trap bad 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> * refactor(qwen4): read the shared attention RoPE table in the indexer 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> * refactor(qwen4): drop the frozen-geometry rail the review found redundant 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> * test(qwen4): make the trap gates prove the launch reached the kernel 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> * fix(qwen4): trap indexer rows that see past their own token and correct the gate notes Signed-off-by: Feathbow <feathbow@gmail.com> * test(qwen4): require the kernel's trap in the QSA trap gates and drop stale notes 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>
2 participants
Summary
Stage 1 (P1) of #1105: the QSA token indexer reference golden that the scoring/selection kernel gate replays.
tools/accuracy/dump_qwen38fn_qsa_indexer_golden.py— dumps per-block scores + selected token indices from the reference (transformers 5.16.1qwen4_exp).test_data/qwen38fn-qsa-indexer-golden.safetensors— 6 cases: short/medium, budget-boundary truncation (shrunk budget → block_topk 4), and the always-kept tail.Validation
torch.equal(ref_mask, my_mask): the re-derivation reproduces the reference module's own attention mask — already passes on all 6 cases (a mismatch aborts the dump).No runtime code changes; fixture + dump script only.
Refs #1105