Add request-owned eight-way MTP serving for Qwen 35B

[davidtai]

Summary

This PR adds a real eight-request MTP path for Qwen 35B A3B.

Two to eight ready requests run in one fixed B8, depth-one MTP cohort. One request still uses the existing solo MTP path. This path does not switch to AR.

The batch is lockstep at the model-operation level. Each cycle executes one draft and one target verification for the active rows. The requests do not share a logical context.

What was broken

The visible symptom was that eight concurrent prompts looked like one growing context. Requests advanced together, memory kept growing, and long runs could end in resource-limit or invalid counter behavior.

There were several separate causes behind that symptom.

1. There was no complete request-owned MTP cohort path

The old concurrent route did not carry the whole solo MTP contract into a batch. A correct request needs its own target cache, committed MTP history, RNG stream, recurrent state, token budget, stop state, cancellation state, stream, and terminal result.

Grouping requests without all of those owned parts is not true MTP concurrency.

This PR adds a dedicated mtp_batch scheduler and service. A single sealed request uses unchanged solo MTP. Two to eight compatible requests use the fixed B8 lane. Later requests wait for the next cohort.

2. Physical cache growth was being confused with committed context growth

A depth-one MTP verification temporarily writes two target positions. A row may commit two tokens, one token, or no tokens if it is inactive. The physical ragged cache had already observed the temporary write.

If its host capacity bound was not restored after commit or rollback, it grew at close to two positions per cycle even when logical progress was only one. Finished and padded rows could also keep advancing while another row was active.

That is how eight individually valid prompts could still create oversized shared allocations. The fix tracks committed target and MTP offsets per row and resets the host capacity bound after every commit or restore. Inactive rows use keep=0 and preserve their pre-verification recurrent state.

3. MTP history was not just another target KV cache

Qwen's depth-one draft uses committed shifted history. For a prompt of length N, the target cache owns all N prompt tokens and the MTP history owns the shifted N-1 transitions.

The batch lane now prefills that history separately for each request, merges it by row, and advances it with the same accepted/rejected semantics as solo generate_mtpk. A rejected row restores the correct pending-primary position instead of carrying speculative state into the next cycle.

4. Finished and cancelled rows were not fully inert

It is not enough to stop emitting tokens. A row must stop changing target offsets, MTP offsets, recurrent state, and capacity bounds.

The final commit path handles keep=0, keep=1, and keep=2. It snapshots the pre-verification recurrent leaves so keep=0 restores the real base state. A mixed-row construction check advances one more cycle to catch incorrect recurrent-state selection.

If a long request is cancelled during prefill, its large cache is replaced with a one-token dummy row before merge. A final cancellation poll runs at the prefill/merge boundary so a last-row race cannot retain a cancelled 131K-token allocation.

5. Cache merge could keep both source and destination allocations alive

Eight scalar prefills are merged into one row-owned B8 cache. Keeping all scalar source caches alive after the B8 arrays materialize can nearly double KV residency.

The merge now materializes the destination and releases source layers as ownership moves. The all-empty MTP-history case, such as eight one-token prompts, also installs the correct empty KV geometry before the first real attention update.

6. The optimized B8 kernel route was not the same as a B1 kernel called eight times

The target verification shape is [B=8, T=2], which flattens to M=16 for projection and MoE work. The original post-convolution callable was built for [B=1, T=2].

This PR installs separate B8/T2 GDN post-convolution callables. The grid expands from 32 head rows to 8 * 32 = 256, and output/state shapes carry all eight rows. Startup compares optimized B8, eager B8, stock B8, and B1 references before installing the lane.

The mathematical work is row-separable. B1 and B8 can still have small BF16 reduction-order differences because their tiling is different. Those differences are bounded at startup and are not context sharing.

7. Cancellation and cleanup could race surviving rows

An early-completed request can return before its cohort peers finish. Request-thread MLX cache clearing must not run while the model-owner thread is still decoding those peers.

The service now stages successful results until cohort cleanup succeeds. Cancellation can close its own future early, but MLX finalization runs once on the model owner after the cohort. Cleanup failure fails the cohort, drains queued jobs, and poisons the lane instead of allowing more model work.

