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, 02:30 UTC
[Model Runner V2] Minor simplification for DCP (#34786)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Feb 18, 2026 at 11:04 UTC
95be2a7f22ab89348bd79056f9050268915c3e7f
5 files changed
+111
-95
vllm/v1/worker/gpu/attn_utils.py
-24
@@ -12,7 +12,6 @@ from vllm.v1.attention.backend import (
12
AttentionMetadataBuilder,
13
CommonAttentionMetadata,
14
)
15
-from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
15
from vllm.v1.kv_cache_interface import (
16
AttentionSpec,
17
KVCacheConfig,
@@ -144,28 +143,6 @@ def build_slot_mappings_by_layer(
143
return slot_mappings_by_layer
144
145
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
-
146
def build_attn_metadata(
147
attn_metadata_builders: list[AttentionMetadataBuilder],
148
num_reqs: int,
@@ -181,7 +158,6 @@ def build_attn_metadata(
158
dcp_local_seq_lens: torch.Tensor | None = None,
159
) -> dict[str, Any]:
160
seq_lens = seq_lens[:num_reqs]
184
-
161
if dcp_local_seq_lens is not None:
162
dcp_local_seq_lens = dcp_local_seq_lens[:num_reqs]
163
vllm/v1/worker/gpu/block_table.py
+30
-40
@@ -4,7 +4,6 @@ from collections.abc import Iterable
4
5
import torch
6
7
-from vllm.distributed import get_dcp_group
7
from vllm.triton_utils import tl, triton
8
from vllm.utils.math_utils import cdiv
9
from vllm.v1.attention.backends.utils import PAD_SLOT_ID
@@ -19,36 +18,29 @@ class BlockTables:
18
max_num_batched_tokens: int,
19
max_model_len: int,
20
device: torch.device,
22
- cp_kv_cache_interleave_size: int = 1,
21
+ cp_size: int = 1,
22
+ cp_rank: int = 0,
23
+ cp_interleave: int = 1,
24
):
25
self.block_sizes = block_sizes
26
self.max_num_reqs = max_num_reqs
27
self.max_num_batched_tokens = max_num_batched_tokens
28
self.max_model_len = max_model_len
29
self.device = device
29
- assert cp_kv_cache_interleave_size >= 1
30
- self.cp_kv_cache_interleave_size = cp_kv_cache_interleave_size
30
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
31
+ self.cp_size = cp_size
32
+ self.cp_rank = cp_rank
33
+ self.cp_interleave = cp_interleave
34
35
self.num_kv_cache_groups = len(self.block_sizes)
36
# num_kv_cache_groups x [max_num_reqs, max_num_blocks]
37
self.block_tables: list[StagedWriteTensor] = []
38
for i in range(self.num_kv_cache_groups):
39
block_size = self.block_sizes[i]
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)
40
+ # When using DCP, each request's KV cache is sharded among different ranks.
41
+ # As a result, one block on the current rank covers `block_size * cp_size`
42
+ # tokens in the full, global (unsharded) sequence.
43
+ max_num_blocks = cdiv(self.max_model_len, block_size * self.cp_size)
44
block_table = StagedWriteTensor(
45
(self.max_num_reqs, max_num_blocks),
46
dtype=torch.int32,
@@ -149,9 +141,9 @@ class BlockTables:
141
self.block_sizes_tensor,
142
self.slot_mappings,
143
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,
144
+ self.cp_rank,
145
+ CP_SIZE=self.cp_size,
146
+ CP_INTERLEAVE=self.cp_interleave,
147
PAD_ID=PAD_SLOT_ID,
148
TRITON_BLOCK_SIZE=1024, # type: ignore
149
)
@@ -204,9 +196,9 @@ def _compute_slot_mappings_kernel(
196
block_sizes, # [num_kv_cache_groups]
197
slot_mappings_ptr, # [num_kv_cache_groups, max_num_tokens]
198
slot_mappings_stride,
207
- TOTAL_CP_WORLD_SIZE: tl.constexpr,
208
- TOTAL_CP_RANK: tl.constexpr,
209
- CP_KV_CACHE_INTERLEAVE_SIZE: tl.constexpr,
199
+ cp_rank,
200
+ CP_SIZE: tl.constexpr,
201
+ CP_INTERLEAVE: tl.constexpr,
202
PAD_ID: tl.constexpr,
203
TRITON_BLOCK_SIZE: tl.constexpr,
204
):
@@ -225,7 +217,6 @@ def _compute_slot_mappings_kernel(
217
block_table_ptr = _load_ptr(block_table_ptrs + group_id, tl.int32)
218
block_table_stride = tl.load(block_table_strides + group_id)
219
block_size = tl.load(block_sizes + group_id)
228
- virtual_block_size = block_size * TOTAL_CP_WORLD_SIZE
220
221
req_state_idx = tl.load(idx_mapping + batch_idx)
222
start_idx = tl.load(query_start_loc + batch_idx)
@@ -233,26 +224,25 @@ def _compute_slot_mappings_kernel(
224
for i in range(start_idx, end_idx, TRITON_BLOCK_SIZE):
225
offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
226
positions = tl.load(pos + offset, mask=offset < end_idx, other=0)
236
- block_indices = positions // virtual_block_size
227
+
228
+ block_indices = positions // (block_size * CP_SIZE)
229
+ block_offsets = positions % (block_size * CP_SIZE)
230
block_numbers = tl.load(
231
block_table_ptr + req_state_idx * block_table_stride + block_indices
232
)
240
- virtual_block_offsets = positions - block_indices * virtual_block_size
233
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
- )
234
+ if CP_SIZE == 1:
235
+ # Common case: Context parallelism is not used.
236
+ slot_ids = block_numbers * block_size + block_offsets
237
+ else:
238
+ # Context parallelism is used.
239
+ is_local = block_offsets // CP_INTERLEAVE % CP_SIZE == cp_rank
240
+ rounds = block_offsets // (CP_INTERLEAVE * CP_SIZE)
241
+ remainder = block_offsets % CP_INTERLEAVE
242
+ local_offsets = rounds * CP_INTERLEAVE + remainder
243
+ slot_ids = block_numbers * block_size + local_offsets
244
+ slot_ids = tl.where(is_local, slot_ids, PAD_ID)
245
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)
246
tl.store(slot_mapping_ptr + offset, slot_ids, mask=offset < end_idx)
247
248
vllm/v1/worker/gpu/cp_utils.py
new
+61
@@ -0,0 +1,61 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+
5
+from vllm.triton_utils import tl, triton
6
+
7
+
8
+def prepare_dcp_local_seq_lens(
9
+ dcp_local_seq_lens: torch.Tensor,
10
+ seq_lens: torch.Tensor,
11
+ num_reqs: int,
12
+ dcp_size: int,
13
+ dcp_rank: int,
14
+ cp_interleave: int,
15
+) -> None:
16
+ """Populate the persistent DCP local seq_lens buffer (CUDA graph safe)."""
17
+ if dcp_size == 1:
18
+ return
19
+
20
+ max_num_reqs = dcp_local_seq_lens.shape[0]
21
+ BLOCK_SIZE = 128
22
+ num_blocks = triton.cdiv(max_num_reqs, BLOCK_SIZE)
23
+ _dcp_local_seq_lens_kernel[(num_blocks,)](
24
+ dcp_local_seq_lens,
25
+ seq_lens,
26
+ dcp_size,
27
+ dcp_rank,
28
+ cp_interleave,
29
+ num_reqs,
30
+ max_num_reqs,
31
+ BLOCK_SIZE,
32
+ )
33
+
34
+
35
+@triton.jit
36
+def _dcp_local_seq_lens_kernel(
37
+ out_ptr,
38
+ seq_lens_ptr,
39
+ dcp_size,
40
+ dcp_rank,
41
+ cp_interleave,
42
+ num_reqs,
43
+ max_num_reqs,
44
+ BLOCK_SIZE: tl.constexpr,
45
+):
46
+ pid = tl.program_id(0)
47
+ block = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
48
+
49
+ seq_lens = tl.load(seq_lens_ptr + block, mask=block < num_reqs)
50
+
51
+ # Distribute KV cache among different ranks, in a round-robin manner.
52
+ rounds = seq_lens // (dcp_size * cp_interleave)
53
+ remainder = seq_lens % (dcp_size * cp_interleave)
54
+
55
+ remainder = tl.maximum(remainder - dcp_rank * cp_interleave, 0)
56
+ remainder = tl.minimum(remainder, cp_interleave)
57
+ local_seq_lens = rounds * cp_interleave + remainder
58
+
59
+ # For [num_reqs, max_num_reqs), pad with 0
60
+ local_seq_lens = tl.where(block < num_reqs, local_seq_lens, 0)
61
+ tl.store(out_ptr + block, local_seq_lens, mask=block < max_num_reqs)
vllm/v1/worker/gpu/cudagraph_utils.py
+2
-18
@@ -10,7 +10,6 @@ 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
13
from vllm.distributed.parallel_state import graph_capture, is_global_first_rank
14
from vllm.forward_context import set_forward_context
15
from vllm.v1.attention.backend import AttentionMetadataBuilder
@@ -18,7 +17,6 @@ from vllm.v1.kv_cache_interface import KVCacheConfig
17
from vllm.v1.worker.gpu.attn_utils import (
18
build_attn_metadata,
19
build_slot_mappings_by_layer,
21
- prepare_dcp_local_seq_lens,
20
)
21
from vllm.v1.worker.gpu.block_table import BlockTables
22
from vllm.v1.worker.gpu.dp_utils import make_num_tokens_across_dp
@@ -259,22 +257,8 @@ def prepare_inputs_to_capture(
257
input_buffers.seq_lens[:num_reqs] = num_tokens
258
input_buffers.seq_lens[num_reqs:] = 0
259
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
- )
260
+ input_buffers.dcp_local_seq_lens[:num_reqs] = num_tokens
261
+ input_buffers.dcp_local_seq_lens[num_reqs:] = 0
262
263
input_block_tables = [x[:num_reqs] for x in block_tables.input_block_tables]
264
slot_mappings = block_tables.slot_mappings[:, :num_tokens]
vllm/v1/worker/gpu/model_runner.py
+18
-13
@@ -33,10 +33,10 @@ 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,
36
)
37
from vllm.v1.worker.gpu.block_table import BlockTables
38
from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
39
+from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
40
from vllm.v1.worker.gpu.cudagraph_utils import CudaGraphManager
41
from vllm.v1.worker.gpu.dp_utils import (
42
get_cudagraph_and_dp_padding,
@@ -192,6 +192,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
192
self.is_first_pp_rank = True
193
self.is_last_pp_rank = True
194
195
+ # Decode context parallelism.
196
+ self.dcp_size = self.parallel_config.decode_context_parallel_size
197
+ self.use_dcp = self.dcp_size > 1
198
+ self.dcp_rank = get_dcp_group().rank_in_group if self.use_dcp else 0
199
+ self.cp_interleave = self.parallel_config.cp_kv_cache_interleave_size
200
+
201
def update_max_model_len(self, max_model_len: int) -> None:
202
self.max_model_len = max_model_len
203
self.req_states.max_model_len = max_model_len
@@ -251,9 +257,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
257
max_num_batched_tokens=self.max_num_tokens,
258
max_model_len=self.max_model_len,
259
device=self.device,
254
- cp_kv_cache_interleave_size=(
255
- self.parallel_config.cp_kv_cache_interleave_size
256
- ),
260
+ cp_size=self.dcp_size,
261
+ cp_rank=self.dcp_rank,
262
+ cp_interleave=self.cp_interleave,
263
)
264
265
self.attn_backends, self.attn_metadata_builders = init_attn_backend(
@@ -636,18 +642,17 @@ class GPUModelRunner(LoRAModelRunnerMixin):
642
)
643
seq_lens = self.input_buffers.seq_lens[:num_reqs]
644
639
- dcp_size = self.parallel_config.decode_context_parallel_size
640
- if dcp_size > 1:
645
+ if self.use_dcp:
646
+ # Prepare dcp local seq_lens.
647
prepare_dcp_local_seq_lens(
648
self.input_buffers.dcp_local_seq_lens,
643
- seq_lens,
649
+ self.input_buffers.seq_lens,
650
num_reqs,
645
- dcp_size=dcp_size,
646
- dcp_rank=get_dcp_group().rank_in_group,
647
- cp_kv_cache_interleave_size=(
648
- self.parallel_config.cp_kv_cache_interleave_size
649
- ),
651
+ self.dcp_size,
652
+ self.dcp_rank,
653
+ self.cp_interleave,
654
)
655
+ dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens[:num_reqs]
656
657
# Prepare M-RoPE positions.
658
if self.uses_mrope:
@@ -696,7 +701,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
701
block_tables=block_tables,
702
slot_mappings=slot_mappings,
703
kv_cache_config=self.kv_cache_config,
699
- dcp_local_seq_lens=self.input_buffers.dcp_local_seq_lens,
704
+ dcp_local_seq_lens=dcp_local_seq_lens,
705
)
706
707
input_ids = self.input_buffers.input_ids[:num_tokens_after_padding]