@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4f7bde572a
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 10:00 UTC

[Kernel] Pack output and LSE in DCP A2A (#41160)

sungsoo ha committed May 1, 2026 at 06:01 UTC 4f7bde572ad05a6c013dde5f8874898fff3c1253
2 files changed +555 -157
tests/distributed/test_dcp_a2a.py
+306 -3
@@ -10,10 +10,95 @@ Tests cover:
10
11 import math
12
13 +import multiprocess as mp
14 import pytest
15 import torch
16 +import torch.distributed as dist
17
18 from vllm.config.parallel import ParallelConfig
19 +from vllm.utils.network_utils import get_open_port
20 +from vllm.utils.system_utils import update_environment_variables
21 +
22 +mp.set_start_method("spawn", force=True)
23 +
24 +
25 +class _FakeCPGroup:
26 + def __init__(self, world_size: int, device_group: dist.ProcessGroup):
27 + self.world_size = world_size
28 + self.device_group = device_group
29 +
30 +
31 +def _dtype_from_name(dtype_name: str) -> torch.dtype:
32 + return {
33 + "float16": torch.float16,
34 + "bfloat16": torch.bfloat16,
35 + "float32": torch.float32,
36 + }[dtype_name]
37 +
38 +
39 +def _packed_a2a_reference(
40 + cp_attn_out: torch.Tensor,
41 + cp_attn_lse: torch.Tensor,
42 + world_size: int,
43 + h_per_rank: int,
44 + is_lse_base_on_e: bool,
45 +) -> tuple[torch.Tensor, torch.Tensor]:
46 + from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
47 +
48 + B, _H, D = cp_attn_out.shape
49 + outputs = (
50 + cp_attn_out.view(B, world_size, h_per_rank, D)
51 + .permute(1, 0, 2, 3)
52 + .contiguous()
53 + .float()
54 + )
55 + lses = cp_attn_lse.view(B, world_size, h_per_rank).permute(1, 0, 2).contiguous()
56 + return _lse_weighted_combine(
57 + outputs,
58 + lses,
59 + return_lse=True,
60 + is_lse_base_on_e=is_lse_base_on_e,
61 + )
62 +
63 +
64 +def _assert_packed_a2a_close(
65 + actual: torch.Tensor,
66 + expected: torch.Tensor,
67 + dtype: torch.dtype,
68 +) -> None:
69 + if dtype == torch.float32:
70 + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5)
71 + else:
72 + torch.testing.assert_close(
73 + actual.float(), expected.float(), rtol=3e-2, atol=3e-2
74 + )
75 +
76 +
77 +def _distributed_run(fn, world_size: int, extra_env: dict[str, str]) -> None:
78 + port = str(get_open_port())
79 + processes: list[mp.Process] = []
80 + for rank in range(world_size):
81 + env = {
82 + "RANK": str(rank),
83 + "LOCAL_RANK": str(rank),
84 + "WORLD_SIZE": str(world_size),
85 + "LOCAL_WORLD_SIZE": str(world_size),
86 + "MASTER_ADDR": "localhost",
87 + "MASTER_PORT": port,
88 + **extra_env,
89 + }
90 + process = mp.Process(target=fn, args=(env,))
91 + processes.append(process)
92 + process.start()
93 +
94 + for process in processes:
95 + process.join(timeout=120)
96 +
97 + for process in processes:
98 + if process.is_alive():
99 + process.kill()
100 + process.join()
101 + assert process.exitcode == 0
102
103
104 class TestDCPCommBackendConfig:
@@ -38,14 +123,14 @@ class TestDCPCommBackendConfig:
123 """A2A backend is valid when DCP > 1."""
124 config = ParallelConfig(
125 dcp_comm_backend="a2a",
41 - tensor_parallel_size=8,
126 + tensor_parallel_size=4,
127 decode_context_parallel_size=4,
128 )
129 assert config.dcp_comm_backend == "a2a"
130
131 def test_invalid_backend_rejected(self):
132 """Invalid backend values are rejected."""
48 - with pytest.raises(ValueError, match="must be one of"):
133 + with pytest.raises(ValueError, match="must be one of|Input should be"):
134 ParallelConfig(
135 dcp_comm_backend="invalid",
136 )
@@ -134,7 +219,7 @@ class TestLSEWeightedCombine:
219 result = _lse_weighted_combine(outputs, lses)
220
221 assert result.shape == (B, H, D)
137 - torch.testing.assert_close(result, outputs[1].squeeze(0), atol=1e-5, rtol=1e-5)
222 + torch.testing.assert_close(result, outputs[1], atol=1e-5, rtol=1e-5)
223
224 def test_mathematically_correct(self):
225 """Verify mathematical correctness of LSE combination."""
@@ -187,6 +272,224 @@ class TestLSEWeightedCombine:
272 assert global_lse.shape == (B, H)
273 assert abs(global_lse.item() - expected_global_lse) < 1e-5
274
275 + def test_base2_return_lse(self):
276 + """Base-2 LSE mode returns log2-sum-exp2 global LSE."""
277 + from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
278 +
279 + outputs = torch.tensor(
280 + [
281 + [[[1.0, 2.0]]],
282 + [[[3.0, 4.0]]],
283 + ]
284 + )
285 + lses = torch.tensor(
286 + [
287 + [[1.0]],
288 + [[2.0]],
289 + ]
290 + )
291 +
292 + result, global_lse = _lse_weighted_combine(
293 + outputs,
294 + lses,
295 + return_lse=True,
296 + is_lse_base_on_e=False,
297 + )
298 +
299 + expected_global_lse = math.log2(2**1 + 2**2)
300 + w0 = 2**1 / (2**1 + 2**2)
301 + w1 = 2**2 / (2**1 + 2**2)
302 + expected = torch.tensor([[[w0 * 1.0 + w1 * 3.0, w0 * 2.0 + w1 * 4.0]]])
303 +
304 + torch.testing.assert_close(result, expected, rtol=1e-5, atol=1e-5)
305 + torch.testing.assert_close(
306 + global_lse,
307 + torch.tensor([[expected_global_lse]]),
308 + rtol=1e-5,
309 + atol=1e-5,
310 + )
311 +
312 + def test_lse_pack_dim(self):
313 + """Packed A2A stores one fp32 LSE in output-dtype lanes."""
314 + from vllm.v1.attention.ops.dcp_alltoall import _dcp_a2a_lse_pack_dim
315 +
316 + assert _dcp_a2a_lse_pack_dim(torch.bfloat16) == 2
317 + assert _dcp_a2a_lse_pack_dim(torch.float16) == 2
318 + assert _dcp_a2a_lse_pack_dim(torch.float32) == 1
319 +
320 +
321 +class TestPackedA2AKernels:
322 + @pytest.mark.skipif(
323 + torch.accelerator.device_count() < 1, reason="CUDA is required."
324 + )
325 + @pytest.mark.parametrize("dtype_name", ["float16", "bfloat16", "float32"])
326 + @pytest.mark.parametrize("return_lse", [False, True])
327 + @pytest.mark.parametrize("is_lse_base_on_e", [False, True])
328 + def test_pack_unpack_combine_matches_reference(
329 + self,
330 + dtype_name: str,
331 + return_lse: bool,
332 + is_lse_base_on_e: bool,
333 + ):
334 + from vllm.v1.attention.ops.dcp_alltoall import (
335 + _dcp_a2a_lse_pack_dim,
336 + _dcp_a2a_pack_send,
337 + _dcp_a2a_unpack_combine,
338 + )
339 +
340 + torch.manual_seed(0)
341 + dtype = _dtype_from_name(dtype_name)
342 + device = torch.device("cuda")
343 + world_size, B, h_per_rank, D = 4, 7, 2, 32
344 + H = world_size * h_per_rank
345 + cp_attn_out = torch.randn(B, H, D, device=device, dtype=dtype)
346 + cp_attn_lse = torch.randn(B, H, device=device, dtype=torch.float32)
347 + lse_pack_dim = _dcp_a2a_lse_pack_dim(dtype)
348 + send_buffer = torch.empty(
349 + (world_size, B, h_per_rank, D + lse_pack_dim),
350 + device=device,
351 + dtype=dtype,
352 + )
353 +
354 + _dcp_a2a_pack_send(
355 + cp_attn_out,
356 + cp_attn_lse,
357 + send_buffer,
358 + world_size,
359 + h_per_rank,
360 + D,
361 + lse_pack_dim,
362 + )
363 + actual = _dcp_a2a_unpack_combine(
364 + send_buffer, D, lse_pack_dim, return_lse, is_lse_base_on_e
365 + )
366 + expected_out, expected_lse = _packed_a2a_reference(
367 + cp_attn_out, cp_attn_lse, world_size, h_per_rank, is_lse_base_on_e
368 + )
369 +
370 + if return_lse:
371 + actual_out, actual_lse = actual
372 + _assert_packed_a2a_close(actual_out, expected_out, dtype)
373 + torch.testing.assert_close(actual_lse, expected_lse, rtol=1e-4, atol=1e-4)
374 + else:
375 + _assert_packed_a2a_close(actual, expected_out, dtype)
376 +
377 +
378 +def _distributed_packed_a2a_worker(env: dict[str, str]) -> None:
379 + update_environment_variables(env)
380 + local_rank = int(env["LOCAL_RANK"])
381 + torch.accelerator.set_device_index(local_rank)
382 + dist.init_process_group(backend="nccl")
383 + use_workspace = env.get("USE_WORKSPACE") == "1"
384 + if use_workspace:
385 + from vllm.v1.worker.workspace import init_workspace_manager
386 +
387 + init_workspace_manager(torch.device(f"cuda:{local_rank}"))
388 + try:
389 + from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
390 +
391 + dtype = _dtype_from_name(env["TEST_DTYPE"])
392 + return_lse = env["RETURN_LSE"] == "1"
393 + is_lse_base_on_e = env["LSE_BASE_E"] == "1"
394 + rank = dist.get_rank()
395 + world_size = dist.get_world_size()
396 + B, h_per_rank, D = 5, 2, 32
397 + H = world_size * h_per_rank
398 +
399 + generator = torch.Generator(device=f"cuda:{local_rank}")
400 + generator.manual_seed(1234 + rank)
401 + cp_attn_out = torch.randn(
402 + B,
403 + H,
404 + D,
405 + device=f"cuda:{local_rank}",
406 + dtype=dtype,
407 + generator=generator,
408 + )
409 + cp_attn_lse = torch.randn(
410 + B,
411 + H,
412 + device=f"cuda:{local_rank}",
413 + dtype=torch.float32,
414 + generator=generator,
415 + )
416 + actual = dcp_a2a_lse_reduce(
417 + cp_attn_out,
418 + cp_attn_lse,
419 + _FakeCPGroup(world_size, dist.group.WORLD),
420 + return_lse=return_lse,
421 + is_lse_base_on_e=is_lse_base_on_e,
422 + )
423 +
424 + gathered_out = [torch.empty_like(cp_attn_out) for _ in range(world_size)]
425 + gathered_lse = [torch.empty_like(cp_attn_lse) for _ in range(world_size)]
426 + dist.all_gather(gathered_out, cp_attn_out)
427 + dist.all_gather(gathered_lse, cp_attn_lse)
428 + outputs = torch.stack(
429 + [
430 + t[:, rank * h_per_rank : (rank + 1) * h_per_rank, :]
431 + for t in gathered_out
432 + ],
433 + dim=0,
434 + ).float()
435 + lses = torch.stack(
436 + [t[:, rank * h_per_rank : (rank + 1) * h_per_rank] for t in gathered_lse],
437 + dim=0,
438 + )
439 + from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
440 +
441 + expected_out, expected_lse = _lse_weighted_combine(
442 + outputs,
443 + lses,
444 + return_lse=True,
445 + is_lse_base_on_e=is_lse_base_on_e,
446 + )
447 +
448 + if return_lse:
449 + actual_out, actual_lse = actual
450 + _assert_packed_a2a_close(actual_out, expected_out, dtype)
451 + torch.testing.assert_close(actual_lse, expected_lse, rtol=1e-4, atol=1e-4)
452 + else:
453 + _assert_packed_a2a_close(actual, expected_out, dtype)
454 + finally:
455 + if use_workspace:
456 + from vllm.v1.worker.workspace import reset_workspace_manager
457 +
458 + reset_workspace_manager()
459 + dist.destroy_process_group()
460 +
461 +
462 +@pytest.mark.skipif(
463 + torch.accelerator.device_count() < 4, reason="Need at least 4 GPUs."
464 +)
465 +@pytest.mark.parametrize("dtype_name", ["float16", "bfloat16", "float32"])
466 +def test_distributed_packed_a2a_matches_reference(dtype_name: str):
467 + _distributed_run(
468 + _distributed_packed_a2a_worker,
469 + world_size=4,
470 + extra_env={
471 + "TEST_DTYPE": dtype_name,
472 + "RETURN_LSE": "1",
473 + "LSE_BASE_E": "1",
474 + },
475 + )
476 +
477 +
478 +@pytest.mark.skipif(
479 + torch.accelerator.device_count() < 4, reason="Need at least 4 GPUs."
480 +)
481 +def test_distributed_packed_a2a_with_workspace_matches_reference():
482 + _distributed_run(
483 + _distributed_packed_a2a_worker,
484 + world_size=4,
485 + extra_env={
486 + "TEST_DTYPE": "bfloat16",
487 + "RETURN_LSE": "1",
488 + "LSE_BASE_E": "1",
489 + "USE_WORKSPACE": "1",
490 + },
491 + )
492 +
493
494 if __name__ == "__main__":
495 pytest.main([__file__, "-v"])
vllm/v1/attention/ops/dcp_alltoall.py
+249 -154
@@ -9,10 +9,8 @@ Instead of gathering the full Q tensor and scattering partial outputs,
9 A2A exchanges partial attention outputs and their LSE values across
10 ranks, then combines them with exact LSE-weighted reduction.
11
12 -This reduces the number of NCCL calls per attention layer from 3
13 -(AG for Q, AG for K metadata, RS for output) to 2 (A2A for output,
14 -A2A for LSE), lowering per-step communication overhead for long-context
15 -decode where NCCL latency is a significant fraction of step time.
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
@@ -28,6 +26,10 @@ import torch
26 import torch.distributed as dist
27
28 from vllm.triton_utils import tl, triton
29 +from vllm.v1.worker.workspace import (
30 + current_workspace_manager,
31 + is_workspace_manager_initialized,
32 +)
33
34 if TYPE_CHECKING:
35 from vllm.distributed.parallel_state import GroupCoordinator
@@ -44,7 +46,6 @@ def _lse_weighted_combine(
46 CPU reference implementation for LSE-weighted combination.
47
48 This is a pure PyTorch implementation used for testing and validation.
47 - For GPU execution, use dcp_lse_combine_triton instead.
49
50 Args:
51 outputs: Partial attention outputs [N, B, H, D]
@@ -102,57 +103,137 @@ def _lse_weighted_combine(
103 return result
104
105
106 +def _dcp_a2a_lse_pack_dim(output_dtype: torch.dtype) -> int:
107 + bits = torch.finfo(output_dtype).bits
108 + if bits == 16:
109 + return 2
110 + if bits == 32:
111 + return 1
112 + raise ValueError(f"Cannot pack fp32 LSE into output dtype {output_dtype}.")
113 +
114 +
115 +def _dcp_a2a_send_recv_buffers(
116 + shape: tuple[int, ...],
117 + device: torch.device,
118 + dtype: torch.dtype,
119 +) -> tuple[torch.Tensor, torch.Tensor]:
120 + if is_workspace_manager_initialized():
121 + send_buffer, recv_buffer = current_workspace_manager().get_simultaneous(
122 + (shape, dtype),
123 + (shape, dtype),
124 + )
125 + return send_buffer, recv_buffer
126 +
127 + return (
128 + torch.empty(shape, device=device, dtype=dtype),
129 + torch.empty(shape, device=device, dtype=dtype),
130 + )
131 +
132 +
133 @triton.jit
106 -def _dcp_lse_combine_kernel(
107 - # Input pointers
108 - recv_output_ptr,
109 - recv_lse_ptr,
110 - # Output pointers
134 +def _dcp_a2a_pack_send_kernel(
135 out_ptr,
112 - out_lse_ptr,
113 - # Strides for recv_output [N, B, H_local, D]
114 - ro_stride_N,
115 - ro_stride_B,
116 - ro_stride_H,
117 - ro_stride_D,
118 - # Strides for recv_lse [N, B, H_local]
119 - rl_stride_N,
120 - rl_stride_B,
121 - rl_stride_H,
122 - # Strides for output [B, H_local, D]
123 - o_stride_B,
124 - o_stride_H,
125 - o_stride_D,
126 - # Constants
136 + lse_ptr,
137 + send_ptr,
138 + out_stride_B,
139 + out_stride_H,
140 + out_stride_D,
141 + lse_stride_B,
142 + lse_stride_H,
143 + send_stride_N,
144 + send_stride_B,
145 + send_stride_H,
146 + send_stride_D,
147 N: tl.constexpr,
148 HEAD_DIM: tl.constexpr,
129 - IS_BASE_E: tl.constexpr,
130 - RETURN_LSE: tl.constexpr,
149 + H_PER_RANK: tl.constexpr,
150 + LSE_PACK_DIM: tl.constexpr,
151 ):
132 - """
133 - Triton kernel for LSE-weighted combination of partial attention outputs.
152 + batch_idx = tl.program_id(0).to(tl.int64)
153 + local_head_idx = tl.program_id(1).to(tl.int64)
154 + d_offsets = tl.arange(0, HEAD_DIM)
155
135 - After All-to-All, each rank has:
136 - - recv_output [N, B, H_local, D]: partial outputs from all KV shards
137 - - recv_lse [N, B, H_local]: partial LSEs from all KV shards
156 + for rank_idx in tl.static_range(N):
157 + src_head_idx = rank_idx * H_PER_RANK + local_head_idx
158 + send_base = (
159 + rank_idx * send_stride_N
160 + + batch_idx * send_stride_B
161 + + local_head_idx * send_stride_H
162 + )
163
139 - This kernel computes the weighted combination locally (no communication).
164 + out_offsets = (
165 + batch_idx * out_stride_B
166 + + src_head_idx * out_stride_H
167 + + d_offsets * out_stride_D
168 + )
169 + tl.store(
170 + send_ptr + send_base + d_offsets * send_stride_D,
171 + tl.load(out_ptr + out_offsets),
172 + )
173
141 - Grid: (B, H_local)
142 - Each program handles one (batch, head) and processes all D elements.
143 - """
174 + lse_val = tl.load(
175 + lse_ptr + batch_idx * lse_stride_B + src_head_idx * lse_stride_H
176 + )
177 + if LSE_PACK_DIM == 1:
178 + tl.store(
179 + send_ptr + send_base + HEAD_DIM * send_stride_D,
180 + lse_val.to(send_ptr.dtype.element_ty),
181 + )
182 + else:
183 + lse_bits = lse_val.to(tl.uint32, bitcast=True)
184 + lo = (lse_bits & 0xFFFF).to(tl.uint16)
185 + hi = ((lse_bits >> 16) & 0xFFFF).to(tl.uint16)
186 + tl.store(
187 + send_ptr + send_base + HEAD_DIM * send_stride_D,
188 + lo.to(send_ptr.dtype.element_ty, bitcast=True),
189 + )
190 + tl.store(
191 + send_ptr + send_base + (HEAD_DIM + 1) * send_stride_D,
192 + hi.to(send_ptr.dtype.element_ty, bitcast=True),
193 + )
194 +
195 +
196 +@triton.jit
197 +def _dcp_a2a_unpack_combine_kernel(
198 + recv_ptr,
199 + out_ptr,
200 + out_lse_ptr,
201 + recv_stride_N,
202 + recv_stride_B,
203 + recv_stride_H,
204 + recv_stride_D,
205 + out_stride_B,
206 + out_stride_H,
207 + out_stride_D,
208 + out_lse_stride_B,
209 + out_lse_stride_H,
210 + N: tl.constexpr,
211 + HEAD_DIM: tl.constexpr,
212 + IS_BASE_E: tl.constexpr,
213 + RETURN_LSE: tl.constexpr,
214 + LSE_PACK_DIM: tl.constexpr,
215 +):
216 batch_idx = tl.program_id(0).to(tl.int64)
217 head_idx = tl.program_id(1).to(tl.int64)
218 + d_offsets = tl.arange(0, HEAD_DIM)
219
147 - # Base offset for this (batch, head)
148 - base_lse_offset = batch_idx * rl_stride_B + head_idx * rl_stride_H
149 - base_out_offset = batch_idx * ro_stride_B + head_idx * ro_stride_H
150 -
151 - # First pass: find max LSE for numerical stability
220 lse_max = -float("inf")
153 - for n in tl.static_range(N):
154 - lse_offset = n * rl_stride_N + base_lse_offset
155 - lse_val = tl.load(recv_lse_ptr + lse_offset)
221 + for rank_idx in tl.static_range(N):
222 + recv_base = (
223 + rank_idx * recv_stride_N
224 + + batch_idx * recv_stride_B
225 + + head_idx * recv_stride_H
226 + )
227 + if LSE_PACK_DIM == 1:
228 + lse_val = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D).to(
229 + tl.float32
230 + )
231 + else:
232 + lo_raw = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D)
233 + hi_raw = tl.load(recv_ptr + recv_base + (HEAD_DIM + 1) * recv_stride_D)
234 + lo = lo_raw.to(tl.uint16, bitcast=True).to(tl.uint32)
235 + hi = hi_raw.to(tl.uint16, bitcast=True).to(tl.uint32)
236 + lse_val = (lo | (hi << 16)).to(tl.float32, bitcast=True)
237 lse_val = tl.where(
238 (lse_val != lse_val) | (lse_val == float("inf")),
239 -float("inf"),
@@ -162,11 +243,23 @@ def _dcp_lse_combine_kernel(
243
244 lse_max = tl.where(lse_max == -float("inf"), 0.0, lse_max)
245
165 - # Second pass: compute sum of exp(lse - max)
246 lse_sum = 0.0
167 - for n in tl.static_range(N):
168 - lse_offset = n * rl_stride_N + base_lse_offset
169 - lse_val = tl.load(recv_lse_ptr + lse_offset)
247 + for rank_idx in tl.static_range(N):
248 + recv_base = (
249 + rank_idx * recv_stride_N
250 + + batch_idx * recv_stride_B
251 + + head_idx * recv_stride_H
252 + )
253 + if LSE_PACK_DIM == 1:
254 + lse_val = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D).to(
255 + tl.float32
256 + )
257 + else:
258 + lo_raw = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D)
259 + hi_raw = tl.load(recv_ptr + recv_base + (HEAD_DIM + 1) * recv_stride_D)
260 + lo = lo_raw.to(tl.uint16, bitcast=True).to(tl.uint32)
261 + hi = hi_raw.to(tl.uint16, bitcast=True).to(tl.uint32)
262 + lse_val = (lo | (hi << 16)).to(tl.float32, bitcast=True)
263 lse_val = tl.where(
264 (lse_val != lse_val) | (lse_val == float("inf")),
265 -float("inf"),
@@ -177,19 +270,28 @@ def _dcp_lse_combine_kernel(
270 else:
271 lse_sum += tl.exp2(lse_val - lse_max)
272
180 - # Compute global LSE
273 if IS_BASE_E: # noqa: SIM108
274 global_lse = tl.log(lse_sum) + lse_max
275 else:
276 global_lse = tl.log2(lse_sum) + lse_max
277
186 - # Third pass: weighted combination across D dimension
187 - d_offsets = tl.arange(0, HEAD_DIM)
278 acc = tl.zeros([HEAD_DIM], dtype=tl.float32)
189 -
190 - for n in tl.static_range(N):
191 - lse_offset = n * rl_stride_N + base_lse_offset
192 - lse_val = tl.load(recv_lse_ptr + lse_offset)
279 + for rank_idx in tl.static_range(N):
280 + recv_base = (
281 + rank_idx * recv_stride_N
282 + + batch_idx * recv_stride_B
283 + + head_idx * recv_stride_H
284 + )
285 + if LSE_PACK_DIM == 1:
286 + lse_val = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D).to(
287 + tl.float32
288 + )
289 + else:
290 + lo_raw = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D)
291 + hi_raw = tl.load(recv_ptr + recv_base + (HEAD_DIM + 1) * recv_stride_D)
292 + lo = lo_raw.to(tl.uint16, bitcast=True).to(tl.uint32)
293 + hi = hi_raw.to(tl.uint16, bitcast=True).to(tl.uint32)
294 + lse_val = (lo | (hi << 16)).to(tl.float32, bitcast=True)
295 lse_val = tl.where(
296 (lse_val != lse_val) | (lse_val == float("inf")),
297 -float("inf"),
@@ -200,80 +302,89 @@ def _dcp_lse_combine_kernel(
302 else:
303 weight = tl.exp2(lse_val - global_lse)
304 weight = tl.where(weight != weight, 0.0, weight)
305 + acc += (
306 + tl.load(recv_ptr + recv_base + d_offsets * recv_stride_D).to(tl.float32)
307 + * weight
308 + )
309
204 - out_offsets = n * ro_stride_N + base_out_offset + d_offsets * ro_stride_D
205 - out_vals = tl.load(recv_output_ptr + out_offsets)
206 - acc += out_vals.to(tl.float32) * weight
207 -
208 - # Store result
310 final_offsets = (
210 - batch_idx * o_stride_B + head_idx * o_stride_H + d_offsets * o_stride_D
311 + batch_idx * out_stride_B + head_idx * out_stride_H + d_offsets * out_stride_D
312 )
313 tl.store(out_ptr + final_offsets, acc)
314
315 if RETURN_LSE:
215 - tl.store(out_lse_ptr + base_lse_offset, global_lse)
216 -
316 + out_lse_offset = batch_idx * out_lse_stride_B + head_idx * out_lse_stride_H
317 + tl.store(out_lse_ptr + out_lse_offset, global_lse)
318
218 -def dcp_lse_combine_triton(
219 - recv_output: torch.Tensor,
220 - recv_lse: torch.Tensor,
221 - return_lse: bool = False,
222 - is_lse_base_on_e: bool = True,
223 -) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
224 - """
225 - Triton-accelerated LSE-weighted combination for DCP A2A.
319
227 - Args:
228 - recv_output: [N, B, H_local, D] - partial outputs from all KV shards
229 - recv_lse: [N, B, H_local] - partial LSEs from all KV shards
230 - return_lse: If True, also return the global LSE
231 - is_lse_base_on_e: If True, LSE is base e; if False, base 2
320 +def _dcp_a2a_pack_send(
321 + cp_attn_out: torch.Tensor,
322 + cp_attn_lse: torch.Tensor,
323 + send_buffer: torch.Tensor,
324 + world_size: int,
325 + h_per_rank: int,
326 + head_dim: int,
327 + lse_pack_dim: int,
328 +) -> None:
329 + grid = (cp_attn_out.shape[0], h_per_rank, 1)
330 + _dcp_a2a_pack_send_kernel[grid](
331 + cp_attn_out,
332 + cp_attn_lse,
333 + send_buffer,
334 + cp_attn_out.stride(0),
335 + cp_attn_out.stride(1),
336 + cp_attn_out.stride(2),
337 + cp_attn_lse.stride(0),
338 + cp_attn_lse.stride(1),
339 + send_buffer.stride(0),
340 + send_buffer.stride(1),
341 + send_buffer.stride(2),
342 + send_buffer.stride(3),
343 + N=world_size,
344 + HEAD_DIM=head_dim,
345 + H_PER_RANK=h_per_rank,
346 + LSE_PACK_DIM=lse_pack_dim,
347 + )
348
233 - Returns:
234 - Combined output [B, H_local, D]
235 - If return_lse=True, also returns global_lse [B, H_local]
236 - """
237 - N, B, H_local, D = recv_output.shape
349
350 +def _dcp_a2a_unpack_combine(
351 + recv_buffer: torch.Tensor,
352 + head_dim: int,
353 + lse_pack_dim: int,
354 + return_lse: bool,
355 + is_lse_base_on_e: bool,
356 +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
357 + world_size, num_tokens, h_per_rank, _ = recv_buffer.shape
358 out = torch.empty(
240 - (B, H_local, D), device=recv_output.device, dtype=recv_output.dtype
359 + (num_tokens, h_per_rank, head_dim),
360 + device=recv_buffer.device,
361 + dtype=recv_buffer.dtype,
362 )
242 -
243 - if return_lse:
244 - out_lse = torch.empty(
245 - (B, H_local), device=recv_lse.device, dtype=recv_lse.dtype
246 - )
247 - else:
248 - out_lse = torch.empty(1, device=recv_lse.device, dtype=recv_lse.dtype)
249 -
250 - ro_stride_N, ro_stride_B, ro_stride_H, ro_stride_D = recv_output.stride()
251 - rl_stride_N, rl_stride_B, rl_stride_H = recv_lse.stride()
252 - o_stride_B, o_stride_H, o_stride_D = out.stride()
253 -
254 - grid = (B, H_local, 1)
255 -
256 - _dcp_lse_combine_kernel[grid](
257 - recv_output,
258 - recv_lse,
363 + out_lse = torch.empty(
364 + (num_tokens, h_per_rank) if return_lse else (1, 1),
365 + device=recv_buffer.device,
366 + dtype=torch.float32 if return_lse else recv_buffer.dtype,
367 + )
368 + grid = (num_tokens, h_per_rank, 1)
369 + _dcp_a2a_unpack_combine_kernel[grid](
370 + recv_buffer,
371 out,
372 out_lse,
261 - ro_stride_N,
262 - ro_stride_B,
263 - ro_stride_H,
264 - ro_stride_D,
265 - rl_stride_N,
266 - rl_stride_B,
267 - rl_stride_H,
268 - o_stride_B,
269 - o_stride_H,
270 - o_stride_D,
271 - N=N,
272 - HEAD_DIM=D,
373 + recv_buffer.stride(0),
374 + recv_buffer.stride(1),
375 + recv_buffer.stride(2),
376 + recv_buffer.stride(3),
377 + out.stride(0),
378 + out.stride(1),
379 + out.stride(2),
380 + out_lse.stride(0),
381 + out_lse.stride(1),
382 + N=world_size,
383 + HEAD_DIM=head_dim,
384 IS_BASE_E=is_lse_base_on_e,
385 RETURN_LSE=return_lse,
386 + LSE_PACK_DIM=lse_pack_dim,
387 )
276 -
388 if return_lse:
389 return out, out_lse
390 return out
@@ -290,17 +401,8 @@ def dcp_a2a_lse_reduce(
401 """
402 Combine partial attention outputs across DCP ranks using All-to-All.
403
293 - Each rank holds attention output for all heads but only a local shard
294 - of the KV cache. This function:
295 - 1. Exchanges partial outputs across ranks via All-to-All
296 - 2. Exchanges LSE values via All-to-All
297 - 3. Combines them with exact LSE-weighted reduction (Triton kernel)
298 -
299 - Tensor flow:
300 - Input: cp_attn_out [B, H, D] - all heads, local KV shard
301 - Reshape: [N, B, H/N, D] - split heads across ranks
302 - A2A: Two all_to_all_single calls (output and LSE)
303 - Combine: recv [N, B, H/N, D] + lse [N, B, H/N] -> [B, H/N, D]
404 + The output and fp32 LSE are packed into a single output-dtype buffer, sent
405 + with one All-to-All, then unpacked and combined with exact LSE weighting.
406
407 Args:
408 cp_attn_out: [B, H, D] where B=num_tokens, H=total_heads, D=head_dim
@@ -321,43 +423,36 @@ def dcp_a2a_lse_reduce(
423 return cp_attn_out, cp_attn_lse
424 return cp_attn_out
425
324 - local_output = cp_attn_out.contiguous()
325 - local_lse = cp_attn_lse.contiguous()
326 -
327 - B, H, D = local_output.shape
426 + B, H, D = cp_attn_out.shape
427 + if H % world_size != 0:
428 + raise ValueError(f"H={H} must be divisible by DCP world size {world_size}.")
429 H_per_rank = H // world_size
430 + lse_pack_dim = _dcp_a2a_lse_pack_dim(cp_attn_out.dtype)
431
330 - # Reshape for All-to-All: [B, H, D] -> [N, B, H/N, D]
331 - # Split heads into N chunks, each destined for a different rank
332 - send_output = (
333 - local_output.view(B, world_size, H_per_rank, D).permute(1, 0, 2, 3).contiguous()
432 + send_buffer, recv_buffer = _dcp_a2a_send_recv_buffers(
433 + (world_size, B, H_per_rank, D + lse_pack_dim),
434 + device=cp_attn_out.device,
435 + dtype=cp_attn_out.dtype,
436 )
335 - recv_output = torch.empty_like(send_output)
336 -
337 - # Same for LSE: [B, H] -> [N, B, H/N]
338 - send_lse = local_lse.view(B, world_size, H_per_rank).permute(1, 0, 2).contiguous()
339 - recv_lse = torch.empty_like(send_lse)
437
341 - # All-to-All for partial attention outputs and LSE values (async overlap)
342 - work_output = dist.all_to_all_single(
343 - recv_output.view(-1),
344 - send_output.view(-1),
345 - group=cp_group.device_group,
346 - async_op=True,
438 + _dcp_a2a_pack_send(
439 + cp_attn_out,
440 + cp_attn_lse,
441 + send_buffer,
442 + world_size,
443 + H_per_rank,
444 + D,
445 + lse_pack_dim,
446 )
348 - work_lse = dist.all_to_all_single(
349 - recv_lse.view(-1),
350 - send_lse.view(-1),
447 +
448 + work = dist.all_to_all_single(
449 + recv_buffer.view(-1),
450 + send_buffer.view(-1),
451 group=cp_group.device_group,
452 async_op=True,
453 )
354 - work_output.wait()
355 - work_lse.wait()
356 -
357 - # LSE-weighted combination via Triton kernel (local, no communication)
358 - return dcp_lse_combine_triton(
359 - recv_output,
360 - recv_lse,
361 - return_lse=return_lse,
362 - is_lse_base_on_e=is_lse_base_on_e,
454 + work.wait()
455 +
456 + return _dcp_a2a_unpack_combine(
457 + recv_buffer, D, lse_pack_dim, return_lse, is_lse_base_on_e
458 )