8. Sampling made the first correct B8 implementation slower

The first correct B8 lane reduced GPU work, but host code still built and sorted full-vocabulary distributions for every row and phase. Greedy requests paid that cost too.

The final route uses direct argmax for greedy sampling. Default stochastic sampling uses one construction-bound batched top-k route, then exact NumPy float64 top-p arithmetic. Exact BF16 ties use one rule everywhere: higher score first, then lower token ID. Target and draft requests with unsupported large device top-k values fail before prompt work.

9. The mlx-lm cache leak is a separate upstream issue

The related ArraysCache.advance() lazy Metal buffer-object leak already has an upstream fix in mlx-lm PR #1642, commit 985af30.

I did not open a duplicate upstream PR. The local Qwen service uses that commit for now. Its launcher checks for _lp_advance and _len_advance and refuses to start if the environment is replaced with stock mlx-lm 0.31.3.

Ownership model after this PR

The cohort shares physical B8 arrays and one model invocation. It does not share one context.

Each real row owns:

• prompt tokens and response ID
• target KV offsets
• committed MTP-history offsets
• GDN recurrent state
• target and draft sampler configuration
• NumPy RNG state
• completion token counts and stop state
• cancellation event, stream callbacks, and terminal future

Padding rows and terminal rows are inert. No new request can join a cohort after it is sealed.

Benchmark protocol

All performance numbers below are MTP versus MTP. AR is not the baseline.

• Model: Qwen3.6-35B-A3B MTPLX optimized-speed checkpoint
• Generation: depth-one MTP, target_prefix verification
• Server: --generation-mode mtp --scheduler-mode mtp_batch
• Capacity: --max-active-requests 8 --decode-batch-max 8
• Context window: 131,072
• Dependency: mlx-lm PR #1642 at 985af30
• Eight unique marker prompts per arm
• Seeds: 4200 through 4207
• Maximum output: 256 tokens per request
• Serial arm: submit the eight requests one after another, so every request uses B1 solo MTP
• B8 arm: release eight HTTP requests together through a thread barrier
• Metric: total completed output tokens divided by wall time
• Three paired rounds for each sampler setting

The two geometries can take different generation paths because of bounded BF16 accumulation differences. The table therefore reports the actual completed token count for each arm instead of assuming identical output length.

Performance investigation

Before the sampler repair, a paired run measured:

Lane	Aggregate output TPS	Relative
Serialized solo MTP	147.903	1.000x
Fixed B8 MTP	140.479	0.9498x

That result failed the 1.20x promotion gate.

The dispatch census explained why:

Lane	Wall time	GPU work/busy	Capacity wait
B1 serial control	10.651 s	8.575 s	—
B8 MTP	14.141 s	4.313 / 4.301 s	2.434 s

B8 had already cut GPU work by about half, but total wall time was worse. A process sample showed NumPy full-vocabulary argsort, cumulative probability work, and host synchronization dominating the remaining time. That led to the direct-argmax and fixed batched top-k changes.

Original default-sampler runs

Default sampling is temperature=0.6, top_p=0.95, top_k=20.

Round	B1 tokens	B1 wall	B1 TPS	B8 tokens	B8 wall	B8 TPS
1	1,077	9.848 s	109.366	1,061	6.593 s	160.932
2	1,083	8.617 s	125.687	1,061	6.479 s	163.753
3	1,221	10.146 s	120.337	1,061	6.570 s	161.500
Median	—	—	120.337	—	—	161.500

Ratio of medians: 1.342x.

Original greedy runs before the device-ID follow-up

Greedy sampling is temperature=0. It uses direct argmax while keeping the same MTP target/draft kernels.

Round	B1 tokens	B1 wall	B1 TPS	B8 tokens	B8 wall	B8 TPS
1	1,299	9.470 s	137.172	1,154	3.620 s	318.790
2	1,299	9.696 s	133.970	1,154	3.594 s	321.070
3	1,299	9.220 s	140.883	1,154	3.594 s	321.124
Median	—	—	137.172	—	—	321.070

Ratio of medians: 2.341x.

