[Spec decode] Support Kimi-K3 DCP with DSpark by wzhao18 · Pull Request #52188 · vllm-project/vllm

wzhao18 · 2026-08-13T16:10:03Z

Purpose

This PR adds support for running Kimi-K3 decode context parallel with DSpark with FlashinferMLA and Tokenspeed as target causal attention backend and Tokenspeed as the draft non-causal backend.

Test Plan

Kimi K3 GSM8k with the different backend combination.

Test Result

Default (no backend specified):

vllm serve moonshotai/Kimi-K3 \
  --tensor-parallel-size 8 \
  -dcp 8 \
  --load-format fastsafetensors \
  --no-enable-flashinfer-autotune \
  --trust-remote-code \
  --language-model-only \
  --attention-config '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true}' \
  --kv-cache-dtype fp8 \
  --speculative-config '{"model":"Inferact/Kimi-K3-DSpark","method":"dspark","num_speculative_tokens":4,"rejection_sample_method":"standard"}'

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9606|±  |0.0054|
|     |       |strict-match    |     5|exact_match|↑  |0.9606|±  |0.0054|

FlashinferMLA target + Tokenspeed draft:

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9613|±  |0.0053|
|     |       |strict-match    |     5|exact_match|↑  |0.9613|±  |0.0053|

Tokenspeed target + Tokenspeed draft:

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9621|±  |0.0053|
|     |       |strict-match    |     5|exact_match|↑  |0.9621|±  |0.0053|

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

claude

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

GirasoleY

Thanks for adding DCP + DSpark support! Overall the change looks good.

I left a few inline comments about moving the repeated metadata construction out of hot path.

GirasoleY · 2026-08-13T23:35:52Z

+            remainder = torch.clamp(
+                per_query_global_lens % dcp_span - self.dcp_rank * interleave,
+                min=0,
+                max=interleave,
+            )
+            seq_lens = per_query_global_lens // dcp_span * interleave + remainder
+            block_table = block_table.repeat_interleave(query_len, dim=0)
+            q = q.unsqueeze(1)


This is inefficient as it run for every mla forward.

Let's construct the dcp related seqlen/query_start_loc/block table expansion in metadata builder, then reuse them for all MLA layers.

The same metadata can be reused in combine path as well.

Thanks for the suggestion. It is now changed to run only once per forward pass and reused across layers. Please check again.

GirasoleY · 2026-08-13T23:44:35Z

@@ -118,6 +122,22 @@ class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
    # Non-causal DSpark blocks are flattened to single-token rows in forward_mqa.
    supports_non_causal_multi_token_decode: ClassVar[bool] = True


Enable supports_non_causal_multi_token_dcp for flashinfer_mla as well?

supports_non_causal_multi_token_dcp is not supported by flashinfer_mla by this PR. will let it use tokenspeed backend for the draft model for now.

Make sense. I added a refactor PR to remove flashinfer specific changes.

Thanks! The refactor looks great.

GirasoleY

Added a refactor to move the flattened decode metadata into MLACommonDecodeMetadata. the rest LGTM :)

…fix in separate PR) This reverts commit f9d4516d49113369e4f6ba5b650daae89ac51d8c. Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>

Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>

github-actions · 2026-08-17T16:18:48Z

✅ Triggered Buildkite CI #84210 for commit 073f76abcee4.

wzhao18 · 2026-08-17T17:22:30Z

/ci retry

wzhao18 · 2026-08-17T19:03:05Z

/ci retry

github-actions · 2026-08-17T19:03:23Z

✅ Queued 1 failed job(s) for retry in Buildkite CI #84210.

Preserve the TP-only FlashInfer DSpark path from vllm-project#50000 while removing the FlashInfer-specific causal DCP flattening and cached metadata introduced by vllm-project#52188. Signed-off-by: Summer Yang <girasoleyang@gmail.com>

Upstream vllm-project#52188 landed DCP-aware draft slot mapping for DFlash/DSpark (cp_local_slot + prepare_dcp_local_seq_lens), superseding the localizer this branch carried. Keep upstream's version, which also handles the null block and the rejected context suffix, and reduce this branch to what it still adds: the KV-head-replicated draft cache for dense (GQA/MHA) drafts. The query block's K/V still must not be slotted into a head-replicated cache, so that behavior moves behind PAD_QUERY_SLOTS_UNDER_CP, derived from the draft attention layers rather than forced for every DCP run. Upstream's MLA draft path (Kimi-K3) keeps its real query slots. Signed-off-by: Mikhail Kostryukov <mike@triptrack.net> Co-authored-by: Claude <noreply@anthropic.com>

Signed-off-by: wzhao18 <wzhao18.sz@gmail.com> Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com>

