[refactor] consolidate cp attn ops (#52839)
Signed-off-by: Summer Yang <girasoleyang@gmail.com> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai>
Summer Yang committed
Aug 19, 2026 at 18:04 UTC
58e5ee0158b6a264c3506f00480e108a34b33ee3
18 files changed
+1603
-1585
tests/distributed/test_dcp_a2a.py
+17
-17
@@ -44,7 +44,7 @@ def _packed_a2a_reference(
44
h_per_rank: int,
45
is_lse_base_on_e: bool,
46
) -> tuple[torch.Tensor, torch.Tensor]:
47
- from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
47
+ from vllm.v1.attention.ops.dcp import _lse_weighted_combine
48
49
B, _H, D = cp_attn_out.shape
50
outputs = (
@@ -157,13 +157,13 @@ class TestLSEWeightedCombine:
157
158
def test_importable(self):
159
"""Verify _lse_weighted_combine is importable."""
160
- from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
160
+ from vllm.v1.attention.ops.dcp import _lse_weighted_combine
161
162
assert callable(_lse_weighted_combine)
163
164
def test_single_rank(self):
165
"""Single rank: output unchanged."""
166
- from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
166
+ from vllm.v1.attention.ops.dcp import _lse_weighted_combine
167
168
# N=1, B=2, H=4, D=8
169
outputs = torch.randn(1, 2, 4, 8)
@@ -176,7 +176,7 @@ class TestLSEWeightedCombine:
176
177
def test_equal_lse(self):
178
"""Equal LSE values: outputs averaged equally."""
179
- from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
179
+ from vllm.v1.attention.ops.dcp import _lse_weighted_combine
180
181
_N, B, H, D = 2, 1, 1, 4
182
outputs = torch.tensor(
@@ -200,7 +200,7 @@ class TestLSEWeightedCombine:
200
201
def test_dominant_rank(self):
202
"""Different LSE values: larger LSE gets more weight."""
203
- from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
203
+ from vllm.v1.attention.ops.dcp import _lse_weighted_combine
204
205
B, H, D = 1, 1, 2
206
outputs = torch.tensor(
@@ -222,7 +222,7 @@ class TestLSEWeightedCombine:
222
torch.testing.assert_close(result, outputs[1], atol=1e-5, rtol=1e-5)
223
224
def test_empty_shard_ignores_undefined_output(self):
225
- from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
225
+ from vllm.v1.attention.ops.dcp import _lse_weighted_combine
226
227
outputs = torch.tensor([[[[float("nan")]]], [[[3.0]]]])
228
lses = torch.tensor([[[-float("inf")]], [[0.0]]])
@@ -232,7 +232,7 @@ class TestLSEWeightedCombine:
232
torch.testing.assert_close(result, outputs[1])
233
234
def test_ag_rs_masks_empty_shard_and_padded_lse(self, monkeypatch):
235
- import vllm.v1.attention.ops.common as common
235
+ import vllm.v1.attention.ops.dcp as dcp
236
237
class FakeGroup:
238
world_size = 2
@@ -243,7 +243,7 @@ class TestLSEWeightedCombine:
243
return torch.cat((tensor, tensor), dim=dim)
244
245
monkeypatch.setattr(
246
- common,
246
+ dcp,
247
"correct_attn_out",
248
lambda output, lses, *args, **kwargs: (output, lses[0]),
249
)
@@ -252,7 +252,7 @@ class TestLSEWeightedCombine:
252
seq_lens = torch.tensor([0, 2], dtype=torch.int32)
253
query_start_loc = torch.tensor([0, 1, 5], dtype=torch.int32)
254
255
- _, masked_lse = common._cp_lse_common(
255
+ _, masked_lse = dcp._cp_lse_common(
256
output,
257
lse,
258
FakeGroup(),
@@ -266,7 +266,7 @@ class TestLSEWeightedCombine:
266
267
def test_mathematically_correct(self):
268
"""Verify mathematical correctness of LSE combination."""
269
- from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
269
+ from vllm.v1.attention.ops.dcp import _lse_weighted_combine
270
271
outputs = torch.tensor(
272
[
@@ -291,7 +291,7 @@ class TestLSEWeightedCombine:
291
292
def test_return_lse(self):
293
"""return_lse=True returns global LSE (logsumexp of inputs)."""
294
- from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
294
+ from vllm.v1.attention.ops.dcp import _lse_weighted_combine
295
296
B, H, D = 1, 1, 2
297
outputs = torch.tensor(
@@ -317,7 +317,7 @@ class TestLSEWeightedCombine:
317
318
def test_base2_return_lse(self):
319
"""Base-2 LSE mode returns log2-sum-exp2 global LSE."""
320
- from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
320
+ from vllm.v1.attention.ops.dcp import _lse_weighted_combine
321
322
outputs = torch.tensor(
323
[
@@ -354,7 +354,7 @@ class TestLSEWeightedCombine:
354
355
def test_lse_pack_dim(self):
356
"""Packed A2A stores one fp32 LSE in output-dtype lanes."""
357
- from vllm.v1.attention.ops.dcp_alltoall import _dcp_a2a_lse_pack_dim
357
+ from vllm.v1.attention.ops.dcp import _dcp_a2a_lse_pack_dim
358
359
assert _dcp_a2a_lse_pack_dim(torch.bfloat16) == 2
360
assert _dcp_a2a_lse_pack_dim(torch.float16) == 2
@@ -374,7 +374,7 @@ class TestPackedA2AKernels:
374
return_lse: bool,
375
is_lse_base_on_e: bool,
376
):
377
- from vllm.v1.attention.ops.dcp_alltoall import (
377
+ from vllm.v1.attention.ops.dcp import (
378
_dcp_a2a_lse_pack_dim,
379
_dcp_a2a_pack_send,
380
_dcp_a2a_unpack_combine,
@@ -421,7 +421,7 @@ class TestPackedA2AKernels:
421
torch.accelerator.device_count() < 1, reason="CUDA is required."
422
)
423
def test_empty_seq_lens_ignore_undefined_output(self):
424
- from vllm.v1.attention.ops.dcp_alltoall import (
424
+ from vllm.v1.attention.ops.dcp import (
425
_dcp_a2a_lse_pack_dim,
426
_dcp_a2a_pack_send,
427
_dcp_a2a_unpack_combine,
@@ -505,7 +505,7 @@ def _distributed_packed_a2a_worker(env: dict[str, str]) -> None:
505
506
init_workspace_manager(torch.device(f"cuda:{local_rank}"))
507
try:
508
- from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
508
+ from vllm.v1.attention.ops.dcp import dcp_a2a_lse_reduce
509
510
dtype = _dtype_from_name(env["TEST_DTYPE"])
511
return_lse = env["RETURN_LSE"] == "1"
@@ -555,7 +555,7 @@ def _distributed_packed_a2a_worker(env: dict[str, str]) -> None:
555
[t[:, rank * h_per_rank : (rank + 1) * h_per_rank] for t in gathered_lse],
556
dim=0,
557
)
558
- from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
558
+ from vllm.v1.attention.ops.dcp import _lse_weighted_combine
559
560
expected_out, expected_lse = _lse_weighted_combine(
561
outputs,
tests/distributed/test_dcp_direct_a2a_lse_reduce.py
+41
-40
@@ -10,7 +10,8 @@ import pytest
10
import torch
11
import torch.distributed as dist
12
13
-import vllm.v1.attention.ops.dcp_utils as dcp_utils
13
+import vllm.v1.attention.ops.cp_common as cp_common
14
+import vllm.v1.attention.ops.dcp as dcp
15
from vllm.utils.network_utils import get_open_port
16
from vllm.utils.system_utils import update_environment_variables
17
@@ -142,17 +143,17 @@ class _FakeProcessGroup:
143
class TestDirectDCPGating:
144
def test_env_disabled_returns_none(self, monkeypatch):
145
monkeypatch.setenv("VLLM_USE_DIRECT_DCP_A2A", "0")
145
- dcp_utils.get_direct_dcp_a2a_workspace.cache_clear()
146
- workspace = dcp_utils.get_direct_dcp_a2a_workspace(
146
+ dcp.get_direct_dcp_a2a_workspace.cache_clear()
147
+ workspace = dcp.get_direct_dcp_a2a_workspace(
148
_FakeGroupCoordinator(), torch.device("cpu"), 16, 2, 32, torch.bfloat16, 1
149
)
150
assert workspace is None
151
152
def test_forced_with_unsupported_dtype_raises(self, monkeypatch):
153
monkeypatch.setenv("VLLM_USE_DIRECT_DCP_A2A", "1")
153
- dcp_utils.get_direct_dcp_a2a_workspace.cache_clear()
154
+ dcp.get_direct_dcp_a2a_workspace.cache_clear()
155
with pytest.raises(ValueError, match="does not support"):
155
- dcp_utils.get_direct_dcp_a2a_workspace(
156
+ dcp.get_direct_dcp_a2a_workspace(
157
_FakeGroupCoordinator(),
158
torch.device("cpu"),
159
16,
@@ -164,14 +165,14 @@ class TestDirectDCPGating:
165
166
def test_zero_ubatches_raises(self):
167
with pytest.raises(ValueError, match="ubatch"):
167
- dcp_utils.DirectDCPA2AWorkspace(
168
+ dcp.DirectDCPA2AWorkspace(
169
None, torch.device("cpu"), 16, 2, 32, torch.bfloat16, num_ubatches=0
170
)
171
172
def test_auto_with_unsupported_dtype_returns_none(self, monkeypatch):
173
monkeypatch.delenv("VLLM_USE_DIRECT_DCP_A2A", raising=False)
173
- dcp_utils.get_direct_dcp_a2a_workspace.cache_clear()
174
- workspace = dcp_utils.get_direct_dcp_a2a_workspace(
174
+ dcp.get_direct_dcp_a2a_workspace.cache_clear()
175
+ workspace = dcp.get_direct_dcp_a2a_workspace(
176
_FakeGroupCoordinator(), torch.device("cpu"), 16, 2, 32, torch.float32, 1
177
)
178
assert workspace is None
@@ -179,8 +180,8 @@ class TestDirectDCPGating:
180
def test_q_gather_env_disabled_returns_none(self, monkeypatch):
181
monkeypatch.setenv("VLLM_USE_DIRECT_DCP_Q_GATHER", "0")
182
monkeypatch.setenv("VLLM_USE_DIRECT_DCP_A2A", "1")
182
- dcp_utils.get_direct_dcp_q_gather_workspace.cache_clear()
183
- workspace = dcp_utils.get_direct_dcp_q_gather_workspace(
183
+ dcp.get_direct_dcp_q_gather_workspace.cache_clear()
184
+ workspace = dcp.get_direct_dcp_q_gather_workspace(
185
_FakeGroupCoordinator(),
186
torch.device("cpu"),
187
16,
@@ -194,17 +195,17 @@ class TestDirectDCPGating:
195
def test_q_gather_flag_is_independent(self, monkeypatch):
196
monkeypatch.setenv("VLLM_USE_DIRECT_DCP_Q_GATHER", "1")
197
monkeypatch.setenv("VLLM_USE_DIRECT_DCP_A2A", "0")
197
- monkeypatch.setattr(dcp_utils, "_symm_mem_spans_group", lambda group: True)
198
- dcp_utils.get_direct_dcp_q_gather_workspace.cache_clear()
198
+ monkeypatch.setattr(cp_common, "_symm_mem_spans_group", lambda group: True)
199
+ dcp.get_direct_dcp_q_gather_workspace.cache_clear()
200
workspace = object()
201
init_workspace = MagicMock(return_value=workspace)
202
monkeypatch.setattr(
202
- dcp_utils,
203
+ dcp,
204
"DirectDCPQGatherWorkspace",
205
init_workspace,
206
)
207
207
- result = dcp_utils.get_direct_dcp_q_gather_workspace(
208
+ result = dcp.get_direct_dcp_q_gather_workspace(
209
_FakeGroupCoordinator(),
210
torch.device("cpu"),
211
16,
@@ -220,8 +221,8 @@ class TestDirectDCPGating:
221
def test_kv_gather_env_disabled_returns_none(self, monkeypatch):
222
monkeypatch.setenv("VLLM_USE_DIRECT_DCP_KV_GATHER", "0")
223
monkeypatch.setenv("VLLM_USE_DIRECT_DCP_A2A", "1")
223
- dcp_utils.get_direct_dcp_kv_gather_workspace.cache_clear()
224
- workspace = dcp_utils.get_direct_dcp_kv_gather_workspace(
224
+ dcp.get_direct_dcp_kv_gather_workspace.cache_clear()
225
+ workspace = dcp.get_direct_dcp_kv_gather_workspace(
226
_FakeGroupCoordinator(), torch.device("cpu"), 64, 576, torch.bfloat16, 1
227
)
228
assert workspace is None
@@ -229,17 +230,17 @@ class TestDirectDCPGating:
230
def test_kv_gather_flag_is_independent(self, monkeypatch):
231
monkeypatch.setenv("VLLM_USE_DIRECT_DCP_KV_GATHER", "1")
232
monkeypatch.setenv("VLLM_USE_DIRECT_DCP_A2A", "0")
232
- monkeypatch.setattr(dcp_utils, "_symm_mem_spans_group", lambda group: True)
233
- dcp_utils.get_direct_dcp_kv_gather_workspace.cache_clear()
233
+ monkeypatch.setattr(cp_common, "_symm_mem_spans_group", lambda group: True)
234
+ dcp.get_direct_dcp_kv_gather_workspace.cache_clear()
235
workspace = object()
236
init_workspace = MagicMock(return_value=workspace)
237
monkeypatch.setattr(
237
- dcp_utils,
238
+ dcp,
239
"DirectDCPKVGatherWorkspace",
240
init_workspace,
241
)
242
242
- result = dcp_utils.get_direct_dcp_kv_gather_workspace(
243
+ result = dcp.get_direct_dcp_kv_gather_workspace(
244
_FakeGroupCoordinator(), torch.device("cpu"), 64, 576, torch.bfloat16, 1
245
)
246
@@ -267,9 +268,9 @@ class TestDirectDCPGating:
268
factory_name,
269
factory_args,
270
):
270
- factory = getattr(dcp_utils, factory_name)
271
+ factory = getattr(dcp, factory_name)
272
monkeypatch.setenv(flag_name, "1")
272
- monkeypatch.setattr(dcp_utils, "_symm_mem_spans_group", lambda group: False)
273
+ monkeypatch.setattr(cp_common, "_symm_mem_spans_group", lambda group: False)
274
factory.cache_clear()
275
276
assert (
@@ -283,25 +284,25 @@ class TestDirectDCPGating:
284
285
def test_kv_gather_rejects_invalid_workspace_geometry(self):
286
with pytest.raises(ValueError, match="ubatch"):
286
- dcp_utils.DirectDCPKVGatherWorkspace(
287
+ dcp.DirectDCPKVGatherWorkspace(
288
None, torch.device("cpu"), 64, 576, num_ubatches=0
289
)
290
with pytest.raises(ValueError, match="divide evenly"):
290
- dcp_utils.DirectDCPKVGatherWorkspace(
291
+ dcp.DirectDCPKVGatherWorkspace(
292
_FakeProcessGroup(), torch.device("cpu"), 63, 576
293
)
294
with pytest.raises(ValueError, match="16-byte"):
294
- dcp_utils.DirectDCPKVGatherWorkspace(
295
+ dcp.DirectDCPKVGatherWorkspace(
296
_FakeProcessGroup(), torch.device("cpu"), 64, 3
297
)
298
299
def test_q_gather_rejects_invalid_workspace_geometry(self):
300
with pytest.raises(ValueError, match="ubatch"):
300
- dcp_utils.DirectDCPQGatherWorkspace(
301
+ dcp.DirectDCPQGatherWorkspace(
302
None, torch.device("cpu"), 16, 2, 32, num_ubatches=0
303
)
304
with pytest.raises(ValueError, match="padded heads"):
304
- dcp_utils.DirectDCPQGatherWorkspace(
305
+ dcp.DirectDCPQGatherWorkspace(
306
_FakeProcessGroup(),
307
torch.device("cpu"),
308
16,
@@ -310,7 +311,7 @@ class TestDirectDCPGating:
311
padded_num_heads=7,
312
)
313
with pytest.raises(ValueError, match="16-byte"):
313
- dcp_utils.DirectDCPQGatherWorkspace(
314
+ dcp.DirectDCPQGatherWorkspace(
315
_FakeProcessGroup(),
316
torch.device("cpu"),
317
16,
@@ -332,7 +333,7 @@ def _manager_config(dcp_comm_backend: str = "a2a"):
333
334
335
def test_mla_dcp_manager_selects_direct_backends(monkeypatch):
335
- import vllm.v1.attention.ops.dcp_utils as dcp_manager
336
+ import vllm.v1.attention.ops.dcp as dcp_manager
337
338
group = MagicMock(world_size=2)
339
monkeypatch.setattr(dcp_manager, "get_dcp_group", lambda: group)
@@ -391,7 +392,7 @@ def test_mla_dcp_manager_selects_direct_backends(monkeypatch):
392
393
394
def test_mla_dcp_manager_selects_fallback_backends(monkeypatch):
394
- import vllm.v1.attention.ops.dcp_utils as dcp_manager
395
+ import vllm.v1.attention.ops.dcp as dcp_manager
396
397
group = MagicMock(world_size=2)
398
gathered_query = torch.empty(1, 4, 8)
@@ -460,11 +461,11 @@ def test_dcp_workspace_covers_parallel_drafting():
461
config.num_speculative_tokens = 3
462
config.speculative_config = MagicMock(parallel_drafting=True)
463
463
- assert dcp_utils.get_dcp_workspace_max_num_tokens(config) == 28
464
+ assert dcp.get_dcp_workspace_max_num_tokens(config) == 28
465
466
467
def test_mla_dcp_manager_selects_pcp_combine(monkeypatch):
467
- import vllm.v1.attention.ops.dcp_utils as dcp_manager
468
+ import vllm.v1.attention.ops.dcp as dcp_manager
469
470
monkeypatch.setattr(dcp_manager, "get_dcp_group", lambda: MagicMock(world_size=2))
471
manager = dcp_manager.MLADCPManager(
@@ -521,7 +522,7 @@ def test_sparse_mla_builder_initializes_dcp_manager(monkeypatch):
522
lambda _: MagicMock(kv_lora_rank=8, qk_rope_head_dim=4),
523
)
524
524
- manager = object.__new__(dcp_utils.MLADCPManager)
525
+ manager = object.__new__(dcp.MLADCPManager)
526
manager.init_kv_gather = MagicMock()
527
layer = MagicMock(dcp_manager=manager)
528
config = MagicMock()
@@ -580,12 +581,12 @@ def _distributed_direct_q_gather_worker(env: dict[str, str]) -> None:
581
heads_per_rank, head_dim, max_num_tokens = 6, 576, 128
582
padded_num_heads = 128 if world_size == 4 else None
583
active_ubatch = [0]
583
- dcp_utils.dbo_current_ubatch_id = lambda: active_ubatch[0]
584
+ dcp.dbo_current_ubatch_id = lambda: active_ubatch[0]
585
for dtype_idx, dtype_name in enumerate(
586
("bfloat16", "float8_e4m3fn", "float32")
587
):
588
dtype = _dtype_from_name(dtype_name)
588
- workspace = dcp_utils.DirectDCPQGatherWorkspace(
589
+ workspace = dcp.DirectDCPQGatherWorkspace(
590
dist.group.WORLD,
591
device,
592
max_num_tokens,
@@ -747,11 +748,11 @@ def _distributed_direct_kv_gather_worker(env: dict[str, str]) -> None:
748
token_dim = 576
749
max_gathered_tokens = 128 * world_size
750
active_ubatch = [0]
750
- dcp_utils.dbo_current_ubatch_id = lambda: active_ubatch[0]
751
+ dcp.dbo_current_ubatch_id = lambda: active_ubatch[0]
752
753
for dtype_idx, dtype_name in enumerate(("bfloat16", "float16")):
754
dtype = _dtype_from_name(dtype_name)
754
- workspace = dcp_utils.DirectDCPKVGatherWorkspace(
755
+ workspace = dcp.DirectDCPKVGatherWorkspace(
756
dist.group.WORLD,
757
device,
758
max_gathered_tokens,
@@ -818,7 +819,7 @@ def _distributed_direct_a2a_worker(env: dict[str, str]) -> None:
819
torch.accelerator.set_device_index(local_rank)
820
dist.init_process_group(backend="nccl")
821
try:
821
- from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
822
+ from vllm.v1.attention.ops.dcp import _lse_weighted_combine
823
824
rank = dist.get_rank()
825
world_size = dist.get_world_size()
@@ -829,8 +830,8 @@ def _distributed_direct_a2a_worker(env: dict[str, str]) -> None:
830
heads_per_rank, head_dim, max_num_tokens = 6, 512, 128
831
total_heads = world_size * heads_per_rank
832
active_ubatch = [0]
832
- dcp_utils.dbo_current_ubatch_id = lambda: active_ubatch[0]
833
- workspace = dcp_utils.DirectDCPA2AWorkspace(
833
+ dcp.dbo_current_ubatch_id = lambda: active_ubatch[0]
834
+ workspace = dcp.DirectDCPA2AWorkspace(
835
dist.group.WORLD,
836
device,
837
max_num_tokens,
tests/v1/attention/test_flashinfer_mla_dcp.py
+1
-1
@@ -15,7 +15,7 @@ requires_flashinfer_mla = pytest.mark.skipif(
15
16
17
def test_mla_dcp_gathered_query_reserves_backend_head_storage():
18
- from vllm.v1.attention.ops.dcp_utils import reserve_query_head_storage
18
+ from vllm.v1.attention.ops.dcp import reserve_query_head_storage
19
20
query = torch.randn(3, 24, 576, dtype=torch.bfloat16)
21
tests/v1/attention/test_indexer_dcp_localize.py
+1
-1
@@ -12,7 +12,7 @@ from vllm.v1.attention.backends.mla.sparse_utils import (
12
triton_filter_and_convert_dcp_index,
13
)
14
from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
15
-from vllm.v1.attention.ops.common import CPTritonContext, correct_attn_out
15
+from vllm.v1.attention.ops.dcp import CPTritonContext, correct_attn_out
16
17
18
def _local_count(length: int, rank: int, world: int, interleave: int) -> int:
vllm/model_executor/layers/attention/mla_attention.py
+5
-5
@@ -242,10 +242,6 @@ from vllm.model_executor.layers.attention.attention import (
242
from vllm.model_executor.layers.attention.kv_transfer_utils import (
243
maybe_transfer_kv_layer,
244
)
245
-from vllm.model_executor.layers.attention.pcp import (
246
- finalize_mla_pcp_decode,
247
- maybe_gather_mla_latent_cache_inputs,
248
-)
245
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
246
from vllm.model_executor.layers.linear import (
247
ColumnParallelLinear,
@@ -293,8 +289,12 @@ from vllm.v1.attention.backends.utils import (
289
get_num_attention_heads_from_layers,
290
split_decodes_and_prefills,
291
)
296
-from vllm.v1.attention.ops.dcp_utils import MLADCPManager
292
+from vllm.v1.attention.ops.dcp import MLADCPManager
293
from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
294
+from vllm.v1.attention.ops.pcp import (
295
+ finalize_mla_pcp_decode,
296
+ maybe_gather_mla_latent_cache_inputs,
297
+)
298
from vllm.v1.attention.selector import get_attn_backend
299
from vllm.v1.kv_cache_interface import (
300
AttentionSpec,
vllm/model_executor/layers/attention/sparse_mla_attention.py
+1
-1
@@ -32,7 +32,7 @@ from vllm.utils.torch_utils import is_quantized_kv_cache, np_to_pinned_tensor
32
from vllm.v1.attention.backend import AttentionMetadata, AttentionMetadataBuilder
33
from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
34
from vllm.v1.attention.backends.utils import split_decodes_and_prefills
35
-from vllm.v1.attention.ops.dcp_utils import MLADCPManager
35
+from vllm.v1.attention.ops.dcp import MLADCPManager
36
from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
37
38
if TYPE_CHECKING:
vllm/model_executor/layers/sparse_attn_indexer.py
+1
-1
@@ -13,7 +13,6 @@ from vllm.distributed import get_dcp_group, get_pcp_group
13
from vllm.forward_context import get_forward_context
14
from vllm.logger import init_logger
15
from vllm.model_executor.custom_op import CustomOp
16
-from vllm.model_executor.layers.attention.pcp import maybe_gather_indexer_k
16
from vllm.model_executor.layers.quantization.utils.quant_utils import (
17
get_fp8_min_max,
18
)
@@ -35,6 +34,7 @@ from vllm.v1.attention.backends.mla.indexer import (
34
DeepseekV32IndexerMetadata,
35
)
36
from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
37
+from vllm.v1.attention.ops.pcp import maybe_gather_indexer_k
38
from vllm.v1.worker.workspace import current_workspace_manager
39
40
logger = init_logger(__name__)
vllm/models/deepseek_v32/attention.py
+4
-4
@@ -13,10 +13,6 @@ from vllm.distributed.parallel_state import get_tp_group
13
from vllm.forward_context import get_forward_context
14
from vllm.model_executor.layers.attention import MLAAttention
15
from vllm.model_executor.layers.attention.attention import get_attention_context
16
-from vllm.model_executor.layers.attention.pcp import (
17
- finalize_mla_pcp_decode,
18
- maybe_gather_mla_latent_cache_inputs,
19
-)
16
from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm
17
from vllm.model_executor.layers.linear import (
18
ColumnParallelLinear,
@@ -42,6 +38,10 @@ from vllm.model_executor.models.utils import extract_layer_index
38
from vllm.models.deepseek_v32.common.kernels import fused_norm_rope, fused_q
39
from vllm.platforms import current_platform
40
from vllm.utils.torch_utils import is_quantized_kv_cache
41
+from vllm.v1.attention.ops.pcp import (
42
+ finalize_mla_pcp_decode,
43
+ maybe_gather_mla_latent_cache_inputs,
44
+)
45
46
if TYPE_CHECKING:
47
from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata
vllm/models/kimi_k3/nvidia/mla.py
+1
-1
@@ -92,7 +92,7 @@ from vllm.v1.attention.backend import (
92
MLAAttentionImpl,
93
)
94
from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
95
-from vllm.v1.attention.ops.dcp_utils import MLADCPManager
95
+from vllm.v1.attention.ops.dcp import MLADCPManager
96
from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
97
from vllm.v1.attention.selector import get_attn_backend
98
from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec, get_kv_quant_mode
vllm/v1/attention/backend.py
+1
-1
@@ -886,7 +886,7 @@ class AttentionImplBase(ABC, Generic[T]):
886
# False => base 2 (lse = log2(sum(exp(qk))))
887
# -- e.g. FlashInfer trtllm-gen MLA
888
# The DCP combine kernel (cp_lse_ag_out_rs / dcp_a2a_lse_reduce in
889
- # vllm/v1/attention/ops/common.py) branches on this via its IS_BASE_E
889
+ # vllm/v1/attention/ops/dcp.py) branches on this via its IS_BASE_E
890
# constexpr; getting it wrong silently corrupts the cross-shard
891
# softmax denominator.
892
lse_base_on_e: bool = True
vllm/v1/attention/backends/flash_attn.py
+4
-2
@@ -34,8 +34,10 @@ from vllm.v1.attention.backends.utils import (
34
fill_mm_prefix_query_ranges,
35
get_dcp_local_seq_lens,
36
)
37
-from vllm.v1.attention.ops.common import cp_lse_ag_out_rs
38
-from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
37
+from vllm.v1.attention.ops.dcp import (
38
+ cp_lse_ag_out_rs,
39
+ dcp_a2a_lse_reduce,
40
+)
41
from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
42
from vllm.v1.worker.workspace import current_workspace_manager
43
vllm/v1/attention/backends/flashinfer.py
+4
-2
@@ -75,8 +75,10 @@ from vllm.v1.attention.backends.utils import (
75
infer_global_hyperparameters,
76
split_decodes_and_prefills,
77
)
78
-from vllm.v1.attention.ops.common import cp_lse_ag_out_rs
79
-from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
78
+from vllm.v1.attention.ops.dcp import (
79
+ cp_lse_ag_out_rs,
80
+ dcp_a2a_lse_reduce,
81
+)
82
from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
83
from vllm.v1.kv_cache_interface import (
84
AttentionSpec,
vllm/v1/attention/ops/common.py
-299
@@ -2,308 +2,9 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import torch
4
5
-from vllm.distributed.parallel_state import GroupCoordinator
5
from vllm.triton_utils import tl, triton
6
7
9
-def mask_dcp_empty_shards_(
10
- lse: torch.Tensor,
11
- seq_lens: torch.Tensor | None,
12
- query_start_loc: torch.Tensor | None,
13
-) -> None:
14
- if seq_lens is None and query_start_loc is None:
15
- return
16
- if seq_lens is None or query_start_loc is None:
17
- raise ValueError("seq_lens and query_start_loc must be provided together")
18
- if (
19
- seq_lens.ndim != 1
20
- or query_start_loc.ndim != 1
21
- or query_start_loc.shape[0] != seq_lens.shape[0] + 1
22
- ):
23
- raise ValueError("query_start_loc must contain one boundary per sequence")
24
-
25
- row_indices = torch.arange(
26
- lse.shape[0], device=lse.device, dtype=query_start_loc.dtype
27
- )
28
- sequence_indices = torch.searchsorted(
29
- query_start_loc[1:], row_indices, right=True
30
- ).clamp_max(seq_lens.shape[0] - 1)
31
- empty_rows = (row_indices >= query_start_loc[-1]) | (
32
- seq_lens[sequence_indices] == 0
33
- )
34
- lse.masked_fill_(empty_rows[:, None], float("-inf"))
35
-
36
-
37
-@triton.jit
38
-def _correct_attn_cp_out_kernel(
39
- outputs_ptr,
40
- new_output_ptr,
41
- lses_ptr,
42
- vlse_ptr,
43
- outputs_stride_B,
44
- outputs_stride_H,
45
- outputs_stride_D,
46
- lses_stride_N,
47
- lses_stride_B,
48
- lses_stride_H,
49
- lse_idx,
50
- HEAD_DIM: tl.constexpr,
51
- N_ROUNDED: tl.constexpr,
52
- IS_BASE_E: tl.constexpr,
53
-):
54
- """
55
- Apply the all-gathered lses to correct each local rank's attention
56
- output. we still need perform a cross-rank reduction to obtain the
57
- final attention output.
58
-
59
- Args:
60
- outputs_ptr (triton.PointerType):
61
- Pointer to input tensor of shape [ B, H, D ]
62
- lses_ptr (triton.PointerType):
63
- Pointer to input tensor of shape [ N, B, H ]
64
- new_output_ptr (triton.PointerType):
65
- Pointer to output tensor of shape [ B, H, D ]
66
- vlse_ptr (triton.PointerType):
67
- Pointer to output tensor of shape [ B, H ]
68
- """
69
- batch_idx = tl.program_id(axis=0).to(tl.int64)
70
- head_idx = tl.program_id(axis=1).to(tl.int64)
71
- d_offsets = tl.arange(0, HEAD_DIM)
72
- num_n_offsets = tl.arange(0, N_ROUNDED)
73
-
74
- # shape = [N]
75
- lse_offsets = (
76
- num_n_offsets * lses_stride_N
77
- + batch_idx * lses_stride_B
78
- + head_idx * lses_stride_H
79
- )
80
-
81
- # calc final lse
82
- lse = tl.load(lses_ptr + lse_offsets).to(tl.float32)
83
- lse = tl.where((lse != lse) | (lse == float("inf")), -float("inf"), lse)
84
- lse_max = tl.max(lse, axis=0)
85
- lse_max = tl.where(lse_max == -float("inf"), 0, lse_max)
86
- lse -= lse_max
87
- if IS_BASE_E:
88
- lse_exp = tl.exp(lse)
89
- lse_acc = tl.sum(lse_exp, axis=0)
90
- lse = tl.log(lse_acc)
91
- else:
92
- lse_exp = tl.exp2(lse)
93
- lse_acc = tl.sum(lse_exp, axis=0)
94
- lse = tl.log2(lse_acc)
95
- lse += lse_max
96
-
97
- lse_offsets = batch_idx * lses_stride_B + head_idx * lses_stride_H
98
- tl.store(vlse_ptr + lse_offsets, lse)
99
-
100
- # shape = [D]
101
- output_offsets = (
102
- batch_idx * outputs_stride_B
103
- + head_idx * outputs_stride_H
104
- + d_offsets * outputs_stride_D
105
- )
106
-
107
- # correct output
108
- lse_offset = (
109
- lse_idx * lses_stride_N + batch_idx * lses_stride_B + head_idx * lses_stride_H
110
- )
111
- lse_tmp = tl.load(lses_ptr + lse_offset).to(tl.float32)
112
- lse_finally = lse_tmp - lse
113
- lse_finally = tl.where(
114
- (lse_finally != lse_finally) | (lse_finally == float("inf")),
115
- -float("inf"),
116
- lse_finally,
117
- )
118
- factor = tl.exp(lse_finally) if IS_BASE_E else tl.exp2(lse_finally)
119
- output = tl.load(outputs_ptr + output_offsets)
120
- output = output * factor
121
- output = tl.where(factor == 0.0, 0.0, output)
122
-
123
- tl.store(new_output_ptr + output_offsets, output)
124
-
125
-
126
-class CPTritonContext:
127
- """The CPTritonContext is used to avoid recompilation of the Triton JIT."""
128
-
129
- def __init__(self):
130
- self.inner_kernel = None
131
-
132
- def call_kernel(self, kernel, grid, *regular_args, **const_args):
133
- if self.inner_kernel is None:
134
- self.inner_kernel = kernel[grid](*regular_args, **const_args)
135
- else:
136
- self.inner_kernel[grid](*regular_args)
137
-
138
-
139
-def correct_attn_out(
140
- out: torch.Tensor,
141
- lses: torch.Tensor,
142
- cp_rank: int,
143
- ctx: CPTritonContext,
144
- is_lse_base_on_e: bool = True,
145
-) -> tuple[torch.Tensor, torch.Tensor]:
146
- """Correct the attention output using the all-gathered lses.
147
-
148
- Args:
149
- out: Tensor of shape [ B, H, D ]
150
- lses: Tensor of shape [ N, B, H ]
151
- cp_rank: Current rank in the context-parallel group
152
- ctx: Triton context to avoid recompilation
153
-
154
- Returns:
155
- Tuple of (out, lse) with corrected attention and final log-sum-exp.
156
- """
157
- if ctx is None:
158
- ctx = CPTritonContext()
159
-
160
- # --- Normalize to 3D views ---
161
- if out.ndim == 4 and out.shape[1] == 1:
162
- out = out.squeeze(1)
163
- assert out.ndim == 3, f"expected out [B,H,D] or [B,1,H,D], got {tuple(out.shape)}"
164
-
165
- if lses.ndim == 4 and lses.shape[-1] == 1:
166
- lses = lses.squeeze(-1)
167
- if lses.ndim == 4 and lses.shape[1] == 1:
168
- lses = lses.squeeze(1)
169
- assert lses.ndim == 3, (
170
- f"expected lses [N,B,H] (optionally with a 1-sized extra dim), "
171
- f"got {tuple(lses.shape)}"
172
- )
173
-
174
- B, H, D = out.shape
175
- N = lses.shape[0]
176
-
177
- # Strides after we normalized shapes to 3-D views. The kernel computes
178
- # offsets for `vlse_ptr` using lses_stride_B/H, so the output buffer must
179
- # have the same B/H stride layout as a slice of `lses`.
180
- o_sB, o_sH, o_sD = out.stride()
181
- l_sN, l_sB, l_sH = lses.stride()
182
-
183
- # Allocate LSE with the same B/H strides as `lses` so writes land correctly
184
- # even when `lses` is a non-contiguous view (e.g., 4-D to 3-D squeeze).
185
- lse = torch.empty_strided(
186
- (B, H), (l_sB, l_sH), device=lses.device, dtype=lses.dtype
187
- )
188
-
189
- # Kernel launch config
190
- grid = (B, H, 1)
191
-
192
- regular_args = (
193
- out,
194
- out,
195
- lses,
196
- lse,
197
- o_sB,
198
- o_sH,
199
- o_sD,
200
- l_sN,
201
- l_sB,
202
- l_sH,
203
- cp_rank,
204
- )
205
- const_args = {"HEAD_DIM": D, "N_ROUNDED": N, "IS_BASE_E": is_lse_base_on_e}
206
- ctx.call_kernel(_correct_attn_cp_out_kernel, grid, *regular_args, **const_args)
207
- return out, lse
208
-
209
-
210
-def _cp_lse_common(
211
- cp_attn_out: torch.Tensor,
212
- cp_attn_lse: torch.Tensor,
213
- cp_group: GroupCoordinator,
214
- ctx: CPTritonContext | None = None,
215
- is_lse_base_on_e=True,
216
- seq_lens: torch.Tensor | None = None,
217
- query_start_loc: torch.Tensor | None = None,
218
-):
219
- """
220
- cp_attn_out: [ B, H, D ]
221
- cp_attn_lse: [ B, H ]
222
- """
223
- if cp_group.world_size == 1:
224
- return cp_attn_out
225
-
226
- if ctx is None:
227
- ctx = CPTritonContext()
228
-
229
- cp_attn_lse = cp_attn_lse.contiguous()
230
- mask_dcp_empty_shards_(cp_attn_lse, seq_lens, query_start_loc)
231
- lses = cp_group.all_gather(cp_attn_lse, dim=0).reshape(
232
- (cp_group.world_size,) + cp_attn_lse.shape
233
- )
234
- out, lse = correct_attn_out(
235
- cp_attn_out,
236
- lses,
237
- cp_group.rank_in_group,
238
- ctx,
239
- is_lse_base_on_e=is_lse_base_on_e,
240
- )
241
- return out, lse
242
-
243
-
244
-def cp_lse_ag_out_rs(
245
- cp_attn_out: torch.Tensor,
246
- cp_attn_lse: torch.Tensor,
247
- cp_group: GroupCoordinator,
248
- ctx: CPTritonContext | None = None,
249
- return_lse: bool = False,
250
- is_lse_base_on_e=True,
251
- seq_lens: torch.Tensor | None = None,
252
- query_start_loc: torch.Tensor | None = None,
253
-):
254
- """
255
- cp_attn_out: [ B, H, D ]
256
- cp_attn_lse: [ B, H ]
257
- """
258
- out, lse = _cp_lse_common(
259
- cp_attn_out,
260
- cp_attn_lse,
261
- cp_group,
262
- ctx=ctx,
263
- is_lse_base_on_e=is_lse_base_on_e,
264
- seq_lens=seq_lens,
265
- query_start_loc=query_start_loc,
266
- )
267
- out = cp_group.reduce_scatter(out, dim=1)
268
-
269
- if return_lse:
270
- cp_num_heads = lse.shape[1] // cp_group.world_size
271
- cp_rank = cp_group.rank_in_group
272
- lse = lse[:, cp_num_heads * cp_rank : cp_num_heads * (cp_rank + 1)]
273
- return out, lse
274
- return out
275
-
276
-
277
-def cp_lse_ag_out_ar(
278
- cp_attn_out: torch.Tensor,
279
- cp_attn_lse: torch.Tensor,
280
- cp_group: GroupCoordinator,
281
- ctx: CPTritonContext | None = None,
282
- return_lse: bool = False,
283
- is_lse_base_on_e=True,
284
- seq_lens: torch.Tensor | None = None,
285
- query_start_loc: torch.Tensor | None = None,
286
-):
287
- """
288
- cp_attn_out: [ B, H, D ]
289
- cp_attn_lse: [ B, H ]
290
- """
291
- out, lse = _cp_lse_common(
292
- cp_attn_out,
293
- cp_attn_lse,
294
- cp_group,
295
- ctx=ctx,
296
- is_lse_base_on_e=is_lse_base_on_e,
297
- seq_lens=seq_lens,
298
- query_start_loc=query_start_loc,
299
- )
300
- out = cp_group.all_reduce(out)
301
-
302
- if return_lse:
303
- return out, lse
304
- return out
305
-
306
-
8
@triton.jit
9
def _pack_seq_kernel(
10
x_ptr, # [N, D]
vllm/v1/attention/ops/cp_common.py
new
+154
@@ -0,0 +1,154 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Shared symmetric-memory infrastructure for context-parallel attention."""
4
+
5
+from __future__ import annotations
6
+
7
+import functools
8
+from typing import TYPE_CHECKING, Any
9
+
10
+import torch
11
+
12
+from vllm.distributed.parallel_state import in_the_same_node_as
13
+from vllm.logger import init_logger
14
+from vllm.platforms import current_platform
15
+
16
+if TYPE_CHECKING:
17
+ from torch.distributed import ProcessGroup
18
+
19
+ from vllm.distributed.parallel_state import GroupCoordinator
20
+
21
+logger = init_logger(__name__)
22
+
23
+try:
24
+ import torch.distributed._symmetric_memory as symm_mem
25
+
26
+ symm_mem_available = True
27
+except ImportError:
28
+ symm_mem = None # type: ignore[assignment]
29
+ symm_mem_available = False
30
+
31
+
32
+@functools.cache
33
+def _symm_mem_spans_group(group: GroupCoordinator) -> bool:
34
+ """Probe whether the group has NVLS symmetric memory."""
35
+ if not symm_mem_available:
36
+ return False
37
+ try:
38
+ from torch._C._autograd import DeviceType
39
+ from torch._C._distributed_c10d import _SymmetricMemory
40
+
41
+ device = torch.device("cuda", torch.accelerator.current_device_index())
42
+ if not _SymmetricMemory.has_multicast_support(DeviceType.CUDA, device.index):
43
+ return False
44
+ probe = symm_mem.empty(8, dtype=torch.uint8, device=device)
45
+ probe.zero_()
46
+ torch.accelerator.synchronize()
47
+ handle = symm_mem.rendezvous(probe, group.device_group.group_name)
48
+ spans = handle is not None and handle.multicast_ptr != 0
49
+ except Exception as error:
50
+ logger.debug("Direct CP symmetric-memory probe failed: %s", error)
51
+ return False
52
+ logger.debug_once(
53
+ "Direct CP symmetric memory across %d ranks: %s",
54
+ group.world_size,
55
+ "available" if spans else "unavailable",
56
+ )
57
+ return spans
58
+
59
+
60
+def direct_cp_enabled(
61
+ group: GroupCoordinator,
62
+ dtype: torch.dtype,
63
+ use_direct: bool | None,
64
+ supported_dtypes: tuple[torch.dtype, ...] | None = None,
65
+) -> bool:
66
+ if use_direct is not None:
67
+ return use_direct
68
+ return (
69
+ symm_mem_available
70
+ and current_platform.is_cuda()
71
+ and (supported_dtypes is None or dtype in supported_dtypes)
72
+ and (
73
+ all(in_the_same_node_as(group.cpu_group, source_rank=0))
74
+ or _symm_mem_spans_group(group)
75
+ )
76
+ )
77
+
78
+
79
+def direct_cp_multicast_enabled(
80
+ group: GroupCoordinator,
81
+ dtype: torch.dtype,
82
+ use_direct: bool | None,
83
+ supported_dtypes: tuple[torch.dtype, ...] | None = None,
84
+) -> bool:
85
+ return direct_cp_enabled(
86
+ group, dtype, use_direct, supported_dtypes
87
+ ) and _symm_mem_spans_group(group)
88
+
89
+
90
+class DirectCPWorkspace:
91
+ def __init__(
92
+ self,
93
+ group: ProcessGroup,
94
+ device: torch.device,
95
+ num_ubatches: int,
96
+ ) -> None:
97
+ self.group = group
98
+ self.world_size = group.size()
99
+ self.rank = group.rank()
100
+ self.device = torch.device(device)
101
+ self.num_ubatches = num_ubatches
102
+ self.epoch = torch.zeros(num_ubatches, dtype=torch.int64, device=self.device)
103
+ self._allocations: list[tuple[torch.Tensor, Any, list[torch.Tensor]]] = []
104
+
105
+ def _allocate(
106
+ self, shape: tuple[int, ...], dtype: torch.dtype
107
+ ) -> tuple[torch.Tensor, torch.Tensor]:
108
+ storage = symm_mem.empty(shape, device=self.device, dtype=dtype)
109
+ storage.zero_()
110
+ torch.accelerator.synchronize()
111
+ handle = symm_mem.rendezvous(storage, self.group.group_name)
112
+ assert handle is not None, "CP symmetric memory rendezvous returned None"
113
+ handle.barrier()
114
+ views = [
115
+ handle.get_buffer(peer, list(shape), dtype, 0)
116
+ for peer in range(self.world_size)
117
+ ]
118
+ self.device = storage.device
119
+ peer_ptrs = torch.tensor(
120
+ [
121
+ [view[ubatch].data_ptr() for view in views]
122
+ for ubatch in range(self.num_ubatches)
123
+ ],
124
+ dtype=torch.int64,
125
+ device=self.device,
126
+ )
127
+ self._allocations.append((storage, handle, views))
128
+ return storage, peer_ptrs
129
+
130
+ def _multicast_ptrs(self, storage: torch.Tensor) -> list[int]:
131
+ disabled = [0] * self.num_ubatches
132
+ for allocated, handle, _ in self._allocations:
133
+ if allocated is storage:
134
+ break
135
+ else:
136
+ return disabled
137
+ try:
138
+ from torch._C._autograd import DeviceType
139
+ from torch._C._distributed_c10d import _SymmetricMemory
140
+
141
+ if not _SymmetricMemory.has_multicast_support(
142
+ DeviceType.CUDA, storage.device.index
143
+ ):
144
+ return disabled
145
+ multicast_base = handle.multicast_ptr
146
+ except Exception:
147
+ return disabled
148
+ if not multicast_base:
149
+ return disabled
150
+ storage_base = storage.data_ptr()
151
+ return [
152
+ multicast_base + (storage[ubatch].data_ptr() - storage_base)
153
+ for ubatch in range(self.num_ubatches)
154
+ ]
vllm/v1/attention/ops/dcp.py
new
+1368
@@ -0,0 +1,1368 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""MLA DCP collective selection and direct symmetric-memory implementations."""
4
+
5
+from __future__ import annotations
6
+
7
+import functools
8
+from collections.abc import Callable
9
+from typing import TYPE_CHECKING, Protocol
10
+
11
+import torch
12
+import torch.distributed as dist
13
+
14
+import vllm.envs as envs
15
+from vllm.config import VllmConfig
16
+from vllm.distributed import get_dcp_group
17
+from vllm.logger import init_logger
18
+from vllm.triton_utils import tl, triton
19
+from vllm.v1.attention.ops.cp_common import (
20
+ DirectCPWorkspace,
21
+ direct_cp_enabled,
22
+ direct_cp_multicast_enabled,
23
+)
24
+from vllm.v1.worker.ubatching import dbo_current_ubatch_id
25
+
26
+logger = init_logger(__name__)
27
+
28
+if TYPE_CHECKING:
29
+ from torch.distributed import ProcessGroup
30
+
31
+ from vllm.distributed.parallel_state import GroupCoordinator
32
+
33
+
34
+# LSE/output combine
35
+
36
+
37
+def mask_dcp_empty_shards_(
38
+ lse: torch.Tensor,
39
+ seq_lens: torch.Tensor | None,
40
+ query_start_loc: torch.Tensor | None,
41
+) -> None:
42
+ if seq_lens is None and query_start_loc is None:
43
+ return
44
+ if seq_lens is None or query_start_loc is None:
45
+ raise ValueError("seq_lens and query_start_loc must be provided together")
46
+ if (
47
+ seq_lens.ndim != 1
48
+ or query_start_loc.ndim != 1
49
+ or query_start_loc.shape[0] != seq_lens.shape[0] + 1
50
+ ):
51
+ raise ValueError("query_start_loc must contain one boundary per sequence")
52
+
53
+ row_indices = torch.arange(
54
+ lse.shape[0], device=lse.device, dtype=query_start_loc.dtype
55
+ )
56
+ sequence_indices = torch.searchsorted(
57
+ query_start_loc[1:], row_indices, right=True
58
+ ).clamp_max(seq_lens.shape[0] - 1)
59
+ empty_rows = (row_indices >= query_start_loc[-1]) | (
60
+ seq_lens[sequence_indices] == 0
61
+ )
62
+ lse.masked_fill_(empty_rows[:, None], float("-inf"))
63
+
64
+
65
+# AG + RS/AR implementation
66
+
67
+
68
+@triton.jit
69
+def _correct_attn_cp_out_kernel(
70
+ outputs_ptr,
71
+ new_output_ptr,
72
+ lses_ptr,
73
+ vlse_ptr,
74
+ outputs_stride_B,
75
+ outputs_stride_H,
76
+ outputs_stride_D,
77
+ lses_stride_N,
78
+ lses_stride_B,
79
+ lses_stride_H,
80
+ lse_idx,
81
+ HEAD_DIM: tl.constexpr,
82
+ N_ROUNDED: tl.constexpr,
83
+ IS_BASE_E: tl.constexpr,
84
+):
85
+ """
86
+ Apply the all-gathered lses to correct each local rank's attention
87
+ output. we still need perform a cross-rank reduction to obtain the
88
+ final attention output.
89
+
90
+ Args:
91
+ outputs_ptr (triton.PointerType):
92
+ Pointer to input tensor of shape [ B, H, D ]
93
+ lses_ptr (triton.PointerType):
94
+ Pointer to input tensor of shape [ N, B, H ]
95
+ new_output_ptr (triton.PointerType):
96
+ Pointer to output tensor of shape [ B, H, D ]
97
+ vlse_ptr (triton.PointerType):
98
+ Pointer to output tensor of shape [ B, H ]
99
+ """
100
+ batch_idx = tl.program_id(axis=0).to(tl.int64)
101
+ head_idx = tl.program_id(axis=1).to(tl.int64)
102
+ d_offsets = tl.arange(0, HEAD_DIM)
103
+ num_n_offsets = tl.arange(0, N_ROUNDED)
104
+
105
+ # shape = [N]
106
+ lse_offsets = (
107
+ num_n_offsets * lses_stride_N
108
+ + batch_idx * lses_stride_B
109
+ + head_idx * lses_stride_H
110
+ )
111
+
112
+ # calc final lse
113
+ lse = tl.load(lses_ptr + lse_offsets).to(tl.float32)
114
+ lse = tl.where((lse != lse) | (lse == float("inf")), -float("inf"), lse)
115
+ lse_max = tl.max(lse, axis=0)
116
+ lse_max = tl.where(lse_max == -float("inf"), 0, lse_max)
117
+ lse -= lse_max
118
+ if IS_BASE_E:
119
+ lse_exp = tl.exp(lse)
120
+ lse_acc = tl.sum(lse_exp, axis=0)
121
+ lse = tl.log(lse_acc)
122
+ else:
123
+ lse_exp = tl.exp2(lse)
124
+ lse_acc = tl.sum(lse_exp, axis=0)
125
+ lse = tl.log2(lse_acc)
126
+ lse += lse_max
127
+
128
+ lse_offsets = batch_idx * lses_stride_B + head_idx * lses_stride_H
129
+ tl.store(vlse_ptr + lse_offsets, lse)
130
+
131
+ # shape = [D]
132
+ output_offsets = (
133
+ batch_idx * outputs_stride_B
134
+ + head_idx * outputs_stride_H
135
+ + d_offsets * outputs_stride_D
136
+ )
137
+
138
+ # correct output
139
+ lse_offset = (
140
+ lse_idx * lses_stride_N + batch_idx * lses_stride_B + head_idx * lses_stride_H
141
+ )
142
+ lse_tmp = tl.load(lses_ptr + lse_offset).to(tl.float32)
143
+ lse_finally = lse_tmp - lse
144
+ lse_finally = tl.where(
145
+ (lse_finally != lse_finally) | (lse_finally == float("inf")),
146
+ -float("inf"),
147
+ lse_finally,
148
+ )
149
+ factor = tl.exp(lse_finally) if IS_BASE_E else tl.exp2(lse_finally)
150
+ output = tl.load(outputs_ptr + output_offsets)
151
+ output = output * factor
152
+ output = tl.where(factor == 0.0, 0.0, output)
153
+
154
+ tl.store(new_output_ptr + output_offsets, output)
155
+
156
+
157
+class CPTritonContext:
158
+ """The CPTritonContext is used to avoid recompilation of the Triton JIT."""
159
+
160
+ def __init__(self):
161
+ self.inner_kernel = None
162
+
163
+ def call_kernel(self, kernel, grid, *regular_args, **const_args):
164
+ if self.inner_kernel is None:
165
+ self.inner_kernel = kernel[grid](*regular_args, **const_args)
166
+ else:
167
+ self.inner_kernel[grid](*regular_args)
168
+
169
+
170
+def correct_attn_out(
171
+ out: torch.Tensor,
172
+ lses: torch.Tensor,
173
+ cp_rank: int,
174
+ ctx: CPTritonContext,
175
+ is_lse_base_on_e: bool = True,
176
+) -> tuple[torch.Tensor, torch.Tensor]:
177
+ """Correct the attention output using the all-gathered lses.
178
+
179
+ Args:
180
+ out: Tensor of shape [ B, H, D ]
181
+ lses: Tensor of shape [ N, B, H ]
182
+ cp_rank: Current rank in the context-parallel group
183
+ ctx: Triton context to avoid recompilation
184
+
185
+ Returns:
186
+ Tuple of (out, lse) with corrected attention and final log-sum-exp.
187
+ """
188
+ if ctx is None:
189
+ ctx = CPTritonContext()
190
+
191
+ # --- Normalize to 3D views ---
192
+ if out.ndim == 4 and out.shape[1] == 1:
193
+ out = out.squeeze(1)
194
+ assert out.ndim == 3, f"expected out [B,H,D] or [B,1,H,D], got {tuple(out.shape)}"
195
+
196
+ if lses.ndim == 4 and lses.shape[-1] == 1:
197
+ lses = lses.squeeze(-1)
198
+ if lses.ndim == 4 and lses.shape[1] == 1:
199
+ lses = lses.squeeze(1)
200
+ assert lses.ndim == 3, (
201
+ f"expected lses [N,B,H] (optionally with a 1-sized extra dim), "
202
+ f"got {tuple(lses.shape)}"
203
+ )
204
+
205
+ B, H, D = out.shape
206
+ N = lses.shape[0]
207
+
208
+ # Strides after we normalized shapes to 3-D views. The kernel computes
209
+ # offsets for `vlse_ptr` using lses_stride_B/H, so the output buffer must
210
+ # have the same B/H stride layout as a slice of `lses`.
211
+ o_sB, o_sH, o_sD = out.stride()
212
+ l_sN, l_sB, l_sH = lses.stride()
213
+
214
+ # Allocate LSE with the same B/H strides as `lses` so writes land correctly
215
+ # even when `lses` is a non-contiguous view (e.g., 4-D to 3-D squeeze).
216
+ lse = torch.empty_strided(
217
+ (B, H), (l_sB, l_sH), device=lses.device, dtype=lses.dtype
218
+ )
219
+
220
+ # Kernel launch config
221
+ grid = (B, H, 1)
222
+
223
+ regular_args = (
224
+ out,
225
+ out,
226
+ lses,
227
+ lse,
228
+ o_sB,
229
+ o_sH,
230
+ o_sD,
231
+ l_sN,
232
+ l_sB,
233
+ l_sH,
234
+ cp_rank,
235
+ )
236
+ const_args = {"HEAD_DIM": D, "N_ROUNDED": N, "IS_BASE_E": is_lse_base_on_e}
237
+ ctx.call_kernel(_correct_attn_cp_out_kernel, grid, *regular_args, **const_args)
238
+ return out, lse
239
+
240
+
241
+def _cp_lse_common(
242
+ cp_attn_out: torch.Tensor,
243
+ cp_attn_lse: torch.Tensor,
244
+ cp_group: GroupCoordinator,
245
+ ctx: CPTritonContext | None = None,
246
+ is_lse_base_on_e=True,
247
+ seq_lens: torch.Tensor | None = None,
248
+ query_start_loc: torch.Tensor | None = None,
249
+):
250
+ """
251
+ cp_attn_out: [ B, H, D ]
252
+ cp_attn_lse: [ B, H ]
253
+ """
254
+ if cp_group.world_size == 1:
255
+ return cp_attn_out
256
+
257
+ if ctx is None:
258
+ ctx = CPTritonContext()
259
+
260
+ cp_attn_lse = cp_attn_lse.contiguous()
261
+ mask_dcp_empty_shards_(cp_attn_lse, seq_lens, query_start_loc)
262
+ lses = cp_group.all_gather(cp_attn_lse, dim=0).reshape(
263
+ (cp_group.world_size,) + cp_attn_lse.shape
264
+ )
265
+ out, lse = correct_attn_out(
266
+ cp_attn_out,
267
+ lses,
268
+ cp_group.rank_in_group,
269
+ ctx,
270
+ is_lse_base_on_e=is_lse_base_on_e,
271
+ )
272
+ return out, lse
273
+
274
+
275
+def cp_lse_ag_out_rs(
276
+ cp_attn_out: torch.Tensor,
277
+ cp_attn_lse: torch.Tensor,
278
+ cp_group: GroupCoordinator,
279
+ ctx: CPTritonContext | None = None,
280
+ return_lse: bool = False,
281
+ is_lse_base_on_e=True,
282
+ seq_lens: torch.Tensor | None = None,
283
+ query_start_loc: torch.Tensor | None = None,
284
+):
285
+ """
286
+ cp_attn_out: [ B, H, D ]
287
+ cp_attn_lse: [ B, H ]
288
+ """
289
+ out, lse = _cp_lse_common(
290
+ cp_attn_out,
291
+ cp_attn_lse,
292
+ cp_group,
293
+ ctx=ctx,
294
+ is_lse_base_on_e=is_lse_base_on_e,
295
+ seq_lens=seq_lens,
296
+ query_start_loc=query_start_loc,
297
+ )
298
+ out = cp_group.reduce_scatter(out, dim=1)
299
+
300
+ if return_lse:
301
+ cp_num_heads = lse.shape[1] // cp_group.world_size
302
+ cp_rank = cp_group.rank_in_group
303
+ lse = lse[:, cp_num_heads * cp_rank : cp_num_heads * (cp_rank + 1)]
304
+ return out, lse
305
+ return out
306
+
307
+
308
+def cp_lse_ag_out_ar(
309
+ cp_attn_out: torch.Tensor,
310
+ cp_attn_lse: torch.Tensor,
311
+ cp_group: GroupCoordinator,
312
+ ctx: CPTritonContext | None = None,
313
+ return_lse: bool = False,
314
+ is_lse_base_on_e=True,
315
+ seq_lens: torch.Tensor | None = None,
316
+ query_start_loc: torch.Tensor | None = None,
317
+):
318
+ """
319
+ cp_attn_out: [ B, H, D ]
320
+ cp_attn_lse: [ B, H ]
321
+ """
322
+ out, lse = _cp_lse_common(
323
+ cp_attn_out,
324
+ cp_attn_lse,
325
+ cp_group,
326
+ ctx=ctx,
327
+ is_lse_base_on_e=is_lse_base_on_e,
328
+ seq_lens=seq_lens,
329
+ query_start_loc=query_start_loc,
330
+ )
331
+ out = cp_group.all_reduce(out)
332
+
333
+ if return_lse:
334
+ return out, lse
335
+ return out
336
+
337
+
338
+# Standard A2A implementation
339
+
340
+
341
+def _lse_weighted_combine(
342
+ outputs: torch.Tensor,
343
+ lses: torch.Tensor,
344
+ return_lse: bool = False,
345
+ is_lse_base_on_e: bool = True,
346
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
347
+ """
348
+ CPU reference implementation for LSE-weighted combination.
349
+
350
+ This is a pure PyTorch implementation used for testing and validation.
351
+
352
+ Args:
353
+ outputs: Partial attention outputs [N, B, H, D]
354
+ N = number of KV shards (ranks)
355
+ B = batch size (num_tokens)
356
+ H = number of heads per rank
357
+ D = head dimension
358
+ lses: Log-sum-exp values [N, B, H]
359
+ return_lse: If True, also return the global LSE
360
+ is_lse_base_on_e: If True, LSE is base e; if False, base 2
361
+
362
+ Returns:
363
+ Combined output [B, H, D], and optionally global LSE [B, H]
364
+ """
365
+ N, B, H, D = outputs.shape
366
+
367
+ # Handle NaN and inf in LSEs
368
+ lses = torch.where(
369
+ torch.isnan(lses) | torch.isinf(lses),
370
+ torch.tensor(float("-inf"), device=lses.device, dtype=lses.dtype),
371
+ lses,
372
+ )
373
+
374
+ # Compute max LSE for numerical stability
375
+ lse_max, _ = lses.max(dim=0) # [B, H]
376
+ lse_max = torch.where(
377
+ lse_max == float("-inf"),
378
+ torch.zeros_like(lse_max),
379
+ lse_max,
380
+ )
381
+
382
+ # Compute weights: softmax over the N dimension
383
+ if is_lse_base_on_e:
384
+ weights = torch.exp(lses - lse_max.unsqueeze(0)) # [N, B, H]
385
+ else:
386
+ weights = torch.pow(2.0, lses - lse_max.unsqueeze(0)) # [N, B, H]
387
+
388
+ # Handle NaN weights
389
+ weights = torch.where(torch.isnan(weights), torch.zeros_like(weights), weights)
390
+
391
+ # Normalize weights
392
+ weight_sum = weights.sum(dim=0, keepdim=True) # [1, B, H]
393
+ weights = weights / weight_sum.clamp(min=1e-10) # [N, B, H]
394
+
395
+ # Weighted combination: sum over N dimension
396
+ weights = weights.unsqueeze(-1)
397
+ outputs = torch.where(weights == 0, torch.zeros_like(outputs), outputs)
398
+ result = (outputs * weights).sum(dim=0) # [B, H, D]
399
+
400
+ if return_lse:
401
+ if is_lse_base_on_e:
402
+ global_lse = torch.log(weight_sum.squeeze(0)) + lse_max # [B, H]
403
+ else:
404
+ global_lse = torch.log2(weight_sum.squeeze(0)) + lse_max # [B, H]
405
+ return result, global_lse
406
+
407
+ return result
408
+
409
+
410
+def _dcp_a2a_lse_pack_dim(output_dtype: torch.dtype) -> int:
411
+ bits = torch.finfo(output_dtype).bits
412
+ if bits == 16:
413
+ return 2
414
+ if bits == 32:
415
+ return 1
416
+ raise ValueError(f"Cannot pack fp32 LSE into output dtype {output_dtype}.")
417
+
418
+
419
+def _dcp_a2a_send_recv_buffers(
420
+ shape: tuple[int, ...],
421
+ device: torch.device,
422
+ dtype: torch.dtype,
423
+) -> tuple[torch.Tensor, torch.Tensor]:
424
+ # Don't use the shared WorkspaceManager here. A FULL cudagraph bakes in the
425
+ # buffer address at capture, but the workspace is growable and sized only to
426
+ # the largest *captured* batch (the cudagraph capture cap). Any eager a2a
427
+ # with a bigger batch regrows it, freeing that address and poisoning every
428
+ # captured graph -> illegal memory access on replay. This bites the very
429
+ # first request: the post-capture warmup runs an eager decode at
430
+ # max_num_seqs (> the cap), so the graphs are already dangling before the
431
+ # server is ready. torch.empty buffers instead live in the graph's private
432
+ # pool and stay valid for its lifetime (as _dcp_a2a_unpack_combine and the
433
+ # AG+RS combine path already rely on).
434
+ return (
435
+ torch.empty(shape, device=device, dtype=dtype),
436
+ torch.empty(shape, device=device, dtype=dtype),
437
+ )
438
+
439
+
440
+@triton.jit
441
+def _dcp_a2a_pack_send_kernel(
442
+ out_ptr,
443
+ lse_ptr,
444
+ send_ptr,
445
+ out_stride_B,
446
+ out_stride_H,
447
+ out_stride_D,
448
+ lse_stride_B,
449
+ lse_stride_H,
450
+ send_stride_N,
451
+ send_stride_B,
452
+ send_stride_H,
453
+ send_stride_D,
454
+ N: tl.constexpr,
455
+ HEAD_DIM: tl.constexpr,
456
+ H_PER_RANK: tl.constexpr,
457
+ LSE_PACK_DIM: tl.constexpr,
458
+):
459
+ batch_idx = tl.program_id(0).to(tl.int64)
460
+ local_head_idx = tl.program_id(1).to(tl.int64)
461
+ d_offsets = tl.arange(0, HEAD_DIM)
462
+
463
+ for rank_idx in tl.static_range(N):
464
+ src_head_idx = rank_idx * H_PER_RANK + local_head_idx
465
+ send_base = (
466
+ rank_idx * send_stride_N
467
+ + batch_idx * send_stride_B
468
+ + local_head_idx * send_stride_H
469
+ )
470
+
471
+ out_offsets = (
472
+ batch_idx * out_stride_B
473
+ + src_head_idx * out_stride_H
474
+ + d_offsets * out_stride_D
475
+ )
476
+ tl.store(
477
+ send_ptr + send_base + d_offsets * send_stride_D,
478
+ tl.load(out_ptr + out_offsets),
479
+ )
480
+
481
+ lse_val = tl.load(
482
+ lse_ptr + batch_idx * lse_stride_B + src_head_idx * lse_stride_H
483
+ ).to(tl.float32)
484
+ if LSE_PACK_DIM == 1:
485
+ tl.store(
486
+ send_ptr + send_base + HEAD_DIM * send_stride_D,
487
+ lse_val.to(send_ptr.dtype.element_ty),
488
+ )
489
+ else:
490
+ lse_bits = lse_val.to(tl.uint32, bitcast=True)
491
+ lo = (lse_bits & 0xFFFF).to(tl.uint16)
492
+ hi = ((lse_bits >> 16) & 0xFFFF).to(tl.uint16)
493
+ tl.store(
494
+ send_ptr + send_base + HEAD_DIM * send_stride_D,
495
+ lo.to(send_ptr.dtype.element_ty, bitcast=True),
496
+ )
497
+ tl.store(
498
+ send_ptr + send_base + (HEAD_DIM + 1) * send_stride_D,
499
+ hi.to(send_ptr.dtype.element_ty, bitcast=True),
500
+ )
501
+
502
+
503
+@triton.jit
504
+def _dcp_a2a_unpack_combine_kernel(
505
+ recv_ptr,
506
+ out_ptr,
507
+ out_lse_ptr,
508
+ recv_stride_N,
509
+ recv_stride_B,
510
+ recv_stride_H,
511
+ recv_stride_D,
512
+ out_stride_B,
513
+ out_stride_H,
514
+ out_stride_D,
515
+ out_lse_stride_B,
516
+ out_lse_stride_H,
517
+ N: tl.constexpr,
518
+ HEAD_DIM: tl.constexpr,
519
+ IS_BASE_E: tl.constexpr,
520
+ RETURN_LSE: tl.constexpr,
521
+ LSE_PACK_DIM: tl.constexpr,
522
+):
523
+ batch_idx = tl.program_id(0).to(tl.int64)
524
+ head_idx = tl.program_id(1).to(tl.int64)
525
+ d_offsets = tl.arange(0, HEAD_DIM)
526
+
527
+ lse_max = -float("inf")
528
+ for rank_idx in tl.static_range(N):
529
+ recv_base = (
530
+ rank_idx * recv_stride_N
531
+ + batch_idx * recv_stride_B
532
+ + head_idx * recv_stride_H
533
+ )
534
+ if LSE_PACK_DIM == 1:
535
+ lse_val = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D).to(
536
+ tl.float32
537
+ )
538
+ else:
539
+ lo_raw = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D)
540
+ hi_raw = tl.load(recv_ptr + recv_base + (HEAD_DIM + 1) * recv_stride_D)
541
+ lo = lo_raw.to(tl.uint16, bitcast=True).to(tl.uint32)
542
+ hi = hi_raw.to(tl.uint16, bitcast=True).to(tl.uint32)
543
+ lse_val = (lo | (hi << 16)).to(tl.float32, bitcast=True)
544
+ lse_val = tl.where(
545
+ (lse_val != lse_val) | (lse_val == float("inf")),
546
+ -float("inf"),
547
+ lse_val,
548
+ )
549
+ lse_max = tl.maximum(lse_max, lse_val)
550
+
551
+ lse_max = tl.where(lse_max == -float("inf"), 0.0, lse_max)
552
+
553
+ lse_sum = 0.0
554
+ for rank_idx in tl.static_range(N):
555
+ recv_base = (
556
+ rank_idx * recv_stride_N
557
+ + batch_idx * recv_stride_B
558
+ + head_idx * recv_stride_H
559
+ )
560
+ if LSE_PACK_DIM == 1:
561
+ lse_val = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D).to(
562
+ tl.float32
563
+ )
564
+ else:
565
+ lo_raw = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D)
566
+ hi_raw = tl.load(recv_ptr + recv_base + (HEAD_DIM + 1) * recv_stride_D)
567
+ lo = lo_raw.to(tl.uint16, bitcast=True).to(tl.uint32)
568
+ hi = hi_raw.to(tl.uint16, bitcast=True).to(tl.uint32)
569
+ lse_val = (lo | (hi << 16)).to(tl.float32, bitcast=True)
570
+ lse_val = tl.where(
571
+ (lse_val != lse_val) | (lse_val == float("inf")),
572
+ -float("inf"),
573
+ lse_val,
574
+ )
575
+ if IS_BASE_E:
576
+ lse_sum += tl.exp(lse_val - lse_max)
577
+ else:
578
+ lse_sum += tl.exp2(lse_val - lse_max)
579
+
580
+ if IS_BASE_E: # noqa: SIM108
581
+ global_lse = tl.log(lse_sum) + lse_max
582
+ else:
583
+ global_lse = tl.log2(lse_sum) + lse_max
584
+
585
+ acc = tl.zeros([HEAD_DIM], dtype=tl.float32)
586
+ for rank_idx in tl.static_range(N):
587
+ recv_base = (
588
+ rank_idx * recv_stride_N
589
+ + batch_idx * recv_stride_B
590
+ + head_idx * recv_stride_H
591
+ )
592
+ if LSE_PACK_DIM == 1:
593
+ lse_val = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D).to(
594
+ tl.float32
595
+ )
596
+ else:
597
+ lo_raw = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D)
598
+ hi_raw = tl.load(recv_ptr + recv_base + (HEAD_DIM + 1) * recv_stride_D)
599
+ lo = lo_raw.to(tl.uint16, bitcast=True).to(tl.uint32)
600
+ hi = hi_raw.to(tl.uint16, bitcast=True).to(tl.uint32)
601
+ lse_val = (lo | (hi << 16)).to(tl.float32, bitcast=True)
602
+ lse_val = tl.where(
603
+ (lse_val != lse_val) | (lse_val == float("inf")),
604
+ -float("inf"),
605
+ lse_val,
606
+ )
607
+ if IS_BASE_E:
608
+ weight = tl.exp(lse_val - global_lse)
609
+ else:
610
+ weight = tl.exp2(lse_val - global_lse)
611
+ weight = tl.where(weight != weight, 0.0, weight)
612
+ partial = tl.load(recv_ptr + recv_base + d_offsets * recv_stride_D).to(
613
+ tl.float32
614
+ )
615
+ partial = tl.where(weight == 0.0, 0.0, partial)
616
+ acc += partial * weight
617
+
618
+ final_offsets = (
619
+ batch_idx * out_stride_B + head_idx * out_stride_H + d_offsets * out_stride_D
620
+ )
621
+ tl.store(out_ptr + final_offsets, acc)
622
+
623
+ if RETURN_LSE:
624
+ out_lse_offset = batch_idx * out_lse_stride_B + head_idx * out_lse_stride_H
625
+ tl.store(out_lse_ptr + out_lse_offset, global_lse)
626
+
627
+
628
+def _dcp_a2a_pack_send(
629
+ cp_attn_out: torch.Tensor,
630
+ cp_attn_lse: torch.Tensor,
631
+ send_buffer: torch.Tensor,
632
+ world_size: int,
633
+ h_per_rank: int,
634
+ head_dim: int,
635
+ lse_pack_dim: int,
636
+ seq_lens: torch.Tensor | None = None,
637
+ query_start_loc: torch.Tensor | None = None,
638
+) -> None:
639
+ mask_dcp_empty_shards_(cp_attn_lse, seq_lens, query_start_loc)
640
+ grid = (cp_attn_out.shape[0], h_per_rank, 1)
641
+ _dcp_a2a_pack_send_kernel[grid](
642
+ cp_attn_out,
643
+ cp_attn_lse,
644
+ send_buffer,
645
+ cp_attn_out.stride(0),
646
+ cp_attn_out.stride(1),
647
+ cp_attn_out.stride(2),
648
+ cp_attn_lse.stride(0),
649
+ cp_attn_lse.stride(1),
650
+ send_buffer.stride(0),
651
+ send_buffer.stride(1),
652
+ send_buffer.stride(2),
653
+ send_buffer.stride(3),
654
+ N=world_size,
655
+ HEAD_DIM=head_dim,
656
+ H_PER_RANK=h_per_rank,
657
+ LSE_PACK_DIM=lse_pack_dim,
658
+ )
659
+
660
+
661
+def _dcp_a2a_unpack_combine(
662
+ recv_buffer: torch.Tensor,
663
+ head_dim: int,
664
+ lse_pack_dim: int,
665
+ return_lse: bool,
666
+ is_lse_base_on_e: bool,
667
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
668
+ world_size, num_tokens, h_per_rank, _ = recv_buffer.shape
669
+ out = torch.empty(
670
+ (num_tokens, h_per_rank, head_dim),
671
+ device=recv_buffer.device,
672
+ dtype=recv_buffer.dtype,
673
+ )
674
+ out_lse = torch.empty(
675
+ (num_tokens, h_per_rank) if return_lse else (1, 1),
676
+ device=recv_buffer.device,
677
+ dtype=torch.float32 if return_lse else recv_buffer.dtype,
678
+ )
679
+ grid = (num_tokens, h_per_rank, 1)
680
+ _dcp_a2a_unpack_combine_kernel[grid](
681
+ recv_buffer,
682
+ out,
683
+ out_lse,
684
+ recv_buffer.stride(0),
685
+ recv_buffer.stride(1),
686
+ recv_buffer.stride(2),
687
+ recv_buffer.stride(3),
688
+ out.stride(0),
689
+ out.stride(1),
690
+ out.stride(2),
691
+ out_lse.stride(0),
692
+ out_lse.stride(1),
693
+ N=world_size,
694
+ HEAD_DIM=head_dim,
695
+ IS_BASE_E=is_lse_base_on_e,
696
+ RETURN_LSE=return_lse,
697
+ LSE_PACK_DIM=lse_pack_dim,
698
+ )
699
+ if return_lse:
700
+ return out, out_lse
701
+ return out
702
+
703
+
704
+def dcp_a2a_lse_reduce(
705
+ cp_attn_out: torch.Tensor,
706
+ cp_attn_lse: torch.Tensor,
707
+ cp_group: GroupCoordinator,
708
+ ctx: CPTritonContext | None = None,
709
+ return_lse: bool = False,
710
+ is_lse_base_on_e: bool = True,
711
+ seq_lens: torch.Tensor | None = None,
712
+ query_start_loc: torch.Tensor | None = None,
713
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
714
+ """
715
+ Combine partial attention outputs across DCP ranks using All-to-All.
716
+
717
+ The output and LSE are packed into a single output-dtype buffer, sent
718
+ with one All-to-All, then unpacked and combined with exact LSE weighting.
719
+
720
+ Args:
721
+ cp_attn_out: [B, H, D] where B=num_tokens, H=total_heads, D=head_dim
722
+ cp_attn_lse: [B, H] floating-point log-sum-exp values
723
+ cp_group: GroupCoordinator for DCP communication
724
+ ctx: CPTritonContext (unused, for signature compatibility)
725
+ return_lse: If True, also return the combined global LSE
726
+ is_lse_base_on_e: If True, LSE is base e; if False, base 2
727
+ seq_lens: Local KV lengths. Empty shards contribute zero weight.
728
+ query_start_loc: Cumulative query-token offsets for each request.
729
+
730
+ Returns:
731
+ Combined output [B, H/N, D] (head-scattered)
732
+ If return_lse=True, also returns global_lse [B, H/N]
733
+ """
734
+ world_size = cp_group.world_size
735
+
736
+ if world_size == 1:
737
+ if return_lse:
738
+ return cp_attn_out, cp_attn_lse
739
+ return cp_attn_out
740
+
741
+ B, H, D = cp_attn_out.shape
742
+ if H % world_size != 0:
743
+ raise ValueError(f"H={H} must be divisible by DCP world size {world_size}.")
744
+ H_per_rank = H // world_size
745
+ lse_pack_dim = _dcp_a2a_lse_pack_dim(cp_attn_out.dtype)
746
+
747
+ send_buffer, recv_buffer = _dcp_a2a_send_recv_buffers(
748
+ (world_size, B, H_per_rank, D + lse_pack_dim),
749
+ device=cp_attn_out.device,
750
+ dtype=cp_attn_out.dtype,
751
+ )
752
+
753
+ _dcp_a2a_pack_send(
754
+ cp_attn_out,
755
+ cp_attn_lse,
756
+ send_buffer,
757
+ world_size,
758
+ H_per_rank,
759
+ D,
760
+ lse_pack_dim,
761
+ seq_lens=seq_lens,
762
+ query_start_loc=query_start_loc,
763
+ )
764
+
765
+ work = dist.all_to_all_single(
766
+ recv_buffer.view(-1),
767
+ send_buffer.view(-1),
768
+ group=cp_group.device_group,
769
+ async_op=True,
770
+ )
771
+ work.wait()
772
+
773
+ return _dcp_a2a_unpack_combine(
774
+ recv_buffer, D, lse_pack_dim, return_lse, is_lse_base_on_e
775
+ )
776
+
777
+
778
+def get_dcp_workspace_max_num_tokens(vllm_config: VllmConfig) -> int:
779
+ scheduler_config = vllm_config.scheduler_config
780
+ speculative_config = vllm_config.speculative_config
781
+ speculative_tokens = vllm_config.num_speculative_tokens
782
+ tokens_per_seq = (
783
+ 1
784
+ + (
785
+ 2
786
+ if speculative_config is not None and speculative_config.parallel_drafting
787
+ else 1
788
+ )
789
+ * speculative_tokens
790
+ )
791
+ return min(
792
+ scheduler_config.max_num_batched_tokens,
793
+ max(
794
+ scheduler_config.max_num_seqs * tokens_per_seq,
795
+ vllm_config.compilation_config.max_cudagraph_capture_size or 0,
796
+ ),
797
+ )
798
+
799
+
800
+def reserve_query_head_storage(
801
+ query: torch.Tensor, padded_num_heads: int
802
+) -> torch.Tensor:
803
+ """Reserve backing storage for fixed-head decode kernels."""
804
+ assert query.ndim == 3
805
+ assert query.shape[1] <= padded_num_heads
806
+ padded = query.new_empty((query.shape[0], padded_num_heads, query.shape[2]))
807
+ padded.resize_(query.shape)
808
+ padded.copy_(query)
809
+ return padded
810
+
811
+
812
+# Symmetric-memory A2A implementation
813
+
814
+
815
+_A2A_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16)
816
+
817
+
818
+class DirectDCPA2AWorkspace(DirectCPWorkspace):
819
+ """Persistent symmetric buffers for direct DCP output exchange."""
820
+
821
+ def __init__(
822
+ self,
823
+ group: ProcessGroup,
824
+ device: torch.device,
825
+ max_num_tokens: int,
826
+ heads_per_rank: int,
827
+ head_dim: int,
828
+ dtype: torch.dtype = torch.bfloat16,
829
+ num_ubatches: int = 1,
830
+ ) -> None:
831
+ if dtype not in _A2A_SUPPORTED_DTYPES:
832
+ raise ValueError(f"Direct DCP A2A does not support {dtype}")
833
+ if num_ubatches < 1:
834
+ raise ValueError(
835
+ f"Direct DCP A2A requires at least one ubatch slot, got {num_ubatches}"
836
+ )
837
+ super().__init__(group, device, num_ubatches)
838
+ self.max_num_tokens = max_num_tokens
839
+ self.heads_per_rank = heads_per_rank
840
+ self.head_dim = head_dim
841
+
842
+ output_shape = (
843
+ num_ubatches,
844
+ 2,
845
+ self.world_size,
846
+ max_num_tokens,
847
+ heads_per_rank,
848
+ head_dim,
849
+ )
850
+ lse_shape = (
851
+ num_ubatches,
852
+ 2,
853
+ self.world_size,
854
+ max_num_tokens,
855
+ heads_per_rank,
856
+ )
857
+ signal_shape = (num_ubatches, 2, self.world_size)
858
+ self.received_output, self.peer_output_ptrs = self._allocate(
859
+ output_shape, dtype
860
+ )
861
+ self.received_lse, self.peer_lse_ptrs = self._allocate(lse_shape, torch.float32)
862
+ self.received_signal, self.peer_signal_ptrs = self._allocate(
863
+ signal_shape, torch.int32
864
+ )
865
+
866
+ def lse_reduce(
867
+ self,
868
+ partial_output: torch.Tensor,
869
+ partial_lse: torch.Tensor,
870
+ is_lse_base_on_e: bool,
871
+ seq_lens: torch.Tensor | None = None,
872
+ query_start_loc: torch.Tensor | None = None,
873
+ ) -> torch.Tensor:
874
+ ubatch = dbo_current_ubatch_id()
875
+ num_tokens = partial_output.shape[0]
876
+ output = partial_output.new_empty(
877
+ (num_tokens, self.heads_per_rank, self.head_dim)
878
+ )
879
+ torch.ops._C.direct_dcp_a2a_lse_reduce(
880
+ partial_output,
881
+ partial_lse,
882
+ seq_lens,
883
+ query_start_loc,
884
+ self.peer_output_ptrs[ubatch],
885
+ self.peer_lse_ptrs[ubatch],
886
+ self.peer_signal_ptrs[ubatch],
887
+ self.received_output[ubatch],
888
+ self.received_lse[ubatch],
889
+ self.received_signal[ubatch],
890
+ self.epoch[ubatch : ubatch + 1],
891
+ output,
892
+ self.world_size,
893
+ self.rank,
894
+ self.max_num_tokens,
895
+ is_lse_base_on_e,
896
+ )
897
+ return output
898
+
899
+
900
+@functools.cache
901
+def get_direct_dcp_a2a_workspace(
902
+ group: GroupCoordinator,
903
+ device: torch.device,
904
+ max_num_tokens: int,
905
+ heads_per_rank: int,
906
+ head_dim: int,
907
+ dtype: torch.dtype,
908
+ num_ubatches: int,
909
+) -> DirectDCPA2AWorkspace | None:
910
+ if not direct_cp_enabled(
911
+ group, dtype, envs.VLLM_USE_DIRECT_DCP_A2A, _A2A_SUPPORTED_DTYPES
912
+ ):
913
+ return None
914
+ return DirectDCPA2AWorkspace(
915
+ group.device_group,
916
+ device,
917
+ max_num_tokens,
918
+ heads_per_rank,
919
+ head_dim,
920
+ dtype,
921
+ num_ubatches,
922
+ )
923
+
924
+
925
+# Q gather
926
+
927
+# Symmetric-memory implementation
928
+
929
+
930
+def _q_gather_layout_supported(
931
+ world_size: int,
932
+ heads_per_rank: int,
933
+ head_dim: int,
934
+ dtype: torch.dtype,
935
+ padded_num_heads: int | None,
936
+) -> bool:
937
+ element_size = torch.empty((), dtype=dtype).element_size()
938
+ gathered_num_heads = world_size * heads_per_rank
939
+ storage_num_heads = (
940
+ gathered_num_heads if padded_num_heads is None else padded_num_heads
941
+ )
942
+ return (
943
+ heads_per_rank * head_dim * element_size % 16 == 0
944
+ and storage_num_heads * head_dim * element_size % 16 == 0
945
+ )
946
+
947
+
948
+class DirectDCPQGatherWorkspace(DirectCPWorkspace):
949
+ """Publish query shards directly into the consumer-final symmetric buffer.
950
+
951
+ The final buffer is reusable after the downstream DCP output combine. That
952
+ combine orders all ranks after attention has consumed the gathered query.
953
+ """
954
+
955
+ def __init__(
956
+ self,
957
+ group: ProcessGroup,
958
+ device: torch.device,
959
+ max_num_tokens: int,
960
+ heads_per_rank: int,
961
+ head_dim: int,
962
+ dtype: torch.dtype = torch.bfloat16,
963
+ num_ubatches: int = 1,
964
+ padded_num_heads: int | None = None,
965
+ ) -> None:
966
+ if num_ubatches < 1:
967
+ raise ValueError(
968
+ "Direct DCP q-gather requires at least one ubatch slot, "
969
+ f"got {num_ubatches}"
970
+ )
971
+ if max_num_tokens < 1 or heads_per_rank < 1 or head_dim < 1:
972
+ raise ValueError(
973
+ "Direct DCP q-gather dimensions must be positive, got "
974
+ f"T={max_num_tokens}, H={heads_per_rank}, D={head_dim}"
975
+ )
976
+ gathered_num_heads = group.size() * heads_per_rank
977
+ if not _q_gather_layout_supported(
978
+ group.size(), heads_per_rank, head_dim, dtype, padded_num_heads
979
+ ):
980
+ raise ValueError("Direct DCP q-gather requires 16-byte-aligned query rows.")
981
+ super().__init__(group, device, num_ubatches)
982
+ if self.world_size <= 1:
983
+ raise ValueError("Direct DCP q-gather requires at least two ranks")
984
+ self.max_num_tokens = max_num_tokens
985
+ self.heads_per_rank = heads_per_rank
986
+ self.gathered_num_heads = gathered_num_heads
987
+ self.padded_num_heads = (
988
+ self.gathered_num_heads if padded_num_heads is None else padded_num_heads
989
+ )
990
+ if self.padded_num_heads < self.gathered_num_heads:
991
+ raise ValueError(
992
+ "Direct DCP q-gather padded heads must cover gathered heads: "
993
+ f"{self.padded_num_heads} < {self.gathered_num_heads}"
994
+ )
995
+ self.head_dim = head_dim
996
+
997
+ query_shape = (
998
+ num_ubatches,
999
+ max_num_tokens,
1000
+ self.padded_num_heads,
1001
+ head_dim,
1002
+ )
1003
+ signal_shape = (num_ubatches, 2, self.world_size)
1004
+ self.final_query, _ = self._allocate(query_shape, dtype)
1005
+ self.received_signal, _ = self._allocate(signal_shape, torch.int32)
1006
+ query_multicast_ptrs = self._multicast_ptrs(self.final_query)
1007
+ signal_multicast_ptrs = self._multicast_ptrs(self.received_signal)
1008
+ self.multicast_ptrs = list(
1009
+ zip(query_multicast_ptrs, signal_multicast_ptrs, strict=True)
1010
+ )
1011
+ if not all(
1012
+ query_ptr and signal_ptr for query_ptr, signal_ptr in self.multicast_ptrs
1013
+ ):
1014
+ raise RuntimeError(
1015
+ "Direct DCP q-gather requires NVLS symmetric-memory multicast."
1016
+ )
1017
+ self.completion = self.received_signal.new_zeros((num_ubatches, 1))
1018
+ torch.accelerator.synchronize()
1019
+
1020
+ def gather(self, local_query: torch.Tensor) -> torch.Tensor:
1021
+ ubatch = dbo_current_ubatch_id()
1022
+ if not 0 <= ubatch < self.num_ubatches:
1023
+ raise ValueError(
1024
+ f"DCP q-gather ubatch {ubatch} exceeds {self.num_ubatches} slots"
1025
+ )
1026
+ if local_query.ndim == 3 and local_query.shape[1] != self.heads_per_rank:
1027
+ raise ValueError(
1028
+ f"DCP q-gather expected {self.heads_per_rank} local query heads, "
1029
+ f"got {local_query.shape[1]}"
1030
+ )
1031
+
1032
+ num_tokens = local_query.shape[0]
1033
+ output = torch.as_strided(
1034
+ self.final_query[ubatch],
1035
+ size=(num_tokens, self.gathered_num_heads, self.head_dim),
1036
+ stride=(
1037
+ self.gathered_num_heads * self.head_dim,
1038
+ self.head_dim,
1039
+ 1,
1040
+ ),
1041
+ )
1042
+ query_multicast_ptr, signal_multicast_ptr = self.multicast_ptrs[ubatch]
1043
+ torch.ops._C.direct_dcp_q_gather(
1044
+ local_query,
1045
+ output,
1046
+ self.received_signal[ubatch],
1047
+ self.completion[ubatch],
1048
+ self.epoch[ubatch : ubatch + 1],
1049
+ self.world_size,
1050
+ self.rank,
1051
+ self.max_num_tokens,
1052
+ self.padded_num_heads,
1053
+ query_multicast_ptr,
1054
+ signal_multicast_ptr,
1055
+ )
1056
+ return output
1057
+
1058
+
1059
+@functools.cache
1060
+def get_direct_dcp_q_gather_workspace(
1061
+ group: GroupCoordinator,
1062
+ device: torch.device,
1063
+ max_num_tokens: int,
1064
+ heads_per_rank: int,
1065
+ head_dim: int,
1066
+ dtype: torch.dtype,
1067
+ num_ubatches: int,
1068
+ padded_num_heads: int | None = None,
1069
+) -> DirectDCPQGatherWorkspace | None:
1070
+ if not direct_cp_multicast_enabled(group, dtype, envs.VLLM_USE_DIRECT_DCP_Q_GATHER):
1071
+ return None
1072
+ if not _q_gather_layout_supported(
1073
+ group.world_size, heads_per_rank, head_dim, dtype, padded_num_heads
1074
+ ):
1075
+ return None
1076
+ return DirectDCPQGatherWorkspace(
1077
+ group.device_group,
1078
+ device,
1079
+ max_num_tokens,
1080
+ heads_per_rank,
1081
+ head_dim,
1082
+ dtype,
1083
+ num_ubatches,
1084
+ padded_num_heads,
1085
+ )
1086
+
1087
+
1088
+# KV gather
1089
+
1090
+# Symmetric-memory implementation
1091
+
1092
+
1093
+_KV_GATHER_SUPPORTED_DTYPES = (
1094
+ torch.float16,
1095
+ torch.bfloat16,
1096
+ torch.float8_e4m3fn,
1097
+)
1098
+
1099
+
1100
+def _kv_gather_layout_supported(token_dim: int, dtype: torch.dtype) -> bool:
1101
+ return token_dim * torch.empty((), dtype=dtype).element_size() % 16 == 0
1102
+
1103
+
1104
+class DirectDCPKVGatherWorkspace(DirectCPWorkspace):
1105
+ """Persistent symmetric buffers for direct DCP KV gather."""
1106
+
1107
+ def __init__(
1108
+ self,
1109
+ group: ProcessGroup,
1110
+ device: torch.device,
1111
+ max_gathered_tokens: int,
1112
+ token_dim: int,
1113
+ dtype: torch.dtype = torch.bfloat16,
1114
+ num_ubatches: int = 1,
1115
+ ) -> None:
1116
+ if dtype not in _KV_GATHER_SUPPORTED_DTYPES:
1117
+ raise ValueError(f"Direct DCP kv-gather does not support {dtype}")
1118
+ if num_ubatches < 1:
1119
+ raise ValueError(
1120
+ "Direct DCP kv-gather requires at least one ubatch slot, "
1121
+ f"got {num_ubatches}"
1122
+ )
1123
+ if max_gathered_tokens < 1 or token_dim < 1:
1124
+ raise ValueError(
1125
+ "Direct DCP kv-gather dimensions must be positive, got "
1126
+ f"T={max_gathered_tokens}, D={token_dim}"
1127
+ )
1128
+ if not _kv_gather_layout_supported(token_dim, dtype):
1129
+ raise ValueError("Direct DCP kv-gather requires 16-byte-aligned KV rows.")
1130
+ super().__init__(group, device, num_ubatches)
1131
+ if self.world_size <= 1:
1132
+ raise ValueError("Direct DCP kv-gather requires at least two ranks")
1133
+ if max_gathered_tokens % self.world_size != 0:
1134
+ raise ValueError(
1135
+ "Direct DCP kv-gather capacity must divide evenly across "
1136
+ f"ranks: {max_gathered_tokens} % {self.world_size} != 0"
1137
+ )
1138
+ self.max_gathered_tokens = max_gathered_tokens
1139
+
1140
+ kv_shape = (num_ubatches, 2, max_gathered_tokens, token_dim)
1141
+ signal_shape = (num_ubatches, 2, self.world_size)
1142
+ self.received_kv, _ = self._allocate(kv_shape, dtype)
1143
+ self.received_signal, _ = self._allocate(signal_shape, torch.int32)
1144
+ kv_multicast_ptrs = self._multicast_ptrs(self.received_kv)
1145
+ signal_multicast_ptrs = self._multicast_ptrs(self.received_signal)
1146
+ self.multicast_ptrs = list(
1147
+ zip(kv_multicast_ptrs, signal_multicast_ptrs, strict=True)
1148
+ )
1149
+ if not all(kv_ptr and signal_ptr for kv_ptr, signal_ptr in self.multicast_ptrs):
1150
+ raise RuntimeError(
1151
+ "Direct DCP kv-gather requires NVLS symmetric-memory multicast."
1152
+ )
1153
+ self.completion = self.received_signal.new_zeros((num_ubatches, 2))
1154
+ torch.accelerator.synchronize()
1155
+
1156
+ def gather(self, gathered_kv: torch.Tensor, local_kv: torch.Tensor) -> None:
1157
+ ubatch = dbo_current_ubatch_id()
1158
+ if not 0 <= ubatch < self.num_ubatches:
1159
+ raise ValueError(
1160
+ f"DCP kv-gather ubatch {ubatch} exceeds {self.num_ubatches} slots"
1161
+ )
1162
+ kv_multicast_ptr, signal_multicast_ptr = self.multicast_ptrs[ubatch]
1163
+ torch.ops._C.direct_dcp_kv_gather(
1164
+ local_kv,
1165
+ self.received_kv[ubatch],
1166
+ self.received_signal[ubatch],
1167
+ self.completion[ubatch],
1168
+ self.epoch[ubatch : ubatch + 1],
1169
+ gathered_kv,
1170
+ self.world_size,
1171
+ self.rank,
1172
+ self.max_gathered_tokens,
1173
+ kv_multicast_ptr,
1174
+ signal_multicast_ptr,
1175
+ )
1176
+
1177
+
1178
+@functools.cache
1179
+def get_direct_dcp_kv_gather_workspace(
1180
+ group: GroupCoordinator,
1181
+ device: torch.device,
1182
+ max_gathered_tokens: int,
1183
+ token_dim: int,
1184
+ dtype: torch.dtype,
1185
+ num_ubatches: int,
1186
+) -> DirectDCPKVGatherWorkspace | None:
1187
+ if not direct_cp_multicast_enabled(
1188
+ group,
1189
+ dtype,
1190
+ envs.VLLM_USE_DIRECT_DCP_KV_GATHER,
1191
+ _KV_GATHER_SUPPORTED_DTYPES,
1192
+ ):
1193
+ return None
1194
+ if not _kv_gather_layout_supported(token_dim, dtype):
1195
+ return None
1196
+ return DirectDCPKVGatherWorkspace(
1197
+ group.device_group,
1198
+ device,
1199
+ max_gathered_tokens,
1200
+ token_dim,
1201
+ dtype,
1202
+ num_ubatches,
1203
+ )
1204
+
1205
+
1206
+# MLA DCP backend selection
1207
+
1208
+
1209
+class DCPCombine(Protocol):
1210
+ def __call__(
1211
+ self,
1212
+ partial_output: torch.Tensor,
1213
+ partial_lse: torch.Tensor,
1214
+ *,
1215
+ seq_lens: torch.Tensor,
1216
+ query_start_loc: torch.Tensor,
1217
+ ) -> torch.Tensor: ...
1218
+
1219
+
1220
+class MLADCPManager:
1221
+ """Select and own layer-level collective implementations for MLA DCP."""
1222
+
1223
+ _kv_gather: Callable[[torch.Tensor, torch.Tensor], object]
1224
+
1225
+ def __init__(
1226
+ self,
1227
+ vllm_config: VllmConfig,
1228
+ device: torch.device,
1229
+ num_heads: int,
1230
+ query_head_dim: int,
1231
+ output_head_dim: int,
1232
+ query_dtype: torch.dtype,
1233
+ output_dtype: torch.dtype,
1234
+ padded_num_heads: int | None,
1235
+ is_lse_base_on_e: bool,
1236
+ use_pcp: bool,
1237
+ ) -> None:
1238
+ parallel_config = vllm_config.parallel_config
1239
+ self.group = get_dcp_group()
1240
+ self.device = torch.device(device)
1241
+ self.num_ubatches = max(parallel_config.num_ubatches, 1)
1242
+ self.max_num_tokens = get_dcp_workspace_max_num_tokens(vllm_config)
1243
+ self.use_a2a = parallel_config.dcp_comm_backend == "a2a"
1244
+ self.padded_num_heads = padded_num_heads
1245
+
1246
+ self.combine = self._init_combine(
1247
+ num_heads,
1248
+ output_head_dim,
1249
+ output_dtype,
1250
+ is_lse_base_on_e,
1251
+ use_pcp,
1252
+ )
1253
+ self.query_gather = (
1254
+ None
1255
+ if use_pcp
1256
+ else self._init_query_gather(
1257
+ num_heads,
1258
+ query_head_dim,
1259
+ query_dtype,
1260
+ )
1261
+ )
1262
+
1263
+ def _init_combine(
1264
+ self,
1265
+ num_heads: int,
1266
+ head_dim: int,
1267
+ dtype: torch.dtype,
1268
+ is_lse_base_on_e: bool,
1269
+ use_pcp: bool,
1270
+ ) -> DCPCombine:
1271
+ direct_workspace = None
1272
+ if self.use_a2a:
1273
+ direct_workspace = get_direct_dcp_a2a_workspace(
1274
+ self.group,
1275
+ self.device,
1276
+ self.max_num_tokens,
1277
+ num_heads,
1278
+ head_dim,
1279
+ dtype,
1280
+ self.num_ubatches,
1281
+ )
1282
+ if direct_workspace is not None:
1283
+ logger.info_once("Using direct symmetric-memory DCP A2A for MLA.")
1284
+ return functools.partial(
1285
+ direct_workspace.lse_reduce,
1286
+ is_lse_base_on_e=is_lse_base_on_e,
1287
+ )
1288
+
1289
+ combine_fn = (
1290
+ dcp_a2a_lse_reduce
1291
+ if self.use_a2a
1292
+ else cp_lse_ag_out_ar
1293
+ if use_pcp
1294
+ else cp_lse_ag_out_rs
1295
+ )
1296
+ return functools.partial(
1297
+ combine_fn,
1298
+ cp_group=self.group,
1299
+ is_lse_base_on_e=is_lse_base_on_e,
1300
+ )
1301
+
1302
+ def _init_query_gather(
1303
+ self,
1304
+ num_heads: int,
1305
+ head_dim: int,
1306
+ dtype: torch.dtype,
1307
+ ) -> Callable[[torch.Tensor], torch.Tensor]:
1308
+ direct_workspace = get_direct_dcp_q_gather_workspace(
1309
+ self.group,
1310
+ self.device,
1311
+ self.max_num_tokens,
1312
+ num_heads,
1313
+ head_dim,
1314
+ dtype,
1315
+ self.num_ubatches,
1316
+ self.padded_num_heads,
1317
+ )
1318
+ if direct_workspace is not None:
1319
+ logger.info_once("Using direct symmetric-memory DCP query gather for MLA.")
1320
+ return direct_workspace.gather
1321
+ return self._gather_query
1322
+
1323
+ def _gather_query(self, query: torch.Tensor) -> torch.Tensor:
1324
+ query = self.group.all_gather(query, dim=1)
1325
+ if self.padded_num_heads is not None:
1326
+ query = reserve_query_head_storage(query, self.padded_num_heads)
1327
+ return query
1328
+
1329
+ def init_kv_gather(
1330
+ self,
1331
+ workspace: torch.Tensor,
1332
+ max_gathered_tokens: int,
1333
+ ) -> None:
1334
+ world_size = self.group.world_size
1335
+ assert max_gathered_tokens > 0
1336
+ assert max_gathered_tokens % world_size == 0
1337
+ assert workspace.ndim == 2
1338
+ assert workspace.is_contiguous()
1339
+ assert workspace.shape[0] == (
1340
+ max_gathered_tokens + max_gathered_tokens // world_size
1341
+ )
1342
+ assert workspace.shape[1] > 0
1343
+
1344
+ direct_workspace = get_direct_dcp_kv_gather_workspace(
1345
+ self.group,
1346
+ workspace.device,
1347
+ max_gathered_tokens,
1348
+ workspace.shape[1],
1349
+ workspace.dtype,
1350
+ self.num_ubatches,
1351
+ )
1352
+ if direct_workspace is not None:
1353
+ logger.info_once(
1354
+ "Using direct symmetric-memory DCP chunked-context KV gather for MLA."
1355
+ )
1356
+ self._kv_gather = direct_workspace.gather
1357
+ else:
1358
+ self._kv_gather = functools.partial(
1359
+ torch.distributed.all_gather_into_tensor,
1360
+ group=self.group.device_group,
1361
+ )
1362
+
1363
+ def kv_gather(
1364
+ self,
1365
+ gathered_kv: torch.Tensor,
1366
+ local_kv: torch.Tensor,
1367
+ ) -> object:
1368
+ return self._kv_gather(gathered_kv, local_kv)
vllm/v1/attention/ops/dcp_alltoall.py
deleted
-470
@@ -1,470 +0,0 @@
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 by exchanging
13
-the partial output and LSE in a single packed All-to-All payload.
14
-
15
-Usage:
16
- vllm serve model --tp 16 --dcp 16 --dcp-comm-backend a2a
17
-
18
-Reference: https://arxiv.org/abs/2507.07120
19
-"""
20
-
21
-from __future__ import annotations
22
-
23
-from typing import TYPE_CHECKING
24
-
25
-import torch
26
-import torch.distributed as dist
27
-
28
-from vllm.triton_utils import tl, triton
29
-from vllm.v1.attention.ops.common import mask_dcp_empty_shards_
30
-
31
-if TYPE_CHECKING:
32
- from vllm.distributed.parallel_state import GroupCoordinator
33
- from vllm.v1.attention.ops.common import CPTritonContext
34
-
35
-
36
-def _lse_weighted_combine(
37
- outputs: torch.Tensor,
38
- lses: torch.Tensor,
39
- return_lse: bool = False,
40
- is_lse_base_on_e: bool = True,
41
-) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
42
- """
43
- CPU reference implementation for LSE-weighted combination.
44
-
45
- This is a pure PyTorch implementation used for testing and validation.
46
-
47
- Args:
48
- outputs: Partial attention outputs [N, B, H, D]
49
- N = number of KV shards (ranks)
50
- B = batch size (num_tokens)
51
- H = number of heads per rank
52
- D = head dimension
53
- lses: Log-sum-exp values [N, B, H]
54
- return_lse: If True, also return the global LSE
55
- is_lse_base_on_e: If True, LSE is base e; if False, base 2
56
-
57
- Returns:
58
- Combined output [B, H, D], and optionally global LSE [B, H]
59
- """
60
- N, B, H, D = outputs.shape
61
-
62
- # Handle NaN and inf in LSEs
63
- lses = torch.where(
64
- torch.isnan(lses) | torch.isinf(lses),
65
- torch.tensor(float("-inf"), device=lses.device, dtype=lses.dtype),
66
- lses,
67
- )
68
-
69
- # Compute max LSE for numerical stability
70
- lse_max, _ = lses.max(dim=0) # [B, H]
71
- lse_max = torch.where(
72
- lse_max == float("-inf"),
73
- torch.zeros_like(lse_max),
74
- lse_max,
75
- )
76
-
77
- # Compute weights: softmax over the N dimension
78
- if is_lse_base_on_e:
79
- weights = torch.exp(lses - lse_max.unsqueeze(0)) # [N, B, H]
80
- else:
81
- weights = torch.pow(2.0, lses - lse_max.unsqueeze(0)) # [N, B, H]
82
-
83
- # Handle NaN weights
84
- weights = torch.where(torch.isnan(weights), torch.zeros_like(weights), weights)
85
-
86
- # Normalize weights
87
- weight_sum = weights.sum(dim=0, keepdim=True) # [1, B, H]
88
- weights = weights / weight_sum.clamp(min=1e-10) # [N, B, H]
89
-
90
- # Weighted combination: sum over N dimension
91
- weights = weights.unsqueeze(-1)
92
- outputs = torch.where(weights == 0, torch.zeros_like(outputs), outputs)
93
- result = (outputs * weights).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
-def _dcp_a2a_lse_pack_dim(output_dtype: torch.dtype) -> int:
106
- bits = torch.finfo(output_dtype).bits
107
- if bits == 16:
108
- return 2
109
- if bits == 32:
110
- return 1
111
- raise ValueError(f"Cannot pack fp32 LSE into output dtype {output_dtype}.")
112
-
113
-
114
-def _dcp_a2a_send_recv_buffers(
115
- shape: tuple[int, ...],
116
- device: torch.device,
117
- dtype: torch.dtype,
118
-) -> tuple[torch.Tensor, torch.Tensor]:
119
- # Don't use the shared WorkspaceManager here. A FULL cudagraph bakes in the
120
- # buffer address at capture, but the workspace is growable and sized only to
121
- # the largest *captured* batch (the cudagraph capture cap). Any eager a2a
122
- # with a bigger batch regrows it, freeing that address and poisoning every
123
- # captured graph -> illegal memory access on replay. This bites the very
124
- # first request: the post-capture warmup runs an eager decode at
125
- # max_num_seqs (> the cap), so the graphs are already dangling before the
126
- # server is ready. torch.empty buffers instead live in the graph's private
127
- # pool and stay valid for its lifetime (as _dcp_a2a_unpack_combine and the
128
- # AG+RS combine path already rely on).
129
- return (
130
- torch.empty(shape, device=device, dtype=dtype),
131
- torch.empty(shape, device=device, dtype=dtype),
132
- )
133
-
134
-
135
-@triton.jit
136
-def _dcp_a2a_pack_send_kernel(
137
- out_ptr,
138
- lse_ptr,
139
- send_ptr,
140
- out_stride_B,
141
- out_stride_H,
142
- out_stride_D,
143
- lse_stride_B,
144
- lse_stride_H,
145
- send_stride_N,
146
- send_stride_B,
147
- send_stride_H,
148
- send_stride_D,
149
- N: tl.constexpr,
150
- HEAD_DIM: tl.constexpr,
151
- H_PER_RANK: tl.constexpr,
152
- LSE_PACK_DIM: tl.constexpr,
153
-):
154
- batch_idx = tl.program_id(0).to(tl.int64)
155
- local_head_idx = tl.program_id(1).to(tl.int64)
156
- d_offsets = tl.arange(0, HEAD_DIM)
157
-
158
- for rank_idx in tl.static_range(N):
159
- src_head_idx = rank_idx * H_PER_RANK + local_head_idx
160
- send_base = (
161
- rank_idx * send_stride_N
162
- + batch_idx * send_stride_B
163
- + local_head_idx * send_stride_H
164
- )
165
-
166
- out_offsets = (
167
- batch_idx * out_stride_B
168
- + src_head_idx * out_stride_H
169
- + d_offsets * out_stride_D
170
- )
171
- tl.store(
172
- send_ptr + send_base + d_offsets * send_stride_D,
173
- tl.load(out_ptr + out_offsets),
174
- )
175
-
176
- lse_val = tl.load(
177
- lse_ptr + batch_idx * lse_stride_B + src_head_idx * lse_stride_H
178
- ).to(tl.float32)
179
- if LSE_PACK_DIM == 1:
180
- tl.store(
181
- send_ptr + send_base + HEAD_DIM * send_stride_D,
182
- lse_val.to(send_ptr.dtype.element_ty),
183
- )
184
- else:
185
- lse_bits = lse_val.to(tl.uint32, bitcast=True)
186
- lo = (lse_bits & 0xFFFF).to(tl.uint16)
187
- hi = ((lse_bits >> 16) & 0xFFFF).to(tl.uint16)
188
- tl.store(
189
- send_ptr + send_base + HEAD_DIM * send_stride_D,
190
- lo.to(send_ptr.dtype.element_ty, bitcast=True),
191
- )
192
- tl.store(
193
- send_ptr + send_base + (HEAD_DIM + 1) * send_stride_D,
194
- hi.to(send_ptr.dtype.element_ty, bitcast=True),
195
- )
196
-
197
-
198
-@triton.jit
199
-def _dcp_a2a_unpack_combine_kernel(
200
- recv_ptr,
201
- out_ptr,
202
- out_lse_ptr,
203
- recv_stride_N,
204
- recv_stride_B,
205
- recv_stride_H,
206
- recv_stride_D,
207
- out_stride_B,
208
- out_stride_H,
209
- out_stride_D,
210
- out_lse_stride_B,
211
- out_lse_stride_H,
212
- N: tl.constexpr,
213
- HEAD_DIM: tl.constexpr,
214
- IS_BASE_E: tl.constexpr,
215
- RETURN_LSE: tl.constexpr,
216
- LSE_PACK_DIM: tl.constexpr,
217
-):
218
- batch_idx = tl.program_id(0).to(tl.int64)
219
- head_idx = tl.program_id(1).to(tl.int64)
220
- d_offsets = tl.arange(0, HEAD_DIM)
221
-
222
- lse_max = -float("inf")
223
- for rank_idx in tl.static_range(N):
224
- recv_base = (
225
- rank_idx * recv_stride_N
226
- + batch_idx * recv_stride_B
227
- + head_idx * recv_stride_H
228
- )
229
- if LSE_PACK_DIM == 1:
230
- lse_val = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D).to(
231
- tl.float32
232
- )
233
- else:
234
- lo_raw = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D)
235
- hi_raw = tl.load(recv_ptr + recv_base + (HEAD_DIM + 1) * recv_stride_D)
236
- lo = lo_raw.to(tl.uint16, bitcast=True).to(tl.uint32)
237
- hi = hi_raw.to(tl.uint16, bitcast=True).to(tl.uint32)
238
- lse_val = (lo | (hi << 16)).to(tl.float32, bitcast=True)
239
- lse_val = tl.where(
240
- (lse_val != lse_val) | (lse_val == float("inf")),
241
- -float("inf"),
242
- lse_val,
243
- )
244
- lse_max = tl.maximum(lse_max, lse_val)
245
-
246
- lse_max = tl.where(lse_max == -float("inf"), 0.0, lse_max)
247
-
248
- lse_sum = 0.0
249
- for rank_idx in tl.static_range(N):
250
- recv_base = (
251
- rank_idx * recv_stride_N
252
- + batch_idx * recv_stride_B
253
- + head_idx * recv_stride_H
254
- )
255
- if LSE_PACK_DIM == 1:
256
- lse_val = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D).to(
257
- tl.float32
258
- )
259
- else:
260
- lo_raw = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D)
261
- hi_raw = tl.load(recv_ptr + recv_base + (HEAD_DIM + 1) * recv_stride_D)
262
- lo = lo_raw.to(tl.uint16, bitcast=True).to(tl.uint32)
263
- hi = hi_raw.to(tl.uint16, bitcast=True).to(tl.uint32)
264
- lse_val = (lo | (hi << 16)).to(tl.float32, bitcast=True)
265
- lse_val = tl.where(
266
- (lse_val != lse_val) | (lse_val == float("inf")),
267
- -float("inf"),
268
- lse_val,
269
- )
270
- if IS_BASE_E:
271
- lse_sum += tl.exp(lse_val - lse_max)
272
- else:
273
- lse_sum += tl.exp2(lse_val - lse_max)
274
-
275
- if IS_BASE_E: # noqa: SIM108
276
- global_lse = tl.log(lse_sum) + lse_max
277
- else:
278
- global_lse = tl.log2(lse_sum) + lse_max
279
-
280
- acc = tl.zeros([HEAD_DIM], dtype=tl.float32)
281
- for rank_idx in tl.static_range(N):
282
- recv_base = (
283
- rank_idx * recv_stride_N
284
- + batch_idx * recv_stride_B
285
- + head_idx * recv_stride_H
286
- )
287
- if LSE_PACK_DIM == 1:
288
- lse_val = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D).to(
289
- tl.float32
290
- )
291
- else:
292
- lo_raw = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D)
293
- hi_raw = tl.load(recv_ptr + recv_base + (HEAD_DIM + 1) * recv_stride_D)
294
- lo = lo_raw.to(tl.uint16, bitcast=True).to(tl.uint32)
295
- hi = hi_raw.to(tl.uint16, bitcast=True).to(tl.uint32)
296
- lse_val = (lo | (hi << 16)).to(tl.float32, bitcast=True)
297
- lse_val = tl.where(
298
- (lse_val != lse_val) | (lse_val == float("inf")),
299
- -float("inf"),
300
- lse_val,
301
- )
302
- if IS_BASE_E:
303
- weight = tl.exp(lse_val - global_lse)
304
- else:
305
- weight = tl.exp2(lse_val - global_lse)
306
- weight = tl.where(weight != weight, 0.0, weight)
307
- partial = tl.load(recv_ptr + recv_base + d_offsets * recv_stride_D).to(
308
- tl.float32
309
- )
310
- partial = tl.where(weight == 0.0, 0.0, partial)
311
- acc += partial * weight
312
-
313
- final_offsets = (
314
- batch_idx * out_stride_B + head_idx * out_stride_H + d_offsets * out_stride_D
315
- )
316
- tl.store(out_ptr + final_offsets, acc)
317
-
318
- if RETURN_LSE:
319
- out_lse_offset = batch_idx * out_lse_stride_B + head_idx * out_lse_stride_H
320
- tl.store(out_lse_ptr + out_lse_offset, global_lse)
321
-
322
-
323
-def _dcp_a2a_pack_send(
324
- cp_attn_out: torch.Tensor,
325
- cp_attn_lse: torch.Tensor,
326
- send_buffer: torch.Tensor,
327
- world_size: int,
328
- h_per_rank: int,
329
- head_dim: int,
330
- lse_pack_dim: int,
331
- seq_lens: torch.Tensor | None = None,
332
- query_start_loc: torch.Tensor | None = None,
333
-) -> None:
334
- mask_dcp_empty_shards_(cp_attn_lse, seq_lens, query_start_loc)
335
- grid = (cp_attn_out.shape[0], h_per_rank, 1)
336
- _dcp_a2a_pack_send_kernel[grid](
337
- cp_attn_out,
338
- cp_attn_lse,
339
- send_buffer,
340
- cp_attn_out.stride(0),
341
- cp_attn_out.stride(1),
342
- cp_attn_out.stride(2),
343
- cp_attn_lse.stride(0),
344
- cp_attn_lse.stride(1),
345
- send_buffer.stride(0),
346
- send_buffer.stride(1),
347
- send_buffer.stride(2),
348
- send_buffer.stride(3),
349
- N=world_size,
350
- HEAD_DIM=head_dim,
351
- H_PER_RANK=h_per_rank,
352
- LSE_PACK_DIM=lse_pack_dim,
353
- )
354
-
355
-
356
-def _dcp_a2a_unpack_combine(
357
- recv_buffer: torch.Tensor,
358
- head_dim: int,
359
- lse_pack_dim: int,
360
- return_lse: bool,
361
- is_lse_base_on_e: bool,
362
-) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
363
- world_size, num_tokens, h_per_rank, _ = recv_buffer.shape
364
- out = torch.empty(
365
- (num_tokens, h_per_rank, head_dim),
366
- device=recv_buffer.device,
367
- dtype=recv_buffer.dtype,
368
- )
369
- out_lse = torch.empty(
370
- (num_tokens, h_per_rank) if return_lse else (1, 1),
371
- device=recv_buffer.device,
372
- dtype=torch.float32 if return_lse else recv_buffer.dtype,
373
- )
374
- grid = (num_tokens, h_per_rank, 1)
375
- _dcp_a2a_unpack_combine_kernel[grid](
376
- recv_buffer,
377
- out,
378
- out_lse,
379
- recv_buffer.stride(0),
380
- recv_buffer.stride(1),
381
- recv_buffer.stride(2),
382
- recv_buffer.stride(3),
383
- out.stride(0),
384
- out.stride(1),
385
- out.stride(2),
386
- out_lse.stride(0),
387
- out_lse.stride(1),
388
- N=world_size,
389
- HEAD_DIM=head_dim,
390
- IS_BASE_E=is_lse_base_on_e,
391
- RETURN_LSE=return_lse,
392
- LSE_PACK_DIM=lse_pack_dim,
393
- )
394
- if return_lse:
395
- return out, out_lse
396
- return out
397
-
398
-
399
-def dcp_a2a_lse_reduce(
400
- cp_attn_out: torch.Tensor,
401
- cp_attn_lse: torch.Tensor,
402
- cp_group: GroupCoordinator,
403
- ctx: CPTritonContext | None = None,
404
- return_lse: bool = False,
405
- is_lse_base_on_e: bool = True,
406
- seq_lens: torch.Tensor | None = None,
407
- query_start_loc: torch.Tensor | None = None,
408
-) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
409
- """
410
- Combine partial attention outputs across DCP ranks using All-to-All.
411
-
412
- The output and LSE are packed into a single output-dtype buffer, sent
413
- with one All-to-All, then unpacked and combined with exact LSE weighting.
414
-
415
- Args:
416
- cp_attn_out: [B, H, D] where B=num_tokens, H=total_heads, D=head_dim
417
- cp_attn_lse: [B, H] floating-point log-sum-exp values
418
- cp_group: GroupCoordinator for DCP communication
419
- ctx: CPTritonContext (unused, for signature compatibility)
420
- return_lse: If True, also return the combined global LSE
421
- is_lse_base_on_e: If True, LSE is base e; if False, base 2
422
- seq_lens: Local KV lengths. Empty shards contribute zero weight.
423
- query_start_loc: Cumulative query-token offsets for each request.
424
-
425
- Returns:
426
- Combined output [B, H/N, D] (head-scattered)
427
- If return_lse=True, also returns global_lse [B, H/N]
428
- """
429
- world_size = cp_group.world_size
430
-
431
- if world_size == 1:
432
- if return_lse:
433
- return cp_attn_out, cp_attn_lse
434
- return cp_attn_out
435
-
436
- B, H, D = cp_attn_out.shape
437
- if H % world_size != 0:
438
- raise ValueError(f"H={H} must be divisible by DCP world size {world_size}.")
439
- H_per_rank = H // world_size
440
- lse_pack_dim = _dcp_a2a_lse_pack_dim(cp_attn_out.dtype)
441
-
442
- send_buffer, recv_buffer = _dcp_a2a_send_recv_buffers(
443
- (world_size, B, H_per_rank, D + lse_pack_dim),
444
- device=cp_attn_out.device,
445
- dtype=cp_attn_out.dtype,
446
- )
447
-
448
- _dcp_a2a_pack_send(
449
- cp_attn_out,
450
- cp_attn_lse,
451
- send_buffer,
452
- world_size,
453
- H_per_rank,
454
- D,
455
- lse_pack_dim,
456
- seq_lens=seq_lens,
457
- query_start_loc=query_start_loc,
458
- )
459
-
460
- work = dist.all_to_all_single(
461
- recv_buffer.view(-1),
462
- send_buffer.view(-1),
463
- group=cp_group.device_group,
464
- async_op=True,
465
- )
466
- work.wait()
467
-
468
- return _dcp_a2a_unpack_combine(
469
- recv_buffer, D, lse_pack_dim, return_lse, is_lse_base_on_e
470
- )
vllm/v1/attention/ops/dcp_utils.py
deleted
-740
@@ -1,740 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""MLA DCP collective selection and direct symmetric-memory implementations."""
4
-
5
-from __future__ import annotations
6
-
7
-import functools
8
-from collections.abc import Callable
9
-from typing import TYPE_CHECKING, Any, Protocol
10
-
11
-import torch
12
-
13
-import vllm.envs as envs
14
-from vllm.config import VllmConfig
15
-from vllm.distributed import get_dcp_group
16
-from vllm.distributed.parallel_state import in_the_same_node_as
17
-from vllm.logger import init_logger
18
-from vllm.platforms import current_platform
19
-from vllm.v1.attention.ops.common import cp_lse_ag_out_ar, cp_lse_ag_out_rs
20
-from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
21
-from vllm.v1.worker.ubatching import dbo_current_ubatch_id
22
-
23
-logger = init_logger(__name__)
24
-
25
-if TYPE_CHECKING:
26
- from torch.distributed import ProcessGroup
27
-
28
- from vllm.distributed.parallel_state import GroupCoordinator
29
-
30
-try:
31
- import torch.distributed._symmetric_memory as symm_mem
32
-
33
- symm_mem_available = True
34
-except ImportError:
35
- symm_mem = None # type: ignore[assignment]
36
- symm_mem_available = False
37
-
38
-
39
-@functools.cache
40
-def _symm_mem_spans_group(group: GroupCoordinator) -> bool:
41
- """Probe whether the group has NVLS symmetric memory."""
42
- if not symm_mem_available:
43
- return False
44
- try:
45
- from torch._C._autograd import DeviceType
46
- from torch._C._distributed_c10d import _SymmetricMemory
47
-
48
- device = torch.device("cuda", torch.accelerator.current_device_index())
49
- if not _SymmetricMemory.has_multicast_support(DeviceType.CUDA, device.index):
50
- return False
51
- probe = symm_mem.empty(8, dtype=torch.uint8, device=device)
52
- probe.zero_()
53
- torch.accelerator.synchronize()
54
- handle = symm_mem.rendezvous(probe, group.device_group.group_name)
55
- spans = handle is not None and handle.multicast_ptr != 0
56
- except Exception as error:
57
- logger.debug("Direct DCP symmetric-memory probe failed: %s", error)
58
- return False
59
- logger.debug_once(
60
- "Direct DCP symmetric memory across %d ranks: %s",
61
- group.world_size,
62
- "available" if spans else "unavailable",
63
- )
64
- return spans
65
-
66
-
67
-def _direct_dcp_enabled(
68
- group: GroupCoordinator,
69
- dtype: torch.dtype,
70
- use_direct: bool | None,
71
- supported_dtypes: tuple[torch.dtype, ...] | None = None,
72
-) -> bool:
73
- if use_direct is not None:
74
- return use_direct
75
- return (
76
- symm_mem_available
77
- and current_platform.is_cuda()
78
- and (supported_dtypes is None or dtype in supported_dtypes)
79
- and (
80
- all(in_the_same_node_as(group.cpu_group, source_rank=0))
81
- or _symm_mem_spans_group(group)
82
- )
83
- )
84
-
85
-
86
-def _direct_dcp_multicast_enabled(
87
- group: GroupCoordinator,
88
- dtype: torch.dtype,
89
- use_direct: bool | None,
90
- supported_dtypes: tuple[torch.dtype, ...] | None = None,
91
-) -> bool:
92
- return _direct_dcp_enabled(
93
- group, dtype, use_direct, supported_dtypes
94
- ) and _symm_mem_spans_group(group)
95
-
96
-
97
-def get_dcp_workspace_max_num_tokens(vllm_config: VllmConfig) -> int:
98
- scheduler_config = vllm_config.scheduler_config
99
- speculative_config = vllm_config.speculative_config
100
- speculative_tokens = vllm_config.num_speculative_tokens
101
- tokens_per_seq = (
102
- 1
103
- + (
104
- 2
105
- if speculative_config is not None and speculative_config.parallel_drafting
106
- else 1
107
- )
108
- * speculative_tokens
109
- )
110
- return min(
111
- scheduler_config.max_num_batched_tokens,
112
- max(
113
- scheduler_config.max_num_seqs * tokens_per_seq,
114
- vllm_config.compilation_config.max_cudagraph_capture_size or 0,
115
- ),
116
- )
117
-
118
-
119
-class _DirectDCPWorkspace:
120
- def __init__(
121
- self,
122
- group: ProcessGroup,
123
- device: torch.device,
124
- num_ubatches: int,
125
- ) -> None:
126
- self.group = group
127
- self.world_size = group.size()
128
- self.rank = group.rank()
129
- self.device = torch.device(device)
130
- self.num_ubatches = num_ubatches
131
- self.epoch = torch.zeros(num_ubatches, dtype=torch.int64, device=self.device)
132
- self._allocations: list[tuple[torch.Tensor, Any, list[torch.Tensor]]] = []
133
-
134
- def _allocate(
135
- self, shape: tuple[int, ...], dtype: torch.dtype
136
- ) -> tuple[torch.Tensor, torch.Tensor]:
137
- storage = symm_mem.empty(shape, device=self.device, dtype=dtype)
138
- storage.zero_()
139
- torch.accelerator.synchronize()
140
- handle = symm_mem.rendezvous(storage, self.group.group_name)
141
- assert handle is not None, "DCP symmetric memory rendezvous returned None"
142
- handle.barrier()
143
- views = [
144
- handle.get_buffer(peer, list(shape), dtype, 0)
145
- for peer in range(self.world_size)
146
- ]
147
- self.device = storage.device
148
- peer_ptrs = torch.tensor(
149
- [
150
- [view[ubatch].data_ptr() for view in views]
151
- for ubatch in range(self.num_ubatches)
152
- ],
153
- dtype=torch.int64,
154
- device=self.device,
155
- )
156
- self._allocations.append((storage, handle, views))
157
- return storage, peer_ptrs
158
-
159
- def _multicast_ptrs(self, storage: torch.Tensor) -> list[int]:
160
- disabled = [0] * self.num_ubatches
161
- for allocated, handle, _ in self._allocations:
162
- if allocated is storage:
163
- break
164
- else:
165
- return disabled
166
- try:
167
- from torch._C._autograd import DeviceType
168
- from torch._C._distributed_c10d import _SymmetricMemory
169
-
170
- if not _SymmetricMemory.has_multicast_support(
171
- DeviceType.CUDA, storage.device.index
172
- ):
173
- return disabled
174
- multicast_base = handle.multicast_ptr
175
- except Exception:
176
- return disabled
177
- if not multicast_base:
178
- return disabled
179
- storage_base = storage.data_ptr()
180
- return [
181
- multicast_base + (storage[ubatch].data_ptr() - storage_base)
182
- for ubatch in range(self.num_ubatches)
183
- ]
184
-
185
-
186
-def reserve_query_head_storage(
187
- query: torch.Tensor, padded_num_heads: int
188
-) -> torch.Tensor:
189
- """Reserve backing storage for fixed-head decode kernels."""
190
- assert query.ndim == 3
191
- assert query.shape[1] <= padded_num_heads
192
- padded = query.new_empty((query.shape[0], padded_num_heads, query.shape[2]))
193
- padded.resize_(query.shape)
194
- padded.copy_(query)
195
- return padded
196
-
197
-
198
-_A2A_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16)
199
-
200
-
201
-class DirectDCPA2AWorkspace(_DirectDCPWorkspace):
202
- """Persistent symmetric buffers for direct DCP output exchange."""
203
-
204
- def __init__(
205
- self,
206
- group: ProcessGroup,
207
- device: torch.device,
208
- max_num_tokens: int,
209
- heads_per_rank: int,
210
- head_dim: int,
211
- dtype: torch.dtype = torch.bfloat16,
212
- num_ubatches: int = 1,
213
- ) -> None:
214
- if dtype not in _A2A_SUPPORTED_DTYPES:
215
- raise ValueError(f"Direct DCP A2A does not support {dtype}")
216
- if num_ubatches < 1:
217
- raise ValueError(
218
- f"Direct DCP A2A requires at least one ubatch slot, got {num_ubatches}"
219
- )
220
- super().__init__(group, device, num_ubatches)
221
- self.max_num_tokens = max_num_tokens
222
- self.heads_per_rank = heads_per_rank
223
- self.head_dim = head_dim
224
-
225
- output_shape = (
226
- num_ubatches,
227
- 2,
228
- self.world_size,
229
- max_num_tokens,
230
- heads_per_rank,
231
- head_dim,
232
- )
233
- lse_shape = (
234
- num_ubatches,
235
- 2,
236
- self.world_size,
237
- max_num_tokens,
238
- heads_per_rank,
239
- )
240
- signal_shape = (num_ubatches, 2, self.world_size)
241
- self.received_output, self.peer_output_ptrs = self._allocate(
242
- output_shape, dtype
243
- )
244
- self.received_lse, self.peer_lse_ptrs = self._allocate(lse_shape, torch.float32)
245
- self.received_signal, self.peer_signal_ptrs = self._allocate(
246
- signal_shape, torch.int32
247
- )
248
-
249
- def lse_reduce(
250
- self,
251
- partial_output: torch.Tensor,
252
- partial_lse: torch.Tensor,
253
- is_lse_base_on_e: bool,
254
- seq_lens: torch.Tensor | None = None,
255
- query_start_loc: torch.Tensor | None = None,
256
- ) -> torch.Tensor:
257
- ubatch = dbo_current_ubatch_id()
258
- num_tokens = partial_output.shape[0]
259
- output = partial_output.new_empty(
260
- (num_tokens, self.heads_per_rank, self.head_dim)
261
- )
262
- torch.ops._C.direct_dcp_a2a_lse_reduce(
263
- partial_output,
264
- partial_lse,
265
- seq_lens,
266
- query_start_loc,
267
- self.peer_output_ptrs[ubatch],
268
- self.peer_lse_ptrs[ubatch],
269
- self.peer_signal_ptrs[ubatch],
270
- self.received_output[ubatch],
271
- self.received_lse[ubatch],
272
- self.received_signal[ubatch],
273
- self.epoch[ubatch : ubatch + 1],
274
- output,
275
- self.world_size,
276
- self.rank,
277
- self.max_num_tokens,
278
- is_lse_base_on_e,
279
- )
280
- return output
281
-
282
-
283
-@functools.cache
284
-def get_direct_dcp_a2a_workspace(
285
- group: GroupCoordinator,
286
- device: torch.device,
287
- max_num_tokens: int,
288
- heads_per_rank: int,
289
- head_dim: int,
290
- dtype: torch.dtype,
291
- num_ubatches: int,
292
-) -> DirectDCPA2AWorkspace | None:
293
- if not _direct_dcp_enabled(
294
- group, dtype, envs.VLLM_USE_DIRECT_DCP_A2A, _A2A_SUPPORTED_DTYPES
295
- ):
296
- return None
297
- return DirectDCPA2AWorkspace(
298
- group.device_group,
299
- device,
300
- max_num_tokens,
301
- heads_per_rank,
302
- head_dim,
303
- dtype,
304
- num_ubatches,
305
- )
306
-
307
-
308
-def _q_gather_layout_supported(
309
- world_size: int,
310
- heads_per_rank: int,
311
- head_dim: int,
312
- dtype: torch.dtype,
313
- padded_num_heads: int | None,
314
-) -> bool:
315
- element_size = torch.empty((), dtype=dtype).element_size()
316
- gathered_num_heads = world_size * heads_per_rank
317
- storage_num_heads = (
318
- gathered_num_heads if padded_num_heads is None else padded_num_heads
319
- )
320
- return (
321
- heads_per_rank * head_dim * element_size % 16 == 0
322
- and storage_num_heads * head_dim * element_size % 16 == 0
323
- )
324
-
325
-
326
-class DirectDCPQGatherWorkspace(_DirectDCPWorkspace):
327
- """Publish query shards directly into the consumer-final symmetric buffer.
328
-
329
- The final buffer is reusable after the downstream DCP output combine. That
330
- combine orders all ranks after attention has consumed the gathered query.
331
- """
332
-
333
- def __init__(
334
- self,
335
- group: ProcessGroup,
336
- device: torch.device,
337
- max_num_tokens: int,
338
- heads_per_rank: int,
339
- head_dim: int,
340
- dtype: torch.dtype = torch.bfloat16,
341
- num_ubatches: int = 1,
342
- padded_num_heads: int | None = None,
343
- ) -> None:
344
- if num_ubatches < 1:
345
- raise ValueError(
346
- "Direct DCP q-gather requires at least one ubatch slot, "
347
- f"got {num_ubatches}"
348
- )
349
- if max_num_tokens < 1 or heads_per_rank < 1 or head_dim < 1:
350
- raise ValueError(
351
- "Direct DCP q-gather dimensions must be positive, got "
352
- f"T={max_num_tokens}, H={heads_per_rank}, D={head_dim}"
353
- )
354
- gathered_num_heads = group.size() * heads_per_rank
355
- if not _q_gather_layout_supported(
356
- group.size(), heads_per_rank, head_dim, dtype, padded_num_heads
357
- ):
358
- raise ValueError("Direct DCP q-gather requires 16-byte-aligned query rows.")
359
- super().__init__(group, device, num_ubatches)
360
- if self.world_size <= 1:
361
- raise ValueError("Direct DCP q-gather requires at least two ranks")
362
- self.max_num_tokens = max_num_tokens
363
- self.heads_per_rank = heads_per_rank
364
- self.gathered_num_heads = gathered_num_heads
365
- self.padded_num_heads = (
366
- self.gathered_num_heads if padded_num_heads is None else padded_num_heads
367
- )
368
- if self.padded_num_heads < self.gathered_num_heads:
369
- raise ValueError(
370
- "Direct DCP q-gather padded heads must cover gathered heads: "
371
- f"{self.padded_num_heads} < {self.gathered_num_heads}"
372
- )
373
- self.head_dim = head_dim
374
-
375
- query_shape = (
376
- num_ubatches,
377
- max_num_tokens,
378
- self.padded_num_heads,
379
- head_dim,
380
- )
381
- signal_shape = (num_ubatches, 2, self.world_size)
382
- self.final_query, _ = self._allocate(query_shape, dtype)
383
- self.received_signal, _ = self._allocate(signal_shape, torch.int32)
384
- query_multicast_ptrs = self._multicast_ptrs(self.final_query)
385
- signal_multicast_ptrs = self._multicast_ptrs(self.received_signal)
386
- self.multicast_ptrs = list(
387
- zip(query_multicast_ptrs, signal_multicast_ptrs, strict=True)
388
- )
389
- if not all(
390
- query_ptr and signal_ptr for query_ptr, signal_ptr in self.multicast_ptrs
391
- ):
392
- raise RuntimeError(
393
- "Direct DCP q-gather requires NVLS symmetric-memory multicast."
394
- )
395
- self.completion = self.received_signal.new_zeros((num_ubatches, 1))
396
- torch.accelerator.synchronize()
397
-
398
- def gather(self, local_query: torch.Tensor) -> torch.Tensor:
399
- ubatch = dbo_current_ubatch_id()
400
- if not 0 <= ubatch < self.num_ubatches:
401
- raise ValueError(
402
- f"DCP q-gather ubatch {ubatch} exceeds {self.num_ubatches} slots"
403
- )
404
- if local_query.ndim == 3 and local_query.shape[1] != self.heads_per_rank:
405
- raise ValueError(
406
- f"DCP q-gather expected {self.heads_per_rank} local query heads, "
407
- f"got {local_query.shape[1]}"
408
- )
409
-
410
- num_tokens = local_query.shape[0]
411
- output = torch.as_strided(
412
- self.final_query[ubatch],
413
- size=(num_tokens, self.gathered_num_heads, self.head_dim),
414
- stride=(
415
- self.gathered_num_heads * self.head_dim,
416
- self.head_dim,
417
- 1,
418
- ),
419
- )
420
- query_multicast_ptr, signal_multicast_ptr = self.multicast_ptrs[ubatch]
421
- torch.ops._C.direct_dcp_q_gather(
422
- local_query,
423
- output,
424
- self.received_signal[ubatch],
425
- self.completion[ubatch],
426
- self.epoch[ubatch : ubatch + 1],
427
- self.world_size,
428
- self.rank,
429
- self.max_num_tokens,
430
- self.padded_num_heads,
431
- query_multicast_ptr,
432
- signal_multicast_ptr,
433
- )
434
- return output
435
-
436
-
437
-@functools.cache
438
-def get_direct_dcp_q_gather_workspace(
439
- group: GroupCoordinator,
440
- device: torch.device,
441
- max_num_tokens: int,
442
- heads_per_rank: int,
443
- head_dim: int,
444
- dtype: torch.dtype,
445
- num_ubatches: int,
446
- padded_num_heads: int | None = None,
447
-) -> DirectDCPQGatherWorkspace | None:
448
- if not _direct_dcp_multicast_enabled(
449
- group, dtype, envs.VLLM_USE_DIRECT_DCP_Q_GATHER
450
- ):
451
- return None
452
- if not _q_gather_layout_supported(
453
- group.world_size, heads_per_rank, head_dim, dtype, padded_num_heads
454
- ):
455
- return None
456
- return DirectDCPQGatherWorkspace(
457
- group.device_group,
458
- device,
459
- max_num_tokens,
460
- heads_per_rank,
461
- head_dim,
462
- dtype,
463
- num_ubatches,
464
- padded_num_heads,
465
- )
466
-
467
-
468
-_KV_GATHER_SUPPORTED_DTYPES = (
469
- torch.float16,
470
- torch.bfloat16,
471
- torch.float8_e4m3fn,
472
-)
473
-
474
-
475
-def _kv_gather_layout_supported(token_dim: int, dtype: torch.dtype) -> bool:
476
- return token_dim * torch.empty((), dtype=dtype).element_size() % 16 == 0
477
-
478
-
479
-class DirectDCPKVGatherWorkspace(_DirectDCPWorkspace):
480
- """Persistent symmetric buffers for direct DCP KV gather."""
481
-
482
- def __init__(
483
- self,
484
- group: ProcessGroup,
485
- device: torch.device,
486
- max_gathered_tokens: int,
487
- token_dim: int,
488
- dtype: torch.dtype = torch.bfloat16,
489
- num_ubatches: int = 1,
490
- ) -> None:
491
- if dtype not in _KV_GATHER_SUPPORTED_DTYPES:
492
- raise ValueError(f"Direct DCP kv-gather does not support {dtype}")
493
- if num_ubatches < 1:
494
- raise ValueError(
495
- "Direct DCP kv-gather requires at least one ubatch slot, "
496
- f"got {num_ubatches}"
497
- )
498
- if max_gathered_tokens < 1 or token_dim < 1:
499
- raise ValueError(
500
- "Direct DCP kv-gather dimensions must be positive, got "
501
- f"T={max_gathered_tokens}, D={token_dim}"
502
- )
503
- if not _kv_gather_layout_supported(token_dim, dtype):
504
- raise ValueError("Direct DCP kv-gather requires 16-byte-aligned KV rows.")
505
- super().__init__(group, device, num_ubatches)
506
- if self.world_size <= 1:
507
- raise ValueError("Direct DCP kv-gather requires at least two ranks")
508
- if max_gathered_tokens % self.world_size != 0:
509
- raise ValueError(
510
- "Direct DCP kv-gather capacity must divide evenly across "
511
- f"ranks: {max_gathered_tokens} % {self.world_size} != 0"
512
- )
513
- self.max_gathered_tokens = max_gathered_tokens
514
-
515
- kv_shape = (num_ubatches, 2, max_gathered_tokens, token_dim)
516
- signal_shape = (num_ubatches, 2, self.world_size)
517
- self.received_kv, _ = self._allocate(kv_shape, dtype)
518
- self.received_signal, _ = self._allocate(signal_shape, torch.int32)
519
- kv_multicast_ptrs = self._multicast_ptrs(self.received_kv)
520
- signal_multicast_ptrs = self._multicast_ptrs(self.received_signal)
521
- self.multicast_ptrs = list(
522
- zip(kv_multicast_ptrs, signal_multicast_ptrs, strict=True)
523
- )
524
- if not all(kv_ptr and signal_ptr for kv_ptr, signal_ptr in self.multicast_ptrs):
525
- raise RuntimeError(
526
- "Direct DCP kv-gather requires NVLS symmetric-memory multicast."
527
- )
528
- self.completion = self.received_signal.new_zeros((num_ubatches, 2))
529
- torch.accelerator.synchronize()
530
-
531
- def gather(self, gathered_kv: torch.Tensor, local_kv: torch.Tensor) -> None:
532
- ubatch = dbo_current_ubatch_id()
533
- if not 0 <= ubatch < self.num_ubatches:
534
- raise ValueError(
535
- f"DCP kv-gather ubatch {ubatch} exceeds {self.num_ubatches} slots"
536
- )
537
- kv_multicast_ptr, signal_multicast_ptr = self.multicast_ptrs[ubatch]
538
- torch.ops._C.direct_dcp_kv_gather(
539
- local_kv,
540
- self.received_kv[ubatch],
541
- self.received_signal[ubatch],
542
- self.completion[ubatch],
543
- self.epoch[ubatch : ubatch + 1],
544
- gathered_kv,
545
- self.world_size,
546
- self.rank,
547
- self.max_gathered_tokens,
548
- kv_multicast_ptr,
549
- signal_multicast_ptr,
550
- )
551
-
552
-
553
-@functools.cache
554
-def get_direct_dcp_kv_gather_workspace(
555
- group: GroupCoordinator,
556
- device: torch.device,
557
- max_gathered_tokens: int,
558
- token_dim: int,
559
- dtype: torch.dtype,
560
- num_ubatches: int,
561
-) -> DirectDCPKVGatherWorkspace | None:
562
- if not _direct_dcp_multicast_enabled(
563
- group,
564
- dtype,
565
- envs.VLLM_USE_DIRECT_DCP_KV_GATHER,
566
- _KV_GATHER_SUPPORTED_DTYPES,
567
- ):
568
- return None
569
- if not _kv_gather_layout_supported(token_dim, dtype):
570
- return None
571
- return DirectDCPKVGatherWorkspace(
572
- group.device_group,
573
- device,
574
- max_gathered_tokens,
575
- token_dim,
576
- dtype,
577
- num_ubatches,
578
- )
579
-
580
-
581
-class DCPCombine(Protocol):
582
- def __call__(
583
- self,
584
- partial_output: torch.Tensor,
585
- partial_lse: torch.Tensor,
586
- *,
587
- seq_lens: torch.Tensor,
588
- query_start_loc: torch.Tensor,
589
- ) -> torch.Tensor: ...
590
-
591
-
592
-class MLADCPManager:
593
- """Select and own layer-level collective implementations for MLA DCP."""
594
-
595
- _kv_gather: Callable[[torch.Tensor, torch.Tensor], object]
596
-
597
- def __init__(
598
- self,
599
- vllm_config: VllmConfig,
600
- device: torch.device,
601
- num_heads: int,
602
- query_head_dim: int,
603
- output_head_dim: int,
604
- query_dtype: torch.dtype,
605
- output_dtype: torch.dtype,
606
- padded_num_heads: int | None,
607
- is_lse_base_on_e: bool,
608
- use_pcp: bool,
609
- ) -> None:
610
- parallel_config = vllm_config.parallel_config
611
- self.group = get_dcp_group()
612
- self.device = torch.device(device)
613
- self.num_ubatches = max(parallel_config.num_ubatches, 1)
614
- self.max_num_tokens = get_dcp_workspace_max_num_tokens(vllm_config)
615
- self.use_a2a = parallel_config.dcp_comm_backend == "a2a"
616
- self.padded_num_heads = padded_num_heads
617
-
618
- self.combine = self._init_combine(
619
- num_heads,
620
- output_head_dim,
621
- output_dtype,
622
- is_lse_base_on_e,
623
- use_pcp,
624
- )
625
- self.query_gather = (
626
- None
627
- if use_pcp
628
- else self._init_query_gather(
629
- num_heads,
630
- query_head_dim,
631
- query_dtype,
632
- )
633
- )
634
-
635
- def _init_combine(
636
- self,
637
- num_heads: int,
638
- head_dim: int,
639
- dtype: torch.dtype,
640
- is_lse_base_on_e: bool,
641
- use_pcp: bool,
642
- ) -> DCPCombine:
643
- direct_workspace = None
644
- if self.use_a2a:
645
- direct_workspace = get_direct_dcp_a2a_workspace(
646
- self.group,
647
- self.device,
648
- self.max_num_tokens,
649
- num_heads,
650
- head_dim,
651
- dtype,
652
- self.num_ubatches,
653
- )
654
- if direct_workspace is not None:
655
- logger.info_once("Using direct symmetric-memory DCP A2A for MLA.")
656
- return functools.partial(
657
- direct_workspace.lse_reduce,
658
- is_lse_base_on_e=is_lse_base_on_e,
659
- )
660
-
661
- combine_fn = (
662
- dcp_a2a_lse_reduce
663
- if self.use_a2a
664
- else cp_lse_ag_out_ar
665
- if use_pcp
666
- else cp_lse_ag_out_rs
667
- )
668
- return functools.partial(
669
- combine_fn,
670
- cp_group=self.group,
671
- is_lse_base_on_e=is_lse_base_on_e,
672
- )
673
-
674
- def _init_query_gather(
675
- self,
676
- num_heads: int,
677
- head_dim: int,
678
- dtype: torch.dtype,
679
- ) -> Callable[[torch.Tensor], torch.Tensor]:
680
- direct_workspace = get_direct_dcp_q_gather_workspace(
681
- self.group,
682
- self.device,
683
- self.max_num_tokens,
684
- num_heads,
685
- head_dim,
686
- dtype,
687
- self.num_ubatches,
688
- self.padded_num_heads,
689
- )
690
- if direct_workspace is not None:
691
- logger.info_once("Using direct symmetric-memory DCP query gather for MLA.")
692
- return direct_workspace.gather
693
- return self._gather_query
694
-
695
- def _gather_query(self, query: torch.Tensor) -> torch.Tensor:
696
- query = self.group.all_gather(query, dim=1)
697
- if self.padded_num_heads is not None:
698
- query = reserve_query_head_storage(query, self.padded_num_heads)
699
- return query
700
-
701
- def init_kv_gather(
702
- self,
703
- workspace: torch.Tensor,
704
- max_gathered_tokens: int,
705
- ) -> None:
706
- world_size = self.group.world_size
707
- assert max_gathered_tokens > 0
708
- assert max_gathered_tokens % world_size == 0
709
- assert workspace.ndim == 2
710
- assert workspace.is_contiguous()
711
- assert workspace.shape[0] == (
712
- max_gathered_tokens + max_gathered_tokens // world_size
713
- )
714
- assert workspace.shape[1] > 0
715
-
716
- direct_workspace = get_direct_dcp_kv_gather_workspace(
717
- self.group,
718
- workspace.device,
719
- max_gathered_tokens,
720
- workspace.shape[1],
721
- workspace.dtype,
722
- self.num_ubatches,
723
- )
724
- if direct_workspace is not None:
725
- logger.info_once(
726
- "Using direct symmetric-memory DCP chunked-context KV gather for MLA."
727
- )
728
- self._kv_gather = direct_workspace.gather
729
- else:
730
- self._kv_gather = functools.partial(
731
- torch.distributed.all_gather_into_tensor,
732
- group=self.group.device_group,
733
- )
734
-
735
- def kv_gather(
736
- self,
737
- gathered_kv: torch.Tensor,
738
- local_kv: torch.Tensor,
739
- ) -> object:
740
- return self._kv_gather(gathered_kv, local_kv)