@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 95ed0feaa5

DCP supports hybrid attention (#40996)

Signed-off-by: YanXu <yancey.yx@alibaba-inc.com> Signed-off-by: Jingyi Yang <girasoleyang@gmail.com> Co-authored-by: Jingyi Yang <girasoleyang@gmail.com>

Yan Xu committed Jul 10, 2026 at 12:34 UTC 95ed0feaa5cd7fb16d72c53ce04950aaf07c4698

26 files changed +785 -127

tests/distributed/test_context_parallel.py

@@ -33,6 +33,7 @@ CP_TEST_MODELS = [

33 # [LANGUAGE GENERATION]

34 "deepseek-ai/DeepSeek-V2-Lite-Chat",

35 "Qwen/Qwen2.5-1.5B-Instruct",

36 + "Qwen/Qwen3.5-0.8B", # hybrid attention model

37 ]

39 # GSM8K eval configuration

@@ -46,6 +47,7 @@ MIN_ACCURACY = {

47 "deepseek-ai/DeepSeek-V2-Lite-Chat": 0.64,

48 # .buildkite/lm-eval-harness/configs/Qwen2.5-1.5B-Instruct.yaml

49 "Qwen/Qwen2.5-1.5B-Instruct": 0.52,

50 + "Qwen/Qwen3.5-0.8B": 0.33,

51 }

@@ -151,6 +153,12 @@ else:

153 cp_kv_cache_interleave_size=16, attn_backend="FLASHINFER"

154 ),

155 ],

156 + "Qwen/Qwen3.5-0.8B": [

157 + CPTestSettings.detailed(

158 + cp_kv_cache_interleave_size=16,

159 + attn_backend="FLASH_ATTN",

160 + ),

161 + ],

162 }

163

164

tests/distributed/test_pynccl.py

+47

@@ -16,6 +16,7 @@ from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator

16 from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary

17 from vllm.distributed.parallel_state import (

18 ensure_model_parallel_initialized,

19 + get_tp_group,

20 get_world_group,

21 graph_capture,

22 init_distributed_environment,

@@ -199,6 +200,52 @@ def test_pynccl_all_gather():

200 distributed_run(all_gather_worker_fn, 2)

201

202

203 +@worker_fn_wrapper

204 +def cuda_communicator_all_gather_dim_worker_fn():

205 + with ensure_current_vllm_config():

206 + ensure_model_parallel_initialized(2, 1)

207 +

208 + tp_group = get_tp_group()

209 + comm = tp_group.device_communicator

210 + assert comm is not None

211 +

212 + rank = tp_group.rank_in_group

213 + world_size = tp_group.world_size

214 + device = tp_group.device

215 +

216 + shape = (2, 3, 4)

217 + num_elems = 1

218 + for size in shape:

219 + num_elems *= size

220 +

221 + for dim in (1, -1):

222 + tensor = (

223 + torch.arange(num_elems, dtype=torch.float32, device=device).reshape(shape)

224 + + rank * num_elems

225 + )

226 + expected = torch.cat(

227 + [

228 + torch.arange(num_elems, dtype=torch.float32, device=device).reshape(

229 + shape

230 + )

231 + + r * num_elems

232 + for r in range(world_size)

233 + ],

234 + dim=dim,

235 + )

236 +

237 + result = comm.all_gather(tensor, dim=dim)

238 + torch.accelerator.synchronize()

239 + torch.testing.assert_close(result, expected, rtol=1e-5, atol=1e-8)

240 +

241 +

242 +@pytest.mark.skipif(

243 + torch.accelerator.device_count() < 2, reason="Need at least 2 GPUs to run the test."

244 +)

245 +def test_cuda_communicator_all_gather_dim_not_zero():

246 + distributed_run(cuda_communicator_all_gather_dim_worker_fn, 2)

247 +

248 +

249 @worker_fn_wrapper

250 def all_gatherv_worker_fn():

251 pynccl_comm = PyNcclCommunicator(

tests/models/language/generation/test_hybrid.py

+13 -1

@@ -43,6 +43,11 @@ HYBRID_MODELS = [

43 "tiny-random/qwen3-next-moe",

44 ]

46 +HYBRID_MODELS_REQUIRING_CHUNKED_PREFILL = {

47 + "LiquidAI/LFM2-1.2B",

48 + "tiny-random/qwen3-next-moe",

49 +}

50 +

51 FULL_CUDA_GRAPH_MODELS = [

52 "ai21labs/Jamba-tiny-dev",

53 "pfnet/plamo-2-1b",

@@ -92,8 +97,15 @@ def test_models(

97 example_prompts, max_tokens, num_logprobs

98 )

100 + extra_kwargs = {}

101 + if model in HYBRID_MODELS_REQUIRING_CHUNKED_PREFILL:

102 + extra_kwargs["enable_chunked_prefill"] = True

103 +

104 with vllm_runner(

96 - model, max_num_seqs=MAX_NUM_SEQS, attention_backend=ATTN_BACKEND

105 + model,

106 + max_num_seqs=MAX_NUM_SEQS,

107 + attention_backend=ATTN_BACKEND,

108 + **extra_kwargs,

109 ) as vllm_model:

110 vllm_outputs = vllm_model.generate_greedy_logprobs(

111 example_prompts, max_tokens, num_logprobs

tests/models/multimodal/generation/test_vit_cudagraph.py

@@ -171,6 +171,7 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {

171 "Describe this video in one sentence."

172 ),

173 needs_video_metadata=True,

174 + vllm_runner_kwargs={"enable_chunked_prefill": True},

175 marks=[pytest.mark.core_model],

176 ),

177 "internvl": VitCudagraphTestConfig(

tests/test_config.py

+4 -4

@@ -1121,14 +1121,14 @@ def test_is_chunked_prefill_supported(

1121 (

1122 "Qwen/Qwen3-Next-80B-A3B-Instruct",

1123 "hybrid",

1124 - False,

1125 - "Hybrid models do not support prefix caching since the feature is still experimental.", # noqa: E501

1124 + True,

1125 + "Generative hybrid models support prefix caching.", # noqa: E501

1126 ),

1127 (

1128 "ibm-granite/granite-4.0-h-small",

1129 "hybrid",

1130 - False,

1131 - "Hybrid models do not support prefix caching since the feature is still experimental.", # noqa: E501

1130 + True,

1131 + "Generative hybrid models support prefix caching.", # noqa: E501

1132 ),

1133 (

1134 "state-spaces/mamba-130m-hf",

tests/v1/streaming_input/test_gpu_model_runner_streaming.py

@@ -38,6 +38,7 @@ def mock_model_runner_with_input_batch():

38 vocab_size=32000,

39 block_sizes=[16],

40 kernel_block_sizes=[16],

41 + max_num_blocks_per_req=[64],

42 logitsprocs=None,

43 is_pooling_model=False,

44 )

tests/v1/worker/test_cp_utils.py new

+45

@@ -0,0 +1,45 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

3 +import pytest

4 +import torch

5 +

6 +from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens

7 +from vllm.v1.worker.cp_utils import should_skip_dcp_context_attention

8 +

9 +

10 +def test_skip_gate_only_for_zero_context():

11 + assert should_skip_dcp_context_attention(torch.zeros(3, dtype=torch.int32))

12 + assert not should_skip_dcp_context_attention(

13 + torch.tensor([0, 5, 0], dtype=torch.int32)

14 + )

15 +

16 +

17 +@pytest.mark.parametrize(

18 + "dcp_world_size,interleave_size,context_len",

19 + [(2, 16, 10), (4, 16, 10), (8, 16, 10), (4, 1, 2)],

20 +)

21 +def test_skip_gate_rank_invariant_with_divergent_local_context(

22 + dcp_world_size: int, interleave_size: int, context_len: int

23 +):

24 + """Contexts shorter than a full interleave round land entirely on a

25 + subset of DCP ranks, so the per-rank local context lengths diverge:

26 + some ranks hold zero local context while others hold all of it. Ranks

27 + with zero local context must still take the collective (non-skip) path,

28 + otherwise the query all-gather in _forward_with_dcp deadlocks across

29 + ranks. The skip gate must therefore depend only on the rank-invariant

30 + global context lengths, never on get_dcp_local_seq_lens output.

31 + """

32 + context_kv_lens = torch.tensor([context_len], dtype=torch.int32)

33 + local_maxes = [

34 + int(

35 + get_dcp_local_seq_lens(

36 + context_kv_lens, dcp_world_size, rank, interleave_size

37 + ).max()

38 + )

39 + for rank in range(dcp_world_size)

40 + ]

41 + # Precondition: the local view diverges across ranks.

42 + assert 0 in local_maxes

43 + assert max(local_maxes) > 0

44 + # The batch still has context globally, so no rank may skip.

45 + assert not should_skip_dcp_context_attention(context_kv_lens)

tests/v1/worker/test_gpu_input_batch.py

@@ -238,6 +238,7 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int):

238 vocab_size=1024,

239 block_sizes=[1],

240 kernel_block_sizes=[1],

241 + max_num_blocks_per_req=[1024],

242 )

243 reqs: list[CachedRequestState] = []

244 req_id_reqs = {}

@@ -332,6 +333,7 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis

333 vocab_size=1024,

334 block_sizes=[1],

335 kernel_block_sizes=[1],

336 + max_num_blocks_per_req=[1024],

337 )

338 ref_input_batch: InputBatch = InputBatch(

339 max_num_reqs=batch_size,

@@ -341,6 +343,7 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis

343 vocab_size=1024,

344 block_sizes=[1],

345 kernel_block_sizes=[1],

346 + max_num_blocks_per_req=[1024],

347 )

348

349 reqs: list[CachedRequestState] = []

@@ -409,6 +412,7 @@ def test_pooling_prompt_lens_not_aliased(device: str):

412 vocab_size=VOCAB_SIZE,

413 block_sizes=[16],

414 kernel_block_sizes=[16],

415 + max_num_blocks_per_req=[64],

416 is_pooling_model=True,

417 )

418

@@ -444,6 +448,7 @@ def test_placeholder_spec_token_ids_written_verbatim():

448 vocab_size=VOCAB_SIZE,

449 block_sizes=[16],

450 kernel_block_sizes=[16],

451 + max_num_blocks_per_req=[1],

452 )

