perf(gemma4): regular decode steps stop rebuilding what did not change

docs/models/gemma4/serving.md

@@ -27,6 +27,16 @@ The register router accepts exactly 128 experts and from 1 through 32 picks. A n
 
 The checkpoint-backed `the_routed_block_matches_the_reference_formulas` gate owns the scratch-capacity, companion-route and coarse-block evidence, including a narrow block replay after a coarse block on one scratch. On shared rows, it proves that the 16-row and 64-row block pick the same experts with the same weight bits and produce the same gate, weighted-down and block bits. `router_topk_matches_the_exact_128_expert_contract` owns the register-router boundary and non-finite rows. The kernels-owned `kimi_marlin_align_boundary_matches_vllm_contract` oracle owns stable counts, offsets, padding and expert-local order on both sides of the alignment dispatch boundary. `scripts/gemma4_gates.sh` owns the Gemma crate's ignored gates and the kernels crate's Gemma router contract (the device-only test under the `gemma4` feature), holding both crates' ignored sets against its manifest; the Kimi alignment oracle needs the `kimi-k2` feature and an `sm_90` device and is run by hand.
 
+## The decode pipeline
+
+Greedy decode rounds run a depth-two software pipeline on the base stream: a step's argmax writes its picks straight into the id buffer the next step's embedding reads, its readback lands in one of two pinned slots, and the emitted token stream lags compute by one step. The production invariant is that token ids, finish reasons and token counts are identical to collecting every step at the same batch composition, and that the kernels inside a step and their order do not change.
+
+A batch is eligible only when every row is greedy, scores no logprobs, sits at least two tokens from its length cap and is not already stopping. While a successor step is in flight the active row order is pinned: the successor's ids were written on the device for the current order, so no row may be retired, reordered or added until that step is collected. A stop found late marks the row stopping and it retires when the pipeline drains. The pipeline drains before a synchronous or mixed admission that will use the step arena or change the roster, a lane join, a cancelled row, or a batch that stopped being eligible. An admission that turns out not viable — a closed or invalid request, a page-shortage requeue, or a failed Scheduled send — and an async lane launch keep it; the lane join drains. A solo admission into an idle engine has nothing to drain but must drop the fingerprint the retired roster left, since it starts a new roster with new page identities. The collect-every-step path takes over after a drain.
+
+A regular decode step, every row one token further with its page, chunk and split structure untouched, also skips its metadata rebuild and every upload: a kernel captured at the decode graph's tail advances the per-row tables in place, and a fingerprint of the previous step proves the advanced device state is what a full upload would write. Mixed steps, page turns, chunk boundaries, admissions and the precapture warm pass invalidate the fingerprint and rebuild as before.
+
+The staged sampler chain — suppression, argmax and the device copy of the picks — is captured per bucket at startup and launched as one graph beside the decode replay; the pinned readback stays outside so the collector keeps the copy's own event. The staged path runs on the base stream only and refuses a stream override. If a step fails after its device work was enqueued, the engine synchronizes the stream before it drops the batch and returns its pages, and aborts if the device cannot be synchronized, the policy the prefill lane already applies.
+
 ## The two pools
 
 Gemma 4 runs two attention families with different KV shapes, so the budget is two budgets. With 16-token pages, `C = ceil(8192/16) = 512` context pages and `W = ceil(1024/16) + 1 = 65` window pages:

pegainfer-core/src/ops.rs

@@ -24,6 +24,7 @@ pub use pegainfer_kernels::ops::SuppressIds;
 pub use pegainfer_kernels::ops::accumulate_bf16_token_scaled_to_f32_into;
 pub use pegainfer_kernels::ops::add_batch;
 pub use pegainfer_kernels::ops::add_batch_into;
+pub use pegainfer_kernels::ops::advance_decode_metadata;
 pub use pegainfer_kernels::ops::argmax;
 pub use pegainfer_kernels::ops::argmax_batch_bf16_into;
 pub use pegainfer_kernels::ops::batch_prefill_paged_hd512_into;

pegainfer-core/src/ops/paged_plan.rs

@@ -223,6 +223,9 @@ impl PrefillPagedPlan {
     pub fn kv_chunk_size_d(&self) -> &CudaSlice<i32> {
         self.inner.kv_chunk_size_d()
     }
+    pub fn decode_metadata_d_mut(&mut self) -> (&mut CudaSlice<i32>, &mut CudaSlice<i32>) {
+        self.inner.decode_metadata_d_mut()
+    }
     pub fn total_num_rows_d(&self) -> &CudaSlice<u32> {
         self.inner.total_num_rows_d()
     }