@tduncan / Nvidia-Optimization-Stack-VLLM / commits / d1e3eee6fb

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

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

Wei Zhao committed Aug 17, 2026 at 16:08 UTC d1e3eee6fb8ed3623241ef5c8e3ac533f775bff9

15 files changed +341 -51

tests/transformers_utils/test_dspark_mla_config.py

-23

@@ -140,26 +140,3 @@ def test_dspark_mla_speculative_config_preserves_architecture(tmp_path):

140 assert speculative_config.draft_model_config.architectures == ["K3DSparkModel"]

141 assert speculative_config.draft_model_config.hf_config.model_type == "k3_dspark"

142 assert speculative_config.draft_model_config.use_mla

143 -

144 -

145 -def test_dspark_mla_rejects_decode_context_parallelism(tmp_path):

146 - target_path = tmp_path / "target"

147 - draft_path = tmp_path / "draft"

148 - _write_target_config(target_path)

149 - _write_dspark_config(draft_path)

150 - target_config = ModelConfig(

151 - model=str(target_path), tokenizer_mode="skip", max_model_len=32768

152 - )

153 -

154 - with pytest.raises(ValueError, match="does not currently support decode context"):

155 - SpeculativeConfig(

156 - model=str(draft_path),

157 - method="dspark",

158 - num_speculative_tokens=8,

159 - target_model_config=target_config,

160 - target_parallel_config=ParallelConfig(

161 - tensor_parallel_size=2,

162 - decode_context_parallel_size=2,

163 - distributed_executor_backend="external_launcher",

164 - ),

165 - )

tests/v1/attention/test_flashinfer_mla_dcp.py

@@ -45,6 +45,7 @@ def test_flashinfer_mla_forward_uses_gathered_head_count(monkeypatch):

45 impl.bmm1_scale = 1.0

46 impl.bmm2_scale = 1.0

47 impl.need_to_return_lse_for_decode = True

48 + impl.dcp_world_size = 2

49 impl.num_heads = 6

50 impl.qk_nope_head_dim = 128

51 impl.kv_lora_rank = 512

tests/v1/attention/test_mla_backends.py

+77