453 req = CachedRequestState(

454 req_id="req",

@@ -491,6 +496,7 @@ def test_pooling_metadata_token_id_buffers(

496 vocab_size=VOCAB_SIZE,

497 block_sizes=[16],

498 kernel_block_sizes=[16],

499 + max_num_blocks_per_req=[64],

500 is_pooling_model=True,

501 )

502 req = _construct_pooling_request(0, PoolingParams(**pooling_params))

tests/v1/worker/test_gpu_model_runner.py

@@ -89,6 +89,7 @@ def initialize_kv_cache(runner: GPUModelRunner):

89 kernel_block_sizes=[

90 kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size

91 ],

92 + max_num_blocks_per_req=[NUM_BLOCKS],

93 )

94 runner.initialize_attn_backend(kv_cache_config)

@@ -1397,6 +1398,7 @@ def test_input_batch_with_kernel_block_sizes():

1398 vocab_size=vocab_size,

1399 block_sizes=block_sizes,

1400 kernel_block_sizes=kernel_block_sizes,

1401 + max_num_blocks_per_req=[16, 8],

1402 )

1403

1404 # Verify that block tables were created with kernel block sizes

@@ -1457,6 +1459,7 @@ def test_hybrid_cache_integration(default_vllm_config, dist_init):

1459 vocab_size=runner.model_config.get_vocab_size(),

1460 block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size],

1461 kernel_block_sizes=[16],

1462 + max_num_blocks_per_req=[NUM_BLOCKS],

1463 ) # Use kernel block size

1464

1465 runner.initialize_attn_backend(kv_cache_config)

vllm/config/model.py

+2 -5

@@ -1881,11 +1881,8 @@ class ModelConfig:

1881 else:

1882 # for generative models

1883 if attn_type == "hybrid":

1884 - logger.debug(

1885 - "Hybrid models do not support prefix caching since the feature "

1886 - "is still experimental."

1887 - )

1888 - return False

1884 + logger.debug("Generative hybrid models support prefix caching.")

1885 + return True

1886 elif attn_type == "attention_free":

1887 logger.debug(

1888 "Attention free models do not support prefix caching since the "

vllm/distributed/device_communicators/cuda_communicator.py

+20 -3

@@ -341,13 +341,30 @@ class CudaCommunicator(DeviceCommunicatorBase):

341 def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:

342 # Route uniform dim-0 all-gathers through NVLS symmetric memory when

343 # enabled (mirrors reduce_scatter); otherwise fall back to the

344 - # base-class ring all-gather. Sequence parallelism's gather-before-GEMM

345 - # uses dim=0 with tp-aligned (uniform) shards.

344 + # PyNccl/base-class all-gather. Sequence parallelism's

345 + # gather-before-GEMM uses dim=0 with tp-aligned (uniform) shards.

346 if dim < 0:

347 dim += input_.dim()

348 if dim == 0 and should_nccl_symm_mem_ag_rs():

349 return self._all_gather_symm_mem(input_.contiguous())

350 - return super().all_gather(input_, dim)

350 +

351 + pynccl_comm = self.pynccl_comm

352 + if pynccl_comm is None or pynccl_comm.disabled:

353 + return super().all_gather(input_, dim)

354 +

355 + input_size = input_.size()

356 + output_size = (input_size[0] * self.world_size,) + input_size[1:]

357 + output_tensor = torch.empty(

358 + output_size, dtype=input_.dtype, device=input_.device

359 + )

360 + pynccl_comm.all_gather(output_tensor, input_.contiguous())

361 + output_tensor = output_tensor.reshape((self.world_size,) + input_size)

362 + output_tensor = output_tensor.movedim(0, dim)

363 + return output_tensor.reshape(

364 + input_size[:dim]

365 + + (self.world_size * input_size[dim],)

366 + + input_size[dim + 1 :]

367 + )

368

369 def reduce_scatter(self, input_: torch.Tensor, dim: int = -1):

370 world_size = self.world_size

vllm/engine/arg_utils.py

+5 -1

@@ -2493,7 +2493,11 @@ class EngineArgs:

2493 self, model_config: ModelConfig

2494 ) -> None:

2495 default_chunked_prefill = model_config.is_chunked_prefill_supported

2496 - default_prefix_caching = model_config.is_prefix_caching_supported

2496 + # Hybrid models support prefix caching but keep it opt-in for now

2497 + # while the feature matures.

2498 + default_prefix_caching = (

2499 + model_config.is_prefix_caching_supported and not model_config.is_hybrid

2500 + )

2501

2502 if self.enable_chunked_prefill is None:

2503 self.enable_chunked_prefill = default_chunked_prefill

vllm/model_executor/layers/attention_layer_base.py

@@ -19,6 +19,7 @@ class AttentionLayerBase(ABC):

