Pull requests / #2985

#2985 batch-groups: MTP drafts in the pipelined slot groups (STRATA_PIPED_MTP, opt-in)

open · @JacobSussan · 0 コメント · GitHub で見る

BenchmarksMulti-GPUAMD / HIPNVIDIA / CUDADocumentationWindowsLinux

本文

Closes #2973.

## Summary

With a layer split, `--batch` pipelines the slots in groups (one group per stage), but the pipelined path runs no MTP drafts, and `--batch-mtp` with serial windows leaves all but one card waiting (on four stages it lost 30-40% of the four streams' total against the pipelined default).

`STRATA_PIPED_MTP=N` (with `--batch-mtp`) keeps the pipeline and drafts in it: each group's window carries its slot's token and up to N MTP proposals (at most `--spec` - 1; the first always, each further one while the draft layer was at least `--spec-min-p` sure of the one before, the solo path's rule). It needs one slot per group (as many slots as stages, which `--batch-groups auto` gives for that count) and is left off, with a line at start, otherwise. Default off: nothing changes without the variable.

## How it works

How many rows to keep is known only when the last stage has picked, and by then the earlier stages already run other groups' windows, which reuse the scratch the commit replays from.

- `Verifier::batch_launch_rows`: the group's window, then a small save graph that copies the rows the commit replays (each GDN layer's conv and recurrence inputs, each QSA layer's raw indexer keys, the PLE history snapshot) to the group's own hand-off rows of a per-stage buffer (`init_deferred_commit`). No commit.
- `Verifier::batch_commit_rows`: after the verdict, every stage queues the commit of the kept rows on its own stream (no host wait). It runs before that group's next window on the stage, and the GDN state and indexer tail are untouched until then, as for a solo window. The accepted prefixes come from a per-slot mapped area that is written again only at that slot's next verdict, after its next window has passed the stage.
- The slot's drafter (the `--batch-mtp` slot drafters) catches up on the window's final rows and drafts the next proposals as an asynchronous chain (`chain_launch`, so the pipeline goes on meanwhile). The rows are copied on the last stage's stream, ahead of its next window, and the chain waits on an event after the copy. A group whose chain has not landed lets another group start first.
- A drain (and the moment every slot is done) waits for the queued commits and chains, so an admission, a parked conversation or the solo path taking a slot back sees the final state.
- With it on, the stages keep the window size they have without `--batch-mtp`. `--batch-mtp` alone sizes them for 8-row serial windows, which also lets a lone request's windows grow; here a lone file edit then measured 160 against 178 tok/s.
- `batch_groups::resolve` gains `mtp_piped` (with unit tests), and docs/BATCHING.md a section.

## Exactness

`tools/batch_test.py` (4 slots on a 4-GPU split, 200 and 300 tokens, `STRATA_IQ_MT_MIN=1 --pcie-frac 0 --adapt-swaps 0`): every slot decodes exactly the tokens it decodes alone, for `STRATA_PIPED_MTP` = 1, 2 and 3, and for the default path with the same binary. Windows per slot for 1131 tokens: 1131 (default), 652 (N = 1), 506 (N = 2), 455 (N = 3).

## Measured

4x RTX 3090 (PCIe 3.0 x8, no P2P), 2x Xeon Gold 6126, Linux, CUDA 13.2, UD-Q4_K_XL, split 11,24,36, `--trim-stage-weights`, `"parallel": 4` (4 groups of 1), `--spec 4 --spec-min-p 0.70`, FP16 K/V, `STRATA_PIPED_CONCURRENT=1`. Four concurrent 900-token answers through the HTTP server, reasoning off, total output tok/s:

| | default | `=1` | `=2` | `=3` |
| --- | ---: | ---: | ---: | ---: |
| code (median of 3) | 203 | 244-269 (2 runs) | 270 | 259 (1 run) |
| prose (median of 3) | 209 | 230-240 (2 runs) | 216 | 215 (1 run) |
| a file returned with one rename (1 run; mostly its 3K-token prompts) | 129 | 140-142 | 144 | 153 |
| one request alone: code / edit / prose | 134 / 175 / 96 | unchanged | unchanged | unchanged |

Proposals kept at N = 1: code 88%, the edit 97%, prose 62%. On this rig a two-row window costs about 1.5x a one-row window (`STRATA_VERIFY_PROFILE`: the routed experts of two tokens, 0.44 -> 0.74 ms, and the host's plan on the stages that miss experts), which is why prose gains little. Cost: the saved rows (25 MiB per stage at N = 1 with 4 slots, 51 MiB at N = 3), the slot drafters as `--batch-mtp` has them, and the extra graphs per row count.

## Not run

AMD (HIP), the draft layer on the first stage (`mtp_first`: the row copy and its event would then cross devices), helper expert caches, more than one slot per group.

関連リンク

インストール・モデル・リリースへの站内リンク。