gemma4: serving prep needs pool-write forms — hd256 QKV in one pass, hd512 V alongside K

[FeathBow]

Description

The read side of two-family KV serving is stocked: windowed and paged attention entries exist for both head dims. The write side is not. The hd256 plain prep (#853) has only its oracle form — Q and K land in contiguous token-major buffers — so a serving path would run norm+RoPE, then a separate per-token scatter into the paged pool. And the hd512 prep writes K into the pool but not V: the K=V fork resolved in #877 (value is k_proj's raw output, weightless-normed, never rotated) would still cost a D2D fork copy plus a separate norm pass every step.

Both gaps close in the prep kernels themselves. For hd512 the V block is the weightless RMS norm of the same raw K the kernel is already holding — it can share the K norm's denominator and write both blocks in one pass. For hd256, where local layers carry a real v_proj, one kernel can norm+rotate Q to a contiguous output while norming+rotating K and weightless-norming V straight into the layer's pool blocks.

Proposed Plan

• Add qkv_norm_rope_paged_prefill_hd256_plain_into: Q normed + rotated into contiguous q_out; K normed + rotated and V weightless-normed (never rotated) scattered directly into the layer's K/V blocks — one kernel, no intermediate scatter.
• Extend the hd512 prefill prep to write V alongside K: the weightless norm of the same raw K, sharing its inv_rms, in the same pass — no separate v_proj input, no fork copy.
• Validate at the wrapper through one checked geometry helper: layout self-consistency (block/stride arithmetic re-derived with overflow-checked multiplies), layer bound, whole-page pool, page coverage of [start_pos, start_pos + seq_len), position range against the rope tables, table lengths, vec-backed weights. The kernel __trap()s on out-of-range positions or page ids as the second layer of that defence.
• Pin the slot contract on the entry: slot positions derive from start_pos + token, so absolute and cache-relative coordinates coincide below the sliding window — window-crossing callers need a separate slot mapping first.
• Extend both smokes: the hd512 suite asserts the V blocks and rewrites its negative coverage; the hd256 suite gains a serving-form case at the 12B local geometry — Q against the contiguous oracle, K and V blocks at their layout-derived pool offsets, and an untouched-regions sweep.

Acceptance Criteria

• The hd512 smoke gates the V block values (the K=V fork under the shared denominator) and the untouched regions, with offsets derived from the layout rather than hand-picked.
• The existing hd256 and hd512 suites stay green; the trap contracts hold.
• The hd256 smoke's pool-write case gates the serving entry end to end: Q pairing, K rotation and V weightless-norm placement, layer offset, and every region the kernel must never write.