gahow/xtrain
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

docs: M2b — batched decode results (token-identical + ~1.7x rollout, device-cache next)

Browse Source 
Implementation log (docs/18) + Phase-3 row (evolution.md): rope_pos primitive +
gate, the batched engine (decode_attention/repeat_kv reused), the token-
identical batch gate, and the measured ~1.7x rollout-inclusive step speedup +
memory stabilization. Closes the M2 decode engine (M2a single-seq + M2b
batched); names the device-side cache as the remaining lever.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-06-30 17:20:01 +08:00
parent 361c5290fa
commit 0f76c0fdb0
2 changed files with 34 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
docs/18-post-training-rl-sft.md
View File

@@ -522,3 +522,35 @@ leash wired, format held); the held-out flatness + the two memory/throughput wal
reported findings. The honest end-state of the post-training arc: **a complete, correctness-gated
SFT → KV-cache → DPO → GRPO stack** — the infrastructure learned in full, with measured, honest
limits on what alignment can do for a capability the base model lacks.
### M2b — batched KV-cache decode (landed; completes the M2 engine, fixes the rollout long-pole)
Built after M4 (where the rollout long-pole bit hardest): decode the **G samples of one prompt in
lockstep** — one forward per step over the whole group → G× fewer kernel launches, the deferred
fix from M2a.
**One new primitive:** `rope_pos(x, positions[])` — RoPE with a *per-row* absolute position (new
forward-only kernel), since the G batched rows share one decode position (M2a's `rope_at` does
`pos0 + row`, wrong for a batch at a single position). **Gate:** bit-identical to the full rope
for positions `[0..n]`, and to `rope_at(P)` per row for a uniform `P`.
**Engine (`generate_cached_batch`):** `BatchKVCache` carries a G dimension (`[T, G·num_kv, hd]`
host-accumulated → `[G·num_kv, T, hd]`); the batched `decode_step` threads G through embed /
projections / QK-norm / `rope_pos` / cache. Two M2a pieces drop in unchanged: `decode_attention`
is already batch-agnostic (`bh = G·nh`), and `repeat_kv(nh, batch=G)` broadcasts per group. No
finished-mask (all G generate `max_new`; the caller cuts at EOS) and no ragged-length prompts yet
— both perf-only follow-ups.
**Gate (token-identical):** all G **greedy** rows are byte-identical to the single-sequence decode
(`tests/decode_batch.rs`, 8 query / 2 kv heads → exercises the `repeat_kv` batching) — pins that
G-way batching indexes each sequence's K/V with no cross-row contamination.
**Throughput (v12 1.05B, G=6·B=6, easy task, rollout wired into `train_grpo`):** ~8.5 s/step vs
~14–16 s/step for the single-seq cached rollout — **~1.7×**, rollout-inclusive. Short of the full
G× because (a) the per-token-logp forwards + the PG update also cost, and (b) the M2a per-layer
**host round-trip** is still there (now G× the data in one transfer, not removed). The full
device-side cache (no host round-trip) is the remaining decode-engine optimization. Batching also
**stabilises memory**: one batched forward per step vs G separate allocations that fragmented the
caching allocator (the M4 OOM). So M2b closes the decode-engine milestone (M2a single-seq + M2b
batched) and turns the rollout long-pole from "OOM/unbounded" into a bounded ~1.7× win — measured,
with the device-cache as the named next lever.