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

docs: M2c — device KV cache + the bottleneck-shift finding

Browse Source 
Implementation log (docs/18) + Phase-3 row (evolution.md): cat_seq device cache,
gates hold (token-identical), and the profile-first finding — ~10% single-seq
decode but no GRPO-step change because the long pole shifted to the per-sample
logp/PG forwards after M2b batching. Names ragged batched prefill as the next
decode lever.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-06-30 17:39:10 +08:00
parent 3a3425960c
commit 41d46208a6
2 changed files with 24 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -554,3 +554,25 @@ device-side cache (no host round-trip) is the remaining decode-engine optimizati
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.
### M2c — device-side KV cache (landed; the bottleneck moved, a profile-first finding)
The named M2b follow-up: keep K/V on the GPU (`[bh,T,hd]`, an `Option<Tensor>` per layer) and
grow it by one token per step via a new `cat_seq` kernel (concat along the seq dim) — removing the
M2a/M2b per-layer **host round-trip** (`to_cpu`/`from_slice`/re-upload) *and* the `transpose_3d01`.
Both single-seq and batched decode refactored to it (cleaner than the host `Vec` + rebuild).
**Gates hold:** `cat_seq == host concat`; `decode_kv` single-seq + `decode_batch` G-way both still
**token-identical**; GQA training path unaffected.
**The finding (why this is a measure-first lesson, not a speedup story):** removing the host
round-trip buys **~10%** on *pure* single-seq decode (133 → 147 tok/s @128) but **does not move the
GRPO step** (~8.5 s/step, unchanged). Because after M2b batching, the rollout is no longer the
step's bottleneck — the per-sample **`per_token_logp` captures** (2 forwards/sample) and the
**PG-update** forwards+backwards (`model.forward`, full-sequence, per sample) now dominate. So the
long pole **shifted** from the rollout to the training-side forwards (cf. T11/T17/M2a: profile
before optimizing — the bottleneck you fixed is not the one that remains). The device cache is
still a real, correctness-gated improvement (cleaner code, less PCIe, ~10% decode); the honest
headline is that the *next* decode lever is **ragged batched prefill of the per-sample forwards**,
not the cache. The M2 decode engine is now M2a (single-seq) + M2b (batched) + M2c (device cache),
all token-identical-gated; the post-training stack remains complete with its bottleneck mapped.