minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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
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