minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

[DCP] Fuse the a2a pack/unpack copies in the MLA LSE reduce (#34614)

Browse Source 
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Khoa Pham
2026-08-12 16:04:46 -07:00
committed by GitHub
co-authored by Claude Opus 5
parent c05eb856f7
commit 6c6294b7be
4 changed files with 134 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/__init__.py
+2
View File
@@ -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"),
]:
python/sglang/kernels/ops/attention/dcp_kernels.py
+81
View File
@@ -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,
python/sglang/srt/layers/dcp/comm.py
+4 -40
View File
@@ -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,10 +491,7 @@ 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.
@@ -526,18 +500,8 @@ def dcp_a2a_lse_reduce(
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
test/registered/kernels/test_dcp_lse_combine.py
+41
View File
@@ -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