@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6cb901093f
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 7, 06:30 UTC

[Core] Add All-to-All communication backend for DCP (#34883)

Signed-off-by: Sungsoo Ha <sungsooh@nvidia.com> Signed-off-by: sungsoo ha <hasungsoo@gmail.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

sungsoo ha committed Mar 4, 2026 at 07:01 UTC 6cb901093f3df8e26cbc0a8a0e1a884f4dbaa5ea
8 files changed +658 -17
tests/distributed/test_dcp_a2a.py new
+192
@@ -0,0 +1,192 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for DCP A2A communication backend (no GPU required).
4 +
5 +Tests cover:
6 +1. DCP A2A config validation (--dcp-comm-backend)
7 +2. KVP group function exists
8 +3. LSE-weighted combination correctness
9 +"""
10 +
11 +import math
12 +
13 +import pytest
14 +import torch
15 +
16 +from vllm.config.parallel import ParallelConfig
17 +
18 +
19 +class TestDCPCommBackendConfig:
20 + """Test --dcp-comm-backend config validation."""
21 +
22 + def test_default_is_ag_rs(self):
23 + """Default comm backend is ag_rs."""
24 + config = ParallelConfig()
25 + assert config.dcp_comm_backend == "ag_rs"
26 +
27 + def test_a2a_requires_dcp_greater_than_1(self):
28 + """A2A backend requires decode_context_parallel_size > 1."""
29 + with pytest.raises(
30 + ValueError, match="requires decode_context_parallel_size > 1"
31 + ):
32 + ParallelConfig(
33 + dcp_comm_backend="a2a",
34 + decode_context_parallel_size=1,
35 + )
36 +
37 + def test_a2a_with_dcp_valid(self):
38 + """A2A backend is valid when DCP > 1."""
39 + config = ParallelConfig(
40 + dcp_comm_backend="a2a",
41 + tensor_parallel_size=8,
42 + decode_context_parallel_size=4,
43 + )
44 + assert config.dcp_comm_backend == "a2a"
45 +
46 + def test_invalid_backend_rejected(self):
47 + """Invalid backend values are rejected."""
48 + with pytest.raises(ValueError, match="must be one of"):
49 + ParallelConfig(
50 + dcp_comm_backend="invalid",
51 + )
52 +
53 + def test_ag_rs_with_dcp_1_valid(self):
54 + """ag_rs backend is valid with DCP=1 (no DCP)."""
55 + config = ParallelConfig(
56 + dcp_comm_backend="ag_rs",
57 + decode_context_parallel_size=1,
58 + )
59 + assert config.dcp_comm_backend == "ag_rs"
60 +
61 +
62 +class TestLSEWeightedCombine:
63 + """Test LSE-weighted combination logic (CPU only, no GPU).
64 +
65 + The _lse_weighted_combine function is the reference implementation
66 + that verifies the Triton kernel's correctness. It computes:
67 +
68 + result[b,h,d] = sum_n(w_n * output_n[b,h,d])
69 +
70 + where w_n = softmax(lse_n) = exp(lse_n) / sum_k(exp(lse_k))
71 + """
72 +
73 + def test_importable(self):
74 + """Verify _lse_weighted_combine is importable."""
75 + from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
76 +
77 + assert callable(_lse_weighted_combine)
78 +
79 + def test_single_rank(self):
80 + """Single rank: output unchanged."""
81 + from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
82 +
83 + # N=1, B=2, H=4, D=8
84 + outputs = torch.randn(1, 2, 4, 8)
85 + lses = torch.randn(1, 2, 4)
86 +
87 + result = _lse_weighted_combine(outputs, lses)
88 +
89 + assert result.shape == (2, 4, 8)
90 + torch.testing.assert_close(result, outputs.squeeze(0), rtol=1e-5, atol=1e-5)
91 +
92 + def test_equal_lse(self):
93 + """Equal LSE values: outputs averaged equally."""
94 + from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
95 +
96 + _N, B, H, D = 2, 1, 1, 4
97 + outputs = torch.tensor(
98 + [
99 + [[[1.0, 2.0, 3.0, 4.0]]], # Rank 0
100 + [[[5.0, 6.0, 7.0, 8.0]]], # Rank 1
101 + ]
102 + )
103 + lses = torch.tensor(
104 + [
105 + [[0.0]], # Rank 0
106 + [[0.0]], # Rank 1
107 + ]
108 + )
109 +
110 + result = _lse_weighted_combine(outputs, lses)
111 +
112 + expected = (outputs[0] + outputs[1]) / 2
113 + assert result.shape == (B, H, D)
114 + torch.testing.assert_close(result, expected, rtol=1e-5, atol=1e-5)
115 +
116 + def test_dominant_rank(self):
117 + """Different LSE values: larger LSE gets more weight."""
118 + from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
119 +
120 + B, H, D = 1, 1, 2
121 + outputs = torch.tensor(
122 + [
123 + [[[0.0, 0.0]]], # Rank 0
124 + [[[1.0, 1.0]]], # Rank 1
125 + ]
126 + )
127 + lses = torch.tensor(
128 + [
129 + [[-100.0]], # Rank 0: negligible contribution
130 + [[0.0]], # Rank 1: dominant
131 + ]
132 + )
133 +
134 + result = _lse_weighted_combine(outputs, lses)
135 +
136 + assert result.shape == (B, H, D)
137 + torch.testing.assert_close(result, outputs[1].squeeze(0), atol=1e-5, rtol=1e-5)
138 +
139 + def test_mathematically_correct(self):
140 + """Verify mathematical correctness of LSE combination."""
141 + from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
142 +
143 + outputs = torch.tensor(
144 + [
145 + [[[2.0, 4.0]]],
146 + [[[6.0, 8.0]]],
147 + ]
148 + )
149 + lses = torch.tensor(
150 + [
151 + [[1.0]], # exp(1) ≈ 2.718
152 + [[2.0]], # exp(2) ≈ 7.389
153 + ]
154 + )
155 +
156 + result = _lse_weighted_combine(outputs, lses)
157 +
158 + w0 = math.exp(1) / (math.exp(1) + math.exp(2))
159 + w1 = math.exp(2) / (math.exp(1) + math.exp(2))
160 + expected = torch.tensor([[[w0 * 2.0 + w1 * 6.0, w0 * 4.0 + w1 * 8.0]]])
161 +
162 + torch.testing.assert_close(result, expected, rtol=1e-4, atol=1e-4)
163 +
164 + def test_return_lse(self):
165 + """return_lse=True returns global LSE (logsumexp of inputs)."""
166 + from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
167 +
168 + B, H, D = 1, 1, 2
169 + outputs = torch.tensor(
170 + [
171 + [[[1.0, 2.0]]],
172 + [[[3.0, 4.0]]],
173 + ]
174 + )
175 + lses = torch.tensor(
176 + [
177 + [[1.0]],
178 + [[2.0]],
179 + ]
180 + )
181 +
182 + result, global_lse = _lse_weighted_combine(outputs, lses, return_lse=True)
183 +
184 + expected_global_lse = math.log(math.exp(1) + math.exp(2))
185 +
186 + assert result.shape == (B, H, D)
187 + assert global_lse.shape == (B, H)
188 + assert abs(global_lse.item() - expected_global_lse) < 1e-5
189 +
190 +
191 +if __name__ == "__main__":
192 + pytest.main([__file__, "-v"])
vllm/config/parallel.py
+14
@@ -36,6 +36,7 @@ ExpertPlacementStrategy = Literal["linear", "round_robin"]
36 DistributedExecutorBackend = Literal["ray", "mp", "uni", "external_launcher"]
37 DataParallelBackend = Literal["ray", "mp"]
38 EPLBPolicyOption = Literal["default"]
39 +DCPCommBackend = Literal["ag_rs", "a2a"]
40 All2AllBackend = Literal[
41 "naive",
42 "pplx",
@@ -287,6 +288,14 @@ class ParallelConfig:
288 and will be deprecated when PCP is fully supported.
289
290 """
291 + dcp_comm_backend: DCPCommBackend = "ag_rs"
292 + """Communication backend for Decode Context Parallel (DCP).
293 + - "ag_rs": AllGather + ReduceScatter (default, existing behavior)
294 + - "a2a": All-to-All exchange of partial outputs + LSE, then
295 + combine with Triton kernel. Reduces NCCL calls from 3 to 2
296 + per layer for MLA models.
297 + """
298 +
299 cp_kv_cache_interleave_size: int = 1
300 """Interleave size of kv_cache storage while using DCP or PCP.
301 For `total_cp_rank = pcp_rank * dcp_world_size + dcp_rank`,
@@ -392,6 +401,11 @@ class ParallelConfig:
401 f"dcp_size={self.decode_context_parallel_size}."
402 )
403
404 + if self.dcp_comm_backend == "a2a" and self.decode_context_parallel_size <= 1:
405 + raise ValueError(
406 + "dcp_comm_backend='a2a' requires decode_context_parallel_size > 1."
407 + )
408 +
409 return self
410
411 @property
vllm/config/vllm.py
+2
@@ -1645,6 +1645,8 @@ class VllmConfig:
1645 f"tensor_parallel_size={self.parallel_config.tensor_parallel_size}, " # noqa
1646 f"pipeline_parallel_size={self.parallel_config.pipeline_parallel_size}, " # noqa
1647 f"data_parallel_size={self.parallel_config.data_parallel_size}, " # noqa
1648 + f"decode_context_parallel_size={self.parallel_config.decode_context_parallel_size}, " # noqa
1649 + f"dcp_comm_backend={self.parallel_config.dcp_comm_backend}, " # noqa
1650 f"disable_custom_all_reduce={self.parallel_config.disable_custom_all_reduce}, " # noqa
1651 f"quantization={self.model_config.quantization}, "
1652 f"enforce_eager={self.model_config.enforce_eager}, "
vllm/engine/arg_utils.py
+7
@@ -85,6 +85,7 @@ from vllm.config.observability import DetailedTraceModules
85 from vllm.config.parallel import (
86 All2AllBackend,
87 DataParallelBackend,
88 + DCPCommBackend,
89 DistributedExecutorBackend,
90 ExpertPlacementStrategy,
91 )
@@ -405,6 +406,7 @@ class EngineArgs:
406 tensor_parallel_size: int = ParallelConfig.tensor_parallel_size
407 prefill_context_parallel_size: int = ParallelConfig.prefill_context_parallel_size
408 decode_context_parallel_size: int = ParallelConfig.decode_context_parallel_size
409 + dcp_comm_backend: DCPCommBackend = ParallelConfig.dcp_comm_backend
410 dcp_kv_cache_interleave_size: int = ParallelConfig.dcp_kv_cache_interleave_size
411 cp_kv_cache_interleave_size: int = ParallelConfig.cp_kv_cache_interleave_size
412 data_parallel_size: int = ParallelConfig.data_parallel_size
@@ -820,6 +822,10 @@ class EngineArgs:
822 "-dcp",
823 **parallel_kwargs["decode_context_parallel_size"],
824 )
825 + parallel_group.add_argument(
826 + "--dcp-comm-backend",
827 + **parallel_kwargs["dcp_comm_backend"],
828 + )
829 parallel_group.add_argument(
830 "--dcp-kv-cache-interleave-size",
831 **parallel_kwargs["dcp_kv_cache_interleave_size"],
@@ -1720,6 +1726,7 @@ class EngineArgs:
1726 worker_cls=self.worker_cls,
1727 worker_extension_cls=self.worker_extension_cls,
1728 decode_context_parallel_size=self.decode_context_parallel_size,
1729 + dcp_comm_backend=self.dcp_comm_backend,
1730 dcp_kv_cache_interleave_size=self.dcp_kv_cache_interleave_size,
1731 cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
1732 _api_process_count=self._api_process_count,
vllm/model_executor/layers/attention/mla_attention.py
+33 -8
@@ -203,8 +203,17 @@ from tqdm import tqdm
203 import vllm.envs as envs
204 from vllm import _custom_ops as ops
205 from vllm._aiter_ops import rocm_aiter_ops
206 -from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config
207 -from vllm.distributed.parallel_state import get_dcp_group, is_global_first_rank
206 +from vllm.config import (
207 + CacheConfig,
208 + ModelConfig,
209 + VllmConfig,
210 + get_current_vllm_config,
211 + get_current_vllm_config_or_none,
212 +)
213 +from vllm.distributed.parallel_state import (
214 + get_dcp_group,
215 + is_global_first_rank,
216 +)
217 from vllm.forward_context import ForwardContext, get_forward_context
218 from vllm.logger import init_logger
219 from vllm.model_executor.custom_op import CustomOp
@@ -253,6 +262,7 @@ from vllm.v1.attention.backends.utils import (
262 split_decodes_and_prefills,
263 )
264 from vllm.v1.attention.ops.common import cp_lse_ag_out_rs
265 +from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
266 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
267 from vllm.v1.attention.selector import get_attn_backend
268 from vllm.v1.kv_cache_interface import (
@@ -393,6 +403,13 @@ class MLAAttention(nn.Module, AttentionLayerBase):
403
404 self.use_sparse = use_sparse
405
406 + vllm_config = get_current_vllm_config_or_none()
407 + self.dcp_a2a = (
408 + vllm_config is not None
409 + and vllm_config.parallel_config.decode_context_parallel_size > 1
410 + and vllm_config.parallel_config.dcp_comm_backend == "a2a"
411 + )
412 +
413 # Initialize q/k/v range constants.
414 self.q_range = torch.tensor(envs.Q_SCALE_CONSTANT, dtype=torch.float32)
415 self.k_range = torch.tensor(envs.K_SCALE_CONSTANT, dtype=torch.float32)
@@ -647,12 +664,20 @@ class MLAAttention(nn.Module, AttentionLayerBase):
664
665 # correct dcp attn_out with lse.
666 if self.impl.dcp_world_size > 1:
650 - attn_out = cp_lse_ag_out_rs(
651 - attn_out,
652 - lse,
653 - get_dcp_group(),
654 - is_lse_base_on_e=not getattr(self, "_use_fi_prefill", False),
655 - )
667 + if self.dcp_a2a:
668 + attn_out = dcp_a2a_lse_reduce(
669 + attn_out,
670 + lse,
671 + get_dcp_group(),
672 + is_lse_base_on_e=not getattr(self, "_use_fi_prefill", False),
673 + )
674 + else:
675 + attn_out = cp_lse_ag_out_rs(
676 + attn_out,
677 + lse,
678 + get_dcp_group(),
679 + is_lse_base_on_e=not getattr(self, "_use_fi_prefill", False),
680 + )
681
682 # v_up projection
683 self._v_up_proj(attn_out, out=mqa_output_slice)
vllm/v1/attention/backends/flash_attn.py
+17 -3
@@ -23,6 +23,7 @@ from vllm.v1.attention.backends.fa_utils import (
23 is_flash_attn_varlen_func_available,
24 )
25 from vllm.v1.attention.ops.common import cp_lse_ag_out_rs
26 +from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
27 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
28
29 if is_flash_attn_varlen_func_available():
@@ -32,7 +33,12 @@ if is_flash_attn_varlen_func_available():
33 get_scheduler_metadata,
34 reshape_and_cache_flash,
35 )
35 -from vllm.config import VllmConfig, get_current_vllm_config, get_layers_from_vllm_config
36 +from vllm.config import (
37 + VllmConfig,
38 + get_current_vllm_config,
39 + get_current_vllm_config_or_none,
40 + get_layers_from_vllm_config,
41 +)
42 from vllm.config.cache import CacheDType
43 from vllm.distributed.parallel_state import get_dcp_group
44 from vllm.logger import init_logger
@@ -609,6 +615,14 @@ class FlashAttentionImpl(AttentionImpl):
615
616 self.supports_quant_query_input = True
617
618 + vllm_config = get_current_vllm_config_or_none()
619 + dcp_a2a = (
620 + vllm_config is not None
621 + and vllm_config.parallel_config.decode_context_parallel_size > 1
622 + and vllm_config.parallel_config.dcp_comm_backend == "a2a"
623 + )
624 + self.dcp_combine = dcp_a2a_lse_reduce if dcp_a2a else cp_lse_ag_out_rs
625 +
626 def forward(
627 self,
628 layer: torch.nn.Module,
@@ -857,8 +871,8 @@ class FlashAttentionImpl(AttentionImpl):
871 v_descale=v_descale,
872 num_splits=attn_metadata.max_num_splits,
873 )
860 - # FA returns LSE in shape [ H, B ] but cp_lse_ag_out_rs wants [ B, H ]
861 - context_attn_out_cor, context_lse_cor = cp_lse_ag_out_rs(
874 + # FA returns LSE in shape [ H, B ] but DCP combine wants [ B, H ]
875 + context_attn_out_cor, context_lse_cor = self.dcp_combine(
876 context_attn_out,
877 context_lse.transpose(0, 1),
878 get_dcp_group(),
vllm/v1/attention/backends/flashinfer.py
+30 -6
@@ -3,6 +3,7 @@
3 """Attention layer with FlashInfer."""
4
5 from dataclasses import dataclass
6 +from functools import partial
7 from typing import ClassVar
8
9 import numpy as np
@@ -19,7 +20,11 @@ from flashinfer.utils import FP4Tensor
20 from typing_extensions import override
21
22 from vllm import envs
22 -from vllm.config import CUDAGraphMode, VllmConfig, get_current_vllm_config
23 +from vllm.config import (
24 + CUDAGraphMode,
25 + VllmConfig,
26 + get_current_vllm_config_or_none,
27 +)
28 from vllm.config.cache import CacheDType
29 from vllm.distributed.parallel_state import get_dcp_group
30 from vllm.logger import init_logger
@@ -59,6 +64,7 @@ from vllm.v1.attention.backends.utils import (
64 split_decodes_and_prefills,
65 )
66 from vllm.v1.attention.ops.common import cp_lse_ag_out_rs
67 +from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
68 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
69 from vllm.v1.kv_cache_interface import AttentionSpec, UniformTypeKVCacheSpecs
70 from vllm.v1.utils import CpuGpuBuffer
@@ -170,7 +176,12 @@ class BatchDCPPrefillWrapper:
176 def __init__(
177 self,
178 workspace_buffer: torch.Tensor | None = None,
179 + dcp_a2a: bool = False,
180 ):
181 + if dcp_a2a:
182 + self._dcp_combine = partial(dcp_a2a_lse_reduce, is_lse_base_on_e=False)
183 + else:
184 + self._dcp_combine = partial(cp_lse_ag_out_rs, is_lse_base_on_e=False)
185 self._context = BatchPrefillWithPagedKVCacheWrapper(
186 workspace_buffer, get_kv_cache_layout()
187 )
@@ -249,12 +260,11 @@ class BatchDCPPrefillWrapper:
260 v_scale=layer._v_scale_float,
261 return_lse=True,
262 )
252 - output_context, lse_context = cp_lse_ag_out_rs(
263 + output_context, lse_context = self._dcp_combine(
264 output_context_tmp,
265 lse_context_tmp,
266 get_dcp_group(),
267 return_lse=True,
257 - is_lse_base_on_e=False,
268 )
269 lse_context = lse_context.transpose(0, 1).contiguous()
270
@@ -550,6 +560,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
560 self.dcp_rank = 0
561 self.dcp_kv_cache_interleave_size = 1
562 self.use_dcp = self.dcp_world_size > 1
563 + self.dcp_a2a = (
564 + self.use_dcp and vllm_config.parallel_config.dcp_comm_backend == "a2a"
565 + )
566
567 self.num_qo_heads = self.model_config.get_num_attention_heads(
568 self.vllm_config.parallel_config
@@ -699,6 +712,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
712 if self.use_dcp:
713 self._prefill_wrapper = BatchDCPPrefillWrapper(
714 workspace_buffer=self._get_workspace_buffer(),
715 + dcp_a2a=self.dcp_a2a,
716 )
717 else:
718 self._prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper(
@@ -1208,15 +1222,26 @@ class FlashInferImpl(AttentionImpl):
1222 self.sinks = sinks
1223
1224 self.support_trtllm_attn = can_use_trtllm_attention(num_heads, num_kv_heads)
1211 - vllm_config = get_current_vllm_config()
1225 + vllm_config = get_current_vllm_config_or_none()
1226 self.supports_quant_query_input = (
1227 self.support_trtllm_attn
1228 + and vllm_config is not None
1229 and not vllm_config.attention_config.disable_flashinfer_q_quantization
1230 )
1231 self.bmm1_scale: float | None = None
1232 self.bmm2_scale: float | None = None
1233 self.o_sf_scale: float | None = None
1234
1235 + dcp_a2a = (
1236 + vllm_config is not None
1237 + and vllm_config.parallel_config.decode_context_parallel_size > 1
1238 + and vllm_config.parallel_config.dcp_comm_backend == "a2a"
1239 + )
1240 + if dcp_a2a:
1241 + self.dcp_combine = partial(dcp_a2a_lse_reduce, is_lse_base_on_e=False)
1242 + else:
1243 + self.dcp_combine = partial(cp_lse_ag_out_rs, is_lse_base_on_e=False)
1244 +
1245 def fused_output_quant_supported(self, quant_key: QuantKey):
1246 return (
1247 self.support_trtllm_attn
@@ -1503,11 +1528,10 @@ class FlashInferImpl(AttentionImpl):
1528 lse=lse,
1529 return_lse=True,
1530 )
1506 - output[:num_decode_tokens] = cp_lse_ag_out_rs(
1531 + output[:num_decode_tokens] = self.dcp_combine(
1532 output_tmp,
1533 lse,
1534 get_dcp_group(),
1510 - is_lse_base_on_e=False,
1535 )
1536 else:
1537 decode_wrapper.run(
vllm/v1/attention/ops/dcp_alltoall.py new
+363
@@ -0,0 +1,363 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +DCP All-to-All communication backend for attention.
5 +
6 +Provides All-to-All (A2A) communication as an alternative to
7 +AllGather + ReduceScatter (AG+RS) for Decode Context Parallel (DCP).
8 +Instead of gathering the full Q tensor and scattering partial outputs,
9 +A2A exchanges partial attention outputs and their LSE values across
10 +ranks, then combines them with exact LSE-weighted reduction.
11 +
12 +This reduces the number of NCCL calls per attention layer from 3
13 +(AG for Q, AG for K metadata, RS for output) to 2 (A2A for output,
14 +A2A for LSE), lowering per-step communication overhead for long-context
15 +decode where NCCL latency is a significant fraction of step time.
16 +
17 +Usage:
18 + vllm serve model --tp 16 --dcp 16 --dcp-comm-backend a2a
19 +
20 +Reference: https://arxiv.org/abs/2507.07120
21 +"""
22 +
23 +from __future__ import annotations
24 +
25 +from typing import TYPE_CHECKING
26 +
27 +import torch
28 +import torch.distributed as dist
29 +
30 +from vllm.triton_utils import tl, triton
31 +
32 +if TYPE_CHECKING:
33 + from vllm.distributed.parallel_state import GroupCoordinator
34 + from vllm.v1.attention.ops.common import CPTritonContext
35 +
36 +
37 +def _lse_weighted_combine(
38 + outputs: torch.Tensor,
39 + lses: torch.Tensor,
40 + return_lse: bool = False,
41 + is_lse_base_on_e: bool = True,
42 +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
43 + """
44 + CPU reference implementation for LSE-weighted combination.
45 +
46 + This is a pure PyTorch implementation used for testing and validation.
47 + For GPU execution, use dcp_lse_combine_triton instead.
48 +
49 + Args:
50 + outputs: Partial attention outputs [N, B, H, D]
51 + N = number of KV shards (ranks)
52 + B = batch size (num_tokens)
53 + H = number of heads per rank
54 + D = head dimension
55 + lses: Log-sum-exp values [N, B, H]
56 + return_lse: If True, also return the global LSE
57 + is_lse_base_on_e: If True, LSE is base e; if False, base 2
58 +
59 + Returns:
60 + Combined output [B, H, D], and optionally global LSE [B, H]
61 + """
62 + N, B, H, D = outputs.shape
63 +
64 + # Handle NaN and inf in LSEs
65 + lses = torch.where(
66 + torch.isnan(lses) | torch.isinf(lses),
67 + torch.tensor(float("-inf"), device=lses.device, dtype=lses.dtype),
68 + lses,
69 + )
70 +
71 + # Compute max LSE for numerical stability
72 + lse_max, _ = lses.max(dim=0) # [B, H]
73 + lse_max = torch.where(
74 + lse_max == float("-inf"),
75 + torch.zeros_like(lse_max),
76 + lse_max,
77 + )
78 +
79 + # Compute weights: softmax over the N dimension
80 + if is_lse_base_on_e:
81 + weights = torch.exp(lses - lse_max.unsqueeze(0)) # [N, B, H]
82 + else:
83 + weights = torch.pow(2.0, lses - lse_max.unsqueeze(0)) # [N, B, H]
84 +
85 + # Handle NaN weights
86 + weights = torch.where(torch.isnan(weights), torch.zeros_like(weights), weights)
87 +
88 + # Normalize weights
89 + weight_sum = weights.sum(dim=0, keepdim=True) # [1, B, H]
90 + weights = weights / weight_sum.clamp(min=1e-10) # [N, B, H]
91 +
92 + # Weighted combination: sum over N dimension
93 + result = (outputs * weights.unsqueeze(-1)).sum(dim=0) # [B, H, D]
94 +
95 + if return_lse:
96 + if is_lse_base_on_e:
97 + global_lse = torch.log(weight_sum.squeeze(0)) + lse_max # [B, H]
98 + else:
99 + global_lse = torch.log2(weight_sum.squeeze(0)) + lse_max # [B, H]
100 + return result, global_lse
101 +
102 + return result
103 +
104 +
105 +@triton.jit
106 +def _dcp_lse_combine_kernel(
107 + # Input pointers
108 + recv_output_ptr,
109 + recv_lse_ptr,
110 + # Output pointers
111 + out_ptr,
112 + out_lse_ptr,
113 + # Strides for recv_output [N, B, H_local, D]
114 + ro_stride_N,
115 + ro_stride_B,
116 + ro_stride_H,
117 + ro_stride_D,
118 + # Strides for recv_lse [N, B, H_local]
119 + rl_stride_N,
120 + rl_stride_B,
121 + rl_stride_H,
122 + # Strides for output [B, H_local, D]
123 + o_stride_B,
124 + o_stride_H,
125 + o_stride_D,
126 + # Constants
127 + N: tl.constexpr,
128 + HEAD_DIM: tl.constexpr,
129 + IS_BASE_E: tl.constexpr,
130 + RETURN_LSE: tl.constexpr,
131 +):
132 + """
133 + Triton kernel for LSE-weighted combination of partial attention outputs.
134 +
135 + After All-to-All, each rank has:
136 + - recv_output [N, B, H_local, D]: partial outputs from all KV shards
137 + - recv_lse [N, B, H_local]: partial LSEs from all KV shards
138 +
139 + This kernel computes the weighted combination locally (no communication).
140 +
141 + Grid: (B, H_local)
142 + Each program handles one (batch, head) and processes all D elements.
143 + """
144 + batch_idx = tl.program_id(0).to(tl.int64)
145 + head_idx = tl.program_id(1).to(tl.int64)
146 +
147 + # Base offset for this (batch, head)
148 + base_lse_offset = batch_idx * rl_stride_B + head_idx * rl_stride_H
149 + base_out_offset = batch_idx * ro_stride_B + head_idx * ro_stride_H
150 +
151 + # First pass: find max LSE for numerical stability
152 + lse_max = -float("inf")
153 + for n in tl.static_range(N):
154 + lse_offset = n * rl_stride_N + base_lse_offset
155 + lse_val = tl.load(recv_lse_ptr + lse_offset)
156 + lse_val = tl.where(
157 + (lse_val != lse_val) | (lse_val == float("inf")),
158 + -float("inf"),
159 + lse_val,
160 + )
161 + lse_max = tl.maximum(lse_max, lse_val)
162 +
163 + lse_max = tl.where(lse_max == -float("inf"), 0.0, lse_max)
164 +
165 + # Second pass: compute sum of exp(lse - max)
166 + lse_sum = 0.0
167 + for n in tl.static_range(N):
168 + lse_offset = n * rl_stride_N + base_lse_offset
169 + lse_val = tl.load(recv_lse_ptr + lse_offset)
170 + lse_val = tl.where(
171 + (lse_val != lse_val) | (lse_val == float("inf")),
172 + -float("inf"),
173 + lse_val,
174 + )
175 + if IS_BASE_E:
176 + lse_sum += tl.exp(lse_val - lse_max)
177 + else:
178 + lse_sum += tl.exp2(lse_val - lse_max)
179 +
180 + # Compute global LSE
181 + if IS_BASE_E: # noqa: SIM108
182 + global_lse = tl.log(lse_sum) + lse_max
183 + else:
184 + global_lse = tl.log2(lse_sum) + lse_max
185 +
186 + # Third pass: weighted combination across D dimension
187 + d_offsets = tl.arange(0, HEAD_DIM)
188 + acc = tl.zeros([HEAD_DIM], dtype=tl.float32)
189 +
190 + for n in tl.static_range(N):
191 + lse_offset = n * rl_stride_N + base_lse_offset
192 + lse_val = tl.load(recv_lse_ptr + lse_offset)
193 + lse_val = tl.where(
194 + (lse_val != lse_val) | (lse_val == float("inf")),
195 + -float("inf"),
196 + lse_val,
197 + )
198 + if IS_BASE_E:
199 + weight = tl.exp(lse_val - global_lse)
200 + else:
201 + weight = tl.exp2(lse_val - global_lse)
202 + weight = tl.where(weight != weight, 0.0, weight)
203 +
204 + out_offsets = n * ro_stride_N + base_out_offset + d_offsets * ro_stride_D
205 + out_vals = tl.load(recv_output_ptr + out_offsets)
206 + acc += out_vals.to(tl.float32) * weight
207 +
208 + # Store result
209 + final_offsets = (
210 + batch_idx * o_stride_B + head_idx * o_stride_H + d_offsets * o_stride_D
211 + )
212 + tl.store(out_ptr + final_offsets, acc)
213 +
214 + if RETURN_LSE:
215 + tl.store(out_lse_ptr + base_lse_offset, global_lse)
216 +
217 +
218 +def dcp_lse_combine_triton(
219 + recv_output: torch.Tensor,
220 + recv_lse: torch.Tensor,
221 + return_lse: bool = False,
222 + is_lse_base_on_e: bool = True,
223 +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
224 + """
225 + Triton-accelerated LSE-weighted combination for DCP A2A.
226 +
227 + Args:
228 + recv_output: [N, B, H_local, D] - partial outputs from all KV shards
229 + recv_lse: [N, B, H_local] - partial LSEs from all KV shards
230 + return_lse: If True, also return the global LSE
231 + is_lse_base_on_e: If True, LSE is base e; if False, base 2
232 +
233 + Returns:
234 + Combined output [B, H_local, D]
235 + If return_lse=True, also returns global_lse [B, H_local]
236 + """
237 + N, B, H_local, D = recv_output.shape
238 +
239 + out = torch.empty(
240 + (B, H_local, D), device=recv_output.device, dtype=recv_output.dtype
241 + )
242 +
243 + if return_lse:
244 + out_lse = torch.empty(
245 + (B, H_local), device=recv_lse.device, dtype=recv_lse.dtype
246 + )
247 + else:
248 + out_lse = torch.empty(1, device=recv_lse.device, dtype=recv_lse.dtype)
249 +
250 + ro_stride_N, ro_stride_B, ro_stride_H, ro_stride_D = recv_output.stride()
251 + rl_stride_N, rl_stride_B, rl_stride_H = recv_lse.stride()
252 + o_stride_B, o_stride_H, o_stride_D = out.stride()
253 +
254 + grid = (B, H_local, 1)
255 +
256 + _dcp_lse_combine_kernel[grid](
257 + recv_output,
258 + recv_lse,
259 + out,
260 + out_lse,
261 + ro_stride_N,
262 + ro_stride_B,
263 + ro_stride_H,
264 + ro_stride_D,
265 + rl_stride_N,
266 + rl_stride_B,
267 + rl_stride_H,
268 + o_stride_B,
269 + o_stride_H,
270 + o_stride_D,
271 + N=N,
272 + HEAD_DIM=D,
273 + IS_BASE_E=is_lse_base_on_e,
274 + RETURN_LSE=return_lse,
275 + )
276 +
277 + if return_lse:
278 + return out, out_lse
279 + return out
280 +
281 +
282 +def dcp_a2a_lse_reduce(
283 + cp_attn_out: torch.Tensor,
284 + cp_attn_lse: torch.Tensor,
285 + cp_group: GroupCoordinator,
286 + ctx: CPTritonContext | None = None,
287 + return_lse: bool = False,
288 + is_lse_base_on_e: bool = True,
289 +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
290 + """
291 + Combine partial attention outputs across DCP ranks using All-to-All.
292 +
293 + Each rank holds attention output for all heads but only a local shard
294 + of the KV cache. This function:
295 + 1. Exchanges partial outputs across ranks via All-to-All
296 + 2. Exchanges LSE values via All-to-All
297 + 3. Combines them with exact LSE-weighted reduction (Triton kernel)
298 +
299 + Tensor flow:
300 + Input: cp_attn_out [B, H, D] - all heads, local KV shard
301 + Reshape: [N, B, H/N, D] - split heads across ranks
302 + A2A: Two all_to_all_single calls (output and LSE)
303 + Combine: recv [N, B, H/N, D] + lse [N, B, H/N] -> [B, H/N, D]
304 +
305 + Args:
306 + cp_attn_out: [B, H, D] where B=num_tokens, H=total_heads, D=head_dim
307 + cp_attn_lse: [B, H] log-sum-exp values (fp32)
308 + cp_group: GroupCoordinator for DCP communication
309 + ctx: CPTritonContext (unused, for signature compatibility)
310 + return_lse: If True, also return the combined global LSE
311 + is_lse_base_on_e: If True, LSE is base e; if False, base 2
312 +
313 + Returns:
314 + Combined output [B, H/N, D] (head-scattered)
315 + If return_lse=True, also returns global_lse [B, H/N]
316 + """
317 + world_size = cp_group.world_size
318 +
319 + if world_size == 1:
320 + if return_lse:
321 + return cp_attn_out, cp_attn_lse
322 + return cp_attn_out
323 +
324 + local_output = cp_attn_out.contiguous()
325 + local_lse = cp_attn_lse.contiguous()
326 +
327 + B, H, D = local_output.shape
328 + H_per_rank = H // world_size
329 +
330 + # Reshape for All-to-All: [B, H, D] -> [N, B, H/N, D]
331 + # Split heads into N chunks, each destined for a different rank
332 + send_output = (
333 + local_output.view(B, world_size, H_per_rank, D).permute(1, 0, 2, 3).contiguous()
334 + )
335 + recv_output = torch.empty_like(send_output)
336 +
337 + # Same for LSE: [B, H] -> [N, B, H/N]
338 + send_lse = local_lse.view(B, world_size, H_per_rank).permute(1, 0, 2).contiguous()
339 + recv_lse = torch.empty_like(send_lse)
340 +
341 + # All-to-All for partial attention outputs and LSE values (async overlap)
342 + work_output = dist.all_to_all_single(
343 + recv_output.view(-1),
344 + send_output.view(-1),
345 + group=cp_group.device_group,
346 + async_op=True,
347 + )
348 + work_lse = dist.all_to_all_single(
349 + recv_lse.view(-1),
350 + send_lse.view(-1),
351 + group=cp_group.device_group,
352 + async_op=True,
353 + )
354 + work_output.wait()
355 + work_lse.wait()
356 +
357 + # LSE-weighted combination via Triton kernel (local, no communication)
358 + return dcp_lse_combine_triton(
359 + recv_output,
360 + recv_lse,
361 + return_lse=return_lse,
362 + is_lse_base_on_e=is_lse_base_on_e,
363 + )