19 """

21 impl: "AttentionImpl"

22 + supports_dcp: bool = True

24 @abstractmethod

25 def get_attn_backend(self) -> type[AttentionBackend]:

vllm/model_executor/layers/mamba/abstract.py

@@ -22,6 +22,7 @@ class MambaBase(AttentionLayerBase):

22 # Contains the KV cache (mamba state) for the layer

23 # in the shape specified by `self.get_state_shape`.

24 kv_cache: tuple[torch.Tensor, ...]

25 + supports_dcp: bool = False

27 @abstractmethod

28 def get_state_shape(self) -> Iterable[tuple[int, ...]]:

vllm/utils/cpu_triton_utils.py

+9 -5

@@ -29,13 +29,17 @@ def _compute_slot_mapping_kernel_impl(

29 block_table_stride: int, # max_num_blocks_per_req

30 block_size: int,

31 slot_mapping: torch.Tensor, # [max_num_tokens], int64

32 - TOTAL_CP_WORLD_SIZE: int,

33 - TOTAL_CP_RANK: int,

34 - CP_KV_CACHE_INTERLEAVE_SIZE: int,

35 - PAD_ID: int,

36 - BLOCK_SIZE: int,

32 + KV_CACHE_BLOCK_SIZE: int | None = None,

33 + BLOCKS_PER_KV_BLOCK: int = 1,

34 + TOTAL_CP_WORLD_SIZE: int = 1,

35 + TOTAL_CP_RANK: int = 0,

36 + CP_KV_CACHE_INTERLEAVE_SIZE: int = 1,

37 + PAD_ID: int = -1,

38 + BLOCK_SIZE: int = 1024,

39 ) -> None:

40 assert TOTAL_CP_WORLD_SIZE == 1, "Context Parallelism is not supported on CPU."

41 + if BLOCKS_PER_KV_BLOCK != 1:

42 + assert block_size * BLOCKS_PER_KV_BLOCK == KV_CACHE_BLOCK_SIZE

43 torch.ops._C.compute_slot_mapping_kernel_impl(

44 query_start_loc,

45 positions,

vllm/v1/attention/backend.py

@@ -796,6 +796,8 @@ class AttentionImplBase(ABC, Generic[T]):

796

797 # Whether the attention impl supports Prefill Context Parallelism.

798 supports_pcp: bool = False

799 + # Whether the attention impl supports Decode Context Parallelism.

800 + supports_dcp: bool = True

801 # Whether the attention impl(or ops) supports MTP

802 # when cp_kv_cache_interleave_size > 1

803 supports_mtp_with_cp_non_trivial_interleave_size: bool = False

vllm/v1/attention/backends/flash_attn.py

+176 -45

@@ -56,10 +56,14 @@ from vllm.v1.attention.backend import (

56 AttentionMetadataBuilder,

57 CommonAttentionMetadata,

58 )

59 -from vllm.v1.attention.backends.utils import (

60 - get_kv_cache_layout,

61 -)

59 +from vllm.v1.attention.backends.utils import get_kv_cache_layout

60 from vllm.v1.kv_cache_interface import AttentionSpec

61 +from vllm.v1.worker.cp_utils import (

62 + run_split_fa2_dcp_context_attention,

63 + should_skip_dcp_context_attention,

64 + should_split_fa2_dcp_context_attention,

65 + split_dcp_context_queries,

66 +)

68 logger = init_logger(__name__)

@@ -245,6 +249,13 @@ class FlashAttentionMetadata:

249 max_dcp_context_kv_len: int | None = None

250 dcp_context_kv_lens: torch.Tensor | None = None

251

252 + # Split counts for FA2 DCP context attention. num_prefill_* tracks

253 + # context-bearing extend rows; pure prefills do not attend to DCP context.

254 + num_decode_reqs: int = 0

255 + num_prefill_reqs: int = 0

256 + num_decode_tokens: int = 0

257 + num_prefill_tokens: int = 0

258 +

259 # Optional aot scheduling

260 scheduler_metadata: torch.Tensor | None = None

261 prefix_scheduler_metadata: torch.Tensor | None = None

@@ -513,6 +524,10 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad

524 use_cascade = common_prefix_len > 0

525 max_dcp_context_kv_len = 0

526 dcp_context_kv_lens = None

527 + num_decode_reqs = 0

528 + num_prefill_reqs = 0

529 + num_decode_tokens = 0

530 + num_prefill_tokens = 0

531

532 cu_prefix_query_lens = None

533 prefix_kv_lens = None

@@ -532,23 +547,54 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad

547 self._dcp_context_kv_lens[num_reqs:] = 0

548 dcp_context_kv_lens = self._dcp_context_kv_lens[:num_reqs]

549

550 + skip_dcp_context_attention = False

551 + if common_attn_metadata.seq_lens_cpu_upper_bound is not None:

552 + query_lens_cpu = (

553 + common_attn_metadata.query_start_loc_cpu[1 : num_reqs + 1]

554 + - common_attn_metadata.query_start_loc_cpu[:num_reqs]

555 + )

556 + context_kv_lens_cpu = (

557 + common_attn_metadata.seq_lens_cpu_upper_bound[:num_reqs]

558 + - query_lens_cpu

559 + )

560 + skip_dcp_context_attention = should_skip_dcp_context_attention(

561 + context_kv_lens_cpu

562 + )

563 +

564 + if max_query_len > 1:

565 + (

566 + num_decode_reqs,

567 + num_prefill_reqs,

568 + num_decode_tokens,

569 + num_prefill_tokens,

570 + ) = split_dcp_context_queries(

571 + common_attn_metadata.query_start_loc_cpu,

572 + common_attn_metadata.seq_lens_cpu_upper_bound,

573 + max_query_len,

574 + num_actual_tokens,

575 + )

576 +

577 # After DCP distribution, the maximum number of tokens for any rank is

578 # ceil(L / (N * I)) * I, where L is max_seq_len, N is dcp_world_size,

579 # and I is cp_kv_cache_interleave_size.

580 # This eliminates GPU->CPU sync while minimizing workspace over-allocation.

539 - num_partitions = self.dcp_world_size * self.cp_kv_cache_interleave_size

540 - max_dcp_context_kv_len = (

541 - (max_seq_len + num_partitions - 1) // num_partitions

542 - ) * self.cp_kv_cache_interleave_size

543 -

544 - scheduler_metadata = schedule(

545 - batch_size=num_reqs,

546 - cu_query_lens=query_start_loc,

547 - max_query_len=max_query_len,

548 - seqlens=dcp_context_kv_lens,

549 - max_seq_len=max_dcp_context_kv_len,

550 - causal=False,

551 - )

581 + if skip_dcp_context_attention:

582 + max_dcp_context_kv_len = 0

583 + scheduler_metadata = None

584 + else:

585 + num_partitions = self.dcp_world_size * self.cp_kv_cache_interleave_size

586 + max_dcp_context_kv_len = (

587 + (max_seq_len + num_partitions - 1) // num_partitions

588 + ) * self.cp_kv_cache_interleave_size

589 +

590 + scheduler_metadata = schedule(

591 + batch_size=num_reqs,

592 + cu_query_lens=query_start_loc,

593 + max_query_len=max_query_len,

594 + seqlens=dcp_context_kv_lens,

595 + max_seq_len=max_dcp_context_kv_len,

596 + causal=False,

597 + )

598 elif use_cascade:

599 cu_prefix_query_lens = torch.tensor(

600 [0, num_actual_tokens], dtype=torch.int32, device=self.device

@@ -614,6 +660,10 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad

660 slot_mapping=slot_mapping,

661 max_dcp_context_kv_len=max_dcp_context_kv_len,

662 dcp_context_kv_lens=dcp_context_kv_lens,

663 + num_decode_reqs=num_decode_reqs,

664 + num_prefill_reqs=num_prefill_reqs,

665 + num_decode_tokens=num_decode_tokens,

666 + num_prefill_tokens=num_prefill_tokens,

667 use_cascade=use_cascade,

668 common_prefix_len=common_prefix_len,

669 scheduler_metadata=scheduler_metadata,

@@ -743,8 +793,12 @@ class FlashAttentionImpl(AttentionImpl):

793 self.dcp_combine = dcp_a2a_lse_reduce if dcp_a2a else cp_lse_ag_out_rs

794

795 self._dcp_dtype: torch.dtype | None = None

796 + self._dcp_max_num_tokens: int = 0

797 if vllm_config is not None and self.dcp_world_size > 1:

798 self._dcp_dtype = vllm_config.model_config.dtype

799 + self._dcp_max_num_tokens = (

800 + vllm_config.scheduler_config.max_num_batched_tokens

801 + )

802

803 def forward(

804 self,

@@ -1063,40 +1117,120 @@ class FlashAttentionImpl(AttentionImpl):

1117 block_table = attn_metadata.block_table

1118

1119 query = query.contiguous()

1120 + if attn_metadata.max_dcp_context_kv_len == 0:

1121 + flash_attn_varlen_func(

1122 + q=query,

1123 + k=key,

1124 + v=value,

1125 + out=output,

1126 + cu_seqlens_q=cu_seqlens_q,

1127 + max_seqlen_q=max_seqlen_q,

1128 + cu_seqlens_k=cu_seqlens_q,

1129 + max_seqlen_k=max_seqlen_q,

1130 + softmax_scale=self.scale,

1131 + causal=attn_metadata.causal,

1132 + alibi_slopes=self.alibi_slopes,

1133 + window_size=list(self.sliding_window)

1134 + if self.sliding_window is not None

1135 + else None,

1136 + softcap=self.logits_soft_cap,

1137 + return_softmax_lse=True,

1138 + fa_version=self.vllm_flash_attn_version,

1139 + q_descale=q_descale,

1140 + k_descale=k_descale,

1141 + v_descale=v_descale,

1142 + num_splits=attn_metadata.max_num_splits,

1143 + )

1144 + return output

1145 +

1146 query_across_dcp = get_dcp_group().all_gather(query, dim=1)

1147 sliding_window_size = (

1148 list(self.sliding_window) if self.sliding_window is not None else None

1149 )

1150 n = query_across_dcp.shape[0]

1071 - (dcp_context_out,) = current_workspace_manager().get_simultaneous(

1151 + num_reqs = cu_seqlens_q.shape[0] - 1

1152 + num_decodes = attn_metadata.num_decode_reqs

1153 + num_context_prefills = attn_metadata.num_prefill_reqs

1154 + num_decode_tokens = attn_metadata.num_decode_tokens

1155 + num_context_prefill_tokens = attn_metadata.num_prefill_tokens

1156 + split_dcp_context = should_split_fa2_dcp_context_attention(

1157 + self.vllm_flash_attn_version,

1158 + max_seqlen_q,

1159 + num_reqs,

1160 + num_decodes,

1161 + num_context_prefills,

1162 + )

1163 + dcp_context_out_tokens = max(n, self._dcp_max_num_tokens)

1164 + dcp_context_out_spec = (

1165 (

1073 - (n, self.num_heads * self.dcp_world_size, self.head_size),

1074 - self._dcp_dtype,

1166 + dcp_context_out_tokens,

1167 + self.num_heads * self.dcp_world_size,

1168 + self.head_size,

1169 ),

1170 + self._dcp_dtype,

1171 )

1077 - context_attn_out, context_lse = flash_attn_varlen_func(

1078 - q=query_across_dcp,

1079 - k=key_cache,

1080 - v=value_cache,

1081 - out=dcp_context_out,

1082 - cu_seqlens_q=cu_seqlens_q,

1083 - max_seqlen_q=max_seqlen_q,

1084 - seqused_k=attn_metadata.dcp_context_kv_lens,

1085 - max_seqlen_k=attn_metadata.max_dcp_context_kv_len,

1086 - softmax_scale=self.scale,

1087 - causal=False,

1088 - alibi_slopes=self.alibi_slopes,

1089 - window_size=sliding_window_size,

1090 - block_table=block_table,

1091 - softcap=self.logits_soft_cap,

1092 - return_softmax_lse=True,

1093 - scheduler_metadata=attn_metadata.scheduler_metadata,

1094 - fa_version=self.vllm_flash_attn_version,

1095 - q_descale=q_descale,

1096 - k_descale=k_descale,

1097 - v_descale=v_descale,

1098 - num_splits=attn_metadata.max_num_splits,

1172 + (dcp_context_out_workspace,) = current_workspace_manager().get_simultaneous(

1173 + dcp_context_out_spec,

1174 )

1175 + dcp_context_out = dcp_context_out_workspace[:n]

1176 +

1177 + if split_dcp_context:

1178 + # TODO: Remove this DCP + FA2 mixed decode/prefill workaround once

1179 + # FA4 supports this Qwen3.5 shape.

1180 + assert attn_metadata.dcp_context_kv_lens is not None

1181 + assert attn_metadata.max_dcp_context_kv_len is not None

1182 + assert self.vllm_flash_attn_version is not None

1183 + context_attn_out, context_lse = run_split_fa2_dcp_context_attention(

1184 + flash_attn_varlen_func,

1185 + query_across_dcp,

1186 + key_cache,

1187 + value_cache,

1188 + dcp_context_out,

1189 + cu_seqlens_q,

1190 + max_seqlen_q,

1191 + attn_metadata.dcp_context_kv_lens,

1192 + attn_metadata.max_dcp_context_kv_len,

1193 + self.scale,

1194 + self.alibi_slopes,

1195 + sliding_window_size,

1196 + block_table,

1197 + self.logits_soft_cap,

1198 + self.vllm_flash_attn_version,

1199 + q_descale,

1200 + k_descale,

1201 + v_descale,

1202 + attn_metadata.max_num_splits,

1203 + self.num_heads,

1204 + self.dcp_world_size,

1205 + num_decodes,

1206 + num_context_prefills,

1207 + num_decode_tokens,

1208 + num_context_prefill_tokens,

1209 + )

1210 + else:

1211 + context_attn_out, context_lse = flash_attn_varlen_func(

1212 + q=query_across_dcp,

1213 + k=key_cache,

1214 + v=value_cache,

1215 + out=dcp_context_out,

1216 + cu_seqlens_q=cu_seqlens_q,

1217 + max_seqlen_q=max_seqlen_q,

1218 + seqused_k=attn_metadata.dcp_context_kv_lens,

1219 + max_seqlen_k=attn_metadata.max_dcp_context_kv_len,

1220 + softmax_scale=self.scale,

1221 + causal=False,

1222 + alibi_slopes=self.alibi_slopes,

1223 + window_size=sliding_window_size,

1224 + block_table=block_table,

1225 + softcap=self.logits_soft_cap,

1226 + return_softmax_lse=True,

1227 + scheduler_metadata=attn_metadata.scheduler_metadata,

1228 + fa_version=self.vllm_flash_attn_version,

1229 + q_descale=q_descale,

1230 + k_descale=k_descale,

1231 + v_descale=v_descale,

1232 + num_splits=attn_metadata.max_num_splits,

1233 + )

1234 # FA returns LSE in shape [ H, B ] but DCP combine wants [ B, H ]

1235 context_attn_out_cor, context_lse_cor = self.dcp_combine(

1236 context_attn_out,

@@ -1106,14 +1240,11 @@ class FlashAttentionImpl(AttentionImpl):

1240 )

1241 context_lse_cor = context_lse_cor.transpose(0, 1).contiguous()

1242

1109 - (dcp_query_out,) = current_workspace_manager().get_simultaneous(

1110 - ((query.shape[0], self.num_heads, self.head_size), self._dcp_dtype),

1111 - )

1243 query_attn_out, query_lse = flash_attn_varlen_func(

1244 q=query,

1245 k=key,

1246 v=value,

1116 - out=dcp_query_out,

1247 + out=output,

1248 cu_seqlens_q=cu_seqlens_q,

1249 max_seqlen_q=max_seqlen_q,

1250 cu_seqlens_k=cu_seqlens_q,

vllm/v1/core/kv_cache_coordinator.py

+37 -12

@@ -550,12 +550,28 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

550 # different KV cache groups have different block sizes, the actual block size

551 # can be a multiple of hash_block_size.

552 self.hash_block_size = hash_block_size

553 + self.dcp_world_size = dcp_world_size

554 + group_block_sizes = [

555 + manager.block_size for manager in self.single_type_managers

556 + ]

557 assert all(

554 - g.kv_cache_spec.block_size % hash_block_size == 0

555 - for g in kv_cache_config.kv_cache_groups

556 - ), "block_size must be divisible by hash_block_size"

557 - assert dcp_world_size == 1, "DCP not support hybrid attn now."

558 + block_size % hash_block_size == 0 for block_size in group_block_sizes

559 + ), (

560 + "Each KV cache group's real block_size must be divisible by "

561 + f"hash_block_size. block_sizes={group_block_sizes}, "

562 + f"hash_block_size={hash_block_size}"

563 + )

564 assert pcp_world_size == 1, "PCP not support hybrid attn now."

565 + if dcp_world_size > 1:

566 + # DCP shards full-attention KV across ranks and replicates Mamba

567 + # state; other spec types (e.g. sliding window) have no DCP-aware

568 + # handling yet, so reject them explicitly.

569 + for g in kv_cache_config.kv_cache_groups:

570 + assert isinstance(g.kv_cache_spec, (FullAttentionSpec, MambaSpec)), (

571 + "DCP with hybrid KV cache layouts only supports "

572 + "full-attention and Mamba groups, got: "

573 + f"{type(g.kv_cache_spec).__name__}."

574 + )

575 self.verify_and_split_kv_cache_groups()

576

577 def verify_and_split_kv_cache_groups(self) -> None:

@@ -651,11 +667,11 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

667 - The number of tokens of the longest cache hit.

668 """

