@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2108a571d7

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 13, 14:30 UTC

[DCP] Support dcp kv_cache interleave size > 1 (#26696)

Signed-off-by: zhangsicheng5 <zhangsicheng5@huawei.com> Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com> Signed-off-by: Qiu <qiuchunshuo@huawei.com> Co-authored-by: QiuChunshuo <qiuchunshuo@huawei.com>

zhangsicheng5 committed Nov 9, 2025 at 03:45 UTC 2108a571d7ee522346d309ed902e9da8e7cdc97b

12 files changed +202 -79

tests/distributed/test_context_parallel.py

@@ -30,6 +30,7 @@ class ParallelSetup(NamedTuple):

30 tp_size: int

31 pp_size: int

32 dcp_size: int

33 + dcp_kv_cache_interleave_size: int

34 eager_mode: bool

35 chunked_prefill: bool

@@ -52,6 +53,7 @@ class CPTestSettings:

53 tp_base: int = 4,

54 pp_base: int = 1,

55 dcp_base: int = 1,

56 + dcp_kv_cache_interleave_size: int = 1,

57 multi_node_only: bool = False,

58 runner: RunnerOption = "auto",

59 load_format: str | None = None,

@@ -66,6 +68,7 @@ class CPTestSettings:

68 tp_size=tp_base,

69 pp_size=pp_multiplier * pp_base,

70 dcp_size=int(dcp_multiplier * tp_base),

71 + dcp_kv_cache_interleave_size=dcp_kv_cache_interleave_size,

72 eager_mode=eager_mode_val,

73 chunked_prefill=chunked_prefill_val,

74 )

@@ -108,6 +111,7 @@ def _compare_cp_with_tp(

111 tp_size,

112 pp_size,

113 dcp_size,

114 + dcp_kv_cache_interleave_size,

115 eager_mode,

116 chunked_prefill,

117 ) = parallel_setup

@@ -180,6 +184,8 @@ def _compare_cp_with_tp(

184 str(pp_size),

185 "--decode-context-parallel-size",

186 str(dcp_size),

187 + "--dcp-kv-cache-interleave-size",

188 + str(dcp_kv_cache_interleave_size),

189 "--distributed-executor-backend",

190 distributed_backend,

191 ]