### What this PR does / why we need it? The PR adapts vllm-ascend for compatibility with the latest vLLM main (commit `ba07e4a4`). | Files | Upstream vLLM change | vllm-ascend adaptation | |-------|---------------------|------------------------| | `.github/vllm-main-verified.commit` | — | Updated verified main commit hash from `cdc4824a21` to `ba07e4a48` | | `tests/e2e/conftest.py` | [vllm#53272](vllm-project/vllm#53272) — upstream plans to remove native Hunyuan V1/VL; [vllm#51665](vllm-project/vllm#51665) — dropped HunYuanVL `lm_head` workaround | Added `skip` condition for HunyuanVL e2e when `vllm_version_is("0.27.1")` is False (vLLM main) | | `tests/ut/core/test_profiling_chunk.py` | — vLLM main `Scheduler.__init__` reads `model_config.uses_mrope`, which infinitely recurses on a bare MagicMock | Version-gated: sets `type(model_config).uses_mrope = PropertyMock(return_value=False)` on main | | `tests/ut/core/test_recompute_scheduler.py` | — vLLM main `Scheduler.add_request` reads `spec_decode_metrics_level` | Version-gated: sets `scheduler.spec_decode_metrics_level = "none"` on main | | `tests/ut/patch/platform/test_patch_structured_output.py` | — Upstream changed structured output validation error type from `ValueError` to `VLLMValidationError` | Version-gated: `error_type = ValueError if vllm_version_is("0.27.1") else VLLMValidationError` used in all three fake validation functions and `pytest.raises` | | `vllm_ascend/attention/attention_v1.py` | [vllm#52839](vllm-project/vllm#52839) — moved `pcp.py` from `vllm.model_executor.layers.attention.pcp` to `vllm.v1.attention.ops.pcp` | Version-gated `_gather_prefill_cache_inputs` import path | | `vllm_ascend/compilation/acl_graph.py` | [vllm#49134](vllm-project/vllm#49134) — `get_current_vllm_config()` now raises `AssertionError` when called outside `set_current_vllm_config()` context | `update_full_graph_params` version-gated: main branch wraps `get_impl_cls()`/`update_graph_params` in `with set_current_vllm_config(vllm_config):` | | `vllm_ascend/models/deepseek_mtp.py` | [vllm#53106](vllm-project/vllm#53106) — removed `skip_prefixes` kwarg from `AutoWeightsLoader.__init__` | `AscendGlmMoeDsaForCausalLM.load_weights` version-gated: 0.27.1 uses `skip_prefixes=["rot."]`; main uses `WeightsMapper(orig_to_new_prefix={"rot.": None})` passed via `load_weights(weights, mapper=mapper)` | | `vllm_ascend/spec_decode/llm_base_proposer.py` | [vllm#52861](vllm-project/vllm#52861) — added `DeepseekV32MTPModel` to MTP architecture set in `model_returns_tuple()` | `model_returns_tuple` version-gated: 0.27.1 checks `{"DeepSeekMTPModel", "KimiK3MTPModel"}`; main also includes `"DeepseekV32MTPModel"` | | `vllm_ascend/worker/v2/spec_decode/dflash/speculator.py` | [vllm#52188](vllm-project/vllm#52188) — added `cp_rank`, `CP_SIZE`, `CP_INTERLEAVE` params to `_prepare_dflash_inputs_kernel` for DCP support | Entire `_prepare_dflash_inputs_kernel_ascend` kernel duplicated under `vllm_version_is("0.27.1")` gate: 0.27.1 uses 30 pos + 3 constexpr (no DCP params); main uses 31 pos + 5 constexpr | ### Does this PR introduce _any_ user-facing change? ### How was this patch tested? - vLLM version: v0.27.1 - vLLM main: vllm-project/vllm@cdc4824 --------- Signed-off-by: hfadzxy <starmoon_zhang@163.com>

