diff --git a/docs/models/qwen4/bring-up.md b/docs/models/qwen4/bring-up.md index 1ce5b308a..296dffba8 100644 --- a/docs/models/qwen4/bring-up.md +++ b/docs/models/qwen4/bring-up.md @@ -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). @@ -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 diff --git a/pegainfer-kernels/Cargo.toml b/pegainfer-kernels/Cargo.toml index 4e75d1268..5f852b5c1 100644 --- a/pegainfer-kernels/Cargo.toml +++ b/pegainfer-kernels/Cargo.toml @@ -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"] diff --git a/pegainfer-kernels/csrc/qwen4/qwen4_qsa_indexer_select.cu b/pegainfer-kernels/csrc/qwen4/qwen4_qsa_indexer_select.cu new file mode 100644 index 000000000..bb62b1d5f --- /dev/null +++ b/pegainfer-kernels/csrc/qwen4/qwen4_qsa_indexer_select.cu @@ -0,0 +1,299 @@ +// Qwen3.8-Flash-Next (Qwen4 line) QSA token-indexer block selection. +// +// Turns the per-(row, block) scores of `qwen4_qsa_indexer_score.cu` into the +// token-level visibility mask the reference builds, frozen semantics from +// transformers 5.16.1 `qwen4_exp` (`Qwen4ExpTextQSAIndexer.forward`). For a +// query row at position `p` (0-based, tokens 0..p visible): +// +// blocks = (p + 1) / ratio complete micro-blocks, each 4 tokens +// tail = (p + 1) - 4 * blocks tokens 4 * blocks ..= p +// keep = min(budget / ratio, blocks) +// visible = the tokens of the `keep` highest block scores, ties broken by +// the LOWER block index, plus the tail +// +// The selection order is descending score, ascending block index. `torch.topk` +// leaves the pick at the cut open, so a block tied with the last kept one is an +// equally valid pick; this kernel's canonical tie rule is the one the parity +// gate holds the fixture to. +// +// GRID CONTRACT, inherited from the score pass: the visible list starts at +// token 0, so a row at `p` sees exactly tokens [0, p] and block b covers +// [4b, 4b + 4). A caller whose visible list starts elsewhere (an evicted +// prefix) has to hand both passes a shifted key array; nothing here can detect +// one. +// +// A non-finite score inside `[0, blocks)` or an out-of-range `rows` entry is an +// upstream contract violation and traps. The score pass leaves the columns past +// `blocks` untouched, and this kernel never reads them. +// +// One CTA per query row. The target is the portable subset (no TMA, no cluster, +// no wgmma). A row under the budget keeps every complete block and writes the +// visible prefix in one pass. Otherwise an exact 8-bit radix select over the +// row's sortable f32 keys (four passes) finds the `keep`-th score, and the +// blocks equal to that threshold are marked in ascending block-index order to +// preserve the tie rule, one bit per block in dynamic shared memory sized from +// the widest row this launch can see (`tokens / ratio` blocks). The host +// refuses a launch whose bitmap plus the block's static shared state exceeds +// the default per-block ceiling. +// + +#include "../common.cuh" +#include "../shared/ffi_guard.cuh" + +#include +#include + +#include + +namespace { + +constexpr int kRatio = 4; // indexer_compress_ratio +// The budget's block ceiling: 2048 tokens / 4. `tokens` itself is not bounded +// by this (a row can see far more blocks than the budget keeps). +constexpr int kMaxKeptBlocks = 512; +constexpr int kThreads = 256; +constexpr int kWarps = kThreads / WARP_SIZE; +constexpr int kRadixBins = 256; +// Per-warp radix histograms, the threshold state and the equality prefix, +// allocated statically beside the dynamic bitmap and therefore counted against +// the same per-block ceiling. +constexpr int kStaticSharedBytes = + kWarps * kRadixBins * static_cast(sizeof(unsigned int)) + + 2 * static_cast(sizeof(unsigned int)) + + (kThreads + 1) * static_cast(sizeof(int)); +// The default per-block shared-memory ceiling on every supported arch; the +// bitmap plus the block's static state must fit under it without an opt-in. +constexpr int kMaxSharedBytes = 48 * 1024; + +// f32 -> monotonic u32 for the radix select. -0.0 and +0.0 compare equal in the +// selection order, so fold both zeros to +0.0 before taking the bits. +__device__ __forceinline__ unsigned int sortable_key(float score) { + const unsigned int bits = __float_as_uint(score == 0.0f ? 0.0f : score); + return (bits & 0x80000000u) ? ~bits : (bits ^ 0x80000000u); +} + +// One bit per block, set once a threshold-equal block has been taken; every +// block above the threshold is kept without a mark. +__device__ __forceinline__ bool block_taken(const unsigned int* taken, int block) { + return ((taken[block >> 5] >> (block & 31)) & 1u) != 0u; +} + +__device__ __forceinline__ void mark_taken(unsigned int* taken, int block) { + // The equality pass marks from several threads, and equal keys share a word. + atomicOr(&taken[block >> 5], 1u << (block & 31)); +} + +// Per query row: validate the row, check its scores, select the kept blocks and +// write the row's visibility mask. One CTA per row, the warps split the blocks. +__global__ void qwen4_qsa_select_kernel( + const float* __restrict__ scores, // [query_rows, score_cols] + const int* __restrict__ rows, // [query_rows] positions + unsigned char* __restrict__ mask, // [query_rows, tokens] + int tokens, + int score_cols, + int budget) { + extern __shared__ unsigned int taken[]; + const int row = blockIdx.x; + const int tid = threadIdx.x; + + const int pos = rows[row]; + if (pos < 0 || pos >= tokens) { + __trap(); + } + const int visible = pos + 1; + const int blocks = visible / kRatio; + const int tail = visible - blocks * kRatio; + + const float* row_scores = scores + static_cast(row) * score_cols; + for (int block = tid; block < blocks; block += blockDim.x) { + if (!isfinite(row_scores[block])) { + __trap(); + } + } + + const int keep_cap = budget / kRatio; + const int keep = keep_cap < blocks ? keep_cap : blocks; + unsigned char* row_mask = mask + static_cast(row) * tokens; + + // Under the budget every complete block is kept: one pass writes the whole + // visible prefix and the search is skipped. + if (keep == blocks) { + for (int token = tid; token < tokens; token += blockDim.x) { + row_mask[token] = token < visible ? 1 : 0; + } + return; + } + + // Exact radix select of the keep-th key: four 8-bit passes narrow the + // threshold one key byte at a time. + __shared__ unsigned int hist[kWarps * kRadixBins]; + __shared__ unsigned int radix_prefix; + __shared__ int radix_remaining; + const int warp = tid / WARP_SIZE; + if (tid == 0) { + radix_prefix = 0u; + radix_remaining = keep; + } + __syncthreads(); + for (int shift = 24; shift >= 0; shift -= 8) { + for (int i = tid; i < kWarps * kRadixBins; i += blockDim.x) { + hist[i] = 0u; + } + __syncthreads(); + const unsigned int prefix = radix_prefix; + const unsigned int high_mask = + shift == 24 ? 0u : (0xFFFFFFFFu << (shift + 8)); + for (int block = tid; block < blocks; block += blockDim.x) { + const unsigned int key = sortable_key(row_scores[block]); + if ((key & high_mask) != (prefix & high_mask)) { + continue; + } + atomicAdd(&hist[warp * kRadixBins + ((key >> shift) & 0xFFu)], 1u); + } + __syncthreads(); + if (tid == 0) { + int remaining = radix_remaining; + int chosen = -1; + for (int bin = kRadixBins - 1; bin >= 0; --bin) { + int count = 0; + for (int w = 0; w < kWarps; ++w) { + count += static_cast(hist[w * kRadixBins + bin]); + } + if (count >= remaining) { + chosen = bin; + break; + } + remaining -= count; + } + if (chosen < 0) { + __trap(); // `keep` exceeds the row's block count + } + radix_prefix = prefix | (static_cast(chosen) << shift); + radix_remaining = remaining; + } + __syncthreads(); + } + const unsigned int threshold = radix_prefix; + const int need = radix_remaining; + + // Zero the mask row and the bitmap that marks the threshold-equal blocks. + const int max_blocks = tokens / kRatio; + const int words = (max_blocks + 31) >> 5; + for (int token = tid; token < tokens; token += blockDim.x) { + row_mask[token] = 0; + } + for (int word = tid; word < words; word += blockDim.x) { + taken[word] = 0u; + } + __syncthreads(); + + // Equal keys are ordered by block index, matching the canonical tie rule. + // Contiguous ranges make a per-thread prefix the global order. + __shared__ int equal_offsets[kThreads + 1]; + const int start = static_cast(static_cast(tid) * blocks / blockDim.x); + const int end = + static_cast(static_cast(tid + 1) * blocks / blockDim.x); + int equal_count = 0; + for (int block = start; block < end; ++block) { + equal_count += sortable_key(row_scores[block]) == threshold; + } + equal_offsets[tid] = equal_count; + __syncthreads(); + if (tid == 0) { + int running = 0; + for (int i = 0; i < kThreads; ++i) { + const int count = equal_offsets[i]; + equal_offsets[i] = running; + running += count; + } + equal_offsets[kThreads] = running; + if (running < need) { + __trap(); // the radix threshold must cover the kept rank + } + } + __syncthreads(); + int take = need - equal_offsets[tid]; + for (int block = start; block < end && take > 0; ++block) { + if (sortable_key(row_scores[block]) == threshold) { + mark_taken(taken, block); + --take; + } + } + __syncthreads(); + + // Every block above the threshold is kept, the chosen equals fill the rest, + // and the tail is unconditional. + for (int block = tid; block < blocks; block += blockDim.x) { + const unsigned int key = sortable_key(row_scores[block]); + const bool kept = + key > threshold || (key == threshold && block_taken(taken, block)); + if (!kept) { + continue; + } + const int base = block * kRatio; +#pragma unroll + for (int i = 0; i < kRatio; ++i) { + row_mask[base + i] = 1; + } + } + for (int i = tid; i < tail; i += blockDim.x) { + row_mask[blocks * kRatio + i] = 1; + } +} + +CUresult consume_last_cuda_error() { return map_cuda_error(cudaGetLastError()); } + +} // namespace + +extern "C" { + +// Select the kept blocks of every query row and write its token visibility +// mask. Row `j` covers tokens `[0, rows[j]]`; `mask` is one byte per (row, +// token), 0 or 1, with row stride `tokens`. `scores` is one f32 row of +// `score_cols` per query row, of which `(rows[j] + 1) / kRatio` are read. +// +// `rows` carries the query rows' positions in the token order (also the row +// order of `scores` and `mask`); it is not assumed to be every position. +// `budget` is the token budget, a positive multiple of kRatio whose block count +// is at most kMaxKeptBlocks. +// +// The compiled geometry is kRatio = 4 and one byte per mask cell; a caller +// whose buffers are laid out differently cannot be detected here. An empty +// launch is refused; `qsa_indexer_select_into` treats one as a no-op before it +// reaches this entry. +CUresult qwen4_qsa_indexer_select_cuda( + const float* scores, + const int* rows, + unsigned char* mask, + int query_rows, + int tokens, + int score_cols, + int budget, + cudaStream_t stream) { + if (scores == nullptr || rows == nullptr || mask == nullptr) { + return CUDA_ERROR_INVALID_VALUE; + } + if (query_rows <= 0 || tokens <= 0 || score_cols < 0) { + return CUDA_ERROR_INVALID_VALUE; + } + if (budget <= 0 || budget % kRatio != 0 || budget / kRatio > kMaxKeptBlocks) { + return CUDA_ERROR_INVALID_VALUE; + } + // Every row reads at most `tokens / kRatio` blocks from its score row. + if (tokens / kRatio > score_cols) { + return CUDA_ERROR_INVALID_VALUE; + } + const int max_blocks = tokens / kRatio; + const int words = (max_blocks + 31) >> 5; + const int64_t shared_bytes = + static_cast(words) * static_cast(sizeof(unsigned int)); + if (shared_bytes + kStaticSharedBytes > kMaxSharedBytes) { + return CUDA_ERROR_INVALID_VALUE; + } + qwen4_qsa_select_kernel<<(shared_bytes), stream>>>( + scores, rows, mask, tokens, score_cols, budget); + return consume_last_cuda_error(); +} + +} // extern "C" diff --git a/pegainfer-kernels/src/ffi/qwen4.rs b/pegainfer-kernels/src/ffi/qwen4.rs index 623b78ebc..a881d7d60 100644 --- a/pegainfer-kernels/src/ffi/qwen4.rs +++ b/pegainfer-kernels/src/ffi/qwen4.rs @@ -43,4 +43,24 @@ unsafe extern "C" { eps: f32, stream: CUstream, ) -> CUresult; + + /// Select the kept micro-blocks of every query row plus its incomplete + /// tail, and write the row's token visibility mask (one byte per cell, + /// `0` or `1`). + /// + /// `rows[j]` is the row's position, `scores` its `score_cols`-wide f32 + /// block-score row and `mask` its `tokens`-wide byte row. `budget` is the + /// token budget. The compiled geometry is four tokens per block and one + /// byte per mask cell; a caller whose buffers are laid out differently + /// cannot be detected here. + pub fn qwen4_qsa_indexer_select_cuda( + scores: *const f32, + rows: *const i32, + mask: *mut u8, + query_rows: i32, + tokens: i32, + score_cols: i32, + budget: i32, + stream: CUstream, + ) -> CUresult; } diff --git a/pegainfer-kernels/src/ops.rs b/pegainfer-kernels/src/ops.rs index 4d63fffee..cea47982c 100644 --- a/pegainfer-kernels/src/ops.rs +++ b/pegainfer-kernels/src/ops.rs @@ -240,6 +240,8 @@ pub use norm::rms_norm_rows_into; #[cfg(feature = "qwen4")] pub use qwen4::qsa_indexer_block_scores_into; #[cfg(feature = "qwen4")] +pub use qwen4::qsa_indexer_select_into; +#[cfg(feature = "qwen4")] pub use qwen4::rms_norm_grouped_offset_batch_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 2efbbe9c2..f16dc300c 100644 --- a/pegainfer-kernels/src/ops/qwen4.rs +++ b/pegainfer-kernels/src/ops/qwen4.rs @@ -1,13 +1,15 @@ //! Qwen3.8-Flash-Next (Qwen4 line) kernels. //! -//! QSA token-indexer block scoring turns the fused `index_qk_proj` output into -//! one score per (query row, visible micro-block); selection — top-k plus the -//! unconditional tail — is a later stage, deliberately not here. Frozen -//! semantics live in `csrc/qwen4/qwen4_qsa_indexer_score.cu` (transformers -//! 5.16.1 `qwen4_exp` `Qwen4ExpTextQSAIndexer.forward`); the kernel is compiled -//! for one geometry 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. +//! [`qsa_indexer_block_scores_into`] turns the fused `index_qk_proj` output +//! into one score per (query row, visible micro-block); +//! [`qsa_indexer_select_into`] turns those scores into the row's token-level +//! visibility mask (top-k plus the unconditional tail). Frozen semantics live +//! in `csrc/qwen4/qwen4_qsa_indexer_score.cu` and +//! `csrc/qwen4/qwen4_qsa_indexer_select.cu` (transformers 5.16.1 `qwen4_exp` +//! `Qwen4ExpTextQSAIndexer.forward`); the kernels are compiled for one geometry +//! and these wrappers size their buffers from it. The checkpoint's own indexer +//! fields are checked where the model crate reads its config, once that caller +//! lands — the score stage takes only `rms_eps` from the caller. //! //! The grouped `(1+weight)` RMSNorm serves the hyper-connection and PLE streams //! (`csrc/qwen4/qwen4_grouped_rmsnorm.cu`); its contract is on the wrapper. @@ -231,3 +233,81 @@ pub fn rms_norm_grouped_offset_batch_into( result.result()?; Ok(()) } + +/// The budget's block ceiling: the checkpoint's 2048-token budget at four +/// tokens per block. +const MAX_KEPT_BLOCKS: usize = 512; + +/// Turn one score per (query row, complete block) into the row's token-level +/// visibility mask (semantics: `csrc/qwen4/qwen4_qsa_indexer_select.cu`). +/// +/// `scores` is `[rows.len(), score_cols]` f32 with row stride `score_cols`; the +/// buffer may hold spare capacity (`scores.len() >= rows.len() * score_cols`). +/// Only `(rows[j] + 1) / 4` entries of row `j` are read. `mask` is one byte per +/// (row, token), row stride `tokens`. `budget` is the token budget, a positive +/// multiple of four keeping at most 512 blocks, and the score row must hold +/// `tokens / 4` blocks. +/// +/// An empty launch (`rows.is_empty()` or `tokens == 0`) is a no-op. A malformed +/// budget and a short buffer are refused before any pointer is formed. +pub fn qsa_indexer_select_into( + ctx: &DeviceContext, + scores: &CudaSlice, + score_cols: usize, + rows: &CudaSlice, + budget: usize, + tokens: usize, + mask: &mut CudaSlice, +) -> Result<()> { + if rows.is_empty() || tokens == 0 { + return Ok(()); + } + ensure!( + budget > 0 + && budget.is_multiple_of(COMPRESS_RATIO) + && budget / COMPRESS_RATIO <= MAX_KEPT_BLOCKS, + "QSA indexer budget {budget} must be a positive multiple of {COMPRESS_RATIO} \ + keeping at most {MAX_KEPT_BLOCKS} blocks" + ); + ensure!( + scores.len() >= rows.len().saturating_mul(score_cols), + "QSA indexer score buffer {} f32 cannot hold {} rows of {score_cols}", + scores.len(), + rows.len() + ); + ensure!( + score_cols >= tokens / COMPRESS_RATIO, + "QSA indexer score row holds {score_cols} blocks, but a mask over \ + {tokens} tokens needs {}", + tokens / COMPRESS_RATIO + ); + ensure!( + mask.len() >= rows.len().saturating_mul(tokens), + "QSA indexer mask holds {} bytes, needs {} for {} x {tokens}", + mask.len(), + rows.len().saturating_mul(tokens), + rows.len() + ); + // The kernel takes extents as i32; checked before any pointer is formed. + let query_rows = checked_i32(rows.len(), "QSA indexer query rows")?; + let tokens = checked_i32(tokens, "QSA indexer tokens")?; + let score_cols = checked_i32(score_cols, "QSA indexer score columns")?; + let budget = checked_i32(budget, "QSA indexer budget")?; + let (scores_ptr, _scores_guard) = scores.device_ptr(&ctx.stream); + let (rows_ptr, _rows_guard) = rows.device_ptr(&ctx.stream); + let (mask_ptr, _mask_guard) = mask.device_ptr_mut(&ctx.stream); + unsafe { + ffi::qwen4_qsa_indexer_select_cuda( + scores_ptr as *const f32, + rows_ptr as *const i32, + mask_ptr as *mut u8, + query_rows, + tokens, + score_cols, + budget, + active_cu_stream(ctx), + ) + } + .result() + .map_err(|err| anyhow::anyhow!("QSA indexer select launch failed: {err}")) +} diff --git a/pegainfer-kernels/tests/qwen4_qsa_indexer_score_parity.rs b/pegainfer-kernels/tests/qwen4_qsa_indexer_score_parity.rs index fe81a9212..de68584df 100644 --- a/pegainfer-kernels/tests/qwen4_qsa_indexer_score_parity.rs +++ b/pegainfer-kernels/tests/qwen4_qsa_indexer_score_parity.rs @@ -8,7 +8,8 @@ //! scores, so this gate replays those inputs through the kernel and compares. //! Cases cover the tail (5), a block-aligned short input (8), a medium one //! (64), the truncation branch under a shrunk budget (14, 16, 20) and the real -//! budget boundary at 2049..2054 visible tokens. +//! budget boundary at 2049..2054 visible tokens. The `select` module replays +//! the same fixture's recorded block scores through the selection kernel. //! //! Provenance is asserted, not assumed: the fixture records the reference //! version and the sha256 of the frozen config it was dumped from. The gate @@ -751,3 +752,485 @@ fn a_poisoned_key_keeps_its_nan_score() { assert!(saw_nan > 0, "no row saw the poisoned block"); println!("poisoned key: block 0 is NaN on {saw_nan} rows, later blocks finite"); } + +mod select { + //! `test_data/qwen38fn-qsa-indexer-golden.safetensors` (with the scoring gate's + //! fixture) stores, per case, the reference's per-row block scores and the final + //! `[rows, tokens]` visibility mask. This gate replays those recorded scores + //! through the selection kernel and compares the mask. Selection is + //! comparison-only — no arithmetic touches a score — so there is no tolerance + //! here: on a row whose cut is strict the mask must match **bitwise**, and the + //! only rows held to less are the ones where `torch.topk` left a block tied at + //! the cut (see below). + //! + //! Ties: for each row the gate computes the cut score; if no other block shares + //! it the row is compared bitwise, and otherwise both masks are required to be + //! structurally valid (a causal prefix, only complete four-token blocks plus the + //! exact tail, exactly `min(budget / 4, blocks) * 4 + tail` bits set), to + //! select only scores at least as high as every block they exclude, and to keep + //! the same number of tied-at-cut blocks. **Ungated boundary:** every tie here + //! is at zero; the fixture has no row tied at a non-zero cut, so that case is + //! not gated. + + // Exact float comparison is the point: a tied cut is defined by score-equal + // blocks, so a margin here would erase the distinction the gate exists for. + #![allow(clippy::float_cmp)] + + use std::cmp::Ordering; + + use pegainfer_kernels::ops::qsa_indexer_select_into; + use pegainfer_kernels::tensor::DeviceContext; + + use super::*; + + /// The four tokens per micro-block the kernel compiles in. The config's + /// `indexer_compress_ratio` must equal it; a drift is caught in `load_case` + /// rather than silently selecting blocks of another width. + const RATIO: usize = 4; + /// Written over the mask buffer before every launch, to prove the kernel writes + /// every cell of every row. + const SENTINEL: u8 = 0xAA; + + /// The reference's mask is `torch.bool`, one byte per cell. + fn mask_tensor(fixture: &SafeTensors<'_>, name: &str) -> (Vec, Vec) { + let view = fixture + .tensor(name) + .unwrap_or_else(|e| panic!("{name}: {e}")); + assert_eq!(view.dtype(), safetensors::Dtype::BOOL, "{name} dtype"); + (view.shape().to_vec(), view.data().to_vec()) + } + + fn budget_tensor(fixture: &SafeTensors<'_>, name: &str) -> i32 { + let (shape, host) = i32_tensor(fixture, name); + assert_eq!(shape, vec![1], "{name} shape"); + host[0] + } + + /// One case's recorded scores, rows, budget and the reference's mask. + struct Case { + index: usize, + tokens: usize, + budget: i32, + rows: Vec, + score_cols: usize, + /// `[rows, score_cols]`, the NaN padding included. + scores: Vec, + /// `[rows, tokens]`, one byte per cell. + golden: Vec, + } + + fn load_case(fixture: &SafeTensors<'_>, index: usize, scenario: &Scenario) -> Case { + let (rows_shape, rows) = i32_tensor(fixture, &format!("case{index}_rows")); + let (score_shape, scores) = f32_tensor(fixture, &format!("case{index}_scores")); + let (mask_shape, golden) = mask_tensor(fixture, &format!("case{index}_mask")); + let budget = budget_tensor(fixture, &format!("case{index}_budget")); + + assert_eq!(rows_shape, vec![rows.len()], "case {index} rows shape"); + assert_eq!(score_shape.len(), 2, "case {index} score rank"); + assert_eq!(mask_shape.len(), 2, "case {index} mask rank"); + assert_eq!(score_shape[0], rows.len(), "case {index} score rows"); + assert_eq!(mask_shape[0], rows.len(), "case {index} mask rows"); + let tokens = mask_shape[1]; + assert_eq!(golden.len(), rows.len() * tokens, "case {index} mask cells"); + assert!(tokens >= 1, "case {index} has no tokens"); + let score_cols = score_shape[1]; + + // The config's ratio and the compiled one must agree before any shape is + // read through them: the mask's block structure is the reference's, and the + // kernel selects blocks of exactly RATIO tokens. + assert_eq!( + scenario.ratio, RATIO, + "config indexer_compress_ratio {} is not the compiled {RATIO}", + scenario.ratio + ); + assert!( + budget > 0 && (budget as usize).is_multiple_of(scenario.ratio), + "case {index} budget {budget} is not a positive multiple of the ratio" + ); + let budget_blocks = budget as usize / scenario.ratio; + assert!( + budget_blocks <= scenario.max_budget_blocks, + "case {index} budget keeps {budget_blocks} blocks, past the {} the config ceiling allows", + scenario.max_budget_blocks + ); + // One row can see at most `tokens / ratio` blocks; the score row must hold + // them, and every cell past the row's own prefix is the reference's NaN. + let block_ceiling = tokens / scenario.ratio; + assert!( + score_cols >= block_ceiling, + "case {index}: {score_cols} score columns cannot hold {block_ceiling} blocks" + ); + for (slot, row) in rows.iter().enumerate() { + assert!( + *row >= 0 && (*row as usize) < tokens, + "case {index} row {row} is outside the {tokens} tokens" + ); + let blocks = (*row as usize + 1) / scenario.ratio; + let row_scores = &scores[slot * score_cols..][..score_cols]; + for (block, score) in row_scores.iter().take(blocks).enumerate() { + assert!( + score.is_finite(), + "case {index} row {row} block {block} is not finite: {score}" + ); + } + for (block, score) in row_scores.iter().enumerate().skip(blocks) { + assert!( + score.is_nan(), + "case {index} row {row} column {block} is not the reference's padding" + ); + } + } + Case { + index, + tokens, + budget, + rows, + score_cols, + scores, + golden, + } + } + + /// What the frozen config says about this gate's geometry. + struct Scenario { + ratio: usize, + /// The config's own budget, in blocks; the fixture must exercise it. + default_budget: i32, + /// The compiled budget ceiling in blocks. + max_budget_blocks: usize, + } + + impl Scenario { + fn from_config(config: &serde_json::Value) -> Scenario { + let text = &config["text_config"]; + let ratio = text["indexer_compress_ratio"] + .as_u64() + .expect("indexer_compress_ratio") as usize; + let default_budget = text["indexer_budget"].as_u64().expect("indexer_budget") as i32; + // The selection kernel refuses a budget past 512 blocks; the config's + // own budget must stay inside that. + let max_budget_blocks = 512; + assert!(ratio > 0, "indexer_compress_ratio"); + assert!(default_budget > 0, "indexer_budget"); + assert!(default_budget as usize / ratio <= max_budget_blocks); + Scenario { + ratio, + default_budget, + max_budget_blocks, + } + } + + fn blocks_of(&self, budget: i32) -> usize { + budget as usize / self.ratio + } + } + + impl Case { + fn row_scores(&self, slot: usize) -> &[f32] { + &self.scores[slot * self.score_cols..][..self.score_cols] + } + + fn blocks(&self, slot: usize) -> usize { + (self.rows[slot] as usize + 1) / RATIO + } + + fn tail(&self, slot: usize) -> usize { + (self.rows[slot] as usize + 1) % RATIO + } + + fn keep(&self, slot: usize, budget_blocks: usize) -> usize { + budget_blocks.min(self.blocks(slot)) + } + + fn golden_row(&self, slot: usize) -> &[u8] { + &self.golden[slot * self.tokens..][..self.tokens] + } + } + + /// The canonical order over a score row: descending score, ascending block + /// index, the kernel's tie rule. + fn canonical_order(scores: &[f32]) -> Vec { + let mut order: Vec = (0..scores.len()).collect(); + order.sort_by(|&a, &b| { + let by_score = if scores[a] == scores[b] { + Ordering::Equal + } else { + scores[b].partial_cmp(&scores[a]).expect("finite scores") + }; + by_score.then(a.cmp(&b)) + }); + order + } + + /// A host-built mask for one row: the tokens of `selected` blocks plus the + /// first `tail_bits` tokens of the row's incomplete block, which starts at + /// `blocks * RATIO`. + fn host_mask(tokens: usize, blocks: usize, selected: &[usize], tail_bits: usize) -> Vec { + assert!(blocks * RATIO + tail_bits <= tokens, "row mask bounds"); + let mut mask = vec![0u8; tokens]; + for &block in selected { + mask[block * RATIO..][..RATIO].fill(1); + } + mask[blocks * RATIO..][..tail_bits].fill(1); + mask + } + + fn differing(a: &[u8], b: &[u8]) -> Vec { + assert_eq!(a.len(), b.len(), "compared rows differ in length"); + (0..a.len()).filter(|&i| a[i] != b[i]).collect() + } + + fn assert_same_row(label: &str, got: &[u8], want: &[u8]) { + if let Some(at) = differing(got, want).first() { + panic!( + "{label}: first difference at token {at}: got {}, want {}. \ + Full difference: {:?}", + got[*at], + want[*at], + differing(got, want) + ); + } + } + + /// The structure every valid row mask has, returning its selected blocks: only + /// `0` or `1`, complete four-token blocks below the row's own token, nothing + /// past it, the exact tail set, and exactly `selected * 4 + tail` bits. + fn structure( + label: &str, + row_mask: &[u8], + blocks: usize, + tail: usize, + pos: usize, + ) -> Vec { + assert!( + row_mask.iter().all(|cell| *cell <= 1), + "{label}: the mask holds a cell other than 0 or 1" + ); + let mut selected = Vec::new(); + for block in 0..blocks { + let bits = &row_mask[block * RATIO..][..RATIO]; + if bits.iter().all(|cell| *cell == 1) { + selected.push(block); + } else { + assert!( + bits.iter().all(|cell| *cell == 0), + "{label}: block {block} is only partly set: {bits:?}" + ); + } + } + for (i, cell) in row_mask[blocks * RATIO..=pos].iter().enumerate() { + assert_eq!( + *cell, 1, + "{label}: tail token {i} of {tail} must be visible" + ); + } + for (i, cell) in row_mask[pos + 1..].iter().enumerate() { + assert_eq!( + *cell, + 0, + "{label}: token {} is past the row's own position and must not be visible", + pos + 1 + i + ); + } + let set: usize = row_mask.iter().map(|cell| *cell as usize).sum(); + assert_eq!( + set, + selected.len() * RATIO + tail, + "{label}: {set} set bits is not the selected blocks' tokens plus the tail" + ); + selected + } + + /// Runs one case's mask through the production launch and returns the + /// downloaded cells. + fn run_case(ctx: &DeviceContext, case: &Case) -> Vec { + let scores = ctx.stream.clone_htod(&case.scores).expect("score upload"); + let rows = ctx.stream.clone_htod(&case.rows).expect("rows upload"); + let mut mask = ctx + .stream + .clone_htod(&vec![SENTINEL; case.rows.len() * case.tokens]) + .expect("mask buffer"); + qsa_indexer_select_into( + ctx, + &scores, + case.score_cols, + &rows, + case.budget as usize, + case.tokens, + &mut mask, + ) + .unwrap_or_else(|e| panic!("case {}: launch: {e:#}", case.index)); + ctx.sync().expect("sync"); + let host = ctx.stream.clone_dtoh(&mask).expect("mask download"); + assert!( + host.iter().all(|cell| *cell <= 1), + "case {}: left a cell unwritten or wrote past 1", + case.index + ); + host + } + + /// One case's measured comparison. + #[derive(Default)] + struct CaseReport { + rows: usize, + bitwise: usize, + tie: usize, + truncating: usize, + /// Rows that keep every block and take the all-kept fast path. + all_kept: usize, + } + + fn check_case(ctx: &DeviceContext, case: &Case, scenario: &Scenario) -> CaseReport { + let budget_blocks = scenario.blocks_of(case.budget); + let rows = case.rows.len(); + let tokens = case.tokens; + + let production = run_case(ctx, case); + + let mut report = CaseReport { + rows, + ..CaseReport::default() + }; + for slot in 0..rows { + let row = case.rows[slot]; + let pos = row as usize; + let blocks = case.blocks(slot); + let tail = case.tail(slot); + let keep = case.keep(slot, budget_blocks); + let scores = &case.row_scores(slot)[..blocks]; + let canonical = canonical_order(scores); + let label = format!("case {} row {row}", case.index); + + // The kernel and the host's canonical order must agree on every row. + // On a tie row this still pins the kernel (same total order), while the + // golden is held to validity instead of to its open pick. + let ours = &production[slot * tokens..][..tokens]; + assert_same_row( + &label, + ours, + &host_mask(tokens, blocks, &canonical[..keep], tail), + ); + + let golden = case.golden_row(slot); + let cut = if blocks > keep { + let mut ranked = scores.to_vec(); + ranked.sort_by(|a, b| b.partial_cmp(a).expect("finite scores")); + Some(ranked[keep - 1]) + } else { + None + }; + let tie = cut.is_some_and(|cut| { + let at_cut = scores.iter().filter(|score| **score == cut).count(); + let kept_at_cut = keep - scores.iter().filter(|score| **score > cut).count(); + at_cut > kept_at_cut + }); + if tie { + report.tie += 1; + let cut = cut.expect("a tie needs a cut"); + let at_cut = scores.iter().filter(|score| **score == cut).count(); + println!( + "{label}: tie at cut {cut}, {at_cut} blocks tied, {keep} kept; \ + compared up to the tie" + ); + let ours_selected = structure(&format!("{label} (ours)"), ours, blocks, tail, pos); + let golden_selected = + structure(&format!("{label} (golden)"), golden, blocks, tail, pos); + assert_eq!(ours_selected.len(), keep, "{label}: kept blocks"); + assert_eq!(golden_selected.len(), keep, "{label}: golden kept blocks"); + // Every kept score at least every excluded one, on both masks. + for (name, selected) in [("ours", &ours_selected), ("golden", &golden_selected)] { + let lowest = selected + .iter() + .map(|&block| scores[block]) + .fold(f32::INFINITY, f32::min); + let highest_excluded = (0..blocks) + .filter(|block| !selected.contains(block)) + .map(|block| scores[block]) + .fold(f32::NEG_INFINITY, f32::max); + assert!( + lowest >= highest_excluded, + "{label}: {name} keeps a block scoring {lowest} below the excluded \ + {highest_excluded}" + ); + } + let ours_at_cut = ours_selected + .iter() + .filter(|&&block| scores[block] == cut) + .count(); + let golden_at_cut = golden_selected + .iter() + .filter(|&&block| scores[block] == cut) + .count(); + assert_eq!( + ours_at_cut, golden_at_cut, + "{label}: tied-at-cut picks ({cut}) differ between the kernel and the reference" + ); + } else { + report.bitwise += 1; + assert_same_row(&label, ours, golden); + } + if blocks > keep { + report.truncating += 1; + } else { + report.all_kept += 1; + } + } + report + } + + #[test] + fn qsa_indexer_select_matches_the_reference() { + let Some(ctx) = common::device_or_skip() else { + return; + }; + let bytes = std::fs::read(fixture_path()).expect("read fixture"); + let config = assert_provenance(&bytes); + let fixture = SafeTensors::deserialize(&bytes).expect("parse fixture"); + let scenario = Scenario::from_config(&config); + + let mut total = CaseReport::default(); + let mut budgets = Vec::new(); + for index in 0..CASES { + let case = load_case(&fixture, index, &scenario); + let report = check_case(&ctx, &case, &scenario); + budgets.push(case.budget); + println!( + "case {index}: {} rows, {} bitwise, {} tie, {} truncating, {} all-kept; \ + tokens {}, budget {} ({} blocks)", + report.rows, + report.bitwise, + report.tie, + report.truncating, + report.all_kept, + case.tokens, + case.budget, + scenario.blocks_of(case.budget), + ); + total.rows += report.rows; + total.bitwise += report.bitwise; + total.tie += report.tie; + total.truncating += report.truncating; + total.all_kept += report.all_kept; + } + println!( + "selection parity over {CASES} cases: {} rows, {} bitwise, {} up to ties, \ + {} truncating, {} all-kept", + total.rows, total.bitwise, total.tie, total.truncating, total.all_kept, + ); + + assert!( + budgets.contains(&scenario.default_budget), + "the config's own budget {} is not exercised", + scenario.default_budget + ); + assert!( + total.truncating > 0, + "the fixture must exercise the truncating path" + ); + assert!( + total.all_kept > 0, + "the fixture must exercise the all-kept fast path" + ); + } +}