Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 12 additions & 4 deletions docs/models/qwen4/bring-up.md
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
# 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 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 and block-selection 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.
>
> **Last touched:** 2026-10
> **Last touched:** 2026-11

Tracked in #1103 (roadmap) with slices A #1104 / B #1127 / C #1128 and component bricks #1105 (QSA), #1106 (n-gram), #1107 (MoE), #1108 (hyper-connections).

Expand Down Expand Up @@ -44,14 +44,22 @@ The committed `config.json` (`pegainfer-qwen4/tests/frozen_config.json`) is the

`#1105` splits QSA into three gateable units. The middle one — the token indexer's block scores — is in `pegainfer-kernels/csrc/qwen4/` behind the kernels crate's own `qwen4` feature: the model crate above still has **no CUDA dependency at all**, and this is the line's first device code. The directory is named after the model line like `qwen35` and `glm52`, not after the checkpoint.

Two kernels against the frozen transformers-5.16.1 semantics. One pools every complete micro-block of `indexer_compress_ratio` tokens into a single key — a **float32** mean of the block's raw keys, cast back to bf16 exactly as the reference casts it, then `(1 + w)` RMSNorm and partial RoPE at the block's first token — and the other norms each of a query row's four heads at the row's own position and reduces `sum_h relu(q_h . k_b) / sqrt(indexer_head_dim)`. The pooled half is a separate pass because it does not depend on the query row at all: the reference rebuilds it per row from the same keys through the same bf16 rounding points, so one pass over the keys is the same numbers for every row that can see them. `ratio = 4` divides the 16-token KV page, so a micro-block never crosses one and the selection stage will need no block↔page mapping — but the pool pass does not use that, and the **block grid starts at the first visible token**, which this kernel requires to be token 0 (causal text over an un-evicted key cache; a visible list anchored elsewhere would move the grid, and nothing in the kernel can detect that, so such a caller has to hand it a shifted key array). The target is the portable subset: no TMA, no cluster, no wgmma. Both kernels are SIMT shuffle reductions rather than `mma.sync` bodies — one block's score is a 128-wide dot over 4 query heads, a shape `m16n8k16` cannot tile without padding 4 rows to 16.
Two kernels against the frozen transformers-5.16.1 semantics. One pools every complete micro-block of `indexer_compress_ratio` tokens into a single key — a **float32** mean of the block's raw keys, cast back to bf16 exactly as the reference casts it, then `(1 + w)` RMSNorm and partial RoPE at the block's first token — and the other norms each of a query row's four heads at the row's own position and reduces `sum_h relu(q_h . k_b) / sqrt(indexer_head_dim)`. The pooled half is a separate pass because it does not depend on the query row at all: the reference rebuilds it per row from the same keys through the same bf16 rounding points, so one pass over the keys is the same numbers for every row that can see them. `ratio = 4` divides the 16-token KV page, so a micro-block never crosses one and the selection stage needs no block↔page mapping — but the pool pass does not use that, and the **block grid starts at the first visible token**, which this kernel requires to be token 0 (causal text over an un-evicted key cache; a visible list anchored elsewhere would move the grid, and nothing in the kernel can detect that, so such a caller has to hand it a shifted key array). The target is the portable subset: no TMA, no cluster, no wgmma. Both kernels are SIMT shuffle reductions rather than `mma.sync` bodies — one block's score is a 128-wide dot over 4 query heads, a shape `m16n8k16` cannot tile without padding 4 rows to 16.

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

The gate replays `test_data/qwen38fn-qsa-indexer-golden.safetensors` (PR #1150) and holds the fixture's recorded `format_version`, `transformers_version` and `config_sha256` against the frozen config — and derives its geometry from that config, cross-checked against the fixture's own tensor shapes — so a golden from another checkpoint or another semantics fails as a score mismatch instead of being skipped. It reads the fixture's **validity prefix** per row rather than comparing through the NaN padding the reference leaves, and proves the kernel writes exactly that prefix by checking the untouched tail still holds a sentinel. Six of seven cases land within 9.6e-7 of the reference and are held to a 2e-6 bound; the 2055-token case, the only one with positions past 63, keeps a 5e-3 bound until its 2.4e-3 residual is traced (its shape suggests one flipped bf16 rounding decision). In-range `rows`, and `visible_len` up to the row's own token, are contract: anything else traps on the device, and a poisoned key keeps its NaN score instead of reading as a valid zero. The kernel takes no budget: the reference scores every complete block and truncates only at selection. **Next: that selection** — top-k over the block scores plus the unconditional tail, which is a unit of its own and depends on the same golden, not on this kernel.
The gate replays `test_data/qwen38fn-qsa-indexer-golden.safetensors` (PR #1150) and holds the fixture's recorded `format_version`, `transformers_version` and `config_sha256` against the frozen config — and derives its geometry from that config, cross-checked against the fixture's own tensor shapes — so a golden from another checkpoint or another semantics fails as a score mismatch instead of being skipped. It reads the fixture's **validity prefix** per row rather than comparing through the NaN padding the reference leaves, and proves the kernel writes exactly that prefix by checking the untouched tail still holds a sentinel. Six of seven cases land within 9.6e-7 of the reference and are held to a 2e-6 bound; the 2055-token case, the only one with positions past 63, keeps a 5e-3 bound until its 2.4e-3 residual is traced (its shape suggests one flipped bf16 rounding decision). In-range `rows`, and `visible_len` up to the row's own token, are contract: anything else traps on the device, and a poisoned key keeps its NaN score instead of reading as a valid zero. The kernel takes no budget: the reference scores every complete block and truncates only at selection.

## The QSA indexer's selection half

`csrc/qwen4/qwen4_qsa_indexer_select.cu` (`qsa_indexer_select_into`) turns the block scores into the token-level mask the sparse attention reads; the contract (tie rule, unconditional tail, trap, grid) is on the kernel header, and the wrapper doc carries the Rust-side buffer rules.

The select is an exact 8-bit radix select over the row's f32 keys: one 256k-token row at budget 2048 takes 298 µs on an L20, and the all-kept fast path takes 1.4 µs.

Two boundaries the code cannot state. A row's visible list starts at token 0 (the scoring pass's grid), so a prefix-evicted visible list is out of contract and needs a shifted key array on both stages. And an accuracy gate must not require mask equality with the reference at a tied cut: `torch.topk` leaves that pick open, so the reference's choice can differ from the kernel's canonical one and still be valid.

## The n-gram table is derived, not read

Expand Down
4 changes: 2 additions & 2 deletions pegainfer-kernels/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -38,8 +38,8 @@ deepseek-v2-lite = []
# and the TileLang-generated hd512 global-attention prefill.
gemma4 = []
qwen35 = []
# Qwen4 line (Qwen3.8-Flash-Next): the QSA token-indexer block-scoring kernels
# and the grouped (1+weight) RMSNorm.
# Qwen4 line (Qwen3.8-Flash-Next): the QSA token-indexer block-scoring and
# block-selection kernels and the grouped (1+weight) RMSNorm.
qwen4 = []
# Shared MoE/MLA third-party substrate: DeepEP, DeepGEMM, and FlashMLA.
glm52 = ["moe"]
Expand Down
Loading
Loading