attention: support quantized KV cache in MLA DCP context prefill - #14
Closed
voipmonitor wants to merge 50 commits into
Closed
attention: support quantized KV cache in MLA DCP context prefill#14voipmonitor wants to merge 50 commits into
voipmonitor wants to merge 50 commits into
Conversation
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Tests: .venv/bin/python -m pytest tests/kernels/attention/test_cp_lse_ag_rs.py -v Co-authored-by: OpenAI Codex <codex@openai.com>
Tests: .venv/bin/python -m pytest tests/config/test_virtual_tp.py tests/engine/test_arg_utils.py::test_virtual_tp_sharding_cli_alias -v Co-authored-by: OpenAI Codex <codex@openai.com>
Tests: .venv/bin/python -m ruff check vllm/compilation/b12x_capture.py vllm/compilation/cuda_graph.py vllm/v1/worker/gpu/cudagraph_utils.py vllm/envs.py Co-authored-by: OpenAI Codex <codex@openai.com>
Tests: .venv/bin/python -m ruff check vllm/config/kernel.py Co-authored-by: OpenAI Codex <codex@openai.com>
Tests: .venv/bin/python -m ruff check vllm/model_executor/kernels/linear/__init__.py vllm/model_executor/kernels/linear/scaled_mm/__init__.py vllm/model_executor/kernels/linear/scaled_mm/b12x.py vllm/model_executor/layers/quantization/fp8.py vllm/model_executor/warmup/deep_gemm_warmup.py Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Filter safetensors shards and tensors by checkpoint weight prefixes so draft model loading uses the index instead of scanning unrelated shards. Warm up DeepSeek V4 compressor Triton signatures before KV cache allocation and align sparse MLA decode metadata with speculative rows under DCP. Co-authored-by: Codex <codex@openai.com>
Route B12X sparse-indexer prefill through the paged compressed indexer with row-shared page tables, and size prefill chunks against the B12X scorer supertile so workspace reservations match runtime. Co-authored-by: OpenAI Codex <codex@openai.com>
Pass the active request count into autoregressive draft sampling so padded FULL-cudagraph rows cannot overwrite draft logits for real requests. Co-authored-by: OpenAI Codex <codex@openai.com>
GLM/Kimi DCP decode uses an uncompressed MLA indexer cache (compress_ratio == 1), but the sparse indexer still has to write and read through DCP rank-local cache pages. The previous path only remapped slots when compress_ratio > 1, so DCP4 wrote the indexer K cache with global slot ids while B12X attention consumed it through a DCP-local page table. Use DCP-local slot mapping whenever decode-context parallelism is active, and feed DCP-local seq_lens to the decode sparse-indexer metadata. For MTP/native expanded decode rows, compute the per-token lengths from global seq_lens first, then convert the expanded lengths to DCP-local lengths so query_start_loc arithmetic remains correct. Validated on GLM-5.1 NVFP4 DCP4 nomtp with B12X_MLA_SPARSE, V2 model runner, FULL+PIECEWISE CUDA graphs, 50k context smoke, and cc1 decode bench.
Automatically apply B12X virtual TP padding when supported B12X attention/MoE paths need compatible TP shapes. Remove the public virtual TP sharding and alignment CLI/config knobs, use fixed B12X capability alignments, and warn when automatic padding is enabled. Co-authored-by: OpenAI Codex <codex@openai.com>
Port vllm-project#39945 onto the MiMo Black Benediction branch. The fix keeps same-sized prefill batches from being misclassified by the uniform_batch path, which DFlash relies on during draft proposal.
Port vllm-project#40898 onto the MiMo Black Benediction branch. This preserves DFlash sliding-window layer types, supports multiple draft KV groups, keeps shared raw target/draft KV handling, and normalizes DFlash auxiliary target layer ids while retaining local MiMo/B12X cache and quantization fixes.
Port vllm-project#43081 onto the MiMo Black Benediction branch. This allows non-causal DFlash batches to use FlashInfer native prefill where supported and keeps unsupported DCP/TRTLLM/NVFP4 cases out of that path.
Port the delta from vllm-project#43200 on top of the DFlash SWA stack. FlashInfer metadata builders are now split by backend-owned layer attributes such as window_left, logits soft cap, scale, and sinks. The local runner grouping also keeps the existing num_heads_q split so DFlash/target head-count differences remain isolated.
MiMo-V2.5-Pro-FP4-DFlash stores SWA in dflash_config.use_swa and dflash_config.swa_window_size, while Transformers materializes layer_types as full_attention and sliding_window as None. Map that config to sliding_attention layers so the DFlash drafter actually uses SWA instead of silently running full attention.
FP8 checkpoints with store_dtype=mxfp4 (MiMo V2.5 Pro) routed moe_backend=b12x to GptOssMxfp4MoEMethod, which hardcodes the GPT-OSS clamped SwiGLU (gemm1_alpha=1.702, beta=1.0, swiglu_limit=7.0). The B12X W4A16 kernel honors swiglu_limit and clamped gate at +7 / up at +-7 in all 69 MiMo MoE layers. MiMo is trained with unclamped SiLU, so this systematically distorted expert outputs and pushed long generations into repetition loops; the FlashInfer Cutlass path already used Mxfp4MoEMethod (limit=None) and did not loop. Route b12x to Mxfp4MoEMethod alongside cutlass/trtllm/deep_gemm. Verified: estonia C=6 R=12 probe now 12/12 stop, 0 hit max_tokens; temp-0 6k-token generations show no tail cycling. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
MiMo V2.5 DFlash acceptance was ~1.98-2.0 mean accepted length vs the official 3.18-6.30 claims. Three divergences from the reference implementation (checkpoint dflash/dflash.py + sglang DFLASH): 1. Mask embedding: the FP4 checkpoint's target embed row for mask_token_id 151669 is zeroed (norm 8e-4); the trained vector ships in dflash/mask_embedding.pt (norm 3.9) and was never loaded, so every masked draft slot embedded as ~zero. Load it after embedding sharing and write the row TP-aware into the shared VocabParallelEmbedding. 2. Block masking: all-SWA draft layers (MiMo: all 5) were forced to causal attention because the triton kernel could not combine non-causal with sliding window. Reference DFlash drafts the block bidirectionally with the sliding window only bounding context lookback. The triton mask now applies the window as a backward-only bound in non-causal mode, and dflash.py no longer special-cases SWA layers to causal. 3. Last aux feature: the reference extracts HF hidden_states[id+1] for target_layer_ids [0,15,31,47,69]; index 70 is the post-final-norm hidden. Capture the final aux layer after self.norm in mimo_v2 (escape hatch VLLM_DFLASH_PRENORM_LAST_AUX=1). Also folds in the prior overlay WIP this builds on: non-causal triton attention plumbing, DFlash SWA layer typing/sink-bias loading in qwen3_dflash, and draft backend inheritance from the target. Measured on TP8 b12x/fp8-KV with num_speculative_tokens=7 (block 8), temp-0 probes: mean accept length code 4.96, math 4.99, chat 2.74 (official: HumanEval 4.54, Math500 5.56, MT-Bench 3.18); per-position acceptance now decays smoothly (0.89..0.32) instead of collapsing to 0 after position 2. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
DFlash decode collapsed from 169 tok/s at ctx 0 to 30-44 tok/s at 128k-190k context (below the 61 tok/s no-MTP baseline). Two causes: 1. Non-causal sliding-window tile pruning was missing: the draft's SWA layers scanned the full context every step because tile pruning in compute_tile_loop_bounds was gated on CAUSAL. Add a non-causal branch that prunes leading tiles via the backward window bound while keeping tile_end at num_tiles (the rest of the draft block stays visible), restoring the constant-in-context draft cost the DFlash design promises. 2. Spec-decode verify (query_len 8) was locked out of the 3D segment-parallel softmax path (gated on max_seqlen_q == 1) and ran the 2D kernel with a grid of ~9 CTAs for long KV scans. Gate the 3D path on total query tokens fitting the segment buffers instead; single-token decode behavior is unchanged (q tokens == num_seqs), prefills stay 2D via the same bound. Measured (TP8, fp8 KV, dflash ntok7, temp 0, ~190k ctx): step time 97ms -> 30.1ms, decode 44 -> 153.8 tok/s with mean accept length 4.65; ctx-0 decode 175 tok/s and acceptance probes unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Attention.get_kv_cache_spec asserts against model_config.use_mla for sliding-window layers, which rejected DFlash drafts on MLA targets (Kimi K2.6: 'MLA is not supported for slidingwindow') even though the draft layer itself is plain GQA. DFlashAttention now builds its FullAttentionSpec directly for SWA layers (it widens them anyway since context KVs must persist) instead of converting the parent's SlidingWindowSpec. Verified: Kimi K2.6 (KimiK25ForConditionalGeneration, TRITON_MLA, fp8 KV, TP8) + SubSir/Kimi-K2.6-DFlash-tmp draft (qwen3 DFlash, 5x SWA-2048 + 1 full layer, YaRN rope) starts and serves; temp-0 acceptance probe: mean accept length math 4.81 (official SGLang claim 4.8), code 3.28, chat 2.02. MiMo DFlash acceptance probes unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The grouped MLA decode kernel (_fwd_grouped_kernel_stage1) with Lk=576 (DeepSeek/Kimi latent KV) and num_stages=2 compiles to ~100KB of shared memory, which exceeds the 101376-byte per-block limit on sm86/sm89/sm120 (RTX PRO 6000 Blackwell: required 102400). Kimi K2.6 with TRITON_MLA crashed on the first pure-decode step; spec-decode verify steps (q_len > 1) take the extend path and were unaffected. Drop to a single pipeline stage when the device's opt-in shared memory limit (or known capability) is at or below the overflow threshold. H100/B200-class GPUs (>200KB smem) keep num_stages=2. Verified: Kimi K2.6 TP8 TRITON_MLA fp8-KV no-MTP now serves and decodes coherently on RTX PRO 6000 Blackwell. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Black-benediction carried only the generic upstream TRITON_MLA decode path; the voipmonitor glm51-v6-awq-mxfp8 branch (kimi v8 runtime, via 868600a) has a tuned rewrite: dynamic per-batch num_kv_splits selection, per-bucket kernel configs from bench/tune_triton_mla.py (tuned on RTX PRO 6000 Blackwell), CG-bucket attn_logits pooling, and direct stage1 kernel launches. The tuned configs also stay within the 99KB sm120 shared-memory limit, so the generic-path num_stages fallback no longer binds for decode. Kimi K2.6 TP8 fp8-KV no-MTP 0k cc1 decode: 79.2 -> 91.7 tok/s (v8-image reference: 98.5; the remainder tracks the cuBLAS 13.4.1.2-vs-13.4.0.1 and toolchain delta of that image). Only the import of get_current_vllm_config_or_none needed no change; both files are verbatim from the source branch. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Adds a DFlashSpeculator for the V2 runner's speculator framework. The draft proposes a whole block [bonus, MASK x k] in ONE eager non-causal forward; context comes from the target's aux hidden states, fc-combined and projected into the draft KV cache at the target step's slots (the draft never runs autoregressively over the context, so per-step draft cost is constant). Pieces: - gpu/spec_decode/dflash/speculator.py: the speculator. Reuses the model-side DFlash machinery (precompute_and_store_context_kv, mask embedding loading, non-causal sliding-window attention). Drafts run eagerly (no draft CUDA graphs yet). - gpu/spec_decode/__init__.py: register method 'dflash'. - eagle/eagle3_utils.py: derive aux layers from dflash_config target_layer_ids (+1, HF hidden_states[i+1] semantics). - gpu/model_runner.py: enable aux hidden state outputs for dflash. - gpu/attn_utils.py: optional causal flag on build_attn_metadata. - config/vllm.py: allow method 'dflash' + its parallel_drafting on V2. - mla/triton_mla.py: oversized eager batches (e.g. spec-verify warmup above the CG capture limit) fall back to transient buffers instead of asserting. Bug found on the way: load_eagle_model resolves lm_head on the top-level target module, which silently leaves weightless drafts with a random head behind multimodal wrappers (KimiK25ForConditionalGeneration). The DFlash loader re-shares from get_language_model(). Measured (Kimi K2.6 TP8 fp8-KV, dflash ntok7, temp 0, 0k cc1): 160.8 tok/s vs 114.2 on the V1 runner (+41%); mean accept length math 5.15 / code 3.46 / chat 3.09 (official SGLang claim: math 4.8). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Register the B12X_ATTN backend, wire MiMo V2 to use the per-layer V-head paged backend, and add caller-scratch-owned B12X paged attention support. Use the least common multiple when unifying KV cache page sizes so non-divisible page sizes can share a cache tensor, with focused coverage for the new behavior. Add serve helper scripts for Kimi K2.6, MiMo V2.5 Pro DFlash, and RedHatAI Qwen3 VL NVFP4. Tests: git diff --check Tests: .venv/bin/python -m py_compile vllm/v1/attention/backends/b12x_attn.py vllm/model_executor/models/mimo_v2.py vllm/v1/attention/backends/registry.py vllm/v1/core/kv_cache_utils.py tests/v1/core/test_kv_cache_utils.py Tests: .venv/bin/python -m pytest tests/v1/core/test_kv_cache_utils.py::test_unify_kv_cache_spec_page_size_uses_lcm_for_non_divisible_pages tests/v1/core/test_kv_cache_utils.py::test_get_kv_cache_config_one_worker -q Co-authored-by: OpenAI Codex <codex@openai.com>
DCP chunked-context prefill gathered KV pages with cp_gather_cache, which requires src/dst dtype equality and so crashed with 'src_cache and dst must have the same dtype' for quantized (fp8) KV caches. Route the gather through gather_and_maybe_dequant_cache when the cache dtype differs from the workspace dtype, build the padded_local_token_to_seq grid the dequant gather needs, and pass the layer k_scale instead of hardcoding None. Ported from the glm51-v6 lineage where DCP+fp8 KV was validated. Validated on Kimi-K2.7-Code TP8 (TRITON_MLA, kv fp8, V2 runner): DCP 2/4/8 boot, math smoke, and temp-0 needle retrieval at 4k-67k prompt tokens all pass; DCP1 parity confirmed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI 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:
✨ Finishing Touches🧪 Generate unit tests (beta)
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 was referenced Jun 12, 2026
Author
|
Superseded by #16 — same logical change (quantized KV in MLA DCP context prefill) rebased clean onto current |
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
3 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.
What
DCP (
--decode-context-parallel-size > 1) crashed on any chunked-context prefill with a quantized KV cache:_context_parallel_compute_prefill_contextgathered KV pages withcp_gather_cache, which requires src/dst dtype equality — impossible when the cache is fp8 and the gather workspace is in model dtype (the call site also hardcodedk_scale=Noneunder the "DCP not support scaled kvcache" assumption).This routes the gather through
gather_and_maybe_dequant_cachewhen the cache dtype differs from the workspace dtype, builds thepadded_local_token_to_seqgrid that the dequantizing gather needs, and passes the layerk_scalethrough. Unquantized caches keep the existingcp_gather_cachepath. The kernel already supports this signature; the change is Python-only. Ported from the glm51-v6 lineage where DCP+fp8 KV was previously validated.Why not a duplicate
No open PR touches MLA DCP or
cp_gather_cache(checkedgh pr list; only #8 is open, unrelated MTP draft fix).Tests
Kimi-K2.7-Code TP8, TRITON_MLA,
--kv-cache-dtype fp8, V2 runner,max_num_batched_tokens 8192(chunked prefill), PR13 image + this file overlaid:finish_reason=stop).python -m py_compileclean.AI assistance
This change was developed with AI assistance (Claude). I reviewed every changed line and ran the tests above.
🤖 Generated with Claude Code