669

654 - def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList:

655 - if kv_cache_spec.block_size == self.hash_block_size:

670 + def _get_block_hashes(block_size: int) -> BlockHashList:

671 + if block_size == self.hash_block_size:

672 return block_hashes

673 return BlockHashListWithBlockSize(

658 - block_hashes, self.hash_block_size, kv_cache_spec.block_size

674 + block_hashes, self.hash_block_size, block_size

675 )

676

677 num_groups = len(self.kv_cache_config.kv_cache_groups)

@@ -680,13 +696,14 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

696 for idx, (spec, group_ids, manager_cls, use_eagle) in enumerate(

697 self.attention_groups

698 ):

699 + group_block_size = self.single_type_managers[group_ids[0]].block_size

700 cached_blocks = hit_blocks_by_group[group_ids[0]]

701 if isinstance(spec, FullAttentionSpec) and cached_blocks is not None:

702 # Full attention is downward-closed: we only need to look

703 # up cached blocks once; on subsequent iterations just trim

704 # to the (reduced) current hit length.

705 curr_hit_length = (

689 - curr_hit_length // spec.block_size * spec.block_size

706 + curr_hit_length // group_block_size * group_block_size

707 )

708 continue

709

@@ -696,18 +713,23 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

713 if drop_eagle_block:

714 # Eagle needs to match one more block and then pop the last.

715 _max_length = min(

699 - curr_hit_length + spec.block_size, max_cache_hit_length

716 + curr_hit_length + group_block_size, max_cache_hit_length

717 )

718 hit_blocks = manager_cls.find_longest_cache_hit(

702 - block_hashes=_get_block_hashes(spec),

719 + block_hashes=_get_block_hashes(group_block_size),

720 max_length=_max_length,

721 kv_cache_group_ids=group_ids,

722 block_pool=self.block_pool,

723 kv_cache_spec=spec,

724 drop_eagle_block=drop_eagle_block,

725 alignment_tokens=self.scheduler_block_size,

726 + dcp_world_size=(

727 + self.dcp_world_size

728 + if isinstance(spec, FullAttentionSpec)

729 + else 1

730 + ),

731 )

710 - _new_hit_length = len(hit_blocks[0]) * spec.block_size

732 + _new_hit_length = len(hit_blocks[0]) * group_block_size

733 if drop_eagle_block:

734 eagle_verified.add(idx)

735 elif _new_hit_length < curr_hit_length:

@@ -728,7 +750,10 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

750 # Truncate full attention blocks to final hit_length (if present)

751 first_group = self.attention_groups[0]

752 if isinstance(first_group.spec, FullAttentionSpec):

731 - num_blocks = hit_length // first_group.spec.block_size

753 + group_block_size = self.single_type_managers[

754 + first_group.group_ids[0]

755 + ].block_size

756 + num_blocks = hit_length // group_block_size

757 for group_id in first_group.group_ids:

