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 8, 03:45 UTC
[Feature] Decode Context Parallel support for GPU model runner v2 (#34179)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Feb 17, 2026 at 19:27 UTC
ab33d2a629be6eca2dd946b1628af4d23d39c547
5 files changed
+117
-3
vllm/v1/worker/gpu/attn_utils.py
+28
@@ -12,6 +12,7 @@ from vllm.v1.attention.backend import (
12
AttentionMetadataBuilder,
13
CommonAttentionMetadata,
14
)
15
+from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
16
from vllm.v1.kv_cache_interface import (
17
AttentionSpec,
18
KVCacheConfig,
@@ -143,6 +144,28 @@ def build_slot_mappings_by_layer(
144
return slot_mappings_by_layer
145
146
147
+def prepare_dcp_local_seq_lens(
148
+ dcp_local_seq_lens: torch.Tensor,
149
+ seq_lens: torch.Tensor,
150
+ num_reqs: int,
151
+ dcp_size: int,
152
+ dcp_rank: int,
153
+ cp_kv_cache_interleave_size: int,
154
+) -> None:
155
+ """Populate the persistent DCP local seq_lens buffer (CUDA graph safe)."""
156
+ if dcp_size <= 1:
157
+ return
158
+
159
+ local_seq_lens = get_dcp_local_seq_lens(
160
+ seq_lens[:num_reqs],
161
+ dcp_size=dcp_size,
162
+ dcp_rank=dcp_rank,
163
+ cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
164
+ )
165
+ dcp_local_seq_lens[:num_reqs].copy_(local_seq_lens, non_blocking=True)
166
+ dcp_local_seq_lens[num_reqs:].zero_()
167
+
168
+
169
def build_attn_metadata(
170
attn_metadata_builders: list[AttentionMetadataBuilder],
171
num_reqs: int,
@@ -155,9 +178,13 @@ def build_attn_metadata(
178
block_tables: Sequence[torch.Tensor],
179
slot_mappings: torch.Tensor,
180
kv_cache_config: KVCacheConfig,
181
+ dcp_local_seq_lens: torch.Tensor | None = None,
182
) -> dict[str, Any]:
183
seq_lens = seq_lens[:num_reqs]
184
185
+ if dcp_local_seq_lens is not None:
186
+ dcp_local_seq_lens = dcp_local_seq_lens[:num_reqs]
187
+
188
attn_metadata: dict[str, Any] = {}
189
kv_cache_groups = kv_cache_config.kv_cache_groups
190
for i, kv_cache_spec in enumerate(kv_cache_groups):
@@ -175,6 +202,7 @@ def build_attn_metadata(
202
block_table_tensor=block_table,
203
slot_mapping=slot_mapping,
204
causal=True,
205
+ dcp_local_seq_lens=dcp_local_seq_lens,
206
)
207
208
attn_metadata_builder = attn_metadata_builders[i]
vllm/v1/worker/gpu/block_table.py
+43
-3
@@ -4,6 +4,7 @@ from collections.abc import Iterable
4
5
import torch
6
7
+from vllm.distributed import get_dcp_group
8
from vllm.triton_utils import tl, triton
9
from vllm.utils.math_utils import cdiv
10
from vllm.v1.attention.backends.utils import PAD_SLOT_ID
@@ -18,19 +19,36 @@ class BlockTables:
19
max_num_batched_tokens: int,
20
max_model_len: int,
21
device: torch.device,
22
+ cp_kv_cache_interleave_size: int = 1,
23
):
24
self.block_sizes = block_sizes
25
self.max_num_reqs = max_num_reqs
26
self.max_num_batched_tokens = max_num_batched_tokens
27
self.max_model_len = max_model_len
28
self.device = device
29
+ assert cp_kv_cache_interleave_size >= 1
30
+ self.cp_kv_cache_interleave_size = cp_kv_cache_interleave_size
31
+
32
+ try:
33
+ dcp = get_dcp_group()
34
+ self.dcp_world_size, self.dcp_rank = dcp.world_size, dcp.rank_in_group
35
+ except AssertionError:
36
+ self.dcp_world_size, self.dcp_rank = 1, 0
37
+ # TODO(wentao): PCP supprot
38
+ self.total_cp_world_size = self.dcp_world_size
39
+ self.total_cp_rank = self.dcp_rank
40
41
self.num_kv_cache_groups = len(self.block_sizes)
42
# num_kv_cache_groups x [max_num_reqs, max_num_blocks]
43
self.block_tables: list[StagedWriteTensor] = []
44
for i in range(self.num_kv_cache_groups):
45
block_size = self.block_sizes[i]
33
- max_num_blocks = cdiv(self.max_model_len, block_size)
46
+ # with DCP, a request's KV is sharded across
47
+ # ranks, so one physical block on this rank
48
+ # corresponds to `block_size * total_cp_world_size`
49
+ # tokens in the global (unsharded) sequence.
50
+ virtual_block_size = block_size * self.total_cp_world_size
51
+ max_num_blocks = cdiv(self.max_model_len, virtual_block_size)
52
block_table = StagedWriteTensor(
53
(self.max_num_reqs, max_num_blocks),
54
dtype=torch.int32,
@@ -131,6 +149,9 @@ class BlockTables:
149
self.block_sizes_tensor,
150
self.slot_mappings,
151
self.slot_mappings.stride(0),
152
+ TOTAL_CP_WORLD_SIZE=self.total_cp_world_size,
153
+ TOTAL_CP_RANK=self.total_cp_rank,
154
+ CP_KV_CACHE_INTERLEAVE_SIZE=self.cp_kv_cache_interleave_size,
155
PAD_ID=PAD_SLOT_ID,
156
TRITON_BLOCK_SIZE=1024, # type: ignore
157
)
@@ -183,6 +204,9 @@ def _compute_slot_mappings_kernel(
204
block_sizes, # [num_kv_cache_groups]
205
slot_mappings_ptr, # [num_kv_cache_groups, max_num_tokens]
206
slot_mappings_stride,
207
+ TOTAL_CP_WORLD_SIZE: tl.constexpr,
208
+ TOTAL_CP_RANK: tl.constexpr,
209
+ CP_KV_CACHE_INTERLEAVE_SIZE: tl.constexpr,
210
PAD_ID: tl.constexpr,
211
TRITON_BLOCK_SIZE: tl.constexpr,
212
):
@@ -201,6 +225,7 @@ def _compute_slot_mappings_kernel(
225
block_table_ptr = _load_ptr(block_table_ptrs + group_id, tl.int32)
226
block_table_stride = tl.load(block_table_strides + group_id)
227
block_size = tl.load(block_sizes + group_id)
228
+ virtual_block_size = block_size * TOTAL_CP_WORLD_SIZE
229
230
req_state_idx = tl.load(idx_mapping + batch_idx)
231
start_idx = tl.load(query_start_loc + batch_idx)
@@ -208,11 +233,26 @@ def _compute_slot_mappings_kernel(
233
for i in range(start_idx, end_idx, TRITON_BLOCK_SIZE):
234
offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
235
positions = tl.load(pos + offset, mask=offset < end_idx, other=0)
211
- block_indices = positions // block_size
236
+ block_indices = positions // virtual_block_size
237
block_numbers = tl.load(
238
block_table_ptr + req_state_idx * block_table_stride + block_indices
239
)
215
- slot_ids = block_numbers * block_size + positions % block_size
240
+ virtual_block_offsets = positions - block_indices * virtual_block_size
241
+
242
+ # determine whether the token is stored on this CP rank.
243
+ is_local = (
244
+ virtual_block_offsets // CP_KV_CACHE_INTERLEAVE_SIZE
245
+ ) % TOTAL_CP_WORLD_SIZE == TOTAL_CP_RANK
246
+ # mapping virture block offsets to local block offsets.
247
+ local_block_offsets = (
248
+ virtual_block_offsets // (TOTAL_CP_WORLD_SIZE * CP_KV_CACHE_INTERLEAVE_SIZE)
249
+ ) * CP_KV_CACHE_INTERLEAVE_SIZE + (
250
+ virtual_block_offsets % CP_KV_CACHE_INTERLEAVE_SIZE
251
+ )
252
+
253
+ # physical slot index
254
+ slot_ids = block_numbers * block_size + local_block_offsets
255
+ slot_ids = tl.where(is_local, slot_ids, PAD_ID)
256
tl.store(slot_mapping_ptr + offset, slot_ids, mask=offset < end_idx)
257
258
vllm/v1/worker/gpu/cudagraph_utils.py
+20
@@ -10,6 +10,7 @@ from tqdm import tqdm
10
11
from vllm.config import VllmConfig
12
from vllm.config.compilation import CUDAGraphMode
13
+from vllm.distributed import get_dcp_group
14
from vllm.distributed.parallel_state import graph_capture, is_global_first_rank
15
from vllm.forward_context import set_forward_context
16
from vllm.v1.attention.backend import AttentionMetadataBuilder
@@ -17,6 +18,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig
18
from vllm.v1.worker.gpu.attn_utils import (
19
build_attn_metadata,
20
build_slot_mappings_by_layer,
21
+ prepare_dcp_local_seq_lens,
22
)
23
from vllm.v1.worker.gpu.block_table import BlockTables
24
from vllm.v1.worker.gpu.dp_utils import make_num_tokens_across_dp
@@ -257,6 +259,23 @@ def prepare_inputs_to_capture(
259
input_buffers.seq_lens[:num_reqs] = num_tokens
260
input_buffers.seq_lens[num_reqs:] = 0
261
262
+ try:
263
+ dcp_group = get_dcp_group()
264
+ dcp_world_size = dcp_group.world_size
265
+ dcp_rank = dcp_group.rank_in_group
266
+ except AssertionError:
267
+ dcp_world_size = 1
268
+ dcp_rank = 0
269
+ if dcp_world_size > 1:
270
+ prepare_dcp_local_seq_lens(
271
+ input_buffers.dcp_local_seq_lens,
272
+ input_buffers.seq_lens,
273
+ num_reqs,
274
+ dcp_size=dcp_world_size,
275
+ dcp_rank=dcp_rank,
276
+ cp_kv_cache_interleave_size=block_tables.cp_kv_cache_interleave_size,
277
+ )
278
+
279
input_block_tables = [x[:num_reqs] for x in block_tables.input_block_tables]
280
slot_mappings = block_tables.slot_mappings[:, :num_tokens]
281
slot_mappings_by_layer = build_slot_mappings_by_layer(
@@ -275,5 +294,6 @@ def prepare_inputs_to_capture(
294
block_tables=input_block_tables,
295
slot_mappings=slot_mappings,
296
kv_cache_config=kv_cache_config,
297
+ dcp_local_seq_lens=input_buffers.dcp_local_seq_lens,
298
)
299
return attn_metadata, slot_mappings_by_layer
vllm/v1/worker/gpu/input_batch.py
+4
@@ -27,6 +27,10 @@ class InputBuffers:
27
max_num_reqs + 1, dtype=torch.int32, device=device
28
)
29
self.seq_lens = torch.zeros(max_num_reqs, dtype=torch.int32, device=device)
30
+ # DCP: per-request local seq_lens buffer
31
+ self.dcp_local_seq_lens = torch.zeros(
32
+ max_num_reqs, dtype=torch.int32, device=device
33
+ )
34
35
36
@dataclass
vllm/v1/worker/gpu/model_runner.py
+22
@@ -11,6 +11,7 @@ import torch.nn as nn
11
from vllm.config import VllmConfig
12
from vllm.config.compilation import CUDAGraphMode
13
from vllm.distributed.parallel_state import (
14
+ get_dcp_group,
15
get_pp_group,
16
prepare_communication_buffer_for_model,
17
)
@@ -24,6 +25,7 @@ from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
25
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
26
from vllm.v1.kv_cache_interface import KVCacheConfig
27
from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput
28
+from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
29
from vllm.v1.worker.gpu.async_utils import AsyncOutput
30
from vllm.v1.worker.gpu.attn_utils import (
31
build_attn_metadata,
@@ -31,6 +33,7 @@ from vllm.v1.worker.gpu.attn_utils import (
33
get_kv_cache_spec,
34
init_attn_backend,
35
init_kv_cache,
36
+ prepare_dcp_local_seq_lens,
37
)
38
from vllm.v1.worker.gpu.block_table import BlockTables
39
from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
@@ -248,11 +251,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
251
max_num_batched_tokens=self.max_num_tokens,
252
max_model_len=self.max_model_len,
253
device=self.device,
254
+ cp_kv_cache_interleave_size=(
255
+ self.parallel_config.cp_kv_cache_interleave_size
256
+ ),
257
)
258
259
self.attn_backends, self.attn_metadata_builders = init_attn_backend(
260
self.kv_cache_config, self.vllm_config, self.device
261
)
262
+ check_attention_cp_compatibility(self.vllm_config)
263
if self.do_spec_decode:
264
# HACK(woosuk)
265
self.speculator.set_attn(
@@ -294,6 +301,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
301
block_tables=block_tables,
302
slot_mappings=slot_mappings,
303
kv_cache_config=self.kv_cache_config,
304
+ dcp_local_seq_lens=self.input_buffers.dcp_local_seq_lens,
305
)
306
input_batch.attn_metadata = attn_metadata
307
input_batch.slot_mappings = slot_mappings_by_layer
@@ -627,6 +635,19 @@ class GPUModelRunner(LoRAModelRunnerMixin):
635
)
636
seq_lens = self.input_buffers.seq_lens[:num_reqs]
637
638
+ dcp_size = self.parallel_config.decode_context_parallel_size
639
+ if dcp_size > 1:
640
+ prepare_dcp_local_seq_lens(
641
+ self.input_buffers.dcp_local_seq_lens,
642
+ seq_lens,
643
+ num_reqs,
644
+ dcp_size=dcp_size,
645
+ dcp_rank=get_dcp_group().rank_in_group,
646
+ cp_kv_cache_interleave_size=(
647
+ self.parallel_config.cp_kv_cache_interleave_size
648
+ ),
649
+ )
650
+
651
# Prepare M-RoPE positions.
652
if self.uses_mrope:
653
self.mrope_states.prepare_mrope_positions(
@@ -674,6 +695,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
695
block_tables=block_tables,
696
slot_mappings=slot_mappings,
697
kv_cache_config=self.kv_cache_config,
698
+ dcp_local_seq_lens=self.input_buffers.dcp_local_seq_lens,
699
)
700
701
input_ids = self.input_buffers.input_ids[:num_tokens_after_padding]