@@ -837,9 +837,86 @@ def test_mock_mla_dcp_fp8_decode_gathers_quantized_query(

837 def test_tokenspeed_mla_noncausal_capability():

838 builder = tokenspeed_mla_module.TokenspeedMLAMetadataBuilder

839 assert builder.supports_non_causal_multi_token_decode

840 + assert builder.supports_non_causal_multi_token_dcp

841 assert tokenspeed_mla_module.TokenspeedMLABackend.supports_non_causal()

842

843

844 +def test_flashinfer_mla_dcp_multi_token_decode_uses_per_query_bounds(monkeypatch):

845 + flashinfer_mla_module = pytest.importorskip(

846 + "vllm.v1.attention.backends.mla.flashinfer_mla"

847 + )

848 +

849 + decode_call = None

850 +

851 + def fake_decode(**kwargs):

852 + nonlocal decode_call

853 + decode_call = kwargs

854 + query = kwargs["query"]

855 + output = torch.empty(*query.shape[:-1], 512, dtype=torch.bfloat16)

856 + lse = torch.empty(query.shape[0], query.shape[-2], dtype=torch.float32)

857 + return output, lse

858 +

859 + monkeypatch.setattr(

860 + flashinfer_mla_module,

861 + "trtllm_batch_decode_with_kv_cache_mla",

862 + fake_decode,

863 + )

864 + monkeypatch.setattr(

865 + flashinfer_mla_module,

866 + "_get_workspace_buffer",

867 + lambda return_lse: torch.empty(1, dtype=torch.int8),

868 + )

869 +

870 + impl = object.__new__(flashinfer_mla_module.FlashInferMLAImpl)

871 + impl.dcp_world_size = 2

872 + impl.dcp_rank = 1

873 + impl.cp_kv_cache_interleave_size = 1

874 + impl.need_to_return_lse_for_decode = True

875 + impl.kv_lora_rank = 512

876 + impl.qk_nope_head_dim = 128

877 + impl.qk_rope_head_dim = 64

878 + impl.bmm1_scale = 1.0

879 + impl.bmm2_scale = 1.0

880 +

881 + block_table = torch.tensor([[1], [2]], dtype=torch.int32)

882 + metadata = SimpleNamespace(

883 + num_decodes=2,

884 + num_decode_tokens=6,

885 + max_seq_len=7,

886 + causal=True,

887 + decode=SimpleNamespace(

888 + block_table=block_table,

889 + seq_lens=torch.tensor([5, 6], dtype=torch.int32),

890 + dcp_tot_seq_lens=torch.tensor([10, 13], dtype=torch.int32),

891 + flattened_block_table=None,

892 + flattened_seq_lens=None,

893 + query_len=0,

894 + ),

895 + )

896 + query = torch.empty(6, 2, 576, dtype=torch.bfloat16)

897 + kv_cache = torch.empty(3, 16, 576, dtype=torch.bfloat16)

898 +

899 + output, lse = impl.forward_mqa(

900 + query,

901 + kv_cache,

902 + metadata,

903 + SimpleNamespace(),

904 + )

905 +

906 + assert output.shape == (6, 2, 512)

907 + assert lse is not None

908 + assert lse.shape == (6, 2)

909 + assert decode_call is not None

910 + assert decode_call["query"].shape == (6, 1, 2, 576)

911 + torch.testing.assert_close(

912 + decode_call["seq_lens"],

913 + torch.tensor([4, 4, 5, 5, 6, 6], dtype=torch.int32),

914 + )

915 + torch.testing.assert_close(

916 + decode_call["block_tables"], block_table.repeat_interleave(3, dim=0)

917 + )

918 +

919 +

920 @pytest.mark.parametrize(

921 ("causal", "tokens_per_decode", "dcp_world_size", "dcp_rank"),

922 [

tests/v1/spec_decode/test_dflash_prepare_inputs.py

+25 -1

@@ -17,7 +17,14 @@ pytestmark = pytest.mark.skipif(

17 )

20 -def _run_prepare(*, target_positions: list[int], block_table_values: list[int]):

20 +def _run_prepare(

21 + *,

22 + target_positions: list[int],

23 + block_table_values: list[int],

24 + cp_rank: int = 0,

25 + cp_size: int = 1,

26 + cp_interleave: int = 1,

27 +):

28 device = torch.device("cuda")

29 max_num_reqs = 4

30 max_num_tokens = 16

@@ -86,6 +93,9 @@ def _run_prepare(*, target_positions: list[int], block_table_values: list[int]):

93 input_seeds,

94 block_table,

95 4,

96 + cp_rank,

97 + cp_size,

98 + cp_interleave,

99 123,

100 num_speculative_steps,

101 num_speculative_steps,

@@ -131,6 +141,20 @@ def test_prepare_dflash_inputs_excludes_rejected_context_suffix():

141 assert out.seeds[2].item() == 17

142

143

144 +def test_prepare_dflash_inputs_excludes_rejected_context_suffix_with_dcp():

145 + out = _run_prepare(

146 + target_positions=[10, 11, 12, 13],

147 + block_table_values=[0, 7, 8, 9],

148 + cp_rank=1,

149 + cp_size=2,

150 + cp_interleave=2,

151 + )

152 +

153 + assert out.context_positions[:4].tolist() == [10, 11, 0, 0]

154 + assert out.context_slot_mapping[:4].tolist() == [28, 29, PAD_SLOT_ID, PAD_SLOT_ID]

155 + assert out.query_slot_mapping[:3].tolist() == [PAD_SLOT_ID, PAD_SLOT_ID, 30]

156 +

157 +

158 def test_prepare_dflash_inputs_never_writes_the_null_block():

159 # The valid context uses logical block 0 and the replacement query uses

160 # logical block 1. Both map to the null block and must remain unwritable.

vllm/config/speculative.py

-10

@@ -1059,16 +1059,6 @@ class SpeculativeConfig:

1059 if self.method in ("dflash", "dspark"):

1060 self.parallel_drafting = True

1061

1062 - if (

1063 - self.method == "dspark"

1064 - and "K3DSparkModel" in self.draft_model_config.architectures

1065 - and self.target_parallel_config.decode_context_parallel_size > 1

1066 - ):

1067 - raise ValueError(

1068 - "MLA DSpark does not currently support decode context "

1069 - "parallelism; set decode_context_parallel_size=1."

1070 - )

1071 -

1062 if self.num_speculative_tokens is not None and hasattr(

1063 self.draft_model_config.hf_config, "num_lookahead_tokens"

1064 ):

vllm/model_executor/layers/attention/mla_attention.py

+29

@@ -1984,6 +1984,9 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1984 # Whether this builder can flatten a non-causal query block into decode rows.

1985 supports_non_causal_multi_token_decode: ClassVar[bool] = False

1986

1987 + # Whether can support non-causal multi-token decode with DCP KV cache.

1988 + supports_non_causal_multi_token_dcp: ClassVar[bool] = False

1989 +

1990 # The threshold for reordering the batch into decode and prefill requests.

1991 # If > 1, the batch will be reordered such that requests with

1992 # query length <= threshold are classified as decode requests.

@@ -1991,6 +1994,31 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1994 # when speculative decoding is enabled.

1995 reorder_batch_threshold: int = 1

1996

1997 + def _validate_dspark_dcp_support(self, supports_dcp_with_varlen: bool) -> None:

1998 + speculative_config = getattr(self.vllm_config, "speculative_config", None)

1999 + parallel_config = self.vllm_config.parallel_config

2000 + if (

2001 + speculative_config is None

2002 + or getattr(speculative_config, "method", None) != "dspark"

2003 + or parallel_config.decode_context_parallel_size <= 1

2004 + ):

2005 + return

2006 +

2007 + if self.non_causal_multi_token_decode:

2008 + supported = self.supports_non_causal_multi_token_dcp

2009 + query_mode = "non-causal draft"

2010 + else:

2011 + supported = supports_dcp_with_varlen

2012 + query_mode = "causal multi-token"

2013 +

2014 + if not supported:

2015 + raise ValueError(

2016 + f"{type(self).__name__} does not support {query_mode} MLA "

2017 + "attention for DSpark with decode context parallelism. Select "

2018 + "a backend with explicit DSpark DCP support or set "

2019 + "decode_context_parallel_size=1."

2020 + )

2021 +

2022 @staticmethod

2023 def determine_chunked_prefill_workspace_size(vllm_config: VllmConfig) -> int:

2024 scheduler_config = vllm_config.scheduler_config

@@ -2084,6 +2112,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

2112 self.non_causal_multi_token_decode = getattr(

2113 kv_cache_spec, "non_causal_multi_token_decode", False

2114 )

2115 + self._validate_dspark_dcp_support(supports_dcp_with_varlen)

2116

2117 # A draft cache group can have a different head count from the target.

2118 self.num_heads = get_num_attention_heads_from_layers(

vllm/models/kimi_k3/nvidia/mla.py

-5

@@ -341,11 +341,6 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):

341 "parallelism."

342 )

343 self.dcp_world_size = parallel_config.decode_context_parallel_size

344 - assert self.dcp_world_size <= 1 or self.rotary_emb is None, (

345 - "Kimi-K3 MultiHeadLatentAttention does not support RoPE with decode "

346 - "context parallelism because gathered queries require gathered "

347 - "positions."

348 - )

344 self.dcp_manager: MLADCPManager | None = None

345 if self.dcp_world_size > 1:

346 query_dtype = (

vllm/v1/attention/backend.py

@@ -330,6 +330,11 @@ class AttentionBackend(ABC):

330 except NotImplementedError:

331 return False

332

333 + @classmethod

334 + def supports_non_causal_dcp(cls) -> bool:

335 + builder_cls = cls.get_builder_cls()

336 + return bool(getattr(builder_cls, "supports_non_causal_multi_token_dcp", False))

337 +

338 @classmethod

339 def supports_attn_type(cls, attn_type: str) -> bool:

340 """Check if backend supports a given attention type.

@@ -378,6 +383,7 @@ class AttentionBackend(ABC):

383 use_kv_connector: bool = False,

384 use_pcp: bool = False,

385 use_adaptive_verification: bool = False,

386 + use_dcp: bool = False,

387 ) -> list[str]:

388 invalid_reasons = []

389 if not cls.supports_head_size(head_size):

@@ -414,6 +420,8 @@ class AttentionBackend(ABC):

420 invalid_reasons.append("sliding window not supported")

421 if use_non_causal and not cls.supports_non_causal():

422 invalid_reasons.append("non-causal attention not supported")

423 + if use_mla and use_non_causal and use_dcp and not cls.supports_non_causal_dcp():

424 + invalid_reasons.append("non-causal MLA attention with DCP not supported")

425 if use_batch_invariant and not cls.supports_batch_invariance():

426 invalid_reasons.append("batch invariance not supported")

427 if use_kv_connector and not cls.supports_kv_connector():

vllm/v1/attention/backends/mla/flashinfer_mla.py

+115 -7

@@ -1,7 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -from typing import ClassVar

4 +from dataclasses import dataclass

5 +from typing import TYPE_CHECKING, ClassVar

7 import torch

8 from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla

@@ -15,6 +16,7 @@ from vllm.config.cache import CacheDType

16 from vllm.logger import init_logger

17 from vllm.model_executor.layers.attention.mla_attention import (

18 MLACommonBackend,

19 + MLACommonDecodeMetadata,

20 MLACommonImpl,

21 MLACommonMetadata,

22 MLACommonMetadataBuilder,

@@ -30,6 +32,10 @@ from vllm.v1.attention.backend import (

32 )

33 from vllm.v1.attention.backends.utils import KVCacheLayoutType

35 +if TYPE_CHECKING:

36 + from vllm.config import VllmConfig

37 + from vllm.v1.kv_cache_interface import AttentionSpec

38 +

39 logger = init_logger(__name__)

@@ -112,12 +118,56 @@ def _get_multi_ctas_kv_counter_buffer(

118 return _fi_multi_ctas_kv_counter

119

120

115 -class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):

121 +@dataclass

122 +class FlashInferMLADecodeMetadata(MLACommonDecodeMetadata):

123 + flattened_block_table: torch.Tensor | None = None

124 + flattened_seq_lens: torch.Tensor | None = None

125 + query_len: int = 0

126 +

127 +

128 +@dataclass

129 +class FlashInferMLAMetadata(MLACommonMetadata[FlashInferMLADecodeMetadata]):

130 + pass

131 +

132 +

133 +class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[FlashInferMLAMetadata]):

134 _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH

135 query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM

136 # Non-causal DSpark blocks are flattened to single-token rows in forward_mqa.

137 supports_non_causal_multi_token_decode: ClassVar[bool] = True

138

139 + def __init__(

140 + self,

141 + kv_cache_spec: "AttentionSpec",

142 + layer_names: list[str],

143 + vllm_config: "VllmConfig",

144 + device: torch.device,

145 + ) -> None:

146 + super().__init__(

147 + kv_cache_spec,

148 + layer_names,

149 + vllm_config,

150 + device,

151 + FlashInferMLAMetadata,

152 + supports_dcp_with_varlen=True,

153 + )

154 +

155 + def _build_decode(

156 + self,

157 + block_table_tensor: torch.Tensor,

158 + seq_lens_device: torch.Tensor,

159 + max_seq_len: int,

160 + query_start_loc_cpu: torch.Tensor,

161 + query_start_loc_device: torch.Tensor,

162 + num_decode_tokens: int,

163 + dcp_tot_seq_lens_device: torch.Tensor | None,

164 + ) -> FlashInferMLADecodeMetadata:

165 + return FlashInferMLADecodeMetadata(

166 + block_table=block_table_tensor,

167 + seq_lens=seq_lens_device,

168 + dcp_tot_seq_lens=dcp_tot_seq_lens_device,

169 + )

170 +

171

172 class FlashInferMLABackend(MLACommonBackend):

173 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

@@ -193,7 +243,7 @@ class FlashInferMLABackend(MLACommonBackend):

243 return "HND"

244

245

196 -class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):

246 +class FlashInferMLAImpl(MLACommonImpl[FlashInferMLAMetadata]):

247 can_return_lse_for_decode: bool = True

248 # trtllm-gen MLA decode emits LSE in log2 (per flashinfer's own

249 # reference at flashinfer/trace/templates/attention.py:81:

@@ -264,7 +314,7 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):

314 self,

315 q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

316 kv_c_and_k_pe_cache: torch.Tensor,

267 - attn_metadata: MLACommonMetadata,

317 + attn_metadata: FlashInferMLAMetadata,

318 layer: AttentionLayer,

319 ) -> tuple[torch.Tensor, torch.Tensor | None]:

320 assert kv_c_and_k_pe_cache.numel() > 0

@@ -276,16 +326,28 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):

326

327 block_table = attn_metadata.decode.block_table

328 seq_lens = attn_metadata.decode.seq_lens

329 + query_len = attn_metadata.num_decode_tokens // attn_metadata.num_decodes

330

331 if not attn_metadata.causal:

332 # Non-causal DSpark block: flatten to single-token decode rows with

333 # per-row context seq_lens (trtllm-gen has no causal flag and would

334 # otherwise mask the block causally).

284 - query_len = attn_metadata.num_decode_tokens // attn_metadata.num_decodes

335 q = q.unsqueeze(1)

336 if query_len > 1:

287 - block_table = block_table.repeat_interleave(query_len, dim=0)

288 - seq_lens = seq_lens.repeat_interleave(query_len)

337 + block_table, seq_lens = self._prepare_flattened_decode_metadata(

338 + attn_metadata,

339 + query_len,

340 + causal=False,

341 + )

342 + elif self.dcp_world_size > 1 and query_len > 1:

343 + # Causal DCP block: flatten to single-token decode rows with

344 + # per-row rank-local seq_lens for each query's visible prefix.

345 + block_table, seq_lens = self._prepare_flattened_decode_metadata(

346 + attn_metadata,

347 + query_len,

348 + causal=True,

349 + )

350 + q = q.unsqueeze(1)

351 # trtllm API requires extra dimension q_len_per_request for MTP

352 elif attn_metadata.num_decode_tokens % attn_metadata.num_decodes != 0:

353 logger.warning_once(

@@ -358,3 +420,49 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):

420 o = o.view(-1, o.shape[-2], o.shape[-1])

421

422 return o, lse

423 +

424 + def _prepare_flattened_decode_metadata(

425 + self,

426 + attn_metadata: FlashInferMLAMetadata,

427 + query_len: int,

428 + *,

429 + causal: bool,

430 + ) -> tuple[torch.Tensor, torch.Tensor]:

431 + """Prepare flattened decode metadata once for all layers in the group."""

432 + decode = attn_metadata.decode

433 + assert decode is not None

434 + if decode.query_len:

435 + assert decode.query_len == query_len

436 + assert decode.flattened_block_table is not None

437 + assert decode.flattened_seq_lens is not None

438 + return decode.flattened_block_table, decode.flattened_seq_lens

439 +

440 + block_table = decode.block_table.repeat_interleave(query_len, dim=0)

441 + if causal:

442 + global_seq_lens = decode.dcp_tot_seq_lens

443 + assert global_seq_lens is not None

444 + offsets = torch.arange(

445 + query_len - 1,

446 + -1,

447 + -1,

448 + device=global_seq_lens.device,

449 + dtype=global_seq_lens.dtype,

450 + )

451 + per_query_global_lens = torch.clamp(

452 + (global_seq_lens.unsqueeze(1) - offsets).reshape(-1), min=0

453 + )

454 + interleave = self.cp_kv_cache_interleave_size

455 + dcp_span = self.dcp_world_size * interleave

456 + remainder = torch.clamp(

457 + per_query_global_lens % dcp_span - self.dcp_rank * interleave,

458 + min=0,

459 + max=interleave,

460 + )

461 + seq_lens = per_query_global_lens // dcp_span * interleave + remainder

462 + else:

463 + seq_lens = decode.seq_lens.repeat_interleave(query_len)

464 +

465 + decode.flattened_block_table = block_table

466 + decode.flattened_seq_lens = seq_lens

467 + decode.query_len = query_len

468 + return block_table, seq_lens

vllm/v1/attention/backends/mla/tokenspeed_mla.py

@@ -61,6 +61,7 @@ class TokenspeedMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):

61 # The kernel accepts an explicit causal mask, so a non-causal DSpark

62 # block can remain fused instead of being flattened to single tokens.

63 supports_non_causal_multi_token_decode: ClassVar[bool] = True

64 + supports_non_causal_multi_token_dcp: ClassVar[bool] = True

66 def __init__(

67 self,

vllm/v1/attention/selector.py

+4 -1

@@ -38,6 +38,7 @@ class AttentionSelectorConfig(NamedTuple):

38 use_kv_connector: bool = False

39 use_pcp: bool = False

40 use_adaptive_verification: bool = False

41 + use_dcp: bool = False

43 def __repr__(self):

44 return (

@@ -56,7 +57,8 @@ class AttentionSelectorConfig(NamedTuple):

57 f"use_batch_invariant={self.use_batch_invariant}, "

58 f"use_kv_connector={self.use_kv_connector}, "

59 f"use_adaptive_verification={self.use_adaptive_verification}, "

59 - f"use_pcp={self.use_pcp})"

60 + f"use_pcp={self.use_pcp}, "

61 + f"use_dcp={self.use_dcp})"

62 )

@@ -168,6 +170,7 @@ def get_attn_backend(

170 use_kv_connector=use_kv_connector,

171 use_pcp=vllm_config.parallel_config.prefill_context_parallel_size > 1,

172 use_adaptive_verification=use_adaptive_verification,

173 + use_dcp=vllm_config.parallel_config.decode_context_parallel_size > 1,

174 )

175

176 # A per-KV-group override (keyed by KVCacheSpecKind) takes precedence over

vllm/v1/worker/gpu/cp_utils.py

+21

@@ -59,3 +59,24 @@ def _dcp_local_seq_lens_kernel(

59 # For [num_reqs, max_num_reqs), pad with 0

60 local_seq_lens = tl.where(block < num_reqs, local_seq_lens, 0)

61 tl.store(out_ptr + block, local_seq_lens, mask=block < max_num_reqs)

62 +

63 +

64 +@triton.jit

65 +def cp_local_slot(

66 + positions,

67 + block_numbers,

68 + block_size,

69 + cp_rank,

70 + CP_SIZE: tl.constexpr,

71 + CP_INTERLEAVE: tl.constexpr,

72 + PAD_ID: tl.constexpr,

73 +):

74 + """Return rank-local KV slots, or PAD_ID for positions not owned by this rank."""

75 + block_offsets = positions % (block_size * CP_SIZE)

76 + if CP_SIZE == 1:

77 + return block_numbers * block_size + block_offsets

78 + is_local = block_offsets // CP_INTERLEAVE % CP_SIZE == cp_rank

79 + rounds = block_offsets // (CP_INTERLEAVE * CP_SIZE)

80 + remainder = block_offsets % CP_INTERLEAVE

81 + local_offsets = rounds * CP_INTERLEAVE + remainder

82 + return tl.where(is_local, block_numbers * block_size + local_offsets, PAD_ID)

vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py

+13

@@ -11,6 +11,7 @@ from vllm.v1.worker.gpu.attn_utils import (

11 build_slot_mappings_by_layer,

12 )

13 from vllm.v1.worker.gpu.block_table import BlockTables

14 +from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens

15 from vllm.v1.worker.gpu.cudagraph_utils import (

16 AttentionState,

17 BatchExecutionDescriptor,

@@ -41,6 +42,17 @@ def _prepare_dflash_inputs_to_capture(

42 attn_metadata = None

43 if not skip_attn:

44 query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np)

45 + dcp_local_seq_lens = None

46 + if block_tables.cp_size > 1:

47 + prepare_dcp_local_seq_lens(

48 + input_buffers.dcp_local_seq_lens,

49 + input_buffers.seq_lens,

50 + num_reqs,

51 + block_tables.cp_size,

52 + block_tables.cp_rank,

53 + block_tables.cp_interleave,

54 + )

55 + dcp_local_seq_lens = input_buffers.dcp_local_seq_lens

56 attn_metadata = build_attn_metadata(

57 attn_groups=attn_groups,

58 num_reqs=num_reqs,

@@ -49,6 +61,7 @@ def _prepare_dflash_inputs_to_capture(

61 query_start_loc_cpu=query_start_loc_cpu,

62 max_query_len=num_tokens // num_reqs,

63 seq_lens=input_batch.seq_lens,

64 + dcp_local_seq_lens=dcp_local_seq_lens,

65 max_seq_len=max_model_len,

66 block_tables=input_block_tables,

67 slot_mappings=slot_mappings,

vllm/v1/worker/gpu/spec_decode/dflash/speculator.py

+41 -4

@@ -17,6 +17,7 @@ from vllm.v1.attention.backends.utils import PAD_SLOT_ID

17 from vllm.v1.kv_cache_interface import KVCacheConfig

18 from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer

19 from vllm.v1.worker.gpu.block_table import BlockTables

20 +from vllm.v1.worker.gpu.cp_utils import cp_local_slot, prepare_dcp_local_seq_lens

21 from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp

22 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers

23 from vllm.v1.worker.gpu.model_states.interface import ModelState

@@ -285,10 +286,21 @@ class DFlashSpeculator(DraftModelSpeculator):

286 num_query_per_req: int | None = None,

287 causal: bool | Mapping[int, bool] = False,

288 query_start_loc_np: np.ndarray | None = None,

289 + dcp_local_seq_lens: torch.Tensor | None = None,

290 ) -> dict[str, Any] | None:

291 if not self.draft_attn_layer_names:

292 return None

293 assert num_query_per_req is None # Omitted for DFlash, read from self instead

294 + if dcp_local_seq_lens is None and self.block_tables.cp_size > 1:

295 + prepare_dcp_local_seq_lens(

296 + self.input_buffers.dcp_local_seq_lens,

297 + self.input_buffers.seq_lens,

298 + num_reqs,

299 + self.block_tables.cp_size,

300 + self.block_tables.cp_rank,

301 + self.block_tables.cp_interleave,

302 + )

303 + dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens

304 return super()._build_draft_attn_metadata(

305 num_reqs,

306 num_reqs_padded,

@@ -298,6 +310,7 @@ class DFlashSpeculator(DraftModelSpeculator):

310 num_query_per_req=self.num_query_per_req,

311 causal=causal,

312 query_start_loc_np=query_start_loc_np,

313 + dcp_local_seq_lens=dcp_local_seq_lens,

314 )

315

316 @torch.inference_mode()

@@ -393,6 +406,9 @@ class DFlashSpeculator(DraftModelSpeculator):

406 seeds,

407 self.block_tables.input_block_tables[gid],

408 self.block_tables.kernel_block_sizes[gid],

409 + self.block_tables.cp_rank,

410 + self.block_tables.cp_size,

411 + self.block_tables.cp_interleave,

412 self.parallel_drafting_token_id,

413 self.num_query_per_req,

414 self.num_speculative_steps,

@@ -507,8 +523,11 @@ def _prepare_dflash_inputs_kernel(

523 max_num_reqs,

524 max_num_tokens,

525 max_model_len,

526 + cp_rank,

527 SAMPLE_FROM_ANCHOR: tl.constexpr,

528 PAD_SLOT_ID: tl.constexpr,

529 + CP_SIZE: tl.constexpr,

530 + CP_INTERLEAVE: tl.constexpr,

531 BLOCK_SIZE: tl.constexpr,

532 ):

533 req_idx = tl.program_id(0)

@@ -543,7 +562,7 @@ def _prepare_dflash_inputs_kernel(

562 # --- Context positions / slots ---

563 ctx_pos_idx = ctx_start + tl.where(is_ctx, j, 0)

564 ctx_pos = tl.load(target_positions_ptr + ctx_pos_idx, mask=is_valid_ctx, other=0)

546 - ctx_block_num = ctx_pos // block_size

565 + ctx_block_num = ctx_pos // (block_size * CP_SIZE)

566 ctx_block_num = tl.minimum(ctx_block_num, block_table_stride - 1)

567 ctx_block_id = tl.load(

568 block_table_ptr + req_idx * block_table_stride + ctx_block_num,

@@ -554,9 +573,12 @@ def _prepare_dflash_inputs_kernel(

573 # to it after eviction; rejected suffix rows are invalid context as well.

574 # Neither kind of row may write draft KV into physical block 0.

575 ctx_resident = is_valid_ctx & (ctx_block_id != 0)

576 + local_ctx_slot = cp_local_slot(

577 + ctx_pos, ctx_block_id, block_size, cp_rank, CP_SIZE, CP_INTERLEAVE, PAD_SLOT_ID

578 + )

579 ctx_slot = tl.where(

580 ctx_resident,

559 - ctx_block_id * block_size + (ctx_pos % block_size),

581 + local_ctx_slot,

582 PAD_SLOT_ID,

583 )

584 # Stored over the full [0, num_ctx) span while the loads above are masked to

@@ -573,7 +595,7 @@ def _prepare_dflash_inputs_kernel(

595 is_bonus = is_query & (query_off == 0)

596 input_id = tl.where(is_bonus, bonus_token, parallel_drafting_token_id)

597

576 - q_block_num = query_pos // block_size

598 + q_block_num = query_pos // (block_size * CP_SIZE)

599 q_block_num = tl.minimum(q_block_num, block_table_stride - 1)

600 q_block_id = tl.load(

601 block_table_ptr + req_idx * block_table_stride + q_block_num,

@@ -583,9 +605,18 @@ def _prepare_dflash_inputs_kernel(

605 # A null block is never a writable cache slot. This can occur when a

606 # sliding-window block table contains evicted/global padding entries.

607 q_resident = is_query & (q_block_id != 0)

608 + local_q_slot = cp_local_slot(

609 + query_pos,

610 + q_block_id,

611 + block_size,

612 + cp_rank,

613 + CP_SIZE,

614 + CP_INTERLEAVE,

615 + PAD_SLOT_ID,

616 + )

617 q_slot = tl.where(

618 q_resident,

588 - q_block_id * block_size + (query_pos % block_size),

619 + local_q_slot,

620 PAD_SLOT_ID,

621 )

622

@@ -681,6 +712,9 @@ def prepare_dflash_inputs(

712 # [max_num_reqs, max_num_blocks]

713 block_table: torch.Tensor,

714 block_size: int,

715 + cp_rank: int,

716 + cp_size: int,

717 + cp_interleave: int,

718 parallel_drafting_token_id: int,

719 num_query_per_req: int,

720 num_speculative_steps: int,

@@ -728,7 +762,10 @@ def prepare_dflash_inputs(

762 max_num_reqs,

763 max_num_tokens,

764 max_model_len,

765 + cp_rank,

766 SAMPLE_FROM_ANCHOR=sample_from_anchor,

767 PAD_SLOT_ID=PAD_SLOT_ID,

768 + CP_SIZE=cp_size,

769 + CP_INTERLEAVE=cp_interleave,

770 BLOCK_SIZE=BLOCK_SIZE,

771 )

vllm/v1/worker/gpu/spec_decode/speculator.py

@@ -234,6 +234,7 @@ class DraftModelSpeculator(BaseSpeculator):

234 num_query_per_req: int = 1,

235 causal: bool | Mapping[int, bool] = True,

236 query_start_loc_np: np.ndarray | None = None,

237 + dcp_local_seq_lens: torch.Tensor | None = None,

238 ) -> dict[str, Any] | None:

239 if query_start_loc_np is not None:

240 # Non-uniform query layout (e.g. multi-module MTP's mixed

@@ -278,6 +279,11 @@ class DraftModelSpeculator(BaseSpeculator):

279 query_start_loc_cpu=query_start_loc_cpu,

280 max_query_len=max_query_len,

281 seq_lens=self.input_buffers.seq_lens[:num_reqs_padded],

282 + dcp_local_seq_lens=(

283 + None

284 + if dcp_local_seq_lens is None

285 + else dcp_local_seq_lens[:num_reqs_padded]

286 + ),

287 max_seq_len=self.draft_max_seq_len,

288 block_tables=block_tables,

289 slot_mappings=slot_mappings,