@@ -207,6 +213,7 @@ CP_TEXT_GENERATION_MODELS = {

213 "deepseek-ai/DeepSeek-V2-Lite-Chat": [

214 CPTestSettings.detailed(),

215 CPTestSettings.detailed(tp_base=2),

216 + CPTestSettings.detailed(tp_base=2, dcp_kv_cache_interleave_size=64),

217 ],

218 "bigcode/gpt_bigcode-santacoder": [

219 CPTestSettings.detailed(),

tests/v1/worker/test_gpu_model_runner.py

@@ -951,6 +951,7 @@ def test_hybrid_block_table_initialization():

951 max_num_reqs = 10

952 max_num_blocks_per_req = 20

953 max_num_batched_tokens = 512

954 + dcp_kv_cache_interleave_size = 8

955

956 block_table = BlockTable(

957 block_size=block_size,

@@ -960,6 +961,7 @@ def test_hybrid_block_table_initialization():

961 pin_memory=False,

962 device=torch.device(DEVICE),

963 kernel_block_size=kernel_block_sizes[0],

964 + dcp_kv_cache_interleave_size=dcp_kv_cache_interleave_size,

965 )

966

967 # Verify hybrid block configuration

vllm/attention/ops/common.py

@@ -53,6 +53,7 @@ def _correct_attn_cp_out_kernel(

53 lse = tl.load(lses_ptr + lse_offsets)

54 lse = tl.where((lse != lse) | (lse == float("inf")), -float("inf"), lse)

55 lse_max = tl.max(lse, axis=0)

56 + lse_max = tl.where(lse_max == -float("inf"), 0, lse_max)

57 lse -= lse_max

58 lse_exp = tl.exp(lse)

59 lse_acc = tl.sum(lse_exp, axis=0)

vllm/config/parallel.py

+11

@@ -227,6 +227,17 @@ class ParallelConfig:

227 not change by dcp, it simply reuse the GPUs of TP group, and tp_size

228 needs to be divisible by dcp_size."""

229

230 + dcp_kv_cache_interleave_size: int = 1

231 + """Interleave size of kv_cache storage while using dcp or cp > 1,

232 + store interleave_size tokens on (d)cp i,

233 + then store next interleave_size tokens on (d)cp i+1.

234 + Interleave_size=1: token-level align, token i is stored on rank i % (d)cp_size.

235 + Interleave_size=block_size: block-level align, first fill the block on first rank,

236 + token is stored on rank i+1 block j after rank i block j is full.

237 + Block_size should be greater than or equal to dcp_kv_cache_interleave_size.

238 + Block_size should be divisible by dcp_kv_cache_interleave_size.

239 + """

240 +

241 _api_process_count: int = Field(default=1, gt=0)

242 """

243 The number of API processes initialized.

vllm/config/vllm.py

+17

@@ -608,6 +608,23 @@ class VllmConfig:

608 )

609 current_platform.check_and_update_config(self)

610

611 + assert (

612 + self.parallel_config.dcp_kv_cache_interleave_size

613 + <= self.cache_config.block_size

614 + and self.cache_config.block_size

615 + % self.parallel_config.dcp_kv_cache_interleave_size

616 + == 0

617 + ), (

618 + f"Block_size({self.cache_config.block_size}) should be "

619 + "greater than or equal to and divisible by dcp_kv_cache_interleave_size "

620 + f"({self.parallel_config.dcp_kv_cache_interleave_size})."

621 + )

622 +

623 + assert (

624 + self.parallel_config.dcp_kv_cache_interleave_size == 1

625 + or self.speculative_config is None

626 + ), "MTP with dcp_kv_cache_interleave_size > 1 is not supported now."

627 +

628 # Do this after all the updates to compilation_config.mode

629 if self.compilation_config.mode == CompilationMode.VLLM_COMPILE:

630 self.compilation_config.set_splitting_ops_for_v1()

vllm/engine/arg_utils.py

@@ -385,6 +385,7 @@ class EngineArgs:

385 pipeline_parallel_size: int = ParallelConfig.pipeline_parallel_size

386 tensor_parallel_size: int = ParallelConfig.tensor_parallel_size

387 decode_context_parallel_size: int = ParallelConfig.decode_context_parallel_size

388 + dcp_kv_cache_interleave_size: int = ParallelConfig.dcp_kv_cache_interleave_size

389 data_parallel_size: int = ParallelConfig.data_parallel_size

390 data_parallel_rank: int | None = None

391 data_parallel_start_rank: int | None = None

@@ -750,6 +751,10 @@ class EngineArgs:

751 "-dcp",

752 **parallel_kwargs["decode_context_parallel_size"],

753 )

754 + parallel_group.add_argument(

755 + "--dcp-kv-cache-interleave-size",

756 + **parallel_kwargs["dcp_kv_cache_interleave_size"],

757 + )

758 parallel_group.add_argument(

759 "--data-parallel-size", "-dp", **parallel_kwargs["data_parallel_size"]

760 )

@@ -1518,6 +1523,7 @@ class EngineArgs:

1523 worker_cls=self.worker_cls,

1524 worker_extension_cls=self.worker_extension_cls,

1525 decode_context_parallel_size=self.decode_context_parallel_size,

1526 + dcp_kv_cache_interleave_size=self.dcp_kv_cache_interleave_size,

1527 _api_process_count=self._api_process_count,

1528 _api_process_rank=self._api_process_rank,

1529 )

vllm/v1/attention/backends/flash_attn.py

+11 -2

