Problem
The HD256 and HD512 QK/RoPE smoke suites maintain a second Rust implementation of substantial GPU semantics.
At 69c4e149:
pegainfer-kernels/tests/hd256_qk_rope_plain_smoke.rs is 783 lines and has four closed-form tests covering RMSNorm, RoPE pairing/tail behavior, paged-KV address calculation, and pool writes.pegainfer-kernels/tests/hd512_qk_rope_smoke.rs is 951 lines and has three analogous closed-form tests.- Both files independently define helpers such as
inv_rms, normed, expected_prep, expected_full, expected_pool, and assert_pool.
These tests have value, but seven large host-side reconstructions create a parallel implementation that must evolve with the CUDA contract. A shared specification mistake can also be repeated in both the kernel and its local oracle.
The newer row-offset tests are a better shape: compare the offset GPU arm against the same production GPU path at zero offset, then verify untouched prefix/suffix/KV sentinels. They should remain.
Direction
Keep a small numerical anchor for each genuinely distinct operator contract. Move broader numerical correctness to externally generated fixed fixtures or model-level HF golden gates. Use GPU-vs-GPU metamorphic checks and sentinels for row ownership, offset, and untouched-memory invariants.
Do not replace the current duplication with a large shared Rust oracle module; that moves the second implementation rather than removing it.
Acceptance criteria
- Inventory the seven current closed-form tests and state which independent failure each retained test catches.
- Reduce the host-side RMSNorm/RoPE/paged-KV reconstruction to the minimum numerical anchors needed for HD256 and HD512.
- Any replacement fixed fixture records the external producer and exact geometry; expected outputs are not regenerated by Rust code that mirrors the CUDA operator.
- Keep the decode/prefill row-offset metamorphic gates, including untouched sentinel and KV-pool checks.
- Keep host-contract refusal coverage for invalid rotary dimensions, positions, and undersized pools without adding operator math to those tests.
- Do not add a shared general-purpose CPU QK/RoPE implementation.
- Run the retained HD256 and HD512 GPU gates on one supported GPU and include the exact test list and results in the PR.
The target is less duplicated semantic code with equal or better confidence, not a file-only reorganization.
Contributor guardrails
- The PR description must include a before/after table: closed-form tests, host-oracle helpers, lines removed, retained invariants, and real GPU gates executed.
- Deleting a test counts only when its independent failure surface is covered by a retained external golden, production-path E2E, refusal contract, or GPU-vs-GPU metamorphic gate. A lower line count alone is not completion.
- Do not move the existing formulas into
common, a new crate, generated Rust, or a parameterized mega-helper. The semantic implementation must be removed, not relocated.
- Do not add mock kernels, fake device buffers, source-string assertions, or tests that only validate the new helper plumbing.
- A replacement fixture must have reproducible external provenance. Hand-entered constants or values generated by the Rust implementation under test are not an independent oracle.
- Keep the change limited to test semantics and private test support. Production dispatch, kernel code, and performance tuning are out of scope.
Problem
The HD256 and HD512 QK/RoPE smoke suites maintain a second Rust implementation of substantial GPU semantics.
At
69c4e149:pegainfer-kernels/tests/hd256_qk_rope_plain_smoke.rsis 783 lines and has four closed-form tests covering RMSNorm, RoPE pairing/tail behavior, paged-KV address calculation, and pool writes.pegainfer-kernels/tests/hd512_qk_rope_smoke.rsis 951 lines and has three analogous closed-form tests.inv_rms,normed,expected_prep,expected_full,expected_pool, andassert_pool.These tests have value, but seven large host-side reconstructions create a parallel implementation that must evolve with the CUDA contract. A shared specification mistake can also be repeated in both the kernel and its local oracle.
The newer row-offset tests are a better shape: compare the offset GPU arm against the same production GPU path at zero offset, then verify untouched prefix/suffix/KV sentinels. They should remain.
Direction
Keep a small numerical anchor for each genuinely distinct operator contract. Move broader numerical correctness to externally generated fixed fixtures or model-level HF golden gates. Use GPU-vs-GPU metamorphic checks and sentinels for row ownership, offset, and untouched-memory invariants.
Do not replace the current duplication with a large shared Rust oracle module; that moves the second implementation rather than removing it.
Acceptance criteria
The target is less duplicated semantic code with equal or better confidence, not a file-only reorganization.
Contributor guardrails
common, a new crate, generated Rust, or a parameterized mega-helper. The semantic implementation must be removed, not relocated.