diff --git a/docs/index.md b/docs/index.md index cbe3d91a5..f8a832c03 100644 --- a/docs/index.md +++ b/docs/index.md @@ -191,6 +191,7 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | `subsystems/kernels/pegainfer-bench-hawk-audit.md` | Hawk retained `pegainfer-bench` on two report consumers, reduced workspace findings 576 → 361, and exposed dead state that was deleted; 332 remaining findings are in the excluded kvbm fork. | | `subsystems/kernels/build-rs-submodule-init.md` | `pegainfer-kernels/build.rs` initializes missing git submodules automatically for first-time builds before checking vendored third-party kernel headers. | | `subsystems/kernels/kernel-op-reports.md` | Qwen3 kernel/report tooling is feature-gated: `qwen3_kernel_report` covers per-op kernel reports, and `qwen3_model_report` emits runtime-traced eager-DAG decode operator rollups with TensorSpec `KernelCall`s, latency stats, tables, and Graphviz DOT; measured FA2 `CTA_TILE_Q=64` prefill default in place. | +| `subsystems/kernels/qk-rope-smoke-oracles.md` | The hd256/hd512 QK-RoPE device gates keep three closed-form numerical anchors and derive every other expectation from a GPU arm an anchor already certifies (anchor → flat kernel → paged prefill → paged decode); no host-side second implementation of RMSNorm/RoPE/paged addressing. Records the gate command, the `PEGAINFER_REQUIRE_GPU=1` skip trap, and the negative controls. | | `subsystems/kernels/typed-forward-pipeline.md` | Reusable typed tensor pipeline macro in `pegainfer-kernels` so model crates can express common `typed_ops` chains without model-specific wrapper macros. | | `subsystems/kernels/tvm-ffi-mvp.md` | Optional `tvm-ffi-triton-cubin` bridge in `pegainfer-kernels` plus a packed TVM wrapper for the Qwen3.5 GDR solve Triton AOT CUBIN launcher. | diff --git a/docs/subsystems/kernels/qk-rope-smoke-oracles.md b/docs/subsystems/kernels/qk-rope-smoke-oracles.md new file mode 100644 index 000000000..dadc8a15a --- /dev/null +++ b/docs/subsystems/kernels/qk-rope-smoke-oracles.md @@ -0,0 +1,176 @@ +# HD256 / HD512 QK-RoPE Smoke Oracles + +**Created**: 2026-08-27 +**Last touched**: 2026-08 +**TL;DR**: The hd256/hd512 QK-RoPE device gates keep exactly three closed-form numerical anchors +(hd256 full-rotation and partial-tail, hd512 partial-rope) and derive every other expectation from +a GPU arm that an anchor already certifies. The chain is anchor → flat kernel → paged prefill → +paged decode: each link catches a failure the link above it cannot, and no link re-implements +RMSNorm, RoPE, or the paged address formula a second time. Adding a fourth dispatch variant means +extending the chain, not writing a new host oracle. + +## Why the chain exists + +Before #943 the two gate files carried seven closed-form tests and two independent copies of +`inv_rms`, `normed`, `expected_prep`, `expected_full`, `expected_pool`, and `assert_pool`. That is a +second implementation of the CUDA operator's substance, living in Rust, that has to evolve in +lockstep with the `.cu` contract. It fails in two ways that a passing test cannot show you: + +- The host oracle drifts from the kernel and the gate goes red for a reason that is not a kernel bug. +- A specification mistake is made once and copied into both the kernel and its local oracle, so both + are wrong and the gate is green. + +The fix is not a shared oracle module — that relocates the second implementation instead of removing +it. The fix is to stop producing expected values on the host wherever a GPU arm that is already +certified can produce them. + +## The chain + +``` + closed-form anchor (host) the only retained RMSNorm/RoPE formulas + │ certifies + ▼ + flat kernel hd256: qk_norm_rope_prefill_hd256_plain + (contiguous, no pool) hd512: qk_norm_partial_rope_batched_decode_hd512 + │ is the value oracle for + ▼ + paged prefill the only retained address derivation + (values from flat, addresses from + PagedKvLayout, every other slot 0) + │ is the bitwise oracle for + ▼ + paged decode zero host math, zero addressing + (equivalent metadata ⇒ bit-identical) +``` + +Each link is load-bearing against a different defect: + +| Defect | Caught by | +| --- | --- | +| RMSNorm or RoPE pairing/tail formula wrong | the closed-form anchor | +| Paged address formula wrong (right value, wrong slot) | paged prefill vs flat, plus the exact-zero sentinels | +| Stray pool write outside the request's positions | the exact-zero sentinels on unreferenced pages and the other layer | +| Per-token metadata routing wrong (CSR window, per-token origin, `positions[]`) | paged decode vs paged prefill, bit for bit | + +What makes the chain sound rather than circular: the flat kernel and the paged template share +`rms_norm_elem_*` and `apply_rope_pair_*` as device functions, so a bug in the shared math surfaces +at the anchor, which tests the flat kernel directly. The lower links then only have to answer +"did the certified value reach the right place, for the right row". + +## Retained coverage + +Three closed-form anchors, and nothing else on the host reconstructs operator math. + +| Anchor | Contract it pins | +| --- | --- | +| `hd256 full_rotation_matches_closed_form` | `rotary_dim == HD`: the Gemma 4 local-layer production config, where the pass-through tail is empty and no thread takes the `d >= rotary_dim` branch | +| `hd256 partial_rotation_exercises_tail` | `rotary_dim < HD`: the only config that exercises the pass-through tail | +| `hd512 decode_prep_matches_closed_form` | hd512 rotary_dim is always 128 < 512, so one run covers rope-lo, rope-hi and tail together — a second hd512 anchor would be redundant | + +Two host expressions survive outside the anchors, both one line, both the weightless V norm +(`x * inv_rms(x)`): hd256's V band reduces a separate `v_batch` input and hd512's V is the K=V fork, +so neither flat kernel emits a V the paged arm could be compared against. One multiply is not a +second implementation, and the alternative — driving a flat kernel with all-ones norm weights and an +identity RoPE row to synthesise a V oracle on device — couples the test to the period-3 cos/sin table +being unchanged, which is a worse trade for one multiply. + +## Why the pool assertions are bitwise + +`assert_pool_bits` compares bf16 bit patterns, not floats within a tolerance. Every non-zero +expectation is either a value another GPU arm produced from the same input at the same position — so +the two cannot differ at all — or the one-line V norm, which is reproducible exactly because three +things hold at once: + +- `csrc/shared/*.cu` compiles without `--use_fast_math`. That flag is scoped to the vendored + `k3_flash_kda` unit in `build.rs`, so `1.0f / sqrtf(x)` here is the correctly-rounded IEEE + division and square root the host also computes, not the `rsqrt` approximation. +- Test inputs are constant across each head vector, so summing `HD` copies of `x²` in the kernel's + reduction tree is exact in f32 and `total / HD` lands on `x²` with no rounding. +- Neither `x * inv_rms` nor `total / HD + eps` gives nvcc a multiply-add to contract into an FMA. + +A pool failure that is off by one ulp therefore means one of those three stopped holding — a new +nvcc flag, a non-constant test input, a restructured reduction — not that the tolerance was too +tight. Loosening the assertion is the wrong fix; find which one changed. + +The refusal gates (`rejects_bad_rotary_dim`, `rejects_position_beyond_cos_table`, +`prefill_rejects_undersized_kv_pool`) and the row-offset metamorphic gates +(`decode_prep_row_offset_serves_only_the_suffix`, `prefill_prep_row_offset_serves_only_the_suffix`, +`split_read_row_offset_serves_only_the_suffix`) carry no operator math and are untouched by this +discipline. Do not add operator math to them. + +## What the model-level gate does and does not cover + +`pegainfer-gemma4/src/layer_oracle.rs` replays the HF golden fixture's layer probes through the real +implementation on the real checkpoint, covering both layer types — hd256 sliding and hd512 global. +It is genuine external corroboration for the production configuration, and it is **not** a substitute +for the anchors: it is `#[ignore]`d, needs the pinned 12B checkpoint via `PEGAINFER_TEST_MODEL_PATH` +plus a GPU, and never runs by default. It also only ever sees the production `rotary_dim`, so the +partial-tail contract has no coverage there at all. + +Retiring an anchor in favour of that gate would move a routinely-runnable check onto one that nobody +executes. See `docs/conventions/migration-defense.md`: a deleted defence needs a named successor that +actually runs. + +## Running the gates + +CI compiles these targets but never runs them — they need a device. Both test targets are +auto-discovered with no `required-features`, so the default feature set is enough: + +```bash +PEGAINFER_REQUIRE_GPU=1 cargo test --release -p pegainfer-kernels --test hd256_qk_rope_plain_smoke --test hd512_qk_rope_smoke -- --nocapture +``` + +`PEGAINFER_REQUIRE_GPU=1` is not optional for a formal gate. Without it, +`tests/common/mod.rs::device_or_skip` treats a missing device as a skip and the whole suite passes +green without executing anything. + +The trap binaries (`hd512_qk_rope_trap`, `hd512_qk_rope_trap_page`, `hd256_decode_csr_trap`) stay in +their own targets on purpose: `__trap()` poisons the CUDA context for whatever runs next in the same +binary. + +## Negative controls + +A gate built out of GPU-vs-GPU comparisons can pass because both arms are equally broken, so changing +one of these gates means re-earning confidence in it. Perturb the kernel, confirm red, revert. These +three were run on an A40 (sm_86, CUDA 12.8) against the gates as they stand: + +| # | Perturbation | Turns red | How it reports | +| --- | --- | --- | --- | +| 1 | `+ 1` on the `kv_head * HD` term in `paged_kv_offset_hd256_plain` / `paged_kv_offset_hd512` | both `*_lands_*_at_layout_addresses` | `pool (a zero expectation is a slot the kernel must not touch)[59392]: got 0, expected 129` (hd256), `[14848]: got 0, expected 65` (hd512) | +| 2 | `page_origins[token]` → `0` in the `PER_TOKEN_META` branch | both `paged_decode_equals_paged_prefill_*` | `CUDA_ERROR_LAUNCH_FAILED` at the first D2H | +| 2b | `csr_page_row_checked(..., token, ...)` → `..., 0, ...` | both `paged_decode_equals_paged_prefill_*` | `per-token pool writes vs the whole-window run[57344]: got -129, expected 0` (hd256), `[14336]: got 1, expected 0` (hd512) | + +Control 2 and 2b are both needed, and 2b is the one that matters. With the origin ignored, the last +row computes a row index past the end of its own window, so the kernel's own `__trap()` fires and the +launch dies before anything is compared — red, but it is the device guard doing the work, not the +assertion. Pinning every row to row 0's window instead stays in bounds, no guard fires, and the only +thing that can catch it is the bitwise pool comparison. A control that only ever trips a `__trap()` +proves the kernel guards itself, not that the gate would notice a wrong-but-in-bounds page. + +Control 2b also shows why the per-token arm compresses each row's window to that row's own page, +giving the rows distinct non-zero origins (0, 1, 1, 2) over four different table spans. Handing every +row the full table at origin 0 would make the decode arm agree with the prefill arm whether or not +the per-token metadata is read at all, and the gate would pass vacuously. Keep the origins distinct +if this test is ever reworked. + +**Pass `--no-fail-fast`.** `cargo test` stops after the first test binary that fails, so a control run +without it reports hd256 red and never executes hd512 at all — which looks exactly like a control +that passed on both. Filter to the single test under control as well (`-- --test-threads=1 `): +control 2 traps, and a trap poisons the CUDA context for every other test sharing that binary. + +## The assumption to watch + +The bottom link assumes paged prefill and paged decode remain two instantiations of one template +(`qkv_norm_rope_paged_prefill_hd256_plain_kernel`, +`qk_norm_partial_rope_paged_prefill_hd512_kernel`), differing only in how `pos`, the +page window, and the origin are fetched. If a future change forks decode into a standalone kernel for +performance, the bitwise comparison stops being "did the metadata route correctly" and becomes "do +two independent implementations happen to agree" — still useful, but no longer equivalent to the +closed-form coverage it replaced. Whoever forks it owns re-deciding what certifies the decode arm. + +## Next step + +Extend the chain rather than the host oracle. A new dispatch variant gets compared against the arm +one link above it; a genuinely new operator contract — a rotary regime with different control flow, +a new norm — gets its own small anchor, and says in its doc comment which control flow makes it +distinct from the existing three. diff --git a/pegainfer-kernels/tests/hd256_qk_rope_plain_smoke.rs b/pegainfer-kernels/tests/hd256_qk_rope_plain_smoke.rs index 58883b858..182f13b66 100644 --- a/pegainfer-kernels/tests/hd256_qk_rope_plain_smoke.rs +++ b/pegainfer-kernels/tests/hd256_qk_rope_plain_smoke.rs @@ -3,6 +3,22 @@ //! Manual gate: CI compiles this but never runs it. Run on a GPU box with //! PEGAINFER_REQUIRE_GPU=1, which turns a missing device into a failure //! rather than a skip. +//! +//! Two closed-form anchors, then a chain. `full_rotation_matches_closed_form` +//! and `partial_rotation_exercises_tail` are the only tests here that restate +//! the operator, and they are two because rotary_dim == HD and rotary_dim < HD +//! are different control flow: at full width no thread reaches the +//! pass-through tail at all, and full width is the Gemma 4 local-layer +//! production config. From there each gate is measured against an arm the +//! anchors certify — the contiguous kernel certifies the values the paged +//! prefill must land at its layout addresses, and the paged prefill certifies +//! the bits the paged decode must reproduce. Nothing below the anchors +//! recomputes the norm, the RoPE, or the paged address formula. +//! +//! Extend the chain, do not grow a host oracle: a new dispatch variant is +//! compared against the arm one link above it. See +//! docs/subsystems/kernels/qk-rope-smoke-oracles.md for why, and for the +//! negative controls that keep a GPU-vs-GPU gate honest. mod common; @@ -202,67 +218,82 @@ const NUM_LAYERS: usize = 2; const PAGE_INDICES: [i32; 4] = [3, 7, 5, 9]; const POOL_PAGES: usize = 8; -/// K and V blocks at their layout-derived offsets; everything else stays -/// 0.0. V is the weightless norm of the v input — never rotated, no weight. -fn expected_pool(layout: &PagedKvLayout, layer: usize, kw: &[bf16]) -> Vec { - let mut exp = vec![0.0f32; layout.page_stride * POOL_PAGES]; - let layer_offset = (layer * layout.layer_stride) as i64; - for t in 0..SEQ_LEN { - let pos = START_POS + t; - let page = PAGE_INDICES[pos / PAGE_SIZE] as i64; - for h in 0..NUM_KV_HEADS { - let k_x = signed(K_BASE, h, t); - let k_inv = inv_rms(k_x); - let v_x = signed(V_BASE, h, t); - let v_val = bf16::from_f32(v_x * inv_rms(v_x)).to_f32(); - let base = page * layout.page_stride as i64 - + layer_offset - + (pos % PAGE_SIZE) as i64 * KV_DIM as i64 - + h as i64 * HD as i64; - for d in 0..HD { - exp[(base + d as i64) as usize] = expected_prep(k_x, kw, k_inv, d, pos, HD); - exp[(base + layout.kv_block_len as i64 + d as i64) as usize] = v_val; - } - } - } - exp +/// The pool slot a (position, kv head, element) triple owns, read straight +/// off the production `PagedKvLayout` rather than restated from the page +/// geometry. This is the only address arithmetic left in the file: the +/// decode gate below compares two pools element for element and needs none. +fn pool_k_offset( + layout: &PagedKvLayout, + layer: usize, + page: usize, + pos: usize, + kv_head: usize, +) -> usize { + page * layout.page_stride + + layer * layout.layer_stride + + (pos % layout.page_size) * layout.num_kv_heads * layout.head_dim + + kv_head * layout.head_dim } -/// Exact zero is the assertion, not sloppiness: it marks a slot the kernel -/// must never have written — unreferenced pages, the other layer, and the -/// slots outside the request's positions. -#[allow(clippy::float_cmp)] -fn assert_pool(got: &[f32], expected: &[f32]) { - assert_eq!(got.len(), expected.len()); +/// Bitwise, not approximate: every expectation here is a value another GPU +/// arm produced from the same input at the same position, so the two cannot +/// differ at all. The lone host-computed expectation (the weightless V norm) +/// is exact too, for reasons worth reading before loosening this: see +/// docs/subsystems/kernels/qk-rope-smoke-oracles.md. +/// +/// A zero expectation in a pool comparison is a slot the kernel must never +/// have written — unreferenced pages, the other layer, the slots outside the +/// request's positions — so it is asserted as hard as any other value. +/// +/// Reports the first differing element. `assert_eq!` on the whole slice +/// would dump two vectors of tens of thousands of values and bury it. +fn assert_bits_eq(got: &[u16], expected: &[u16], what: &str) { + assert_eq!(got.len(), expected.len(), "{what}: length differs"); for (i, (&g, &e)) in got.iter().zip(expected).enumerate() { - if e == 0.0 { - assert_eq!(g, 0.0, "pool[{i}]: expected untouched, got {g}"); - } else { - assert!( - (g - e).abs() < 0.02, - "pool[{i}]: got {g}, expected {e} (tolerance 0.02)" - ); - } + assert_eq!( + g, + e, + "{what}[{i}]: got {}, expected {}", + bf16::from_bits(g), + bf16::from_bits(e) + ); } } -#[test] -fn pool_write_matches_closed_form_and_touches_nothing_else() { - let Some(ctx) = common::device_or_skip() else { - return; - }; - let ctx = &ctx; - let qw = q_norm_weights(); - let kw = k_norm_weights(); - let layer = 1; - let layout = PagedKvLayout::new(NUM_LAYERS, NUM_KV_HEADS, HD, PAGE_SIZE); +/// bf16 bits, not the widened f32: the pool stores bf16, so reading both +/// sides the same way lets one comparison cover Q rows and pool slots alike. +fn row_bits(rows: &[f32]) -> Vec { + rows.iter().map(|&v| bf16::from_f32(v).to_bits()).collect() +} + +/// Norm weights and RoPE tables on device: the setup both paged gates share. +fn device_fixture(ctx: &DeviceContext) -> (DeviceVec, DeviceVec, DeviceVec, DeviceVec) { + let (cos_dev, sin_dev) = cos_sin_tables(ctx, COS_MAX_POS, HD); + ( + DeviceVec::from_host(ctx, &q_norm_weights()).expect("q_norm_weight H2D"), + DeviceVec::from_host(ctx, &k_norm_weights()).expect("k_norm_weight H2D"), + cos_dev, + sin_dev, + ) +} + +/// The paged prefill over the fixture window, read back as (Q rows, pool) +/// bit patterns. Both gates below need this arm — one as the subject, one as +/// the oracle the decode form must reproduce — and the wrapper takes +/// twenty-two arguments, so it is spelled once here rather than twice. +fn run_paged_prefill( + ctx: &DeviceContext, + layout: &PagedKvLayout, + layer: usize, + qn: &DeviceVec, + kn: &DeviceVec, + cos_dev: &DeviceVec, + sin_dev: &DeviceVec, +) -> (Vec, Vec) { let q = hidden_input(ctx, Q_BASE, NUM_Q_HEADS, SEQ_LEN); let k = hidden_input(ctx, K_BASE, NUM_KV_HEADS, SEQ_LEN); let v = hidden_input(ctx, V_BASE, NUM_KV_HEADS, SEQ_LEN); let mut q_out = HiddenStates::zeros(ctx, Q_DIM, SEQ_LEN).expect("q_out alloc"); - let (cos_dev, sin_dev) = cos_sin_tables(ctx, COS_MAX_POS, HD); - let qn = DeviceVec::from_host(ctx, &qw).expect("q_norm_weight H2D"); - let kn = DeviceVec::from_host(ctx, &kw).expect("k_norm_weight H2D"); let pool: CudaSlice = ctx .stream .alloc_zeros(layout.page_stride * POOL_PAGES) @@ -271,7 +302,6 @@ fn pool_write_matches_closed_form_and_touches_nothing_else() { .stream .clone_htod(&PAGE_INDICES) .expect("page_indices H2D"); - qkv_norm_rope_paged_prefill_hd256_plain_into( ctx, &q, @@ -280,11 +310,11 @@ fn pool_write_matches_closed_form_and_touches_nothing_else() { &mut q_out, 0, &pool, - &layout, - &qn, - &kn, - &cos_dev, - &sin_dev, + layout, + qn, + kn, + cos_dev, + sin_dev, layer, &page_indices, 0, @@ -296,17 +326,85 @@ fn pool_write_matches_closed_form_and_touches_nothing_else() { HD, EPS, ) - .expect("pool prep launch"); + .expect("paged prefill launch"); + let pool_host: Vec = ctx.stream.clone_dtoh(&pool).expect("pool D2H"); + ( + row_bits(&q_out.to_host(ctx).expect("q_out D2H")), + pool_host.iter().map(|x| x.to_bits()).collect(), + ) +} - let qo = q_out.to_host(ctx).expect("q_out D2H"); - assert_close( - &qo, - &expected_full(Q_BASE, &qw, Q_DIM, HD), - "pool-write Q pairing", +/// The paged prefill's addressing, against the contiguous arm rather than a +/// second copy of the operator. `full_rotation_matches_closed_form` +/// certifies that arm's norm and RoPE at this rotary width, so driving it on +/// the same rows from the same start_pos yields the Q and K this prefill +/// must produce, and the only open question left here is whether the paged +/// form lands those values in the slots `PagedKvLayout` names and leaves +/// every other slot alone. +/// +/// V has no contiguous counterpart — the flat kernel carries no V band — so +/// its weightless norm stays as a one-line expectation. That is one multiply +/// over a reduction the anchor already pins, not a second operator. +#[test] +fn paged_prefill_lands_flat_values_at_layout_addresses() { + let Some(ctx) = common::device_or_skip() else { + return; + }; + let ctx = &ctx; + let layer = 1; + let layout = PagedKvLayout::new(NUM_LAYERS, NUM_KV_HEADS, HD, PAGE_SIZE); + let (qn, kn, cos_dev, sin_dev) = device_fixture(ctx); + + // Oracle arm: the same rows over the same position window, through the + // contiguous kernel the closed-form anchors certify. + let oracle_q = hidden_input(ctx, Q_BASE, NUM_Q_HEADS, SEQ_LEN); + let oracle_k = hidden_input(ctx, K_BASE, NUM_KV_HEADS, SEQ_LEN); + let mut oracle_q_out = HiddenStates::zeros(ctx, Q_DIM, SEQ_LEN).expect("oracle q_out alloc"); + let mut oracle_k_out = HiddenStates::zeros(ctx, KV_DIM, SEQ_LEN).expect("oracle k_out alloc"); + qk_norm_rope_prefill_hd256_plain_into( + ctx, + &oracle_q, + &oracle_k, + &mut oracle_q_out, + &mut oracle_k_out, + &qn, + &kn, + &cos_dev, + &sin_dev, + START_POS, + COS_MAX_POS, + NUM_Q_HEADS, + NUM_KV_HEADS, + HD, + EPS, + ) + .expect("oracle flat prep launch"); + let oracle_q_bits = row_bits(&oracle_q_out.to_host(ctx).expect("oracle q_out D2H")); + let oracle_k_vals = oracle_k_out.to_host(ctx).expect("oracle k_out D2H"); + + let (q_bits, pool_bits) = run_paged_prefill(ctx, &layout, layer, &qn, &kn, &cos_dev, &sin_dev); + assert_bits_eq( + &q_bits, + &oracle_q_bits, + "paged Q vs the contiguous arm over the same window", ); - let pool_host: Vec = ctx.stream.clone_dtoh(&pool).expect("pool D2H"); - let pool_f: Vec = pool_host.iter().map(|x| x.to_f32()).collect(); - assert_pool(&pool_f, &expected_pool(&layout, layer, &kw)); + + let mut expected = vec![0u16; layout.page_stride * POOL_PAGES]; + for t in 0..SEQ_LEN { + let pos = START_POS + t; + let page = PAGE_INDICES[pos / PAGE_SIZE] as usize; + for h in 0..NUM_KV_HEADS { + let v_x = bf16::from_f32(signed(V_BASE, h, t)).to_f32(); + let v_val = bf16::from_f32(v_x * inv_rms(v_x)).to_bits(); + let base = pool_k_offset(&layout, layer, page, pos, h); + for d in 0..HD { + expected[base + d] = + bf16::from_f32(oracle_k_vals[t * KV_DIM + h * HD + d]).to_bits(); + expected[base + layout.kv_block_len + d] = v_val; + } + } + } + assert_bits_eq(&pool_bits, &expected, "pool"); } /// rotary_dim = 256 is the Gemma 4 local-layer case: the full head rotates @@ -439,47 +537,74 @@ fn rejects_position_beyond_cos_table() { ); } +/// The per-token metadata form against the whole-window form. Both are +/// instantiations of one template, differing only in where `pos`, the page +/// window and the window's origin come from, so driving the decode arm with +/// a table replicated per row and the origin the prefill arm was handed must +/// reproduce that arm bit for bit. +/// +/// Everything the two share cancels: the norm, the RoPE, the V band, the +/// pool addressing. What is left under test is exactly the routing — +/// `positions[]`, the CSR window, the per-token origin — which is the only +/// code `PER_TOKEN_META` switches on. No host expectation is computed here +/// and no address is derived. #[test] -fn batched_decode_prep_matches_closed_form() { +fn paged_decode_equals_paged_prefill_over_the_same_positions() { let Some(ctx) = common::device_or_skip() else { return; }; let ctx = &ctx; - let qw = q_norm_weights(); - let kw = k_norm_weights(); let layer = 1; let layout = PagedKvLayout::new(NUM_LAYERS, NUM_KV_HEADS, HD, PAGE_SIZE); - let batch = 2usize; - // [7] at origin 1 maps pos 3 to page 7; [3, 5] at origin 0 maps it to page 5. - let positions: [i32; 2] = [3, 3]; - let origins: [i32; 2] = [1, 0]; - let pages_cat: [i32; 3] = [7, 3, 5]; - let indptr: [i32; 3] = [0, 1, 3]; - - let q = hidden_input(ctx, Q_BASE, NUM_Q_HEADS, batch); - let k = hidden_input(ctx, K_BASE, NUM_KV_HEADS, batch); - let v = hidden_input(ctx, V_BASE, NUM_KV_HEADS, batch); - let mut q_out = HiddenStates::zeros(ctx, Q_DIM, batch).expect("q_out alloc"); - let (cos_dev, sin_dev) = cos_sin_tables(ctx, COS_MAX_POS, HD); - let qn = DeviceVec::from_host(ctx, &qw).expect("q_norm_weight H2D"); - let kn = DeviceVec::from_host(ctx, &kw).expect("k_norm_weight H2D"); - let pool: CudaSlice = ctx + let (qn, kn, cos_dev, sin_dev) = device_fixture(ctx); + + // Whole-window arm: one table for the whole prompt, positions derived + // from start_pos, one scalar origin for every row. + let (prefill_q, prefill_kv) = + run_paged_prefill(ctx, &layout, layer, &qn, &kn, &cos_dev, &sin_dev); + + // Per-token arm: equivalent metadata in the form the decode path takes. + // Each row's window is compressed to start at the page holding that row's + // own position — a caller that released the front passes exactly this — + // so the origins differ across rows (0, 1, 1, 2) and no two rows share a + // table span. The per-token origin and the CSR window are both + // load-bearing here, not defaulted: a form that ignored + // page_origins[token] would take row pos/page_size into a window that no + // longer starts at page 0 and land somewhere else. Page 9 stays past + // every reachable row in both arms, so dereferencing the out-of-range + // sentinel traps rather than passing quietly. + let mut per_token_pages: Vec = Vec::new(); + let mut indptr: Vec = vec![0]; + let mut origins: Vec = Vec::new(); + let mut positions: Vec = Vec::new(); + for t in 0..SEQ_LEN { + let pos = START_POS + t; + let origin = pos / PAGE_SIZE; + per_token_pages.extend_from_slice(&PAGE_INDICES[origin..]); + indptr.push(per_token_pages.len() as i32); + origins.push(origin as i32); + positions.push(pos as i32); + } + let q = hidden_input(ctx, Q_BASE, NUM_Q_HEADS, SEQ_LEN); + let k = hidden_input(ctx, K_BASE, NUM_KV_HEADS, SEQ_LEN); + let v = hidden_input(ctx, V_BASE, NUM_KV_HEADS, SEQ_LEN); + let mut decode_q_out = HiddenStates::zeros(ctx, Q_DIM, SEQ_LEN).expect("q_out alloc"); + let decode_pool: CudaSlice = ctx .stream .alloc_zeros(layout.page_stride * POOL_PAGES) .expect("pool alloc"); - let pages_d: CudaSlice = ctx.stream.clone_htod(&pages_cat).expect("pages H2D"); + let pages_d: CudaSlice = ctx.stream.clone_htod(&per_token_pages).expect("pages H2D"); let indptr_d: CudaSlice = ctx.stream.clone_htod(&indptr).expect("indptr H2D"); let origins_d: CudaSlice = ctx.stream.clone_htod(&origins).expect("origins H2D"); let positions_d: CudaSlice = ctx.stream.clone_htod(&positions).expect("positions H2D"); - qkv_norm_rope_paged_decode_hd256_plain_into( ctx, &q, &k, &v, - &mut q_out, + &mut decode_q_out, 0, - &pool, + &decode_pool, &layout, &qn, &kn, @@ -496,43 +621,17 @@ fn batched_decode_prep_matches_closed_form() { HD, EPS, ) - .expect("batched decode prep launch"); - - let qo = q_out.to_host(ctx).expect("q_out D2H"); - let mut q_exp = vec![0.0f32; Q_DIM * batch]; - for (row, &pos) in positions.iter().enumerate() { - for h in 0..NUM_Q_HEADS { - let x = signed(Q_BASE, h, row); - let inv = inv_rms(x); - for d in 0..HD { - q_exp[row * Q_DIM + h * HD + d] = expected_prep(x, &qw, inv, d, pos as usize, HD); - } - } - } - assert_close(&qo, &q_exp, "batched decode pairing/tail"); - - let pool_host: Vec = ctx.stream.clone_dtoh(&pool).expect("pool D2H"); - let got: Vec = pool_host.iter().map(|x| x.to_f32()).collect(); - let mut exp = vec![0.0f32; layout.page_stride * POOL_PAGES]; - let layer_offset = (layer * layout.layer_stride) as i64; - for (row, (&pos, &page)) in positions.iter().zip([7i32, 5i32].iter()).enumerate() { - for h in 0..NUM_KV_HEADS { - let k_x = signed(K_BASE, h, row); - let k_inv = inv_rms(k_x); - let v_x = signed(V_BASE, h, row); - let v_val = bf16::from_f32(v_x * inv_rms(v_x)).to_f32(); - let base = page as i64 * layout.page_stride as i64 - + layer_offset - + (pos as usize % PAGE_SIZE) as i64 * KV_DIM as i64 - + h as i64 * HD as i64; - for d in 0..HD { - exp[(base + d as i64) as usize] = - expected_prep(k_x, &kw, k_inv, d, pos as usize, HD); - exp[(base + layout.kv_block_len as i64 + d as i64) as usize] = v_val; - } - } - } - assert_pool(&got, &exp); + .expect("paged decode prep launch"); + + let decode_q = row_bits(&decode_q_out.to_host(ctx).expect("q_out D2H")); + let decode_pool_host: Vec = ctx.stream.clone_dtoh(&decode_pool).expect("pool D2H"); + let decode_kv: Vec = decode_pool_host.iter().map(|x| x.to_bits()).collect(); + assert_bits_eq(&decode_q, &prefill_q, "per-token Q vs the whole-window run"); + assert_bits_eq( + &decode_kv, + &prefill_kv, + "per-token pool writes vs the whole-window run", + ); } /// The row-offset suffix contract: with `row_offset = 1` over three rows, diff --git a/pegainfer-kernels/tests/hd512_qk_rope_smoke.rs b/pegainfer-kernels/tests/hd512_qk_rope_smoke.rs index 7644021fa..ba61b0155 100644 --- a/pegainfer-kernels/tests/hd512_qk_rope_smoke.rs +++ b/pegainfer-kernels/tests/hd512_qk_rope_smoke.rs @@ -4,6 +4,21 @@ //! PEGAINFER_REQUIRE_GPU=1, which turns a missing device into a failure //! rather than a skip. Traps are in their own binaries — __trap() poisons //! the context for whatever runs next. +//! +//! One closed-form anchor, then a chain. `decode_prep_matches_closed_form` +//! is the only test here that restates the operator: it pins RMSNorm and the +//! partial-RoPE pairing/tail on the batched arm, and rotary_dim 128 < 512 +//! puts rope-lo, rope-hi and the pass-through tail in one run, so a second +//! anchor would be redundant. From there each gate is measured against an +//! arm the one above it certifies — batched decode certifies the values the +//! paged prefill must land at its layout addresses, and the paged prefill +//! certifies the bits the paged decode must reproduce. Nothing below the +//! anchor recomputes the norm, the RoPE, or the paged address formula. +//! +//! Extend the chain, do not grow a host oracle: a new dispatch variant is +//! compared against the arm one link above it. See +//! docs/subsystems/kernels/qk-rope-smoke-oracles.md for why, and for the +//! negative controls that keep a GPU-vs-GPU gate honest. mod common; @@ -80,6 +95,15 @@ fn cos_sin_tables(ctx: &DeviceContext, rows: usize) -> (DeviceVec, DeviceVec) { ) } +/// Per-row distinct values so a row swap cannot hide behind a shared +/// constant, each row constant across its head vector so the 512-way +/// reduction stays exact. +fn row_major(base: f32, dim: usize) -> Vec { + (0..SEQ_LEN) + .flat_map(|t| vec![bf16::from_f32(base + t as f32); dim]) + .collect() +} + fn expected_prep(x: f32, w: &[bf16], inv: f32, d: usize, row: usize) -> f32 { if d < HALF_ROTARY { let lo = normed(x, w, inv, d); @@ -102,16 +126,10 @@ fn expected_prep(x: f32, w: &[bf16], inv: f32, d: usize, row: usize) -> f32 { } } -fn expected_full( - x: f32, - w: &[bf16], - inv: f32, - dim: usize, - row_of: impl Fn(usize) -> usize, -) -> Vec { +fn expected_full(x: f32, w: &[bf16], inv: f32, dim: usize) -> Vec { let mut full = vec![0.0f32; dim * SEQ_LEN]; for t in 0..SEQ_LEN { - let row = row_of(t); + let row = POSITIONS[t] as usize; for d in 0..dim { full[t * dim + d] = expected_prep(x, w, inv, d % HD, row); } @@ -119,29 +137,21 @@ fn expected_full( full } -/// K and V blocks; everything else stays 0.0. The offsets are derived from -/// the layout, so oracle and kernel share no hand-picked raw offset. V is -/// the K=V fork: the same raw vector under the shared inv_rms, weightless -/// and un-rotated. -fn expected_pool(x: f32, w: &[bf16], inv: f32, layout: &PagedKvLayout, layer: usize) -> Vec { - let mut exp = vec![0.0f32; POOL_LEN]; - let layer_offset = (layer * layout.layer_stride) as i64; - let v_val = bf16::from_f32(x * inv).to_f32(); - for t in 0..SEQ_LEN { - let pos = START_POS + t; - let page = PAGE_INDICES[pos / PAGE_SIZE] as i64; - for h in 0..NUM_KV_HEADS { - let base = page * PAGE_STRIDE - + layer_offset - + (pos % PAGE_SIZE) as i64 * KV_DIM as i64 - + h as i64 * HD as i64; - for d in 0..HD { - exp[(base + d as i64) as usize] = expected_prep(x, w, inv, d, pos); - exp[(base + layout.kv_block_len as i64 + d as i64) as usize] = v_val; - } - } - } - exp +/// The pool slot a (position, kv head, element) triple owns, read straight +/// off the production `PagedKvLayout` rather than restated from the page +/// geometry. This is the only address arithmetic left in the file: the +/// decode gate below compares two pools element for element and needs none. +fn pool_k_offset( + layout: &PagedKvLayout, + layer: usize, + page: usize, + pos: usize, + kv_head: usize, +) -> usize { + page * layout.page_stride + + layer * layout.layer_stride + + (pos % layout.page_size) * layout.num_kv_heads * layout.head_dim + + kv_head * layout.head_dim } fn assert_close(got: &[f32], expected: &[f32], what: &str) { @@ -154,26 +164,39 @@ fn assert_close(got: &[f32], expected: &[f32], what: &str) { } } -/// Exact zero is the assertion, not sloppiness: it marks a slot the kernel -/// must never have written — unreferenced pages, the other layer, and the -/// slots outside the request's positions. -#[allow(clippy::float_cmp)] -fn assert_pool(got: &[f32], expected: &[f32]) { - assert_eq!(got.len(), expected.len()); +/// Bitwise, not approximate: every expectation here is a value another GPU +/// arm produced from the same input at the same position, so the two cannot +/// differ at all. The lone host-computed expectation (the weightless V norm) +/// is exact too, for reasons worth reading before loosening this: see +/// docs/subsystems/kernels/qk-rope-smoke-oracles.md. +/// +/// A zero expectation in a pool comparison is a slot the kernel must never +/// have written — unreferenced pages, the other layer, the slots outside the +/// request's positions — so it is asserted as hard as any other value. +/// +/// Reports the first differing element. `assert_eq!` on the whole slice +/// would dump two vectors of tens of thousands of values and bury it. +fn assert_bits_eq(got: &[u16], expected: &[u16], what: &str) { + assert_eq!(got.len(), expected.len(), "{what}: length differs"); for (i, (&g, &e)) in got.iter().zip(expected).enumerate() { - if e == 0.0 { - assert_eq!(g, 0.0, "pool[{i}]: expected untouched, got {g}"); - } else { - assert!( - (g - e).abs() < 0.02, - "pool[{i}]: got {g}, expected {e} (tolerance 0.02)" - ); - } + assert_eq!( + g, + e, + "{what}[{i}]: got {}, expected {}", + bf16::from_bits(g), + bf16::from_bits(e) + ); } } -/// Starts at 1: w[0] = 0 would make dim 0 normalise to 0.0, which -/// assert_pool cannot tell from an untouched slot. +/// bf16 bits, not the widened f32: the pool stores bf16, so reading both +/// sides the same way lets one comparison cover Q rows and pool slots alike. +fn row_bits(rows: &[f32]) -> Vec { + rows.iter().map(|&v| bf16::from_f32(v).to_bits()).collect() +} + +/// Starts at 1: w[0] = 0 would make dim 0 normalise to 0.0, which the pool +/// comparison cannot tell from an untouched slot. fn q_norm_weights() -> Vec { (1..=HD).map(|d| bf16::from_f32(d as f32)).collect() } @@ -185,35 +208,40 @@ fn k_norm_weights() -> Vec { (1..=HD).map(|d| bf16::from_f32(-(d as f32))).collect() } -#[test] -fn prefill_prep_matches_closed_form() { - let Some(ctx) = common::device_or_skip() else { - return; - }; - let ctx = &ctx; - let qw = q_norm_weights(); - let kw = k_norm_weights(); - let layer = 1; - let layout = PagedKvLayout::new(NUM_LAYERS, NUM_KV_HEADS, HD, PAGE_SIZE); - assert_eq!( - layout.page_stride as i64, PAGE_STRIDE, - "test geometry drift: layout page_stride must match PAGE_STRIDE" - ); - let ones_q = vec![bf16::from_f32(Q_INPUT); Q_DIM * SEQ_LEN]; - let q = HiddenStates::from_host(ctx, &ones_q, Q_DIM, SEQ_LEN).expect("q H2D"); - let ones_k = vec![bf16::from_f32(K_INPUT); KV_DIM * SEQ_LEN]; - let k = HiddenStates::from_host(ctx, &ones_k, KV_DIM, SEQ_LEN).expect("k H2D"); - let mut q_out = HiddenStates::zeros(ctx, Q_DIM, SEQ_LEN).expect("q_out alloc"); - +/// Norm weights and RoPE tables on device: the setup both paged gates share. +fn device_fixture(ctx: &DeviceContext) -> (DeviceVec, DeviceVec, DeviceVec, DeviceVec) { let (cos_dev, sin_dev) = cos_sin_tables(ctx, 8); - let qn = DeviceVec::from_host(ctx, &qw).expect("q_norm_weight H2D"); - let kn = DeviceVec::from_host(ctx, &kw).expect("k_norm_weight H2D"); + ( + DeviceVec::from_host(ctx, &q_norm_weights()).expect("q_norm_weight H2D"), + DeviceVec::from_host(ctx, &k_norm_weights()).expect("k_norm_weight H2D"), + cos_dev, + sin_dev, + ) +} + +/// The paged prefill over the fixture window, read back as (Q rows, pool) +/// bit patterns. Both gates below need this arm — one as the subject, one as +/// the oracle the decode form must reproduce — and the wrapper takes twenty +/// arguments, so it is spelled once here rather than twice. +fn run_paged_prefill( + ctx: &DeviceContext, + layout: &PagedKvLayout, + layer: usize, + qn: &DeviceVec, + kn: &DeviceVec, + cos_dev: &DeviceVec, + sin_dev: &DeviceVec, +) -> (Vec, Vec) { + let q = + HiddenStates::from_host(ctx, &row_major(Q_INPUT, Q_DIM), Q_DIM, SEQ_LEN).expect("q H2D"); + let k = + HiddenStates::from_host(ctx, &row_major(K_INPUT, KV_DIM), KV_DIM, SEQ_LEN).expect("k H2D"); + let mut q_out = HiddenStates::zeros(ctx, Q_DIM, SEQ_LEN).expect("q_out alloc"); let pool: CudaSlice = ctx.stream.alloc_zeros(POOL_LEN).expect("pool alloc"); let page_indices: CudaSlice = ctx .stream .clone_htod(&PAGE_INDICES) .expect("page_indices H2D"); - qk_norm_partial_rope_paged_prefill_hd512_into( ctx, &q, @@ -221,11 +249,11 @@ fn prefill_prep_matches_closed_form() { &mut q_out, 0, &pool, - &layout, - &qn, - &kn, - &cos_dev, - &sin_dev, + layout, + qn, + kn, + cos_dev, + sin_dev, layer, &page_indices, 0, @@ -236,59 +264,166 @@ fn prefill_prep_matches_closed_form() { ROTARY_DIM, EPS, ) - .expect("prefill prep launch"); + .expect("paged prefill launch"); + let pool_host: Vec = ctx.stream.clone_dtoh(&pool).expect("pool D2H"); + ( + row_bits(&q_out.to_host(ctx).expect("q_out D2H")), + pool_host.iter().map(|x| x.to_bits()).collect(), + ) +} - let qo = q_out.to_host(ctx).expect("q_out D2H"); - assert_close( - &qo, - &expected_full(Q_INPUT, &qw, inv_rms(Q_INPUT), Q_DIM, |t| START_POS + t), - "prefill pairing/tail L1", +/// The paged prefill's addressing, against the batched-decode arm rather +/// than a second copy of the operator. `decode_prep_matches_closed_form` +/// certifies that arm's norm and RoPE; driving it at the same positions +/// yields the Q and K this prefill must produce, so the only open question +/// left here is whether the paged form lands those values in the slots +/// `PagedKvLayout` names and leaves every other slot alone. +/// +/// Rows carry distinct values so a row swap cannot hide behind a shared +/// constant, and each row stays constant across its head vector, which is +/// what keeps the 512-way reduction exact. +#[test] +fn prefill_prep_lands_batched_values_at_layout_addresses() { + let Some(ctx) = common::device_or_skip() else { + return; + }; + let ctx = &ctx; + let layer = 1; + let layout = PagedKvLayout::new(NUM_LAYERS, NUM_KV_HEADS, HD, PAGE_SIZE); + assert_eq!( + layout.page_stride as i64, PAGE_STRIDE, + "test geometry drift: layout page_stride must match PAGE_STRIDE" ); - let pool_host: Vec = ctx.stream.clone_dtoh(&pool).expect("pool D2H"); - let pool_f: Vec = pool_host.iter().map(|x| x.to_f32()).collect(); - assert_pool( - &pool_f, - &expected_pool(K_INPUT, &kw, inv_rms(K_INPUT), &layout, layer), + let q_host = row_major(Q_INPUT, Q_DIM); + let k_host = row_major(K_INPUT, KV_DIM); + let (qn, kn, cos_dev, sin_dev) = device_fixture(ctx); + + // Oracle arm: the same rows at the same absolute positions the prefill + // will derive from start_pos, through the kernel the closed-form anchor + // certifies. K comes back updated in place. + let oracle_positions: Vec = (0..SEQ_LEN).map(|t| (START_POS + t) as i32).collect(); + let positions_d: CudaSlice = ctx + .stream + .clone_htod(&oracle_positions) + .expect("positions H2D"); + let oracle_q = HiddenStates::from_host(ctx, &q_host, Q_DIM, SEQ_LEN).expect("oracle q H2D"); + let mut oracle_k = + HiddenStates::from_host(ctx, &k_host, KV_DIM, SEQ_LEN).expect("oracle k H2D"); + let mut oracle_q_out = HiddenStates::zeros(ctx, Q_DIM, SEQ_LEN).expect("oracle q_out alloc"); + qk_norm_partial_rope_batched_decode_hd512_into( + ctx, + &oracle_q, + &mut oracle_q_out, + &mut oracle_k, + &qn, + &kn, + &cos_dev, + &sin_dev, + &positions_d, + 8, // cos_max_pos + NUM_Q_HEADS, + NUM_KV_HEADS, + ROTARY_DIM, + EPS, + ) + .expect("oracle batched decode launch"); + let oracle_q_bits = row_bits(&oracle_q_out.to_host(ctx).expect("oracle q_out D2H")); + let oracle_k_vals = oracle_k.to_host(ctx).expect("oracle k D2H"); + + let (q_bits, pool_bits) = run_paged_prefill(ctx, &layout, layer, &qn, &kn, &cos_dev, &sin_dev); + assert_bits_eq( + &q_bits, + &oracle_q_bits, + "prefill Q vs the batched-decode arm at the same positions", ); + + // K into the pool's K block at its layout address, V into the V block a + // kv_block_len further on. V is the K=V fork: the weightless norm of the + // same raw row, sharing inv_rms, never rotated. + let mut expected = vec![0u16; POOL_LEN]; + for t in 0..SEQ_LEN { + let pos = START_POS + t; + let page = PAGE_INDICES[pos / PAGE_SIZE] as usize; + let raw = bf16::from_f32(K_INPUT + t as f32).to_f32(); + let v_val = bf16::from_f32(raw * inv_rms(raw)).to_bits(); + for h in 0..NUM_KV_HEADS { + let base = pool_k_offset(&layout, layer, page, pos, h); + for d in 0..HD { + expected[base + d] = + bf16::from_f32(oracle_k_vals[t * KV_DIM + h * HD + d]).to_bits(); + expected[base + layout.kv_block_len + d] = v_val; + } + } + } + assert_bits_eq(&pool_bits, &expected, "pool"); } +/// The per-token metadata form against the whole-window form. Both are +/// instantiations of one template, differing only in where `pos`, the page +/// window and the window's origin come from, so driving the decode arm with +/// a table replicated per row and origins pinned to 0 — what the prefill +/// arm hardcodes — must reproduce the prefill arm bit for bit. +/// +/// Everything the two share cancels: the norm, the RoPE, the V fork, the +/// pool addressing. What is left under test is exactly the routing — +/// `positions[]`, the CSR window, the per-token origin — which is the only +/// code `PER_TOKEN_META` switches on. No host expectation is computed here +/// and no address is derived. #[test] -fn paged_decode_prep_matches_closed_form() { +fn paged_decode_equals_paged_prefill_over_the_same_positions() { let Some(ctx) = common::device_or_skip() else { return; }; let ctx = &ctx; - let qw = q_norm_weights(); - let kw = k_norm_weights(); let layer = 1; let layout = PagedKvLayout::new(NUM_LAYERS, NUM_KV_HEADS, HD, PAGE_SIZE); - let batch = 2usize; - let positions: [i32; 2] = [3, 1]; - let pages_cat: [i32; 3] = [3, 7, 5]; - let indptr: [i32; 3] = [0, 2, 3]; + let (qn, kn, cos_dev, sin_dev) = device_fixture(ctx); - let ones_q = vec![bf16::from_f32(Q_INPUT); Q_DIM * batch]; - let q = HiddenStates::from_host(ctx, &ones_q, Q_DIM, batch).expect("q H2D"); - let ones_k = vec![bf16::from_f32(K_INPUT); KV_DIM * batch]; - let k = HiddenStates::from_host(ctx, &ones_k, KV_DIM, batch).expect("k H2D"); - let mut q_out = HiddenStates::zeros(ctx, Q_DIM, batch).expect("q_out alloc"); + // Whole-window arm: one table for the whole prompt, positions derived + // from start_pos, origin hardcoded to 0 by the PER_TOKEN_META = false + // instantiation. + let (prefill_q, prefill_kv) = + run_paged_prefill(ctx, &layout, layer, &qn, &kn, &cos_dev, &sin_dev); - let (cos_dev, sin_dev) = cos_sin_tables(ctx, 8); - let qn = DeviceVec::from_host(ctx, &qw).expect("q_norm_weight H2D"); - let kn = DeviceVec::from_host(ctx, &kw).expect("k_norm_weight H2D"); - let pool: CudaSlice = ctx.stream.alloc_zeros(POOL_LEN).expect("pool alloc"); - let pages_d: CudaSlice = ctx.stream.clone_htod(&pages_cat).expect("pages H2D"); + // Per-token arm: equivalent metadata in the form the decode path takes. + // Each row's window is compressed to start at the page holding that row's + // own position, which is what the global family is allowed to do, so the + // origins differ across rows (0, 1, 1, 2) and no two rows share a table + // span — the per-token origin and the CSR window are both load-bearing + // here, not defaulted. A form that ignored page_origins[token] would take + // row pos/page_size into a window that no longer starts at page 0 and + // land somewhere else. Page 9 stays past every reachable row in both + // arms, so dereferencing the out-of-range sentinel traps rather than + // passing quietly. + let mut per_token_pages: Vec = Vec::new(); + let mut indptr: Vec = vec![0]; + let mut origins: Vec = Vec::new(); + let mut positions: Vec = Vec::new(); + for t in 0..SEQ_LEN { + let pos = START_POS + t; + let origin = pos / PAGE_SIZE; + per_token_pages.extend_from_slice(&PAGE_INDICES[origin..]); + indptr.push(per_token_pages.len() as i32); + origins.push(origin as i32); + positions.push(pos as i32); + } + let q = + HiddenStates::from_host(ctx, &row_major(Q_INPUT, Q_DIM), Q_DIM, SEQ_LEN).expect("q H2D"); + let k = + HiddenStates::from_host(ctx, &row_major(K_INPUT, KV_DIM), KV_DIM, SEQ_LEN).expect("k H2D"); + let mut decode_q_out = HiddenStates::zeros(ctx, Q_DIM, SEQ_LEN).expect("q_out alloc"); + let decode_pool: CudaSlice = ctx.stream.alloc_zeros(POOL_LEN).expect("pool alloc"); + let pages_d: CudaSlice = ctx.stream.clone_htod(&per_token_pages).expect("pages H2D"); let indptr_d: CudaSlice = ctx.stream.clone_htod(&indptr).expect("indptr H2D"); - let origins_zero_d: CudaSlice = ctx.stream.clone_htod(&[0i32; 2]).expect("origins H2D"); + let origins_d: CudaSlice = ctx.stream.clone_htod(&origins).expect("origins H2D"); let positions_d: CudaSlice = ctx.stream.clone_htod(&positions).expect("positions H2D"); - qk_norm_partial_rope_paged_decode_hd512_into( ctx, &q, &k, - &mut q_out, + &mut decode_q_out, 0, - &pool, + &decode_pool, &layout, &qn, &kn, @@ -297,9 +432,9 @@ fn paged_decode_prep_matches_closed_form() { layer, &pages_d, &indptr_d, - &origins_zero_d, + &origins_d, &positions_d, - 8, // cos_max_pos + 8, NUM_Q_HEADS, NUM_KV_HEADS, ROTARY_DIM, @@ -307,39 +442,15 @@ fn paged_decode_prep_matches_closed_form() { ) .expect("paged decode prep launch"); - let qo = q_out.to_host(ctx).expect("q_out D2H"); - let q_inv = inv_rms(Q_INPUT); - let mut q_exp = vec![0.0f32; Q_DIM * batch]; - for (t, &pos) in positions.iter().enumerate() { - for h in 0..NUM_Q_HEADS { - for d in 0..HD { - q_exp[t * Q_DIM + h * HD + d] = expected_prep(Q_INPUT, &qw, q_inv, d, pos as usize); - } - } - } - assert_close(&qo, &q_exp, "paged decode pairing/tail"); - - let inv = inv_rms(K_INPUT); - let v_val = bf16::from_f32(K_INPUT * inv).to_f32(); - let mut exp = vec![0.0f32; POOL_LEN]; - let layer_offset = (layer * layout.layer_stride) as i64; - // (row, its page, its position): row 0's window [3, 7] covers pos 3 in - // page 7; row 1's window [5] covers pos 1 in page 5. - for (page, pos) in [(7i64, 3usize), (5, 1)] { - for h in 0..NUM_KV_HEADS { - let base = page * PAGE_STRIDE - + layer_offset - + (pos % PAGE_SIZE) as i64 * KV_DIM as i64 - + h as i64 * HD as i64; - for d in 0..HD { - exp[(base + d as i64) as usize] = expected_prep(K_INPUT, &kw, inv, d, pos); - exp[(base + layout.kv_block_len as i64 + d as i64) as usize] = v_val; - } - } - } - let pool_host: Vec = ctx.stream.clone_dtoh(&pool).expect("pool D2H"); - let pool_f: Vec = pool_host.iter().map(|x| x.to_f32()).collect(); - assert_pool(&pool_f, &exp); + let decode_q = row_bits(&decode_q_out.to_host(ctx).expect("q_out D2H")); + let decode_pool_host: Vec = ctx.stream.clone_dtoh(&decode_pool).expect("pool D2H"); + let decode_kv: Vec = decode_pool_host.iter().map(|x| x.to_bits()).collect(); + assert_bits_eq(&decode_q, &prefill_q, "per-token Q vs the whole-window run"); + assert_bits_eq( + &decode_kv, + &prefill_kv, + "per-token pool writes vs the whole-window run", + ); } #[test] @@ -382,17 +493,13 @@ fn decode_prep_matches_closed_form() { let qo = q_out.to_host(ctx).expect("q_out D2H"); assert_close( &qo, - &expected_full(Q_INPUT, &qw, inv_rms(Q_INPUT), Q_DIM, |t| { - POSITIONS[t] as usize - }), + &expected_full(Q_INPUT, &qw, inv_rms(Q_INPUT), Q_DIM), "decode pairing/tail", ); let k_host = k.to_host(ctx).expect("k D2H"); assert_close( &k_host, - &expected_full(K_INPUT, &kw, inv_rms(K_INPUT), KV_DIM, |t| { - POSITIONS[t] as usize - }), + &expected_full(K_INPUT, &kw, inv_rms(K_INPUT), KV_DIM), "decode pairing/tail", ); }