feat(gemma4): paged KV serving across both attention families

[FeathBow]

Description

Closes #884.

Serving the 12B tower needs two KV families, not one: local layers are sliding-window attention at head_dim 256 with 8 KV heads, global layers are full attention at head_dim 512 with 1 KV head. One pool with one page geometry cannot hold both, and a request holds pages in both at once. This adds a pool per family, a per-request state in each, and the step machinery over them — prefill, greedy decode, and the admission that keeps the two pools consistent. The pool-write prep (#881) and the windowed attention wrapper (#883) it rides on are already in.

• Admission is atomic across both pools. The page count is the exact frontier account — a ceiling over the post-step KV length per family, minus what the request already holds. Both reservations are taken before either is committed; a refused side drops the other, and both pools sit at their pre-request occupancy. A half-admitted request would otherwise strand the granted family's pages for the request's lifetime.
• The layer forwards take the plan and the pools, never a request's state. Both families' attention plans are built over the same token span; the buffer and page geometry come from the pools the server owns, and everything request-scoped reaches the kernels as plan metadata. Request state is touched only at the step boundary, so moving the resident window later changes the plan builder, not the layer forwards.
• Decode runs the prefill entry at seq_len 1. Correct first; the dedicated windowed decode form is a measurement-gated decision for the batching stage. The greedy loop that drives the reference gate lives with the gate, not in the serving module — production sampling is the frontend's.
• Nothing changes state before the step is admissible. The tokens are checked against the embedding's rows (the kernel indexes it without validating), both families must sit at the same frontier and hold exactly the pages the frontier accounts for, and the window and rope bounds are checked — all before either family advances, because a KvState only ever moves forward.
• Page ids only mean something against the pool that issued them. A request's state is checked to belong to this server's pools before admission commits or the executor indexes a buffer with its pages.
• Model identity is structural, not checked. The server holds the weights its pools, rope tables and layer numbering were built from, so a step cannot be run against a different checkpoint — pages written under one are meaningless under another, and a same-shape mismatch would pass every geometry check. Device identity is pinned on both sides: construction refuses weights that live on a different GPU than the one the pools are allocated on, and a step refuses a context from another device.
• Fail-closed at the window. Cache slots and absolute positions coincide only below the sliding window, and this path refuses to cross.

Test Env

Single GPU (sm_89, x86_64), CUDA 12.9.

Verification

• cargo fmt --all -- --check; cargo build/cargo clippy --release -p pegainfer-gemma4 --features gemma4 --all-targets -- -D warnings; the same pair featureless — all clean. 26 library tests pass.
• Light GPU test: atomic dual-pool admission rollback — a forced partial failure leaves both pools at their pre-request occupancy.
• KV A/B (paged serving vs the no-KV oracle forward, same weights and tokens): prefill argmax equal at every one of 9 positions, max |dlogit| 1.3125; four decode steps fed the oracle's continuation, argmax equal, max |dlogit| 0.90625/0.953125/0.875/0.625 — all under the calibrated 2.0 gate.
• Greedy generation matches HF generate() token for token on three prompts, 50 tokens each. The dumper refuses anything that is not reachable rather than recording it: sdpa and eager must agree on every token, a replay must reproduce the run, and each case is truncated at the first step whose top1-top2 margin falls to the gate. The three admitted cases measured minimum margins of 2.69/5.62/6.81 at dump time — over 26 candidate prompts, exactly three stayed decisive for 50 straight steps, because near-tie steps flip between independent implementations even where HF agrees with itself (flips observed up to margin 1.69). The fixture records the gate and the agreement claim, not the per-step margins themselves.

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: 572758b175

ℹ️ 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

Here are some automated review suggestions for this pull request.

Reviewed commit: c1d74119ae

ℹ️ 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".

[chatgpt-codex-connector]

Require the serving stream, not just its device

When step receives a different DeviceContext for the same GPU—or successive decode steps alternate such contexts—this check passes even though DeviceContext::new_with_device creates a fresh CUDA stream and explicitly disables cross-stream event tracking (pegainfer-kernels/src/tensor.rs:287-297). Work on the new stream can therefore race the constructor's RoPE/pool initialization or a preceding step's KV writes, producing logits from incomplete cache state; require the construction stream/context or add explicit synchronization.

Useful? React with 👍 / 👎.

[FeathBow]

Good catch. Device equality does not establish stream ordering here because event tracking is disabled. I’ll bind the serving path to the construction stream so pool/RoPE initialization and successive KV writes remain ordered without adding per-step synchronization.

[FeathBow]

@codex check again please

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 0edb4dbe83

ℹ️ 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".

[chatgpt-codex-connector]

Commit KV frontiers only after a successful step

When any fallible operation after these lines returns Err—starting with plan_step's GPU metadata allocations and continuing through H2D, scratch allocation, kernel launches, and the logits tail—both frontiers already claim the new tokens even though their KV rows may be unwritten or only partially written. If the caller retries or otherwise reuses this GemmaKv, the next start_pos skips the failed span and attention reads invalid cache contents; defer the frontier commit until success or explicitly roll back/poison the state on every error path.

Useful? React with 👍 / 👎.

[FeathBow]

Great. The current plan construction depends on a prematurely advanced frontier, and an Err leaves reusable-looking state that may describe unwritten KV.

[FeathBow]

@codex review again please

[chatgpt-codex-connector]

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

Reviewed commit: 00ca5cdbdc

ℹ️ 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".