Default sampling is slower because it must compute full-distribution mass, deterministic top-k support, top-p truncation, RNG draws, MTP acceptance ratios, and residual correction. Greedy sampling mostly reduces that work to argmax.

Follow-up: restored B1 control and device-resident greedy sampling

The earlier 159.7 TPS serial figure is not the optimized PR #174 B1 decode
rate. It submits eight separate HTTP requests and includes eight prefills and
HTTP overhead. The exact PR #174 K1 harness was rerun unchanged before making
another B8 comparison:

Optimized B1 K1 control	Repeat 1	Repeat 2
long-code natural stop	197.321 TPS	198.883 TPS
short coding prompt (143 tokens with the current tokenizer)	207.207 TPS	207.044 TPS

The current branch proved compiled target-prefix, device draft input,
whole-MoE, GDN post-convolution, packed projections, row-owned routing, and
combine-tail were installed for that B1 run. These values are single-request
decode TPS. They are not the same metric as served aggregate B8 output TPS.

The B8 audit found a real missing optimization. Greedy B8 was copying full
[8,V] draft logits and [8,2,V] verify logits to NumPy every cycle. Commit
a798551e adds a cohort-selected greedy route that keeps the eight draft IDs on
the device, feeds those IDs directly into the B8/T2 verify graph, and transfers
only small token-ID arrays. Default stochastic sampling keeps the existing
batched sparse route. Penalty-bearing requests keep the dense compatibility
route.

Served B8 workload	Before	After	Change	Output check
coding, 8 x 192 tokens	414.852 TPS	452.413 TPS	+9.05%	all 8 hashes exact
legacy greedy	324.019 TPS	349.064 TPS	+7.73%	all 8 hashes exact
legacy default sampler	170.406 TPS	166.743 TPS	-2.15%	all 8 hashes exact

The default sampler does not select the new greedy route. Its small timing
change is recorded as run-to-run drift, not a performance claim. The coding
receipt SHA-256 is
b0683bab11fabab4f544f3d34292d92ceb559cf034b5adf5c2cc72fc7a40ef49.
The legacy greedy receipt SHA-256 is
8f64743bc745a61f198dc370103d175e6c4b713c6d0573d863323aa69f6a8e47.

Two nearby candidates were measured and rejected:

• forcing unsorted MoE gather reduced legacy greedy B8 from 324.019 to 301.040
  TPS (-7.09%) and did not change the eight B8 output hashes;
• a construction-only M16 whole-MoE probe improved one real target block from
  0.5308 to 0.5112 ms (1.038x), but changed BF16 block output
  (max_abs=0.0491). That was too little gain for a new 40-layer arithmetic
  route.

The full local suite passes with the same two unchanged cached vLLM-Metal ABI
tests deselected. All changed Python files pass Ruff and git diff --check.

Correctness and resource benchmarks

Eight-way marker isolation

• Eight HTTP 200 responses
• Eight distinct response IDs
• Every response contained its own marker
• No response contained another row's marker
• Health reported real width 8 on qwen35b_a3b_mtp_batch_b8_t2_m16

Active cancellation

• Maximum observed active requests: 8
• Cancelled rows: 1 and 6
• Six surviving requests completed
• All six survivors contained only their own marker
• Eight distinct response IDs were observed
• Final pending/active count: 0/0
• Final scheduler error: null

Long-context ownership

• Eight concurrent prompts
• 13,239 prompt tokens per request
• 32 completion tokens per request
• Total elapsed time: 27.583 seconds
• Eight HTTP 200 responses and eight distinct IDs
• No foreign markers
• Final pending/active count: 0/0
• Peak MLX memory: 32,557,292,816 bytes
• No negative scheduler values, context overflow, or Metal resource failure

Final live health

• generation mode: mtp
• scheduler mode: mtp_batch
• real-width histogram: B1 = 48, B8 = 8
• physical fixed-width cycles: B8 = 735
• last real width: 8
• last route: qwen35b_a3b_mtp_batch_b8_t2_m16
• last error: null
• MLX finalization scope: cohort_owner_after_decode

Numerical parity statement

Full B1 and B8 output hashes are not required to match. Their different BF16 tiling can make a near-tied logit cross an argmax or sampling boundary, after which generation follows a different path.

