From 6c6294b7be23f751683f63b2c52775ec950b9a30 Mon Sep 17 00:00:00 2001 From: Khoa Pham Date: Wed, 12 Aug 2026 16:04:46 -0700 Subject: [PATCH] [DCP] Fuse the a2a pack/unpack copies in the MLA LSE reduce (#34614) Co-authored-by: Claude Opus 5 (1M context) --- .../sglang/kernels/ops/attention/__init__.py | 2 + .../kernels/ops/attention/dcp_kernels.py | 81 +++++++++++++++++++ python/sglang/srt/layers/dcp/comm.py | 56 +++---------- .../kernels/test_dcp_lse_combine.py | 41 ++++++++++ 4 files changed, 134 insertions(+), 46 deletions(-) diff --git a/python/sglang/kernels/ops/attention/__init__.py b/python/sglang/kernels/ops/attention/__init__.py index 4d77dc1bc..4030ea3bc 100644 --- a/python/sglang/kernels/ops/attention/__init__.py +++ b/python/sglang/kernels/ops/attention/__init__.py @@ -116,6 +116,8 @@ for _mod, _fn in [ ("flash_mla_sm120", "flash_mla_with_kvcache_sm120"), ("dcp_kernels", "create_dcp_kv_indices"), ("dcp_kernels", "correct_attn_out"), + ("dcp_kernels", "dcp_lse_combine_triton"), + ("dcp_kernels", "dcp_pack_a2a_send"), ("pa_page_table", "_build_pa_page_table"), ("nsa_triton_decode", "triton_sparse_attn_decode"), ]: diff --git a/python/sglang/kernels/ops/attention/dcp_kernels.py b/python/sglang/kernels/ops/attention/dcp_kernels.py index f65c1777c..771f3ad96 100644 --- a/python/sglang/kernels/ops/attention/dcp_kernels.py +++ b/python/sglang/kernels/ops/attention/dcp_kernels.py @@ -382,6 +382,87 @@ def _lse_pack_dim(output_dtype: torch.dtype) -> int: return torch.finfo(torch.float32).bits // torch.finfo(output_dtype).bits +@triton.jit +def _dcp_pack_a2a_send_kernel( + out_ptr, + lse_ptr, + send_ptr, + out_stride_B, + out_stride_H, + lse_stride_B, + lse_stride_H, + send_stride_N, + send_stride_B, + send_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.""" + b = tl.program_id(0).to(tl.int64) + h = tl.program_id(1).to(tl.int64) + + 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 + ) + + for start in tl.range(0, WORDS, BLOCK): + offs = start + tl.arange(0, BLOCK) + 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)) + + +def dcp_pack_a2a_send( + cp_attn_out: torch.Tensor, + cp_attn_lse: torch.Tensor, + send_combined: torch.Tensor, +) -> None: + """Pack ``[B, H, D]`` partials + ``[B, H]`` LSE into the a2a send buffer. + + ``send_combined`` is ``[N, B_max, H // N, D + lse_pack_dim]``; 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) + + if cp_attn_lse.dtype != torch.float32: + raise ValueError(f"cp_attn_lse must be float32, got {cp_attn_lse.dtype}") + 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: + raise ValueError( + f"send buffer {tuple(send_combined.shape)} does not match " + f"out {tuple(cp_attn_out.shape)}" + ) + + out_words = cp_attn_out.view(torch.float32) + send_words = send_combined.view(torch.float32) + words = D // lpd + + _dcp_pack_a2a_send_kernel[(B, H)]( + out_words, + cp_attn_lse, + send_words, + 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), + H_PER_RANK=H_per_rank, + WORDS=words, + BLOCK=min(1024, triton.next_power_of_2(words)), + ) + + @triton.jit def _dcp_lse_combine_kernel( recv_output_ptr, diff --git a/python/sglang/srt/layers/dcp/comm.py b/python/sglang/srt/layers/dcp/comm.py index fc21ea922..279c615fb 100644 --- a/python/sglang/srt/layers/dcp/comm.py +++ b/python/sglang/srt/layers/dcp/comm.py @@ -30,6 +30,7 @@ from sglang.kernels.ops.attention.dcp_kernels import ( _lse_pack_dim, correct_attn_out, dcp_lse_combine_triton, + dcp_pack_a2a_send, ) from sglang.srt.distributed.device_communicators.pynccl_allocator import ( use_symmetric_memory, @@ -476,34 +477,10 @@ def dcp_a2a_lse_reduce( out_dtype = cp_attn_out.dtype lpd = _lse_pack_dim(out_dtype) # 2 for bf16/fp16 - # Reshape [B, H, D] -> [N, B, H/N, D] — split heads across ranks - reshaped_out = cp_attn_out.view(B, N, H_per_rank, D).permute(1, 0, 2, 3) - reshaped_lse = cp_attn_lse.view(B, N, H_per_rank).permute(1, 0, 2) - if cuda_graph_buffers is not None: - # CUDA graph path with pre-allocated fused buffers. send_combined = cuda_graph_buffers["send_combined"] recv_combined = cuda_graph_buffers["recv_combined"] - send_lse_stg = cuda_graph_buffers["send_lse"] - recv_lse_stg = cuda_graph_buffers["recv_lse"] - - send_combined[:, :B, :, :D].copy_(reshaped_out) - send_lse_stg[:, :B, :].copy_(reshaped_lse) - send_combined[:, :, :, D:].copy_( - send_lse_stg.view(out_dtype).view(N, -1, H_per_rank, lpd) - ) - - cp_group.all_to_all_single( - recv_combined.reshape(-1).view(torch.uint8), - send_combined.reshape(-1).view(torch.uint8), - ) - recv_output = recv_combined[:, :B, :, :D] - recv_lse_stg.view(out_dtype).view(N, -1, H_per_rank, lpd).copy_( - recv_combined[:, :, :, D:] - ) - recv_lse = recv_lse_stg[:, :B, :] else: - send_lse_contig = reshaped_lse.contiguous() # [N, B, H_per_rank] fp32 send_combined = torch.empty( N, B, @@ -514,30 +491,17 @@ def dcp_a2a_lse_reduce( ) recv_combined = torch.empty_like(send_combined) - send_combined[:, :, :, :D].copy_(reshaped_out) - send_combined[:, :, :, D:].copy_( - send_lse_contig.view(out_dtype).view(N, B, H_per_rank, lpd) - ) + dcp_pack_a2a_send(cp_attn_out, cp_attn_lse, send_combined) - # 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. - cp_group.all_to_all_single( - recv_combined.reshape(-1).view(torch.uint8), - send_combined.reshape(-1).view(torch.uint8), - ) + # 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. + cp_group.all_to_all_single( + recv_combined.reshape(-1).view(torch.uint8), + send_combined.reshape(-1).view(torch.uint8), + ) - recv_output = recv_combined[:, :, :, :D] - recv_lse_stg = torch.empty( - N, - B, - H_per_rank, - dtype=torch.float32, - device=cp_attn_out.device, - ) - recv_lse_stg.view(out_dtype).view(N, B, H_per_rank, lpd).copy_( - recv_combined[:, :, :, D:] - ) - recv_lse = recv_lse_stg + recv_output = recv_combined[:, :B, :, :D] + recv_lse = recv_combined.view(torch.float32)[:, :B, :, D // lpd] 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 0249b8a1e..b05102b86 100644 --- a/test/registered/kernels/test_dcp_lse_combine.py +++ b/test/registered/kernels/test_dcp_lse_combine.py @@ -452,6 +452,47 @@ class TestDCPA2AReduceWithCUDAGraphBuffers(CustomTestCase): self.assertEqual(result.shape, (B, H_per_rank, D)) self.assertFalse(torch.isnan(result).any()) + def test_pack_matches_the_copy_formulation_it_replaces(self): + from sglang.kernels.ops.attention.dcp_kernels import ( + _lse_pack_dim, + dcp_pack_a2a_send, + ) + + for N, B, H_per_rank, D, dtype in ( + (2, 4, 8, 128, torch.bfloat16), + (4, 3, 2, 64, torch.float16), + (8, 1, 12, 512, torch.bfloat16), + ): + with self.subTest(N=N, B=B, H_per_rank=H_per_rank, D=D, dtype=dtype): + H = H_per_rank * N + lpd = _lse_pack_dim(dtype) + max_bs = B + 5 + out = torch.randn(B, H, D, device=self.device, dtype=dtype) + lse = torch.randn(B, H, device=self.device, dtype=torch.float32) + + got = torch.zeros( + N, max_bs, H_per_rank, D + lpd, dtype=dtype, device=self.device + ) + dcp_pack_a2a_send(out, lse, got) + + want = torch.zeros_like(got) + want[:, :B, :, :D] = out.view(B, N, H_per_rank, D).permute(1, 0, 2, 3) + want[:, :B, :, D:] = ( + lse.view(B, N, H_per_rank) + .permute(1, 0, 2) + .contiguous() + .view(dtype) + .view(N, B, H_per_rank, lpd) + ) + self.assertTrue( + torch.equal(got.view(torch.uint8), want.view(torch.uint8)) + ) + + lane = got.view(torch.float32)[:, :B, :, D // lpd] + self.assertTrue( + torch.equal(lane, lse.view(B, N, H_per_rank).permute(1, 0, 2)) + ) + 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