Forward Cuda Gemma

pub fn decodeStep(self: *ForwardGemma, token: u32, pos: u32, run_layers: bool) !u32

One greedy decode step for `token` at sequence position `pos`.

pub fn prefillStep(self: *ForwardGemma, token: u32, pos: u32) !void

Prefill helper (mirrors ForwardCuda.prefillStep): run every layer to build the KV cache, but SKIP the tail rms_norm + LM head + argmax — a prompt-internal token's logits are never read.

Saves the vocab-sized head matvec on T-1 of the T prompt tokens; the MoE gemma model (small active forward, full head) benefits most. Async layer ops drained here; bit-identical generation.

pub fn prefillBatched(self: *ForwardGemma, tokens: []const u32) !u32

Batched dense-gemma prefill: process ALL T prompt tokens at once, reading each weight ONCE for all tokens via the gemm_*_tiled_v2 register-blocked GEMMs (5.9× over the per-token matvec). Returns the last token's argmax (= the first generated token), exactly like running prefillStep on tokens [0..T-1] then decodeStep on the last. ADDITIVE: builds its own token-major scratch and never touches the single-token decode path.

Phase-1 scope is the DENSE gemma-31b (n_experts==0). Phase-2 cycle 1 adds the 26b MoE: its ATTENTION block is the SAME structure as the dense model, so it shares the batched attention path (GEMM Q/K/V/O + batched causal attn + batched norm/RoPE/KV-write) — bit-identical to the per-token attentionLayer. Its routed-expert FFN is still LOOPED per token (the existing single-token moeFfnBlock, fed each token's hidden slice via an alias-swap) because the FFN is position-independent → looping it is output-identical; full batched-expert routing (group T tokens by expert) is a later cycle.

pub fn decodeBatch(self: *ForwardGemma, tokens: []const u32, positions: []const u32, slots: []const u32, out_tokens: []u32) !void

Batched DECODE step (DENSE gemma only): advance B independent sequences by ONE token each in a SINGLE B-row forward. The projections + FFN reuse the batched-prefill GEMM path — `gemmDispatch`/`ffnBlockBatched` read each weight ONCE for all B rows (the launch / weight-bandwidth amortization that makes batching a throughput win). The attention inner (per-head V/Q/K norm + RoPE + KV-write + causal softmax) LOOPS per row through the proven single-token kernels (`rms_norm_rope_qkv` / `gemma_attention`) against each sequence's OWN KV slot at its OWN position — so row b attends over slot `slots[b]`'s history [0..positions[b]] only. Fusing that per-row loop into one batched-seq attention kernel is sub-step 1c (a perf step; the looped form here is already correct for mixed positions). ADDITIVE — the production `decodeStep` / `prefillBatched` paths are untouched.

Requires `allocSlotKv(n_slots, slot_ctx)` first, with n_slots > max(slots) and slot_ctx > max(positions). Each sequence b's slot must already hold its KV history for [0..positions[b]-1] (written by prior `decodeBatch` calls — e.g. a per-token B=1 prefill into the slot). Writes b's new K/V at its slot.

tokens[b] input token for sequence b this step positions[b] sequence b's current position (slot holds [0..pos-1]) slots[b] sequence b's KV slot index out_tokens[b] greedy argmax for sequence b (caller advances pos + feeds back)

pub fn allocSlotKv(self: *ForwardGemma, n_slots: u32, slot_ctx: u32) !void

Allocate slot-based KV for `n_slots` concurrent sequences of up to `slot_ctx` positions each.

ADDITIVE: a fresh per-layer allocation that never aliases or touches the production single-sequence kv_k/kv_v. The batched-decode path (sub-steps 1b/1c) writes/reads it via `slotKvOffsetBytes`; the production decodeStep is unchanged. Idempotent — frees any prior slot KV first. Sub-step 1a only allocates + smoke-tests it.

pub fn slotKvOffsetBytes(self: *const ForwardGemma, L: u32, slot: u32, pos: u32) usize

Byte offset of sequence-slot `slot`'s K/V for position `pos` in layer `L`'s slot KV buffer: (slot*slot_ctx + pos)*kv_dim(L)*sizeof(f32).

This is the exact indexing the 1c per-sequence kv-write + slot attention kernels use.

pub fn slotKvSmoke(self: *ForwardGemma) !bool

Sub-step 1a plumbing smoke: prove the slot-KV offset arithmetic round-trips and that distinct (slot,pos) pairs map to NON-overlapping device regions.

Writes a sentinel into (slot 0, pos 0) and a distinct pattern into the LAST (slot, pos) of layer 0's K cache, reads both back, and checks neither write clobbered the other. Returns true on success. Requires allocSlotKv first.

pub fn ffnLayerPub(self: *ForwardGemma, L: u32) !void

FFN block dispatched exactly as decodeStep: routed MoE when this layer carries a router, dense GeGLU otherwise.