metal: fix the async-output completion-event wedge - #3
Open
auroter wants to merge 4 commits into
Open
Conversation
added 4 commits
August 14, 2026 21:50
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.
…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).
The engine could park forever in MPSEvent::synchronize on the async-output copy event -- GPU idle, signal never delivered -- most often on a boot's first multi-chunk prefill (deterministic 2/2 in the 08-14 bisect, again on 08-15), forcing a boot-ramp ops protocol. Root cause: MPS exposes exactly one stream (torch.Stream(mps) always returns stream_id 0), so AsyncOutput's CUDA cross-stream choreography (set_stream, copy_stream.wait_stream(main_stream), generic torch.Event recorded on the fake copy stream) buys no overlap on Metal and routes step completion through an event that can lose its signal on timing-sensitive boots. On Metal, skip the stream context, wait_stream, and generic Event; enqueue the same non-blocking D2H copies on the only stream and record a native torch.mps.Event there for get_output() to wait on. VLLM_QC_ASYNC_OUT_DRAIN=1 swaps the event for a full torch.mps.synchronize() drain as an ops fallback (measured -3.7% decode, so not the default). CUDA/ROCm path is byte-identical. Validated over five fresh dsv4-xxs-1 boots: all pinned anchors bit-exact (8-tok, off1-2000, 2500x64), decode 32.0 tok/s (best wall of the day), both historical wedge-trigger protocols clean, kill-switch boot clean.
auroter
pushed a commit
to auroter/SlimServe
that referenced
this pull request
Aug 17, 2026
…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).
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
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.
Fixes the Metal serving bug where the engine could hang forever on a boot's first large prefill request, which until now required a boot-ramp workaround (serve a short request with real decode steps before any multi-chunk prefill, or the boot was lost and needed a restart).
The bug
On timing-sensitive boots the engine parked forever inside
MPSEvent::synchronizewaiting on the async-output copy event: GPU idle, no command-buffer errors, the completion signal simply never arrived. Reproduced deterministically in the 08-14 bisect (2/2 boots) and again during the 08-15 merge gate. Once parked, the boot was unrecoverable.Root cause
Apple exposes one GPU queue, and
torch.Stream(mps)always returns the same stream (id 0). The shared async-output path still ran the full CUDA cross-stream choreography — switch to a "copy stream",wait_streamon the "main stream", record a generictorch.Eventon the copy stream. On Metal that dance buys zero overlap (there is only one stream) and routes every step's completion through an event that can lose its signal.The fix
On Metal, skip the choreography entirely: the same non-blocking device-to-host copies enqueue on the one real stream, and completion is a native
torch.mps.Eventrecorded there — no stream switching, no cross-stream waits. The CUDA/ROCm path is byte-identical. An ops fallback (VLLM_QC_ASYNC_OUT_DRAIN=1) replaces the event with a full stream drain; it costs about 4% decode throughput, so it is not the default.Validation (dsv4-xxs-1, M1 Ultra, five fresh boots)
One honest caveat: the wedge is a host-timing race and had drifted out of its trigger window on the test box this week (the unfixed build also ran the triggers clean today). The fix rests on the structural argument — after it, there is no cross-stream event on the completion path left to lose — plus the earlier deterministic bisect evidence. The boot-ramp ops protocol stays recommended until a longer soak confirms it is no longer needed.