feat(qwen4): score QSA token-indexer blocks in native CUDA

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 it is CPU-only. 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 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
 
@@ -40,6 +40,19 @@ The committed `config.json` (`pegainfer-qwen4/tests/frozen_config.json`) is the
 
 **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.
 
+## The QSA indexer's scoring half
+
+`#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.
+
+```
+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 n-gram table is derived, not read
 
 `config.json` states none of the table's dimensions; they follow from five fields:
@@ -75,6 +88,7 @@ Still to verify is the reuse on this line's activations. The tensor contract der
 - **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.
 - **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.
 - **Do not use `load_shard_info_fixed`.** It is a Qwen3.5-specific shard-*filename* workaround living in shared `pegainfer-core`; this checkpoint uses standard `model-0000N-of-00131` naming and wants the plain loader.
 
 ## Still open

pegainfer-kernels/Cargo.toml

@@ -15,6 +15,11 @@ tvm-ffi = { version = "0.1.0-alpha.0", optional = true }
 [dev-dependencies]
 # RLIMIT_FSIZE, to produce a real short write in the algo-store test.
 libc = { workspace = true }
+# The Qwen3.8-Flash-Next QSA indexer gate replays a committed golden fixture
+# and holds its provenance against the frozen config.
+safetensors = { workspace = true }
+serde_json = { workspace = true }
+sha2 = { workspace = true }
 
 [build-dependencies]
 cc = { workspace = true }
@@ -30,6 +35,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.
+qwen4 = []
 # Shared MoE/MLA third-party substrate: DeepEP, DeepGEMM, and FlashMLA.
 glm52 = ["moe"]
 # Kimi K3: TileLang-generated kernels, FP8xFP4 masked grouped GEMM, and the

pegainfer-kernels/build.rs

@@ -768,6 +768,17 @@ fn is_qwen35_source(csrc_dir: &Path, path: &Path) -> bool {
     }
 }
 
+/// Qwen4 line (Qwen3.8-Flash-Next) model-local CUDA: the QSA token indexer.
+/// Named after the model line, like qwen35 and glm52.
+fn is_qwen4_source(csrc_dir: &Path, path: &Path) -> bool {
+    match path.strip_prefix(csrc_dir) {
+        Ok(relative) => relative
+            .components()
+            .any(|part| part.as_os_str() == "qwen4"),
+        Err(_) => false,
+    }
+}
+
 /// GLM5.2 model-local CUDA shims over DeepGEMM/FlashMLA contracts.
 fn is_glm52_source(csrc_dir: &Path, path: &Path) -> bool {
     match path.strip_prefix(csrc_dir) {
@@ -2212,6 +2223,7 @@ fn main() {
     // --- k3: DeepGEMM-only, no DeepEP/NCCL dependency ---
     let k3_enabled = cfg!(feature = "k3");
     let qwen35_enabled = cfg!(feature = "qwen35");
+    let qwen4_enabled = cfg!(feature = "qwen4");
     if glm52_enabled {
         generate_glm52_trtllm_fmha_cubins(&crate_root(), &out_dir);
         build_glm52_cutedsl_fp8_dsl(&crate_root(), &out_dir, &cuda_include);
@@ -2281,6 +2293,9 @@ fn main() {
             if !qwen35_enabled && is_qwen35_source(&csrc_dir, path) {
                 return None;
             }
+            if !qwen4_enabled && is_qwen4_source(&csrc_dir, path) {
+                return None;
+            }
             // --- k3 ---
             if !k3_enabled && is_k3_source(&csrc_dir, path) {
                 return None;
@@ -2717,6 +2732,11 @@ fn main() {
             "cargo:warning=Kimi-K2 CUDA kernels disabled; enable the pegainfer-kernels `kimi-k2` feature to build them"
         );
     }
+    if !qwen4_enabled {
+        println!(
+            "cargo:warning=Qwen4 QSA indexer kernels disabled; enable the pegainfer-kernels `qwen4` feature to build them"
+        );
+    }
 
     // tilelang: BEGIN
     if cfg!(feature = "gemma4") {