Skip to content

Enable serial multi-GPU batch MTP - #1253

Closed
ilumn wants to merge 12 commits into
Niko1221:mainfrom
ilumn:feat/multi-gpu-batch-mtp
Closed

ilumn wants to merge 12 commits into
Niko1221:mainfrom
ilumn:feat/multi-gpu-batch-mtp

Conversation

@ilumn

@ilumn ilumn commented Oct 6, 2026 •

Copy link
Copy Markdown

This enables per-slot multi-token prediction (MTP) when the target model is split across CUDA GPUs. Previously, --layer-split 24 --batch 4 --batch-mtp disabled batch MTP and used ordinary batching. The enabled path verifies one proposal per slot through the stages serially, with --batch-groups 1, and commits the same accepted prefix on every stage.

A slot drafter needs the target model's final residual and output head, which reside on the last GPU in a layer split. Slot drafters now bind that GPU's weights and head, retain their residual buffers and draft-KV copies there, and use the first stage's PLE history. Each target stage supports eight-row verification windows with bounded graph caching. Allocations, idle waits and cleanup select the GPU that owns the corresponding buffers. Shared draft-vocabulary metadata supports the top-two diagnostic.

This depends on #1242, which isolates concurrent slots' image positions and corrects shared draft-head metadata. That dependency is included in this main-based PR. The serial multi-GPU feature commit is 0de9365f18a2ddcc7e428cc70a639ffcdfcc3b55 (7 files, 210 additions, 25 deletions), based on dependency commit 8fc9490e9cb62e4d3c3ab5f09c1ca8d42f4a8623; validation used main 82f46a8c8f475f001ad76d92f58f4a4f8ffb0253.

Validation on Talos

CUDA 13 / GNU 13.3 Release build, 2× RTX 5090 32 GiB, Swift IQ3_XXS, native Q5_K head, 106299-token draft subset, 8K allocated context, split 24, visible device order 1 then 0, 3 GiB reserve per GPU. Exactness controls: STRATA_IQ_MT_MIN=1 --pcie-frac 0 --adapt-every 1000000 --no-prefill-borrow. Requests generated up to 64 tokens; long prompt 1210 tokens. Independent review covered this serial implementation and its validation evidence.

Check Coverage Result
CUDA build, whitespace, Python syntax Full native build and new harness Pass
CPU regressions 27 checks Pass
Host native regressions 13 on GPU1, 3 on GPU0 Pass
Dual batch 2 / spec 2 Four speculative rows across both stages Exact solo parity
Five-slot rotation Eight-row windows, greedy and seed 1234 / temperature 0.7 / top_k 20 Exact solo parity
Q4 subset head, top-two --mtp-q4 head, STRATA_MTP_TOP2=1 Exact solo parity
Single-GPU regression Four-slot batch MTP Exact solo parity
Plain split/fallback Ordinary batching; groups 2 disable MTP with expected warning Exact solo parity
Staggered lifecycle Long admission, completed next turn, completed/cancelled slot → solo, reuse; greedy/seeded 7/7 each
Synthetic mixed images/text Different 4×6/6×4 grids, malformed admission, cancel/reuse; greedy/seeded 10/10
Real mixed images/text Existing CPU encoder, 11×5/5×11 grids, same lifecycle; greedy/seeded 10/10
Equivalent-mode performance 18 runs, 3 repetitions per mode/profile, identical per-stage cache counts and cross-mode tokens Pass

Speculative runs were checked for grouped-row captures and no disabled-option warning; lifecycle logs prove cache and turn-checkpoint restore paths ran.

Performance and limits

Four fixed prompts, 64 tokens each, same binary/settings, 3 repetitions per batch mode/profile. Rates below are median tokens/s. Normal MTP solo controls ran in each invocation; its rates separate decode from prompt reading. Batch overall includes admissions; the second batch rate starts at the first batch token and still includes overlapping admissions.

Experts per stage Normal MTP solo decode / with prefill Plain serial overall / from first batch token Plain pipeline overall / from first batch token Batch MTP overall / from first batch token
4000 70.9 / 55.8 55.6 / 59.5 61.0 / 66.2 54.7 / 58.7
11000 180.1 / 148.0 150.1 / 164.4 176.7 / 196.8 148.0 / 161.8

Batch MTP was slower than ordinary pipelined batching in both measured profiles. It adds opt-in compatibility for split models; the default is unchanged. These 64-token measurements do not establish long-context or sustained server-load performance.

Explicit fallbacks remain for --batch-groups > 1, same-device split stages, helper caches, remote expert optimization and peer devices. --pipeline-windows remains off with batch slots. Validated scope is two CUDA GPUs and the existing model/runtime above. HIP, more than two GPUs, more than five slots, worst-case context/memory pressure, other model formats and speculative pipeline groups are unrun. The existing idle split BYIELD path is outside this validation; the new harness checks supported staggered admission. Penalty sampling and coupled draft sampling were not validated.

With a compatible two-GPU model config, the included tools reproduce rotating batches and lifecycle checks:

CUDA_VISIBLE_DEVICES=1,0 python tools/batch_test.py --exe build/strata --config strata-model.json --batch 5 --n 5 --max-new 64 --extra "--batch-mtp --batch-groups 1 --layer-split 24 --trim-stage-weights --vision --pcie-frac 0 --adapt-every 1000000 --no-prefill-borrow"
CUDA_VISIBLE_DEVICES=1,0 python tools/batch_mtp_split_test.py --exe build/strata --config strata-model.json --keys "temperature=0.7 top_k=20 seed=1234" --extra "--batch-mtp --batch-groups 1 --layer-split 24 --trim-stage-weights --vision --pcie-frac 0 --adapt-every 1000000 --no-prefill-borrow --prompt-cache 4096"

Mixed-image/performance harnesses, exact local command records and logs are retained separately in the validation bundle.

Longer-run exactness investigation

A separate experimental implementation overlaps speculative slot groups across the GPU stages. Its 32-request workload at concurrency 4 and 512 generated tokens per request produced intermittent exact-token mismatches, including with overlap disabled while retaining the experiment's private verifier/drafter structure. Its benchmark matrix is paused. The cause and whether the published serial implementation is affected remain unresolved; the scoped results above do not establish long-run exactness for that experiment.

The existing QFUSE activation-quantization issue (#1139) was investigated as a possible explanation: batched GDN can consume stale quantized activations when STRATA_QFUSE=1. The QFUSE correction in the concurrent-generation PR (#1209) addresses that path, with runtime validation reported on gfx1151. The Talos CUDA diagnostic configurations leave QFUSE disabled, so this specific mechanism is inactive there. Other underlying causes remain under investigation.

ilumn added 2 commits October 6, 2026 15:36
Port the original slot-isolation fix onto the rewritten upstream main,
reconcile verifier fusion guards, and scope batch-MTP graph recordings.
Validate vision records and retain shared draft-head format metadata.

Original-author-commit: 10114f8
Superseded-PR: Niko1221#994
Talos-validation: CUDA build, focused regressions and 13 model checks passed
@ilumn
ilumn marked this pull request as ready for review October 6, 2026 23:38
@blange48

blange48 commented Oct 7, 2026

Copy link
Copy Markdown
Contributor

Tested on a 4-stage split. Rig: 4× RTX 5080 16 GB (PCIe Gen3), i9-10920X, Qwen3.8-Flash-Next IQ3_S, 262K, int8 KV, --layer-split 12,24,36 --trim-stage-weights, MTP spec 4, --batch 8 --batch-groups 4, conversation cache 8 GiB, on v0.1.40.1 plus our serve patches; HTTP server, T=0.7, 300 tokens. For this PR: --batch-groups 1 --batch-mtp against --batch-groups 1 alone, same image. The engine log shows the 8 slot drafters ("shared draft weights, 50 MiB of private state" ×8).

  • Exactness: batch_test with --batch-mtp: 8/8 and 3/3 identical to solo (--pcie-frac 0 --adapt-every 1000000).
  • Throughput (total tok/s, clients 1 / 2 / 3 / 4 / 8):
    • groups 1, no batch MTP: 123 / 125 / 148 / 174 / 202
    • groups 1 + --batch-mtp: 120 / 125 / 149 / 173 / 200
    • our usual groups 4 (pipelined), no batch MTP: 121 / 154 / 218 / 258 / 378
  • Greedy 1 / 4 clients: 152 / 179 vs 148 / 180.

So it works and stays exact, but on 4 stages batch MTP gives no gain over serial batching, and pipelining the stages is worth far more (~1.9× at 8 clients). This matches your own table (batch MTP 54.7 vs plain serial 55.6 vs pipeline 61.0). A pipelined variant (groups > 1) would be the interesting next step for splits.

@noon-at-cgn

Copy link
Copy Markdown
Contributor

@ilumn I have been running the same design on 2x RTX 3080 20 GB (UD-Q4_K_XL, two slots) and opened #1636 for it, based on current main. It conflicts with this PR in generate.cpp, verify.cpp and docs/BATCHING.md, so only one can go first. The PR text compares the two in detail: what is the same (drafters on the last stage's GPU, serial windows, eight-row windows), what is only here (your OnDevice guards, late_bind accounting, tools/batch_mtp_split_test.py) and what is only there (a tested row-layout function and, for 0.1.41, a fix so --batch-groups auto does not turn batch MTP off). I did not copy anything from this PR. If you would rather go first I will rebase #1636 onto yours; if not, I will keep your harness and guards in mind and credit them.

@Niko1221

Copy link
Copy Markdown
Owner

Thanks. Your own table and blange48's four-stage run both show that batch MTP over a serial split is no faster than the pipelined groups, which are the default on a split, and we already carry a batch-MTP path for two stages from #1636. So we are not taking it. A pipelined batch-MTP variant with a measured gain would be welcome.

@Niko1221

Copy link
Copy Markdown
Owner

Thanks for the pull request.

We're resetting the tracker for 0.2.0-beta. There are more than 1,300 open issues and pull requests, too many for us to review properly, so we're closing everything that isn't being worked on right now. Closed here does not mean rejected.

To reopen it takes about two minutes:

  1. Rebase on main if it no longer applies, and re-run your checks on 0.2.0-beta.
  2. Start the title with its area: [Engine], [GUI], [TUI], [Benchmark], [Docs] or [Other].
  3. Fill in the pull request template: the problem and result, and the validation (what you ran, on which revision, what you saw).
  4. Comment /reopen. A check runs and either reopens the pull request or lists what is still missing.

The rules are in CONTRIBUTING.md.

@Niko1221 Niko1221 closed this Oct 19, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants