test(kernels): reduce the duplicated CPU operator oracles in the HD256/HD512 QK-RoPE smokes - #1002
Open
BreezyB1n wants to merge 3 commits into
Open
Conversation
The two QK-RoPE device gates carry seven closed-form tests backed by two independent host copies of the RMSNorm/RoPE/paged-address math. That second implementation has to track the .cu contract, and a specification mistake made once can be copied into both the kernel and its local oracle, leaving the gate green while both are wrong. Record the discipline that replaces it: three closed-form anchors, and every other expectation derived from a GPU arm an anchor already certifies (anchor -> flat kernel -> paged prefill -> paged decode). The rest of the doc is the part that is easy to get wrong when reworking these gates, all of it measured on an A40 rather than argued: - why the gemma4 layer oracle is corroboration and not a successor: it is #[ignore]d and needs the pinned 12B checkpoint, so it never runs by default (docs/conventions/migration-defense.md wants a successor that actually runs) - why PEGAINFER_REQUIRE_GPU=1 is mandatory: without it a missing device is a skip and the suite passes having executed nothing - why the pool assertions are bitwise, and the three conditions that make the one host-computed value exact - three negative controls with their observed output, and why the in-bounds one (2b) is the one that proves the assertion is load-bearing rather than the kernel's own __trap() guard - why the decode arm must keep distinct non-zero per-token origins, or it agrees with the prefill arm whether or not the metadata is read - why control runs need --no-fail-fast: cargo stops after the first failing binary, so hd512 silently never runs Refs pegainfer-project#943 Signed-off-by: B1n_FreeToRoAm <z2662728664@gmail.com>
The file carried three closed-form tests over a host copy of the operator (inv_rms/normed/expected_prep/expected_full/expected_pool). Keep one anchor and derive the rest from a GPU arm it certifies. decode_prep_matches_closed_form stays as the anchor: rotary_dim 128 < 512 puts rope-lo, rope-hi and the pass-through tail in a single run, so a second anchor would restate the same three paths. prefill_prep_matches_closed_form becomes prefill_prep_lands_batched_values_at_layout_addresses: the batched arm the anchor certifies produces the Q and K at the same positions, and the test asks only whether the paged form lands them at the addresses PagedKvLayout names, with exact-zero sentinels everywhere else. paged_decode_prep_matches_closed_form becomes paged_decode_equals_paged_prefill_over_the_same_positions: both are instantiations of one template, so equivalent metadata must give identical bits. Everything shared cancels and only the per-token routing is left under test. The per-token arm compresses each row's window to start at that row's own page - what a caller that released the front passes - so the rows carry distinct non-zero origins (0, 1, 1, 2) over four different table spans. Handing every row the full table at origin 0 would agree whether or not page_origins[token] is read at all; the replaced test used origins [1, 0] for exactly this reason and that coverage is kept. assert_pool/expected_pool give way to one assert_bits_eq that takes a label and reports the first differing element. assert_eq! on the whole slice dumps two vectors of tens of thousands of values, which makes a real failure unreadable - control 2b below was invisible until this landed. The paged prefill arm is spelled once, in run_paged_prefill: both gates need it - one as the subject, one as the oracle the decode form must reproduce - and the wrapper takes twenty-plus arguments, so a second copy is a second place to edit when the signature moves. device_fixture does the same for the norm weights and RoPE tables. Each file is now down to three kernel call sites, one per distinct kernel, which is the floor. assert_bits_eq's label is just "pool"; the zero-means-untouched semantics moved into the helper's doc comment, where it is not appended to every unrelated failure message. Successors for what was deleted: - expected_pool -> the batched arm's K values, at layout addresses - inline q_exp/pool reconstruction -> the bitwise prefill/decode gate - 0.02 pool tolerance -> assert_bits_eq, exact - per-token origin coverage -> the compressed windows described above Refusal gates and all three row-offset metamorphic gates are byte-identical to before; the anchor's only edit is dropping expected_full's now-unused row_of closure. Verified on one A40 (sm_86, CUDA 12.8, driver 570.211.01): - PEGAINFER_REQUIRE_GPU=1 cargo test --release -p pegainfer-kernels --test hd512_qk_rope_smoke -> 9 passed; 0 failed - control 1 (paged address +1) -> prefill_prep_lands_batched_values_at_ layout_addresses FAILED at pool[14848]: got 0, expected 65 - control 2b (CSR window pinned to row 0) -> paged_decode_equals_paged_ prefill FAILED at pool[14336]: got 1, expected 0 - cargo fmt --all --check and cargo clippy --release -p pegainfer-kernels --all-targets -- -D warnings both clean Refs pegainfer-project#943 Signed-off-by: B1n_FreeToRoAm <z2662728664@gmail.com>
Same reduction as the hd512 sibling: four closed-form tests over a host copy of the operator become two anchors plus two gates measured against a GPU arm those anchors certify. Two anchors rather than one because rotary_dim == HD and rotary_dim < HD are different control flow, not different data: at full width no thread reaches the pass-through tail at all, and full width is the Gemma 4 local-layer production config. full_rotation_matches_closed_form and partial_rotation_exercises_tail both stay, byte-identical. pool_write_matches_closed_form_and_touches_nothing_else becomes paged_prefill_lands_flat_values_at_layout_addresses, measured against the contiguous kernel the anchors certify. batched_decode_prep_matches_closed_form becomes paged_decode_equals_paged_prefill_over_the_same_positions - a bitwise comparison of the two instantiations of one template, with each row's window compressed to start at its own page so the rows carry distinct non-zero origins (0, 1, 1, 2). That keeps what the replaced test proved with origins [1, 0]: the per-token origin is read per row, not defaulted. assert_pool/expected_pool give way to one assert_bits_eq that takes a label and reports the first differing element, rather than dumping two vectors of tens of thousands of values on failure. The paged prefill arm is spelled once, in run_paged_prefill: both gates need it - one as the subject, one as the oracle the decode form must reproduce - and the wrapper takes twenty-plus arguments, so a second copy is a second place to edit when the signature moves. device_fixture does the same for the norm weights and RoPE tables. Each file is now down to three kernel call sites, one per distinct kernel, which is the floor. assert_bits_eq's label is just "pool"; the zero-means-untouched semantics moved into the helper's doc comment, where it is not appended to every unrelated failure message. Successors for what was deleted: - expected_pool -> the contiguous arm's K values, at layout addresses - inline q_exp/pool reconstruction -> the bitwise prefill/decode gate - 0.02 pool tolerance -> assert_bits_eq, exact - per-token origin coverage -> the compressed windows described above V keeps a one-line weightless norm: the contiguous kernel has no V band, so there is no arm to borrow it from. Both refusal gates and both row-offset metamorphic gates are byte-identical to before. Verified on one A40 (sm_86, CUDA 12.8, driver 570.211.01): - PEGAINFER_REQUIRE_GPU=1 cargo test --release -p pegainfer-kernels --test hd256_qk_rope_plain_smoke -> 8 passed; 0 failed - control 1 (paged address +1) -> paged_prefill_lands_flat_values_at_ layout_addresses FAILED at pool[59392]: got 0, expected 129 - control 2b (CSR window pinned to row 0) -> paged_decode_equals_paged_ prefill FAILED at pool[57344]: got -129, expected 0 - cargo fmt --all --check and cargo clippy --release -p pegainfer-kernels --all-targets -- -D warnings both clean Refs pegainfer-project#943 Signed-off-by: B1n_FreeToRoAm <z2662728664@gmail.com>
Collaborator
|
Please rebase this PR onto the current |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Fixes #943
The HD256/HD512 QK-RoPE device gates carried seven closed-form tests over two independent host copies of the RMSNorm / RoPE / paged-address math (
inv_rms,normed,expected_prep,expected_full,expected_pool,assert_pool). That second implementation has to track the.cucontract, and a specification mistake made once can be copied into both the kernel and its local oracle — leaving the gate green while both are wrong.This keeps three closed-form anchors and derives every other expectation from a GPU arm an anchor already certifies:
Each link catches a failure the link above it cannot: a wrong formula surfaces at the anchor (which tests the flat kernel directly); a wrong address surfaces as flat values landing in the wrong pool slot; wrong per-token routing surfaces as the two instantiations of one template disagreeing bit for bit.
No production code is touched — kernels, dispatch, and
src/are untouched; the only.cuedits in this branch are the temporary negative-control perturbations, reverted before every commit.Why the line count goes up
hd256783 → 882,hd512951 → 1058. The metric the issue asks for is duplicated semantic code, and that is down; the line count is up because the metamorphic method the issue asks for runs two real kernel arms per gate, and these wrappers take 20–22 arguments, one per line under rustfmt. Roughly 31–37% of each new gate is bare argument lines. Each file is now at three kernel call sites — one per distinct kernel, which is the floor — afterrun_paged_prefillanddevice_fixtureremoved the duplicated arm and setup.Type of Change
The template has no box for a test-semantics change, which is what this primarily is;
Documentation updateis checked for the new subsystem doc that ships with it.Inventory: the seven closed-form tests
full_rotation_matches_closed_formrotary_dim == HD: the Gemma 4 local-layer production config, where the tail is empty and no thread takes thed >= rotary_dimbranchpartial_rotation_exercises_tailrotary_dim < HD: the only config that exercises the pass-through tailpool_write_matches_closed_form_and_touches_nothing_elsepaged_prefill_lands_flat_values_at_layout_addressesbatched_decode_prep_matches_closed_formpositions[], CSR window, per-token originpaged_decode_equals_paged_prefill_over_the_same_positionsdecode_prep_matches_closed_formrotary_dim128 < 512 puts rope-lo, rope-hi and the tail in one run, so a second hd512 anchor would restate the same three pathsprefill_prep_matches_closed_formprefill_prep_lands_batched_values_at_layout_addressespaged_decode_prep_matches_closed_formpaged_decode_equals_paged_prefill_over_the_same_positionshd256 keeps two anchors rather than one because
rotary_dim == HDandrotary_dim < HDare different control flow, not different data.The nine non-closed-form tests (three refusal gates, four row-offset metamorphic gates,
split_read_row_offset_serves_only_the_suffix,prefill_rejects_undersized_kv_pool) contain no operator math and are untouched.Before / after
expected_pool,assert_pool)mainThe two surviving functions that are not byte-identical are
expected_full(dropped a now-unusedrow_ofclosure parameter) and its only callerdecode_prep_matches_closed_form.Deleted defences and their successors
Per
docs/conventions/migration-defense.md:expected_pool(×2)PagedKvLayout, with exact-zero sentinels elsewhereq_exp/ pool reconstruction (×3)assert_pool's 0.02 toleranceassert_bits_eq, exactorigins = [1, 0]in test #40, 1, 1, 2over four different table spans. Handing every row the full table at origin 0 would agree whether or not the metadata is read at all; control 2b below is what proves this is not vacuous.Verification
One NVIDIA A40 (sm_86), CUDA 12.8, driver 570.195.03.
PEGAINFER_REQUIRE_GPU=1 cargo test --release -p pegainfer-kernels \ --test hd256_qk_rope_plain_smoke --test hd512_qk_rope_smokePEGAINFER_REQUIRE_GPU=1is not optional: without it a missing device is a skip and the suite passes having executed nothing.cargo fmt --all --checkclean;cargo clippy --release -p pegainfer-kernels --all-targets -- -D warningsexit 0.Negative controls
A GPU-vs-GPU gate can pass because both arms are equally broken, so each retained gate was made to fail on purpose. All three perturbations were reverted afterwards (
git statusoncsrc/clean).+ 1on thekv_head * HDterm inpaged_kv_offset_hd256_plain/paged_kv_offset_hd512pool[59392]: got 0, expected 129(hd256),pool[14848]: got 0, expected 65(hd512)page_origins[token]→0in thePER_TOKEN_METAbranchCUDA_ERROR_LAUNCH_FAILEDat the first D2Hcsr_page_row_checked(..., token, ...)→..., 0, ...per-token pool writes vs the whole-window run[57344]: got -129, expected 0(hd256),[14336]: got 1, expected 0(hd512)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 stays in bounds, no guard fires, and only the bitwise pool comparison can catch it. 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 runs need
--no-fail-fast:cargo teststops after the first failing binary, so without it hd512 never executes and the run looks exactly like a control that passed on both.Notes for review
docs/subsystems/kernels/qk-rope-smoke-oracles.mdrecords the chain, the exactness argument behind the bitwise assertions, and the negative controls with their observed output, so the next person to touch these gates does not have to re-derive why they are shaped this way.pegainfer-gemma4/src/layer_oracle.rsis corroboration, not a successor. It replays the HF golden through the real implementation for both layer types, but it is#[ignore]d and needs the pinned 12B checkpoint, so it never runs by default, and it only ever sees the productionrotary_dim. Retiring an anchor in its favour would move a routinely-runnable check onto one nobody executes.expected_poolis gone from both files.pool_k_offsetdeliberately stays duplicated per file rather than moving totests/common/: paged-KV addressing is one of the three things the issue names, and relocating it would be the move the issue rules out.assert_bits_eq/row_bitsalso stay per file.tests/common/mod.rsis compiled into ten test binaries and all ten usedevice_or_skip; helpers used by two of them would need#[allow(dead_code)]in a module shared by the other eight, which is worse than 15 duplicated lines.v_batchinput and hd512's V is the K=V fork, so neither flat kernel emits a V to compare against.Checklist
docs/conventions/coding-style.md).CLAUDE.md).