Pull requests / #1599

#1599 Layer split: batch the draft layer's prompt K/V on the last stage, +3 to 8% prompt speed (opt-in, `STRATA_SPLIT_MTP_BATCH=1`)

open · @noon-at-cgn · 0 comments · View on GitHub

BenchmarksMulti-GPUAMD / HIPNVIDIA / CUDADocumentation

Description

## Title
Issue: Related: none found (searched open PRs 2026-10-08; closed #453 added E-9 for a ring K/V on one GPU).

## Summary
**Why it matters:** on a layer split, every prompt chunk ends with the draft layer's own pass, which blocks the pipeline (0.9-1.3 s per request in the host's "after each chunk" time). One GPU already avoids that with the batched E-9 pass (`Prefill::draft_kv`), but `generate.cpp` only calls it when `!multi_gpu`. With `STRATA_SPLIT_MTP_BATCH=1` the last stage runs E-9 too: that time drops to 0.1-0.25 s, and long prompts read faster with the same greedy output.

| 2x RTX 3080, layer split 23, pooled medians, 2 restarts per arm | env off | env on |
|---|---|---|
| prompt read 25k tokens | 1996 tok/s | 2129 (+6.7%) |
| 51k | 2448 | 2553 (+4.3%) |
| 104k | 2846 | 2940 (+3.3%) |
| greedy output, 6 prompts (5 x 256 tokens + one 25.5k prompt) | - | identical, 6/6 in all five comparisons, A/A included |
| drafts accepted | 826 of 1210 (68.3%) | 826 of 1210 |

Default off; the default path is unchanged.

## What changed
- `src/program/generate.cpp`: `on_chunk` asks `draft_kv_plan::plan()` which path to try; on a split with the variable on it calls `draft_kv` on the drafter's stage. One log line per process says what the variable did.
- `include/strata/program/draft_kv_plan.hpp` (new): the choice as a pure function. `src/program/draft_kv_plan_test.cpp` + `CMakeLists.txt`: CPU test (18 checks).
- `src/prefill/prefill.cpp`, `include/strata/prefill/prefill.hpp`, `sycl/src/prefill/prefill.cpp`: `draft_kv` makes the stage's device current and takes an optional `why` that names the reason when it declines.
- `tools/split_mtp_batch_parity.py` (new), `docs/MULTI_GPU.md`: one paragraph.

## Extra Notes
Measured on engine 0.1.40.3 (d5ea713) plus this change; the branch itself is cut from main fb58e0d.
Machine: 2x RTX 3080 20 GB (220 W), Xeon E5-2696 v4, UD-Q4_K_XL, `--batch-mtp`, 2 slots, `--prefill auto:16384`. Same binary in both arms, only the env differs; full restarts in the order off/on/off/on. All 8 on/off prompt-read ratios are above 1 (1.03 to 1.08); the two off restarts differ by at most 2% at 25k and 0.3% elsewhere. Parity ran with `--adapt-every 0 --suffix-draft 0` (with the adaptive tier on, greedy output is not repeatable across restarts).
**Open question, stated plainly:** decode is not shown to be unchanged. The on medians are about 3% lower in both pairs (pooled solo 82.6 -> 80.2 tok/s, c2 93.2 -> 90.8), but the restart-to-restart control moves up to 4.7% and no mechanism is known (the patch touches only the prompt path). With two restarts per arm it can be neither confirmed nor excluded.
The drafter's K/V come from Q8_1 x Q8_0 MMQ instead of its own mmvq, so they are not bit-identical; no acceptance difference was visible at this n. Needs a `--native` pack, otherwise it logs why and keeps the old pass. Checked: `draft_kv_plan_test`, CUDA build. Not built: HIP, SYCL (only the new parameter). Not tested: AMD, three or more cards, without `--batch-mtp`, the decline path.

Related on strata.com

Editorial links to help you install, pick models, or read release notes — not part of the upstream thread.