@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4516d44b7f

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, 05:45 UTC

[DCP] Support Decode Context Parallel (DCP) for GQA with Flashinfer (#25438)

Signed-off-by: gaojc <1055866782@qq.com> Signed-off-by: Jingchun Gao <gaojingchun1@huawei.com> Signed-off-by: Jingchun Gao <63247409+gjc0824@users.noreply.github.com> Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com> Co-authored-by: gaojingchun (A) <g00955623@china.huawei.com> Co-authored-by: Jingchun Gao <gaojingchun1@huawei.com> Co-authored-by: QiuChunshuo <qiuchunshuo@huawei.com>

Jingchun Gao committed Nov 14, 2025 at 19:24 UTC 4516d44b7f990b8f92450e73720b89cc8ac155ca

5 files changed +331 -51

tests/distributed/test_context_parallel.py

+15 -2

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

39 class CPTestOptions(NamedTuple):

40 multi_node_only: bool

41 load_format: str | None = None

42 + attn_backend: str | None = None

45 @dataclass

@@ -58,6 +59,7 @@ class CPTestSettings:

59 multi_node_only: bool = False,

60 runner: RunnerOption = "auto",

61 load_format: str | None = None,

62 + attn_backend: str | None = None,

63 ):

64 parallel_setups = []

65 for eager_mode_val in [False]:

@@ -79,7 +81,9 @@ class CPTestSettings:

81 distributed_backends=["mp"],

82 runner=runner,

83 test_options=CPTestOptions(

82 - multi_node_only=multi_node_only, load_format=load_format

84 + multi_node_only=multi_node_only,

85 + load_format=load_format,

86 + attn_backend=attn_backend,

87 ),

88 )

@@ -117,7 +121,7 @@ def _compare_cp_with_tp(

121 chunked_prefill,

122 ) = parallel_setup

123

120 - multi_node_only, load_format = test_options

124 + multi_node_only, load_format, attn_backend = test_options

125

126 model_info = HF_EXAMPLE_MODELS.find_hf_info(model_id)

127 model_info.check_transformers_version(on_fail="skip")

