metal: DSV4 M1 Ultra prefill/decode campaign - #2
Open
auroter wants to merge 6 commits into
Open
Conversation
|
Important Review available on request
Reviews should be triggered manually for repositories with fewer than 10 stars. Select Trigger review above or comment ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Prefill and decode optimization for dsv4-xxs-1 on Apple Metal, plus the cleanup pass that made the tree PR-ready. Prefill: 2048-token chunk 79 -> 371.6 tok/s (4.7x; 134% of the ds4 reference's 277). Tiled MoE prefill GEMM (qc_moe_mm_map0 + 64-slot dual-half iq2_xxs tile, q2_K SoA down twin), dense-causal MMA prefill FA over pre-decoded scratches, fused compressor front (cr=4/cr=128), wide-tile transposed-store q8_0 qgemm, prefill-width mhc, and a max_num_batched_tokens fix (2048 -> 2176) that stops the scheduler splitting a 2048-token prompt into two chunks. Decode: 15 -> 30.5 tok/s end-to-end on the 1k-in/2k-out anchor. Per-step marshalling memos, fused mhc pre/post, dsv4 router top-k, sum-folded q2_K down projection, split-K sparse-MLA, load-time q2_K SoA repack, fp16-direct kv insert, single-dispatch inverse RoPE. Gates: 8-tok, off1-2000, and 2500-token anchors bit-exact (shas plus spec-decode counters) across the cleanup; six kernel oracles pass; prefill walls at the standing baseline. Method and evidence in perf/optimization_status.md and perf/baseline_status.md. Known open issues, documented in-code and in perf/: the async-output completion-event race at multi-chunk request completion (boot ramp protocol in perf/prefill_handoff.md; fix planned as a separate PR) and the step-tape diagnostic pending route re-validation.
auroter
force-pushed
the
metal-m1ultra-campaign
branch
from
August 15, 2026 04:50
f6a152b to
9c2fde1
Compare
added 2 commits
August 14, 2026 22:56
…ltra campaign Semantic resolutions: - qc_mmvq: campaign mb fast path keeps precedence for its gated envelope (batch 2-8, q8_0/q2_K/iq2_xxs/q6_K); Muse-Glimmer's qgemv_mm row-block walk becomes the general path, with the q8_0-small summation-order carve-out extended to it. - rms_norm: single binding dispatches bf16 contiguous inputs to the Muse-Glimmer fixed-D kernels and everything else (DSV4 fp16 strided q/k splits) to the w32/strided variants; duplicate binding and Python wrapper removed; Muse encodes labeled for the CB census. - compressor: NaN-hunt write-site debug hooks composed with the comp_full_compress phase bracket. - tk_launch.h: both launch_qgemv_mb and launch_qgemv_mm retained. - dsv4-xxs-1 metal profile: main's 262144/16 GiB sizing retained with the campaign's fp16 dtype and 2176 batched-token reserve; long-context note rewritten (Metal indexer chain now implemented, boot-ramp caveat). - benchmark harness: --dump-completions and chars_per_token both kept. - notebooks: both campaign histories merged chronologically. - metallib rebuilt from merged kernel sources; extension rebuilt; all six Metal kernel oracles pass.
… M1 Ultra campaign Two conflicts, both trivial: build_attn_metadata grew both sides' new kwargs (our num_computed_tokens_cpu, main's steady_cache), and the notebook keeps both appended sections in chronological order. The steady uniform-decode path is inert on Metal (requires FULL CUDA-graph mode and the VLLM_STEADY_DECODE_META opt-in). Profile tests pass (48/48).
added 3 commits
August 17, 2026 12:03
…fill host guards Verified all 8 findings from the QuixiCore-Metal PR QuixiAI#3 review against this tree; applied what holds, documented what was already guarded, rejected one. - mla fp8 insert (all three kernels): build 2^-e from the float bit pattern instead of fast-math exp2, which measures 2 ulps low at negative integer inputs on M1 Ultra (probe in perf/results/2026-08-17/coderabbit_fixes/). Matches the indexer kernels and the exact fp32 reference; cached codes change only for blocks with amax > 448. Scale byte derives from the same clamped ei (bit-identical in the reachable exponent range). Anchors must be re-gated on next boot (see notebook). - multi-row MoE GEMV hosts: N-divisibility now enforced — a8_vec routes non-multiple N to the one-row kernel, swiglu/sum TORCH_CHECK. The mr grid ceil-divides N and tail simdgroups read weight rows past N before the store guards; every DSV4 dim divides, so serving behavior is unchanged. - prefill FA: nc/ns % 32 TORCH_CHECK enforcing the _pad_slots contract; contract documented on the kernel and both launchers. - tk_launch.h: host-checked contracts documented on the router (<=1024 experts, topk <= 8), indexer compress (cr == 4 vs HISTORY_MAX), compress front (cr in {4,128}) and prefill launchers. - moe_mm_id: AoS qs base is 4-byte aligned (84*blk+16), not 16 — comment corrected before it justified a bad widening. - rms_norm: 256-thread dispatch contract documented (shm slots are only written by simdgroups that run). - save_partial_states: bf16 score+ape add documented as required for bit parity with the Triton reference (review suggested widening; rejected). Rebuilt metallib; full kernel suites pass (prefill FA oracle, 36 tiled-GEMM checks, SoA/sum6 bit-identical, compress-front c128 bitwise, indexer topk).
Fast-math metal::exp2 is 2 ulps low at negative integer inputs, so decode was reconstructing every typical KV scale slightly below what the insert kernels used. Reconstruct 2^(e-127) exactly from the float bit pattern via a shared mla_ue8m0_scale helper. Metallib rebuilt; all six Metal kernel suites pass. Anchors re-gated with the insert-scale fix in one boot.
Fresh ramped boot on the fixed build: 8-tok 573db39598e7, off1-2000 bb83cc3054a3 (57.3 s, step ms unchanged), 2500x64 f75e1d41ac3d. All deterministic across repeat runs; wall shifts are acceptance-mix effects of the re-rolled trajectories. baseline_status UPDATE 30.
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
1 participant
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.
Prefill and decode optimization for DeepSeek V4 Flash 0731 (
dsv4-xxs-1profile) on Apple Metal, M1 Ultra.The point of this campaign: match and beat antirez's ds4 — the optimized reference engine for this model — on the same box, with the full SlimServe/vLLM serving stack (OpenAI API, DSpark speculative decoding, TurboQuant draft KV, prefix caching). Both engines measured on the same M1 Ultra with the exact-token harness.
SlimServe vs ds4 (same box, same model, same workloads)
Campaign start was 79 tok/s prefill and 15 tok/s decode, so this is 4.7× / 2.1× over where the stack began.
Where the speed comes from
Decode (15 → 31.6 tok/s)
dsv4_router_topk, sum-folded q2_K down projection (GEMV + weighted-sum in one kernel), split-K sparse-MLA decode, fp16-direct KV insert, single-dispatch inverse RoPE.Prefill (79 → 369 tok/s)
qc_moe_mm_map0work queue feeds a 64-slot dual-half iq2_xxs simdgroup-MMA tile and a q2_K SoA down twin — MoE layers stop being token-at-a-time GEMVs at prefill widths. Biggest single prefill win.max_num_batched_tokens2048 → 2176 so the scheduler's draft-token reserve stops splitting a 2048-token prompt into two chunks.Also in this PR
main(A100 Q4_K fused MoE, NaN split-K fix, Muse-Glimmer Metal): the campaign's fast paths and Muse-Glimmer's coexist behind format/dtype dispatch, and the profile ships main's 262144-token sizing — which this branch's Metal indexer port actually makes reachable, and which measures slightly faster at the decode anchor.tests/kernels/(runnable directly on the box, no pytest needed).perf/(method, per-wave evidence, raw-artifact locations).Validation
Zero-numeric-change discipline for the cleanup and the main merge: pinned to the benchmark config, the 8-tok, off1-2000, and 2500-token serving anchors are bit-exact (output sha256 + spec-decode counters) against the pre-cleanup build, and all six kernel oracles pass. The table above is measured on the merged tree under the shipping profile.
Known issues (documented in-code; fixes coming as separate PRs)
perf/prefill_handoff.mdavoids it deterministically, andvllm/platforms/metal_compat.pydocuments the details.VLLM_QC_STEP_TAPE, default off) needs route re-validation before its tape-only mode is trusted; seevllm/models/deepseek_v4/metal_tape.py.