Description
Gemma 4 caches at two head dims — 256 for its sliding layers, 512 for its full-attention layers — and the shared KvLayout applies one head dim uniformly to every layer. Heterogeneous buffers over one block pool is already an established pattern here, so that part is not the problem. What is new is that the two groups want different lifetimes: the sliding group needs to release pages the full-attention group is still reading, and a request holds one assignment list, so a block is held or released for both groups at once.
Two kernel prerequisites block correct serving rather than one. The sliding mask is required because sliding layers must not attend past 1024 positions. Paged attention instantiated at head_dim 512 is required because the full-attention group's pages are otherwise unreadable — the shared translation unit compiles nothing at that width. Both must exist before the model serves correctly, and the mask needs an explicit conversion — Gemma counts the current token inside its window while FlashInfer's window_left is an inclusive distance, so 1024 has to be passed as 1023 or the model attends one token too far. With both in place and no reclamation the model is correct and budget-limited, reaching tens of thousands of positions; the declared 262144 is what it cannot reach, at 84 GiB for one unreclaimed request at the smallest supported size. Reclamation is not purely a capacity exercise either: it splits an array that today carries two meanings, which is a correctness problem of its own.
Several decisions here constrain work that follows and are cheaper to settle now than to discover: whether K and V can share storage, how tensor parallelism shards a group with a single KV head, what CUDA graph capture has to duplicate, and how the prefill chunk size becomes a capacity parameter rather than a latency knob.
Proposed Plan
- State the two group layouts with their per-token and per-request costs at all three sizes, derived from the published configs rather than from one size.
- Separate the two kernel prerequisites from the capacity work, and say what a correct engine can serve before any eviction exists.
- Enumerate what page eviction breaks, each read in the tree — including what it does not break, so the implementation does not chase phantom obstacles.
- Decide the pool topology with the failure mode that rules out the simpler option stated structurally rather than as a race.
- Settle
attention_k_eq_v, the per-group admission and reservation arithmetic, the tensor-parallel sharding rule, what graph capture duplicates, the RoPE budget, and the explicit exclusions.
Acceptance Criteria
- Every structural number is stated for all three sizes and is reproducible from the published configs.
- Every claim about existing engine behaviour cites the code it was read from; claims that do not survive checking are corrected rather than repeated, including the ones the document previously made.
- Correctness prerequisites and capacity work are costed separately, so a reader can tell which is which.
- Where the contract says an existing mechanism is missing, it is missing — not merely unfamiliar.
- The document says what it cannot settle and what would settle it.
Description
Gemma 4 caches at two head dims — 256 for its sliding layers, 512 for its full-attention layers — and the shared
KvLayoutapplies one head dim uniformly to every layer. Heterogeneous buffers over one block pool is already an established pattern here, so that part is not the problem. What is new is that the two groups want different lifetimes: the sliding group needs to release pages the full-attention group is still reading, and a request holds one assignment list, so a block is held or released for both groups at once.Two kernel prerequisites block correct serving rather than one. The sliding mask is required because sliding layers must not attend past 1024 positions. Paged attention instantiated at head_dim 512 is required because the full-attention group's pages are otherwise unreadable — the shared translation unit compiles nothing at that width. Both must exist before the model serves correctly, and the mask needs an explicit conversion — Gemma counts the current token inside its window while FlashInfer's
window_leftis an inclusive distance, so 1024 has to be passed as 1023 or the model attends one token too far. With both in place and no reclamation the model is correct and budget-limited, reaching tens of thousands of positions; the declared 262144 is what it cannot reach, at 84 GiB for one unreclaimed request at the smallest supported size. Reclamation is not purely a capacity exercise either: it splits an array that today carries two meanings, which is a correctness problem of its own.Several decisions here constrain work that follows and are cheaper to settle now than to discover: whether K and V can share storage, how tensor parallelism shards a group with a single KV head, what CUDA graph capture has to duplicate, and how the prefill chunk size becomes a capacity parameter rather than a latency knob.
Proposed Plan
attention_k_eq_v, the per-group admission and reservation arithmetic, the tensor-parallel sharding rule, what graph capture duplicates, the RoPE budget, and the explicit exclusions.Acceptance Criteria