lablup / lablup/mlxcel

perf(speculative): re-measure the batch-capable B=1 MTP hardware gate on M1 Ultra and M3 Ultra and adjust mtp_b1_default

Open
#1,217 1 comment 0 reactions 0 assignees View on GitHub
area:benchmark area:inference platform:macos priority:medium status:ready type:performance
Dominant language
Rust
Stars
467
Forks
54
Avg merge
4h 25m
Merged PRs (30d)
310

Description

## Summary

The static per-hardware gate for the B=1 MTP burst, `mtp_b1_default` in `src/server/batch/speculative_burst.rs`, decides as follows: an `MLXCEL_ENABLE_MTP_B1` override wins in both directions; otherwise non-batchable targets always run the burst, and batch-capable targets run it only when `has_neural_accelerator` is true (an M5+ chip-generation proxy from `mlxcel_core::hardware`). The pre-M5 decline for batch-capable targets rests on measurements taken before several substantial MTP changes landed, and M3 Ultra has never been measured on the batch-capable pairing at all: the binary NA proxy lumps it with M1 Ultra even though every recent measurement places it much closer to M5 Max. This issue asks to re-measure the batch-capable B=1 pairing on M1 Ultra and M3 Ultra on current main and to adjust the gate according to the results.

## Current behavior

