Make speculative decoding honor structured-output grammars - #4
Draft
ropoctl wants to merge 14 commits into
Draft
Conversation
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).
…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.
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.
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
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 |
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.
This is a dependency-stacked draft for local DSV4 Metal validation. It currently includes the open work from #2 (M1 Ultra campaign), #3 (Metal async-output completion), and #1 (Responses custom tools); review the DSpark commits after those dependencies land.
The new work:
Local DSV4 evidence: the original reasoning-boundary grammar crash was eliminated and grammar-aware DSpark reached mean acceptance 6.00 with observed 80-100% draft acceptance. The run then exposed the tokenizer metadata defect: BOS token 0 was legal inside filename /(.+)/. The offline real-tokenizer regression now rejects BOS while preserving ordinary filename tokens. The server was stopped at the user request before a post-fix end-to-end rerun, so this remains a draft rather than claiming full apply_patch acceptance.
Validation:
No model-envelope literals or DeepSeek token IDs are added to the inference core; DSML, Qwen, Kimi, and Muse syntax remains in their structural-tag/chat-format adapters.