758 if (blks := hit_blocks_by_group[group_id]) is not None:

759 del blks[num_blocks:]

vllm/v1/core/kv_cache_utils.py

+8 -8

@@ -627,7 +627,8 @@ def resolve_kv_cache_block_sizes(

627 - ``scheduler_block_size`` is the token-alignment invariant used by the

628 scheduler (e.g. for ``num_computed_tokens`` rounding). Single group:

629 ``cache_config.block_size * dcp * pcp``. Multiple groups: LCM of every

630 - group's block size — context parallelism is not supported here.

630 + group's effective block size. Attention groups are scaled by DCP/PCP;

631 + Mamba groups keep their full per-rank state and are not scaled.

632 - ``hash_block_size`` is the granularity at which ``Request.block_hashes``

633 is computed. Single group: equals scheduler block size. Multiple groups:

634 ``cache_config.hash_block_size`` override if set, else the GCD of group

@@ -645,13 +646,12 @@ def resolve_kv_cache_block_sizes(

646 bs = cache_config.block_size * dcp * pcp

647 return bs, bs

648

648 - if dcp != 1 or pcp != 1:

649 - raise ValueError(

650 - "Hybrid KV cache groups with multiple block sizes do not "

651 - "support context parallelism (dcp_world_size/pcp_world_size > 1)."

652 - )

653 -

654 - group_block_sizes = [g.kv_cache_spec.block_size for g in groups]

649 + group_block_sizes = [

650 + g.kv_cache_spec.block_size * dcp * pcp

651 + if isinstance(g.kv_cache_spec, AttentionSpec)

652 + else g.kv_cache_spec.block_size

653 + for g in groups

654 + ]

655 scheduler_block_size = math.lcm(*group_block_sizes)

656

657 # Block hashes are only consumed by prefix caching and KV connectors

vllm/v1/core/single_type_kv_cache_manager.py

@@ -1031,6 +1031,10 @@ class MambaManager(SingleTypeKVCacheManager):

1031 self, kv_cache_spec: MambaSpec, block_pool: BlockPool, **kwargs

1032 ) -> None:

1033 super().__init__(kv_cache_spec, block_pool, **kwargs)

1034 + # Mamba layers use TP instead of DCP, so each rank holds the full

1035 + # recurrent state. Undo the DCP/PCP block_size scaling that the base

1036 + # class applies for attention groups whose KV cache is partitioned.

1037 + self.block_size = kv_cache_spec.block_size

1038 self.cached_blocks_this_step: set[BlockHashWithGroupId] = set()

1039 self.mamba_cache_mode = kv_cache_spec.mamba_cache_mode

1040 self.num_speculative_blocks: int = kv_cache_spec.num_speculative_blocks

vllm/v1/kv_cache_interface.py

+34

@@ -128,6 +128,18 @@ class KVCacheSpec:

128 """

129 raise NotImplementedError

130

131 + def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:

132 + """

133 + The number of block table entries needed per request, i.e. the row

134 + length of the worker-side block table for this cache group.

135 +

136 + Args:

137 + vllm_config: The vllm config.

138 + max_len: The maximum sequence length to size for, including the

139 + encoder length for encoder-decoder models.

140 + """

141 + return cdiv(max_len, self.block_size)

142 +

143 def copy_with_new_block_size(self, block_size: int) -> Self:

144 """

145 Create a new KVCacheSpec from self but replacing the block size.

@@ -201,6 +213,16 @@ class AttentionSpec(KVCacheSpec):

213 * get_dtype_size(self.dtype)

214 )

215

216 + def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:

217 + # Attention KV is token-interleaved across DCP/PCP ranks, so each rank

218 + # only stores max_len // (dcp * pcp) tokens per request.

219 + parallel_config = vllm_config.parallel_config

220 + total_cp_size = (

221 + parallel_config.decode_context_parallel_size

222 + * parallel_config.prefill_context_parallel_size

223 + )

224 + return cdiv(max_len, self.block_size * total_cp_size)

225 +

226

227 @dataclass(frozen=True, kw_only=True)

228 class FullAttentionSpec(AttentionSpec):

@@ -699,6 +721,18 @@ class MambaSpec(KVCacheSpec):

721 else:

722 return self.page_size_bytes * (1 + self.num_speculative_blocks)

723

724 + def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:

725 + # Mamba state is replicated across DCP/PCP ranks, never sharded, so

726 + # no CP scaling applies.

727 + if vllm_config.cache_config.mamba_cache_mode == "align":

728 + # Block table rows are position-indexed over the full sequence

729 + # even though only 2 + num_speculative_blocks state blocks are

730 + # resident at a time (earlier states are nulled out by

731 + # remove_skipped_blocks), so the row length must cover max_len

732 + # rather than max_memory_usage_bytes.

733 + return cdiv(max_len, self.block_size) + self.num_speculative_blocks

734 + return cdiv(self.max_memory_usage_bytes(vllm_config), self.page_size_bytes)

735 +

736 def is_uniform_with_collection(

737 self, kv_cache_specs: dict[str, KVCacheSpec]

738 ) -> bool:

vllm/v1/worker/block_table.py

+54 -23

@@ -1,6 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

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

4 +from enum import Enum

5 +

6 import numpy as np

7 import torch

@@ -10,11 +12,15 @@ from vllm.triton_utils import tl, triton

12 from vllm.utils.math_utils import cdiv

13 from vllm.v1.attention.backends.utils import PAD_SLOT_ID

14 from vllm.v1.utils import CpuGpuBuffer

13 -from vllm.v1.worker.cp_utils import get_total_cp_world_size

16 logger = init_logger(__name__)

19 +class SlotMappingMode(Enum):

20 + TOKEN_TO_KV_SLOT = "token_to_kv_slot"

21 + NONE = "none"

22 +

23 +

24 class BlockTable:

25 def __init__(

26 self,

@@ -26,6 +32,7 @@ class BlockTable:

32 device: torch.device,

33 kernel_block_size: int,

34 cp_kv_cache_interleave_size: int,

35 + slot_mapping_mode: SlotMappingMode = SlotMappingMode.TOKEN_TO_KV_SLOT,

36 ):

37 """

38 Args:

@@ -38,11 +45,15 @@ class BlockTable:

45 kernel_block_size: The block_size of underlying attention kernel.

46 Will be the same as `block_size` if `block_size` is supported

47 by the attention kernel.

48 + slot_mapping_mode: How this cache group maps scheduled tokens to

49 + cache slots. Mamba-like state caches do not use token slot

50 + mappings and should use SlotMappingMode.NONE.

51 """

52 self.max_num_reqs = max_num_reqs

53 self.max_num_batched_tokens = max_num_batched_tokens

54 self.pin_memory = pin_memory

55 self.device = device

56 + self.kv_cache_block_size = block_size

58 if kernel_block_size == block_size:

59 # Standard case: allocation and computation use same block size

@@ -98,6 +109,7 @@ class BlockTable:

109 self.dcp_world_size = 1

110 self.dcp_rank = 0

111 self.cp_kv_cache_interleave_size = cp_kv_cache_interleave_size

112 + self.slot_mapping_mode = slot_mapping_mode

113

114 def append_row(

115 self,

@@ -145,6 +157,12 @@ class BlockTable:

157 positions: torch.Tensor,

158 ) -> None:

159 num_tokens = positions.shape[0]

160 + if self.slot_mapping_mode == SlotMappingMode.NONE:

161 + # Mamba/GDN groups consume the block table as recurrent state

162 + # indices and do not use per-token slot mappings.

163 + return

164 + assert self.slot_mapping_mode == SlotMappingMode.TOKEN_TO_KV_SLOT

165 +

166 total_cp_world_size = self.pcp_world_size * self.dcp_world_size

167 total_cp_rank = self.pcp_rank * self.dcp_world_size + self.dcp_rank

