diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f43bd969e..583c95ccc 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -151,6 +151,10 @@ jobs: if: ${{ !cancelled() && steps.deps.outcome == 'success' }} run: cargo test --release --locked -p pegainfer-sim --lib + - name: Run pegainfer-qwen4 tests + if: ${{ !cancelled() && steps.deps.outcome == 'success' }} + run: cargo test --release --locked -p pegainfer-qwen4 --tests + - name: Run kvbm-logical unit tests if: ${{ !cancelled() && steps.deps.outcome == 'success' }} run: cargo test --release --locked -p kvbm-logical --lib @@ -188,6 +192,7 @@ jobs: cargo clippy --release --locked -p pegainfer-build -p pegainfer-frontend + -p pegainfer-qwen4 -p pegainfer-sim -p kvbm-logical --all-targets -- -D warnings diff --git a/Cargo.lock b/Cargo.lock index 4c69a4b20..b72d6a87d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3901,6 +3901,7 @@ name = "pegainfer-qwen4" version = "0.1.0" dependencies = [ "anyhow", + "half", "pegainfer-frontend", "safetensors", "serde", diff --git a/docs/index.md b/docs/index.md index d4be427aa..7985351f2 100644 --- a/docs/index.md +++ b/docs/index.md @@ -65,7 +65,7 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | Path | TL;DR | | --- | --- | -| `models/qwen4/bring-up.md` | New line (`qwen4_exp`, first checkpoint Qwen3.8-Flash-Next), not a Qwen3.5 variant; frozen at `de4b8e4d`. Landed: an identity probe and one typed config validation that pins structure rather than widths, so a smaller `qwen4_exp` checkpoint passes it, models the router's upstream default (`norm_topk_prob`) and derives the n-gram table from five config fields via 16 per-head primes. Measured traps: untied LM head, **no final-norm tensor** (the readout is the hyper-connection mixer), `ple_layer_ids` is 1-indexed so the n-gram block is on `layers.1`, the n-gram table is **102.4 GB at BF16** (the quoted ~51 GB is fp8), `output_gate_type: "sigmoid"` is a **live** field selecting the GDN gated-norm activation (the inverse of the Qwen3.8-27B lesson), two RMSNorm conventions in one layer, and the PLE conv is dilated despite an identical tensor shape. Tracked in #1103 / #1104. | +| `models/qwen4/bring-up.md` | New line (`qwen4_exp`, first checkpoint Qwen3.8-Flash-Next), not a Qwen3.5 variant; frozen at `de4b8e4d`. Landed: an identity probe and one typed config validation that pins structure rather than widths, so a smaller `qwen4_exp` checkpoint passes it, models the router's upstream default (`norm_topk_prob`) and derives the n-gram table from five config fields via 16 per-head primes, plus the PLE n-gram block golden and its two readers — a device-free gate (hash exact under i64 equality, block output bit-identical, four negative controls) whose hash half is production host code, and the block's first device kernels in `pegainfer-kernels` behind a new `qwen4` feature (204799 of the golden's 204800 stored outputs come back bit-for-bit, three device-side controls). Measured traps: untied LM head, **no final-norm tensor** (the readout is the hyper-connection mixer), `ple_layer_ids` is 1-indexed so the n-gram block is on `layers.1`, the n-gram table is **102.4 GB at BF16** (the quoted ~51 GB is fp8), `output_gate_type: "sigmoid"` is a **live** field selecting the GDN gated-norm activation (the inverse of the Qwen3.8-27B lesson), two RMSNorm conventions in one layer, and the PLE conv is dilated despite an identical tensor shape. Tracked in #1103 / #1104. | ## models / gemma4 diff --git a/docs/models/qwen4/bring-up.md b/docs/models/qwen4/bring-up.md index 1ce5b308a..2fb05359a 100644 --- a/docs/models/qwen4/bring-up.md +++ b/docs/models/qwen4/bring-up.md @@ -1,6 +1,6 @@ # Qwen4 bring-up -> **TL;DR:** `Qwen/Qwen3.8-Flash-Next`, the first checkpoint of the Qwen4 architecture preview, is a **new line** (`qwen4_exp`), not a Qwen3.5 variant — its `model_type` differs, so the Qwen3.5 probe refuses it by construction. Frozen at revision `de4b8e4d43b917e7706784d8bb445c9af86a3540`. Landed so far: an identity probe and one typed config validation that pins *structure rather than widths* and derives what `config.json` leaves implicit, plus a tensor manifest that derives **exactly** the checkpoint's 1294 text tensors, tested against a snapshot of every shard header and not yet called at runtime. All of the model crate is CPU-only; its first device code lives beside it in `pegainfer-kernels` behind that crate's own `qwen4` feature — the QSA token indexer's block-scoring kernels, gated against a committed reference golden. Nothing runs — detection succeeds and `launch` refuses, because the text graph is slice B. Four facts this line must not get wrong, all measured rather than assumed: the embedding is **untied**; there is **no final-norm tensor** (the readout is the hyper-connection mixer); `ple_layer_ids: [2]` is **1-indexed**, so the n-gram block lives on `layers.1`; and `output_gate_type: "sigmoid"` is a **live** field selecting the GDN gated-norm activation — the inverse of the Qwen3.8-27B lesson, where the same field name is inert. +> **TL;DR:** `Qwen/Qwen3.8-Flash-Next`, the first checkpoint of the Qwen4 architecture preview, is a **new line** (`qwen4_exp`), not a Qwen3.5 variant — its `model_type` differs, so the Qwen3.5 probe refuses it by construction. Frozen at revision `de4b8e4d43b917e7706784d8bb445c9af86a3540`. Landed so far: an identity probe, one typed config validation that pins *structure rather than widths* and derives what `config.json` leaves implicit, a tensor manifest that derives **exactly** the checkpoint's 1294 text tensors, tested against a snapshot of every shard header and not yet called at runtime, and a **device-free replay gate** for the PLE n-gram block golden (hash exact under i64 equality, block output bit-identical) whose hash half is production host code, `pegainfer_qwen4::ngram::NgramHasher`. The PLE block also has **device kernels** now, in `pegainfer-kernels` behind the line's `qwen4` feature, and they return 204799 of the golden's 204800 stored outputs bit-for-bit. The crate itself is still CPU-only and nothing runs — detection succeeds and `launch` refuses, because the text graph is slice B. Four facts this line must not get wrong, all measured rather than assumed: the embedding is **untied**; there is **no final-norm tensor** (the readout is the hyper-connection mixer); `ple_layer_ids: [2]` is **1-indexed**, so the n-gram block lives on `layers.1`; and `output_gate_type: "sigmoid"` is a **live** field selecting the GDN gated-norm activation — the inverse of the Qwen3.8-27B lesson, where the same field name is inert. > > **Last touched:** 2026-10 @@ -37,6 +37,61 @@ The committed `config.json` (`pegainfer-qwen4/tests/frozen_config.json`) is the - **Config** (`src/config.rs`) — every other rule, behind one typed parse (`RawConfig` → `Config`), so detection and loading cannot disagree. The line's `ModelLine::probe` builds the whole `Config`, so `detect` refuses an unservable `qwen4_exp` config before the frontend starts. It follows the Gemma 4 pattern of pinning *structure rather than widths*. Pinned because an alternative would need a kernel or serving path this line does not have: the GDN head dims (they are the Triton AOT specialization keys), `head_dim` 256, the conv kernel ceiling, `rope_type: default` (no YaRN path), an even rotary width (the reference rotates in pairs, so an odd one would rotate one more dim than computed), `output_gate_type: sigmoid`, `mamba_ssm_dtype: float32`, `dtype: bfloat16` with no `quantization_config`, and untied embeddings. The official FP8 sibling and nvidia's NVFP4 keep `dtype: bfloat16` and differ from this checkpoint only by `quantization_config`, so that block is what refuses them; the dtype alone would not. The GDN and attention limits come from the kernels' shared wrappers (`pegainfer_kernels::ops`, behind the kernels' `qwen35` feature): this line's 36 `linear_attention` layers will reuse the Qwen3.5 decode kernels through them once slice B gives the line a device path, and `src/config.rs` mirrors the three GDN constants because the config layer must still test with no device. Deliberately *not* pinned: expert count, hidden size, layer count, head counts, `hc_count` — so a smaller `qwen4_exp` checkpoint passes the same rules, which is what the L20 small-model deployment case needs (`smaller_widths_still_pass`). `mrope_interleaved` and `mrope_section` are accepted rather than refused: a text-only sequence is their degenerate case in the reference itself, which expands one position per token across all three axes before interleaving. The decision is tracked in #1105. Beyond the pins, the config mirrors the rules the reference's own `validate_architecture` enforces — `hc_count > 1`, QSA's `indexer_kv_heads == 1` with the budget a whole multiple of the compress ratio, the rotary width fitting the index head, `rope_theta` present, and PLE confined to a **single** `linear_attention` layer — plus the two head-group integrality rules Qwen3.5's config enforces (`AttentionHeadDivisibility`, `LinearHeadDivisibility`). Multi-layer PLE is refused for now: the reference offsets each layer's prime run by its position in `ple_layer_ids` (`global_head_idx = L * ngram_heads + i`), and this crate derives one table. `norm_topk_prob` is **absent from `config.json`** and comes from the upstream default (`true`, and it decides whether the router renormalizes its top-k), so it is modelled with that default: a port that reads only `config.json` silently drops the router renormalization — this is #1068's mirror image. The other absent field, `seed`, only matters if the n-gram hash multipliers are recomputed instead of read from the checkpoint, so it is not modelled. - **Manifest** (`src/manifest/`, test only until slice B's loader calls it) — the tensor contract, ported from `pegainfer-gemma4/src/manifest/`: `schema.rs` derives it from `Config` and `validate.rs` classifies headers against it. It derives **exactly** the frozen text set, 1658 keys − 333 `model.visual.*` − 31 `mtp.*` = **1294 tensors**, with `lm_head.weight` the only one outside `model.language_model.`. Beyond the untied head, it differs from Gemma 4's in rank-3 shapes (the depthwise conv and two fused expert stacks per layer, which carry no `.weight` suffix) and in a mandatory skip set, because all 333 vision tensors share shard 1 with 16 text tensors. `tests/frozen_tensors.json` (148 KB, written by `tools/accuracy/dump_qwen4_frozen_tensors.py`) holds every shard header of the frozen revision, and `the_contract_is_the_frozen_checkpoint` checks all 1658 tensors against the contract. - **Registration** — the server binary hand-mirrors every line's probe identities in a hint table that cannot be derived away. With the feature off, `feature_gate_hint` names `--features qwen4` for both `qwen4_exp` and `qwen4_exp_text`; with it on, the hint stays silent. A test under the feature binds that row to this crate's `MODEL_TYPE` / `TEXT_MODEL_TYPE`, so the mirror cannot drift from the probe. The hint is **identity-only** — it cannot run a compiled-out probe — so it also fires for the quantized siblings and the tied/YaRN variants of `qwen4_exp`, which the config refuses once the feature is on; the hint names the feature to rebuild with, not a promise of service. `pegainfer-server/tests/detect_gate.rs` then drives the shipped binary against a model directory holding this crate's frozen `config.json`, both ways: the stock build must name the feature, the enabled build must detect the line and refuse to serve. That is the only coverage of `main`'s CLI → registry → hint → `launch` wiring, and neither outcome reaches CUDA — `launch` bails before the engine is built. +- **PLE n-gram golden and its host reader** (`test_data/qwen38fn-ple-golden.safetensors` + recorded by `tools/accuracy/dump_qwen38fn_ple_golden.py`, read by + `tests/ple_ngram_replay.rs`) — the per-layer-embedding block captured from the + reference at this config's real geometry (hidden 2560, `hc_count` 4 → a 10240-wide + stream, 16 heads × 160, conv kernel 4 dilated by 3), plus a **device-free** gate that + replays it: the mul-xor n-gram hash under i64 equality, zero tolerance (128/128 and + 192/192 ids exact on the S=8 and S=12 cases, the latter carrying a mid-sequence EOS so + the window reset has a boundary on both sides), and the block itself from the recorded + `key`/`value` projections — bit-identical on all 81920 + 122880 stored bf16 outputs, + with the tolerance sized from a perturbation study of the chain's only libm-dependent + op rather than guessed. Four negative controls (drop the EOS window reset, dilation 1, + unsigned sqrt, swapped bigram/trigram head order) are each asserted to fail. Two + boundaries the fixture cannot carry: the n-gram table (102.4 GB at BF16) is not stored, + so the gate starts at the block's post-projection inputs, and the recording is + whole-sequence, so the decode-side states (the 2-token id window, the 9-slot dilated + conv state) stay ungated. The hash reads `layer_multipliers` from the fixture instead of + recomputing them from `seed`, which is what keeps `seed` unmodelled above. The device + side — the table gather and the dilated conv kernel — is #1106. The crate's tests and + clippy now run in CI's CPU jobs, which is what keeps the no-device property honest. +- **PLE n-gram golden and its two readers** (`test_data/qwen38fn-ple-golden.safetensors` + recorded by `tools/accuracy/dump_qwen38fn_ple_golden.py`) — the per-layer-embedding + block captured from the reference at this config's real geometry (hidden 2560, + `hc_count` 4 → a 10240-wide stream, 16 heads × 160, conv kernel 4 dilated by 3). + Two consumers, one golden: + - **Host** (`pegainfer-qwen4/tests/ple_ngram_replay.rs`, device-free) — the mul-xor + n-gram hash under i64 equality, zero tolerance (128/128 and 192/192 ids exact on the + S=8 and S=12 cases, the latter carrying a mid-sequence EOS so the window reset has a + boundary on both sides), and an f32 CPU replay of the block from the recorded + `key`/`value` projections — **bit-identical** on all 81920 + 122880 stored bf16 + outputs, with the tolerance sized from a perturbation study of the chain's only + libm-dependent op rather than guessed. Four negative controls (drop the EOS window + reset, dilation 1, unsigned sqrt, swapped bigram/trigram head order) are each + asserted to fail. The hash itself is **production code**, `pegainfer_qwen4::ngram:: + NgramHasher` — the gate calls it, and its two controls compose the same primitives + with one knob each, so a control cannot drift onto a private copy of the hash. + - **Device** (`pegainfer-kernels/tests/qwen4_ple_block_parity.rs`, behind that + crate's `qwen4` feature) — the block's three kernels in + `csrc/qwen4/qwen4_ple_block.cu` (branch-major split, per-branch gate, + dilated conv + residual), driven through `ops::ple_block_into`, with the three + grouped norms delegated to the existing batched `(1+w)` entry point. Measured on + sm_120 / CUDA 12.9.1: **204799 of the 204800 stored outputs come back bit-for-bit** + and the remaining one is a single bf16 ulp (1.221e-4 at 0.0244), against a bound of + 1.6e-2 — the host gate's own number, so both halves of the consumer read against one + scale. Three device-side controls (unsigned sqrt 3.32e-1, a full-stream gate dot + 2.75e-1, dilation 1 7.89e-1) are each asserted to leave that bound, at 21x, 17x and + 49x outside it. + Two boundaries the fixture cannot carry: the n-gram table (102.4 GB at BF16) is not + stored, so both gates start at the block's post-projection inputs, and the recording + is whole-sequence, so the decode-side states (the 2-token id window, the 9-slot + dilated conv state) stay ungated. The hash reads `layer_multipliers` from the fixture + instead of recomputing them from `seed`, which is what keeps `seed` unmodelled above. + What remains for #1106 is the table gather and its loader. The crate's tests and + clippy run in CI's CPU jobs, which is what keeps the host half's no-device property + honest; **no CI job enables the kernels `qwen4` feature**, so the device half runs + only under the manual command in that gate's header. **1-D lengths are part of the contract** (259 vectors), because the shared loader takes them on trust (#1068), and here many carry meaning: `A_log[48]`, `dt_bias[48]`, `norm.weight[128]`, `q_norm[256]`, the indexer layernorms at `[128]`, `layer_multipliers[3]` and the two `[16]` n-gram partition vectors. In the Qwen3.5 GDN loader this line would reuse, `A_log` goes through `tensor_f32_cow`, which accepts BF16 or F32 but never checks the length, and on a single rank `dt_bias` goes through `load_tensor_1d`, which only checks that the byte count is even before reading the bytes as BF16. @@ -83,9 +138,34 @@ Still to verify is the reuse on this line's activations. The tensor contract der ## Traps -- **`output_gate_type` has three unrelated consumers.** On Qwen3.8-27B it is inert and must *not* be implemented (`docs/models/qwen35/support-qwen38.md`). Here it is live and selects the GDN gated-norm activation; the shared `csrc/shared/norm.cu` gated norm **became selectable in #1129** (a sigmoid entry point beside Qwen3.5's SiLU, merged before this line's detection PR), so what remains is this line's call site, which lands with the text graph. It is *also* the marker the frontend's `reasoning_effort` layer arms on. Read the field's consumer, never its name. -- **Two RMSNorm conventions in one layer.** `(1+w)` for `q_norm`/`k_norm`, the indexer layernorms, `hc_norm` and the PLE norms; plain `w` only for the GDN gated norm. The four 10240-wide ones are additionally **grouped by 2560** — likely a reshape of the existing batched `(1+w)` kernel rather than a new kernel, but unverified. -- **The PLE conv is dilated** (kernel 4, dilation `ngram_size` = 3, state length 9) even though its weight tensor `[10240, 1, 4]` is shaped exactly like the GDN conv's. The in-tree `conv1d.cu` is non-dilated and caps its state at four taps, so it cannot serve this. +- **`output_gate_type` has three unrelated consumers.** On Qwen3.8-27B it is inert and + must *not* be implemented (`docs/models/qwen35/support-qwen38.md`). Here it is live and + selects the GDN gated-norm activation; the shared `csrc/shared/norm.cu` gated norm + **became selectable in #1129** (a sigmoid entry point beside Qwen3.5's SiLU, merged + before this line's detection PR), so what remains is this line's call site, which + lands with the text graph. It is *also* the marker the frontend's `reasoning_effort` + layer arms on. Read the field's consumer, never its name. +- **Two RMSNorm conventions in one layer.** `(1+w)` for `q_norm`/`k_norm`, the indexer + layernorms, `hc_norm` and the PLE norms; plain `w` only for the GDN gated norm. The + four 10240-wide ones are additionally **grouped by 2560**, and the grouping is not a + free reshape: `rms_norm_batched_offset_cuda` takes its row stride *from* its + `hidden_dim`, so it can only normalize contiguous 2560 rows, while this stream's four + groups are interleaved row-major. Serving one grouped norm therefore costs a + branch-major transpose plus four calls, and the PLE block's three of them cost twelve. + #1158's single-launch grouped entry point removes both halves; until it lands, the + transpose is the price. The PLE block's device gate measures the cost of reusing that + kernel rather than writing an exact one: its `rsqrt.approx.ftz.f32` reciprocal flips + the bf16 rounding of a couple of normed elements in 10^5, and one of them reaches the + block output as a single ulp. +- **The PLE conv is dilated** (kernel 4, dilation `ngram_size` = 3, state length 9) even + though its weight tensor `[10240, 1, 4]` is shaped exactly like the GDN conv's. The + in-tree `conv1d.cu` is non-dilated and caps its state at four taps, so it cannot serve + this; `csrc/qwen4/qwen4_ple_block.cu` carries the block's own, whose nine-row + zero state is a `source >= 0` bound rather than a buffer, so nothing has to be sized or + cleared per sequence. Two rounding points travel with it: the four taps sum in f32 and + the sum rounds to bf16 **before** the silu, and the silu rounds again before the + residual add — the same double-rounding shape the gate has around its + `1 / sqrt(hidden_size)` divide, and collapsing either pair costs ~1 ulp everywhere. - **Attention writes back at `hidden_size`.** `o_proj [2560, 6144]` returns to the 2560-wide hidden size, and `attn_hyper_connection.block_inject_weight [4, 10240]` carries it into the 10240-wide hyper-connection stream. - **GQA group 12 is not a compiled decode group.** `SUPPORTED_GQA_GROUP_SIZES` is `[1, 2, 3, 4, 6, 8]`, so the FlashInfer batch decode path cannot serve this attention as it stands. QSA is sparse and needs kernels of its own, so whether its decode is captured in a graph is slice B's decision, made with those kernels. - **The QSA gate's bound is measured per case, and its floor was measured too.** Cases 0–5 sit at ≤9.6e-7 and carry a 2e-6 bound; the long case's 2.4e-3 residual is not traced yet; its shape suggests one flipped bf16 rounding decision, and the reference moves by that size across environments (re-deriving its case 2 on x86 from the fixture's own inputs moves one of its 496 scores by 2.2e-3), so it keeps 5e-3 until traced. What the bound must catch was measured on this fixture: keeping the pooled mean in f32 moves scores by 1.55e-2, 100x `rms_eps` by 1.25e-2, dropping `rms_eps` or pooling by sum by 3.3e-3, the pooled key's RoPE by 2.7, and a post-sum relu by 0.83 — every fault sits two to six orders beyond the tight bound. Architecture variance does not loosen it: the gate prints identical numbers on sm_89 and sm_120, and the kernel reads the shared attention cos/sin table (built by `pegainfer_core::rope` with `frequency_dim` equal to the 64 rotated dims, as the reference spaces its frequencies) rather than computing trig, so there is nothing device-specific left to vary. `#1105`'s selection stage must still treat a block tied with the last kept one at the cut as an equally valid pick, because the long case's noise can flip that order. @@ -99,6 +179,7 @@ Still to verify is the reuse on this line's activations. The tensor contract der | GDN decode reuse, executed | slice B: the wrappers are shared in `pegainfer-kernels::ops` and the contract pins their geometry and tensor names, but nothing has run this line's GDN block against the reference. The model-graph fixture that does so is slice B's M0 exit (#1127); the GDN prefill path and the gated norm composition land with it. | | Bookend gate (embedding → mixer → lm_head → argmax) | 1 GPU, three shards (1 + 130 + 131, 5.29 GB, ~2.6 GB of weights on device) and a CPU f64 reference for the mixer. No layer weights needed. | | HF golden fixtures (logits + intermediate probes) | the fleet — 360 GB, plus ≥102.4 GB host RAM if the n-gram table is materialized | +| PLE block served end to end | the table's loader and gather (#1106) — 102.4 GB at BF16 over 128 shards, the one piece between `NgramHasher::ids` and the projections this gate starts from — plus the decode-side states (the 2-token id window, the 9-slot conv state), which a whole-sequence recording cannot pin. The block's own arithmetic is done on both sides. | | Tokenizer / chat-template ID fixtures | transformers **5.16.0+** — the checkpoint's own tokenizer is `Qwen2Tokenizer`, which is older, but fixtures are compared against the reference implementation; nothing in the tree commits token-ID fixtures today | Note that `pegainfer-server` cannot be compiled on a CPU-only box, so the registration above is verifiable only where a CUDA build works. diff --git a/pegainfer-kernels/csrc/qwen4/qwen4_ple_block.cu b/pegainfer-kernels/csrc/qwen4/qwen4_ple_block.cu new file mode 100644 index 000000000..f299f3581 --- /dev/null +++ b/pegainfer-kernels/csrc/qwen4/qwen4_ple_block.cu @@ -0,0 +1,384 @@ +// Qwen3.8-Flash-Next PLE block (per-layer embedding), the device half. +// +// Frozen semantics: transformers 5.16.1 `qwen4_exp`, +// `Qwen4ExpTextPLELayer.forward`. From the two recorded projections and the +// residual stream, over `hc_count` hyper-connection branches of `hidden_size` +// channels each (`width = hc_count * hidden_size`): +// +// k = (1 + w_key) RMSNorm_2560(key) // grouped: one norm per branch +// q = (1 + w_query) RMSNorm_2560(hidden) // grouped, same split +// dot_rg = sum_j bf16(k[r,g,j] * q[r,g,j]) // bf16 per product +// pre_rg = bf16(bf16(dot_rg) / sqrt(hidden_size)) // TWO roundings +// gate_rg = bf16(bf16(max(|pre_rg|, 1e-6)).sqrt() * sign(pre_rg)) +// gated = bf16(sigmoid(gate_rg) * value[r]) // value broadcast over g +// c = (1 + w_conv) RMSNorm_2560(gated) +// out[r,ch]= bf16(gated[r,ch] + silu(bf16(dilated_causal_conv1d(c)[r,ch]))) +// +// The conv is depthwise over `width` channels, kernel 4, dilation `ngram_size` +// (3), so its receptive field is `(kernel - 1) * dilation + 1 = 10` rows and its +// zero initial state is the 9 rows before the sequence. +// +// Three things here are easy to get wrong and are the reason the rounding points +// are spelled out above: +// +// * The gate rounds TWICE around the `1 / sqrt(hidden_size)` divide: the group +// sum to bf16 first, the quotient to bf16 second. Collapsing them is worth +// ~1 ulp on every gate. It is a divide by `sqrt_hidden`, not a multiply by a +// reciprocal — the two round differently, which is why `sqrt_hidden` is an +// argument rather than a compiled constant: the caller computes it the way +// the reference does. +// * The sqrt is SIGNED. `abs().clamp_min(1e-6).sqrt()` loses the sign, and the +// reference restores it with `torch.sign` — which is 0 at 0, not 1, so a zero +// gate sigmoid's to exactly 0.5 rather than to `sigmoid(1e-3)`. +// * The dot is per BRANCH (2560), not per stream (10240). One gate per +// branch, and all four branches then scale the SAME `value` row. +// +// All three grouped norms are the existing batched `(1 + w)` RMSNorm entry point, +// `rms_norm_batched_offset_cuda`, called once per branch from the Rust ops layer: +// it normalizes `seq_len` contiguous rows of `hidden_dim`, and this stream's +// branches are interleaved row-major, so the ops layer splits into branch-major +// scratch first (kernel 1) and the norms then run on contiguous rows. That split +// is what the first kernel below is for, and it is why the block's three grouped +// norms cost twelve launches rather than three. +// +// Target: the sm_89-first portable subset — no TMA, no cluster, no wgmma. All +// three kernels are SIMT: two are element-wise copies/maps and the third is a +// 2560-wide dot with a shuffle-and-shared-memory block reduction, a shape an +// `m16n8k16` body cannot tile without padding one row to sixteen. +// +// Negative-control surface: `gate_span` and `signed_sqrt` are arguments rather +// than compiled constants so the parity gate can turn exactly one knob on the +// device path — a full-stream dot, or an unsigned sqrt — and assert the output +// leaves its bound. `dilation` is genuine geometry (`ngram_size`), and the parity +// gate drives it to 1 for the same reason. A serving path reaches none of them: +// `ops::ple_block_into` is the production entry and passes the reference values, +// and only `ops::ple_block_variant_into` — whose sole callers are the gate's +// three controls — turns a knob. + +#include "../common.cuh" +#include "../shared/ffi_guard.cuh" + +#include +#include +#include + +namespace { + +constexpr int kHidden = 2560; // hidden_size == ple_embed_dim +constexpr int kHc = 4; // hc_count +constexpr int kWidth = kHidden * kHc; +constexpr int kConvKernel = 4; // ple_conv_kernel_size +// The reference's `clamp_min(1e-6)` on the gate's magnitude. +constexpr float kGateClampMin = 1.0e-6f; +// One CTA per (row, branch) for the gate; 256 threads over 2560 channels is 10 +// elements per thread, and 8 warps keeps the cross-warp reduction one round. +constexpr int kGateThreads = 256; +constexpr int kGateWarps = kGateThreads / WARP_SIZE; +// Element-wise passes: 256 lanes over a width that is a multiple of 32, so a +// warp always covers consecutive channels. +constexpr int kMapThreads = 256; + +static_assert(kWidth % WARP_SIZE == 0, "a warp must cover whole channels"); +static_assert(kHidden % kGateThreads == 0, "the gate CTA must tile the group"); + +// Round to bf16 and back, which is what the reference does every time a +// submodule boundary is crossed. Written out rather than left to the compiler +// because the rounding IS the semantics. +__device__ __forceinline__ float bf16_round(float value) { + return __bfloat162float(__float2bfloat16(value)); +} + +// `torch.sign`: 0 at zero (and at -0.0), unlike a `copysign`-style restore. +__device__ __forceinline__ float torch_sign(float value) { + return static_cast(value > 0.0f) - static_cast(value < 0.0f); +} + +// Sum over the whole CTA, left in every thread: the gate scalar is uniform +// across the CTA and each thread then writes its own channels, so broadcasting +// it through shared memory costs one barrier and saves a second reduce. +__device__ __forceinline__ float block_all_sum(float value, float* shared) { + value = warp_reduce_sum(value); + const int lane = threadIdx.x % WARP_SIZE; + const int warp = threadIdx.x / WARP_SIZE; + if (lane == 0) { + shared[warp] = value; + } + __syncthreads(); + float total = 0.0f; +#pragma unroll + for (int w = 0; w < kGateWarps; ++w) { + total += shared[w]; + } + return total; +} + +// Row-major `[seq_len, kWidth]` -> branch-major `[kHc][seq_len][kHidden]`. +// The grouped norms run on contiguous `kHidden` rows, and the block's branches +// are interleaved in the stream, so this is a transpose of the two middle axes. +__global__ void qwen4_ple_split_branches_kernel( + const __nv_bfloat16* __restrict__ row_major, + __nv_bfloat16* __restrict__ branch_major, int seq_len) { + const int64_t total = static_cast(seq_len) * kWidth; + const int64_t index = + static_cast(blockIdx.x) * kMapThreads + threadIdx.x; + if (index >= total) { + return; + } + const int64_t row = index / kWidth; + const int64_t channel = index - row * kWidth; + const int64_t branch = channel / kHidden; + const int64_t within = channel - branch * kHidden; + branch_major[(branch * seq_len + row) * kHidden + within] = row_major[index]; +} + +// One gate scalar per (row, branch), then `sigmoid(gate) * value` over that +// branch's channels. Inputs are branch-major `[kHc][seq_len][kHidden]`; `value` +// is `[seq_len, kHidden]` and is broadcast over the branches that share it. +// +// `gate_span` is `kHidden` for the reference's per-branch dot and `kWidth` for +// the full-stream control, which folds all four branches into one dot — every +// branch then gets the same gate, which is the fault the control stands for. +__global__ void qwen4_ple_gate_kernel( + const __nv_bfloat16* __restrict__ key_normed, + const __nv_bfloat16* __restrict__ query_normed, + const __nv_bfloat16* __restrict__ value, + __nv_bfloat16* __restrict__ gated, int seq_len, float sqrt_hidden, + int gate_span, int signed_sqrt) { + __shared__ float shared[kGateWarps]; + + const int block = blockIdx.x; + const int branch = block % kHc; + const int row = block / kHc; + const int branches = (gate_span == kWidth) ? kHc : 1; + const int first = (gate_span == kWidth) ? 0 : branch; + + // Each product rounds to bf16 before it is summed, as the reference's does: + // the multiply happens inside the submodule that produces it. + float partial = 0.0f; + for (int j = threadIdx.x; j < kHidden; j += kGateThreads) { + for (int b = 0; b < branches; ++b) { + const int64_t at = + (static_cast(first + b) * seq_len + row) * kHidden + j; + const float k = __bfloat162float(key_normed[at]); + const float q = __bfloat162float(query_normed[at]); + partial += bf16_round(k * q); + } + } + const float dot = block_all_sum(partial, shared); + + // Two bf16 roundings around the divide, then the signed sqrt, then the + // sigmoid — every one of them a rounding point of the reference. + const float pre_gate = bf16_round(bf16_round(dot) / sqrt_hidden); + const float magnitude = bf16_round(fmaxf(fabsf(pre_gate), kGateClampMin)); + const float root = bf16_round(sqrtf(magnitude)); + const float restore = signed_sqrt != 0 ? torch_sign(pre_gate) : 1.0f; + const float gate = bf16_round(root * restore); + const float weight = bf16_round(1.0f / (1.0f + expf(-gate))); + + const __nv_bfloat16* value_row = + value + static_cast(row) * kHidden; + __nv_bfloat16* gated_row = + gated + (static_cast(branch) * seq_len + row) * kHidden; + for (int j = threadIdx.x; j < kHidden; j += kGateThreads) { + gated_row[j] = __float2bfloat16(weight * __bfloat162float(value_row[j])); + } +} + +// `out = gated + silu(bf16(dilated causal depthwise conv(conv_in)))`. +// +// `conv_in` and `gated` are branch-major `[kHc][seq_len][kHidden]`; `out` is +// row-major `[seq_len, kWidth]`, which is the stream the block injects into. +// `conv_weight` is `[kWidth, 1, kConvKernel]` flattened. The zero initial state +// of `(kConvKernel - 1) * dilation` rows is the `source >= 0` bound, not a +// buffer: a tap that reaches before the sequence contributes nothing. +// +// The four taps accumulate in tap order in f32, which is the reference's own +// order and short enough that a tree would only change which way the last bit +// falls. +__global__ void qwen4_ple_conv_residual_kernel( + const __nv_bfloat16* __restrict__ conv_in, + const __nv_bfloat16* __restrict__ gated, + const __nv_bfloat16* __restrict__ conv_weight, + __nv_bfloat16* __restrict__ out, int seq_len, int dilation) { + const int64_t total = static_cast(seq_len) * kWidth; + const int64_t index = + static_cast(blockIdx.x) * kMapThreads + threadIdx.x; + if (index >= total) { + return; + } + const int64_t row = index / kWidth; + const int64_t channel = index - row * kWidth; + const int64_t branch = channel / kHidden; + const int64_t within = channel - branch * kHidden; + + float taps[kConvKernel]; +#pragma unroll + for (int t = 0; t < kConvKernel; ++t) { + taps[t] = __bfloat162float(conv_weight[channel * kConvKernel + t]); + } + + // Branch-major base for this channel's column: `(branch * seq_len) * kHidden + // + within`, then walked by `source * kHidden`. Branch 0 is the only one where + // dropping the group factor would still land on the right element, which is + // exactly why the parity gate compares every branch rather than the first. + const int64_t column = branch * seq_len * kHidden + within; + const __nv_bfloat16* conv_column = conv_in + column; + const __nv_bfloat16* gated_column = gated + column; + + float acc = 0.0f; +#pragma unroll + for (int t = 0; t < kConvKernel; ++t) { + const int64_t source = + row - static_cast(kConvKernel - 1 - t) * dilation; + if (source >= 0) { + acc += __bfloat162float(conv_column[source * kHidden]) * taps[t]; + } + } + const float convolved = bf16_round(acc); + const float activated = + bf16_round(convolved / (1.0f + expf(-convolved))); + out[index] = __float2bfloat16( + __bfloat162float(gated_column[row * kHidden]) + activated); +} + +CUresult consume_last_cuda_error() { return map_cuda_error(cudaGetLastError()); } + +// Grid for the element-wise passes, one thread per stream element, or -1 if it +// does not fit an `int`. Formed in 64 bits: `seq_len * kWidth` overflows an +// `int` at ~209k rows, long before the grid dimension's own 2^31 - 1 limit, and +// a wrapped grid would launch nothing at all rather than too little. +int map_grid(int seq_len) { + const int64_t total = static_cast(seq_len) * kWidth; + const int64_t grid = (total + kMapThreads - 1) / kMapThreads; + return grid > static_cast(INT_MAX) ? -1 : static_cast(grid); +} + +} // namespace + +extern "C" { + +// Row-major `[seq_len, width]` -> branch-major `[hc_count][seq_len][hidden]`, +// for each of the two streams the grouped norms read. +// +// `width` is the caller's row stride, checked here: the body splits with the +// width and the group size it was compiled for, so a caller whose rows are +// strided differently would be read out of bounds rather than computed wrongly. +CUresult qwen4_ple_split_branches_cuda(const __nv_bfloat16* row_major, + __nv_bfloat16* branch_major, + int seq_len, int width, + int hidden, int hc_count, + cudaStream_t stream) { + if (row_major == nullptr || branch_major == nullptr) { + return CUDA_ERROR_INVALID_VALUE; + } + if (seq_len < 0) { + return CUDA_ERROR_INVALID_VALUE; + } + if (width != kWidth || hidden != kHidden || hc_count != kHc) { + return CUDA_ERROR_INVALID_VALUE; + } + if (seq_len == 0) { + return CUDA_SUCCESS; + } + const int grid = map_grid(seq_len); + if (grid < 0) { + return CUDA_ERROR_INVALID_VALUE; + } + qwen4_ple_split_branches_kernel<<>>( + row_major, branch_major, seq_len); + return consume_last_cuda_error(); +} + +// One gate scalar per (row, branch) and `sigmoid(gate) * value` over that +// branch's channels. All four buffers are branch-major `[hc_count][seq_len] +// [hidden]`, except `value`, which is `[seq_len, hidden]` and shared by the +// branches. +// +// `sqrt_hidden` is `sqrt(hidden_size)` in the f32 the reference divides by; it +// is an argument because a divide and a reciprocal multiply round differently. +// `gate_span` is `hidden` for the reference and `width` for the full-stream +// control; `signed_sqrt` is 1 for the reference and 0 for the unsigned control. +CUresult qwen4_ple_gate_cuda(const __nv_bfloat16* key_normed, + const __nv_bfloat16* query_normed, + const __nv_bfloat16* value, + __nv_bfloat16* gated, int seq_len, int hidden, + int hc_count, int gate_span, float sqrt_hidden, + int signed_sqrt, cudaStream_t stream) { + if (key_normed == nullptr || query_normed == nullptr || value == nullptr || + gated == nullptr) { + return CUDA_ERROR_INVALID_VALUE; + } + if (seq_len < 0) { + return CUDA_ERROR_INVALID_VALUE; + } + if (hidden != kHidden || hc_count != kHc) { + return CUDA_ERROR_INVALID_VALUE; + } + if (gate_span != kHidden && gate_span != kWidth) { + return CUDA_ERROR_INVALID_VALUE; + } + if (signed_sqrt != 0 && signed_sqrt != 1) { + return CUDA_ERROR_INVALID_VALUE; + } + // A non-finite or non-positive divisor would put NaN in every gate; refuse it + // rather than poison the stream silently. + if (!(sqrt_hidden > 0.0f)) { + return CUDA_ERROR_INVALID_VALUE; + } + if (seq_len == 0) { + return CUDA_SUCCESS; + } + // One CTA per (row, branch): the grid is `seq_len * kHc`, which has to be + // formed in 64 bits before it is narrowed, or a long enough sequence would + // wrap into a negative grid and launch nothing at all. + const int64_t gate_grid = static_cast(seq_len) * kHc; + if (gate_grid > static_cast(INT_MAX)) { + return CUDA_ERROR_INVALID_VALUE; + } + const int grid = static_cast(gate_grid); + qwen4_ple_gate_kernel<<>>( + key_normed, query_normed, value, gated, seq_len, sqrt_hidden, gate_span, + signed_sqrt); + return consume_last_cuda_error(); +} + +// The dilated causal depthwise conv, its silu, and the residual add back onto +// the gated value. `conv_in` and `gated` are branch-major, `out` is row-major +// `[seq_len, hc_count * hidden]`. +// +// `dilation` is `ngram_size`; the state it implies is `(conv_kernel - 1) * +// dilation` rows and is zero by the `source >= 0` bound rather than by a +// buffer, so any positive dilation is served without resizing anything. +CUresult qwen4_ple_conv_residual_cuda(const __nv_bfloat16* conv_in, + const __nv_bfloat16* gated, + const __nv_bfloat16* conv_weight, + __nv_bfloat16* out, int seq_len, + int hidden, int hc_count, + int conv_kernel, int dilation, + cudaStream_t stream) { + if (conv_in == nullptr || gated == nullptr || conv_weight == nullptr || + out == nullptr) { + return CUDA_ERROR_INVALID_VALUE; + } + if (seq_len < 0) { + return CUDA_ERROR_INVALID_VALUE; + } + if (hidden != kHidden || hc_count != kHc || conv_kernel != kConvKernel) { + return CUDA_ERROR_INVALID_VALUE; + } + if (dilation < 1) { + return CUDA_ERROR_INVALID_VALUE; + } + if (seq_len == 0) { + return CUDA_SUCCESS; + } + const int grid = map_grid(seq_len); + if (grid < 0) { + return CUDA_ERROR_INVALID_VALUE; + } + qwen4_ple_conv_residual_kernel<<>>( + conv_in, gated, conv_weight, out, seq_len, dilation); + return consume_last_cuda_error(); +} + +} // extern "C" diff --git a/pegainfer-kernels/src/ffi/qwen4.rs b/pegainfer-kernels/src/ffi/qwen4.rs index 9ff1270ba..a23c07387 100644 --- a/pegainfer-kernels/src/ffi/qwen4.rs +++ b/pegainfer-kernels/src/ffi/qwen4.rs @@ -31,3 +31,59 @@ unsafe extern "C" { stream: CUstream, ) -> CUresult; } + +unsafe extern "C" { + /// Row-major `[seq_len, width]` into branch-major + /// `[hc_count][seq_len][hidden]`, so the grouped norms run on contiguous + /// groups. `width` is the caller's row stride and is refused unless it is + /// the compiled one. + pub fn qwen4_ple_split_branches_cuda( + row_major: *const Half, + branch_major: *mut Half, + seq_len: i32, + width: i32, + hidden: i32, + hc_count: i32, + stream: CUstream, + ) -> CUresult; + + /// One gate scalar per (row, branch) — the per-branch dot, its double + /// rounding around the `1 / sqrt_hidden` divide, the signed sqrt and the + /// sigmoid — times `value`, broadcast over the branches that share it. + /// + /// `key_normed`, `query_normed` and `gated` are branch-major; `value` is + /// `[seq_len, hidden]`. `gate_span` is `hidden` for the reference's + /// per-branch dot and `width` for the full-stream negative control; + /// `signed_sqrt` is 1 for the reference and 0 for the unsigned control. + pub fn qwen4_ple_gate_cuda( + key_normed: *const Half, + query_normed: *const Half, + value: *const Half, + gated: *mut Half, + seq_len: i32, + hidden: i32, + hc_count: i32, + gate_span: i32, + sqrt_hidden: f32, + signed_sqrt: i32, + stream: CUstream, + ) -> CUresult; + + /// `out = gated + silu(bf16(dilated causal depthwise conv(conv_in)))`. + /// `conv_in` and `gated` are branch-major, `out` is row-major + /// `[seq_len, width]`, `conv_weight` is `[width, 1, conv_kernel]` flattened. + /// The zero initial state of `(conv_kernel - 1) * dilation` rows is a bound + /// inside the kernel, not a buffer. + pub fn qwen4_ple_conv_residual_cuda( + conv_in: *const Half, + gated: *const Half, + conv_weight: *const Half, + out: *mut Half, + seq_len: i32, + hidden: i32, + hc_count: i32, + conv_kernel: i32, + dilation: i32, + stream: CUstream, + ) -> CUresult; +} diff --git a/pegainfer-kernels/src/ops.rs b/pegainfer-kernels/src/ops.rs index 21789ecf3..3c622a05b 100644 --- a/pegainfer-kernels/src/ops.rs +++ b/pegainfer-kernels/src/ops.rs @@ -231,6 +231,14 @@ pub use norm::rms_norm_into; pub use norm::rms_norm_offset_into; pub use norm::rms_norm_rows_into; #[cfg(feature = "qwen4")] +pub use qwen4::PleBlockGeometry; +#[cfg(feature = "qwen4")] +pub use qwen4::PleBlockVariant; +#[cfg(feature = "qwen4")] +pub use qwen4::ple_block_into; +#[cfg(feature = "qwen4")] +pub use qwen4::ple_block_variant_into; +#[cfg(feature = "qwen4")] pub use qwen4::qsa_indexer_block_scores_into; #[cfg(feature = "qwen35")] pub use qwen35::GDN_AOT_KEY_HEAD_DIM; diff --git a/pegainfer-kernels/src/ops/qwen4.rs b/pegainfer-kernels/src/ops/qwen4.rs index 5f34bae8f..51a7afb8b 100644 --- a/pegainfer-kernels/src/ops/qwen4.rs +++ b/pegainfer-kernels/src/ops/qwen4.rs @@ -8,12 +8,20 @@ //! and this wrapper sizes its buffers from it. The checkpoint's own indexer //! fields are checked where the model crate reads its config, once that caller //! lands — this stage takes only `rms_eps` from the caller. +//! +//! The PLE block below is the device half of the n-gram golden's consumer: the +//! per-layer-embedding block one decoder layer injects, starting at the two +//! recorded projections. Its three grouped `(1 + w)` RMSNorms are calls to the +//! existing batched entry point rather than new kernels. Frozen semantics: +//! `csrc/qwen4/qwen4_ple_block.cu`. use anyhow::Result; +use anyhow::anyhow; use anyhow::ensure; use cudarc::driver::CudaSlice; use cudarc::driver::DevicePtr; use cudarc::driver::DevicePtrMut; +use cudarc::driver::sys::CUstream; use half::bf16; use crate::ffi; @@ -158,3 +166,546 @@ pub fn qsa_indexer_block_scores_into( .result() .map_err(|err| anyhow::anyhow!("QSA indexer block score launch failed: {err}")) } + +/// The checkpoint's PLE block geometry. Frozen with the line: `pegainfer-qwen4` +/// serves one `qwen4_exp` text geometry, and the kernels are compiled for it. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct PleBlockGeometry { + /// `hidden_size`, which the block's `ple_embed_dim` equals: one branch's + /// width, and the group a norm reduces over. + pub hidden: usize, + /// `hc_count`: hyper-connection branches in the stream. + pub hc_count: usize, + /// `ple_conv_kernel_size`: taps of the depthwise conv. + pub conv_kernel: usize, + /// `ngram_size`: the conv's dilation, so its state is `(conv_kernel - 1) * + /// dilation` rows and its receptive field one more. + pub dilation: usize, + /// `rms_norm_eps`, shared by all three `(1 + w)` norms. + pub rms_eps: f32, +} + +impl PleBlockGeometry { + /// The frozen Qwen3.8-Flash-Next text config's PLE block. + pub fn qwen38_flash_next() -> Self { + Self { + hidden: 2560, + hc_count: 4, + conv_kernel: 4, + dilation: 3, + rms_eps: 1.0e-6, + } + } + + /// Columns of the stream the block reads and writes. + pub fn width(&self) -> usize { + self.hc_count * self.hidden + } + + /// `sqrt(hidden_size)` in the f32 the reference's opmath divides by. + /// + /// Computed the way the reference does — an f64 square root narrowed to f32 — + /// and passed to the kernel rather than compiled into it, because the gate + /// *divides* by this: multiplying by its reciprocal rounds differently, and + /// the difference lands on every gate. + pub fn sqrt_hidden(&self) -> f32 { + (self.hidden as f64).sqrt() as f32 + } + + /// Rows the conv reads before the sequence starts: its zero initial state. + pub fn conv_state(&self) -> usize { + self.conv_kernel.saturating_sub(1) * self.dilation + } + + /// Scratch [`ple_block_into`] needs for `seq_len` rows, in bf16 elements. + /// + /// Six regions of `seq_len * width()`, one per intermediate: the two + /// branch-major splits, the two grouped norms over them, the gated stream + /// and the conv's normed input. Every region is a whole stream, because each + /// is written by one launch and read by the next. + pub fn workspace_len(&self, seq_len: usize) -> usize { + 6 * seq_len * self.width() + } + + /// Refuse a geometry the kernels were not compiled for. Called by + /// [`ple_block_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<()> { + ensure!( + self.hidden == 2560 && self.hc_count == 4 && self.conv_kernel == 4, + "PLE block geometry {self:?} is not the geometry \ + csrc/qwen4/qwen4_ple_block.cu was compiled for" + ); + // The dilation is a runtime argument (the state it implies is a bound + // inside the kernel, not a buffer), but only `ngram_size` is the frozen + // semantics: a different one convolves the wrong rows and still produces + // a plausible stream. + ensure!( + self.dilation == 3, + "PLE block dilation {} is not the compiled ngram_size 3", + self.dilation + ); + // Bit equality, not a tolerance: the norms' epsilon enters an rsqrt, so + // a drifted value changes every element by less than the comparison can + // see and more than the block should tolerate. + ensure!( + self.rms_eps.to_bits() == 1.0e-6f32.to_bits(), + "PLE block rms_norm_eps {} is not the compiled 1e-6", + self.rms_eps + ); + Ok(()) + } +} + +/// Which arithmetic the block computes. +/// +/// [`Reference`](PleBlockVariant::Reference) is the frozen +/// `Qwen4ExpTextPLELayer` semantics and the only variant +/// [`ple_block_into`] passes. The other three are the parity gate's negative +/// controls, on the device side: each turns exactly one knob and each must move +/// the output past the gate's bound, which is what proves the *device* path is +/// pinned and not only the host oracle that mirrors it. +/// +/// They live here rather than in a second copy of the kernels because a control +/// that runs different code from the thing it controls proves nothing: the same +/// launches, the same buffers and the same entry points run with one scalar +/// changed, so a control that passed would mean the gate cannot see that fault. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum PleBlockVariant { + /// The frozen semantics. + Reference, + /// Control: drop the sign restore after the sqrt, so a negative gate reaches + /// the sigmoid as its magnitude and inverts that branch's contribution. + UnsignedSqrtGate, + /// Control: one dot over the whole `hc_count * hidden` stream instead of one + /// per branch, so all four branches share a gate scaled by four. + FullStreamGate, + /// Control: conv dilation 1 instead of `ngram_size`, so three of the four + /// taps read rows the reference never looks at. + Dilation1Conv, +} + +impl PleBlockVariant { + /// Channels one gate's dot reduces over. + fn gate_span(self, geometry: &PleBlockGeometry) -> usize { + match self { + Self::FullStreamGate => geometry.width(), + _ => geometry.hidden, + } + } + + /// Whether the sqrt's sign is restored. + fn signed_sqrt(self) -> i32 { + i32::from(self != Self::UnsignedSqrtGate) + } + + /// The conv's dilation. + fn dilation(self, geometry: &PleBlockGeometry) -> usize { + match self { + Self::Dilation1Conv => 1, + _ => geometry.dilation, + } + } +} + +/// Run one PLE block over `seq_len` rows. +/// +/// `hidden`, `key` and `out` are row-major `[seq_len, width()]` — `hidden` is +/// the residual stream the query norm reads, `key` the projection the key norm +/// reads. `value` is `[seq_len, hidden]` and is broadcast over the branches. The +/// three norm weights are one per channel of the stream, `[width()]`, and +/// `conv_weight` is the depthwise `[width(), 1, conv_kernel]` flattened. +/// +/// `workspace` is caller-owned scratch of [`PleBlockGeometry::workspace_len`] +/// elements; see that method for the layout. Nothing is allocated here, so the +/// call is capturable. +/// +/// This is the frozen semantics; see [`ple_block_variant_into`] for the gate's +/// control surface. +pub fn ple_block_into( + ctx: &DeviceContext, + geometry: &PleBlockGeometry, + hidden: &CudaSlice, + key: &CudaSlice, + value: &CudaSlice, + norm_key: &CudaSlice, + norm_query: &CudaSlice, + norm_conv: &CudaSlice, + conv_weight: &CudaSlice, + seq_len: usize, + out: &mut CudaSlice, + workspace: &mut CudaSlice, +) -> Result<()> { + ple_block_variant_into( + ctx, + geometry, + PleBlockVariant::Reference, + hidden, + key, + value, + norm_key, + norm_query, + norm_conv, + conv_weight, + seq_len, + out, + workspace, + ) +} + +/// [`ple_block_into`] with one knob turned. Only the parity gate's negative +/// controls call this with a variant other than +/// [`Reference`](PleBlockVariant::Reference); see that type for why the knobs +/// are here and not in a second copy of the kernels. +pub fn ple_block_variant_into( + ctx: &DeviceContext, + geometry: &PleBlockGeometry, + variant: PleBlockVariant, + hidden: &CudaSlice, + key: &CudaSlice, + value: &CudaSlice, + norm_key: &CudaSlice, + norm_query: &CudaSlice, + norm_conv: &CudaSlice, + conv_weight: &CudaSlice, + seq_len: usize, + out: &mut CudaSlice, + workspace: &mut CudaSlice, +) -> Result<()> { + geometry.validate()?; + let width = geometry.width(); + let group = geometry.hidden; + let stream_elements = seq_len + .checked_mul(width) + .ok_or_else(|| anyhow!("PLE block extent overflow at {seq_len} rows"))?; + ensure!( + hidden.len() >= stream_elements && key.len() >= stream_elements, + "PLE block inputs hold {} and {} bf16, need {stream_elements} each for {seq_len} rows of \ + {width}", + hidden.len(), + key.len() + ); + ensure!( + value.len() >= seq_len.saturating_mul(group), + "PLE block value holds {} bf16, needs {} for {seq_len} rows of {group}", + value.len(), + seq_len.saturating_mul(group) + ); + ensure!( + norm_key.len() >= width && norm_query.len() >= width && norm_conv.len() >= width, + "PLE block norm weights must hold {width} bf16 each, got {}, {} and {}", + norm_key.len(), + norm_query.len(), + norm_conv.len() + ); + ensure!( + conv_weight.len() >= width.saturating_mul(geometry.conv_kernel), + "PLE block conv weight holds {} bf16, needs {} for {width} channels of {} taps", + conv_weight.len(), + width.saturating_mul(geometry.conv_kernel), + geometry.conv_kernel + ); + ensure!( + out.len() >= stream_elements, + "PLE block out holds {} bf16, needs {stream_elements}", + out.len() + ); + let needed = geometry.workspace_len(seq_len); + ensure!( + workspace.len() >= needed, + "PLE block workspace holds {} bf16, needs {needed} for {seq_len} rows", + workspace.len() + ); + if seq_len == 0 { + return Ok(()); + } + // The entry points take extents as i32; checked before any pointer is formed. + let rows = checked_i32(seq_len, "PLE block rows")?; + let width_i = checked_i32(width, "PLE block stream width")?; + let group_i = checked_i32(group, "PLE block group width")?; + let branches_i = checked_i32(geometry.hc_count, "PLE block branch count")?; + let conv_kernel_i = checked_i32(geometry.conv_kernel, "PLE block conv kernel")?; + let dilation = variant.dilation(geometry); + let dilation_i = checked_i32(dilation, "PLE block conv dilation")?; + let gate_span_i = checked_i32(variant.gate_span(geometry), "PLE block gate span")?; + + let (hidden_ptr, _hidden_guard) = hidden.device_ptr(&ctx.stream); + let (key_ptr, _key_guard) = key.device_ptr(&ctx.stream); + let (value_ptr, _value_guard) = value.device_ptr(&ctx.stream); + let (norm_key_ptr, _norm_key_guard) = norm_key.device_ptr(&ctx.stream); + let (norm_query_ptr, _norm_query_guard) = norm_query.device_ptr(&ctx.stream); + let (norm_conv_ptr, _norm_conv_guard) = norm_conv.device_ptr(&ctx.stream); + let (conv_weight_ptr, _conv_weight_guard) = conv_weight.device_ptr(&ctx.stream); + let (out_ptr, _out_guard) = out.device_ptr_mut(&ctx.stream); + let (workspace_ptr, _workspace_guard) = workspace.device_ptr_mut(&ctx.stream); + let stream = active_cu_stream(ctx); + + // Workspace regions, one whole stream each, addressed in bytes. + let element_bytes = size_of::() as u64; + let region = stream_elements as u64 * element_bytes; + let query_split = workspace_ptr; + let key_split = workspace_ptr + region; + let query_normed = workspace_ptr + 2 * region; + let key_normed = workspace_ptr + 3 * region; + let gated = workspace_ptr + 4 * region; + let conv_in = workspace_ptr + 5 * region; + + // 1. Row-major stream -> branch-major, so the grouped norms see contiguous + // groups of `hidden`. + split_branches( + hidden_ptr, + query_split, + rows, + width_i, + group_i, + branches_i, + stream, + "query", + )?; + split_branches( + key_ptr, key_split, rows, width_i, group_i, branches_i, stream, "key", + )?; + + // 2. The two input norms, one launch per branch. The weight for branch `b` + // is that branch's slice of the stream-wide vector. + let branch_stride = (seq_len * group) as u64 * element_bytes; + let weight_stride = group as u64 * element_bytes; + for branch in 0..geometry.hc_count as u64 { + grouped_norm( + key_split + branch * branch_stride, + norm_key_ptr + branch * weight_stride, + key_normed + branch * branch_stride, + group_i, + rows, + geometry.rms_eps, + stream, + ); + grouped_norm( + query_split + branch * branch_stride, + norm_query_ptr + branch * weight_stride, + query_normed + branch * branch_stride, + group_i, + rows, + geometry.rms_eps, + stream, + ); + } + + // 3. Per-branch gate dot, signed sqrt, sigmoid, and the value broadcast. + let status = unsafe { + ffi::qwen4_ple_gate_cuda( + key_normed as *const ffi::Half, + query_normed as *const ffi::Half, + value_ptr as *const ffi::Half, + gated as *mut ffi::Half, + rows, + group_i, + branches_i, + gate_span_i, + geometry.sqrt_hidden(), + variant.signed_sqrt(), + stream, + ) + }; + status + .result() + .map_err(|err| anyhow!("PLE block gate launch failed: {err}"))?; + + // 4. The conv's own grouped norm over the gated stream. + for branch in 0..geometry.hc_count as u64 { + grouped_norm( + gated + branch * branch_stride, + norm_conv_ptr + branch * weight_stride, + conv_in + branch * branch_stride, + group_i, + rows, + geometry.rms_eps, + stream, + ); + } + + // 5. Dilated causal depthwise conv, silu, and the residual add back onto the + // gated stream, landing row-major in the caller's output. + let status = unsafe { + ffi::qwen4_ple_conv_residual_cuda( + conv_in as *const ffi::Half, + gated as *const ffi::Half, + conv_weight_ptr as *const ffi::Half, + out_ptr as *mut ffi::Half, + rows, + group_i, + branches_i, + conv_kernel_i, + dilation_i, + stream, + ) + }; + status + .result() + .map_err(|err| anyhow!("PLE block conv launch failed: {err}"))?; + Ok(()) +} + +fn split_branches( + row_major: u64, + branch_major: u64, + rows: i32, + width: i32, + group: i32, + branches: i32, + stream: CUstream, + what: &str, +) -> Result<()> { + let status = unsafe { + ffi::qwen4_ple_split_branches_cuda( + row_major as *const ffi::Half, + branch_major as *mut ffi::Half, + rows, + width, + group, + branches, + stream, + ) + }; + status + .result() + .map_err(|err| anyhow!("PLE block {what} branch split failed: {err}")) +} + +/// One branch of a grouped `(1 + w)` RMSNorm, through the same entry point +/// `ops::norm::rms_norm_batch_offset_into` wraps. +/// +/// That wrapper returns no status because the C++ side is `void`, and this call +/// inherits that: a failure here surfaces as the next launch's sticky error, +/// which is why every launch in this module checks its own result. +fn grouped_norm( + x: u64, + weight: u64, + out: u64, + hidden_dim: i32, + seq_len: i32, + eps: f32, + stream: CUstream, +) { + unsafe { + ffi::rms_norm_batched_offset_cuda( + x as *const ffi::Half, + weight as *const ffi::Half, + out as *mut ffi::Half, + hidden_dim, + seq_len, + eps, + stream, + ); + } +} + +#[cfg(test)] +mod tests { + use super::PleBlockGeometry; + use super::PleBlockVariant; + + /// The geometry is the one hand-off this block has with the kernels' + /// compiled constants, so a mismatch has to be refused before a pointer is + /// formed. Device-free: no launch is involved. + #[test] + fn only_the_compiled_geometry_passes() { + let compiled = PleBlockGeometry::qwen38_flash_next(); + compiled.validate().expect("the frozen geometry"); + assert_eq!(compiled.width(), 10_240); + assert_eq!(compiled.conv_state(), 9); + // `sqrt(2560)` narrowed from f64, the divisor the gate uses. + // Bit equality: the gate divides by this, so a last-place difference is + // a different divisor and not a rounding of the same one. + assert_eq!( + compiled.sqrt_hidden().to_bits(), + ((2560f64).sqrt() as f32).to_bits() + ); + assert_eq!(compiled.workspace_len(2), 6 * 2 * 10_240); + for (field, mutated) in [ + ( + "hidden", + PleBlockGeometry { + hidden: 1280, + ..compiled + }, + ), + ( + "hc_count", + PleBlockGeometry { + hc_count: 2, + ..compiled + }, + ), + ( + "conv_kernel", + PleBlockGeometry { + conv_kernel: 2, + ..compiled + }, + ), + ( + "dilation", + PleBlockGeometry { + dilation: 1, + ..compiled + }, + ), + ( + "rms_eps", + PleBlockGeometry { + rms_eps: 1.0e-5, + ..compiled + }, + ), + ( + // One f32 ulp: the geometry's own bit-equality check, and what a + // config drifting in the last place hits. + "rms_eps one ulp", + PleBlockGeometry { + rms_eps: f32::from_bits(1.0e-6f32.to_bits() + 1), + ..compiled + }, + ), + ] { + assert!( + mutated.validate().is_err(), + "a geometry with another {field} must be refused" + ); + } + } + + /// Each control turns exactly one knob, and the reference turns none. A + /// control that changed two would not say which one the gate can see. + #[test] + fn every_control_turns_one_knob() { + let geometry = PleBlockGeometry::qwen38_flash_next(); + let reference = PleBlockVariant::Reference; + assert_eq!(reference.gate_span(&geometry), geometry.hidden); + assert_eq!(reference.signed_sqrt(), 1); + assert_eq!(reference.dilation(&geometry), geometry.dilation); + for variant in [ + PleBlockVariant::UnsignedSqrtGate, + PleBlockVariant::FullStreamGate, + PleBlockVariant::Dilation1Conv, + ] { + let knobs = [ + variant.gate_span(&geometry) != reference.gate_span(&geometry), + variant.signed_sqrt() != reference.signed_sqrt(), + variant.dilation(&geometry) != reference.dilation(&geometry), + ]; + assert_eq!( + knobs.iter().filter(|turned| **turned).count(), + 1, + "{variant:?} must differ from the reference in exactly one knob, got {knobs:?}" + ); + } + assert_eq!( + PleBlockVariant::FullStreamGate.gate_span(&geometry), + geometry.width() + ); + assert_eq!(PleBlockVariant::UnsignedSqrtGate.signed_sqrt(), 0); + assert_eq!(PleBlockVariant::Dilation1Conv.dilation(&geometry), 1); + } +} diff --git a/pegainfer-kernels/tests/qwen4_ple_block_parity.rs b/pegainfer-kernels/tests/qwen4_ple_block_parity.rs new file mode 100644 index 000000000..e716d3ac1 --- /dev/null +++ b/pegainfer-kernels/tests/qwen4_ple_block_parity.rs @@ -0,0 +1,1055 @@ +//! Qwen3.8-Flash-Next PLE block against the frozen reference golden. +//! +//! `test_data/qwen38fn-ple-golden.safetensors` was dumped from transformers +//! 5.16.1 `qwen4_exp` by `tools/accuracy/dump_qwen4_ple_golden.py`. Each case +//! carries the block's inputs *after* the two projections the 95 GiB n-gram +//! table feeds — `hidden`, `key` and `value` — plus the reference's own block +//! output `out`, so this gate replays those inputs through the ops layer and +//! compares. The projections and the table gather are not covered here; the +//! golden starts after them. +//! +//! This is the device half of a two-half consumer. The host half, +//! `pegainfer-qwen4/tests/ple_ngram_replay.rs`, runs with no CUDA at all and +//! replays the same stored `out` **bit-identically** in f32 on the CPU +//! (81920/81920 and 122880/122880 elements exact), which is what makes the +//! bound below tight enough to mean something: every rounding point the +//! reference has is known and reproducible, so the only thing left for a device +//! to be wrong about is a knob, not a tolerance. The device path turns out to be +//! nearly as exact — 204799 of the 204800 stored outputs come back bit-for-bit, +//! and the one that does not is a single bf16 ulp — see [`OUT_MAX_ABS`]. +//! +//! Provenance is asserted on every load, never skipped: the manifest's +//! `format_version`, `transformers_version` and `config`/`config_sha256` are +//! held against this build's frozen config, and that config's PLE geometry +//! against `PleBlockGeometry::qwen38_flash_next()`. A fixture from another +//! checkpoint or another semantics fails here instead of being skipped. +//! +//! Without a device the comparisons skip; `PEGAINFER_REQUIRE_GPU=1` turns that +//! into a failure. The two device-free halves — provenance and the entry points' +//! refusal contract — never skip. +//! +//! cargo test --release -p pegainfer-kernels --features qwen4 \ +//! --test qwen4_ple_block_parity -- --nocapture +//! +//! **What this fixture cannot reach**, recorded rather than papered over: +//! +//! - The gate's `clamp_min(1e-6)` floor and its `sign(0)` branch are both inert +//! here: the smallest `|pre_gate|` either case produces is 0.0232, five orders +//! above the floor, so no gate is clamped and none is exactly zero, and a wrong +//! floor constant would pass. Both are transcribed from the reference's own +//! `abs().clamp_min(1e-6).sqrt() * sign()`, and the host oracle in +//! `pegainfer-qwen4` carries the same two branches with the same non-exercise. +//! - Row counts are 8 and 12, the fixture's own. `seq_len` 0 is pinned only as a +//! guard (the zero-extent positive controls below, which form no launch), and +//! `seq_len` 1 — where every conv tap but the last reads the zero state — is +//! compared against nothing, because the golden records no such case and the +//! fixture is committed. +//! - The ops layer's own extent guards (`hidden.len()`, `workspace.len()` and the +//! rest) are not exercised: they run before any pointer is formed but after a +//! `CudaSlice` exists, so reaching one needs a device and a deliberately short +//! buffer. The entry points' guards — the ones that stand between a caller and +//! an out-of-bounds read — are covered below. + +#![cfg(feature = "qwen4")] + +mod common; + +use cudarc::driver::CudaSlice; +use half::bf16; +use pegainfer_kernels::ops::PleBlockGeometry; +use pegainfer_kernels::ops::PleBlockVariant; +use pegainfer_kernels::ops::ple_block_into; +use pegainfer_kernels::ops::ple_block_variant_into; +use pegainfer_kernels::tensor::DeviceContext; +use safetensors::Dtype; +use safetensors::SafeTensors; +use sha2::Digest; +use sha2::Sha256; + +const FIXTURE: &str = "qwen38fn-ple-golden.safetensors"; +const METADATA_KEY: &str = "qwen4_ple_golden"; +/// Frozen config the golden was dumped from, relative to the workspace root. +const CONFIG: &str = "pegainfer-qwen4/tests/frozen_config.json"; +const FORMAT_VERSION: u64 = 2; +/// The reference implementation the golden's semantics are frozen against. +const REF_TRANSFORMERS: &str = "5.16.1"; +const CASES: usize = 2; + +/// Block-output parity bound, absolute, on a stream whose largest stored element +/// is 0.8438. +/// +/// **Measured on one sm_120 card (RTX 5060 Laptop, CUDA 12.9.1): case 0 is +/// bit-identical, 81920/81920, and case 1 differs on exactly one element of +/// 122880, by 1.221e-4 — one bf16 ulp at that element's magnitude (0.024292 +/// against 0.024414), 0.0076 of this bound.** Over both cases, 204799 of 204800 +/// outputs are bit-for-bit the reference's. +/// +/// The host replay of this same block is bit-identical too, so the bound is not +/// covering a semantics gap. It covers the three places the device path cannot +/// be the host path, all of them the same size: +/// +/// 1. **`rsqrt.approx.ftz.f32`.** The three grouped norms are the existing +/// batched `(1 + w)` RMSNorm, whose reciprocal comes from the PTX +/// approximation instruction, where the reference and the host replay both use +/// a correctly rounded `1 / sqrt`. Its ~1.4e-7 relative error against a bf16 +/// step of ~3.9e-3 flips the bf16 rounding of roughly one normed element in +/// 30000; an instrumented run of this same call during bring-up found exactly +/// two such elements across case 0's three norms (245760 outputs), each one +/// bf16 ulp of itself, and the gated stream bit-identical. +/// 2. **A different reduction tree.** The 2560-wide sum of squares and the +/// 2560-wide gate dot are f32 butterflies here, an f32 tree in the reference +/// and an f64 exact sum in the host replay. All three land within an f32 ulp +/// or two of each other, four orders below the bf16 step the result is +/// rounded to, so they flip the same ~1-in-10^4..10^5 of elements. +/// 3. **`expf`**, the one op the host gate already budgets for: its allowance is +/// a whole bf16 ulp injected into *every* `sigmoid` and *every* `silu`, which +/// moves `out` by at most 7.8e-3. +/// +/// One flipped bf16 anywhere upstream reaches `out` as at most about one bf16 ulp +/// of `out` at that element: the gate is a single scalar per (row, branch), so a +/// flipped `sigmoid` moves that branch's 2560 channels by a common relative +/// step; the conv, its silu (derivative bounded by ~1.1) and the residual add +/// cannot amplify it further. One ulp at the largest stored element is 2^-8 = +/// 3.9e-3, and the bound is four of those — which is also exactly the host +/// gate's bound, so the two halves of this consumer share one number and can be +/// read against each other. +/// +/// It cannot be tightened towards the measured 1.2e-4, and the reason is not "it +/// currently passes": `rsqrt.approx`, `expf` and the reduction order are +/// architecture- and libm-defined, so another GPU flips a *different* set of +/// elements, possibly including one near 0.8438 where a single ulp is 3.9e-3. +/// The flip count is a discrete random variable, not a constant. Nor does it need +/// tightening: every structural fault this block can have — a collapsed rounding +/// point, the wrong group, the wrong divisor, the wrong taps — moves elements by +/// orders of this bound, which is what the three controls below measure at 17x, +/// 21x and 49x outside it. +const OUT_MAX_ABS: f32 = 1.6e-2; + +/// Written over the output before every launch, so a kernel that wrote part of +/// the stream cannot pass as one that wrote all of it. +const SENTINEL: bf16 = bf16::from_bits(0x7fc0); + +fn fixture_path() -> String { + format!("{}/../test_data/{FIXTURE}", env!("CARGO_MANIFEST_DIR")) +} + +fn config_path() -> String { + format!("{}/../{CONFIG}", env!("CARGO_MANIFEST_DIR")) +} + +/// Bit equality for the scalars the kernels pin. They refuse any other value, so +/// a tolerance here would let a config drift through. +fn same_f32(a: f32, b: f32) -> bool { + a.to_bits() == b.to_bits() +} + +fn sha256_hex(bytes: &[u8]) -> String { + use std::fmt::Write as _; + Sha256::digest(bytes) + .iter() + .fold(String::with_capacity(64), |mut hex, byte| { + let _ = write!(hex, "{byte:02x}"); + hex + }) +} + +/// The fixture's own record of what it was dumped from, held against the +/// reference version, the frozen config's bytes and the PLE geometry this crate +/// compiles for. Panics rather than skips: provenance is the whole reason a +/// parity number means anything. +fn assert_provenance(bytes: &[u8]) -> serde_json::Value { + let (_, metadata) = SafeTensors::read_metadata(bytes).expect("fixture metadata"); + let recorded = metadata + .metadata() + .as_ref() + .expect("fixture metadata map") + .get(METADATA_KEY) + .expect("fixture manifest key") + .clone(); + let manifest: serde_json::Value = + serde_json::from_str(&recorded).expect("fixture manifest is not JSON"); + + assert_eq!( + manifest["format_version"].as_u64(), + Some(FORMAT_VERSION), + "fixture format_version; this gate reads version {FORMAT_VERSION}" + ); + assert_eq!( + manifest["transformers_version"].as_str(), + Some(REF_TRANSFORMERS), + "fixture was dumped from another reference version, and the frozen \ + semantics are that one's" + ); + assert_eq!( + manifest["config"].as_str(), + Some(CONFIG), + "fixture was dumped from another config" + ); + + let config_bytes = + std::fs::read(config_path()).unwrap_or_else(|e| panic!("read {}: {e}", config_path())); + let expected = sha256_hex(&config_bytes); + assert_eq!( + manifest["config_sha256"].as_str(), + Some(expected.as_str()), + "{CONFIG} is not the config the fixture was dumped from" + ); + + let config: serde_json::Value = + serde_json::from_slice(&config_bytes).expect("frozen config is not JSON"); + let text = &config["text_config"]; + let geometry = PleBlockGeometry::qwen38_flash_next(); + let field = |name: &str| -> usize { + text[name] + .as_u64() + .unwrap_or_else(|| panic!("{name} must be an integer")) as usize + }; + assert_eq!(field("hidden_size"), geometry.hidden); + // The block's gate divides by `sqrt(hidden_size)`, and its value stream is + // `ple_embed_dim` wide; the reference requires the two to be the same width. + assert_eq!(field("ple_embed_dim"), geometry.hidden); + assert_eq!(field("hc_count"), geometry.hc_count); + assert_eq!(field("ple_conv_kernel_size"), geometry.conv_kernel); + // The conv's dilation is `ngram_size`, which is also the widest n-gram. + assert_eq!(field("ngram_size"), geometry.dilation); + assert_eq!(geometry.conv_state(), 9); + let eps = text["rms_norm_eps"] + .as_f64() + .expect("text_config.rms_norm_eps must be a number"); + assert!( + same_f32(eps as f32, geometry.rms_eps), + "rms_norm_eps {eps} is not the compiled {}", + geometry.rms_eps + ); + let ple_layers = text["ple_layer_ids"].as_array().expect("ple_layer_ids"); + assert_eq!( + ple_layers.len(), + 1, + "the reference offsets each PLE layer's prime run by its position, and \ + this crate derives one table" + ); + geometry.validate().expect("the frozen geometry"); + // Recorded, not asserted: the host the golden was dumped on decides the + // order its float32 reductions round in, which is where the host replay's + // bit-identity was measured and not something this gate can reproduce. + println!( + "golden provenance: transformers {}, torch {}, {}, seed {}", + manifest["transformers_version"].as_str().unwrap_or("?"), + manifest["torch_version"].as_str().unwrap_or("?"), + manifest["platform"].as_str().unwrap_or("?"), + manifest["seed"].as_u64().unwrap_or(0), + ); + manifest +} + +fn i32_tensor(fixture: &SafeTensors<'_>, name: &str) -> (Vec, Vec) { + let view = fixture + .tensor(name) + .unwrap_or_else(|e| panic!("{name}: {e}")); + assert_eq!(view.dtype(), Dtype::I32, "{name} must be i32"); + let host = view + .data() + .as_chunks::<4>() + .0 + .iter() + .map(|b| i32::from_le_bytes(*b)) + .collect(); + (view.shape().to_vec(), host) +} + +fn bf16_tensor(fixture: &SafeTensors<'_>, name: &str) -> (Vec, Vec) { + let view = fixture + .tensor(name) + .unwrap_or_else(|e| panic!("{name}: {e}")); + assert_eq!(view.dtype(), Dtype::BF16, "{name} must be bf16"); + let shape = view.shape().to_vec(); + let host = view + .data() + .as_chunks::<2>() + .0 + .iter() + .map(|b| bf16::from_bits(u16::from_le_bytes(*b))) + .collect(); + (shape, host) +} + +/// One recorded sequence: the block's inputs and the reference's output. +struct Case { + index: usize, + seq_len: usize, + hidden: Vec, + key: Vec, + value: Vec, + reference: Vec, +} + +fn load_case(fixture: &SafeTensors<'_>, index: usize, geometry: &PleBlockGeometry) -> Case { + let width = geometry.width(); + let name = format!("case{index}"); + + let (seq_shape, seq_cell) = i32_tensor(fixture, &format!("{name}_seq_len")); + assert_eq!(seq_shape, vec![1], "{name}_seq_len must be one cell"); + let seq_len = usize::try_from(seq_cell[0]) + .ok() + .filter(|&len| len > 0) + .unwrap_or_else(|| panic!("{name}_seq_len must be a positive row count")); + + // The recording is whole-sequence, so the row selection has to be the + // identity: the conv reads nine rows before the sequence with a zero state, + // and a dropped row would need that row's conv and window state to stay + // correct. Asserted rather than assumed, because a gate that quietly + // compared a prefix would still report a number. + let (rows_shape, rows) = i32_tensor(fixture, &format!("{name}_rows")); + assert_eq!(rows_shape, vec![seq_len], "{name}_rows must be one per row"); + assert!( + rows.iter().enumerate().all(|(i, &row)| row == i as i32), + "{name}_rows must be the identity 0..{seq_len}: the block is replayed \ + whole-sequence and the fixture records no cache state" + ); + + let (hidden_shape, hidden) = bf16_tensor(fixture, &format!("{name}_hidden")); + let (key_shape, key) = bf16_tensor(fixture, &format!("{name}_key")); + let (value_shape, value) = bf16_tensor(fixture, &format!("{name}_value")); + let (out_shape, reference) = bf16_tensor(fixture, &format!("{name}_out")); + for (what, shape) in [ + ("hidden", hidden_shape), + ("key", key_shape), + ("out", out_shape), + ] { + assert_eq!(shape, vec![seq_len, width], "{name}_{what} shape"); + } + assert_eq!( + value_shape, + vec![seq_len, geometry.hidden], + "{name}_value shape" + ); + Case { + index, + seq_len, + hidden, + key, + value, + reference, + } +} + +/// One case's parity: how much was compared, the worst element and where, the +/// shape of the residual, and how much of the bound it used. +/// +/// The identical count and the two thresholds are reported rather than asserted: +/// a handful of bf16-rounding flips and a structural fault are different animals +/// and these counts say which one this is. A systematic error shows up as a +/// large identical-count shortfall *and* a large mean, where rounding noise +/// shows up as a small mean with a max at one or two ulps. +struct Summary { + compared: usize, + max_abs: f32, + mean_abs: f32, + identical: usize, + /// `(row, channel, got, want)` + worst: (usize, usize, f32, f32), + over_1e3: usize, + over_1e4: usize, + bound_fraction: f32, +} + +/// Run one case through the ops layer and compare against the stored output. +fn check_case( + ctx: &DeviceContext, + geometry: &PleBlockGeometry, + variant: PleBlockVariant, + weights: &Weights, + case: &Case, +) -> Summary { + let width = geometry.width(); + let elements = case.seq_len * width; + + let hidden = ctx.stream.clone_htod(&case.hidden).expect("hidden upload"); + let key = ctx.stream.clone_htod(&case.key).expect("key upload"); + let value = ctx.stream.clone_htod(&case.value).expect("value upload"); + // Pre-filled with a NaN sentinel: every element of `out` is written by the + // conv pass, so anything left holding the sentinel is an unwritten row. + let mut out = ctx + .stream + .clone_htod(&vec![SENTINEL; elements]) + .expect("out buffer"); + let mut workspace = ctx + .stream + .alloc_zeros::(geometry.workspace_len(case.seq_len)) + .expect("PLE workspace"); + + // The reference variant goes through the production entry point, which is + // the one a serving path calls; the controls go through the variant entry + // it delegates to. + let result = match variant { + PleBlockVariant::Reference => ple_block_into( + ctx, + geometry, + &hidden, + &key, + &value, + &weights.norm_key, + &weights.norm_query, + &weights.norm_conv, + &weights.conv1d, + case.seq_len, + &mut out, + &mut workspace, + ), + control => ple_block_variant_into( + ctx, + geometry, + control, + &hidden, + &key, + &value, + &weights.norm_key, + &weights.norm_query, + &weights.norm_conv, + &weights.conv1d, + case.seq_len, + &mut out, + &mut workspace, + ), + }; + result.expect("PLE block"); + ctx.sync().expect("sync"); + let got = ctx.stream.clone_dtoh(&out).expect("out download"); + + assert_eq!( + got.len(), + case.reference.len(), + "case {}: length", + case.index + ); + let mut max_abs = 0.0f32; + let mut sum_abs = 0.0f64; + let mut identical = 0usize; + let mut over_1e3 = 0usize; + let mut over_1e4 = 0usize; + let mut worst = (0usize, 0usize, 0.0f32, 0.0f32); + for (i, (x, want)) in got.iter().zip(&case.reference).enumerate() { + let g = x.to_f32(); + let w = want.to_f32(); + // Finiteness first: the sentinel is a NaN, so a NaN the kernels + // *produced* would otherwise be reported as an unwritten element. + assert!( + g.is_finite(), + "case {}: element {i} (row {}, channel {}) came back {}", + case.index, + i / width, + i % width, + if x.to_bits() == SENTINEL.to_bits() { + "still holding the sentinel, so the conv pass did not write the whole stream" + } else { + "non-finite" + } + ); + assert!( + w.is_finite(), + "case {}: the golden's element {i} is not finite", + case.index + ); + let delta = (g - w).abs(); + sum_abs += f64::from(delta); + identical += usize::from(x.to_bits() == want.to_bits()); + over_1e3 += usize::from(delta > 1.0e-3); + over_1e4 += usize::from(delta > 1.0e-4); + if delta > max_abs { + max_abs = delta; + worst = (i / width, i % width, g, w); + } + } + Summary { + compared: got.len(), + max_abs, + mean_abs: (sum_abs / got.len() as f64) as f32, + identical, + worst, + over_1e3, + over_1e4, + bound_fraction: max_abs / OUT_MAX_ABS, + } +} + +fn print_summary(label: &str, geometry: &PleBlockGeometry, case: &Case, summary: &Summary) { + println!( + "{label} case {}: S={} stream={}x{} out max_abs={:.3e} at row {} channel {} \ + ({:.6} against {:.6}), mean_abs={:.3e}, bit_identical={}/{}, above 1e-3 {}/{}, \ + above 1e-4 {}/{}, bound fraction {:.4} of {OUT_MAX_ABS:.1e}", + case.index, + case.seq_len, + geometry.hc_count, + geometry.hidden, + summary.max_abs, + summary.worst.0, + summary.worst.1, + summary.worst.2, + summary.worst.3, + summary.mean_abs, + summary.identical, + summary.compared, + summary.over_1e3, + summary.compared, + summary.over_1e4, + summary.compared, + summary.bound_fraction, + ); +} + +/// The block's own weights, uploaded once and shared by every case: the three +/// stream-wide `(1 + w)` norm weights and the depthwise conv. +struct Weights { + norm_key: CudaSlice, + norm_query: CudaSlice, + norm_conv: CudaSlice, + conv1d: CudaSlice, +} + +fn load_weights( + ctx: &DeviceContext, + fixture: &SafeTensors<'_>, + geometry: &PleBlockGeometry, +) -> Weights { + let width = geometry.width(); + // The conv weight is stored rank-3 (`[width, 1, conv_kernel]`, the + // reference's own depthwise shape) and read flat, one tap run per channel. + let upload = |name: &str, shape: &[usize]| -> CudaSlice { + let (stored, host) = bf16_tensor(fixture, name); + assert_eq!(&stored, shape, "{name} must be {shape:?}"); + ctx.stream + .clone_htod(&host) + .unwrap_or_else(|e| panic!("{name} upload: {e}")) + }; + Weights { + norm_key: upload("norm_key", &[width]), + norm_query: upload("norm_query", &[width]), + norm_conv: upload("norm_conv", &[width]), + conv1d: upload("conv1d_weight", &[width, 1, geometry.conv_kernel]), + } +} + +/// The fixture's provenance and its geometry against this crate's compiled +/// constants. Device-free, and never behind `device_or_skip`: a fixture from +/// another reference would make every number below meaningless. +#[test] +fn fixture_records_the_frozen_reference() { + let bytes = std::fs::read(fixture_path()).expect("read fixture"); + assert_provenance(&bytes); + let fixture = SafeTensors::deserialize(&bytes).expect("parse fixture"); + let geometry = PleBlockGeometry::qwen38_flash_next(); + let cases: Vec = (0..CASES) + .map(|index| load_case(&fixture, index, &geometry)) + .collect(); + let lengths: Vec = cases.iter().map(|case| case.seq_len).collect(); + // The two cases the dump records, asserted rather than assumed. The block + // itself is EOS-blind — the hash is not, and the device-free gate covers + // that — but the row counts are what make both ends of the conv's warm-up + // exercised: a case no longer than the nine-row state has rows whose taps are + // all padding, and a case longer than it has rows whose four taps are all + // real. A fixture that recorded only one of the two would still produce a + // parity number. + assert_eq!(lengths, vec![8, 12], "the recorded case lengths"); + assert!( + lengths.iter().any(|&len| len <= geometry.conv_state()), + "a case must be short enough that some rows read the whole zero state, \ + or the conv's warm-up is untested" + ); + assert!( + lengths.iter().any(|&len| len > geometry.conv_state()), + "a case must be long enough that some rows read four real taps" + ); +} + +/// The kernels' own half of the refusal contract, likewise device-free: every +/// entry point rejects a geometry it was not compiled for *before* it forms a +/// launch, so this can be pinned without a device and must never be skippable. +/// +/// The two zero-extent calls per entry point come first as a **positive control** +/// — the compiled geometry with nothing to do, which the guards must let through +/// and which forms no launch. Without them an entry point that refused +/// everything would make every refusal below vacuous. +#[test] +fn the_entry_points_refuse_a_geometry_they_were_not_compiled_for() { + use cudarc::driver::sys::CUresult; + use cudarc::driver::sys::CUstream; + use pegainfer_kernels::ffi::Half; + use pegainfer_kernels::ffi::qwen4_ple_conv_residual_cuda as conv; + use pegainfer_kernels::ffi::qwen4_ple_gate_cuda as gate; + use pegainfer_kernels::ffi::qwen4_ple_split_branches_cuda as split; + + // Dangling and well-aligned, never dereferenced: every call below leaves + // through a guard or through the zero-extent early return, and a launch is + // the only thing that would touch a pointer. + let half: *const Half = std::ptr::NonNull::::dangling().as_ptr(); + // Any positive finite divisor passes the guard; the compiled geometry's own + // `sqrt(2560)` keeps the positive controls honest about what they let through. + let sqrt_hidden = PleBlockGeometry::qwen38_flash_next().sqrt_hidden(); + let invalid = CUresult::CUDA_ERROR_INVALID_VALUE as i32; + let success = CUresult::CUDA_SUCCESS as i32; + let null_stream = CUstream::default(); + + // Positive controls: the compiled geometry with no rows to do. + assert_eq!( + (unsafe { split(half, half.cast_mut(), 0, 10240, 2560, 4, null_stream) }) as i32, + success, + "the compiled split geometry must pass the guards" + ); + assert_eq!( + (unsafe { + gate( + half, + half, + half, + half.cast_mut(), + 0, + 2560, + 4, + 2560, + sqrt_hidden, + 1, + null_stream, + ) + }) as i32, + success, + "the compiled gate geometry must pass the guards" + ); + assert_eq!( + (unsafe { + conv( + half, + half, + half, + half.cast_mut(), + 0, + 2560, + 4, + 4, + 3, + null_stream, + ) + }) as i32, + success, + "the compiled conv geometry must pass the guards" + ); + + for (what, result) in [ + ( + "a split over a stream that is not the compiled 10240 wide", + (unsafe { split(half, half.cast_mut(), 0, 10239, 2560, 4, null_stream) }) as i32, + ), + ( + "a split over a group that is not the compiled 2560 wide", + (unsafe { split(half, half.cast_mut(), 0, 10240, 1280, 4, null_stream) }) as i32, + ), + ( + "a split over a branch count that is not the compiled 4", + (unsafe { split(half, half.cast_mut(), 0, 10240, 2560, 2, null_stream) }) as i32, + ), + ( + "a negative row count", + (unsafe { split(half, half.cast_mut(), -1, 10240, 2560, 4, null_stream) }) as i32, + ), + ( + "a gate over a group that is not the compiled 2560 wide", + (unsafe { + gate( + half, + half, + half, + half.cast_mut(), + 0, + 1280, + 4, + 1280, + sqrt_hidden, + 1, + null_stream, + ) + }) as i32, + ), + ( + "a gate over a branch count that is not the compiled 4", + (unsafe { + gate( + half, + half, + half, + half.cast_mut(), + 0, + 2560, + 2, + 2560, + sqrt_hidden, + 1, + null_stream, + ) + }) as i32, + ), + ( + "a gate whose span is neither the group nor the stream", + (unsafe { + gate( + half, + half, + half, + half.cast_mut(), + 0, + 2560, + 4, + 5120, + sqrt_hidden, + 1, + null_stream, + ) + }) as i32, + ), + ( + "a gate whose sign restore is neither 0 nor 1", + (unsafe { + gate( + half, + half, + half, + half.cast_mut(), + 0, + 2560, + 4, + 2560, + sqrt_hidden, + 2, + null_stream, + ) + }) as i32, + ), + ( + // A zero divisor would put an infinity in every gate. + "a gate with a zero sqrt_hidden", + (unsafe { + gate( + half, + half, + half, + half.cast_mut(), + 0, + 2560, + 4, + 2560, + 0.0, + 1, + null_stream, + ) + }) as i32, + ), + ( + // `!(x > 0)` refuses NaN, which would otherwise poison every gate + // silently rather than fail loudly. + "a gate with a NaN sqrt_hidden", + (unsafe { + gate( + half, + half, + half, + half.cast_mut(), + 0, + 2560, + 4, + 2560, + f32::NAN, + 1, + null_stream, + ) + }) as i32, + ), + ( + "a conv over a group that is not the compiled 2560 wide", + (unsafe { + conv( + half, + half, + half, + half.cast_mut(), + 0, + 1280, + 4, + 4, + 3, + null_stream, + ) + }) as i32, + ), + ( + "a conv over a branch count that is not the compiled 4", + (unsafe { + conv( + half, + half, + half, + half.cast_mut(), + 0, + 2560, + 2, + 4, + 3, + null_stream, + ) + }) as i32, + ), + ( + "a conv over a kernel that is not the compiled 4 taps", + (unsafe { + conv( + half, + half, + half, + half.cast_mut(), + 0, + 2560, + 4, + 2, + 3, + null_stream, + ) + }) as i32, + ), + ( + "a conv with a zero dilation", + (unsafe { + conv( + half, + half, + half, + half.cast_mut(), + 0, + 2560, + 4, + 4, + 0, + null_stream, + ) + }) as i32, + ), + ( + "a conv with a negative dilation", + (unsafe { + conv( + half, + half, + half, + half.cast_mut(), + 0, + 2560, + 4, + 4, + -3, + null_stream, + ) + }) as i32, + ), + ] { + assert_eq!(result, invalid, "{what} must be refused"); + } + + // One null per kind of argument, so the guards are pinned to cover the + // pointers and not only the extents. + for (what, result) in [ + ( + "a null split source", + (unsafe { + split( + std::ptr::null(), + half.cast_mut(), + 0, + 10240, + 2560, + 4, + null_stream, + ) + }) as i32, + ), + ( + "a null split destination", + (unsafe { split(half, std::ptr::null_mut(), 0, 10240, 2560, 4, null_stream) }) as i32, + ), + ( + "a null gate value", + (unsafe { + gate( + half, + half, + std::ptr::null(), + half.cast_mut(), + 0, + 2560, + 4, + 2560, + sqrt_hidden, + 1, + null_stream, + ) + }) as i32, + ), + ( + "a null gated output", + (unsafe { + gate( + half, + half, + half, + std::ptr::null_mut(), + 0, + 2560, + 4, + 2560, + sqrt_hidden, + 1, + null_stream, + ) + }) as i32, + ), + ( + "a null conv weight", + (unsafe { + conv( + half, + half, + std::ptr::null(), + half.cast_mut(), + 0, + 2560, + 4, + 4, + 3, + null_stream, + ) + }) as i32, + ), + ( + "a null conv output", + (unsafe { + conv( + half, + half, + half, + std::ptr::null_mut(), + 0, + 2560, + 4, + 4, + 3, + null_stream, + ) + }) as i32, + ), + ] { + assert_eq!(result, invalid, "{what} must be refused"); + } +} + +#[test] +fn ple_block_matches_the_reference() { + let Some(ctx) = common::device_or_skip() else { + return; + }; + let bytes = std::fs::read(fixture_path()).expect("read fixture"); + assert_provenance(&bytes); + let fixture = SafeTensors::deserialize(&bytes).expect("parse fixture"); + let geometry = PleBlockGeometry::qwen38_flash_next(); + let weights = load_weights(&ctx, &fixture, &geometry); + + let mut total = 0usize; + let mut worst_abs = 0.0f32; + let mut worst_fraction = 0.0f32; + let mut identical = 0usize; + for index in 0..CASES { + let case = load_case(&fixture, index, &geometry); + let summary = check_case(&ctx, &geometry, PleBlockVariant::Reference, &weights, &case); + print_summary("reference", &geometry, &case, &summary); + total += summary.compared; + identical += summary.identical; + worst_abs = worst_abs.max(summary.max_abs); + worst_fraction = worst_fraction.max(summary.bound_fraction); + assert!( + summary.max_abs <= OUT_MAX_ABS, + "case {}: the PLE block deviates by {:.3e} > {OUT_MAX_ABS:.3e} at row {} channel {} \ + ({:.6} against {:.6}); see OUT_MAX_ABS for how the bound is derived", + index, + summary.max_abs, + summary.worst.0, + summary.worst.1, + summary.worst.2, + summary.worst.3, + ); + } + println!( + "parity over {CASES} cases and {total} block outputs: max abs {worst_abs:.3e}, worst \ + bound fraction {worst_fraction:.4} of {OUT_MAX_ABS:.1e}, {identical} of {total} \ + bit-identical to the golden" + ); + assert!( + worst_fraction <= 1.0, + "the worst block output is {worst_fraction} of the bound" + ); +} + +/// Control (a): the gate is `sign(g) * sqrt(|g|)`. Dropping the sign restore +/// pushes every negative gate through the sigmoid as its magnitude, which +/// inverts that branch's contribution to the stream. +#[test] +fn an_unsigned_sqrt_breaks_the_block_output() { + assert_control_fails( + PleBlockVariant::UnsignedSqrtGate, + "unsigned sqrt on the gate", + ); +} + +/// Control (b): the conv is dilated by `ngram_size` (3), with a `(kernel - 1) * +/// dilation` = 9 row zero state. Dilation 1 reads three of its four taps from +/// rows the reference never looks at. +#[test] +fn dilation_one_breaks_the_block_output() { + assert_control_fails(PleBlockVariant::Dilation1Conv, "conv dilation 1"); +} + +/// Control (c): the gate dot is per BRANCH — 2560 channels, one gate per +/// branch, all four branches scaling the same `value` row. One dot over the +/// whole 10240-wide stream instead gives every branch the same, larger gate. +#[test] +fn a_full_stream_gate_breaks_the_block_output() { + assert_control_fails(PleBlockVariant::FullStreamGate, "a full-stream gate dot"); +} + +/// A negative control that passes is worse than no control: it means the gate +/// cannot see the fault it claims to pin. Every case must leave the bound, not +/// just one, or a knob that only matters for a long sequence would look pinned +/// by a short one. +fn assert_control_fails(variant: PleBlockVariant, label: &str) { + let Some(ctx) = common::device_or_skip() else { + return; + }; + let bytes = std::fs::read(fixture_path()).expect("read fixture"); + assert_provenance(&bytes); + let fixture = SafeTensors::deserialize(&bytes).expect("parse fixture"); + let geometry = PleBlockGeometry::qwen38_flash_next(); + let weights = load_weights(&ctx, &fixture, &geometry); + + let mut worst = 0.0f32; + for index in 0..CASES { + let case = load_case(&fixture, index, &geometry); + let summary = check_case(&ctx, &geometry, variant, &weights, &case); + print_summary(label, &geometry, &case, &summary); + worst = worst.max(summary.max_abs); + assert!( + summary.max_abs > OUT_MAX_ABS, + "case {}: {label} stayed within {:.3e} of the golden ({:.4} of the bound), so the \ + device path does not pin the knob this control turns", + index, + summary.max_abs, + summary.bound_fraction + ); + } + println!( + "{label}: worst deviation {worst:.3e} = {:.0}x the {OUT_MAX_ABS:.1e} bound", + worst / OUT_MAX_ABS + ); +} diff --git a/pegainfer-qwen4/Cargo.toml b/pegainfer-qwen4/Cargo.toml index 28b09b82d..ffccc5903 100644 --- a/pegainfer-qwen4/Cargo.toml +++ b/pegainfer-qwen4/Cargo.toml @@ -13,6 +13,7 @@ serde = { workspace = true } serde_json = { workspace = true } [dev-dependencies] +half = { workspace = true } safetensors = { workspace = true } sha2 = { workspace = true } diff --git a/pegainfer-qwen4/src/lib.rs b/pegainfer-qwen4/src/lib.rs index 9ccfbb712..e6fa648e7 100644 --- a/pegainfer-qwen4/src/lib.rs +++ b/pegainfer-qwen4/src/lib.rs @@ -14,6 +14,9 @@ mod config; #[cfg(test)] mod manifest; pub mod model_line; +// The PLE block's n-gram addressing: pure integer host math, so it belongs to +// this crate's CUDA-free half and the device gather consumes what it produces. +pub mod ngram; mod probe; pub use probe::MODEL_TYPE; diff --git a/pegainfer-qwen4/src/ngram.rs b/pegainfer-qwen4/src/ngram.rs new file mode 100644 index 000000000..8bb8bfb05 --- /dev/null +++ b/pegainfer-qwen4/src/ngram.rs @@ -0,0 +1,359 @@ +//! Host-side n-gram addressing for the per-layer-embedding (PLE) block. +//! +//! The PLE block injects a hashed n-gram embedding into one decoder layer. Its +//! first half is pure integer arithmetic: turn a token sequence into one table +//! row per (position, head). That half lives here, on the host, in a crate with +//! no CUDA dependency — the second half, the 95 GiB table gather those rows +//! address, is a device concern and consumes what this module produces. +//! +//! The semantics are `Qwen4ExpTextNGramEmbedding.forward`: +//! +//! 1. **An EOS-padded history.** A whole-sequence forward has no cache state, so +//! the history is [`NgramHasher::context`] copies of `eos_token_id` followed +//! by the sequence. The padding is what makes the first rows' windows legal +//! rather than truncated, and it is why the reference drops those warm-up +//! rows from its output. +//! 2. **An EOS-resetting context window** (`_shift_right_ignore_eos`). For a +//! window `shift` positions wide, position `p` reads `history[p - shift]` — +//! unless that read would reach across an EOS boundary or off the front of +//! the history, in which case it reads `eos_token_id` instead. The boundary +//! is the cummax over EOS positions: an n-gram never spans two segments. +//! 3. **A mul-xor fold per (row, order).** Each window position's token is +//! multiplied by that position's `layer_multipliers` entry and the products +//! are XOR-folded. One fold per order, so an order-`k` n-gram folds the +//! `k` widest windows down to the narrowest. +//! 4. **A per-head prime remainder plus a per-head offset.** Each head reduces +//! its order's fold modulo its own prime vocabulary size and adds its own row +//! offset into the flat table, so the heads occupy disjoint row ranges and +//! one id names one (head, row) pair. +//! +//! Heads are laid out **(order, head), order-major**: the `heads_per_ngram` +//! bigram heads first, then the trigram heads, and so on up to `ngram_size`. +//! The table index a head's remainder is taken against is therefore the same +//! number as the column that head's id lands in, which is what makes the layout +//! worth stating explicitly — a transposed (head, order) layout keeps every id +//! in range and produces a completely different embedding. +//! +//! Every id this module returns is a row index into the flat n-gram table whose +//! geometry the config derives: `heads * head_dim` wide rows, `sum(head_vocab_sizes)` +//! of them before padding. A gather turns `ids` into `rows x heads x head_dim` +//! activation, and nothing here needs the table itself to exist. +//! +//! Pinned by `pegainfer-qwen4/tests/ple_ngram_replay.rs` against a golden +//! recorded from the reference at the frozen checkpoint's real geometry: the ids +//! are compared with i64 equality at zero tolerance, so a wrong window, fold, +//! modulus, offset or layout fails rather than drifting. + +use anyhow::Result; +use anyhow::ensure; + +/// The n-gram hash and its inputs. +/// +/// Everything here is checkpoint data, not geometry: the reference derives the +/// per-head primes and their offsets from config (`_find_nth_prime_after` over +/// `ngram_vocab_size_base`) and also stores them, so a loader can build this +/// from either and cross-check. Constructing one validates the inputs, because +/// a bad modulus is a panic in `rem_euclid` and a short multiplier run is a +/// silent read past the window. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct NgramHasher { + eos_token_id: i64, + /// One odd multiplier per context position, `ngram_size` of them. Position + /// 0 is the token at the hashed row itself, so the fold always uses + /// `multipliers[0]` and an order-`k` n-gram adds `k - 1` more. + layer_multipliers: Vec, + /// Per-head prime modulus, `(ngram_size - 1) * heads_per_ngram` of them. + head_vocab_sizes: Vec, + /// Per-head row offset into the flat table, same length and laid out the + /// same way. + head_offsets: Vec, + heads_per_ngram: usize, +} + +impl NgramHasher { + /// Build a hasher from the checkpoint's own vectors. + /// + /// `layer_multipliers` must carry one entry per context position + /// (`ngram_size` of them, at least 2 — an n-gram needs a context token), + /// and both per-head vectors one entry per head. The reference generates its + /// multipliers as positive odd integers bounded by `(2^63 - 1) / + /// vocab_size`, which is what keeps every `token * multiplier` product and + /// their XOR fold non-negative; positivity is required here so that + /// `rem_euclid` and a truncated remainder cannot disagree on a fold that a + /// different checkpoint produced some other way. + /// + /// # Errors + /// + /// If the vectors are not the lengths the head layout implies, if any + /// multiplier is not positive, if any head modulus is not at least 2 (a + /// modulus of 1 hashes every token to one row, and 0 does not divide), if any + /// offset is negative, or if a head's `offset + modulus` overflows i64 — the + /// table row count is a sum of these, so an overflowing range is a table that + /// cannot be sized. Contiguity of the offsets is *not* checked here: that is a + /// property of the table the config derives, not of the hash. + pub fn new( + eos_token_id: i64, + layer_multipliers: Vec, + head_vocab_sizes: Vec, + head_offsets: Vec, + heads_per_ngram: usize, + ) -> Result { + ensure!( + heads_per_ngram > 0, + "qwen4_exp n-gram: heads_per_ngram must be positive" + ); + ensure!( + layer_multipliers.len() >= 2, + "qwen4_exp n-gram: layer_multipliers must carry one entry per context position \ + (ngram_size, at least 2), got {}", + layer_multipliers.len() + ); + let orders = layer_multipliers.len() - 1; + let heads = orders * heads_per_ngram; + ensure!( + head_vocab_sizes.len() == heads && head_offsets.len() == heads, + "qwen4_exp n-gram: {orders} orders x {heads_per_ngram} heads needs {heads} entries in \ + both per-head vectors, got {} moduli and {} offsets", + head_vocab_sizes.len(), + head_offsets.len() + ); + for (position, multiplier) in layer_multipliers.iter().enumerate() { + ensure!( + *multiplier > 0, + "qwen4_exp n-gram: layer_multipliers[{position}] must be positive, got \ + {multiplier}" + ); + } + for head in 0..heads { + let size = head_vocab_sizes[head]; + let offset = head_offsets[head]; + ensure!( + size >= 2, + "qwen4_exp n-gram: head {head} has modulus {size}, which cannot address a table" + ); + ensure!( + offset >= 0, + "qwen4_exp n-gram: head {head} has a negative table offset {offset}" + ); + ensure!( + offset.checked_add(size).is_some(), + "qwen4_exp n-gram: head {head}'s range {offset}..+{size} overflows i64" + ); + } + Ok(Self { + eos_token_id, + layer_multipliers, + head_vocab_sizes, + head_offsets, + heads_per_ngram, + }) + } + + /// The token that both pads the warm-up rows and closes a segment. + pub fn eos_token_id(&self) -> i64 { + self.eos_token_id + } + + /// `ngram_size - 1`: the number of order blocks in a row of + /// [`ids`](Self::ids), and — the same number, which is why + /// [`context`](Self::context) exists rather than a second derivation — of + /// warm-up rows the history is padded with. + pub fn orders(&self) -> usize { + self.layer_multipliers.len() - 1 + } + + /// Heads per order block. + pub fn heads_per_ngram(&self) -> usize { + self.heads_per_ngram + } + + /// Total heads, and the width of one row of [`NgramHasher::ids`]. + pub fn heads(&self) -> usize { + self.orders() * self.heads_per_ngram + } + + /// Warm-up rows the history is padded with: one per context token an + /// n-gram's window can reach back over. Always [`orders`](Self::orders). + pub fn context(&self) -> usize { + self.orders() + } + + /// The order an output column's block carries: slot 0 is the bigram, slot 1 + /// the trigram, and so on. + pub fn order_of_slot(&self, slot: usize) -> usize { + slot + 2 + } + + /// The table head an (order, head) pair's modulus and offset belong to. + /// + /// Under the reference's order-major layout this is also the *output column* + /// the id lands in, `slot * heads_per_ngram + head` for `slot = order - 2`. + /// That identity is worth stating because it is what makes a transposed + /// (head, order) layout silently wrong: every id still names a real table + /// row, so nothing goes out of range and only a value comparison catches it. + pub fn table_head_of(&self, order: usize, head: usize) -> usize { + (order - 2) * self.heads_per_ngram + head + } + + /// The history a whole-sequence forward hashes: [`context`](Self::context) + /// EOS rows, then the sequence. + /// + /// A decode step instead carries its own `ngram_size - 1` token window as + /// state and hands that in; this is the no-state case, which is what the + /// reference's prefill and what the golden record. + pub fn padded_history(&self, input_ids: &[i64]) -> Vec { + let mut history = vec![self.eos_token_id; self.context()]; + history.extend_from_slice(input_ids); + history + } + + /// For every position of `history`, the first position of the EOS-delimited + /// segment that contains it. + /// + /// This is the cummax over EOS positions that `_shift_right_ignore_eos` + /// performs: `segment_starts[p]` is one past the last EOS *strictly before* + /// `p`. An EOS token starts the **next** segment, it does not belong to the + /// one it closes, which is why the search excludes `p` itself. Position 0 is + /// always the start of its own segment. + /// + /// It is returned as data rather than folded into + /// [`shifted_context`](Self::shifted_context) because it is the piece a + /// resumed decode has to carry: a window that reaches back over a segment + /// boundary has to be reset, and the boundary is a property of the history, + /// not of the shift. + pub fn segment_starts(&self, history: &[i64]) -> Vec { + let mut starts = vec![0usize; history.len()]; + let mut segment_start = 0usize; + for (p, &token) in history.iter().enumerate() { + starts[p] = segment_start; + if token == self.eos_token_id { + segment_start = p + 1; + } + } + starts + } + + /// The reference's `_shift_right_ignore_eos` on one history: the token + /// `shift` positions back, or `eos_token_id` where the segment does not + /// reach that far. + /// + /// `segment_starts` is [`segment_starts`](Self::segment_starts) of the same + /// history. Position `p` reads `history[p - shift]` when the read stays + /// inside `p`'s segment (`p - segment_starts[p] >= shift`) and inside the + /// history (`p >= shift`), and the padding token otherwise. `shift` 0 is the + /// identity, so the windows for orders 2..=`ngram_size` are shifts + /// 0..`ngram_size`. + /// + /// # Panics + /// + /// If `segment_starts` is shorter than `history`. + pub fn shifted_context( + &self, + history: &[i64], + shift: usize, + segment_starts: &[usize], + ) -> Vec { + assert!( + segment_starts.len() >= history.len(), + "a shifted context needs one segment start per history position: {} for {}", + segment_starts.len(), + history.len() + ); + if shift == 0 { + return history.to_vec(); + } + (0..history.len()) + .map(|p| { + let in_segment = p - segment_starts[p]; + if in_segment >= shift && p >= shift { + history[p - shift] + } else { + self.eos_token_id + } + }) + .collect() + } + + /// Every window a whole-sequence forward needs, indexed by shift: entry `k` + /// is [`shifted_context`](Self::shifted_context) at shift `k`, for `k` in + /// 0..`ngram_size`. + pub fn windows(&self, history: &[i64], segment_starts: &[usize]) -> Vec> { + (0..self.layer_multipliers.len()) + .map(|shift| self.shifted_context(history, shift, segment_starts)) + .collect() + } + + /// One (row, order) mul-XOR fold: `windows[0][position] * multipliers[0]`, + /// XOR-ed with `windows[k][position] * multipliers[k]` for `1 <= k < order`. + /// + /// `position` is an index into the *history*, not into the output rows: row + /// `r` of a whole-sequence forward is at `r + `[`context`](Self::context). + /// The multiply wraps, which is faithful rather than load-bearing — the + /// reference bounds each multiplier so the products cannot overflow, and an + /// XOR of non-negative i64s keeps the sign bit clear, so the fold is never + /// negative for a hasher that passed [`new`](Self::new). + /// + /// # Panics + /// + /// If `order` exceeds the window count, or `position` exceeds the windows' + /// length. + pub fn fold(&self, windows: &[&[i64]], position: usize, order: usize) -> i64 { + assert!( + order >= 1 && order <= windows.len(), + "an order-{order} fold needs {order} windows, got {}", + windows.len() + ); + let mut mixed = windows[0][position].wrapping_mul(self.layer_multipliers[0]); + for (window, multiplier) in windows[1..order] + .iter() + .zip(&self.layer_multipliers[1..order]) + { + mixed ^= window[position].wrapping_mul(*multiplier); + } + mixed + } + + /// One head's table row for a fold: the fold modulo that head's prime, plus + /// that head's offset. + /// + /// `rem_euclid`, not `%`: the reference uses `torch.remainder`, a floored + /// remainder, which only differs from Rust's truncated one on a negative + /// fold — and a hasher that passed [`new`](Self::new) cannot produce one. + /// Using the Euclidean form keeps the two equal by construction rather than + /// by the accident of the multiplier bound. + /// + /// # Panics + /// + /// If `table_head` is out of range. + pub fn head_id(&self, fold: i64, table_head: usize) -> i64 { + fold.rem_euclid(self.head_vocab_sizes[table_head]) + self.head_offsets[table_head] + } + + /// Hash a whole sequence: one row of [`heads`](Self::heads) table indices per + /// input token, row-major. + /// + /// Column `slot * heads_per_ngram + head` holds the id of the head at + /// (`order_of_slot(slot)`, `head`), so the bigram block occupies the first + /// `heads_per_ngram` columns and the trigram block the next. The warm-up + /// rows are dropped: the output has one row per `input_ids` entry, read from + /// history position `row + `[`context`](Self::context). + pub fn ids(&self, input_ids: &[i64]) -> Vec { + let history = self.padded_history(input_ids); + let starts = self.segment_starts(&history); + let owned = self.windows(&history, &starts); + let windows: Vec<&[i64]> = owned.iter().map(Vec::as_slice).collect(); + let total = self.heads(); + let mut ids = vec![0i64; input_ids.len() * total]; + for slot in 0..self.orders() { + let order = self.order_of_slot(slot); + for (row, id_row) in ids.chunks_mut(total).enumerate() { + let mixed = self.fold(&windows, row + self.context(), order); + for head in 0..self.heads_per_ngram { + id_row[slot * self.heads_per_ngram + head] = + self.head_id(mixed, self.table_head_of(order, head)); + } + } + } + ids + } +} diff --git a/pegainfer-qwen4/tests/ple_ngram_replay.rs b/pegainfer-qwen4/tests/ple_ngram_replay.rs new file mode 100644 index 000000000..cde71cc54 --- /dev/null +++ b/pegainfer-qwen4/tests/ple_ngram_replay.rs @@ -0,0 +1,871 @@ +//! Host-side reader for `test_data/qwen38fn-ple-golden.safetensors` — the +//! consumer half of the Qwen3.8-Flash-Next per-layer-embedding (PLE) golden. +//! +//! The fixture was recorded once from the HF reference (`Qwen4ExpTextPLELayer` + +//! `Qwen4ExpTextNGramEmbedding`, transformers 5.16.1 `models/qwen4_exp`) by +//! `tools/accuracy/dump_qwen38fn_ple_golden.py`, at the frozen checkpoint's real +//! geometry. This gate replays it with integer and f32/f64 host math only — no +//! CUDA, no weights, no Python — so the two things a PLE implementation can get +//! wrong are pinned here and on the device: +//! +//! 1. **The n-gram hash, exactly.** `ngram_ids` is recomputed from `input_ids` +//! by [`pegainfer_qwen4::ngram::NgramHasher`], the production host module the +//! device-side table gather addresses through, and compared with i64 +//! equality at zero tolerance. It is pure integer math, so any deviation is a +//! bug and not noise: the EOS-resetting context window, the bigram/trigram +//! head layout (16 heads = 2 orders x 8) and the per-head mul-xor fold all +//! have to be right for a single id to land. +//! 2. **The PLE block, to the bf16 noise floor.** `out` is recomputed from the +//! stored `hidden`/`key`/`value` (the block's inputs, recorded past the +//! projections that the 95 GiB n-gram table feeds) and compared against the +//! stored block output within [`OUT_MAX_ABS`]. The replay below is a CPU +//! oracle, deliberately kept test-local: it is bit-identical to the stored +//! `out`, which is what makes it the cross-check for the device kernels that +//! compute the same block in +//! `pegainfer-kernels/csrc/qwen38fn/qwen38fn_ple_block.cu` (gated by +//! `pegainfer-kernels/tests/qwen38fn_ple_block_parity.rs`, behind that +//! crate's `qwen4` feature). +//! +//! Provenance is checked on every load, never skipped: the `qwen4_ple_golden` +//! manifest must parse, carry `format_version` 2, and name +//! `pegainfer-qwen4/tests/frozen_config.json` with the sha256 that file actually +//! hashes to. `src/config.rs::the_committed_fixture_is_the_frozen_revision` pins +//! that file to the frozen revision, so the chain is fixture -> config -> +//! checkpoint. The same loader also cross-checks the config's `eos_token_id` +//! against the fixture's stored scalar and every tensor width against the +//! config's geometry. The per-head moduli and offsets it replays with are the +//! fixture's own, checked here to be an increasing run with contiguous offsets — +//! the same table `NgramTable::derive` builds from config, whose endpoints and +//! totals `the_ngram_table_derivation_matches_the_checkpoint` pins. +//! +//! Negative controls: a gate that only ever passes proves nothing, so four +//! single-knob mutations of this same replay are asserted to *fail* — dropping +//! the EOS window reset, dilation 1 instead of `ngram_size`, an unsigned sqrt on +//! the gate, and a swapped bigram/trigram head order. +//! +//! Still ungated here: the n-gram table lookup/gather that turns `ngram_ids` +//! into the 160-dim head vectors, which needs the 95 GiB table's loader +//! (#1106), and the decode-side states (the 2-token id window and the 9-slot +//! conv state), since the recording is whole-sequence. + +use half::bf16; +use pegainfer_qwen4::ngram::NgramHasher; +use safetensors::Dtype; +use safetensors::SafeTensors; +use serde_json::Value; +use sha2::Digest; + +/// Metadata key carrying the fixture manifest. +const META_KEY: &str = "qwen4_ple_golden"; + +/// Layout revision this reader implements. The dump script bumps it on any layout +/// change; a mismatch fails loudly rather than being misread. +const FORMAT_VERSION: u64 = 2; + +/// Repo-relative frozen-config path the manifest must name. +const CONFIG_PATH: &str = "pegainfer-qwen4/tests/frozen_config.json"; + +/// Floor of the gate's signed sqrt, matching the reference's `clamp_min(1e-6)`. +const GATE_CLAMP_MIN: f32 = 1e-6; + +/// Block-output tolerance, absolute. +/// +/// The measured residual on the recording platform is **0.0**: this replay is +/// bit-identical to the stored `out` for both cases (0 of 81920 and 0 of 122880 +/// elements differ), because it emulates each bf16 rounding point of the +/// reference chain — including the gate's *double* rounding, where the group sum +/// is rounded to bf16 before the `1/sqrt(hidden)` divide, and again after it. +/// +/// The allowance covers the one non-portable op in that chain, `expf` (inside +/// `sigmoid` and `silu`); `sqrt`, division and multiplication are IEEE +/// correctly-rounded and the f64 reductions below are this file's own, so both +/// are machine-independent. A perturbation study — nudging *every* `sigmoid` and +/// *every* `silu` output by a whole bf16 ulp, far more than any libm +/// disagreement — moves `out` by at most 7.8e-3, i.e. 2 bf16 ulps at the largest +/// stored |out| (0.8438). 1.6e-2 is 2x that bound, and still 17x below the +/// weakest negative control (unsigned sqrt, 0.27); dilation 1 deviates by 0.79. +const OUT_MAX_ABS: f32 = 1.6e-2; + +/// The frozen checkpoint's geometry, read from the same file the manifest's +/// `config_sha256` is checked against. +struct Geometry { + hidden: usize, + hc: usize, + ngram_size: usize, + heads_per_ngram: usize, + conv_kernel: usize, + eps: f32, + eos: i64, +} + +impl Geometry { + /// Hyper-connection stream width: `hc_count * hidden_size`. + fn width(&self) -> usize { + self.hc * self.hidden + } + + /// Hashed n-gram heads: `(ngram_size - 1) * heads_per_ngram`. + fn heads(&self) -> usize { + (self.ngram_size - 1) * self.heads_per_ngram + } + + /// `sqrt(hidden_size)` in the f32 the reference's opmath divides by. The + /// gate divides by this rather than multiplying by its reciprocal: the two + /// round differently. + fn sqrt_hidden(&self) -> f32 { + (self.hidden as f64).sqrt() as f32 + } +} + +fn geometry() -> Geometry { + let cfg: Value = serde_json::from_str(include_str!("frozen_config.json")) + .expect("frozen_config.json must parse"); + let text = &cfg["text_config"]; + let field = |name: &str| -> usize { + text[name] + .as_u64() + .unwrap_or_else(|| panic!("frozen_config.json text_config.{name} must be an integer")) + as usize + }; + let geo = Geometry { + hidden: field("hidden_size"), + hc: field("hc_count"), + ngram_size: field("ngram_size"), + heads_per_ngram: field("heads_per_ngram"), + conv_kernel: field("ple_conv_kernel_size"), + eps: text["rms_norm_eps"] + .as_f64() + .expect("text_config.rms_norm_eps must be a number") as f32, + eos: text["eos_token_id"] + .as_i64() + .expect("text_config.eos_token_id must be an integer"), + }; + assert!( + geo.ngram_size >= 2 && geo.heads_per_ngram > 0 && geo.conv_kernel > 0, + "degenerate PLE geometry: ngram_size {} heads_per_ngram {} conv_kernel {}", + geo.ngram_size, + geo.heads_per_ngram, + geo.conv_kernel + ); + geo +} + +/// One recorded sequence: the block's inputs and both expected outputs. +struct Case { + name: String, + input_ids: Vec, + ngram_ids: Vec, + hidden: Vec, + key: Vec, + value: Vec, + out: Vec, +} + +impl Case { + fn seq_len(&self) -> usize { + self.input_ids.len() + } +} + +/// The whole fixture, decoded to owned vectors (the safetensors view borrows the +/// file bytes). +struct Fixture { + geo: Geometry, + /// The production hash, built from the fixture's own scalars and per-head + /// vectors. It carries the EOS id, the multipliers and both head runs, so + /// nothing else here needs a copy of them. + hasher: NgramHasher, + norm_key: Vec, + norm_query: Vec, + norm_conv: Vec, + /// Depthwise conv weight, `[width][1][conv_kernel]` flattened. + conv1d: Vec, + cases: Vec, +} + +fn i64_tensor(st: &SafeTensors<'_>, name: &str) -> (Vec, Vec) { + let t = st + .tensor(name) + .unwrap_or_else(|e| panic!("golden missing {name}: {e}")); + assert_eq!(t.dtype(), Dtype::I64, "{name} must be i64"); + let values = t + .data() + .as_chunks::<8>() + .0 + .iter() + .map(|b| i64::from_le_bytes(*b)) + .collect(); + (values, t.shape().to_vec()) +} + +fn bf16_tensor(st: &SafeTensors<'_>, name: &str, width: usize) -> Vec { + let t = st + .tensor(name) + .unwrap_or_else(|e| panic!("golden missing {name}: {e}")); + assert_eq!(t.dtype(), Dtype::BF16, "{name} must be bf16"); + let shape = t.shape().to_vec(); + assert_eq!( + shape.last().copied(), + Some(width), + "{name} must be {} wide, got shape {shape:?}", + width + ); + t.data() + .as_chunks::<2>() + .0 + .iter() + .map(|b| bf16::from_bits(u16::from_le_bytes(*b))) + .collect() +} + +impl Fixture { + fn load() -> Fixture { + let geo = geometry(); + let path = format!( + "{}/../test_data/qwen38fn-ple-golden.safetensors", + env!("CARGO_MANIFEST_DIR") + ); + let bytes = std::fs::read(&path).unwrap_or_else(|e| { + panic!( + "read {path}: {e} — the PLE golden is committed; regenerate it with \ + tools/accuracy/dump_qwen38fn_ple_golden.py only if the layout changes" + ) + }); + let st = SafeTensors::deserialize(&bytes).expect("parse the PLE golden safetensors"); + check_provenance(&bytes); + + let (eos, eos_shape) = i64_tensor(&st, "eos_token_id"); + assert_eq!(eos_shape, vec![1], "eos_token_id must be one scalar"); + assert_eq!( + eos[0], geo.eos, + "the fixture's eos_token_id must be the frozen config's" + ); + + let (multipliers, mult_shape) = i64_tensor(&st, "layer_multipliers"); + assert_eq!( + mult_shape, + vec![geo.ngram_size], + "layer_multipliers must carry one multiplier per context position" + ); + let (head_vocab, vocab_shape) = i64_tensor(&st, "ngram_heads_vocab_sizes"); + let (head_offsets, offset_shape) = i64_tensor(&st, "ngram_heads_offsets"); + assert_eq!( + (vocab_shape, offset_shape), + (vec![geo.heads()], vec![geo.heads()]), + "both n-gram head buffers must be one entry per head" + ); + assert!( + head_vocab.iter().all(|&size| size > 1) + && head_vocab.windows(2).all(|w| w[1] > w[0]) + && head_offsets[0] == 0 + && head_offsets + .windows(2) + .zip(&head_vocab) + .all(|(w, &size)| w[1] == w[0] + size), + "the head moduli must be an increasing run with contiguous offsets" + ); + + let width = geo.width(); + let conv1d = bf16_tensor(&st, "conv1d_weight", geo.conv_kernel); + let cases = load_cases(&st, &geo); + // `NgramHasher::new` validates lengths, multiplier positivity and each + // head's modulus/offset; the fixture's own run shape is checked above, + // because that is a fact about the dump rather than about the hasher. + let hasher = NgramHasher::new( + eos[0], + multipliers, + head_vocab, + head_offsets, + geo.heads_per_ngram, + ) + .expect("the fixture's n-gram vectors must build a hasher"); + Fixture { + geo, + hasher, + norm_key: bf16_tensor(&st, "norm_key", width), + norm_query: bf16_tensor(&st, "norm_query", width), + norm_conv: bf16_tensor(&st, "norm_conv", width), + conv1d, + cases, + } + } +} + +/// The manifest is the fixture's provenance: it must parse, carry the layout this +/// reader implements, and point at the frozen config by content hash. +fn check_provenance(bytes: &[u8]) { + let (_, header) = SafeTensors::read_metadata(bytes).expect("parse the PLE golden header"); + let manifest = header + .metadata() + .as_ref() + .and_then(|meta| meta.get(META_KEY)) + .unwrap_or_else(|| panic!("golden has no {META_KEY:?} metadata manifest")) + .clone(); + let manifest: Value = + serde_json::from_str(&manifest).expect("the golden manifest must be JSON"); + let version = manifest["format_version"] + .as_u64() + .expect("manifest.format_version must be an integer"); + assert_eq!( + version, FORMAT_VERSION, + "fixture layout {version} is not the {FORMAT_VERSION} this reader implements" + ); + assert_eq!( + manifest["config"].as_str(), + Some(CONFIG_PATH), + "manifest.config must name the frozen config this reader hashes" + ); + let digest = sha2::Sha256::digest(include_str!("frozen_config.json").as_bytes()); + let hex = digest + .iter() + .fold(String::with_capacity(64), |mut hex, byte| { + use std::fmt::Write as _; + let _ = write!(hex, "{byte:02x}"); + hex + }); + assert_eq!( + manifest["config_sha256"].as_str(), + Some(hex.as_str()), + "the fixture was recorded against a different {CONFIG_PATH}" + ); +} + +/// Decode every `caseN_*` group, checking the shapes against the config geometry +/// as it goes. +fn load_cases(st: &SafeTensors<'_>, geo: &Geometry) -> Vec { + let width = geo.width(); + let heads = geo.heads(); + let mut cases = Vec::new(); + for index in 0.. { + let name = format!("case{index}"); + if st.tensor(&format!("{name}_seq_len")).is_err() { + break; + } + let (seq_len_cell, seq_shape) = i32_tensor(st, &format!("{name}_seq_len")); + assert_eq!(seq_shape, vec![1], "{name}_seq_len must be one cell"); + let seq_len = usize::try_from(seq_len_cell[0]) + .ok() + .filter(|&len| len > 0) + .unwrap_or_else(|| panic!("{name}_seq_len must be a positive row count")); + + let (rows, rows_shape) = i32_tensor(st, &format!("{name}_rows")); + assert_eq!(rows_shape, vec![seq_len], "{name}_rows must be one per row"); + assert!( + rows.iter().enumerate().all(|(i, &row)| row == i as i32), + "{name}_rows must be the identity 0..{seq_len}: this replay is whole-sequence (the \ + dump records no cache state), so a row selection would need the conv and window \ + state of the dropped rows to stay correct" + ); + + let (input_ids, id_shape) = i64_tensor(st, &format!("{name}_input_ids")); + assert_eq!( + id_shape, + vec![seq_len], + "{name}_input_ids must be one per row" + ); + let (ngram_ids, ngram_shape) = i64_tensor(st, &format!("{name}_ngram_ids")); + assert_eq!( + ngram_shape, + vec![seq_len, heads], + "{name}_ngram_ids must be [rows, heads]" + ); + cases.push(Case { + name, + input_ids, + ngram_ids, + hidden: bf16_tensor(st, &format!("case{index}_hidden"), width), + key: bf16_tensor(st, &format!("case{index}_key"), width), + value: bf16_tensor(st, &format!("case{index}_value"), geo.hidden), + out: bf16_tensor(st, &format!("case{index}_out"), width), + }); + } + assert!( + cases.len() >= 2, + "the dump records two case groups (S=8 without EOS, S=12 with a mid-sequence EOS) but \ + only {} decoded — a group is missing or the naming changed", + cases.len() + ); + for case in &cases { + let rows = case.seq_len(); + assert!( + case.hidden.len() == rows * width + && case.key.len() == rows * width + && case.out.len() == rows * width + && case.value.len() == rows * geo.hidden + && case.ngram_ids.len() == rows * heads, + "{}: an activation does not cover {} rows of the config geometry", + case.name, + rows + ); + } + cases +} + +fn i32_tensor(st: &SafeTensors<'_>, name: &str) -> (Vec, Vec) { + let t = st + .tensor(name) + .unwrap_or_else(|e| panic!("golden missing {name}: {e}")); + assert_eq!(t.dtype(), Dtype::I32, "{name} must be i32"); + let values = t + .data() + .as_chunks::<4>() + .0 + .iter() + .map(|b| i32::from_le_bytes(*b)) + .collect(); + (values, t.shape().to_vec()) +} + +/// Knobs the negative controls turn. The defaults are the reference behaviour. +#[derive(Clone, Copy)] +struct Controls { + /// `_shift_right_ignore_eos`' cummax segment reset: n-grams never cross an + /// EOS boundary. Off = one window for the whole history. + eos_reset: bool, + /// Head layout, bigram block then trigram block. Off = the two swapped. + bigram_first: bool, +} + +impl Controls { + fn reference() -> Self { + Self { + eos_reset: true, + bigram_first: true, + } + } +} + +#[derive(Clone, Copy)] +struct Block { + /// The conv's dilation; the reference uses `ngram_size`. + dilation: usize, + /// Restore the gate's sign after the sqrt, as the reference does. + signed_sqrt: bool, +} + +impl Block { + fn reference(geo: &Geometry) -> Self { + Self { + dilation: geo.ngram_size, + signed_sqrt: true, + } + } +} + +/// Recompute `ngram_ids` — `[seq_len, heads]` flattened, row-major — from the +/// production hasher's own primitives, with the two controls applied. +/// +/// At [`Controls::reference`] this is [`NgramHasher::ids`] element for element, +/// which `the_ngram_hash_replays_exactly` asserts rather than assumes, so the +/// helper cannot drift away from the production path and still look like a +/// control. Each knob turns exactly one piece of the reference's semantics: +/// +/// - `eos_reset: false` hands `shifted_context` a segment start of 0 at every +/// position. That is the cummax replaced by nothing: every window reads its +/// full `shift` back, across any EOS boundary. It is data rather than a second +/// windowing function because the segment start *is* the reset — the primitive +/// takes it as an argument precisely so a resumed decode can carry its own. +/// - `bigram_first: false` walks the order blocks in reverse, so each slot's +/// column keeps its own head's modulus and offset but reduces the other +/// order's fold against them. +/// +/// The (slot, order) nesting below is the only part of the production `ids` this +/// re-derives; the padded history, the window, the fold, the (order, head) table +/// index and the per-head remainder are the production functions, so a control +/// that stopped exercising them would fail the exact test first. +fn hashed_ids(fx: &Fixture, input_ids: &[i64], controls: Controls) -> Vec { + let hasher = &fx.hasher; + let history = hasher.padded_history(input_ids); + let starts = if controls.eos_reset { + hasher.segment_starts(&history) + } else { + vec![0; history.len()] + }; + let owned = hasher.windows(&history, &starts); + let windows: Vec<&[i64]> = owned.iter().map(Vec::as_slice).collect(); + let heads_per_ngram = hasher.heads_per_ngram(); + let orders = hasher.orders(); + let total_heads = hasher.heads(); + let mut ids = vec![0i64; input_ids.len() * total_heads]; + for slot in 0..orders { + let order = if controls.bigram_first { + hasher.order_of_slot(slot) + } else { + hasher.order_of_slot(orders - 1 - slot) + }; + for row in 0..input_ids.len() { + let mixed = hasher.fold(&windows, row + hasher.context(), order); + for head in 0..heads_per_ngram { + ids[row * total_heads + slot * heads_per_ngram + head] = + hasher.head_id(mixed, hasher.table_head_of(order, head)); + } + } + } + ids +} + +/// The reference's grouped RMSNorm: `bf16(x_f32 * rsqrt(mean(x^2) + eps) * (1 + +/// w))` over `hc_count` groups of `hidden_size`, with `x^2` rounded to f32 per +/// element exactly as `pow(2)` does. +/// +/// The mean accumulates in f64. The reference's f32 tree reduction lands within +/// an f32 ulp of that, and the trailing bf16 rounding absorbs the difference — +/// verified bit-identical end to end on both cases — so the exact sum is the +/// portable choice: it does not depend on the shape of the reference's tree. +fn grouped_norm(x: &[bf16], weight: &[bf16], geo: &Geometry) -> Vec { + let width = geo.width(); + assert!( + !x.is_empty() && x.len().is_multiple_of(width) && weight.len() == width, + "a grouped norm needs whole rows of {width} channels and one weight each" + ); + let mut out = vec![bf16::ZERO; x.len()]; + for (row, chunk) in x.chunks(width).enumerate() { + for base in (0..width).step_by(geo.hidden) { + let group = &chunk[base..base + geo.hidden]; + let sum = group + .iter() + .map(|v| { + let f = v.to_f32(); + f64::from(f * f) + }) + .sum::(); + let mean = (sum as f32) / geo.hidden as f32; + let scale = 1.0 / (mean + geo.eps).sqrt(); + for (j, &v) in group.iter().enumerate() { + let normed = v.to_f32() * scale; + out[row * width + base + j] = + bf16::from_f32(normed * (1.0 + weight[base + j].to_f32())); + } + } + } + out +} + +/// `torch.sign`: 0 at zero, unlike Rust's `signum` on `-0.0`. +fn sign(x: f32) -> f32 { + f32::from(x > 0.0) - f32::from(x < 0.0) +} + +fn sigmoid(x: bf16) -> bf16 { + let v = x.to_f32(); + bf16::from_f32(1.0 / (1.0 + (-v).exp())) +} + +fn silu(x: bf16) -> bf16 { + let v = x.to_f32(); + bf16::from_f32(v / (1.0 + (-v).exp())) +} + +/// Replay the PLE block on one case: `[seq_len, hc_count * hidden_size]`. +/// +/// Port of `Qwen4ExpTextPLELayer.forward` from its two recorded projections, with +/// a bf16 rounding wherever the reference's submodule boundaries put one. +fn ple_block(fx: &Fixture, case: &Case, knobs: Block) -> Vec { + let geo = &fx.geo; + let width = geo.width(); + let seq_len = case.seq_len(); + let key_normed = grouped_norm(&case.key, &fx.norm_key, geo); + let query_normed = grouped_norm(&case.hidden, &fx.norm_query, geo); + let sqrt_hidden = geo.sqrt_hidden(); + + // `sigmoid(signed_sqrt(dot(key, query) / sqrt(hidden))) * value`, broadcast + // over the `hc_count` streams that share one value vector. + let mut gated = vec![bf16::ZERO; seq_len * width]; + for row in 0..seq_len { + for group in 0..geo.hc { + let mut acc = 0f64; + for j in 0..geo.hidden { + let i = row * width + group * geo.hidden + j; + let product = bf16::from_f32(key_normed[i].to_f32() * query_normed[i].to_f32()); + acc += f64::from(product.to_f32()); + } + // Two bf16 roundings: the sum, then the divide. Collapsing them into + // one is worth ~1 ulp on every gate and breaks bit-exactness. + let sum = bf16::from_f32(acc as f32); + let pre_gate = bf16::from_f32(sum.to_f32() / sqrt_hidden); + let magnitude = bf16::from_f32(pre_gate.to_f32().abs().max(GATE_CLAMP_MIN)); + let root = bf16::from_f32(magnitude.to_f32().sqrt()); + let restored = if knobs.signed_sqrt { + sign(pre_gate.to_f32()) + } else { + 1.0 + }; + let gate = bf16::from_f32(root.to_f32() * restored); + let weight = sigmoid(gate).to_f32(); + for j in 0..geo.hidden { + let value = case.value[row * geo.hidden + j].to_f32(); + gated[row * width + group * geo.hidden + j] = bf16::from_f32(weight * value); + } + } + } + + // `gated + silu(causal dilated depthwise conv(norm_conv(gated)))` over a zero + // initial state of `(kernel - 1) * dilation` positions. + let conv_in = grouped_norm(&gated, &fx.norm_conv, geo); + let state = (geo.conv_kernel - 1) * knobs.dilation; + let mut out = vec![bf16::ZERO; seq_len * width]; + for channel in 0..width { + for row in 0..seq_len { + let mut acc = 0f32; + for tap in 0..geo.conv_kernel { + let source = row as i64 - (state - knobs.dilation * tap) as i64; + if source >= 0 { + acc += conv_in[source as usize * width + channel].to_f32() + * fx.conv1d[channel * geo.conv_kernel + tap].to_f32(); + } + } + let convolved = bf16::from_f32(acc); + out[row * width + channel] = + bf16::from_f32(gated[row * width + channel].to_f32() + silu(convolved).to_f32()); + } + } + out +} + +/// Element-wise agreement of a replay against the stored tensor. +struct Delta { + max_abs: f32, + mean_abs: f32, + identical: usize, + total: usize, +} + +fn delta(got: &[bf16], want: &[bf16]) -> Delta { + assert_eq!(got.len(), want.len(), "replay shape must match the golden"); + let mut max_abs = 0f32; + let mut sum_abs = 0f64; + let mut identical = 0; + for (g, w) in got.iter().zip(want) { + let d = (g.to_f32() - w.to_f32()).abs(); + max_abs = max_abs.max(d); + sum_abs += f64::from(d); + identical += usize::from(g.to_bits() == w.to_bits()); + } + Delta { + max_abs, + mean_abs: (sum_abs / want.len() as f64) as f32, + identical, + total: want.len(), + } +} + +fn eos_positions(case: &Case, eos: i64) -> Vec { + case.input_ids + .iter() + .enumerate() + .filter(|(_, token)| **token == eos) + .map(|(p, _)| p) + .collect() +} + +#[test] +fn the_ngram_hash_replays_exactly() { + let fx = Fixture::load(); + for case in &fx.cases { + let got = fx.hasher.ids(&case.input_ids); + // The controls' composition helper must be the production path at the + // reference settings, or a control would be measuring its own copy of + // the hash rather than the one this test pins. + assert_eq!( + hashed_ids(&fx, &case.input_ids, Controls::reference()), + got, + "{}: the control helper is not the production hash at the reference settings", + case.name + ); + let mismatched = got + .iter() + .zip(&case.ngram_ids) + .filter(|(g, w)| g != w) + .count(); + println!( + "{}: S={} eos_at={:?} heads={} ids={}/{} exact", + case.name, + case.seq_len(), + eos_positions(case, fx.hasher.eos_token_id()), + fx.geo.heads(), + case.ngram_ids.len() - mismatched, + case.ngram_ids.len() + ); + assert_eq!( + mismatched, + 0, + "{}: {} of {} hashed n-gram ids differ from the golden (zero tolerance: this is \ + integer math)", + case.name, + mismatched, + case.ngram_ids.len() + ); + } +} + +#[test] +fn the_ple_block_replays_within_the_bf16_noise_floor() { + let fx = Fixture::load(); + for case in &fx.cases { + let got = ple_block(&fx, case, Block::reference(&fx.geo)); + let d = delta(&got, &case.out); + println!( + "{}: S={} out max_abs={:.3e} mean_abs={:.3e} bit_identical={}/{} tol={:.1e} \ + dilation={} signed_sqrt=true", + case.name, + case.seq_len(), + d.max_abs, + d.mean_abs, + d.identical, + d.total, + OUT_MAX_ABS, + fx.geo.ngram_size + ); + assert!( + d.max_abs <= OUT_MAX_ABS, + "{}: PLE block output deviates by {:.3e} > {:.3e} (mean {:.3e}, {}/{} elements \ + differ); see OUT_MAX_ABS for how the bound is derived", + case.name, + d.max_abs, + OUT_MAX_ABS, + d.mean_abs, + d.total - d.identical, + d.total + ); + } +} + +/// Control (a): without the cummax EOS reset a window reaches across a segment +/// boundary, so the case with a mid-sequence EOS must change — and the EOS-free +/// case must not, which is what makes this a targeted control rather than a +/// blanket break. The change is narrow by construction: only the row right after +/// the EOS pulls a token from across it, and only for the orders whose window +/// reaches that far (the trigram heads, 8 of the 16). +#[test] +fn dropping_the_eos_window_reset_breaks_the_eos_case() { + let fx = Fixture::load(); + let controls = Controls { + eos_reset: false, + ..Controls::reference() + }; + let mut broke = 0; + for case in &fx.cases { + let positions = eos_positions(case, fx.hasher.eos_token_id()); + let got = hashed_ids(&fx, &case.input_ids, controls); + let mismatched = got + .iter() + .zip(&case.ngram_ids) + .filter(|(g, w)| g != w) + .count(); + println!( + "{}: S={} eos_at={:?} -> {}/{} ids differ without the EOS reset", + case.name, + case.seq_len(), + positions, + mismatched, + case.ngram_ids.len() + ); + if positions.is_empty() { + assert_eq!( + mismatched, 0, + "{} has no EOS, so dropping the reset must leave its ids alone", + case.name + ); + } else { + assert!( + mismatched > 0, + "{} carries an EOS at {positions:?} but dropping the reset changed nothing — the \ + control no longer exercises the segment boundary", + case.name + ); + broke += 1; + } + } + assert!( + broke > 0, + "no case in the fixture carries a mid-sequence EOS, so the reset is untested" + ); +} + +/// Control (d): the head layout is (order, head) with the bigram block first. +/// Swapping the two blocks keeps every id in range, so only a layout assertion +/// catches it. +#[test] +fn swapping_the_bigram_and_trigram_heads_breaks_the_hash() { + let fx = Fixture::load(); + let controls = Controls { + bigram_first: false, + ..Controls::reference() + }; + for case in &fx.cases { + let got = hashed_ids(&fx, &case.input_ids, controls); + let mismatched = got + .iter() + .zip(&case.ngram_ids) + .filter(|(g, w)| g != w) + .count(); + println!( + "{}: {}/{} ids differ with the order blocks swapped", + case.name, + mismatched, + case.ngram_ids.len() + ); + assert_eq!( + mismatched, + case.ngram_ids.len(), + "{}: swapping bigram/trigram must move every id — the two orders occupy disjoint \ + offset ranges and each remainder is smaller than the gap between them, so a swapped \ + id can never land on the stored one", + case.name + ); + } +} + +/// Control (b): the conv is dilated by `ngram_size`, with a `(kernel - 1) * +/// dilation` zero state. Dilation 1 reads the wrong taps. +#[test] +fn dilation_one_breaks_the_block_output() { + let fx = Fixture::load(); + let knobs = Block { + dilation: 1, + ..Block::reference(&fx.geo) + }; + assert_block_control_fails(&fx, knobs, "dilation 1"); +} + +/// Control (c): the gate is `sign(g) * sqrt(|g|)`. Dropping the sign restore +/// pushes every negative gate through `sigmoid` as positive, which inverts the +/// value's contribution for those streams. +#[test] +fn an_unsigned_sqrt_breaks_the_block_output() { + let fx = Fixture::load(); + let knobs = Block { + signed_sqrt: false, + ..Block::reference(&fx.geo) + }; + assert_block_control_fails(&fx, knobs, "unsigned sqrt"); +} + +fn assert_block_control_fails(fx: &Fixture, knobs: Block, label: &str) { + let mut worst = 0f32; + for case in &fx.cases { + let got = ple_block(fx, case, knobs); + let d = delta(&got, &case.out); + worst = worst.max(d.max_abs); + println!( + "{}: out max_abs={:.3e} mean_abs={:.3e} differing={}/{} under {label} (tol {:.1e})", + case.name, + d.max_abs, + d.mean_abs, + d.total - d.identical, + d.total, + OUT_MAX_ABS + ); + assert!( + d.max_abs > OUT_MAX_ABS, + "{}: {label} stayed within {:.3e} of the golden, so the block replay does not pin the \ + knob this control turns", + case.name, + OUT_MAX_ABS + ); + } + println!( + "{label}: worst deviation {worst:.3e} = {:.0}x the tolerance", + worst / OUT_MAX_ABS + ); +} diff --git a/test_data/qwen38fn-ple-golden.safetensors b/test_data/qwen38fn-ple-golden.safetensors new file mode 100644 index 000000000..a3d99bf71 Binary files /dev/null and b/test_data/qwen38fn-ple-golden.safetensors differ diff --git a/tools/accuracy/dump_qwen38fn_ple_golden.py b/tools/accuracy/dump_qwen38fn_ple_golden.py new file mode 100644 index 000000000..c01a67a8e --- /dev/null +++ b/tools/accuracy/dump_qwen38fn_ple_golden.py @@ -0,0 +1,204 @@ +#!/usr/bin/env python3 +"""Generate the PLE (n-gram) block golden at real frozen-config geometry. + +Dumps the Qwen3.8-Flash-Next per-layer-embedding reference (`Qwen4ExpTextPLELayer` ++ `Qwen4ExpTextNGramEmbedding`, transformers 5.16.1 `models/qwen4_exp`) into a +fixture a PLE kernel gate replays. All geometry comes from the real frozen config +(`CONFIG`): hidden 2560, `hc_count` 4 (stream 10240), `ple_embed_dim` 2560, +`ngram_size` 3, `heads_per_ngram` 8 (16 heads x 160), and the real per-head prime +vocabularies/offsets/multipliers. The values cover the mul-xor n-gram hash +(collisions by design), the EOS-resetting context shift, the grouped `(1+w)` norms +(group_size 2560), the signed-sqrt gate, and the **dilated** depthwise conv +(dilation = `ngram_size` = 3, state length 9) the non-dilated in-tree conv1d +cannot serve. + +Three boundaries are moved because the real tensors cannot be stored: + +(a) The n-gram table (~95 GiB bf16) cannot be built, so `nn.Embedding` is capped to + a tiny stub (`_tiny_embedding`) with the real ids folded in via a recorded + `remainder`. The captured ids are the REAL ones and are stored (`ngram_ids`); + the embedding *activation* (a lookup into that stub, weights not stored) is not + in the fixture -- the `key`/`value` projections it feeds are. +(b) `key_proj` [10240, 2560] and `value_proj` [2560, 2560] weights are NOT stored; + forward hooks capture their outputs as `key`/`value`. +(c) Otherwise all shapes are real (norms, conv weight, I64 hash buffers, every + per-case activation); there is no final GEMM -- the block output is terminal. + +Cases: S=8 (no EOS) and S=12 (one mid-sequence EOS, the segment reset). Fixture: +one JSON `qwen4_ple_golden` blob (format_version 2, config path + sha256, seed, +versions), the `eos_token_id` scalar, the storeable weights, and per case the +inputs, `rows` slice, hashed ids, recorded projections and expected block output. +Scope: full-sequence (cache=None) only; the decode states (2-token id window and +dilated conv state) need their own dump. + + python tools/accuracy/dump_qwen38fn_ple_golden.py --out +""" + +from __future__ import annotations + +import argparse +import contextlib +import hashlib +import json +import platform +from pathlib import Path + +import torch +import transformers +from safetensors.torch import save_file + +from transformers.models.qwen4_exp.configuration_qwen4_exp import Qwen4ExpConfig +from transformers.models.qwen4_exp.modeling_qwen4_exp import Qwen4ExpTextPLELayer + +# Frozen checkpoint config. Every geometry below is read from it, not written here. +CONFIG = "pegainfer-qwen4/tests/frozen_config.json" + +SEED = 0x_03_38_F9_5A_11_83 + +# Bumped when the fixture layout changes; the gate asserts it on load. +FORMAT_VERSION = 2 + +# Real n-gram table would be `padded_vocab_size x 160` (~95 GiB bf16). Cap the stub +# table here; real ids are folded in with `remainder(STUB_VOCAB)` (boundary (a)). +STUB_VOCAB = 4096 + +# One PLE layer: config.ple_layer_ids == [2] -> layer_idx 1 -> ple_layer_index 0. +LAYER_IDX = 1 +PLE_LAYER_INDEX = 0 + +# (seq_len, eos_positions). S=8 has no EOS (pure local window); S=12 places one +# mid-sequence EOS so the segment reset has a boundary on both sides of it. +CASES = [(8, ()), (12, (5,))] + + +@contextlib.contextmanager +def _tiny_embedding(cap: int): + """Cap `nn.Embedding` at `cap` rows so the 95 GiB n-gram table builds. + + Scoped to the module constructor; the reference hash math is untouched, so the + ids it computes stay real. + """ + original = torch.nn.Embedding.__init__ + + def patched(self, num_embeddings, embedding_dim, *args, **kwargs): + return original(self, min(num_embeddings, cap), embedding_dim, *args, **kwargs) + + torch.nn.Embedding.__init__ = patched + try: + yield + finally: + torch.nn.Embedding.__init__ = original + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--out", required=True, type=Path) + args = parser.parse_args() + + root = Path(__file__).resolve().parents[2] + raw = (root / CONFIG).read_bytes() + config_sha256 = hashlib.sha256(raw).hexdigest() + cfg = Qwen4ExpConfig(**json.loads(raw)).text_config + + hidden_size = cfg.hidden_size + hc_count = cfg.hc_count + vocab_size = cfg.vocab_size + eos_token_id = cfg.eos_token_id + + torch.manual_seed(SEED) + with _tiny_embedding(STUB_VOCAB): + module = Qwen4ExpTextPLELayer(cfg, LAYER_IDX, PLE_LAYER_INDEX) + with torch.no_grad(): + # RMSNorm weights default to zero; give every weight a non-trivial value so a + # kernel that ignored one of the `(1 + w)` scales would not sail through. + for p in module.parameters(): + p.normal_(0.0, 0.05 if p.dim() >= 2 else 0.5) + module = module.to(torch.bfloat16) + ng = module.ple_embedding + + tensors = { + "eos_token_id": torch.tensor([eos_token_id], dtype=torch.int64), + "layer_multipliers": ng.layer_multipliers.detach().contiguous(), + "ngram_heads_offsets": ng.ngram_heads_offsets.detach().contiguous(), + "ngram_heads_vocab_sizes": ng.ngram_heads_vocab_sizes.detach().contiguous(), + "norm_key": module.norm_key.weight.detach().contiguous(), + "norm_query": module.norm_query.weight.detach().contiguous(), + "norm_conv": module.norm_conv.weight.detach().contiguous(), + "conv1d_weight": module.conv1d.weight.detach().contiguous(), + } + + # Forward hooks capture exactly what the reference computes (boundaries a + b). + captured: dict[str, torch.Tensor] = {} + + def capture_ids(cap: int): + def hook(_mod, inputs): + ids = inputs[0] + captured["ids"] = ids.detach().clone() + return (ids.remainder(cap),) # fold real ids into the stub table + + return hook + + def capture(name: str): + def hook(_mod, _inputs, output): + captured[name] = output.detach().clone() + + return hook + + handles = [ + ng.ngram_embedding.register_forward_pre_hook(capture_ids(STUB_VOCAB)), + module.key_proj.register_forward_hook(capture("key")), + module.value_proj.register_forward_hook(capture("value")), + ] + + for i, (seq_len, eos_positions) in enumerate(CASES): + gen = torch.Generator().manual_seed(SEED + i) + hidden = torch.randn(1, seq_len, hc_count * hidden_size, generator=gen).to(torch.bfloat16) + ids_gen = torch.Generator().manual_seed(SEED + i) + input_ids = torch.randint(1, vocab_size, (1, seq_len), generator=ids_gen) + # Keep random draws off the EOS id, reserved for the explicit EOS positions. + input_ids[input_ids == eos_token_id] = eos_token_id + 1 + for eos_pos in eos_positions: + input_ids[0, eos_pos] = eos_token_id + + with torch.no_grad(): + out = module(hidden, input_ids, None) + ngram_ids = captured["ids"] + key_raw = captured["key"] + value = captured["value"] + + # Row selection: our cases are short, so every row is evaluated, but the + # convention is kept so a longer case can drop rows without a layout change. + rows = torch.arange(seq_len, dtype=torch.int32) + p = f"case{i}_" + tensors[p + "seq_len"] = torch.tensor([seq_len], dtype=torch.int32) + tensors[p + "rows"] = rows.contiguous() + tensors[p + "input_ids"] = input_ids[0].to(torch.int64).contiguous() + tensors[p + "hidden"] = hidden[0].index_select(0, rows.long()).contiguous() + tensors[p + "ngram_ids"] = ngram_ids[0].index_select(0, rows.long()).contiguous() + tensors[p + "key"] = key_raw[0].index_select(0, rows.long()).contiguous() + tensors[p + "value"] = value[0].index_select(0, rows.long()).contiguous() + tensors[p + "out"] = out[0].index_select(0, rows.long()).contiguous() + print( + f"case{i}: S={seq_len} eos={list(eos_positions)} rows={seq_len} " + f"heads={ngram_ids.shape[-1]}" + ) + + for h in handles: + h.remove() + + args.out.parent.mkdir(parents=True, exist_ok=True) + manifest = { + "format_version": FORMAT_VERSION, + "config": CONFIG, + "config_sha256": config_sha256, + "seed": SEED, + "transformers_version": transformers.__version__, + "torch_version": torch.__version__, + "platform": platform.machine(), + } + save_file(tensors, str(args.out), metadata={"qwen4_ple_golden": json.dumps(manifest, sort_keys=True)}) + size_mb = args.out.stat().st_size / (1024 * 1024) + print(f"wrote {args.out} ({size_mb:.2f} MiB, {len(tensors)} tensors)") + +if __name__ == "__main__": + main()