@@ -43,6 +43,7 @@ from vllm.v1.attention.backends.utils import (

43 AttentionCGSupport,

44 AttentionMetadataBuilder,

45 CommonAttentionMetadata,

46 + get_dcp_local_seq_lens,

47 get_kv_cache_layout,

48 )

49 from vllm.v1.kv_cache_interface import AttentionSpec

@@ -238,6 +239,10 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad

239 self.dcp_world_size = 1

240 self.dcp_rank = 0

241

242 + self.dcp_kv_cache_interleave_size = (

243 + self.parallel_config.dcp_kv_cache_interleave_size

244 + )

245 +

246 self.use_full_cuda_graph = (

247 self.compilation_config.cudagraph_mode.has_full_cudagraphs()

248 )

@@ -352,8 +357,12 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad

357 - common_attn_metadata.query_start_loc_cpu[:-1]

358 )

359 dcp_context_kv_lens_cpu = seq_lens_cpu - query_kv_lens_cpu

355 - dcp_context_kv_lens_cpu = dcp_context_kv_lens_cpu // self.dcp_world_size + (

356 - self.dcp_rank <= (dcp_context_kv_lens_cpu - 1) % self.dcp_world_size

360 +

361 + dcp_context_kv_lens_cpu = get_dcp_local_seq_lens(

362 + dcp_context_kv_lens_cpu,

363 + self.dcp_world_size,

364 + self.dcp_rank,

365 + self.dcp_kv_cache_interleave_size,

366 )

367 dcp_context_kv_lens = dcp_context_kv_lens_cpu.to(self.device)

368 max_dcp_context_kv_len = dcp_context_kv_lens.max().item()

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

+77 -75

@@ -225,6 +225,7 @@ from vllm.utils.math_utils import cdiv, round_down

225 from vllm.v1.attention.backends.utils import (

226 AttentionMetadataBuilder,

227 CommonAttentionMetadata,

228 + get_dcp_local_seq_lens,

229 get_per_layer_parameters,

230 infer_global_hyperparameters,

231 split_decodes_and_prefills,

@@ -361,10 +362,9 @@ class MLACommonPrefillMetadata:

362 workspace: torch.Tensor

363

364 # for mla DCP

364 - cp_chunk_seq_lens: list[list[int]] | None = None

365 - origin_context_lens: list[int] | None = None

366 - cp_cu_seq_lens: torch.Tensor | None = None

367 - chunk_size: int | None = None

365 + padded_local_chunk_seq_lens: list[list[int]] | None = None

366 + local_context_lens_allranks: list[list[int]] | None = None

367 + padded_local_cu_seq_lens: torch.Tensor | None = None

368 cu_seq_lens_lst: list[list[int]] | None = None

369

370 block_table: torch.Tensor

@@ -568,6 +568,8 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

568 # DCP might not be initialized in testing

569 self.dcp_world_size = 1

570 self.dcp_rank = 0

571 + self.dcp_local_block_size = parallel_config.dcp_kv_cache_interleave_size

572 + self.dcp_virtual_block_size = self.dcp_local_block_size * self.dcp_world_size

573

574 # Don't try to access the runner on AMD

575 if self.aot_schedule:

@@ -794,15 +796,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

796 )

797 )

798

797 - # Note(hc): update seq_lens of decode reqs under DCP.

798 - if self.dcp_world_size > 1:

799 - assert dcp_local_seq_lens is not None

800 - dcp_local_seq_lens[:num_decodes] = seq_lens[

801 - :num_decodes

802 - ] // self.dcp_world_size + (

803 - self.dcp_rank < seq_lens[:num_decodes] % self.dcp_world_size

804 - )

805 -

799 assert num_decodes + num_prefills == num_reqs

800 assert num_decode_tokens + num_prefill_tokens == num_tokens

801

@@ -811,11 +804,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

804 reqs_start = num_decodes # prefill_start

805

806 context_lens_cpu = num_computed_tokens_cpu[reqs_start:num_reqs]

814 - # Note(hc): The context lengths in the perspective of dcp rank0.