168 _compute_slot_mapping_kernel[(num_reqs + 1,)](

@@ -156,6 +174,8 @@ class BlockTable:

174 self.block_table.gpu.stride(0),

175 self.block_size,

176 self.slot_mapping.gpu,

177 + KV_CACHE_BLOCK_SIZE=self.kv_cache_block_size,

178 + BLOCKS_PER_KV_BLOCK=self.blocks_per_kv_block,

179 TOTAL_CP_WORLD_SIZE=total_cp_world_size,

180 TOTAL_CP_RANK=total_cp_rank,

181 CP_KV_CACHE_INTERLEAVE_SIZE=self.cp_kv_cache_interleave_size,

@@ -226,30 +246,27 @@ class MultiGroupBlockTable:

246 def __init__(

247 self,

248 max_num_reqs: int,

229 - max_model_len: int,

249 max_num_batched_tokens: int,

250 pin_memory: bool,

251 device: torch.device,

252 block_sizes: list[int],

253 kernel_block_sizes: list[int],

235 - max_num_blocks: list[int] | None = None,

254 + max_num_blocks: list[int],

255 cp_kv_cache_interleave_size: int = 1,

256 + slot_mapping_modes: list[SlotMappingMode] | None = None,

257 ) -> None:

258 if len(kernel_block_sizes) != len(block_sizes):

259 raise ValueError(

260 f"kernel_block_sizes length ({len(kernel_block_sizes)}) "

261 f"must match block_sizes length ({len(block_sizes)})"

262 )

243 - if max_num_blocks is None:

244 - # Note(hc): each dcp rank only store

245 - # (max_model_len//dcp_world_size) tokens in kvcache,

246 - # so the block_size which used for calc max_num_blocks_per_req

247 - # must be multiplied by dcp_world_size.

248 - total_cp_world_size = get_total_cp_world_size()

249 - max_num_blocks = [

250 - cdiv(max_model_len, block_size * total_cp_world_size)

251 - for block_size in block_sizes

252 - ]

263 + if slot_mapping_modes is None:

264 + slot_mapping_modes = [SlotMappingMode.TOKEN_TO_KV_SLOT] * len(block_sizes)

265 + if len(slot_mapping_modes) != len(block_sizes):

266 + raise ValueError(

267 + f"slot_mapping_modes length ({len(slot_mapping_modes)}) "

268 + f"must match block_sizes length ({len(block_sizes)})"

269 + )

270

271 if len(max_num_blocks) != len(block_sizes):

272 raise ValueError(

@@ -274,9 +291,15 @@ class MultiGroupBlockTable:

291 device,

292 kernel_block_size,

293 cp_kv_cache_interleave_size,

294 + slot_mapping_mode=slot_mapping_mode,

295 )

278 - for block_size, kernel_block_size, max_num_blocks_per_req in zip(

279 - block_sizes, kernel_block_sizes, max_num_blocks

296 + for (

297 + block_size,

298 + kernel_block_size,

299 + max_num_blocks_per_req,

300 + slot_mapping_mode,

301 + ) in zip(

302 + block_sizes, kernel_block_sizes, max_num_blocks, slot_mapping_modes

303 )

304 ]

305

@@ -332,6 +355,8 @@ def _compute_slot_mapping_kernel(

355 block_table_stride, # max_num_blocks_per_req

356 block_size,

357 slot_mapping_ptr, # [max_num_tokens], int64

358 + KV_CACHE_BLOCK_SIZE: tl.constexpr,

359 + BLOCKS_PER_KV_BLOCK: tl.constexpr,

360 TOTAL_CP_WORLD_SIZE: tl.constexpr,

361 TOTAL_CP_RANK: tl.constexpr,

362 CP_KV_CACHE_INTERLEAVE_SIZE: tl.constexpr,

@@ -354,18 +379,14 @@ def _compute_slot_mapping_kernel(

379 start_idx = tl.load(query_start_loc_ptr + req_idx).to(tl.int64)

380 end_idx = tl.load(query_start_loc_ptr + req_idx + 1).to(tl.int64)

381

357 - virtual_block_size = block_size * TOTAL_CP_WORLD_SIZE

382 + virtual_block_size = KV_CACHE_BLOCK_SIZE * TOTAL_CP_WORLD_SIZE

383 row_offset = req_idx * block_table_stride

384 for i in range(start_idx, end_idx, BLOCK_SIZE):

385 offsets = i + tl.arange(0, BLOCK_SIZE)

386 mask = offsets < end_idx

387 pos = tl.load(positions_ptr + offsets, mask=mask, other=0)

363 - block_indices = pos // virtual_block_size

364 - block_numbers = tl.load(block_table_ptr + row_offset + block_indices).to(

365 - tl.int64

366 - )

367 -

368 - virtual_block_offsets = pos - block_indices * virtual_block_size

388 + virtual_block_indices = pos // virtual_block_size

389 + virtual_block_offsets = pos - virtual_block_indices * virtual_block_size

390 is_local = (

391 virtual_block_offsets // CP_KV_CACHE_INTERLEAVE_SIZE

392 ) % TOTAL_CP_WORLD_SIZE == TOTAL_CP_RANK

@@ -375,6 +396,16 @@ def _compute_slot_mapping_kernel(

396 virtual_block_offsets % CP_KV_CACHE_INTERLEAVE_SIZE

397 )

398

378 - slot_ids = block_numbers * block_size + local_block_offsets

399 + block_indices = (

400 + virtual_block_indices * BLOCKS_PER_KV_BLOCK

401 + + local_block_offsets // block_size

402 + )

403 + block_numbers = tl.load(

404 + block_table_ptr + row_offset + block_indices,

405 + mask=mask & is_local,

406 + other=0,

407 + ).to(tl.int64)

408 + slot_offsets = local_block_offsets % block_size

409 + slot_ids = block_numbers * block_size + slot_offsets

410 slot_ids = tl.where(is_local, slot_ids, PAD_ID)

411 tl.store(slot_mapping_ptr + offsets, slot_ids, mask=mask)

vllm/v1/worker/cp_utils.py

+237

@@ -1,15 +1,23 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 +from types import SimpleNamespace

4 from typing import TYPE_CHECKING, Any, cast

6 +import torch

7 +

8 from vllm.config import VllmConfig, get_layers_from_vllm_config

9 from vllm.distributed import get_dcp_group, get_pcp_group

10 +from vllm.logger import init_logger

11 +from vllm.v1.attention.backend import CommonAttentionMetadata

12 +from vllm.v1.attention.backends.utils import split_decodes_prefills_and_extends

14 if TYPE_CHECKING:

15 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

16 else:

17 AttentionLayerBase = object

19 +logger = init_logger(__name__)

20 +

22 def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None:

23 pcp_size = vllm_config.parallel_config.prefill_context_parallel_size

@@ -56,3 +64,232 @@ def get_total_cp_world_size():

64 # DCP might not be initialized in testing

65 dcp_world_size = 1

66 return dcp_world_size * pcp_world_size

67 +

68 +

69 +def get_dcp_dummy_context_len(

70 + dcp_world_size: int,

71 + cp_kv_cache_interleave_size: int,

72 + has_kv_cache_config: bool,

73 + create_mixed_batch: bool,

74 + is_graph_capturing: bool,

75 + uniform_decode: bool,

76 +) -> int:

77 + if (

78 + dcp_world_size <= 1

79 + or not has_kv_cache_config

80 + or not (create_mixed_batch or (is_graph_capturing and uniform_decode))

81 + ):

82 + return 0

83 + return dcp_world_size * cp_kv_cache_interleave_size

84 +

85 +

86 +def prepare_dcp_dummy_context_metadata(

87 + *,

88 + input_batch: Any,

89 + kv_cache_config: Any,

90 + query_pos: Any,

91 + positions: torch.Tensor,

92 + query_start_loc: Any,

93 + num_reqs: int,

94 + num_tokens_unpadded: int,

95 + dcp_dummy_context_len: int,

96 +) -> None:

97 + """Populate valid fake KV metadata for DCP CUDA graph warmup/capture."""

98 + if dcp_dummy_context_len == 0:

99 + return

100 +

101 + # DCP graph warmup may exercise context attention, so block-table entries

102 + # must point at allocated KV blocks.

103 + assert kv_cache_config is not None

104 + max_valid_block_id = kv_cache_config.num_blocks - 1

105 + assert max_valid_block_id > 0

106 + for blk_table in input_batch.block_table.block_tables:

107 + max_row_blocks = (

108 + blk_table.max_num_blocks_per_req // blk_table.blocks_per_kv_block

109 + )

110 + block_ids = [

111 + (block_idx % max_valid_block_id) + 1 for block_idx in range(max_row_blocks)

112 + ]

113 + for req_idx in range(num_reqs):

114 + blk_table.add_row(block_ids, req_idx)

115 + blk_table.commit_block_table(num_reqs)

116 +

117 + query_pos.copy_to_gpu(num_tokens_unpadded)

118 + positions[:num_tokens_unpadded] = (

119 + query_pos.gpu[:num_tokens_unpadded] + dcp_dummy_context_len

120 + )

121 + input_batch.block_table.compute_slot_mapping(

122 + num_reqs,

123 + query_start_loc.gpu[: num_reqs + 1],

124 + positions[:num_tokens_unpadded],

125 + )

126 +

127 +

128 +def should_skip_dcp_context_attention(context_kv_lens_cpu: torch.Tensor) -> bool:

129 + """Whether DCP context attention can be skipped for this batch.

130 +

131 + Must be computed from rank-invariant inputs only (the global context

132 + lengths, NOT this rank's local share from get_dcp_local_seq_lens): the

133 + non-skip path in _forward_with_dcp issues DCP collectives (query

134 + all-gather + LSE combine), so every DCP rank must take the same branch.

135 + A rank can hold zero local context tokens while other ranks still hold

136 + context for the same batch.

137 + """

138 + return int(context_kv_lens_cpu.max().item()) == 0

139 +

140 +

141 +def split_dcp_context_queries(

142 + query_start_loc: torch.Tensor,

143 + seq_lens_cpu_upper_bound: torch.Tensor | None,

144 + max_query_len: int,

145 + num_actual_tokens: int,

146 +) -> tuple[int, int, int, int]:

147 + """Split reordered DCP context queries into decode and extend regions."""

148 + num_reqs = query_start_loc.shape[0] - 1

149 + if max_query_len <= 1:

150 + return num_reqs, 0, num_actual_tokens, 0

151 + if seq_lens_cpu_upper_bound is None:

152 + return 0, num_reqs, 0, num_actual_tokens

153 +

154 + common_attn_metadata = cast(

155 + CommonAttentionMetadata,

156 + SimpleNamespace(

157 + max_query_len=max_query_len,

158 + num_reqs=num_reqs,

159 + num_actual_tokens=num_actual_tokens,

160 + query_start_loc_cpu=query_start_loc,

161 + seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,

162 + is_prefilling=None,

163 + ),

164 + )

165 + (

166 + num_decodes,

167 + num_extends,

168 + _num_prefills,

169 + num_decode_tokens,

170 + num_extend_tokens,

171 + _num_prefill_tokens,

172 + ) = split_decodes_prefills_and_extends(common_attn_metadata)

173 + return num_decodes, num_extends, num_decode_tokens, num_extend_tokens

174 +

175 +

176 +def should_split_fa2_dcp_context_attention(

177 + fa_version: int | None,

178 + max_query_len: int,

179 + num_reqs: int,

180 + num_decode_reqs: int,

181 + num_context_prefill_reqs: int,

182 +) -> bool:

183 + num_prefills = num_reqs - num_decode_reqs

184 + # TODO: Remove this FA2-only DCP compatibility path once FA4 supports

185 + # the Qwen3.5 head_size=256 shape on Blackwell and can be used here.

186 + # FA2 paged-varlen context attention can fail for DCP mixed batches when

187 + # decode rows, context-bearing extend rows, and zero-context pure prefill

188 + # rows are submitted together.

189 + return (

190 + fa_version == 2

191 + and max_query_len > 1

192 + and num_prefills > 0

193 + and (num_decode_reqs > 0 or num_context_prefill_reqs < num_prefills)

194 + )

195 +

196 +

197 +def run_split_fa2_dcp_context_attention(

198 + flash_attn_varlen_func: Any,

199 + query_across_dcp: torch.Tensor,

200 + key_cache: torch.Tensor,

201 + value_cache: torch.Tensor,

202 + dcp_context_out: torch.Tensor,

203 + cu_seqlens_q: torch.Tensor,

204 + max_seqlen_q: int,

205 + dcp_context_kv_lens: torch.Tensor,

206 + max_dcp_context_kv_len: int,

207 + softmax_scale: float,

208 + alibi_slopes: torch.Tensor | None,

209 + sliding_window_size: list[int] | None,

210 + block_table: torch.Tensor,

211 + softcap: float,

212 + fa_version: int,

213 + q_descale: torch.Tensor | None,

214 + k_descale: torch.Tensor | None,

215 + v_descale: torch.Tensor | None,

216 + max_num_splits: int,

217 + num_heads: int,

218 + dcp_world_size: int,

219 + num_decode_reqs: int,

220 + num_context_prefill_reqs: int,

221 + num_decode_tokens: int,

222 + num_context_prefill_tokens: int,

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

224 + dcp_context_out.zero_()

225 + context_lse = torch.full(

226 + (num_heads * dcp_world_size, query_across_dcp.shape[0]),

227 + -torch.inf,

228 + dtype=torch.float32,

229 + device=query_across_dcp.device,

230 + )

231 +

232 + if num_decode_tokens > 0:

233 + _, decode_context_lse = flash_attn_varlen_func(

234 + q=query_across_dcp[:num_decode_tokens],

235 + k=key_cache,

236 + v=value_cache,

237 + out=dcp_context_out[:num_decode_tokens],

238 + cu_seqlens_q=cu_seqlens_q[: num_decode_reqs + 1],

239 + max_seqlen_q=1,

240 + seqused_k=dcp_context_kv_lens[:num_decode_reqs],

241 + max_seqlen_k=max_dcp_context_kv_len,

242 + softmax_scale=softmax_scale,

243 + causal=False,

244 + alibi_slopes=alibi_slopes,

245 + window_size=sliding_window_size,

246 + block_table=block_table[:num_decode_reqs],

247 + softcap=softcap,

248 + return_softmax_lse=True,

249 + scheduler_metadata=None,

250 + fa_version=fa_version,

251 + q_descale=q_descale[:num_decode_reqs] if q_descale is not None else None,

252 + k_descale=k_descale[:num_decode_reqs] if k_descale is not None else None,

253 + v_descale=v_descale[:num_decode_reqs] if v_descale is not None else None,

254 + num_splits=max_num_splits,

255 + )

256 + context_lse[:, :num_decode_tokens] = decode_context_lse

257 +

258 + if num_context_prefill_tokens > 0:

259 + prefill_start = num_decode_tokens

260 + prefill_end = prefill_start + num_context_prefill_tokens

261 + prefill_query_start_loc = (

262 + cu_seqlens_q[

263 + num_decode_reqs : num_decode_reqs + num_context_prefill_reqs + 1

264 + ]

265 + - num_decode_tokens

266 + )

267 + prefill_req_slice = slice(

268 + num_decode_reqs, num_decode_reqs + num_context_prefill_reqs

269 + )

270 + _, prefill_context_lse = flash_attn_varlen_func(

271 + q=query_across_dcp[prefill_start:prefill_end],

272 + k=key_cache,

273 + v=value_cache,

274 + out=dcp_context_out[prefill_start:prefill_end],

275 + cu_seqlens_q=prefill_query_start_loc,

276 + max_seqlen_q=max_seqlen_q,

277 + seqused_k=dcp_context_kv_lens[prefill_req_slice],

278 + max_seqlen_k=max_dcp_context_kv_len,

279 + softmax_scale=softmax_scale,

280 + causal=False,

281 + alibi_slopes=alibi_slopes,

282 + window_size=sliding_window_size,

283 + block_table=block_table[prefill_req_slice],

284 + softcap=softcap,

285 + return_softmax_lse=True,

286 + scheduler_metadata=None,

287 + fa_version=fa_version,

288 + q_descale=q_descale[prefill_req_slice] if q_descale is not None else None,

289 + k_descale=k_descale[prefill_req_slice] if k_descale is not None else None,

290 + v_descale=v_descale[prefill_req_slice] if v_descale is not None else None,

291 + num_splits=max_num_splits,

292 + )

293 + context_lse[:, prefill_start:prefill_end] = prefill_context_lse

294 +

295 + return dcp_context_out, context_lse

vllm/v1/worker/gpu_input_batch.py

+4 -3

@@ -28,7 +28,7 @@ from vllm.v1.sample.thinking_budget_state import (

28 maybe_create_thinking_budget_state_holder,

29 )

30 from vllm.v1.utils import copy_slice

31 -from vllm.v1.worker.block_table import MultiGroupBlockTable

31 +from vllm.v1.worker.block_table import MultiGroupBlockTable, SlotMappingMode

34 @dataclass

@@ -99,13 +99,14 @@ class InputBatch:

99 vocab_size: int,

100 block_sizes: list[int], # The block_size of each kv cache group

101 kernel_block_sizes: list[int],

102 - max_num_blocks_per_req: list[int] | None = None,

102 + max_num_blocks_per_req: list[int],

103 logitsprocs: LogitsProcessors | None = None,

104 logitsprocs_need_output_token_ids: bool = False,

105 num_spec_tokens: int = 0,

106 is_pooling_model: bool = False,

107 cp_kv_cache_interleave_size: int = 1,

108 reasoning_config: ReasoningConfig | None = None,

109 + slot_mapping_modes: list[SlotMappingMode] | None = None,

110 ):

111 self.thinking_budget_state_holder = maybe_create_thinking_budget_state_holder(

112 reasoning_config,

@@ -171,7 +172,6 @@ class InputBatch:

172 # Block table.

173 self.block_table = MultiGroupBlockTable(

174 max_num_reqs=max_num_reqs,

174 - max_model_len=max_model_len,

175 max_num_batched_tokens=max_num_batched_tokens,

176 pin_memory=PIN_MEMORY,

177 device=device,

@@ -179,6 +179,7 @@ class InputBatch:

179 kernel_block_sizes=kernel_block_sizes,

180 max_num_blocks=max_num_blocks_per_req,

181 cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,

182 + slot_mapping_modes=slot_mapping_modes,

183 )

184

185 # Sampling-related.

vllm/v1/worker/gpu_model_runner.py

+61 -16

@@ -153,10 +153,12 @@ from vllm.v1.kv_cache_interface import (

153 KVCacheConfig,

154 KVCacheGroupSpec,

155 KVCacheSpec,

156 + KVCacheSpecKind,

157 KVQuantMode,

158 MambaSpec,

159 SlidingWindowSpec,

160 UniformTypeKVCacheSpecs,

161 + get_kv_cache_spec_kind,

162 )

163 from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry

164 from vllm.v1.outputs import (

@@ -201,9 +203,11 @@ from vllm.v1.spec_decode.utils import update_num_computed_tokens_for_batch_chang

203 from vllm.v1.structured_output.utils import apply_grammar_bitmask

204 from vllm.v1.utils import CpuGpuBuffer, record_function_or_nullcontext

205 from vllm.v1.worker import mamba_utils

206 +from vllm.v1.worker.block_table import SlotMappingMode

207 from vllm.v1.worker.cp_utils import (

208 check_attention_cp_compatibility,

206 - get_total_cp_world_size,

209 + get_dcp_dummy_context_len,

210 + prepare_dcp_dummy_context_metadata,

211 )

212 from vllm.v1.worker.dp_utils import coordinate_batch_across_dp

213 from vllm.v1.worker.ec_connector_model_runner_mixin import ECConnectorModelRunnerMixin

@@ -681,8 +685,13 @@ class GPUModelRunner(

685 placeholder_block_size = (

686 self.cache_config.block_size or CacheConfig.DEFAULT_BLOCK_SIZE

687 )

688 + placeholder_max_num_blocks = cdiv(

689 + max(self.max_model_len, self.max_encoder_len), placeholder_block_size

690 + )

691 self._init_block_sizes = [placeholder_block_size]

692 self._init_kernel_block_sizes = [placeholder_block_size]

693 + self._init_max_num_blocks = [placeholder_max_num_blocks]

694 + self._init_slot_mapping_modes = [SlotMappingMode.TOKEN_TO_KV_SLOT]

695 self.input_batch = InputBatch(

696 max_num_reqs=self.max_num_reqs,

697 # We need to use the encoder length for encoder-decoder

@@ -693,6 +702,7 @@ class GPUModelRunner(

702 vocab_size=self.model_config.get_vocab_size(),

703 block_sizes=[placeholder_block_size],

704 kernel_block_sizes=[placeholder_block_size],

705 + max_num_blocks_per_req=[placeholder_max_num_blocks],

706 num_spec_tokens=self.num_spec_tokens,

707 logitsprocs=build_logitsprocs(

708 self.vllm_config,

@@ -5858,6 +5868,14 @@ class GPUModelRunner(

5868 num_reqs_padded = (

5869 batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs

5870 )

5871 + dcp_dummy_context_len = get_dcp_dummy_context_len(

5872 + self.dcp_world_size,

5873 + self.parallel_config.cp_kv_cache_interleave_size,

5874 + hasattr(self, "kv_cache_config"),

5875 + create_mixed_batch,

5876 + is_graph_capturing,

5877 + uniform_decode,

5878 + )

5879 ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices(

5880 should_ubatch,

5881 num_scheduled_tokens,

@@ -5900,10 +5918,19 @@ class GPUModelRunner(

5918 # In the mixed batch mode (used for FI warmup), we use

5919 # shorter sequence lengths to run faster.

5920 # TODO(luka) better system for describing dummy batches

5903 - seq_lens = torch.tensor( # type: ignore[assignment]

5904 - [1] * num_decode_tokens + [num_prefill_tokens + 1],

5905 - dtype=torch.int,

5906 - )

5921 + if dcp_dummy_context_len > 0:

5922 + seq_lens = torch.tensor( # type: ignore[assignment]

5923 + [1 + dcp_dummy_context_len] * num_decode_tokens

5924 + + [num_prefill_tokens + dcp_dummy_context_len],

5925 + dtype=torch.int,

5926 + )

5927 + else:

5928 + seq_lens = torch.tensor( # type: ignore[assignment]

5929 + [1] * num_decode_tokens + [num_prefill_tokens + 1],

5930 + dtype=torch.int,

5931 + )

5932 + elif dcp_dummy_context_len > 0:

5933 + seq_lens = max_query_len + dcp_dummy_context_len # type: ignore[assignment]

5934 else:

5935 seq_lens = max_query_len # type: ignore[assignment]

5936 self.optimistic_seq_lens_cpu[:num_reqs] = seq_lens

@@ -5919,6 +5946,17 @@ class GPUModelRunner(

5946 )

5947 self.query_start_loc.copy_to_gpu()

5948

5949 + prepare_dcp_dummy_context_metadata(

5950 + input_batch=self.input_batch,

5951 + kv_cache_config=getattr(self, "kv_cache_config", None),

5952 + query_pos=self.query_pos,

5953 + positions=self.positions,

5954 + query_start_loc=self.query_start_loc,

5955 + num_reqs=num_reqs,

5956 + num_tokens_unpadded=num_tokens_unpadded,

5957 + dcp_dummy_context_len=dcp_dummy_context_len,

5958 + )

5959 +

5960 # Sync block table CPU->GPU so cleared rows from

5961 # remove_request() are visible to the attention metadata

5962 # builder. Without this, stale block IDs from finished

@@ -7029,29 +7067,34 @@ class GPUModelRunner(

7067 """

7068 block_sizes = []

7069 max_num_blocks = []

7070 + slot_mapping_modes = []

7071 max_model_len = max(self.max_model_len, self.max_encoder_len)

7072 for kv_cache_group in kv_cache_config.kv_cache_groups:

7034 - if isinstance(kv_cache_group.kv_cache_spec, EncoderOnlyAttentionSpec):

7073 + kv_cache_spec = kv_cache_group.kv_cache_spec

7074 + kv_cache_spec_kind = get_kv_cache_spec_kind(kv_cache_spec)

7075 + if kv_cache_spec_kind == KVCacheSpecKind.ENCODER_ONLY_ATTENTION:

7076 continue

7036 - block_size = kv_cache_group.kv_cache_spec.block_size

7077 + block_size = kv_cache_spec.block_size

7078 block_sizes.append(block_size)

7038 - max_num_blocks_per_req = cdiv(

7039 - max_model_len, block_size * get_total_cp_world_size()

7040 - )

7041 - if isinstance(kv_cache_group.kv_cache_spec, MambaSpec):

7042 - max_num_blocks_per_req = (

7043 - max_num_blocks_per_req

7044 - if self.cache_config.enable_prefix_caching

7045 - else 1

7046 - ) + kv_cache_group.kv_cache_spec.num_speculative_blocks

7079 + if kv_cache_spec_kind == KVCacheSpecKind.MAMBA:

7080 + slot_mapping_modes.append(SlotMappingMode.NONE)

7081 + else:

7082 + slot_mapping_modes.append(SlotMappingMode.TOKEN_TO_KV_SLOT)

7083 + max_num_blocks_per_req = kv_cache_spec.max_num_blocks_per_req(

7084 + self.vllm_config, max_model_len

7085 + )

7086 max_num_blocks.append(max_num_blocks_per_req)

7087

7088 if (

7089 block_sizes != self._init_block_sizes

7090 or kernel_block_sizes != self._init_kernel_block_sizes

7091 + or max_num_blocks != self._init_max_num_blocks

7092 + or slot_mapping_modes != self._init_slot_mapping_modes

7093 ):

7094 self._init_block_sizes = block_sizes

7095 self._init_kernel_block_sizes = kernel_block_sizes

7096 + self._init_max_num_blocks = max_num_blocks

7097 + self._init_slot_mapping_modes = slot_mapping_modes

7098 self.input_batch = InputBatch(

7099 max_num_reqs=self.max_num_reqs,

7100 max_model_len=max_model_len,

@@ -7065,7 +7108,9 @@ class GPUModelRunner(

7108 logitsprocs=self.input_batch.logitsprocs,

7109 logitsprocs_need_output_token_ids=self.input_batch.logitsprocs_need_output_token_ids,

7110 is_pooling_model=self.is_pooling_model,

7111 + cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size,

7112 reasoning_config=self.vllm_config.reasoning_config,

7113 + slot_mapping_modes=slot_mapping_modes,

7114 )

7115

7116 assert self._init_block_sizes == block_sizes, (

vllm/v1/worker/tpu_input_batch.py

+2 -1

@@ -29,6 +29,7 @@ class InputBatch:

29 vocab_size: int,

30 block_sizes: list[int], # The block_size of each kv cache group

31 kernel_block_sizes: list[int],

32 + max_num_blocks_per_req: list[int],

33 ):

34 self.max_num_reqs = max_num_reqs

35 self.max_model_len = max_model_len

@@ -64,12 +65,12 @@ class InputBatch:

65 # Block table.

66 self.block_table = MultiGroupBlockTable(

67 max_num_reqs=max_num_reqs,

67 - max_model_len=max_model_len,

68 max_num_batched_tokens=max_num_batched_tokens,

69 pin_memory=pin_memory,

70 device=device,

71 block_sizes=block_sizes,

72 kernel_block_sizes=kernel_block_sizes,

73 + max_num_blocks=max_num_blocks_per_req,

74 )

76 # Sampling-related.