PR #40609: [Core] Enable FP8 KV cache with DCP for MLA

• Repository: vllm-project/vllm
• Author: voipmonitor
• State: open | merged: False
• Link: https://github.com/vllm-project/vllm/pull/40609

Description (problem / solution / changelog)

Summary

• enable MLA decode-context-parallel prefill with FP8 KV cache
• gather and dequantize KV cache correctly for the DCP prefill path
• add distributed coverage for kv_cache_dtype=fp8

Context

This PR supersedes the dormant work in #34795 and is tracked from #40608.

Scope

• vllm/model_executor/layers/attention/mla_attention.py
• tests/distributed/test_context_parallel.py

Why

Kimi-style MLA serving on Blackwell relies on TRITON_MLA + DCP + FP8 KV, and the DCP prefill path needs explicit FP8 cache handling to work correctly.

Validation

Validated locally on Blackwell/SM120 with Kimi MLA workloads and added test coverage for DCP + FP8 KV.

Changed files

• tests/distributed/test_context_parallel.py (modified, +10/-1)
• vllm/model_executor/layers/attention/mla_attention.py (modified, +92/-19)

---

PR #40610: [SpecDecode] Fix async proposer synchronization

• Repository: vllm-project/vllm
• Author: voipmonitor
• State: open | merged: False
• Link: https://github.com/vllm-project/vllm/pull/40610

Description (problem / solution / changelog)

Summary

• move prepare_inputs_event.record() so speculative proposer GPU work is included in the synchronization window
• prevent the next async batch from reusing persistent state and block-table-related metadata while proposer work from the previous batch is still in flight

Context

Tracked from #40608.

Root cause

This came from debugging async scheduling + speculative decode ordering, not from a single clean deterministic stacktrace.

The key issue was the lifetime of prepare_inputs_event versus the actual proposer GPU work:

• the event was being recorded from the input-preparation path, which made the step look "finished" too early
• proposer-side GPU work for speculative decode runs later from sample_tokens()
• that ordering allowed the next batch to enter execute_model() / _update_states() and start mutating persistent batch state while the proposer from the previous batch was still reading it on the GPU

In practice this is a concurrency-sensitive race: the failure mode is timing-dependent and can show up as nondeterministic instability / stale state usage rather than one obvious reproducible exception.

The shared state at risk here includes the persistent structures used across batches in the async path, especially block-table-related metadata and other proposer inputs derived from the batch state.

Why this fix

This patch makes prepare_inputs_event represent the full GPU lifetime of the step, not just the input-preparation portion of it.

Concretely, it re-records the event after sample_tokens() finishes, which is the point where speculative proposer GPU work has also completed. That closes the race where the next batch could otherwise observe and mutate shared state too early.

I kept the change intentionally small and local to gpu_model_runner.py because the regression surface is concurrency-sensitive, and this is the least invasive way to restore the intended happens-before relationship.

Scope

• vllm/v1/worker/gpu_model_runner.py

Notes

This is kept as a draft PR because the code change is small but the regression surface is concurrency-sensitive.

Changed files

• vllm/v1/worker/gpu_model_runner.py (modified, +16/-0)

---

PR #40611: [SpecDecode] Allow draft-specific attention backend and KV dtype

• Repository: vllm-project/vllm
• Author: voipmonitor
• State: open | merged: False
• Link: https://github.com/vllm-project/vllm/pull/40611

Description (problem / solution / changelog)

Summary

• allow speculative draft models to use a different attention backend than the target model
• allow speculative draft models to use a different KV-cache dtype than the target model
• plumb the per-draft overrides through EAGLE setup

Context

Tracked from #40608.

Why

Kimi serving needs target/draft backend decoupling:

• the target model runs best with TRITON_MLA
• some draft models are non-MLA and need a different backend or KV-cache dtype

Without this split, otherwise valid target/draft pairings are forced onto a single backend/dtype choice.

Scope

