metal: pre-M5 decode + prefill speedups — 46 t/s decode, +4% prefill on M3 Ultra (bit-exact) - #954
Open
ivanfioravanti wants to merge 16 commits into
Open
Conversation
The end-and-wait decode stage profiler serializes the token, so its stage times mix in per-boundary CPU waits and disable the concurrent shared-expert overlap. Add a diagnostic that keeps the production schedule: each stage boundary commits the open batch command buffer without waiting and records it, and after the token every stage's GPU busy span (GPUEndTime minus GPUStartTime) is printed. Only the serializing profiler now disqualifies concurrent dispatch, so DS4_METAL_STAGE_COUNTERS=1 measures the real schedule; per-token total-cb-busy matches production GPU-busy time and generated tokens are unchanged. Document the attribution it produced on M3 Ultra: dense Q8_0 matvecs already stream at the memory wall, while the per-layer HC pre/post kernels, the Q/KV norm-RoPE-store dispatch, and the routed-MoE window hold the remaining slack.
Non-strict DSpark peaks at 37.8 t/s on M3 Ultra (vs 43.7 plain): the batch verify costs 46-60 ms against a ~29-33 ms two-row floor (uniform 1.5-2.5x per-stage excess, decomposed with the commit-only stage counters), plain decode inside a DSpark session costs ~51 ms/token from hidden-state capture, and the draft proposes on 54-85% of cycles. Scheduler env tuning was swept and documented; making speculation profitable needs the N<=6 microbatch verifier, capture-light plain decode, and a cheaper propose chain.
Token-level profiling inside a DSpark session shows plain decode at a normal ~23.3 ms; the earlier ~51 ms figure misattributed post-accept first-token evals. The real per-cycle costs are the confidence-gated propose (3-8 ms, declining 45-75% of cycles) and ~4-5 ms of bookkeeping. Best swept configuration: 39.5 t/s at --dspark-confidence 0.75 with NO_DRAFT_SKIP=0, still below the 43.7 t/s plain equilibrium.
Per-row routed MoE via the single-token static kernels (bit-identical, verify_layer unchanged) and per-row HC pre via the decode fused producer (valid non-strict output, 39.3 vs 40.2 t/s) both confirm the batch-verify excess is inside the batch kernels' execution rather than dispatch routing. The microbatch verifier needs genuine small-N batched kernels.
The N=2 dual-row producer was bit-exact under the --dspark-strict oracle (after fixing a kernel-argument misbinding the oracle caught) but measured no per-verify gain and collapsed non-strict drafting on the test prompt. The strict-mode oracle is documented as the validation tool for future verify-kernel work; the HC-pre excess was overestimated by serialized stage decomposition.
Per-row MoE, per-row HC pre, and a strict-oracle-verified dual-row HC-pre kernel each recovered nothing measurable, and the batch MoE already runs at its distinct-expert floor. The perfect-sharing verify floor (~29-33 ms) looks undeliverable on this GPU: the N=2 verify near 50 ms is close to its real floor, so speculation is unlikely to beat plain decode on M3 Ultra.
Cold-restart kit: verified baselines with reproduction commands, the 22.6 ms token ledger, the closed-avenue list (8 kernel variants, knob sweeps, 3 microbatch increments), tool usage (commit-only stage profiler, strict-mode oracle, balanced A/B harness), hard-won gotchas, and the go/grind/accept/new-hardware decision that must be made first.
Round 2 (carried in working tree): raw-layer gathered attention parity (+2.7%) and CLI greedy chain decode (+1.75%), 43.46 -> 45.51 t/s, bit-exact. Round 3: ds4-eval decodes in session chain bursts (ds4_session_eval_chain_greedy) outside the think-close window, 44.5 -> 45.2 t/s, traces bit-identical; DS4_DISABLE_GREEDY_CHAIN=1 covers both paths. MoE down-sum6+HC4 tail fusion landed bit-exact, speed-neutral (rollback DS4_METAL_DISABLE_DECODE_MOE_HC_FUSION). KV-staging direct read landed bit-exact but -7% (per-head F32 re-read/re-convert amplification), gated off (opt-in DS4_METAL_ENABLE_DECODE_RAW_DIRECT_KV). make test 44/44; CLI transcript md5 db0c504c... unchanged.
ds4_gpu_router_select_tensor_devtoken still passed 20 arguments after ds4_gpu_encode_router_select gained visual_bias, single_token, vocab_size, and mixed_visual parameters, breaking the build. Pass nil/0/NULL/0/false for the non-visual dev-token path, matching the other three call sites.
Swap the pre-M5 device gate in metal_graph_encode_layer_attention_batch for the shared M5 port helper so resident single-device M5 graphs also skip the four dead indexer-query dispatches on zero-prefix ratio-4 batches while the compressed cache stays at or below top-k. Rollback on M5 is DS4_METAL_DISABLE_M5_BATCH_INDEXER_QUERY_PRUNE; pre-M5 keeps its existing per-feature and aggregate rollbacks. Balanced M5 Max A/B (IQ2_XXS/Q2_K ds4flash.gguf): 753.62/732.83 tok/s at 2048 tokens (+2.84%), 8/8 runs bit-identical (1,034,240 logits). The standard sweep's first frontier measured +8.7% prefill with decode and all ineligible frontiers unchanged. Also measured and rejected on M5 Max (see speed-bench/README.md): indexed prefill RB4 kernels (-0.27..-1.56% vs the MPP dual-heads default) and the Q2 2/32 decode split (-0.6..-0.8% vs the 4/none default).
The line is printed on every session creation, and it reports a build and weights decision the caller did not make and cannot act on. An embedder that creates a session per turn (or per aside, or per sub-agent) has it land in the middle of the user's screen each time. Gated behind DS4_DSPARK_VERBOSE, matching the other DS4_* diagnostic switches. The failure path beside it stays unconditional: that one is news. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Metal performance work for pre-M5 GPUs (M1–M4 resident path) that was benchmarked during the Aug decode campaign but never landed on
main. All changes are guarded by rollback env switches and produce bit-identical output.Main pieces:
metal_graph_greedy_chain): GPU argmax feeds the next token via a device ring, removes the per-tokenwaitUntilCompleted+ 517 KiB logits readback + CPU argmax boundary.Also wired into the session API (
ds4-eval).speed-bench/.Performance (branch vs main, fresh builds, interleaved runs)
Hardware: Apple M3 Ultra, 512 GiB · Model:
DeepSeek-V4-Flash-MXFP4…-chat-v2-mxfp4-0731.ggufGreedy decode
./ds4 -m <mxfp4.gguf> -p "Write a short story about a lighthouse keeper." -c 8192 -n 128 --temp 0+5.0–8.6% tokens/s. Both builds and both runs produced transcript md5
db0c504c8203618552e685bd2c701e4f— the campaign oracle hash, confirming bit-exactness end to end.Prefill
./ds4-bench -m <mxfp4.gguf> --prompt-file speed-bench/promessi_sposi.txt --ctx-start 2048 --ctx-max 8192 --step-incr 2048 --gen-tokens 128First-token latency at ctx 2048: 28.1 ms → 24.3 ms. Steady-state generation at longer contexts is equal or slightly better (39.33 → 39.47 t/s at 8K).
Correctness evidence
db0c504c…) matches on both builds (above).speed-bench/README.mdon this branch).make test44/44 at campaign close; prefix-to-decode checks matched all selected token IDs.Rollback switches
Every optimization is individually disableable, e.g.
DS4_DISABLE_GREEDY_CHAIN=1,DS4_METAL_DISABLE_DECODE_RAW_GATHERED_ATTN,DS4_METAL_DISABLE_DECODE_RAW_PACKED32,DS4_METAL_DISABLE_DECODE_MOE_HC_FUSION,DS4_METAL_DISABLE_PRE_M5_INDEXED_ATTN_PREFILL_RB4,DS4_METAL_DISABLE_PRE_M5_INDEXED_ATTN_PREFILL_HEADS16_DUAL_RB4,DS4_METAL_DISABLE_PRE_M5_BATCH_MOE_SUM6_HC_FUSION,DS4_METAL_DISABLE_PRE_M5_BATCH_ATTN_OUT_HC_FUSION.