This is separate from the sampler tie bug. Given identical BF16 logits, serial and B8 sampling now have identical support, selected token, and next RNG state. The reproduced cutoff case had tokens 50 and 122 tied at 0.9921875; both routes keep token 50, exclude token 122, sample token 104 with seed 5, and preserve the next RNG value.

Test and CI results

• Focused sampler, MTP driver, serving, cancellation, OpenAI, exact-A3B, and sustained-generation tests pass.
• The full local pytest suite passes with four skips after deselecting two unchanged cached vllm-metal ABI tests.
• Changed Python files pass Ruff.
• git diff --check passes.
• GitHub wheel, no-mlx-smoke, and repository-hygiene checks pass.

Default deployment

The local persistent Qwen launcher now uses --generation-mode mtp, --scheduler-mode mtp_batch, and capacity 8. It does not use AR. DeepSeek remains disabled.

EvalPlus quality gate

Date: 2026-08-09
MTPLX commit: f5ece9068580cb4ea1c6a1db5250aad505c9750d
EvalPlus: 0.3.1

Protocol

• Model: Youssofal--Qwen3.6-35B-A3B-MTPLX-Optimized-Speed
• Server: MTP depth 1, mtp_batch, fixed physical B8/T2 kernel
• HumanEval+ 164 tasks, hash fe585eb4df8c88d844eeb463ea4d0302
• MBPP+ 378 tasks, hash ee43ecabebf20deef4bb776a405ac5b1
• One completion per task
• EvalPlus 0.3.1 OpenAI system and user prompts
• Greedy target decoding: temperature 0, top-p 0.95, maximum 768 tokens
• B1 control: one HTTP request at a time, which routes to solo MTP
• B8 candidate: concurrent requests on the fixed physical B8 kernel

The main candidate event audit matched 540 scored responses to mtp_batch
events at real widths 2 through 8. Two HumanEval requests initially arrived
alone. They were regenerated behind an active-8/pending-14 backlog, both
reported real width 8, and both reproduced the old output byte for byte. The
final scored route distribution is B8=488, B7=14, B6=30, and B2=10. All 542
candidate samples therefore came from the physical B8 kernel.

Results

Suite	Metric	B1 solo	B8 kernel	Delta
HumanEval	base	151/164 (92.07%)	151/164 (92.07%)	0.00 points
HumanEval+	base + extra	145/164 (88.41%)	144/164 (87.80%)	-0.61 points
MBPP	base	338/378 (89.42%)	335/378 (88.62%)	-0.79 points
MBPP+	base + extra	289/378 (76.46%)	285/378 (75.40%)	-1.06 points
Combined plus	base + extra	434/542 (80.07%)	429/542 (79.15%)	-0.92 points

Paired plus-test swaps:

• HumanEval+: B1-only 3, B8-only 2, exact McNemar p=1.000
• MBPP+: B1-only 9, B8-only 5, exact McNemar p=0.424
• Combined: B1-only 12, B8-only 7, exact McNemar p=0.359

The point estimate is a small loss, not zero loss. The paired results do not
show a statistically significant systematic regression. B1 and B8 sanitized
programs were byte-identical on 103/164 HumanEval tasks and 200/378 MBPP tasks;
the remaining differences are expected from the different BF16 reduction
geometry.

Commands

/Users/davidtai/projects/evalplus/.venv/bin/python -u \
  evalplus_paired_codegen.py --arm b1 --root ./full \
  --datasets humaneval mbpp

/Users/davidtai/projects/evalplus/.venv/bin/python -u \
  evalplus_paired_codegen.py --arm b8 --root ./full \
  --datasets humaneval mbpp

/Users/davidtai/projects/evalplus/.venv/bin/python -m evalplus.evaluate \
  humaneval --samples <samples.jsonl> --parallel 4

/Users/davidtai/projects/evalplus/.venv/bin/python -m evalplus.evaluate \
  mbpp --samples <samples.jsonl> --parallel 4

The generation run shared the service with unrelated long package-scanning
requests. That was useful row-isolation stress, but it invalidates generation
wall time as a throughput measurement. The clean TPS measurements in PR #245
remain the performance result; this receipt is quality-only.