- `mtp_b1_default` (`src/server/batch/speculative_burst.rs`, pure decision core of `mtp_b1_burst_enabled`): batch-capable target + no neural accelerator = decline to classic decode. Its docstring records the founding measurements: batch-capable 31B + bf16 assistant at ~1.2 to 1.4x on M5 Max, but a consistent 0.75 to 0.96x regression on M1 Ultra (four greedy 160-token prompts), which is why pre-M5 defaults to classic. That policy dates to #165.
- The gate is consulted in `Scheduler::mtp_b1_should_run` (`src/server/batch/scheduler.rs`) when no adaptive policy is attached, i.e. whenever `MLXCEL_MTP_ADAPTIVE=0`.
- It also matters with the adaptive policy (#333) enabled (the default): `MtpPolicy::static_default` (`src/server/batch/mtp_policy.rs`) calls the same `mtp_b1_default` to resolve an ambiguous profiling window. So a stale static default still steers real decisions even in the adaptive path.

## Why the founding evidence is stale

1. The M1 Ultra 0.75 to 0.96x rows behind the gate predate the recent MTP work: #1194 (drafter step attribution and accept-hook counting), #1199 (exactness probe can drop `qmv_wide`), #1203 (drafter projections quantized at load), #1208 (report the effective block width the round loop actually used), #1215 (benchmark methodology: GPU contention guard, round-cost protocol, `scripts/with_indexers_paused.sh`). In particular #1203 is documented in `docs/benchmarks.md` as the change that moved the Qwen pairing on M1 Ultra from 0.59 to 0.70x up to break-even (`docs/benchmark_results/qwen38-mtp-m1ultra-2026-08-16.md`), so its effect on the Gemma 4 31B + bf16 assistant pairing on pre-M5 hosts is plausibly material and has not been measured.
2. M3 Ultra was never measured on the batch-capable pairing. The docstring only cites M5 Max and M1 Ultra. Yet the 2026-08-19 round-cost model in `docs/benchmarks.md` (measured on the non-batchable Gemma 4 12B pairing) puts M3 Ultra at a break-even of ~1.51 emitted tokens per verify at block 4, far closer to M5 Max's ~1.28 than to M1 Ultra's ~2.71, and M3 Ultra measures 1.74 to 2.61x on the B=1 12B pairing.
3. The discriminator identified by the 2026-08-19 measurements is Apple GPU compute generation via the `use_qmv_wide` split (documented in `src/models/speculative_exactness.rs`): generation 15+ (M3, M4, M5) runs a quantized projection at M >= 2 as one wide pass, generation 13 (M1, M2) runs the block as narrow per-position passes. `has_neural_accelerator` (M5+ only) is a coarser proxy than the mechanism, and it misclassifies M3 Ultra to the slow side.

## What to measure

On both M1 Ultra and M3 Ultra, on current main, following the #1215 protocol (`scripts/bench_speculative.sh` under `scripts/with_indexers_paused.sh`, spread and contention guards respected, adaptive profiling window discarded, host and prompt recorded on every row):

1. **Gemma 4 31B + bf16 assistant, B=1** (the batch-capable pairing the gate governs). Force the path with `MLXCEL_ENABLE_MTP_B1=1` and `MLXCEL_MTP_ADAPTIVE=0` so the static gate and the policy do not interfere with the measurement. Run the three standard prompts (enumeration, source code, prose), plus a block-width sweep (`scripts/bench_block_width.sh`) at least on the code row.
2. **Qwen 3.5/3.8 MTP pairing on M3 Ultra** for completeness (M1 Ultra was measured 2026-08-19: ~0.98x untuned and a loss at every width, already recorded in `docs/benchmarks.md`; the target is batch-capable so the same gate governs it).
3. For every row record acceptance, emitted per verify, round cost in classic steps, and spread, so the rows extend the existing round-cost table.

Checkpoints needed: `gemma-4-31b-it-4bit`, `gemma-4-31B-it-assistant-bf16` (already wired into `speculative_bench` `REACHABLE_PAIRINGS`), and the Qwen pairing checkpoints on the M3 Ultra host.

## Adjusting the gate (decision matrix)

- If M3 Ultra gains on the batch-capable pairing(s): replace the `has_neural_accelerator` predicate in `mtp_b1_default` with the finer discriminator the measurements support, e.g. Apple GPU generation >= 15 (the `use_qmv_wide` capability the exactness probe already reads), keeping generation 13 declined.
- If M1 Ultra now reaches break-even or better post-#1203 on the 31B pairing: consider enabling batch-capable targets everywhere and letting the adaptive policy (#333) decline unprofitable pairings per host.
- If both hosts still regress: keep the gate as-is, but refresh the docstring numbers and `docs/benchmarks.md` with the new dated rows so the next reader knows the evidence is current.
- In all cases: update the `mtp_b1_default` unit tests (it is the pure test seam for this decision), update the docstring measurements, and keep the note that the static default also resolves adaptive-policy ambiguity, so its correctness matters even with the adaptive path on.

## Acceptance criteria

- [ ] M1 Ultra and M3 Ultra measurement rows for the Gemma 4 31B + bf16 assistant B=1 pairing recorded under `docs/benchmark_results/` and reflected in `docs/benchmarks.md`, following the #1215 protocol (guards, host and prompt recorded, profiling window discarded).
- [ ] M3 Ultra rows for the Qwen MTP pairing recorded the same way.
- [ ] `mtp_b1_default` adjusted per the decision matrix (or explicitly confirmed unchanged with the dated evidence), with its docstring updated to the new measurements.
- [ ] Unit tests for `mtp_b1_default` updated to the new predicate and passing.
- [ ] The adjusted gate is exercised through the real dispatch path (`Scheduler::mtp_b1_should_run`) on at least one real checkpoint per affected host, not only through unit tests.

## Notes

- Hardware-gated: needs physical M1 Ultra and M3 Ultra hosts. M5 Max rows can be reused from the existing tables.
- `MLXCEL_ENABLE_MTP_B1` remains the both-directions override and is unaffected by this issue.

Contributor guide

Open the contributing guide

Research direction

Start with src/server/batch/speculative_burst.rs and its mtp_b1_default unit tests, then review the dispatch path in src/server/batch/scheduler.rs and MtpPolicy::static_default in src/server/batch/mtp_policy.rs. Run scripts/bench_speculative.sh and scripts/bench_block_width.sh under scripts/with_indexers_paused.sh on the specified M1 Ultra and M3 Ultra hosts, recording the required rows. Done means the dated benchmark results, benchmarks documentation, gate docstring, tests, and real dispatch checks reflect the measurement decision.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust, shell
Domain
performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.