… DFlash/DSpark (vllm-project#48392) Taken from PR vllm-project#48392 head (origin/pr-48392, 8/21) as a net diff against main: the author re-based onto vllm-project#52188 (cp_local_slot in the input-prep kernel), so the earlier in-kernel _pos_to_slot/FORCE_PAD_UNDER_CP approach is superseded by a host-side dcp_kv_head_replicated flag plus a PAD_QUERY_SLOTS_UNDER_CP tl.constexpr, keeping cp_local_slot untouched. - Attention: cache_num_kv_heads decouples stored KV heads per block from the impl's projection head count. - qwen3_dflash: dcp_kv_head_replicas() + precompute all-gather replicate the projected K/V across the DCP group before the cache write. - dflash speculator: PAD query slots under DCP for a replicated draft cache (query block attends fresh K/V, never the cache). - flash_attn: descale=None in _forward_with_dcp for non-quantized caches. - config/vllm.py: allow the sharded case only for Qwen3-family DFlashDraftModel/Qwen3DSparkModel with a non-quantized draft KV cache; keep rejecting v1 runner, DeepSeek-V4 DSpark, other archs, quantized draft caches, and partial replication. - test: 3 new tests in test_dflash_prepare_inputs.py (DCP partition, query-slot PAD for replicated cache, no-DCP no-op). Needed for GLM-5.2 DSpark (Qwen3 draft, MHA 64 KV heads, TP4/DCP4) tests. Signed-off-by: Leoyzen <leoyzen@gmail.com>

wzhao18 requested review from LucasWilkinson, MatthewBonanni, ProExpertProg, WoosukKwon, alexm-redhat, houseroad, mgoin, njhill, pavanimajety, robertgshaw2-redhat, tlrmchlsmth, yewentao256, youkaichao and zhuohan123 as code owners August 13, 2026 16:10

claude Bot reviewed Aug 13, 2026

View reviewed changes

mergify Bot added kimi k3 nvidia mrv2 Model Runner V2 specific labels Aug 13, 2026

github-project-automation Bot added this to NVIDIA Aug 13, 2026

GirasoleY reviewed Aug 13, 2026

View reviewed changes

wzhao18 requested review from ApostaC, heheda12345, ivanium, orozery and ywang96 as code owners August 14, 2026 00:49

wzhao18 force-pushed the wzhao/dcp-dspark branch from eecd761 to 7ecab4c Compare August 14, 2026 14:48

GirasoleY mentioned this pull request Aug 14, 2026

Refine shared DCP decode metadata wzhao18/vllm#1

Merged

GirasoleY approved these changes Aug 14, 2026

View reviewed changes

ivanium added the ready ONLY add when PR is ready to merge/full CI is needed label Aug 14, 2026

wzhao18 added 2 commits August 17, 2026 12:18

Revert "Fix mamba prefix caching with DCP on eagle block drop" (will …

e6ecfdd

…fix in separate PR) This reverts commit f9d4516d49113369e4f6ba5b650daae89ac51d8c. Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>

Fix CI failure due to bad import

073f76a

Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>

wzhao18 force-pushed the wzhao/dcp-dspark branch from 6d6cb7d to 073f76a Compare August 17, 2026 16:18

ivanium merged commit d1e3eee into vllm-project:main Aug 17, 2026
141 checks passed

github-project-automation Bot moved this from Ready to Done in NVIDIA Aug 17, 2026

GirasoleY mentioned this pull request Aug 20, 2026

[Cleanup][MLA] Remove FlashInfer DSpark DCP support #53139

Merged

drakosha mentioned this pull request Aug 21, 2026

[Spec Decode] DFlash/DSpark draft support under decode context parallelism #48392

Open

andyluo7 mentioned this pull request Aug 21, 2026

[ROCm][MLA][DCP] Support causal multi-token verification #51705

Merged

zhangxinyuehfad mentioned this pull request Aug 24, 2026

[Ci] main2main vllm 0821 vllm-project/vllm-ascend#14746

Merged

zufangzhu pushed a commit to zufangzhu/vllm that referenced this pull request Aug 24, 2026

[Spec decode] Support Kimi-K3 DCP with DSpark (vllm-project#52188)

3c6b271

Signed-off-by: wzhao18 <wzhao18.sz@gmail.com> Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com>

fivechenxi mentioned this pull request Aug 25, 2026

[Kimi K3] Integrate B300 DCP and DSpark baseline fixes fivechenxi/vllm#1

Closed

nvyutwu added a commit to nvyutwu/vllm that referenced this pull request Aug 25, 2026

Kimi-K3: materialize shipped DSpark + DCP overlay (vllm-project#52188)

bcd110a

GirasoleY mentioned this pull request Aug 27, 2026

[Attention][DCP] Use FlashInfer native CP for MLA decode #54012

Merged

sunny-infra mentioned this pull request Aug 28, 2026

[RFC]:Train a Dense MLA DSpark draft for GLM-5.2 (DCP-compatible) vllm-project/speculators#1052

Open

seongyun1104 mentioned this pull request Sep 9, 2026

[BugFix] Drafter attention backend auto-selection picks FlashInfer on SM90 fix #48579

Open

tarik-sarac mentioned this pull request Sep 17, 2026

[ROCm][Kimi-K3] Add moonmath kernels for gfx942 MXFP4 SiTU MoE and MLA decode, with decode context parallelism #57412

Open

		@@ -118,6 +122,22 @@ class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
		# Non-causal DSpark blocks are flattened to single-token rows in forward_mqa.
		supports_non_causal_multi_token_decode: ClassVar[bool] = True

Uh oh!

Conversation

wzhao18 commented Aug 13, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Test Result

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

GirasoleY left a comment

Choose a reason for hiding this comment

Uh oh!

GirasoleY Aug 13, 2026

Choose a reason for hiding this comment

Uh oh!

wzhao18 Aug 14, 2026

Choose a reason for hiding this comment

Uh oh!

GirasoleY Aug 13, 2026

Choose a reason for hiding this comment

Uh oh!

wzhao18 Aug 14, 2026

Choose a reason for hiding this comment

Uh oh!

GirasoleY Aug 14, 2026

Choose a reason for hiding this comment

Uh oh!

wzhao18 Aug 14, 2026

Choose a reason for hiding this comment

Uh oh!

GirasoleY left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Aug 17, 2026

Uh oh!

wzhao18 commented Aug 17, 2026

Uh oh!

wzhao18 commented Aug 17, 2026

Uh oh!

github-actions Bot commented Aug 17, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

wzhao18 commented Aug 13, 2026 •

edited by github-actions Bot

Loading