@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ac201a0eaf
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[Feature] Support Decode Context Parallel (DCP) for MLA (#23734)

Signed-off-by: hongchao <hongchao@msh.team> Signed-off-by: youkaichao <youkaichao@gmail.com> Co-authored-by: hongchao <hongchao@msh.team> Co-authored-by: youkaichao <youkaichao@gmail.com>

yzds committed Sep 6, 2025 at 13:24 UTC ac201a0eaf2af032779db3e7ac96f857cbaa5b7a
27 files changed +1000 -231
.buildkite/test-pipeline.yaml
+2 -1
@@ -837,7 +837,7 @@ steps:
837 - pytest -v -s models/test_oot_registration.py # it needs a clean process
838 - pytest -v -s plugins/lora_resolvers # unit tests for in-tree lora resolver plugins
839
840 -- label: Pipeline Parallelism Test # 45min
840 +- label: Pipeline + Context Parallelism Test # 45min
841 timeout_in_minutes: 60
842 mirror_hardwares: [amdexperimental]
843 working_dir: "/vllm-workspace/tests"
@@ -851,6 +851,7 @@ steps:
851 commands:
852 - pytest -v -s distributed/test_pp_cudagraph.py
853 - pytest -v -s distributed/test_pipeline_parallel.py
854 + # - pytest -v -s distributed/test_context_parallel.py # TODO: enable it on Hopper runners or add triton MLA support
855
856 - label: LoRA TP Test (Distributed) # 17 min
857 timeout_in_minutes: 30
csrc/cache.h
-7
@@ -36,13 +36,6 @@ void concat_and_cache_mla(torch::Tensor& kv_c, torch::Tensor& k_pe,
36 const std::string& kv_cache_dtype,
37 torch::Tensor& scale);
38
39 -void cp_fused_concat_and_cache_mla(torch::Tensor& kv_c, torch::Tensor& k_pe,
40 - torch::Tensor& cp_local_token_select_indices,
41 - torch::Tensor& kv_cache,
42 - torch::Tensor& slot_mapping,
43 - const std::string& kv_cache_dtype,
44 - torch::Tensor& scale);
45 -
39 // Just for unittest
40 void convert_fp8(torch::Tensor& dst_cache, torch::Tensor& src_cache,
41 const double scale, const std::string& kv_cache_dtype);
csrc/cache_kernels.cu
-103
@@ -396,51 +396,6 @@ __global__ void concat_and_cache_mla_kernel(
396 copy(k_pe, kv_cache, k_pe_stride, block_stride, pe_dim, kv_lora_rank);
397 }
398
399 -template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
400 -__global__ void cp_fused_concat_and_cache_mla_kernel(
401 - const scalar_t* __restrict__ kv_c, // [num_full_tokens, kv_lora_rank]
402 - const scalar_t* __restrict__ k_pe, // [num_full_tokens, pe_dim]
403 - const int64_t* __restrict__ cp_local_token_select_indices, // [num_tokens]
404 - cache_t* __restrict__ kv_cache, // [num_blocks, block_size, (kv_lora_rank
405 - // + pe_dim)]
406 - const int64_t* __restrict__ slot_mapping, // [num_tokens]
407 - const int block_stride, //
408 - const int entry_stride, //
409 - const int kv_c_stride, //
410 - const int k_pe_stride, //
411 - const int kv_lora_rank, //
412 - const int pe_dim, //
413 - const int block_size, //
414 - const float* scale //
415 -) {
416 - const int64_t token_idx = cp_local_token_select_indices[blockIdx.x];
417 - const int64_t slot_idx = slot_mapping[blockIdx.x];
418 - // NOTE: slot_idx can be -1 if the token is padded
419 - if (slot_idx < 0) {
420 - return;
421 - }
422 - const int64_t block_idx = slot_idx / block_size;
423 - const int64_t block_offset = slot_idx % block_size;
424 -
425 - auto copy = [&](const scalar_t* __restrict__ src, cache_t* __restrict__ dst,
426 - int src_stride, int dst_stride, int size, int offset) {
427 - for (int i = threadIdx.x; i < size; i += blockDim.x) {
428 - const int64_t src_idx = token_idx * src_stride + i;
429 - const int64_t dst_idx =
430 - block_idx * block_stride + block_offset * entry_stride + i + offset;
431 - if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) {
432 - dst[dst_idx] = src[src_idx];
433 - } else {
434 - dst[dst_idx] =
435 - fp8::scaled_convert<cache_t, scalar_t, kv_dt>(src[src_idx], *scale);
436 - }
437 - }
438 - };
439 -
440 - copy(kv_c, kv_cache, kv_c_stride, block_stride, kv_lora_rank, 0);
441 - copy(k_pe, kv_cache, k_pe_stride, block_stride, pe_dim, kv_lora_rank);
442 -}
443 -
399 } // namespace vllm
400
401 // KV_T is the data type of key and value tensors.
@@ -554,20 +509,6 @@ void reshape_and_cache_flash(
509 kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \
510 reinterpret_cast<const float*>(scale.data_ptr()));
511
557 -// KV_T is the data type of key and value tensors.
558 -// CACHE_T is the stored data type of kv-cache.
559 -// KV_DTYPE is the real data type of kv-cache.
560 -#define CALL_CP_FUSED_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE) \
561 - vllm::cp_fused_concat_and_cache_mla_kernel<KV_T, CACHE_T, KV_DTYPE> \
562 - <<<grid, block, 0, stream>>>( \
563 - reinterpret_cast<KV_T*>(kv_c.data_ptr()), \
564 - reinterpret_cast<KV_T*>(k_pe.data_ptr()), \
565 - cp_local_token_select_indices.data_ptr<int64_t>(), \
566 - reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
567 - slot_mapping.data_ptr<int64_t>(), block_stride, entry_stride, \
568 - kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \
569 - reinterpret_cast<const float*>(scale.data_ptr()));
570 -
512 void concat_and_cache_mla(
513 torch::Tensor& kv_c, // [num_tokens, kv_lora_rank]
514 torch::Tensor& k_pe, // [num_tokens, pe_dim]
@@ -606,50 +547,6 @@ void concat_and_cache_mla(
547 CALL_CONCAT_AND_CACHE_MLA);
548 }
549
609 -// Note(hc): cp_fused_concat_and_cache_mla fuses the following three kernel
610 -// calls into one:
611 -// k_c_normed.index_select(0, cp_local_token_select_indices) + \
612 -// k_pe.squeeze(1).index_select(0, cp_local_token_select_indices) + \
613 -// concat_and_cache_mla.
614 -void cp_fused_concat_and_cache_mla(
615 - torch::Tensor& kv_c, // [num_total_tokens, kv_lora_rank]
616 - torch::Tensor& k_pe, // [num_total_tokens, pe_dim]
617 - torch::Tensor& cp_local_token_select_indices, // [num_tokens]
618 - torch::Tensor& kv_cache, // [num_blocks, block_size, (kv_lora_rank +
619 - // pe_dim)]
620 - torch::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens]
621 - const std::string& kv_cache_dtype, torch::Tensor& scale) {
622 - // NOTE(woosuk): In vLLM V1, key.size(0) can be different from
623 - // slot_mapping.size(0) because of padding for CUDA graphs.
624 - // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because
625 - // both include padding.
626 - // In vLLM V1, however, key.size(0) can be larger than slot_mapping.size(0)
627 - // since key includes padding for CUDA graphs, while slot_mapping does not.
628 - // In this case, slot_mapping.size(0) represents the actual number of tokens
629 - // before padding.
630 - // For compatibility with both cases, we use slot_mapping.size(0) as the
631 - // number of tokens.
632 - int num_tokens = slot_mapping.size(0);
633 - int kv_lora_rank = kv_c.size(1);
634 - int pe_dim = k_pe.size(1);
635 - int block_size = kv_cache.size(1);
636 -
637 - TORCH_CHECK(kv_cache.size(2) == kv_lora_rank + pe_dim);
638 -
639 - int kv_c_stride = kv_c.stride(0);
640 - int k_pe_stride = k_pe.stride(0);
641 - int block_stride = kv_cache.stride(0);
642 - int entry_stride = kv_cache.stride(1);
643 -
644 - dim3 grid(num_tokens);
645 - dim3 block(std::min(kv_lora_rank, 512));
646 - const at::cuda::OptionalCUDAGuard device_guard(device_of(kv_c));
647 - const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
648 -
649 - DISPATCH_BY_KV_CACHE_DTYPE(kv_c.dtype(), kv_cache_dtype,
650 - CALL_CP_FUSED_CONCAT_AND_CACHE_MLA);
651 -}
652 -
550 namespace vllm {
551
552 template <typename Tout, typename Tin, Fp8KVCacheDataType kv_dt>
csrc/torch_bindings.cpp
-10
@@ -693,16 +693,6 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
693 " Tensor scale) -> ()");
694 cache_ops.impl("concat_and_cache_mla", torch::kCUDA, &concat_and_cache_mla);
695
696 - cache_ops.def(
697 - "cp_fused_concat_and_cache_mla(Tensor kv_c, Tensor k_pe,"
698 - " Tensor cp_local_token_select_indices,"
699 - " Tensor! kv_cache,"
700 - " Tensor slot_mapping,"
701 - " str kv_cache_dtype,"
702 - " Tensor scale) -> ()");
703 - cache_ops.impl("cp_fused_concat_and_cache_mla", torch::kCUDA,
704 - &cp_fused_concat_and_cache_mla);
705 -
696 // Convert the key and value cache to fp8 data type.
697 cache_ops.def(
698 "convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, "
tests/distributed/test_context_parallel.py new
+263
@@ -0,0 +1,263 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +WARNING: This test runs in both single-node (4 GPUs) and multi-node
5 + (2 node with 2 GPUs each) modes. If the test only uses 2 GPUs, it is
6 + important to set the distributed backend to "mp" to avoid Ray scheduling
7 + all workers in a node other than the head node, which can cause the test
8 + to fail.
9 +"""
10 +import json
11 +import os
12 +from dataclasses import dataclass
13 +from typing import Literal, NamedTuple, Optional
14 +
15 +import pytest
16 +
17 +from vllm.config import RunnerOption
18 +from vllm.logger import init_logger
19 +
20 +from ..models.registry import HF_EXAMPLE_MODELS
21 +from ..utils import compare_two_settings, create_new_process_for_each_test
22 +
23 +logger = init_logger("test_context_parallel")
24 +
25 +VLLM_MULTI_NODE = os.getenv("VLLM_MULTI_NODE", "0") == "1"
26 +
27 +
28 +class ParallelSetup(NamedTuple):
29 + tp_size: int
30 + pp_size: int
31 + dcp_size: int
32 + eager_mode: bool
33 + chunked_prefill: bool
34 +
35 +
36 +class CPTestOptions(NamedTuple):
37 + multi_node_only: bool
38 + load_format: Optional[str] = None
39 +
40 +
41 +@dataclass
42 +class CPTestSettings:
43 + parallel_setups: list[ParallelSetup]
44 + # NOTE: the length of distributed_backends and
45 + # vllm_major_versions should be the same, and they
46 + # are first zipped together to iterate over all
47 + # test settings.
48 + distributed_backends: list[str]
49 + # vllm major version: "0" for V0, "1" for V1
50 + vllm_major_versions: list[str]
51 + runner: RunnerOption
52 + test_options: CPTestOptions
53 +
54 + def __post_init__(self):
55 + if len(self.distributed_backends) != len(self.vllm_major_versions):
56 + raise ValueError(
57 + f"Length mismatch: distributed_backends "
58 + f"({len(self.distributed_backends)}) != "
59 + f"vllm_major_versions ({len(self.vllm_major_versions)})")
60 +
61 + @staticmethod
62 + def detailed(
63 + *,
64 + tp_base: int = 4,
65 + pp_base: int = 1,
66 + dcp_base: int = 1,
67 + multi_node_only: bool = False,
68 + runner: RunnerOption = "auto",
69 + load_format: Optional[str] = None,
70 + ):
71 + parallel_setups = []
72 + for eager_mode_val in [False]:
73 + for pp_multiplier in [1]:
74 + for dcp_multiplier in [2, 4]:
75 + for chunked_prefill_val in [True]:
76 + parallel_setups.append(
77 + ParallelSetup(tp_size=tp_base,
78 + pp_size=pp_multiplier * pp_base,
79 + dcp_size=dcp_multiplier * dcp_base,
80 + eager_mode=eager_mode_val,
81 + chunked_prefill=chunked_prefill_val))
82 + return CPTestSettings(
83 + parallel_setups=parallel_setups,
84 + distributed_backends=["mp"],
85 + vllm_major_versions=["1"],
86 + runner=runner,
87 + test_options=CPTestOptions(multi_node_only=multi_node_only,
88 + load_format=load_format),
89 + )
90 +
91 + def iter_params(self, model_id: str):
92 + opts = self.test_options
93 +
94 + for parallel_setup in self.parallel_setups:
95 + for backend, vllm_major_version in zip(self.distributed_backends,
96 + self.vllm_major_versions):
97 + yield (model_id, parallel_setup, backend, vllm_major_version,
98 + self.runner, opts)
99 +
100 +
101 +def _compare_cp_with_tp(
102 + model_id: str,
103 + parallel_setup: ParallelSetup,
104 + distributed_backend: str,
105 + vllm_major_version: str,
106 + runner: RunnerOption,
107 + test_options: CPTestOptions,
108 + num_gpus_available: int,
109 + *,
110 + method: Literal["generate"],
111 + is_multimodal: bool,
112 +):
113 + (
114 + tp_size,
115 + pp_size,
116 + dcp_size,
117 + eager_mode,
118 + chunked_prefill,
119 + ) = parallel_setup
120 +
121 + multi_node_only, load_format = test_options
122 +
123 + model_info = HF_EXAMPLE_MODELS.find_hf_info(model_id)
124 + model_info.check_transformers_version(on_fail="skip")
125 +
126 + trust_remote_code = model_info.trust_remote_code
127 + tokenizer_mode = model_info.tokenizer_mode
128 + hf_overrides = model_info.hf_overrides
129 +
130 + if load_format == "dummy":
131 + # Avoid OOM
132 + text_overrides = {
133 + "num_hidden_layers": 4,
134 + "hidden_size": 512,
135 + "intermediate_size": 800,
136 + "num_attention_heads": 4,
137 + "num_key_value_heads": 1,
138 + }
139 +
140 + if is_multimodal:
141 + hf_overrides.update({"text_config": text_overrides})
142 + else:
143 + hf_overrides.update(text_overrides)
144 + else:
145 + model_info.check_available_online(on_fail="skip")
146 +
147 + if num_gpus_available < tp_size * pp_size:
148 + pytest.skip(f"Need at least {tp_size} x {pp_size} GPUs")
149 + if VLLM_MULTI_NODE and distributed_backend == "mp":
150 + pytest.skip("Skipping multi-node pipeline parallel test for "
151 + "multiprocessing distributed backend")
152 + if multi_node_only and not VLLM_MULTI_NODE:
153 + pytest.skip("Not in multi-node setting")
154 +
155 + common_args = [
156 + # use half precision for speed and memory savings in CI environment
157 + "--dtype",
158 + "bfloat16",
159 + "--max-model-len",
160 + "2048",
161 + "--max-num-seqs",
162 + "8",
163 + ]
164 + if chunked_prefill:
165 + common_args.append("--enable-chunked-prefill")
166 + if eager_mode:
167 + common_args.append("--enforce-eager")
168 + if runner != "auto":
169 + common_args.extend(["--runner", runner])
170 + if trust_remote_code:
171 + common_args.append("--trust-remote-code")
172 + if tokenizer_mode:
173 + common_args.extend(["--tokenizer-mode", tokenizer_mode])
174 + if load_format:
175 + common_args.extend(["--load-format", load_format])
176 + if hf_overrides:
177 + common_args.extend(["--hf-overrides", json.dumps(hf_overrides)])
178 +
179 + cp_env = tp_env = {
180 + "VLLM_USE_V1":
181 + vllm_major_version, # Note(hc): DCP only support V1 engine only
182 + }
183 +
184 + cp_args = [
185 + *common_args,
186 + "--tensor-parallel-size",
187 + str(tp_size),
188 + "--pipeline-parallel-size",
189 + str(pp_size),
190 + "--decode-context-parallel-size",
191 + str(dcp_size),
192 + "--distributed-executor-backend",
193 + distributed_backend,
194 + ]
195 +
196 + tp_args = [
197 + *common_args,
198 + "--tensor-parallel-size",
199 + str(tp_size),
200 + "--pipeline-parallel-size",
201 + str(pp_size),
202 + "--distributed-executor-backend",
203 + distributed_backend,
204 + ]
205 +
206 + try:
207 + compare_two_settings(model_id,
208 + cp_args,
209 + tp_args,
210 + cp_env,
211 + tp_env,
212 + method=method,
213 + max_wait_seconds=720)
214 + except Exception:
215 + testing_ray_compiled_graph = cp_env is not None
216 + if testing_ray_compiled_graph and vllm_major_version == "0":
217 + # Ray Compiled Graph tests are flaky for V0,
218 + # so we don't want to fail the test
219 + logger.exception("Ray Compiled Graph tests failed")
220 + else:
221 + raise
222 +
223 +
224 +CP_TEXT_GENERATION_MODELS = {
225 + # [MLA attention only]
226 + "deepseek-ai/DeepSeek-V2-Lite-Chat": CPTestSettings.detailed(),
227 +}
228 +
229 +CP_TEST_MODELS = [
230 + # TODO support other models
231 + # [LANGUAGE GENERATION]
232 + "deepseek-ai/DeepSeek-V2-Lite-Chat",
233 +]
234 +
235 +
236 +@pytest.mark.parametrize(
237 + ("model_id", "parallel_setup", "distributed_backend", "vllm_major_version",
238 + "runner", "test_options"),
239 + [
240 + params for model_id, settings in CP_TEXT_GENERATION_MODELS.items()
241 + for params in settings.iter_params(model_id)
242 + if model_id in CP_TEST_MODELS
243 + ],
244 +)
245 +@create_new_process_for_each_test()
246 +def test_cp_generation(
247 + model_id: str,
248 + parallel_setup: ParallelSetup,
249 + distributed_backend: str,
250 + vllm_major_version: str,
251 + runner: RunnerOption,
252 + test_options: CPTestOptions,
253 + num_gpus_available,
254 +):
255 + _compare_cp_with_tp(model_id,
256 + parallel_setup,
257 + distributed_backend,
258 + vllm_major_version,
259 + runner,
260 + test_options,
261 + num_gpus_available,
262 + method="generate",
263 + is_multimodal=False)
vllm/_custom_ops.py
-14
@@ -1625,20 +1625,6 @@ def concat_and_cache_mla(
1625 scale)
1626
1627
1628 -def cp_fused_concat_and_cache_mla(
1629 - kv_c: torch.Tensor,
1630 - k_pe: torch.Tensor,
1631 - cp_local_token_select_indices: torch.Tensor,
1632 - kv_cache: torch.Tensor,
1633 - slot_mapping: torch.Tensor,
1634 - kv_cache_dtype: str,
1635 - scale: torch.Tensor,
1636 -) -> None:
1637 - torch.ops._C_cache_ops.cp_fused_concat_and_cache_mla(
1638 - kv_c, k_pe, cp_local_token_select_indices, kv_cache, slot_mapping,
1639 - kv_cache_dtype, scale)
1640 -
1641 -
1628 def copy_blocks(key_caches: list[torch.Tensor],
1629 value_caches: list[torch.Tensor],
1630 block_mapping: torch.Tensor) -> None:
vllm/attention/ops/common.py new
+139
@@ -0,0 +1,139 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import torch
4 +
5 +from vllm.distributed.parallel_state import GroupCoordinator
6 +from vllm.triton_utils import tl, triton
7 +
8 +
9 +@triton.jit
10 +def _correct_attn_cp_out_kernel(outputs_ptr, new_output_ptr, lses_ptr,
11 + vlse_ptr, outputs_stride_B, outputs_stride_H,
12 + outputs_stride_D, lses_stride_N, lses_stride_B,
13 + lses_stride_H, lse_idx, HEAD_DIM: tl.constexpr,
14 + N_ROUNDED: tl.constexpr):
15 + """
16 + Apply the all-gathered lses to correct each local rank's attention
17 + output. we still need perform a cross-rank reduction to obtain the
18 + final attention output.
19 +
20 + Args:
21 + output: [ B, H, D ]
22 + lses : [ N, B, H ]
23 + cp, batch, q_heads, v_head_dim
24 + Return:
25 + output: [ B, H, D ]
26 + lse : [ B, H ]
27 + """
28 + batch_idx = tl.program_id(axis=0).to(tl.int64)
29 + head_idx = tl.program_id(axis=1).to(tl.int64)
30 + d_offsets = tl.arange(0, HEAD_DIM)
31 + num_n_offsets = tl.arange(0, N_ROUNDED)
32 +
33 + # shape = [N]
34 + lse_offsets = num_n_offsets * lses_stride_N + batch_idx * \
35 + lses_stride_B + head_idx * lses_stride_H
36 +
37 + # calc final lse
38 + lse = tl.load(lses_ptr + lse_offsets)
39 + lse = tl.where((lse != lse) | (lse == float('inf')), -float('inf'), lse)
40 + lse_max = tl.max(lse, axis=0)
41 + lse -= lse_max
42 + lse_exp = tl.exp(lse)
43 + lse_acc = tl.sum(lse_exp, axis=0)
44 + lse = tl.log(lse_acc)
45 + lse += lse_max
46 +
47 + lse_offsets = batch_idx * lses_stride_B + head_idx * lses_stride_H
48 + tl.store(vlse_ptr + lse_offsets, lse)
49 +
50 + # shape = [D]
51 + output_offsets = batch_idx * outputs_stride_B + \
52 + head_idx * outputs_stride_H + \
53 + d_offsets * outputs_stride_D
54 +
55 + # correct output
56 + lse_offset = lse_idx * lses_stride_N + batch_idx * \
57 + lses_stride_B + head_idx * lses_stride_H
58 + lse_tmp = tl.load(lses_ptr + lse_offset)
59 + lse_finally = lse_tmp - lse
60 + lse_finally = tl.where(
61 + (lse_finally != lse_finally) | (lse_finally == float('inf')),
62 + -float('inf'), lse_finally)
63 + factor = tl.exp(lse_finally)
64 + output = tl.load(outputs_ptr + output_offsets)
65 + output = output * factor
66 +
67 + tl.store(new_output_ptr + output_offsets, output)
68 +
69 +
70 +class CPTritonContext:
71 + """ The CPTritonContext is used to avoid recompilation of the Triton JIT.
72 + """
73 +
74 + def __init__(self):
75 + self.inner_kernel = None
76 +
77 + def call_kernel(self, kernel, grid, *regular_args, **const_args):
78 + if self.inner_kernel is None:
79 + self.inner_kernel = kernel[grid](*regular_args, **const_args)
80 + else:
81 + self.inner_kernel[grid](*regular_args)
82 +
83 +
84 +def correct_attn_out(out: torch.Tensor, lses: torch.Tensor, cp_rank: int,
85 + ctx: CPTritonContext):
86 + """
87 + Apply the all-gathered lses to correct each local rank's attention
88 + output. we still need perform a cross-rank reduction to obtain the
89 + final attention output.
90 +
91 + Args:
92 + output: [ B, H, D ]
93 + lses : [ N, B, H ]
94 + Return:
95 + output: [ B, H, D ]
96 + lse : [ B, H ]
97 + """
98 + if ctx is None:
99 + ctx = CPTritonContext()
100 +
101 + lse = torch.empty_like(lses[0])
102 +
103 + grid = (out.shape[0], out.shape[1], 1)
104 + regular_args = (out, out, lses, lse, *out.stride(), *lses.stride(),
105 + cp_rank)
106 + const_args = {
107 + "HEAD_DIM": out.shape[-1],
108 + "N_ROUNDED": lses.shape[0],
109 + }
110 +
111 + ctx.call_kernel(_correct_attn_cp_out_kernel, grid, *regular_args,
112 + **const_args)
113 + return out, lse
114 +
115 +
116 +def cp_lse_ag_out_rs(cp_attn_out: torch.Tensor,
117 + cp_attn_lse: torch.Tensor,
118 + cp_group: GroupCoordinator,
119 + ctx: CPTritonContext = None):
120 + """
121 + cp_attn_out: [ B, H, D ]
122 + cp_attn_lse: [ B, H ]
123 + """
124 + if cp_group.world_size == 1:
125 + return cp_attn_out
126 +
127 + if ctx is None:
128 + ctx = CPTritonContext()
129 +
130 + lses = torch.empty((cp_group.world_size, ) + cp_attn_lse.shape,
131 + dtype=cp_attn_lse.dtype,
132 + device=cp_attn_lse.device)
133 +
134 + cp_attn_lse = cp_attn_lse.contiguous()
135 + lses = cp_group.all_gather(cp_attn_lse, dim=0).view_as(lses)
136 + out, _ = correct_attn_out(cp_attn_out, lses, cp_group.rank_in_group, ctx)
137 + assert out.is_contiguous()
138 + out = cp_group.reduce_scatter(out, dim=1)
139 + return out
vllm/attention/ops/flashmla.py
+3 -1
@@ -105,7 +105,9 @@ def flash_mla_with_kvcache(
105 descale_q,
106 descale_k,
107 )
108 - return out, softmax_lse
108 +
109 + # Note(hc): need revisit when we support DCP with decode query_len > 1.
110 + return out.squeeze(1), softmax_lse.squeeze(-1)
111
112
113 #
vllm/config/parallel.py
+5
@@ -170,6 +170,11 @@ class ParallelConfig:
170 Set to be private as it's not intended to be configured by users.
171 """
172
173 + decode_context_parallel_size: int = 1
174 + """Number of decode context parallel groups, because the world size does
175 + not change by dcp, it simply reuse the GPUs of TP group, and tp_size
176 + needs to be divisible by dcp_size."""
177 +
178 @property
179 def world_size_across_dp(self) -> int:
180 """world_size_across_dp is TPxPPxDP, it is the size of the world
vllm/distributed/parallel_state.py
+48 -1
@@ -904,6 +904,18 @@ def get_tensor_model_parallel_group():
904 return get_tp_group()
905
906
907 +_DCP: Optional[GroupCoordinator] = None
908 +
909 +
910 +def get_dcp_group() -> GroupCoordinator:
911 + assert _DCP is not None, (
912 + "decode context model parallel group is not initialized")
913 + return _DCP
914 +
915 +
916 +# kept for backward compatibility
917 +get_context_model_parallel_group = get_dcp_group
918 +
919 _PP: Optional[GroupCoordinator] = None
920
921 _DP: Optional[GroupCoordinator] = None
@@ -1034,6 +1046,7 @@ def init_distributed_environment(
1046 def initialize_model_parallel(
1047 tensor_model_parallel_size: int = 1,
1048 pipeline_model_parallel_size: int = 1,
1049 + decode_context_model_parallel_size: Optional[int] = 1,
1050 backend: Optional[str] = None,
1051 ) -> None:
1052 """
@@ -1098,6 +1111,23 @@ def initialize_model_parallel(
1111 use_message_queue_broadcaster=True,
1112 group_name="tp")
1113
1114 + # Build the DCP model-parallel groups.
1115 + global _DCP
1116 + assert _DCP is None, (
1117 + "decode context model parallel group is already initialized")
1118 + # Note(hc): In the current implementation of decode context parallel,
1119 + # dcp_size must not exceed tp_size, because the world size does not
1120 + # change by DCP, it simply reuse the GPUs of TP group, and split one
1121 + # TP group into tp_size//dcp_size DCP groups.
1122 + group_ranks = all_ranks.reshape(
1123 + -1, decode_context_model_parallel_size).unbind(0)
1124 + group_ranks = [x.tolist() for x in group_ranks]
1125 + _DCP = init_model_parallel_group(group_ranks,
1126 + get_world_group().local_rank,
1127 + backend,
1128 + use_message_queue_broadcaster=True,
1129 + group_name="dcp")
1130 +
1131 # Build the pipeline model-parallel groups.
1132 global _PP
1133 assert _PP is None, (
@@ -1141,6 +1171,7 @@ def initialize_model_parallel(
1171 def ensure_model_parallel_initialized(
1172 tensor_model_parallel_size: int,
1173 pipeline_model_parallel_size: int,
1174 + decode_context_model_parallel_size: Optional[int] = 1,
1175 backend: Optional[str] = None,
1176 ) -> None:
1177 """Helper to initialize model parallel groups if they are not initialized,
@@ -1151,7 +1182,8 @@ def ensure_model_parallel_initialized(
1182 get_world_group().device_group)
1183 if not model_parallel_is_initialized():
1184 initialize_model_parallel(tensor_model_parallel_size,
1154 - pipeline_model_parallel_size, backend)
1185 + pipeline_model_parallel_size,
1186 + decode_context_model_parallel_size, backend)
1187 return
1188
1189 assert (
@@ -1226,6 +1258,16 @@ def get_tensor_model_parallel_rank():
1258 return get_tp_group().rank_in_group
1259
1260
1261 +def get_decode_context_model_parallel_world_size():
1262 + """Return world size for the decode context model parallel group."""
1263 + return get_dcp_group().world_size
1264 +
1265 +
1266 +def get_decode_context_model_parallel_rank():
1267 + """Return my rank for the decode context model parallel group."""
1268 + return get_dcp_group().rank_in_group
1269 +
1270 +
1271 def get_node_count() -> int:
1272 """Return the total number of nodes in the distributed environment. """
1273 assert _NODE_COUNT is not None, (
@@ -1246,6 +1288,11 @@ def destroy_model_parallel():
1288 _PP.destroy()
1289 _PP = None
1290
1291 + global _DCP
1292 + if _DCP:
1293 + _DCP.destroy()
1294 + _DCP = None
1295 +
1296 global _DP
1297 if _DP:
1298 _DP.destroy()
vllm/engine/arg_utils.py
+17
@@ -306,6 +306,8 @@ class EngineArgs:
306 # number of P/D disaggregation (or other disaggregation) workers
307 pipeline_parallel_size: int = ParallelConfig.pipeline_parallel_size
308 tensor_parallel_size: int = ParallelConfig.tensor_parallel_size
309 + decode_context_parallel_size: int = \
310 + ParallelConfig.decode_context_parallel_size
311 data_parallel_size: int = ParallelConfig.data_parallel_size
312 data_parallel_rank: Optional[int] = None
313 data_parallel_start_rank: Optional[int] = None
@@ -636,6 +638,9 @@ class EngineArgs:
638 **parallel_kwargs["pipeline_parallel_size"])
639 parallel_group.add_argument("--tensor-parallel-size", "-tp",
640 **parallel_kwargs["tensor_parallel_size"])
641 + parallel_group.add_argument(
642 + "--decode-context-parallel-size", "-dcp",
643 + **parallel_kwargs["decode_context_parallel_size"])
644 parallel_group.add_argument("--data-parallel-size", "-dp",
645 **parallel_kwargs["data_parallel_size"])
646 parallel_group.add_argument(
@@ -1156,6 +1161,17 @@ class EngineArgs:
1161 # global layers in interleaved sliding window models.
1162 sliding_window = model_config.get_sliding_window()
1163
1164 + # Note(hc): In the current implementation of decode context
1165 + # parallel(DCP), tp_size needs to be divisible by dcp_size,
1166 + # because the world size does not change by dcp, it simply
1167 + # reuse the GPUs of TP group, and split one TP group into
1168 + # tp_size//dcp_size DCP groups.
1169 + assert self.tensor_parallel_size % self.decode_context_parallel_size \
1170 + == 0, (
1171 + f"tp_size={self.tensor_parallel_size} must be divisible by"
1172 + f"dcp_size={self.decode_context_parallel_size}."
1173 + )
1174 +
1175 cache_config = CacheConfig(
1176 block_size=self.block_size,
1177 gpu_memory_utilization=self.gpu_memory_utilization,
@@ -1306,6 +1322,7 @@ class EngineArgs:
1322 distributed_executor_backend=self.distributed_executor_backend,
1323 worker_cls=self.worker_cls,
1324 worker_extension_cls=self.worker_extension_cls,
1325 + decode_context_parallel_size=self.decode_context_parallel_size,
1326 )
1327
1328 speculative_config = self.create_speculative_config(
vllm/v1/attention/backends/mla/common.py
+309 -25
@@ -201,10 +201,11 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionLayer,
201 AttentionMetadata,
202 MLAAttentionImpl)
203 from vllm.attention.backends.utils import get_mla_dims
204 +from vllm.attention.ops.common import cp_lse_ag_out_rs
205 from vllm.attention.ops.merge_attn_states import merge_attn_states
206 from vllm.attention.utils.fa_utils import get_flash_attn_version
207 from vllm.config import VllmConfig
207 -from vllm.distributed.parallel_state import is_global_first_rank
208 +from vllm.distributed.parallel_state import get_dcp_group, is_global_first_rank
209 from vllm.logger import init_logger
210 from vllm.model_executor.layers.linear import (ColumnParallelLinear,
211 LinearBase,
@@ -323,6 +324,13 @@ class MLACommonPrefillMetadata:
324 seq_lens: torch.Tensor
325 workspace: torch.Tensor
326
327 + # for mla DCP
328 + cp_chunk_seq_lens: Optional[list[list[int]]] = None
329 + origin_context_lens: Optional[list[int]] = None
330 + cp_cu_seq_lens: Optional[torch.Tensor] = None
331 + chunk_size: Optional[int] = None
332 + cu_seq_lens_lst: Optional[list[list[int]]] = None
333 +
334 block_table: torch.Tensor
335 query_start_loc: torch.Tensor
336 max_query_len: int
@@ -444,6 +452,13 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
452 parallel_config)
453 self.mla_dims = get_mla_dims(self.model_config)
454 self.aot_schedule = current_platform.is_cuda()
455 + try:
456 + self.dcp_world_size = get_dcp_group().world_size
457 + self.dcp_rank = get_dcp_group().rank_in_group
458 + except AssertionError:
459 + # DCP might not be initialized in testing
460 + self.dcp_world_size = 1
461 + self.dcp_rank = 0
462
463 # Dont try to access the runner on AMD
464 if self.aot_schedule:
@@ -465,12 +480,27 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
480 128 * 1024)
481 assert self.chunked_prefill_workspace_size >= \
482 scheduler_config.max_num_seqs * cache_config.block_size
468 - self.chunked_prefill_workspace = torch.empty(
469 - (self.chunked_prefill_workspace_size,
470 - self.model_config.get_head_size()),
471 - dtype=self.model_config.dtype,
472 - device=device,
473 - )
483 + if self.dcp_world_size > 1:
484 + # Note(hc): The local kvcache is incomplete when DCP is triggered,
485 + # an additional kvcache allgather across the DCP group is therefore
486 + # required, so the workspace has to be enlarged by 1/DCP relative
487 + # to the original TP allocation.
488 + assert self.chunked_prefill_workspace_size % \
489 + self.dcp_world_size == 0
490 + self.chunked_prefill_workspace = torch.empty(
491 + (self.chunked_prefill_workspace_size +
492 + self.chunked_prefill_workspace_size // self.dcp_world_size,
493 + self.model_config.get_head_size()),
494 + dtype=self.model_config.dtype,
495 + device=device,
496 + )
497 + else:
498 + self.chunked_prefill_workspace = torch.empty(
499 + (self.chunked_prefill_workspace_size,
500 + self.model_config.get_head_size()),
501 + dtype=self.model_config.dtype,
502 + device=device,
503 + )
504
505 self._use_cudnn_prefill = use_cudnn_prefill()
506 self._use_fi_prefill = use_flashinfer_prefill()
@@ -631,6 +661,12 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
661 split_decodes_and_prefills(common_attn_metadata,
662 decode_threshold=self.reorder_batch_threshold)
663
664 + # Note(hc): update seq_lens of decode reqs under DCP.
665 + if self.dcp_world_size > 1:
666 + seq_lens[:num_decodes] = seq_lens[:num_decodes] \
667 + // self.dcp_world_size + (self.dcp_rank <= \
668 + (seq_lens[:num_decodes] - 1) % self.dcp_world_size)
669 +
670 assert num_decodes + num_prefills == num_reqs
671 assert num_decode_tokens + num_prefill_tokens == num_tokens
672
@@ -639,6 +675,10 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
675 reqs_start = num_decodes # prefill_start
676
677 context_lens_cpu = num_computed_tokens_cpu[reqs_start:num_reqs]
678 + # Note(hc): The context lengths in the perspective of dcp rank0.
679 + cp_context_lens_cpu = torch.ceil(context_lens_cpu.float() /
680 + self.dcp_world_size).int()
681 + origin_context_lens = context_lens_cpu.tolist()
682 max_context_len_cpu = context_lens_cpu.max().item()
683 num_prefills_with_context_cpu = (context_lens_cpu > 0).sum().item()
684 prefill_query_start_loc = query_start_loc[
@@ -691,20 +731,66 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
731 out=cu_seq_lens_cpu[:, 1:],
732 dtype=torch.int32)
733
734 + if self.dcp_world_size > 1:
735 + # Note(hc): The above max_context_chunk already enforces
736 + # block_size alignment, DCP just need the block_size can
737 + # be divisible by dcp_world_size, because DCP use
738 + # cp_gather_cache which not require `cp_chunk_starts`
739 + # aligned to page_size.
740 + assert max_context_chunk % self.dcp_world_size == 0
741 + cp_max_context_chunk = max_context_chunk // \
742 + self.dcp_world_size
743 + cp_chunk_starts = \
744 + torch.arange(num_chunks, dtype=torch.int32) \
745 + .unsqueeze(1).expand(-1, num_prefills) \
746 + * cp_max_context_chunk
747 + cp_chunk_ends = torch.min(
748 + cp_context_lens_cpu.unsqueeze(0),
749 + cp_chunk_starts + cp_max_context_chunk)
750 + cp_chunk_seq_lens = (cp_chunk_ends -
751 + cp_chunk_starts).clamp(min=0)
752 +
753 + cp_cu_seq_lens_cpu = torch.zeros(num_chunks,
754 + num_prefills + 1,
755 + dtype=torch.int32,
756 + pin_memory=True)
757 + torch.cumsum(cp_chunk_seq_lens,
758 + dim=1,
759 + out=cp_cu_seq_lens_cpu[:, 1:],
760 + dtype=torch.int32)
761 +
762 chunked_context_metadata_cls = \
763 CudnnPrefillMetadata.ChunkedContextMetadata \
764 if self._use_cudnn_prefill else \
765 MLACommonPrefillMetadata.ChunkedContextMetadata
698 -
699 - chunked_context_metadata = \
700 - chunked_context_metadata_cls(
701 - cu_seq_lens=cu_seq_lens_cpu.to(device, non_blocking=True),
702 - starts=chunk_starts.to(device, non_blocking=True),
703 - seq_tot=chunk_seq_lens.sum(dim=1).tolist(),
704 - max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),
705 - seq_lens=chunk_seq_lens,
706 - workspace=self.chunked_prefill_workspace,
707 - )
766 + if self.dcp_world_size > 1:
767 + chunked_context_metadata = \
768 + chunked_context_metadata_cls(
769 + cu_seq_lens=cu_seq_lens_cpu \
770 + .to(device, non_blocking=True),
771 + starts=cp_chunk_starts.to(device, non_blocking=True),
772 + seq_tot=cp_chunk_seq_lens.sum(dim=1).tolist(),
773 + max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),
774 + seq_lens=chunk_seq_lens,
775 + workspace=self.chunked_prefill_workspace,
776 + cp_chunk_seq_lens=cp_chunk_seq_lens.tolist(),
777 + origin_context_lens=origin_context_lens,
778 + cp_cu_seq_lens=cp_cu_seq_lens_cpu \
779 + .to(device, non_blocking=True),
780 + chunk_size=max_context_chunk,
781 + cu_seq_lens_lst=cu_seq_lens_cpu.tolist(),
782 + )
783 + else:
784 + chunked_context_metadata = \
785 + chunked_context_metadata_cls(
786 + cu_seq_lens=cu_seq_lens_cpu \
787 + .to(device, non_blocking=True),
788 + starts=chunk_starts.to(device, non_blocking=True),
789 + seq_tot=chunk_seq_lens.sum(dim=1).tolist(),
790 + max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),
791 + seq_lens=chunk_seq_lens,
792 + workspace=self.chunked_prefill_workspace,
793 + )
794
795 if self._use_cudnn_prefill:
796 chunked_context_metadata.seq_lens = chunk_seq_lens
@@ -757,6 +843,71 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
843 return attn_metadata
844
845
846 +def reorg_kvcache(
847 + allgatered_kv_c_normed: torch.Tensor,
848 + allgatered_k_pe: torch.Tensor,
849 + cp_chunk_seq_lens_lst: list[int],
850 + origin_context_lens: list[int],
851 + cp_world_size: int,
852 + sum_seq_len: int,
853 + max_seq_len: int,
854 + chunk_size: int,
855 + chunk_idx: int,
856 + toks: int,
857 +) -> tuple[torch.Tensor, torch.Tensor]:
858 + """
859 + reorg kvcache after cp local gather to tp layout for attn kernel.
860 +
861 + Args:
862 + cp_chunk_seq_lens_lst: chunk context lengths under CP.
863 + origin_context_lens: origin full context lengths under CP.
864 + cp_world_size: CP size.
865 + sum_seq_len: the sum of cp_chunk_seq_lens_lst.
866 + max_seq_len: the max value of cp_chunk_seq_lens_lst.
867 + chunk_size: equals to max_context_chunk from
868 + chunked_context_metadata building.
869 + chunk_idx: chunk idx of chunked_prefill.
870 + toks: the number of tokens for local gather cache.
871 + """
872 + kv_c_segments = []
873 + k_pe_segments = []
874 + src_token_idx = 0
875 + max_seq_len_check = 0
876 + for cp_chunk_seq_len, origin_context_len in zip(cp_chunk_seq_lens_lst,
877 + origin_context_lens):
878 + chunk_context_len = chunk_size
879 + if cp_chunk_seq_len != 0:
880 + chunk_context_len = min(
881 + chunk_context_len, origin_context_len - chunk_size * chunk_idx)
882 + cp_target_rank = (chunk_context_len - 1) % cp_world_size
883 + cur_seq_len = 0
884 + for rank in range(cp_world_size):
885 + if rank > cp_target_rank and cp_chunk_seq_len:
886 + real_cp_chunk_seq_len = cp_chunk_seq_len - 1
887 + else:
888 + real_cp_chunk_seq_len = cp_chunk_seq_len
889 + if real_cp_chunk_seq_len:
890 + kv_c_segment = allgatered_kv_c_normed[rank * toks +
891 + src_token_idx:rank *
892 + toks + src_token_idx +
893 + real_cp_chunk_seq_len]
894 + k_pe_segment = allgatered_k_pe[rank * toks +
895 + src_token_idx:rank * toks +
896 + src_token_idx +
897 + real_cp_chunk_seq_len]
898 + kv_c_segments.append(kv_c_segment)
899 + k_pe_segments.append(k_pe_segment)
900 + cur_seq_len += real_cp_chunk_seq_len
901 + max_seq_len_check = max(max_seq_len_check, cur_seq_len)
902 + src_token_idx += cp_chunk_seq_len
903 + reorganized_kv_c_normed = torch.cat(kv_c_segments, dim=0)
904 + reorganized_k_pe = torch.cat(k_pe_segments, dim=0)
905 + assert reorganized_kv_c_normed.shape[0] == sum_seq_len
906 + assert reorganized_k_pe.shape[0] == sum_seq_len
907 + assert max_seq_len_check == max_seq_len
908 + return reorganized_kv_c_normed, reorganized_k_pe
909 +
910 +
911 class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
912 """
913 NOTE: Please read the comment at the top of the file before trying to
@@ -836,6 +987,8 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
987 self.vllm_flash_attn_version == 3
988 and current_platform.get_device_capability()[0] == 9)
989
990 + self.dcp_world_size: Optional[int] = None
991 +
992 def _flash_attn_varlen_diff_headdims(self,
993 q,
994 k,
@@ -1152,6 +1305,108 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
1305
1306 return output, output_lse
1307
1308 + def _context_parallel_compute_prefill_context(
1309 + self,
1310 + q: torch.Tensor,
1311 + kv_c_and_k_pe_cache: torch.Tensor,
1312 + attn_metadata: MLACommonMetadata,
1313 + k_scale: torch.Tensor,
1314 + dcp_world_size: int,
1315 + ):
1316 + assert k_scale is None, "DCP not support sacled kvcache now."
1317 + assert attn_metadata.prefill is not None
1318 + prefill_metadata = attn_metadata.prefill
1319 + assert prefill_metadata.chunked_context is not None
1320 + assert prefill_metadata.chunked_context.cp_chunk_seq_lens is not None
1321 + assert prefill_metadata.chunked_context.origin_context_lens is not None
1322 + assert prefill_metadata.chunked_context.cp_cu_seq_lens is not None
1323 + assert prefill_metadata.chunked_context.chunk_size is not None
1324 + assert prefill_metadata.chunked_context.cu_seq_lens_lst is not None
1325 +
1326 + output = None
1327 + iters = len(prefill_metadata.chunked_context.seq_tot)
1328 + workspace = prefill_metadata.chunked_context.workspace
1329 +
1330 + for i in range(iters):
1331 + toks = prefill_metadata.chunked_context.seq_tot[i]
1332 + ops.cp_gather_cache(
1333 + src_cache=kv_c_and_k_pe_cache,
1334 + dst=workspace,
1335 + block_table=prefill_metadata.block_table,
1336 + cu_seq_lens=prefill_metadata.chunked_context.cp_cu_seq_lens[i],
1337 + batch_size=attn_metadata.num_prefills,
1338 + seq_starts=prefill_metadata.chunked_context.starts[i],
1339 + )
1340 + # workspace
1341 + # |------- N tokens --------|--------- N*dcp_size tokens ----------|
1342 + # |<- use for loca_gather ->|<--------- use for allgather -------->|
1343 + allgather_offset = workspace.shape[0] // (dcp_world_size + 1)
1344 + assert allgather_offset * (dcp_world_size +
1345 + 1) == workspace.shape[0]
1346 + assert toks <= allgather_offset
1347 + local_gathered_kvcache = workspace[:toks]
1348 + cur_allgather_workspace = workspace[
1349 + allgather_offset:allgather_offset * (1 + dcp_world_size)]
1350 + assert toks * dcp_world_size <= cur_allgather_workspace.shape[0]
1351 + cur_allgather_kvcache = cur_allgather_workspace[:toks *
1352 + dcp_world_size]
1353 + cur_allgather_kvcache.copy_(get_dcp_group().all_gather(
1354 + local_gathered_kvcache, dim=0))
1355 + assert cur_allgather_kvcache.shape[
1356 + -1] == self.kv_lora_rank + self.qk_rope_head_dim
1357 + allgatered_kv_c_normed, allgatered_k_pe = \
1358 + cur_allgather_kvcache.unsqueeze(
1359 + 1).split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
1360 +
1361 + kv_c_normed, k_pe = reorg_kvcache(
1362 + allgatered_kv_c_normed,
1363 + allgatered_k_pe,
1364 + cp_chunk_seq_lens_lst=prefill_metadata.chunked_context.
1365 + cp_chunk_seq_lens[i],
1366 + origin_context_lens=prefill_metadata.chunked_context.
1367 + origin_context_lens,
1368 + cp_world_size=dcp_world_size,
1369 + sum_seq_len=prefill_metadata.chunked_context.cu_seq_lens_lst[i]
1370 + [-1],
1371 + max_seq_len=prefill_metadata.chunked_context.max_seq_lens[i],
1372 + chunk_size=prefill_metadata.chunked_context.chunk_size,
1373 + chunk_idx=i,
1374 + toks=toks)
1375 +
1376 + kv_nope = self.kv_b_proj(kv_c_normed)[0].view( \
1377 + -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)
1378 + k_nope, v = kv_nope\
1379 + .split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
1380 + k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))),
1381 + dim=-1)
1382 +
1383 + attn_output, attn_softmax_lse = self._run_prefill_context_chunk(
1384 + prefill=prefill_metadata,
1385 + chunk_idx=i,
1386 + q=q,
1387 + k=k,
1388 + v=v,
1389 + )
1390 +
1391 + if output is None:
1392 + output = attn_output
1393 + output_lse = attn_softmax_lse
1394 + else:
1395 + output_tmp = torch.empty_like(output)
1396 + output_lse_tmp = torch.empty_like(output_lse)
1397 + merge_attn_states(
1398 + output=output_tmp,
1399 + output_lse=output_lse_tmp,
1400 + prefix_output=output,
1401 + prefix_lse=output_lse,
1402 + suffix_output=attn_output,
1403 + suffix_lse=attn_softmax_lse,
1404 + )
1405 + output = output_tmp
1406 + output_lse = output_lse_tmp
1407 +
1408 + return output, output_lse
1409 +
1410 def _forward_prefill(
1411 self,
1412 q: torch.Tensor,
@@ -1162,6 +1417,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
1417 k_scale: torch.Tensor,
1418 ) -> torch.Tensor:
1419 assert attn_metadata.prefill is not None
1420 + assert self.dcp_world_size is not None
1421
1422 has_context = attn_metadata.prefill.chunked_context is not None
1423 kv_nope = self.kv_b_proj(kv_c_normed)[0].view(\
@@ -1181,8 +1437,15 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
1437
1438 if has_context:
1439 suffix_output, suffix_lse = output
1184 - context_output, context_lse = self._compute_prefill_context( \
1185 - q, kv_c_and_k_pe_cache, attn_metadata, k_scale)
1440 + if self.dcp_world_size > 1:
1441 + context_output, context_lse = \
1442 + self._context_parallel_compute_prefill_context(
1443 + q, kv_c_and_k_pe_cache, attn_metadata,
1444 + k_scale=None, dcp_world_size=self.dcp_world_size)
1445 + else:
1446 + context_output, context_lse = \
1447 + self._compute_prefill_context(
1448 + q, kv_c_and_k_pe_cache, attn_metadata, k_scale)
1449
1450 output = torch.empty_like(suffix_output)
1451 merge_attn_states(
@@ -1202,12 +1465,11 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
1465 @abstractmethod
1466 def _forward_decode(
1467 self,
1205 - ql_nope: torch.Tensor,
1206 - q_pe: torch.Tensor,
1468 + q: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],
1469 kv_c_and_k_pe_cache: torch.Tensor,
1470 attn_metadata: M,
1471 layer: AttentionLayer,
1210 - ) -> torch.Tensor:
1472 + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
1473 raise NotImplementedError
1474
1475 def forward(
@@ -1235,6 +1497,9 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
1497 # same expert outputs.
1498 return output.fill_(0)
1499
1500 + if self.dcp_world_size is None:
1501 + self.dcp_world_size = get_dcp_group().world_size
1502 +
1503 fp8_attention = self.kv_cache_dtype.startswith("fp8")
1504
1505 num_actual_toks = attn_metadata.num_actual_tokens
@@ -1313,7 +1578,26 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
1578 layer._q_scale)
1579 decode_q_pe = decode_q_pe.reshape(q_pe_shape)
1580
1316 - output[:num_decode_tokens] = self._forward_decode(
1317 - decode_ql_nope, decode_q_pe, kv_cache, attn_metadata, layer)
1318 -
1581 + decode_q = (decode_ql_nope, decode_q_pe)
1582 + if self.dcp_world_size > 1:
1583 + assert not fp8_attention, "DCP not support fp8 kvcache now."
1584 + # concatenate decode_ql_nope and decode_q_pe -> (B, N, L + P)
1585 + decode_q = torch.cat(decode_q, dim=-1)
1586 + # decode_q do allgather in head dim.
1587 + decode_q = get_dcp_group().all_gather(decode_q, dim=1)
1588 +
1589 + # call decode attn
1590 + attn_out, lse = self._forward_decode(decode_q, kv_cache,
1591 + attn_metadata, layer)
1592 +
1593 + # recorect dcp attn_out with lse.
1594 + if self.dcp_world_size > 1:
1595 + assert lse is not None, (
1596 + "For a mla backend want to enable"
1597 + "DCP, it is mandatory that the corresponding decode attn"
1598 + "kernel return the softmax lse.")
1599 + attn_out = cp_lse_ag_out_rs(attn_out, lse, get_dcp_group())
1600 +
1601 + # v_up projection
1602 + output[:num_decode_tokens] = self._v_up_proj(attn_out)
1603 return output_padded
vllm/v1/attention/backends/mla/cutlass_mla.py
+11 -7
@@ -232,7 +232,7 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
232 self._workspace.get_buf(),
233 self.scale, self._num_kv_splits)
234
235 - return self._v_up_proj(o)
235 + return o
236
237 # TODO: Currently we leave it here only for backup in case something is
238 # wrong with the new SM100 CUTLASS MLA kernel
@@ -265,21 +265,25 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
265 attn_metadata.decode.seq_lens,
266 attn_metadata.decode.block_table, self.scale)
267
268 - return self._v_up_proj(o)
268 + return o
269
270 def _forward_decode(
271 self,
272 - q_nope: torch.Tensor,
273 - q_pe: torch.Tensor,
272 + q: torch.Tensor,
273 kv_c_and_k_pe_cache: torch.Tensor,
274 attn_metadata: MLACommonMetadata,
275 layer: AttentionLayer,
277 - ) -> torch.Tensor:
276 + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
277 + if type(q) is tuple:
278 + q_nope, q_pe = q
279 + else:
280 + q_nope, q_pe = torch.split(
281 + q, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
282 if self._use_old_cutlass_mla:
283 # TODO: Remove the old cutlass MLA kernel after more extensive
284 # testing
285 return self._old_forward_decode(q_nope, q_pe, kv_c_and_k_pe_cache,
282 - attn_metadata)
286 + attn_metadata), None
287
288 return self._sm100_forward_decode(q_nope, q_pe, kv_c_and_k_pe_cache,
285 - attn_metadata)
289 + attn_metadata), None
vllm/v1/attention/backends/mla/flashattn_mla.py
+9 -4
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from dataclasses import dataclass
5 -from typing import ClassVar, Optional
5 +from typing import ClassVar, Optional, Union
6
7 import torch
8
@@ -154,15 +154,20 @@ class FlashAttnMLAImpl(MLACommonImpl[FlashAttnMLAMetadata]):
154
155 def _forward_decode(
156 self,
157 - q_nope: torch.Tensor,
158 - q_pe: torch.Tensor,
157 + q: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],
158 kv_c_and_k_pe_cache: torch.Tensor,
159 attn_metadata: FlashAttnMLAMetadata,
160 layer: AttentionLayer,
162 - ) -> torch.Tensor:
161 + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
162 assert kv_c_and_k_pe_cache.numel() > 0
163 assert attn_metadata.decode is not None
164
165 + if type(q) is tuple:
166 + q_nope, q_pe = q
167 + else:
168 + q_nope, q_pe = torch.split(
169 + q, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
170 +
171 if self.kv_cache_dtype.startswith("fp8"):
172 raise NotImplementedError(
173 "FP8 FlashAttention MLA not yet supported")
vllm/v1/attention/backends/mla/flashmla.py
+9 -9
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from dataclasses import dataclass
5 -from typing import ClassVar, Optional
5 +from typing import ClassVar, Optional, Union
6
7 import torch
8
@@ -169,20 +169,20 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
169
170 def _forward_decode(
171 self,
172 - q_nope: torch.Tensor,
173 - q_pe: torch.Tensor,
172 + q: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],
173 kv_c_and_k_pe_cache: torch.Tensor,
174 attn_metadata: FlashMLAMetadata,
175 layer: AttentionLayer,
177 - ) -> torch.Tensor:
176 + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
177 assert kv_c_and_k_pe_cache.numel() > 0
178 assert attn_metadata.decode is not None
179
181 - q = torch.cat([q_nope, q_pe], dim=-1)\
182 - .unsqueeze(1) # Add seqlen dim of 1 (decode)
180 + if type(q) is tuple:
181 + q = torch.cat(q, dim=-1)
182
184 - o, _ = flash_mla_with_kvcache(
185 - q=q,
183 + assert isinstance(q, torch.Tensor)
184 + o, lse = flash_mla_with_kvcache(
185 + q=q.unsqueeze(1), # Add seqlen dim of 1 (decode)
186 k_cache=kv_c_and_k_pe_cache.unsqueeze(-2), # Add head dim of 1
187 block_table=attn_metadata.decode.block_table,
188 cache_seqlens=attn_metadata.decode.seq_lens,
@@ -196,4 +196,4 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
196 descale_k=layer._k_scale.reshape(1),
197 )
198
199 - return self._v_up_proj(o)
199 + return o, lse
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+8 -7
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from dataclasses import dataclass
5 -from typing import ClassVar, Optional
5 +from typing import ClassVar, Optional, Union
6
7 import torch
8
@@ -220,18 +220,19 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
220
221 def _forward_decode(
222 self,
223 - q_nope: torch.Tensor,
224 - q_pe: torch.Tensor,
223 + q: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],
224 kv_c_and_k_pe_cache: torch.Tensor,
225 attn_metadata: AiterMLAMetadata,
226 layer: AttentionLayer,
228 - ) -> torch.Tensor:
227 + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
228 assert kv_c_and_k_pe_cache.numel() > 0
229 assert attn_metadata.decode is not None
230
232 - B = q_nope.shape[0]
231 + if type(q) is tuple:
232 + q = torch.cat(q, dim=-1)
233
234 - q = torch.cat([q_nope, q_pe], dim=-1)
234 + assert isinstance(q, torch.Tensor)
235 + B = q.shape[0]
236 o = torch.zeros(B,
237 self.num_heads,
238 self.kv_lora_rank,
@@ -249,4 +250,4 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
250 attn_metadata.decode.paged_kv_indices,
251 attn_metadata.decode.paged_kv_last_page_len)
252
252 - return self._v_up_proj(o)
253 + return o, None
vllm/v1/attention/backends/mla/triton_mla.py
+8 -7
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from typing import Optional
4 +from typing import Optional, Union
5
6 import torch
7
@@ -123,21 +123,22 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
123
124 def _forward_decode(
125 self,
126 - q_nope: torch.Tensor,
127 - q_pe: torch.Tensor,
126 + q: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],
127 kv_c_and_k_pe_cache: torch.Tensor,
128 attn_metadata: MLACommonMetadata,
129 layer: AttentionLayer,
131 - ) -> torch.Tensor:
130 + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
131 assert kv_c_and_k_pe_cache.numel() > 0
132 assert attn_metadata.decode is not None
133
134 if self.kv_cache_dtype.startswith("fp8"):
135 raise NotImplementedError("FP8 Triton MLA not yet supported")
136
138 - B = q_nope.shape[0]
137 + if type(q) is tuple:
138 + q = torch.cat(q, dim=-1)
139
140 - q = torch.cat([q_nope, q_pe], dim=-1)
140 + assert isinstance(q, torch.Tensor)
141 + B = q.shape[0]
142 o = torch.zeros(B,
143 self.num_heads,
144 self.kv_lora_rank,
@@ -171,4 +172,4 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
172 attn_metadata.decode.seq_lens, attn_logits,
173 num_kv_splits, self.scale, PAGE_SIZE)
174
174 - return self._v_up_proj(o)
175 + return o, None
vllm/v1/core/kv_cache_coordinator.py
+50 -20
@@ -24,6 +24,7 @@ class KVCacheCoordinator(ABC):
24 use_eagle: bool,
25 enable_caching: bool,
26 enable_kv_cache_events: bool,
27 + dcp_world_size: int,
28 ):
29 self.kv_cache_config = kv_cache_config
30 self.max_model_len = max_model_len
@@ -39,6 +40,7 @@ class KVCacheCoordinator(ABC):
40 kv_cache_spec=kv_cache_group.kv_cache_spec,
41 block_pool=self.block_pool,
42 kv_cache_group_id=i,
43 + dcp_world_size=dcp_world_size,
44 ) for i, kv_cache_group in enumerate(
45 self.kv_cache_config.kv_cache_groups))
46
@@ -197,9 +199,14 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):
199 """
200
201 def __init__(self, kv_cache_config: KVCacheConfig, max_model_len: int,
200 - use_eagle: bool, enable_kv_cache_events: bool):
201 - super().__init__(kv_cache_config, max_model_len, use_eagle, False,
202 - enable_kv_cache_events)
202 + use_eagle: bool, enable_kv_cache_events: bool,
203 + dcp_world_size: int):
204 + super().__init__(kv_cache_config,
205 + max_model_len,
206 + use_eagle,
207 + False,
208 + enable_kv_cache_events,
209 + dcp_world_size=dcp_world_size)
210 self.num_single_type_manager = len(self.single_type_managers)
211
212 def get_num_common_prefix_blocks(self, request_id: str,
@@ -225,12 +232,19 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
232
233 def __init__(self, kv_cache_config: KVCacheConfig, max_model_len: int,
234 use_eagle: bool, enable_caching: bool,
228 - enable_kv_cache_events: bool):
229 - super().__init__(kv_cache_config, max_model_len, use_eagle,
230 - enable_caching, enable_kv_cache_events)
235 + enable_kv_cache_events: bool, dcp_world_size: int):
236 + super().__init__(kv_cache_config,
237 + max_model_len,
238 + use_eagle,
239 + enable_caching,
240 + enable_kv_cache_events,
241 + dcp_world_size=dcp_world_size)
242 self.kv_cache_spec = self.kv_cache_config.kv_cache_groups[
243 0].kv_cache_spec
244 self.block_size = self.kv_cache_spec.block_size
245 + self.dcp_world_size = dcp_world_size
246 + if dcp_world_size > 1:
247 + self.block_size *= dcp_world_size
248 assert len(self.kv_cache_config.kv_cache_groups) == 1, (
249 "UnitaryKVCacheCoordinator assumes only one kv cache group")
250
@@ -246,6 +260,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
260 block_pool=self.block_pool,
261 kv_cache_spec=self.kv_cache_spec,
262 use_eagle=self.use_eagle,
263 + dcp_world_size=self.dcp_world_size,
264 )
265 return hit_blocks, len(hit_blocks[0]) * self.block_size
266
@@ -261,9 +276,14 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
276
277 def __init__(self, kv_cache_config: KVCacheConfig, max_model_len: int,
278 use_eagle: bool, enable_caching: bool,
264 - enable_kv_cache_events: bool):
265 - super().__init__(kv_cache_config, max_model_len, use_eagle,
266 - enable_caching, enable_kv_cache_events)
279 + enable_kv_cache_events: bool, dcp_world_size: int):
280 + super().__init__(kv_cache_config,
281 + max_model_len,
282 + use_eagle,
283 + enable_caching,
284 + enable_kv_cache_events,
285 + dcp_world_size=dcp_world_size)
286 + assert dcp_world_size == 1, "DCP not support hybrid attn now."
287 self.verify_and_split_kv_cache_groups()
288
289 def verify_and_split_kv_cache_groups(self) -> None:
@@ -394,17 +414,27 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
414 return hit_blocks, hit_length
415
416
397 -def get_kv_cache_coordinator(
398 - kv_cache_config: KVCacheConfig, max_model_len: int, use_eagle: bool,
399 - enable_caching: bool,
400 - enable_kv_cache_events: bool) -> KVCacheCoordinator:
417 +def get_kv_cache_coordinator(kv_cache_config: KVCacheConfig,
418 + max_model_len: int, use_eagle: bool,
419 + enable_caching: bool,
420 + enable_kv_cache_events: bool,
421 + dcp_world_size: int) -> KVCacheCoordinator:
422 if not enable_caching:
402 - return KVCacheCoordinatorNoPrefixCache(kv_cache_config, max_model_len,
423 + return KVCacheCoordinatorNoPrefixCache(kv_cache_config,
424 + max_model_len,
425 use_eagle,
404 - enable_kv_cache_events)
426 + enable_kv_cache_events,
427 + dcp_world_size=dcp_world_size)
428 if len(kv_cache_config.kv_cache_groups) == 1:
406 - return UnitaryKVCacheCoordinator(kv_cache_config, max_model_len,
407 - use_eagle, enable_caching,
408 - enable_kv_cache_events)
409 - return HybridKVCacheCoordinator(kv_cache_config, max_model_len, use_eagle,
410 - enable_caching, enable_kv_cache_events)
429 + return UnitaryKVCacheCoordinator(kv_cache_config,
430 + max_model_len,
431 + use_eagle,
432 + enable_caching,
433 + enable_kv_cache_events,
434 + dcp_world_size=dcp_world_size)
435 + return HybridKVCacheCoordinator(kv_cache_config,
436 + max_model_len,
437 + use_eagle,
438 + enable_caching,
439 + enable_kv_cache_events,
440 + dcp_world_size=dcp_world_size)
vllm/v1/core/kv_cache_manager.py
+9
@@ -91,6 +91,7 @@ class KVCacheManager:
91 use_eagle: bool = False,
92 log_stats: bool = False,
93 enable_kv_cache_events: bool = False,
94 + dcp_world_size: int = 1,
95 ) -> None:
96 self.max_model_len = max_model_len
97
@@ -109,12 +110,20 @@ class KVCacheManager:
110 self.block_size = kv_cache_config.kv_cache_groups[
111 0].kv_cache_spec.block_size
112
113 + if dcp_world_size > 1:
114 + assert len(kv_cache_config.kv_cache_groups) == 1
115 + # Note(hc): need revisit. When both DCP and any future
116 + # PCP are enabled, the block_size may need to be scaled
117 + # by a factor of dcp_size × pcp_size?
118 + self.block_size *= dcp_world_size
119 +
120 self.coordinator = get_kv_cache_coordinator(
121 kv_cache_config=kv_cache_config,
122 max_model_len=self.max_model_len,
123 use_eagle=self.use_eagle,
124 enable_caching=self.enable_caching,
125 enable_kv_cache_events=enable_kv_cache_events,
126 + dcp_world_size=dcp_world_size,
127 )
128 self.num_kv_cache_groups = len(kv_cache_config.kv_cache_groups)
129 self.block_pool = self.coordinator.block_pool
vllm/v1/core/kv_cache_utils.py
+6
@@ -846,6 +846,12 @@ def _get_kv_cache_config_uniform_type(vllm_config: VllmConfig,
846 )
847
848 num_tokens = num_blocks * vllm_config.cache_config.block_size
849 + if vllm_config.parallel_config.decode_context_parallel_size > 1:
850 + num_tokens *= vllm_config.parallel_config.decode_context_parallel_size
851 + logger.info(
852 + "Multiplying the GPU KV cache size by the dcp_world_size %d.",
853 + vllm_config.parallel_config.decode_context_parallel_size)
854 +
855 num_tokens_str = f"{num_tokens:,}"
856 logger.info("GPU KV cache size: %s tokens", num_tokens_str)
857 max_model_len_str = f"{vllm_config.model_config.max_model_len:,}"
vllm/v1/core/sched/scheduler.py
+10
@@ -100,6 +100,15 @@ class Scheduler(SchedulerInterface):
100
101 self.block_size = self.cache_config.block_size
102
103 + self.dcp_world_size = \
104 + vllm_config.parallel_config.decode_context_parallel_size
105 + # Note(hc): The scheduler’s block_size must be multiplied
106 + # by dcp_world_size, since block hashes are computed on the
107 + # original full token sequence at a granularity of
108 + # original_block_size × dcp_world_size.
109 + if self.dcp_world_size > 1:
110 + self.block_size *= self.dcp_world_size
111 +
112 # req_id -> Request
113 self.requests: dict[str, Request] = {}
114 # Scheduling policy
@@ -161,6 +170,7 @@ class Scheduler(SchedulerInterface):
170 use_eagle=self.use_eagle,
171 log_stats=self.log_stats,
172 enable_kv_cache_events=self.enable_kv_cache_events,
173 + dcp_world_size=self.dcp_world_size,
174 )
175 self.use_pp = self.parallel_config.pipeline_parallel_size > 1
176
vllm/v1/core/single_type_kv_cache_manager.py
+17 -2
@@ -25,6 +25,7 @@ class SingleTypeKVCacheManager(ABC):
25 kv_cache_spec: KVCacheSpec,
26 block_pool: BlockPool,
27 kv_cache_group_id: int,
28 + dcp_world_size: int = 1,
29 ) -> None:
30 """
31 Initializes the SingleTypeKVCacheManager.
@@ -33,8 +34,10 @@ class SingleTypeKVCacheManager(ABC):
34 block_pool: The block pool.
35 kv_cache_group_id: The id of the kv cache group of this manager.
36 """
36 -
37 self.block_size = kv_cache_spec.block_size
38 + self.dcp_world_size = dcp_world_size
39 + if self.dcp_world_size > 1:
40 + self.block_size *= dcp_world_size
41 self.kv_cache_spec = kv_cache_spec
42 self.block_pool = block_pool
43
@@ -196,6 +199,7 @@ class SingleTypeKVCacheManager(ABC):
199 block_pool: BlockPool,
200 kv_cache_spec: KVCacheSpec,
201 use_eagle: bool,
202 + dcp_world_size: int = 1,
203 ) -> tuple[list[KVCacheBlock], ...]:
204 """
205 Get the longest cache hit prefix of the blocks that is not longer than
@@ -253,6 +257,7 @@ class FullAttentionManager(SingleTypeKVCacheManager):
257 block_pool: BlockPool,
258 kv_cache_spec: KVCacheSpec,
259 use_eagle: bool,
260 + dcp_world_size: int = 1,
261 ) -> tuple[list[KVCacheBlock], ...]:
262 assert isinstance(
263 kv_cache_spec, (FullAttentionSpec, ChunkedLocalAttentionSpec)
@@ -260,7 +265,10 @@ class FullAttentionManager(SingleTypeKVCacheManager):
265 "and chunked local attention groups"
266 computed_blocks: tuple[list[KVCacheBlock], ...] = tuple(
267 [] for _ in range(len(kv_cache_group_ids)))
263 - max_num_blocks = max_length // kv_cache_spec.block_size
268 + block_size = kv_cache_spec.block_size
269 + if dcp_world_size > 1:
270 + block_size *= dcp_world_size
271 + max_num_blocks = max_length // block_size
272 for block_hash in itertools.islice(block_hashes, max_num_blocks):
273 # block_hashes is a chain of block hashes. If a block hash is not
274 # in the cached_block_hash_to_id, the following block hashes are
@@ -310,9 +318,11 @@ class SlidingWindowManager(SingleTypeKVCacheManager):
318 block_pool: BlockPool,
319 kv_cache_spec: KVCacheSpec,
320 use_eagle: bool,
321 + dcp_world_size: int = 1,
322 ) -> tuple[list[KVCacheBlock], ...]:
323 assert isinstance(kv_cache_spec, SlidingWindowSpec), (
324 "SlidingWindowManager can only be used for sliding window groups")
325 + assert dcp_world_size == 1, "DCP not support sliding window attn now."
326
327 # The number of contiguous blocks needed for prefix cache hit.
328 # -1 since the input token itself is also included in the window
@@ -408,6 +418,7 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):
418 block_pool: BlockPool,
419 kv_cache_spec: KVCacheSpec,
420 use_eagle: bool,
421 + dcp_world_size: int = 1,
422 ) -> tuple[list[KVCacheBlock], ...]:
423 """
424 For chunked local attention, we need to find the longest cache hit
@@ -445,6 +456,7 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):
456 "chunked local attention groups")
457 assert use_eagle is False, ("Hybrid KV cache is not supported for " +
458 "eagle + chunked local attention.")
459 + assert dcp_world_size == 1, "DCP not support chunked local attn now."
460 max_num_blocks = max_length // kv_cache_spec.block_size
461 if max_length > 0:
462 local_attention_start_idx = (max_length //
@@ -525,10 +537,12 @@ class MambaManager(SingleTypeKVCacheManager):
537 block_pool: BlockPool,
538 kv_cache_spec: KVCacheSpec,
539 use_eagle: bool,
540 + dcp_world_size: int = 1,
541 ) -> tuple[list[KVCacheBlock], ...]:
542 assert isinstance(
543 kv_cache_spec,
544 MambaSpec), ("MambaManager can only be used for mamba groups")
545 + assert dcp_world_size == 1, "DCP not support mamba now."
546 # Prefix caching is not supported for mamba now. Always return empty
547 # list.
548 computed_blocks: tuple[list[KVCacheBlock], ...] = tuple(
@@ -583,6 +597,7 @@ class CrossAttentionManager(SingleTypeKVCacheManager):
597 block_pool: BlockPool,
598 kv_cache_spec: KVCacheSpec,
599 use_eagle: bool,
600 + dcp_world_size: int = 1,
601 ) -> tuple[list[KVCacheBlock], ...]:
602 assert isinstance(kv_cache_spec, CrossAttentionSpec), (
603 "CrossAttentionManager can only be used for cross-attention groups"
vllm/v1/kv_cache_interface.py
+8
@@ -86,6 +86,12 @@ class FullAttentionSpec(AttentionSpec):
86
87 def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
88 max_model_len = vllm_config.model_config.max_model_len
89 + dcp_world_size = \
90 + vllm_config.parallel_config.decode_context_parallel_size
91 + # Note(hc): each dcp rank only need save
92 + # (max_model_len//dcp_world_size) tokens locally.
93 + if dcp_world_size > 1:
94 + max_model_len = cdiv(max_model_len, dcp_world_size)
95 return cdiv(max_model_len, self.block_size) * self.page_size_bytes
96
97 @classmethod
@@ -162,6 +168,8 @@ class SlidingWindowSpec(AttentionSpec):
168 assert not self.use_mla, "MLA is not supported for sliding window"
169
170 def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
171 + assert vllm_config.parallel_config.decode_context_parallel_size == 1, \
172 + "DCP not support sliding window."
173 max_model_len = vllm_config.model_config.max_model_len
174 max_num_batched_tokens = (
175 vllm_config.scheduler_config.max_num_batched_tokens)
vllm/v1/worker/block_table.py
+50 -9
@@ -4,6 +4,7 @@
4 import numpy as np
5 import torch
6
7 +from vllm.distributed import get_dcp_group
8 from vllm.logger import init_logger
9 from vllm.utils import cdiv
10
@@ -50,6 +51,13 @@ class BlockTable:
51 self.slot_mapping = torch.zeros(self.max_num_batched_tokens,
52 dtype=torch.int64,
53 device=self.device)
54 + try:
55 + self.dcp_world_size = get_dcp_group().world_size
56 + self.dcp_rank = get_dcp_group().rank_in_group
57 + except AssertionError:
58 + # DCP might not be initialized in testing
59 + self.dcp_world_size = 1
60 + self.dcp_rank = 0
61
62 def append_row(
63 self,
@@ -89,13 +97,36 @@ class BlockTable:
97 # NOTE(woosuk): We can't simply use `token_indices // block_size`
98 # here because M (max_model_len) is not necessarily divisible by
99 # block_size.
92 - block_table_indices = (req_indices * self.max_num_blocks_per_req +
93 - positions // self.block_size)
94 - block_numbers = self.block_table_np.ravel()[block_table_indices]
95 - block_offsets = positions % self.block_size
96 - np.add(block_numbers * self.block_size,
97 - block_offsets,
98 - out=self.slot_mapping_np[:req_indices.shape[0]])
100 + if self.dcp_world_size > 1:
101 + # Note(hc): The DCP implement store kvcache with a interleave
102 + # style, the kvcache for the token whose token_idx is i is
103 + # always stored on the GPU whose dcp_rank equals i % cp_world_size:
104 +
105 + # Use a "virtual block" which equals to world_size * block_size
106 + # for block_table_indices calculation.
107 + virtual_block_size = self.block_size * self.dcp_world_size
108 + block_table_indices = (req_indices * self.max_num_blocks_per_req +
109 + positions // virtual_block_size)
110 + block_numbers = self.block_table_np.ravel()[block_table_indices]
111 + # Use virtual_block_size for mask calculation, which marks local
112 + # tokens.
113 + virtual_block_offsets = positions % virtual_block_size
114 + mask = virtual_block_offsets % self.dcp_world_size == self.dcp_rank
115 + # Calcuate local block_offsets
116 + block_offsets = virtual_block_offsets // self.dcp_world_size
117 + # Calcuate slot_mapping
118 + slot_mapping = block_numbers * self.block_size + block_offsets
119 + # Write final slots, use -1 for not-local
120 + self.slot_mapping_np[:req_indices.shape[0]] = np.where(
121 + mask, slot_mapping, -1)
122 + else:
123 + block_table_indices = (req_indices * self.max_num_blocks_per_req +
124 + positions // self.block_size)
125 + block_numbers = self.block_table_np.ravel()[block_table_indices]
126 + block_offsets = positions % self.block_size
127 + np.add(block_numbers * self.block_size,
128 + block_offsets,
129 + out=self.slot_mapping_np[:req_indices.shape[0]])
130
131 def commit_block_table(self, num_reqs: int) -> None:
132 self.block_table[:num_reqs].copy_(self.block_table_cpu[:num_reqs],
@@ -128,9 +159,19 @@ class MultiGroupBlockTable:
159 def __init__(self, max_num_reqs: int, max_model_len: int,
160 max_num_batched_tokens: int, pin_memory: bool,
161 device: torch.device, block_sizes: list[int]) -> None:
162 + # Note(hc): each dcp rank only store
163 + # (max_model_len//dcp_world_size) tokens in kvcache,
164 + # so the block_size which used for calc max_num_blocks_per_req
165 + # must be multiplied by dcp_world_size.
166 + try:
167 + dcp_world_size = get_dcp_group().world_size
168 + except AssertionError:
169 + # DCP might not be initialized in testing
170 + dcp_world_size = 1
171 +
172 self.block_tables = [
132 - BlockTable(block_size, max_num_reqs, cdiv(max_model_len,
133 - block_size),
173 + BlockTable(block_size, max_num_reqs,
174 + cdiv(max_model_len, block_size * dcp_world_size),
175 max_num_batched_tokens, pin_memory, device)
176 for block_size in block_sizes
177 ]
vllm/v1/worker/gpu_model_runner.py
+11
@@ -56,6 +56,7 @@ from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,
56 GiB_bytes, LazyLoader, cdiv, check_use_alibi,
57 get_dtype_size, is_pin_memory_available, round_up,
58 supports_dynamo)
59 +from vllm.v1.attention.backends.mla.flashmla import FlashMLABackend
60 from vllm.v1.attention.backends.utils import (
61 AttentionCGSupport, AttentionMetadataBuilder, CommonAttentionMetadata,
62 create_fast_prefill_custom_backend,
@@ -187,6 +188,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
188 model_config.is_multimodal_raw_input_only_model)
189
190 self.max_model_len = model_config.max_model_len
191 + self.dcp_world_size = self.parallel_config.decode_context_parallel_size
192 self.max_num_tokens = scheduler_config.max_num_batched_tokens
193 self.max_num_reqs = scheduler_config.max_num_seqs
194
@@ -428,6 +430,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
430 return
431
432 if self.reorder_batch_threshold is not None:
433 + if self.dcp_world_size > 1:
434 + assert self.reorder_batch_threshold == 1, \
435 + "DCP not support reorder_batch_threshold > 1 now."
436 reorder_batch_to_split_decodes_and_prefills(
437 self.input_batch,
438 scheduler_output,
@@ -3305,6 +3310,12 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3310 get_kv_transfer_group().set_host_xfer_buffer_ops(
3311 copy_kv_blocks)
3312
3313 + if self.dcp_world_size > 1:
3314 + assert self.attn_groups[0][0].backend is FlashMLABackend, (
3315 + "DCP only support flashmla now."
3316 + "For a mla backend want to enable DCP, it is mandatory that the"
3317 + "corresponding decode attn kernel return the softmax lse.")
3318 +
3319 def may_add_encoder_only_layers_to_kv_cache_config(self) -> None:
3320 """
3321 Add encoder-only layers to the KV cache config.
vllm/v1/worker/gpu_worker.py
+4 -2
@@ -616,7 +616,9 @@ def init_worker_distributed_environment(
616 init_distributed_environment(parallel_config.world_size, rank,
617 distributed_init_method, local_rank, backend)
618
619 - ensure_model_parallel_initialized(parallel_config.tensor_parallel_size,
620 - parallel_config.pipeline_parallel_size)
619 + ensure_model_parallel_initialized(
620 + parallel_config.tensor_parallel_size,
621 + parallel_config.pipeline_parallel_size,
622 + parallel_config.decode_context_parallel_size)
623
624 ensure_kv_transfer_initialized(vllm_config)
vllm/worker/worker.py
+4 -2
@@ -539,8 +539,10 @@ def init_worker_distributed_environment(
539 init_distributed_environment(parallel_config.world_size, rank,
540 distributed_init_method, local_rank,
541 current_platform.dist_backend)
542 - ensure_model_parallel_initialized(parallel_config.tensor_parallel_size,
543 - parallel_config.pipeline_parallel_size)
542 + ensure_model_parallel_initialized(
543 + parallel_config.tensor_parallel_size,
544 + parallel_config.pipeline_parallel_size,
545 + parallel_config.decode_context_parallel_size)
546
547 ensure_kv_transfer_initialized(vllm_config)
548