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

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

Browse Source 
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Khoa Pham
2026-08-13 15:06:02 -07:00
committed by GitHub
co-authored by Claude Opus 5
parent 8bbca87780
commit 81fe452810
3 changed files with 106 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/dcp_kernels.py
+43 -26
View File
@@ -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)),
python/sglang/srt/layers/dcp/comm.py
+22 -12
View File
@@ -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
test/registered/kernels/test_dcp_lse_combine.py
+41 -1
View File
@@ -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