feat(kernels): pool-write prep for Gemma 4 serving at both head dims

[FeathBow]

Description

Fixes #880

The write-side prep entries the Gemma 4 serving path needs, at both head dims:

• hd256 (local layers, real v_proj): qkv_norm_rope_paged_prefill_hd256_plain_into, the serving form of feat(kernels): plain-w QK-norm and RoPE prep at head_dim 256 #853's oracle-form prep — Q normed + rotated into a contiguous q_out, K normed + rotated and V weightless-normed (never rotated) scattered straight into the layer's K/V pool blocks in one kernel, no intermediate scatter.
• hd512 (global layers, K=V fork): the prefill prep now writes the V block alongside K — the weightless RMS norm of the same raw K, sharing its inv_rms denominator, in the same pass. No separate v_proj input and no D2D fork copy; the fork semantics are the ones feat(gemma4): global decoder layer and its HF golden oracle probes #877 pinned against the reference.

Both wrappers derive their pool geometry through one checked helper: the layout's fields are public, so the block and stride arithmetic is re-derived with overflow-checked multiplies, asserted against the layout, and narrowed through checked conversions before anything reaches the unsafe launch; each wrapper additionally checks that its q/kv widths, tensor extents and rope-table length fit the 32-bit arithmetic the kernel indexes with — a contradictory or wrapping layout fails as an Err, never as a release-mode wrap or a divide-by-zero. On top of that sit the layer bound, whole-page pools, page coverage of [start_pos, start_pos + seq_len) (checked addition), position range against the rope tables, table lengths and vec-backed weights; the kernels __trap() on out-of-range positions or page ids as the second layer of that defence. Slot positions derive from start_pos + token — absolute and cache-relative coordinates coincide below the sliding window, which is the entries' documented contract; window-crossing callers need a separate slot mapping first.

Test Env

Single GPU (sm_89, x86_64), CUDA 12.9, PEGAINFER_REQUIRE_GPU=1.

Verification

• The extended hd512 smoke pins the new behaviour: the V block equals the weightless norm of the raw K under the shared denominator, against a host-side bf16 closed-form reference (0.02 tolerance), all offsets derived from the layout rather than hand-picked, and the negative coverage sweeps every region the kernel must never write — unreferenced pages, the other layer, and slots outside the request's positions.
• The hd256 pool-write case pins the serving entry: Q matches the closed-form bf16 reference of the contiguous cases, the K blocks carry the rotated sign patterns and the V blocks the weightless norms at their layout-derived offsets, and the untouched sweep covers unreferenced pages, the other layer and out-of-range slots. The V input's global sign flip makes a V/K source mix-up land a wrong sign in every slot.

Type of Change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation update

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: f42cb4cc84

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[FeathBow]

@codex review please

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. Nice work!

Reviewed commit: f42cb4cc84

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".