• vllm/config/speculative.py
• vllm/v1/spec_decode/eagle.py
• vllm/v1/attention/backends/flash_attn.py
• vllm/v1/worker/cp_utils.py

Notes

Kept as draft for discussion because it broadens speculative-config surface area.

Changed files

• vllm/config/speculative.py (modified, +18/-1)
• vllm/v1/attention/backend.py (modified, +27/-0)
• vllm/v1/attention/backends/flash_attn.py (modified, +3/-2)
• vllm/v1/attention/backends/flashinfer.py (modified, +3/-2)
• vllm/v1/spec_decode/eagle.py (modified, +33/-4)
• vllm/v1/worker/cp_utils.py (modified, +4/-2)

---

PR #40612: [SpecDecode] Add local argmax helper for Llama Eagle3 drafts

• Repository: vllm-project/vllm
• Author: voipmonitor
• State: open | merged: False
• Link: https://github.com/vllm-project/vllm/pull/40612

Description (problem / solution / changelog)

Summary

• add get_top_tokens() support for LlamaForCausalLMEagle3
• enable local argmax reduction for Llama-based EAGLE3 drafts when vocab mapping is identity

Context

Tracked from #40608.

Why

Large-vocab speculative draft models benefit from reducing TP communication with local argmax reduction. This brings the Llama-based EAGLE3 path in line with the same optimization already being discussed for other draft families.

Scope

• vllm/model_executor/models/llama_eagle3.py

Related

• complements existing upstream work in #39419 rather than replacing it

Changed files

• vllm/model_executor/models/llama_eagle3.py (modified, +19/-0)

---

PR #40613: [SpecDecode] Add seq-length gate for speculative decode

• Repository: vllm-project/vllm
• Author: voipmonitor
• State: open | merged: False
• Link: https://github.com/vllm-project/vllm/pull/40613

Description (problem / solution / changelog)

Summary

• add an opt-in seq-length gate that disables speculative decode above a configured threshold
• make the gate effective both in the scheduler path and in the worker-side drafter fit check

Context

Tracked from #40608.

Why

For Kimi MLA long-context workloads, speculative decode can become slower than target-only decoding even when acceptance is healthy. An explicit seq-length gate is a pragmatic way to keep short-context wins without forcing long-context regressions.

Scope

• vllm/v1/core/sched/async_scheduler.py
• vllm/v1/worker/gpu_model_runner.py

Notes

This is intentionally a draft/discussion PR. The current implementation is an environment-based runtime knob used in local deployment recipes, not a proposed final public API shape.

Changed files

• vllm/v1/core/sched/async_scheduler.py (modified, +10/-1)
• vllm/v1/spec_decode/utils.py (modified, +22/-0)
• vllm/v1/worker/gpu_model_runner.py (modified, +38/-1)

---

PR #40614: [Attention] Tune TRITON_MLA for SM120 + FP8 decode

• Repository: vllm-project/vllm
• Author: voipmonitor
• State: open | merged: False
• Link: https://github.com/vllm-project/vllm/pull/40614

Description (problem / solution / changelog)

Summary

• cap num_kv_splits for B=1, SM120, FP8 KV MLA decode
• reduce BLOCK_H for the same narrow configuration

Context

Tracked from #40608.

Why

For Blackwell/SM120 MLA decode with FP8 KV cache, the default split heuristic can overshoot and increase split/merge overhead at large local sequence lengths. A narrower BLOCK_H setting is also useful for the same configuration.

Scope

• vllm/v1/attention/backends/mla/triton_mla.py
• vllm/v1/attention/ops/triton_decode_attention.py

Notes

This is intentionally a draft/discussion PR. The tuning is narrow (B=1, SM120, FP8 KV) and should be evaluated on its own merits rather than mixed with correctness changes.

Changed files

• vllm/v1/attention/backends/mla/triton_mla.py (modified, +13/-0)
• vllm/v1/attention/ops/triton_decode_attention.py (modified, +9/-0)