815 - cp_context_lens_cpu = torch.ceil(

816 - context_lens_cpu.float() / self.dcp_world_size

817 - ).int()

818 - origin_context_lens = context_lens_cpu.tolist()

807 max_context_len_cpu = context_lens_cpu.max().item()

808 num_prefills_with_context_cpu = (context_lens_cpu > 0).sum().item()

809 prefill_query_start_loc = (

@@ -871,32 +859,56 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

859 )

860

861 if self.dcp_world_size > 1:

862 + local_context_lens_allranks = get_dcp_local_seq_lens(

863 + context_lens_cpu,

864 + self.dcp_world_size,

865 + None,

866 + self.dcp_local_block_size,

867 + )

868 + # Note(qcs): The max local context lengths

869 + # padded to `dcp_local_block_size`.

870 + padded_local_context_lens_cpu = (

871 + cdiv(

872 + context_lens_cpu,

873 + self.dcp_virtual_block_size,

874 + )

875 + * self.dcp_local_block_size

876 + )

877 # Note(hc): The above max_context_chunk already enforces

878 # block_size alignment, DCP just need the block_size can

879 # be divisible by dcp_world_size, because DCP use

880 # cp_gather_cache which not require `cp_chunk_starts`

881 # aligned to page_size.

882 assert max_context_chunk % self.dcp_world_size == 0

880 - cp_max_context_chunk = max_context_chunk // self.dcp_world_size

