feat: ragged continuous-batching decode graph (#449 M3 Stage 2a-i)

[inureyes]

Summary

The first half of Stage 2 (#449 M3): a ragged decode_step where each row carries its OWN position and length, so sequences of different lengths can share one batch. This is the graph mechanic continuous batching needs. Validated by reference-equivalence, token-exact on CPU and CUDA.

Scope: this is 2a-i (the graph). The admit/evict scheduler with staggered arrival and slot recycling (2a-ii) is the remaining Stage 2a piece. Spike-only (spike/, outside the Cargo workspace), so CI is unaffected.

What is in it

• Emitter emit_decode_ragged (CLI decode-ragged[-argmax] <B>): per-row pos[B]/cache_len[B]; RoPE cos/sin as a per-row gather [B,d] by pos[B]; key mask per-row [B,S] (valid iff s <= cache_len[b]); the KV write unrolled per row, each row writing at its own pos[b] via dynamic_update_slice (constant row/layer offsets, dynamic pos[r]; scatter stays avoided per the Phase 2a CUDA caveat). The attention contractions and the LM head are identical to the uniform-B graph; the per-row mask carries the raggedness.
• Shim xla_llama_decode_ragged (token/pos/cache_len as [B] arrays, threads the rank-5 KV) plus a host-KV-mirror prefill path: xla_llama_ragged_reset, xla_llama_prefill_slot (single-seq prefill, d2h into the mirror slot), xla_llama_commit (h2d the mirror into the resident rank-5 KV). Reuses the scalar prefill vmfb.
• Driver src/bin/llama_ragged.rs: Phase 1 captures one single-seq reference per prompt; Phase 2 prefills B different-length prompts into slots and runs the ragged decode with each slot at its own position; asserts each slot's stream matches its reference.

Validation (reference-equivalence)

Continuous batching has no single reference, so the gate is: B prompts of DIFFERENT lengths (truncations of the 46-token reference prompt), each run independently through the single-seq path, then run together in one ragged batch. Every slot must match its independent reference, proving per-row state is correct and a sequence is unaffected by its batch-mates.

Device	B	Prompt lengths	Result	ragged ms/step	agg tok/s	vs uniform-B
CPU local-task	4	46,43,40,37	4/4 slots 48/48 EXACT	115	34.7	+4.7%
CUDA GB10	4	46,43,40,37	4/4 slots 48/48 EXACT	89	44.8	+1.9%
CUDA GB10	8	46,43,40,37,34,31,28,25	8/8 slots 48/48 EXACT	96	83.6	+7.5%

Findings

• Ragged decode is correct. Eight sequences of eight different lengths in one batch each reproduce their independent single-seq stream exactly; per-row pos/cache_len/mask and the per-row KV write are all right.
• The per-row KV-write unroll is cheap. B*L dynamic_update_slices per step add only ~2% (B=4) to ~7.5% (B=8) over uniform-B, and lower fine on CPU and CUDA (compile ~5-7.5s).
• CPU large-B validation is reference-bound (Phase 1 runs B single-seq references sequentially), a harness cost, not a graph cost; the B sweep runs on CUDA.

Full tables in spike/iree-ffi/FINDINGS_ragged.md; the full Stage 2 plan in spike/openxla/STAGE2_DESIGN.md.

Reproduce

cd spike/iree-ffi
EMIT=../rust-emitter/target/release/emit; IC=./iree-dist/bin/iree-compile
F="--iree-input-type=stablehlo --iree-hal-target-device=local --iree-hal-local-target-device-backends=llvm-cpu"
$EMIT prefill-argmax p.mlir && $IC $F p.mlir -o p.vmfb
$EMIT decode-argmax s.mlir  && $IC $F s.mlir -o s.vmfb
$EMIT decode-ragged-argmax 4 r.mlir && $IC $F r.mlir -o r.vmfb
IREE_DIST=./iree-dist cargo run --release --bin llama_ragged -- \
  --batch 4 --device local-task --prefill p.vmfb --sdecode s.vmfb --decode r.vmfb

Next

2a-ii: the admit/evict scheduler (request queue, mid-stream prefill into a freed slot, eviction on EOS, slot recycling, staggered-arrival validation). Then 2b productizes the engine into mlxcel-xla.

Refs #449