Mode selection and live command validation

Date: 2026-08-09

I loaded the real Youssofal/Qwen3.6-35B-A3B-MTPLX-Optimized-Speed model once per mode. Each test sent eight requests at the same time. All eight requests completed and returned unique response IDs.

Mode	Observed route	Execution receipt
throughput	qwen35b_a3b_mtp_batch_b8_t2_m16_throughput	real width 8; batch histogram {"8":1}
balanced	qwen35b_a3b_mtp_batch_b8_t2_l0_b1_qkv_z_b_balanced	real width 8; batch histogram {"8":1}
b1-exact	qwen35b_mtp_batch_b1_exact_serial	eight unchanged solo runs; no fixed-width B8 claim

The running production service was then restored to throughput. Its health payload reports MTP depth 1, context 131072, scheduler mtp_batch, real width 8, the throughput route above, and no scheduler error. DeepSeek was not loaded during these tests.

The full mtplx serve command in docs/concurrency/qwen35b-mtp-batch.md was used for the model loads. The saved-config form was also parsed through the real CLI and resolved scheduler_mode=mtp_batch, batching_preset=throughput, width 8, context 131072, and the selected numerics value. mtplx serve does not support --dry-run --json; those options are not shown in the guide.

The generic scheduler contract remains in docs/concurrency.md. Qwen B8 geometry, limits, commands, numerics, health checks, and the live receipt are in the linked backend guide. The guide also states the current mlx-lm PR #1642 prerequisite so a stock dependency resync is not presented as production-safe.

Clean-install public launch verification

Date: 2026-08-09
Commits: 6b04a731, c15903e4

The first clean-package test caught a real documentation and startup gap. The old guide selected the sustained profile and relied on three private launcher exports. The real 35B model therefore failed its construction-time numerical check instead of serving. A normal install also selected released mlx-lm 0.31.3, which lacks the ArraysCache fix from upstream PR #1642.

This follow-up makes that contract public and fail-closed:

• the guide now requires --profile turbo --verify-strategy target_prefix;
• the server installs the measured packed gate/up, GDN TGY4/headquarter, row-owned router, and combine-tail choices at construction, with no hot-path eligibility fallback;
• an incomplete MTP-batch launch contract fails before model weights load;
• stock mlx-lm fails before model load with exact uv pip and pip repair commands for commit 985af30;
• the guide's command is parsed by the real public CLI in tests.

I installed the branch as a wheel in a fresh virtual environment, installed upstream mlx-lm commit 985af30df768a6f4dd2d0c7969d1868ca5dc3e1a, changed to /tmp so the source checkout could not shadow the wheel, and ran the documented command without any MTPLX_* shell exports. All three real-model launches reached MTPLX is ready.

Mode	Eight HTTP results	Unique IDs	Runtime receipt
throughput	8 x 200	8	qwen35b_a3b_mtp_batch_b8_t2_m16_throughput, width 8, {"8":1}
balanced	8 x 200	8	qwen35b_a3b_mtp_batch_b8_t2_l0_b1_qkv_z_b_balanced, width 8, {"8":1}
b1-exact	8 x 200	8	qwen35b_mtp_batch_b1_exact_serial, solo_runs=8

Production was restored to Qwen throughput mode after the test. A fresh eight-request production cohort returned eight HTTP 200 responses and eight unique IDs; health reported width 8, the throughput route, and last_error=null. DeepSeek was not loaded. The GPU lock was released only after that receipt.

The full-suite failures found during this check were also fixed. The native vLLM-Metal extension cache had been keyed only by Python ABI, so another environment's newer but MLX-incompatible binary was reused. The cache filename now includes a SHA-256 fingerprint of the actual libmlx.dylib. The graceful split-SDPA fallback also preserves the native paged kernel's implicit causal mask for multi-token queries.

Verification for the final follow-up:

• the complete local pytest suite passed to 100% with four skips and no deselections;
• both real Metal paged-attention numerical tests passed, including the native partitioned route;
• the MLX-ABI cache-key regression test passed;
• changed Python files passed Ruff;
• git diff --check passed;
• the wheel build completed.