881 - cp_chunk_starts = (

883 + padded_local_max_context_chunk_across_ranks = (

884 + cdiv(

885 + max_context_chunk,

886 + self.dcp_virtual_block_size,

887 + )

888 + * self.dcp_local_block_size

889 + )

890 + local_chunk_starts = (

891 torch.arange(num_chunks, dtype=torch.int32)

892 .unsqueeze(1)

893 .expand(-1, num_prefills)

885 - * cp_max_context_chunk

894 + * padded_local_max_context_chunk_across_ranks

895 )

887 - cp_chunk_ends = torch.min(

888 - cp_context_lens_cpu.unsqueeze(0),

889 - cp_chunk_starts + cp_max_context_chunk,

896 + local_chunk_ends = torch.min(

897 + padded_local_context_lens_cpu.unsqueeze(0),

898 + local_chunk_starts

899 + + padded_local_max_context_chunk_across_ranks,

900 )

891 - cp_chunk_seq_lens = (cp_chunk_ends - cp_chunk_starts).clamp(min=0)

901 + padded_local_chunk_seq_lens = (

902 + local_chunk_ends - local_chunk_starts

903 + ).clamp(min=0)

904

893 - cp_cu_seq_lens_cpu = torch.zeros(

905 + padded_local_cu_chunk_seq_lens_cpu = torch.zeros(

906 num_chunks, num_prefills + 1, dtype=torch.int32, pin_memory=True

907 )

908 torch.cumsum(

897 - cp_chunk_seq_lens,

909 + padded_local_chunk_seq_lens,

910 dim=1,

899 - out=cp_cu_seq_lens_cpu[:, 1:],

911 + out=padded_local_cu_chunk_seq_lens_cpu[:, 1:],

912 dtype=torch.int32,

913 )

914

@@ -908,15 +920,16 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

920 if self.dcp_world_size > 1:

921 chunked_context_metadata = chunked_context_metadata_cls(

922 cu_seq_lens=cu_seq_lens_cpu.to(device, non_blocking=True),

911 - starts=cp_chunk_starts.to(device, non_blocking=True),

912 - seq_tot=cp_chunk_seq_lens.sum(dim=1).tolist(),

923 + starts=local_chunk_starts.to(device, non_blocking=True),

924 + seq_tot=padded_local_chunk_seq_lens.sum(dim=1).tolist(),

925 max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),

926 seq_lens=chunk_seq_lens,

927 workspace=self.chunked_prefill_workspace,

916 - cp_chunk_seq_lens=cp_chunk_seq_lens.tolist(),

917 - origin_context_lens=origin_context_lens,

918 - cp_cu_seq_lens=cp_cu_seq_lens_cpu.to(device, non_blocking=True),

919 - chunk_size=max_context_chunk,

928 + padded_local_chunk_seq_lens=padded_local_chunk_seq_lens.tolist(),

929 + local_context_lens_allranks=local_context_lens_allranks.tolist(),

930 + padded_local_cu_seq_lens=padded_local_cu_chunk_seq_lens_cpu.to(

931 + device, non_blocking=True

932 + ),

933 cu_seq_lens_lst=cu_seq_lens_cpu.tolist(),

934 )

935 else:

@@ -998,64 +1011,52 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1011 def reorg_kvcache(

1012 allgatered_kv_c_normed: torch.Tensor,

1013 allgatered_k_pe: torch.Tensor,

1001 - cp_chunk_seq_lens_lst: list[int],

1002 - origin_context_lens: list[int],

1003 - cp_world_size: int,

1014 + padded_local_chunk_seq_lens_lst: list[int],

1015 + local_context_lens_allranks: list[list[int]],

1016 sum_seq_len: int,

1017 max_seq_len: int,

1006 - chunk_size: int,

1007 - chunk_idx: int,

1018 toks: int,

1019 ) -> tuple[torch.Tensor, torch.Tensor]:

1020 """

1011 - reorg kvcache after cp local gather to tp layout for attn kernel.

1012 -

1021 + reorg and unpad kvcache after cp local gather to tp layout for attn kernel.

1022 + e.g.

1023 + allgatered_kv_c_normed = [T0_0, T0_1, T0_2, T0_3, T1_0, T1_1, ...,

1024 + T0_4, T0_5, pad, pad, T1_2, pad, ...]

1025 + -> reorganized_kv_c_normed = [T0_0, T0_1, T0_2, T0_3, T0_4, T0_5,

1026 + T1_0, T1_1, T1_2, ...]

1027 Args:

1014 - cp_chunk_seq_lens_lst: chunk context lengths under CP.

1015 - origin_context_lens: origin full context lengths under CP.

1016 - cp_world_size: CP size.

1028 + padded_local_chunk_seq_lens_lst: local chunk context lengths

1029 + under current CP rank.

1030 + local_context_lens_allranks: local context lengths on each CP rank.

1031 sum_seq_len: the sum of cp_chunk_seq_lens_lst.

1032 max_seq_len: the max value of cp_chunk_seq_lens_lst.

1019 - chunk_size: equals to max_context_chunk from

1020 - chunked_context_metadata building.

1021 - chunk_idx: chunk idx of chunked_prefill.

1033 toks: the number of tokens for local gather cache.

1034 """

1035 kv_c_segments = []

1036 k_pe_segments = []

1037 src_token_idx = 0

1038 max_seq_len_check = 0

1028 - for cp_chunk_seq_len, origin_context_len in zip(

1029 - cp_chunk_seq_lens_lst, origin_context_lens

1039 + for padded_local_chunk_seq_len, local_context_lens in zip(

1040 + padded_local_chunk_seq_lens_lst, local_context_lens_allranks

1041 ):

1031 - chunk_context_len = chunk_size

1032 - if cp_chunk_seq_len != 0:

1033 - chunk_context_len = min(

1034 - chunk_context_len, origin_context_len - chunk_size * chunk_idx

1035 - )

1036 - cp_target_rank = (chunk_context_len - 1) % cp_world_size

1042 cur_seq_len = 0

1038 - for rank in range(cp_world_size):

1039 - if rank > cp_target_rank and cp_chunk_seq_len:

1040 - real_cp_chunk_seq_len = cp_chunk_seq_len - 1

1041 - else:

1042 - real_cp_chunk_seq_len = cp_chunk_seq_len

1043 - if real_cp_chunk_seq_len:

1043 + for rank, local_context_len in enumerate(local_context_lens):

1044 + if local_context_len != 0:

1045 kv_c_segment = allgatered_kv_c_normed[

1046 rank * toks + src_token_idx : rank * toks

1047 + src_token_idx

1047 - + real_cp_chunk_seq_len

1048 + + local_context_len

1049 ]

1050 k_pe_segment = allgatered_k_pe[

1051 rank * toks + src_token_idx : rank * toks

1052 + src_token_idx

1052 - + real_cp_chunk_seq_len

1053 + + local_context_len

1054 ]

1055 kv_c_segments.append(kv_c_segment)

1056 k_pe_segments.append(k_pe_segment)

1056 - cur_seq_len += real_cp_chunk_seq_len

1057 + cur_seq_len += local_context_len

1058 max_seq_len_check = max(max_seq_len_check, cur_seq_len)

1058 - src_token_idx += cp_chunk_seq_len

1059 + src_token_idx += padded_local_chunk_seq_len

1060 reorganized_kv_c_normed = torch.cat(kv_c_segments, dim=0)

1061 reorganized_k_pe = torch.cat(k_pe_segments, dim=0)

1062 assert reorganized_kv_c_normed.shape[0] == sum_seq_len

@@ -1296,6 +1297,9 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1297 get_current_vllm_config()

1298 )

1299 )

1300 + self.dcp_kv_cache_interleave_size: int = (

1301 + get_current_vllm_config().parallel_config.dcp_kv_cache_interleave_size

1302 + )

1303

1304 def _flash_attn_varlen_diff_headdims(

1305 self, q, k, v, return_softmax_lse=False, softmax_scale=None, **kwargs

@@ -1697,10 +1701,9 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1701 assert attn_metadata.prefill is not None

1702 prefill_metadata = attn_metadata.prefill

1703 assert prefill_metadata.chunked_context is not None

1700 - assert prefill_metadata.chunked_context.cp_chunk_seq_lens is not None

1701 - assert prefill_metadata.chunked_context.origin_context_lens is not None

1702 - assert prefill_metadata.chunked_context.cp_cu_seq_lens is not None

1703 - assert prefill_metadata.chunked_context.chunk_size is not None

1704 + assert prefill_metadata.chunked_context.padded_local_chunk_seq_lens is not None

1705 + assert prefill_metadata.chunked_context.local_context_lens_allranks is not None

1706 + assert prefill_metadata.chunked_context.padded_local_cu_seq_lens is not None

1707 assert prefill_metadata.chunked_context.cu_seq_lens_lst is not None

1708

1709 output = None

@@ -1713,7 +1716,9 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1716 src_cache=kv_c_and_k_pe_cache,

1717 dst=workspace,

1718 block_table=prefill_metadata.block_table,

1716 - cu_seq_lens=prefill_metadata.chunked_context.cp_cu_seq_lens[i],

1719 + cu_seq_lens=prefill_metadata.chunked_context.padded_local_cu_seq_lens[

1720 + i

1721 + ],

1722 batch_size=attn_metadata.num_prefills,

1723 seq_starts=prefill_metadata.chunked_context.starts[i],

1724 )

@@ -1743,15 +1748,12 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1748 kv_c_normed, k_pe = reorg_kvcache(

1749 allgatered_kv_c_normed,

1750 allgatered_k_pe,

1746 - cp_chunk_seq_lens_lst=prefill_metadata.chunked_context.cp_chunk_seq_lens[

1751 + padded_local_chunk_seq_lens_lst=prefill_metadata.chunked_context.padded_local_chunk_seq_lens[

1752 i

1753 ],

1749 - origin_context_lens=prefill_metadata.chunked_context.origin_context_lens,

1750 - cp_world_size=dcp_world_size,

1754 + local_context_lens_allranks=prefill_metadata.chunked_context.local_context_lens_allranks,

1755 sum_seq_len=prefill_metadata.chunked_context.cu_seq_lens_lst[i][-1],

1756 max_seq_len=prefill_metadata.chunked_context.max_seq_lens[i],

1753 - chunk_size=prefill_metadata.chunked_context.chunk_size,

1754 - chunk_idx=i,

1757 toks=toks,

1758 )

1759

vllm/v1/attention/backends/utils.py

+38

@@ -1076,3 +1076,41 @@ def compute_causal_conv1d_metadata(query_start_loc_p: torch.Tensor):

1076 nums_dict[BLOCK_M]["token_chunk_offset_ptr"] = token_chunk_offset_ptr # type: ignore

1077

1078 return nums_dict, batch_ptr, token_chunk_offset_ptr

1079 +

1080 +

1081 +def get_dcp_local_seq_lens(

1082 + seq_lens: torch.Tensor,

1083 + dcp_world_size: int = 1,

1084 + dcp_rank: int | None = None,

1085 + dcp_kv_cache_interleave_size: int = 1,

1086 +) -> torch.Tensor:

1087 + """While using dcp, kv_cache size stored on each rank may be different,

1088 + use this function to calculate split decode seq_lens of each dcp rank.

1089 + Only consider dcp now, we can extend the case of cp based on this.

1090 + """

1091 + num_requests = seq_lens.size(0)

1092 + if dcp_rank is None:

1093 + rank_offsets = (

1094 + torch.arange(dcp_world_size, dtype=torch.int32)

1095 + .unsqueeze(0)

1096 + .repeat(num_requests, 1)

1097 + )

1098 + else:

1099 + rank_offsets = torch.Tensor([[dcp_rank]]).to(dtype=torch.int32)

1100 + seq_lens_tiled = (

1101 + seq_lens.to(torch.int32).unsqueeze(-1).repeat(1, rank_offsets.shape[1])

1102 + )

1103 + base = (

1104 + seq_lens_tiled

1105 + // dcp_kv_cache_interleave_size

1106 + // dcp_world_size

1107 + * dcp_kv_cache_interleave_size

1108 + )

1109 + remainder = seq_lens_tiled - base * dcp_world_size

1110 + remainder = torch.clip(

1111 + remainder - rank_offsets * dcp_kv_cache_interleave_size,

1112 + 0,

1113 + dcp_kv_cache_interleave_size,

1114 + )

1115 + dcp_local_seq_lens = base + remainder

1116 + return dcp_local_seq_lens.squeeze(1)

vllm/v1/worker/block_table.py

+16 -2

@@ -22,6 +22,7 @@ class BlockTable:

22 pin_memory: bool,

23 device: torch.device,

24 kernel_block_size: int,

25 + dcp_kv_cache_interleave_size: int,

26 ):

27 """

28 Args:

@@ -86,6 +87,7 @@ class BlockTable:

87 # DCP might not be initialized in testing

88 self.dcp_world_size = 1

89 self.dcp_rank = 0

90 + self.dcp_kv_cache_interleave_size = dcp_kv_cache_interleave_size

92 def append_row(

93 self,

@@ -144,9 +146,19 @@ class BlockTable:

146 # Use virtual_block_size for mask calculation, which marks local

147 # tokens.

148 virtual_block_offsets = positions % virtual_block_size

147 - mask = virtual_block_offsets % self.dcp_world_size == self.dcp_rank

149 + mask = (

150 + virtual_block_offsets

151 + // self.dcp_kv_cache_interleave_size

152 + % self.dcp_world_size

153 + == self.dcp_rank

154 + )

155 # Calculate local block_offsets

149 - block_offsets = virtual_block_offsets // self.dcp_world_size

156 + block_offsets = (

157 + virtual_block_offsets

158 + // (self.dcp_world_size * self.dcp_kv_cache_interleave_size)

159 + * self.dcp_kv_cache_interleave_size

160 + + virtual_block_offsets % self.dcp_kv_cache_interleave_size

161 + )

162 # Calculate slot_mapping

163 slot_mapping = block_numbers * self.block_size + block_offsets

164 # Write final slots, use -1 for not-local

@@ -234,6 +246,7 @@ class MultiGroupBlockTable:

246 block_sizes: list[int],

247 kernel_block_sizes: list[int],

248 num_speculative_tokens: int = 0,

249 + dcp_kv_cache_interleave_size: int = 1,

250 ) -> None:

251 # Note(hc): each dcp rank only store

252 # (max_model_len//dcp_world_size) tokens in kvcache,

@@ -263,6 +276,7 @@ class MultiGroupBlockTable:

276 pin_memory,

277 device,

278 kernel_block_size,

279 + dcp_kv_cache_interleave_size,

280 )

281 for block_size, kernel_block_size in zip(block_sizes, kernel_block_sizes)

282 ]

vllm/v1/worker/gpu_input_batch.py

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

84 is_spec_decode: bool = False,

85 is_pooling_model: bool = False,

86 num_speculative_tokens: int = 0,

87 + dcp_kv_cache_interleave_size: int = 1,

88 ):

89 self.is_pooling_model = is_pooling_model

90 self.is_spec_decode = is_spec_decode

@@ -137,6 +138,7 @@ class InputBatch:

138 block_sizes=block_sizes,

139 kernel_block_sizes=kernel_block_sizes,

140 num_speculative_tokens=num_speculative_tokens,

141 + dcp_kv_cache_interleave_size=dcp_kv_cache_interleave_size,

142 )

143

144 # Sampling-related.

vllm/v1/worker/gpu_model_runner.py

+14

@@ -35,6 +35,7 @@ from vllm.distributed.eplb.eplb_state import EplbState

35 from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group

36 from vllm.distributed.kv_transfer.kv_connector.utils import copy_kv_blocks

37 from vllm.distributed.parallel_state import (

38 + get_dcp_group,

39 get_pp_group,

40 get_tp_group,

41 graph_capture,

@@ -88,6 +89,7 @@ from vllm.v1.attention.backends.utils import (

89 AttentionMetadataBuilder,

90 CommonAttentionMetadata,

91 create_fast_prefill_custom_backend,

92 + get_dcp_local_seq_lens,

93 reorder_batch_to_split_decodes_and_prefills,

94 split_attn_metadata,

95 )

@@ -275,6 +277,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

277 self.is_multimodal_pruning_enabled = False

278 self.max_model_len = model_config.max_model_len

279 self.dcp_world_size = self.parallel_config.decode_context_parallel_size

280 + self.dcp_rank = 0 if self.dcp_world_size <= 1 else get_dcp_group().rank_in_group

281 self.max_num_tokens = scheduler_config.max_num_batched_tokens

282 self.max_num_reqs = scheduler_config.max_num_seqs

283

@@ -396,6 +399,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

399 # uses output token ids so we set this conservatively.

400 logitsprocs_need_output_token_ids=bool(custom_logitsprocs),

401 is_pooling_model=self.is_pooling_model,

402 + dcp_kv_cache_interleave_size=self.parallel_config.dcp_kv_cache_interleave_size,

403 )

404

405 self.use_async_scheduling = self.scheduler_config.async_scheduling

@@ -1307,6 +1311,16 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1311 logits_indices

1312 )

1313

1314 + # update seq_lens of decode reqs under DCP.

1315 + if self.dcp_world_size > 1:

1316 + self.dcp_local_seq_lens.cpu[:num_reqs] = get_dcp_local_seq_lens(

1317 + self.seq_lens.cpu[:num_reqs],

1318 + self.dcp_world_size,

1319 + self.dcp_rank,

1320 + self.parallel_config.dcp_kv_cache_interleave_size,

1321 + )

1322 + self.dcp_local_seq_lens.copy_to_gpu(num_reqs)

1323 +

1324 attn_metadata: PerLayerAttnMetadata = {}

1325 if ubatch_slices is not None:

1326 attn_metadata = [dict() for _ in range(len(ubatch_slices))]