@@ -177,6 +181,13 @@ def _compare_cp_with_tp(

181 if hf_overrides:

182 common_args.extend(["--hf-overrides", json.dumps(hf_overrides)])

183

184 + if not attn_backend:

185 + cp_env = tp_env = {}

186 + else:

187 + cp_env = tp_env = {

188 + "VLLM_ATTENTION_BACKEND": attn_backend,

189 + }

190 +

191 cp_args = [

192 *common_args,

193 "--tensor-parallel-size",

@@ -205,6 +216,8 @@ def _compare_cp_with_tp(

216 model_id,

217 cp_args,

218 tp_args,

219 + cp_env,

220 + tp_env,

221 method=method,

222 max_wait_seconds=720,

223 )

vllm/config/model.py

@@ -1183,6 +1183,14 @@ class ModelConfig:

1183 f"but got {decode_context_parallel_size}"

1184 )

1185

1186 + num_q_per_kv = total_num_attention_heads // total_num_kv_heads

1187 + assert num_q_per_kv % decode_context_parallel_size == 0, (

1188 + f"Total number of q per kv attn heads ({num_q_per_kv})"

1189 + " must be divisible by dcp world size when enable "

1190 + "decode context parallel for GQA "

1191 + f"({parallel_config.decode_context_parallel_size})."

1192 + )

1193 +

1194 def get_sliding_window(self) -> int | None:

1195 """Get the sliding window size from the HF text config if present."""

1196 return getattr(self.hf_text_config, "sliding_window", None)

vllm/utils/flashinfer.py

@@ -259,6 +259,7 @@ def use_trtllm_attention(

259 num_kv_heads: int,

260 num_tokens: int,

261 max_seq_len: int,

262 + dcp_world_size: int,

263 kv_cache_dtype: str,

264 q_dtype: torch.dtype,

265 is_prefill: bool,

@@ -272,6 +273,14 @@ def use_trtllm_attention(

273 if force_use_trtllm is not None and not force_use_trtllm:

274 return False

275

276 + # Decode context parallel is not supported

277 + if dcp_world_size > 1:

278 + logger.warning_once(

279 + "Trtllm does not support returning LSE and as a result "

280 + "does not support DCP, reverting to FlashInfer"

281 + )

282 + return False

283 +

284 # The platform is not supported

285 if not supports_trtllm_attention():

286 if force_use_trtllm:

vllm/v1/attention/backends/flashinfer.py

+294 -49

@@ -10,6 +10,7 @@ import torch

10 from flashinfer import (

11 BatchDecodeWithPagedKVCacheWrapper,

12 BatchPrefillWithPagedKVCacheWrapper,

13 + BatchPrefillWithRaggedKVCacheWrapper,

14 MultiLevelCascadeAttentionWrapper,

15 )

16 from flashinfer.decode import _get_range_buf, trtllm_batch_decode_with_kv_cache

@@ -24,8 +25,11 @@ from vllm.attention.backends.abstract import (

25 AttentionType,

26 MultipleOf,

27 )

28 +from vllm.attention.ops.common import cp_lse_ag_out_rs

29 +from vllm.attention.ops.merge_attn_states import merge_attn_states

30 from vllm.config import CUDAGraphMode, VllmConfig

31 from vllm.config.cache import CacheDType

32 +from vllm.distributed.parallel_state import get_dcp_group

33 from vllm.logger import init_logger

34 from vllm.model_executor.layers.batch_invariant import (

35 vllm_is_batch_invariant,

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

54 AttentionMetadataBuilder,

55 CommonAttentionMetadata,

56 KVCacheLayoutType,

57 + get_dcp_local_seq_lens,

58 get_kv_cache_layout,

59 get_per_layer_parameters,

60 infer_global_hyperparameters,

@@ -160,6 +165,113 @@ def trtllm_prefill_attn_kvfp8_dequant(

165 return mock_kv_cache, mock_block_table

166

167

168 +class BatchDCPPrefillWrapper:

169 + def __init__(

170 + self,

171 + workspace_buffer: torch.Tensor | None = None,

172 + ):

173 + self._context = BatchPrefillWithPagedKVCacheWrapper(

174 + workspace_buffer, get_kv_cache_layout()

175 + )

176 + self._new_tokens = BatchPrefillWithRaggedKVCacheWrapper(

177 + workspace_buffer, get_kv_cache_layout()

178 + )

179 +

180 + def plan(

181 + self,

182 + qo_indptr_cpu: torch.Tensor,

183 + paged_kv_indptr_cpu: torch.Tensor,

184 + paged_kv_indices: torch.Tensor,

185 + paged_kv_last_page_len_cpu: torch.Tensor,

186 + prefill_start: int,

187 + page_size: int,

188 + num_qo_heads: int,

189 + dcp_world_size: int,

190 + num_kv_heads: int,

191 + head_dim: int,

192 + sm_scale: float,

193 + window_left: int,

194 + logits_soft_cap: float | None,

195 + q_data_type: torch.dtype,

196 + kv_cache_dtype: torch.dtype,

197 + prefill_fixed_split_size: int,

198 + disable_split_kv: bool,

199 + ):

200 + """Plan the prefill operation with given parameters."""

201 + self._context.plan(

202 + qo_indptr_cpu,

203 + paged_kv_indptr_cpu,

204 + paged_kv_indices,

205 + paged_kv_last_page_len_cpu[prefill_start:],

206 + num_qo_heads * dcp_world_size,

207 + num_kv_heads,

208 + head_dim,

209 + page_size,

210 + causal=False, # This is context run

211 + sm_scale=sm_scale,

212 + window_left=window_left,

213 + logits_soft_cap=logits_soft_cap,

214 + q_data_type=q_data_type,

215 + kv_data_type=kv_cache_dtype,

216 + fixed_split_size=prefill_fixed_split_size,

217 + disable_split_kv=disable_split_kv,

218 + )

219 + self._new_tokens.plan(

220 + qo_indptr=qo_indptr_cpu,

221 + kv_indptr=qo_indptr_cpu,

222 + num_qo_heads=num_qo_heads,

223 + num_kv_heads=num_kv_heads,

224 + head_dim_qk=head_dim,

225 + head_dim_vo=head_dim,

226 + causal=True, # This is newtokens run

227 + sm_scale=sm_scale,

228 + window_left=window_left,

229 + logits_soft_cap=logits_soft_cap,

230 + q_data_type=q_data_type,

231 + )

232 +

233 + def run(

234 + self,

235 + layer: torch.nn.Module,

236 + prefill_query: torch.Tensor,

237 + kv_cache_permute: torch.Tensor,

238 + key: torch.Tensor,

239 + value: torch.Tensor,

240 + out: torch.Tensor,

241 + ):

242 + prefill_query_across_dcp = get_dcp_group().all_gather(

243 + prefill_query.contiguous(), dim=1

244 + )

245 + output_context_tmp, lse_context_tmp = self._context.run(

246 + prefill_query_across_dcp,

247 + kv_cache_permute,

248 + k_scale=layer._k_scale_float,

249 + v_scale=layer._v_scale_float,

250 + return_lse=True,

251 + )

252 + output_context, lse_context = cp_lse_ag_out_rs(

253 + output_context_tmp, lse_context_tmp, get_dcp_group(), return_lse=True

254 + )

255 + lse_context = lse_context.transpose(0, 1).contiguous()

256 +

257 + output_query, lse_query = self._new_tokens.run(

258 + prefill_query,

259 + key,

260 + value,

261 + return_lse=True,

262 + )

263 + lse_query = lse_query.transpose(0, 1).contiguous()

264 +

265 + merge_attn_states(

266 + out,

267 + output_context,

268 + lse_context,

269 + output_query,

270 + lse_query,

271 + )

272 + return out

273 +

274 +

275 class FlashInferBackend(AttentionBackend):

276 accept_output_buffer: bool = True

277 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

@@ -281,7 +393,9 @@ class FlashInferMetadata:

393 # For cascade attention (CPU for planning).

394 use_cascade: bool

395

284 - prefill_wrapper: BatchPrefillWithPagedKVCacheWrapper | None = None

396 + prefill_wrapper: (

397 + BatchPrefillWithPagedKVCacheWrapper | BatchDCPPrefillWrapper | None

398 + ) = None

399 decode_wrapper: BatchDecodeWithPagedKVCacheWrapper | None = None

400 cascade_wrapper: MultiLevelCascadeAttentionWrapper | None = None

401

@@ -303,7 +417,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

417 self.cache_config = vllm_config.cache_config

418 self.model_config = vllm_config.model_config

419 self._workspace_buffer = None

306 - self._prefill_wrapper = None # Wrapper for prefill/append

420 + self._prefill_wrapper: (

421 + BatchPrefillWithPagedKVCacheWrapper | BatchDCPPrefillWrapper | None

422 + ) = None # Wrapper for prefill/append

423 self._decode_wrapper = None # Wrapper for decode (general shape)

424

425 if vllm_is_batch_invariant():

@@ -341,9 +457,23 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

457 self.compilation_config.max_cudagraph_capture_size,

458 )

459

344 - self.num_qo_heads = self.model_config.get_num_attention_heads(

345 - self.vllm_config.parallel_config

460 + try:

461 + self.dcp_world_size = get_dcp_group().world_size

462 + self.dcp_rank = get_dcp_group().rank_in_group

463 + self.dcp_kv_cache_interleave_size = (

464 + vllm_config.parallel_config.dcp_kv_cache_interleave_size

465 + )

466 + except AssertionError:

467 + # DCP might not be initialized in testing

468 + self.dcp_world_size = 1

469 + self.dcp_rank = 0

470 + self.dcp_kv_cache_interleave_size = 1

471 +

472 + self.num_qo_heads = (

473 + self.model_config.get_num_attention_heads(self.vllm_config.parallel_config)

474 + * self.dcp_world_size

475 )

476 +

477 self.num_kv_heads = self.kv_cache_spec.num_kv_heads

478 self.head_dim = self.kv_cache_spec.head_size

479 self.page_size = self.kv_cache_spec.block_size

@@ -455,11 +585,19 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

585 )

586 return self._workspace_buffer

587

458 - def _get_prefill_wrapper(self):

588 + def _get_prefill_wrapper(

589 + self,

590 + ) -> BatchPrefillWithPagedKVCacheWrapper | BatchDCPPrefillWrapper:

591 if self._prefill_wrapper is None:

460 - self._prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper(

461 - self._get_workspace_buffer(), get_kv_cache_layout()

462 - )

592 + if self.dcp_world_size > 1:

593 + self._prefill_wrapper = BatchDCPPrefillWrapper(

594 + workspace_buffer=self._get_workspace_buffer(),

595 + )

596 + else:

597 + self._prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper(

598 + self._get_workspace_buffer(), get_kv_cache_layout()

599 + )

600 + assert self._prefill_wrapper is not None

601 return self._prefill_wrapper

602

603 def _get_decode_wrapper(self, batch_size: int, use_cudagraph: bool = False):

@@ -526,9 +664,29 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

664 max_seq_len = common_attn_metadata.max_seq_len

665 seq_lens = common_attn_metadata.seq_lens

666 seq_lens_cpu = common_attn_metadata.seq_lens_cpu

529 - seq_lens_np = seq_lens_cpu.numpy()

667 block_table_tensor = common_attn_metadata.block_table_tensor

668 + qo_indptr_cpu = common_attn_metadata.query_start_loc_cpu

669

670 + if self.dcp_world_size > 1:

671 + if num_prefills > 0:

672 + qo_indptr_prefill_cpu = (

673 + qo_indptr_cpu[num_decodes:] - qo_indptr_cpu[num_decodes]

674 + )

675 + query_lens_prefill_cpu = (

676 + qo_indptr_prefill_cpu[1:] - qo_indptr_prefill_cpu[:-1]

677 + )

678 + seq_lens_cpu[num_decodes:] = (

679 + seq_lens_cpu[num_decodes:] - query_lens_prefill_cpu

680 + )

681 +

682 + seq_lens_cpu = get_dcp_local_seq_lens(

683 + seq_lens_cpu,

684 + self.dcp_world_size,

685 + self.dcp_rank,

686 + self.dcp_kv_cache_interleave_size,

687 + )

688 +

689 + seq_lens_np = seq_lens_cpu.numpy()

690 num_blocks_np = (seq_lens_np + (page_size - 1)) // page_size

691

692 use_cascade = common_prefix_len > 0

@@ -589,7 +747,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

747 # write self.paged_kv_last_page_len_cpu inplace

748 paged_kv_last_page_len_np = seq_lens_np % page_size

749 self.paged_kv_last_page_len_np[:num_reqs] = np.where(

592 - paged_kv_last_page_len_np == 0,

750 + (paged_kv_last_page_len_np == 0) & (seq_lens_np != 0),

751 page_size,

752 paged_kv_last_page_len_np,

753 )

@@ -600,13 +758,16 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

758 self.num_kv_heads,

759 num_prefill_tokens,

760 max_seq_len,

761 + self.dcp_world_size,

762 self.cache_dtype,

763 self.q_data_type,

764 is_prefill=True,

765 has_sinks=self.has_sinks,

766 has_spec=uses_spec_reorder,

767 )

609 - decode_use_trtllm = self.use_trtllm_decode_attention

768 + decode_use_trtllm = (

769 + self.use_trtllm_decode_attention and self.dcp_world_size <= 1

770 + )

771

772 if not (prefill_use_trtllm and decode_use_trtllm):

773 if self.has_sinks:

@@ -651,7 +812,6 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

812 use_cascade=use_cascade,

813 )

814

654 - qo_indptr_cpu = common_attn_metadata.query_start_loc_cpu

815 paged_kv_indptr_cpu = self.paged_kv_indptr_cpu[: 1 + num_reqs]

816 paged_kv_last_page_len_cpu = self.paged_kv_last_page_len_cpu[:num_reqs]

817

@@ -703,24 +863,52 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

863 attn_metadata.max_q_len_prefill = int(query_lens_prefill.max().item())

864

865 if not attn_metadata.prefill_use_trtllm:

706 - attn_metadata.prefill_wrapper.plan(

707 - qo_indptr_cpu,

708 - paged_kv_indptr_cpu,

709 - paged_kv_indices,

710 - paged_kv_last_page_len_cpu[prefill_start:],

711 - self.num_qo_heads,

712 - self.num_kv_heads,

713 - self.head_dim,

714 - self.page_size,

715 - causal=True,

716 - sm_scale=self.sm_scale,

717 - window_left=self.window_left,

718 - logits_soft_cap=self.logits_soft_cap,

719 - q_data_type=self.q_data_type,

720 - kv_data_type=self.kv_cache_dtype,

721 - fixed_split_size=self.prefill_fixed_split_size,

722 - disable_split_kv=self.disable_split_kv,

723 - )

866 + if self.dcp_world_size > 1:

867 + assert isinstance(

868 + attn_metadata.prefill_wrapper, BatchDCPPrefillWrapper

869 + )

870 + attn_metadata.prefill_wrapper.plan(

871 + qo_indptr_cpu=qo_indptr_cpu,

872 + paged_kv_indptr_cpu=paged_kv_indptr_cpu,

873 + paged_kv_indices=paged_kv_indices,

874 + paged_kv_last_page_len_cpu=paged_kv_last_page_len_cpu,

875 + prefill_start=prefill_start,

876 + page_size=self.page_size,

877 + num_qo_heads=self.num_qo_heads,

878 + dcp_world_size=self.dcp_world_size,

879 + num_kv_heads=self.num_kv_heads,

880 + head_dim=self.head_dim,

881 + sm_scale=self.sm_scale,

882 + window_left=self.window_left,

883 + logits_soft_cap=self.logits_soft_cap,

884 + q_data_type=self.q_data_type,

885 + kv_cache_dtype=self.kv_cache_dtype,

886 + prefill_fixed_split_size=self.prefill_fixed_split_size,

887 + disable_split_kv=self.disable_split_kv,

888 + )

889 + else:

890 + assert isinstance(

891 + attn_metadata.prefill_wrapper,

892 + BatchPrefillWithPagedKVCacheWrapper,

893 + )

894 + attn_metadata.prefill_wrapper.plan(

895 + qo_indptr_cpu,

896 + paged_kv_indptr_cpu,

897 + paged_kv_indices,

898 + paged_kv_last_page_len_cpu[prefill_start:],

899 + self.num_qo_heads,

900 + self.num_kv_heads,

901 + self.head_dim,

902 + self.page_size,

903 + causal=True,

904 + sm_scale=self.sm_scale,

905 + window_left=self.window_left,

906 + logits_soft_cap=self.logits_soft_cap,

907 + q_data_type=self.q_data_type,

908 + kv_data_type=self.kv_cache_dtype,

909 + fixed_split_size=self.prefill_fixed_split_size,

910 + disable_split_kv=self.disable_split_kv,

911 + )

912 else:

913 attn_metadata.qo_indptr_gpu = qo_indptr_cpu.to(

914 self.device, non_blocking=True

@@ -770,7 +958,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

958 paged_kv_indices,

959 self.paged_kv_last_page_len_cpu[:num_input_tokens],

960 seq_lens_cpu[:num_input_tokens],

773 - self.num_qo_heads,

961 + self.num_qo_heads * self.dcp_world_size,

962 self.num_kv_heads,

963 self.head_dim,

964 self.page_size,

@@ -797,6 +985,8 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

985

986

987 class FlashInferImpl(AttentionImpl):

988 + can_return_lse_for_decode: bool = True

989 +

990 def __init__(

991 self,

992 num_heads: int,

@@ -989,6 +1179,8 @@ class FlashInferImpl(AttentionImpl):

1179

1180 # Inputs and outputs may be padded for CUDA graphs

1181 query = query[:num_actual_tokens]

1182 + key = key[:num_actual_tokens]

1183 + value = value[:num_actual_tokens]

1184 output_padded = output

1185 output = output[:num_actual_tokens]

1186

@@ -1015,17 +1207,46 @@ class FlashInferImpl(AttentionImpl):

1207 assert prefill_wrapper is not None

1208

1209 if not attn_metadata.prefill_use_trtllm:

1018 - assert prefill_wrapper._causal

1019 - assert prefill_wrapper._window_left == self.window_left

1020 - assert prefill_wrapper._logits_soft_cap == (self.logits_soft_cap or 0.0)

1021 - assert prefill_wrapper._sm_scale == self.scale

1022 - prefill_wrapper.run(

1023 - prefill_query,

1024 - kv_cache_permute,

1025 - k_scale=layer._k_scale_float,

1026 - v_scale=layer._v_scale_float,

1027 - out=output[num_decode_tokens:],

1028 - )

1210 + if self.dcp_world_size > 1:

1211 + assert isinstance(prefill_wrapper, BatchDCPPrefillWrapper)

1212 + assert prefill_wrapper._context._window_left == self.window_left

1213 + assert prefill_wrapper._context._logits_soft_cap == (

1214 + self.logits_soft_cap or 0.0

1215 + )

1216 + assert prefill_wrapper._context._sm_scale == self.scale

1217 + assert not prefill_wrapper._context._causal

1218 + assert prefill_wrapper._new_tokens._window_left == self.window_left

1219 + assert prefill_wrapper._new_tokens._logits_soft_cap == (

1220 + self.logits_soft_cap or 0.0

1221 + )

1222 + assert prefill_wrapper._new_tokens._sm_scale == self.scale

1223 + assert prefill_wrapper._new_tokens._causal

1224 +

1225 + prefill_wrapper.run(

1226 + layer,

1227 + prefill_query,

1228 + kv_cache_permute,

1229 + key[num_decode_tokens:],

1230 + value[num_decode_tokens:],

1231 + out=output[num_decode_tokens:],

1232 + )

1233 + else:

1234 + assert isinstance(

1235 + prefill_wrapper, BatchPrefillWithPagedKVCacheWrapper

1236 + )

1237 + assert prefill_wrapper._window_left == self.window_left

1238 + assert prefill_wrapper._logits_soft_cap == (

1239 + self.logits_soft_cap or 0.0

1240 + )

1241 + assert prefill_wrapper._sm_scale == self.scale

1242 + assert prefill_wrapper._causal

1243 + prefill_wrapper.run(

1244 + prefill_query,

1245 + kv_cache_permute,

1246 + k_scale=layer._k_scale_float,

1247 + v_scale=layer._v_scale_float,

1248 + out=output[num_decode_tokens:],

1249 + )

1250 else:

1251 # prefill_query may be non-contiguous

1252 prefill_query = prefill_query.contiguous()

@@ -1101,13 +1322,37 @@ class FlashInferImpl(AttentionImpl):

1322 assert decode_wrapper._window_left == self.window_left

1323 assert decode_wrapper._logits_soft_cap == (self.logits_soft_cap or 0.0)

1324 assert decode_wrapper._sm_scale == self.scale

1104 - decode_wrapper.run(

1105 - decode_query,

1106 - kv_cache_permute,

1107 - k_scale=layer._k_scale_float,

1108 - v_scale=layer._v_scale_float,

1109 - out=output[:num_decode_tokens],

1110 - )

1325 +

1326 + if self.dcp_world_size > 1:

1327 + decode_query = get_dcp_group().all_gather(

1328 + decode_query.contiguous(), dim=-2

1329 + )

1330 + output_tmp = torch.empty_like(decode_query)

1331 + lse = torch.empty(

1332 + (decode_query.size(0), decode_query.size(1)),

1333 + dtype=torch.float32,

1334 + device=decode_query.device,

1335 + )

1336 + decode_wrapper.run(

1337 + decode_query,

1338 + kv_cache_permute,

1339 + k_scale=layer._k_scale_float,

1340 + v_scale=layer._v_scale_float,

1341 + out=output_tmp,

1342 + lse=lse,

1343 + return_lse=True,

1344 + )

1345 + output[:num_decode_tokens] = cp_lse_ag_out_rs(

1346 + output_tmp, lse, get_dcp_group()

1347 + )

1348 + else:

1349 + decode_wrapper.run(

1350 + decode_query,

1351 + kv_cache_permute,

1352 + k_scale=layer._k_scale_float,

1353 + v_scale=layer._v_scale_float,

1354 + out=output[:num_decode_tokens],

1355 + )

1356 else:

1357 # decode_query may be non-contiguous

1358 decode_query = decode_query.contiguous()

vllm/v1/executor/multiproc_executor.py

@@ -31,6 +31,7 @@ from vllm.distributed import destroy_distributed_environment, destroy_model_para

31 from vllm.distributed.device_communicators.shm_broadcast import Handle, MessageQueue

32 from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator

33 from vllm.distributed.parallel_state import (

34 + get_dcp_group,

35 get_dp_group,

36 get_ep_group,

37 get_pp_group,

@@ -726,6 +727,8 @@ class WorkerProc:

727 pp_rank = get_pp_group().rank_in_group

728 tp_size = get_tp_group().world_size

729 tp_rank = get_tp_group().rank_in_group

730 + dcp_size = get_dcp_group().world_size

731 + dcp_rank = get_dcp_group().rank_in_group

732 process_name = "Worker"

733 if dp_size > 1:

734 process_name += f"_DP{dp_rank}"

@@ -733,6 +736,8 @@ class WorkerProc:

736 process_name += f"_PP{pp_rank}"

737 if tp_size > 1:

738 process_name += f"_TP{tp_rank}"

739 + if dcp_size > 1:

740 + process_name += f"_DCP{dcp_rank}"

741 if enable_ep:

742 ep_rank = get_ep_group().rank_in_group

743 process_name += f"_EP{ep_rank}"