[DCP] Share one pack kernel between both a2a backends (#34651)

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

python/sglang/kernels/ops/attention/dcp_kernels.py

@@ -386,28 +386,34 @@ def _lse_pack_dim(output_dtype: torch.dtype) -> int:
 def _dcp_pack_a2a_send_kernel(
     out_ptr,
     lse_ptr,
-    send_ptr,
+    dst_o_ptr,
+    dst_lse_ptr,
     out_stride_B,
     out_stride_H,
     lse_stride_B,
     lse_stride_H,
-    send_stride_N,
+    dst_o_stride_N,
-    send_stride_B,
+    dst_o_stride_B,
-    send_stride_H,
+    dst_o_stride_H,
+    dst_lse_stride_N,
+    dst_lse_stride_B,
+    dst_lse_stride_H,
     H_PER_RANK: tl.constexpr,
     WORDS: tl.constexpr,
     BLOCK: tl.constexpr,
 ):
-    """Scatter one (batch, head) partial into its peer's a2a send slot."""
+    """Scatter one (batch, head) partial into its peer's send slot."""
     b = tl.program_id(0).to(tl.int64)
     h = tl.program_id(1).to(tl.int64)
+    peer = h // H_PER_RANK
+    h_local = h % H_PER_RANK
 
     src = out_ptr + b * out_stride_B + h * out_stride_H
     dst = (
-        send_ptr
+        dst_o_ptr
-        + (h // H_PER_RANK) * send_stride_N
+        + peer * dst_o_stride_N
-        + b * send_stride_B
+        + b * dst_o_stride_B
-        + (h % H_PER_RANK) * send_stride_H
+        + h_local * dst_o_stride_H
     )
 
     for start in tl.range(0, WORDS, BLOCK):
@@ -415,48 +421,59 @@ def _dcp_pack_a2a_send_kernel(
         mask = offs < WORDS
         tl.store(dst + offs, tl.load(src + offs, mask=mask), mask=mask)
 
-    tl.store(dst + WORDS, tl.load(lse_ptr + b * lse_stride_B + h * lse_stride_H))
+    tl.store(
+        dst_lse_ptr
+        + peer * dst_lse_stride_N
+        + b * dst_lse_stride_B
+        + h_local * dst_lse_stride_H,
+        tl.load(lse_ptr + b * lse_stride_B + h * lse_stride_H),
+    )
 
 
 def dcp_pack_a2a_send(
     cp_attn_out: torch.Tensor,
     cp_attn_lse: torch.Tensor,
-    send_combined: torch.Tensor,
+    dst_o: torch.Tensor,
+    dst_lse: torch.Tensor,
 ) -> None:
-    """Pack ``[B, H, D]`` partials + ``[B, H]`` LSE into the a2a send buffer.
+    """Scatter ``[B, H, D]`` partials + ``[B, H]`` LSE into a transport's send slots.
 
-    ``send_combined`` is ``[N, B_max, H // N, D + lse_pack_dim]``; rows beyond
+    ``dst_o`` is ``[N, B_max, H // N, D]`` and ``dst_lse`` ``[N, B_max, H // N]``,
-    ``B`` are left untouched.
+    in any stride order. Rows beyond ``B`` are left untouched.
     """
     B, H, D = cp_attn_out.shape
-    N, B_max, H_per_rank, row_width = send_combined.shape
+    N, B_max, H_per_rank = dst_lse.shape
-    lpd = _lse_pack_dim(send_combined.dtype)
+    lpd = _lse_pack_dim(cp_attn_out.dtype)
 
-    if cp_attn_lse.dtype != torch.float32:
+    if cp_attn_lse.dtype != torch.float32 or dst_lse.dtype != torch.float32:
-        raise ValueError(f"cp_attn_lse must be float32, got {cp_attn_lse.dtype}")
+        raise ValueError("LSE tensors must be float32")
     if D % lpd:
         raise ValueError(f"head dim {D} must be a multiple of the LSE pack dim {lpd}")
-    if row_width != D + lpd or H_per_rank * N != H or B > B_max:
+    if dst_o.shape != (N, B_max, H_per_rank, D) or H_per_rank * N != H or B > B_max:
         raise ValueError(
-            f"send buffer {tuple(send_combined.shape)} does not match "
+            f"destination {tuple(dst_o.shape)} / {tuple(dst_lse.shape)} does not "
-            f"out {tuple(cp_attn_out.shape)}"
+            f"match out {tuple(cp_attn_out.shape)}"
         )
 
     out_words = cp_attn_out.view(torch.float32)
-    send_words = send_combined.view(torch.float32)
+    dst_o_words = dst_o.view(torch.float32)
     words = D // lpd
 
     _dcp_pack_a2a_send_kernel[(B, H)](
         out_words,
         cp_attn_lse,
-        send_words,
+        dst_o_words,
+        dst_lse,
         out_words.stride(0),
         out_words.stride(1),
         cp_attn_lse.stride(0),
         cp_attn_lse.stride(1),
-        send_words.stride(0),
+        dst_o_words.stride(0),
-        send_words.stride(1),
+        dst_o_words.stride(1),
-        send_words.stride(2),
+        dst_o_words.stride(2),
+        dst_lse.stride(0),
+        dst_lse.stride(1),
+        dst_lse.stride(2),
         H_PER_RANK=H_per_rank,
         WORDS=words,
         BLOCK=min(1024, triton.next_power_of_2(words)),

python/sglang/srt/layers/dcp/comm.py

@@ -491,7 +491,13 @@ def dcp_a2a_lse_reduce(
         )
         recv_combined = torch.empty_like(send_combined)
 
-    dcp_pack_a2a_send(cp_attn_out, cp_attn_lse, send_combined)
+    send_words = send_combined.view(torch.float32)
+    dcp_pack_a2a_send(
+        cp_attn_out,
+        cp_attn_lse,
+        send_combined[:, :, :, :D],
+        send_words[:, :, :, D // lpd],
+    )
 
     # Transport as raw bytes (uint8): the output may be fp8 (fp8 KV cache),
     # which pynccl's dtype enum can't send; byte a2a is exact for equal chunks.
@@ -533,16 +539,21 @@ def _dcp_fi_a2a_lse_reduce(
     assert H % N == 0, f"num_heads ({H}) must be divisible by dcp_size ({N})"
     H_per_rank = H // N
 
-    # FlashInfer sends partial_o[..., peer, :] to `peer`; head h -> peer h//H_per_rank,
+    # Note(kpham-sgl): empty(), not zeros() -- the pack below fills partial_o and
-    # so the peer axis is the outer head split: [B,N,H_pr,D] -> [B,H_pr,N,D].
+    # stats slot 0, and slot 1 is never read by anyone. The a2a moves the stats
-    partial_o = cp_attn_out.view(B, N, H_per_rank, D).permute(0, 2, 1, 3).contiguous()
+    # field as opaque bytes and we only ever read slot 0 back off the wire.
-    # softmax_stats: fp32 [B, H_per_rank, N, S=2] (FI requires S>=2 & even);
+    partial_o = torch.empty(
-    # carry the LSE in lane 0, lane 1 is ignored by the combine.
+        B, H_per_rank, N, D, dtype=cp_attn_out.dtype, device=cp_attn_out.device
-    lse_view = cp_attn_lse.view(B, N, H_per_rank).permute(0, 2, 1)  # [B,H_pr,N]
+    )
-    softmax_stats = torch.zeros(
+    softmax_stats = torch.empty(
         B, H_per_rank, N, 2, dtype=torch.float32, device=cp_attn_out.device
     )
-    softmax_stats[..., 0] = lse_view
+    dcp_pack_a2a_send(
+        cp_attn_out,
+        cp_attn_lse,
+        partial_o.permute(2, 0, 1, 3),
+        softmax_stats[..., 0].permute(2, 0, 1),
+    )
 
     o_out, stats_out = decode_cp_a2a_alltoall(
         partial_o,
@@ -552,9 +563,8 @@ def _dcp_fi_a2a_lse_reduce(
         N,
     )
 
-    # o_out[b,hpr,src] = rank src's partial for local head hpr -> combine layout.
+    recv_output = o_out.permute(2, 0, 1, 3)
-    recv_output = o_out.permute(2, 0, 1, 3).contiguous()  # [N, B, H_per_rank, D]
+    recv_lse = stats_out[..., 0].permute(2, 0, 1)
-    recv_lse = stats_out[..., 0].permute(2, 0, 1).contiguous()  # [N, B, H_per_rank]
 
     combined, _ = dcp_lse_combine_triton(
         recv_output, recv_lse, is_lse_base_on_e=is_lse_base_on_e

test/registered/kernels/test_dcp_lse_combine.py

@@ -473,7 +473,12 @@ class TestDCPA2AReduceWithCUDAGraphBuffers(CustomTestCase):
                 got = torch.zeros(
                     N, max_bs, H_per_rank, D + lpd, dtype=dtype, device=self.device
                 )
-                dcp_pack_a2a_send(out, lse, got)
+                dcp_pack_a2a_send(
+                    out,
+                    lse,
+                    got[:, :, :, :D],
+                    got.view(torch.float32)[:, :, :, D // lpd],
+                )
 
                 want = torch.zeros_like(got)
                 want[:, :B, :, :D] = out.view(B, N, H_per_rank, D).permute(1, 0, 2, 3)
@@ -493,6 +498,41 @@ class TestDCPA2AReduceWithCUDAGraphBuffers(CustomTestCase):
                     torch.equal(lane, lse.view(B, N, H_per_rank).permute(1, 0, 2))
                 )
 
+    def test_pack_serves_the_split_peer_inside_layout(self):
+        from sglang.kernels.ops.attention.dcp_kernels import dcp_pack_a2a_send
+
+        for N, B, H_per_rank, D in ((2, 4, 8, 128), (4, 1, 16, 512)):
+            with self.subTest(N=N, B=B, H_per_rank=H_per_rank, D=D):
+                H = H_per_rank * N
+                out = torch.randn(B, H, D, device=self.device, dtype=torch.bfloat16)
+                lse = torch.randn(B, H, device=self.device, dtype=torch.float32)
+
+                partial_o = torch.empty(
+                    B, H_per_rank, N, D, dtype=torch.bfloat16, device=self.device
+                )
+                stats = torch.zeros(
+                    B, H_per_rank, N, 2, dtype=torch.float32, device=self.device
+                )
+                dcp_pack_a2a_send(
+                    out,
+                    lse,
+                    partial_o.permute(2, 0, 1, 3),
+                    stats[..., 0].permute(2, 0, 1),
+                )
+
+                want_o = out.view(B, N, H_per_rank, D).permute(0, 2, 1, 3)
+                want_lse = lse.view(B, N, H_per_rank).permute(0, 2, 1)
+                self.assertTrue(
+                    torch.equal(
+                        partial_o.view(torch.uint8),
+                        want_o.contiguous().view(torch.uint8),
+                    )
+                )
+                self.assertTrue(torch.equal(stats[..., 0], want_lse))
+                self.assertTrue(
+                    torch.equal(stats[..., 1], torch.zeros_like(stats[..., 1]))
+                )
+
     def test_buffers_have_fixed_data_ptrs(self):
         """Pre-allocated buffer data_ptr must not change -- required for graph replay."""
         from sglang.srt.layers.dcp import dcp_a2a_lse_reduce