From 81fe452810d9256b5d307cad3a3e261c70170591 Mon Sep 17 00:00:00 2001 From: Khoa Pham Date: Thu, 13 Aug 2026 15:06:02 -0700 Subject: [PATCH] [DCP] Share one pack kernel between both a2a backends (#34651) Co-authored-by: Claude Opus 5 (1M context) --- .../kernels/ops/attention/dcp_kernels.py | 69 ++++++++++++------- python/sglang/srt/layers/dcp/comm.py | 34 +++++---- .../kernels/test_dcp_lse_combine.py | 42 ++++++++++- 3 files changed, 106 insertions(+), 39 deletions(-) diff --git a/python/sglang/kernels/ops/attention/dcp_kernels.py b/python/sglang/kernels/ops/attention/dcp_kernels.py index 771f3ad96..caf7bc75c 100644 --- a/python/sglang/kernels/ops/attention/dcp_kernels.py +++ b/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, - send_stride_B, - send_stride_H, + dst_o_stride_N, + dst_o_stride_B, + 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 - + (h // H_PER_RANK) * send_stride_N - + b * send_stride_B - + (h % H_PER_RANK) * send_stride_H + dst_o_ptr + + peer * dst_o_stride_N + + b * dst_o_stride_B + + 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 - ``B`` are left untouched. + ``dst_o`` is ``[N, B_max, H // N, D]`` and ``dst_lse`` ``[N, B_max, H // N]``, + 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 - lpd = _lse_pack_dim(send_combined.dtype) + N, B_max, H_per_rank = dst_lse.shape + lpd = _lse_pack_dim(cp_attn_out.dtype) - if cp_attn_lse.dtype != torch.float32: - raise ValueError(f"cp_attn_lse must be float32, got {cp_attn_lse.dtype}") + if cp_attn_lse.dtype != torch.float32 or dst_lse.dtype != torch.float32: + 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"out {tuple(cp_attn_out.shape)}" + f"destination {tuple(dst_o.shape)} / {tuple(dst_lse.shape)} does not " + 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), - send_words.stride(1), - send_words.stride(2), + dst_o_words.stride(0), + dst_o_words.stride(1), + 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)), diff --git a/python/sglang/srt/layers/dcp/comm.py b/python/sglang/srt/layers/dcp/comm.py index 279c615fb..915a29b7a 100644 --- a/python/sglang/srt/layers/dcp/comm.py +++ b/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, - # so the peer axis is the outer head split: [B,N,H_pr,D] -> [B,H_pr,N,D]. - partial_o = cp_attn_out.view(B, N, H_per_rank, D).permute(0, 2, 1, 3).contiguous() - # softmax_stats: fp32 [B, H_per_rank, N, S=2] (FI requires S>=2 & even); - # carry the LSE in lane 0, lane 1 is ignored by the combine. - lse_view = cp_attn_lse.view(B, N, H_per_rank).permute(0, 2, 1) # [B,H_pr,N] - softmax_stats = torch.zeros( + # Note(kpham-sgl): empty(), not zeros() -- the pack below fills partial_o and + # stats slot 0, and slot 1 is never read by anyone. The a2a moves the stats + # field as opaque bytes and we only ever read slot 0 back off the wire. + partial_o = torch.empty( + B, H_per_rank, N, D, dtype=cp_attn_out.dtype, device=cp_attn_out.device + ) + 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).contiguous() # [N, B, H_per_rank, D] - recv_lse = stats_out[..., 0].permute(2, 0, 1).contiguous() # [N, B, H_per_rank] + recv_output = o_out.permute(2, 0, 1, 3) + recv_lse = stats_out[..., 0].permute(2, 0, 1) combined, _ = dcp_lse_combine_triton( recv_output, recv_lse, is_lse_base_on_e=is_lse_base_on_e diff --git a/test/registered/kernels/test_dcp_lse_combine.py b/test/registered/kernels/test_dcp_lse_combine.py index b05102b86..283986138 100644 --- a/test/registered/kernels/test_dcp_lse_combine.py +++ b/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