467 lines
16 KiB
Python
467 lines
16 KiB
Python
"""Tests for DCP LSE combine kernels.
|
|
|
|
Covers:
|
|
1. Triton LSE combine kernel correctness vs CPU reference (base-e and base-2)
|
|
2. Various DCP world sizes (N=1,2,4,8)
|
|
3. Edge cases: single shard, dominant LSE, equal LSE
|
|
4. return_lse mode
|
|
5. dcp_a2a_lse_reduce with pre-allocated CUDA graph buffers
|
|
"""
|
|
|
|
import unittest
|
|
from unittest.mock import MagicMock
|
|
|
|
import torch
|
|
|
|
from sglang.test.ci.ci_register import register_cuda_ci
|
|
from sglang.test.test_utils import CustomTestCase
|
|
|
|
register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-large")
|
|
|
|
|
|
class TestLSECombineTritonVsCPU(CustomTestCase):
|
|
"""Test Triton LSE combine kernel against CPU reference."""
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
if not torch.cuda.is_available():
|
|
raise unittest.SkipTest("CUDA required for Triton kernel tests")
|
|
cls.device = "cuda"
|
|
|
|
def _run_combine_test(
|
|
self, N, B, H_local, D, is_base_e, dtype=torch.bfloat16, atol=1e-2
|
|
):
|
|
from sglang.kernels.ops.attention.dcp_kernels import (
|
|
_lse_weighted_combine_cpu,
|
|
dcp_lse_combine_triton,
|
|
)
|
|
|
|
torch.manual_seed(42)
|
|
|
|
partial_outputs = torch.randn(N, B, H_local, D, device=self.device, dtype=dtype)
|
|
if is_base_e:
|
|
partial_lses = torch.randn(
|
|
N, B, H_local, device=self.device, dtype=torch.float32
|
|
)
|
|
else:
|
|
partial_lses = (
|
|
torch.randn(N, B, H_local, device=self.device, dtype=torch.float32)
|
|
* 5.0
|
|
)
|
|
|
|
cpu_result = _lse_weighted_combine_cpu(
|
|
partial_outputs.cpu(),
|
|
partial_lses.cpu(),
|
|
is_lse_base_on_e=is_base_e,
|
|
)
|
|
|
|
triton_result, _ = dcp_lse_combine_triton(
|
|
partial_outputs,
|
|
partial_lses,
|
|
is_lse_base_on_e=is_base_e,
|
|
return_lse=False,
|
|
)
|
|
|
|
torch.testing.assert_close(
|
|
triton_result.float().cpu(),
|
|
cpu_result.float(),
|
|
atol=atol,
|
|
rtol=1e-2,
|
|
)
|
|
|
|
def test_n2_base_e(self):
|
|
self._run_combine_test(N=2, B=4, H_local=8, D=64, is_base_e=True)
|
|
|
|
def test_n2_base_2(self):
|
|
self._run_combine_test(N=2, B=4, H_local=8, D=64, is_base_e=False)
|
|
|
|
def test_n4_base_e(self):
|
|
self._run_combine_test(N=4, B=8, H_local=16, D=128, is_base_e=True)
|
|
|
|
def test_n4_base_2(self):
|
|
self._run_combine_test(N=4, B=8, H_local=16, D=128, is_base_e=False)
|
|
|
|
def test_n8_base_e(self):
|
|
self._run_combine_test(N=8, B=4, H_local=8, D=128, is_base_e=True)
|
|
|
|
def test_n8_base_2(self):
|
|
self._run_combine_test(N=8, B=4, H_local=8, D=512, is_base_e=False)
|
|
|
|
def test_n2_large_batch(self):
|
|
self._run_combine_test(N=2, B=64, H_local=16, D=128, is_base_e=False)
|
|
|
|
def test_n4_large_head_dim(self):
|
|
self._run_combine_test(N=4, B=8, H_local=8, D=512, is_base_e=True)
|
|
|
|
def test_flashmla_natural_log_lse_correction(self):
|
|
"""Natural-log LSEs log(2), log(8) give an 8/10 local weight."""
|
|
from sglang.kernels.ops.attention.dcp_kernels import correct_attn_out
|
|
|
|
local_output = torch.tensor(
|
|
[[[10.0]]], device=self.device, dtype=torch.bfloat16
|
|
)
|
|
lses = torch.log(torch.tensor([[[2.0]], [[8.0]]], device=self.device))
|
|
corrected, _ = correct_attn_out(
|
|
local_output,
|
|
lses,
|
|
cp_rank=1,
|
|
ctx=None,
|
|
new_output=torch.empty((1, 1, 1), device=self.device),
|
|
is_lse_base_on_e=True,
|
|
)
|
|
|
|
torch.testing.assert_close(
|
|
corrected.cpu(), torch.tensor([[[8.0]]]), atol=1e-5, rtol=1e-5
|
|
)
|
|
|
|
|
|
class TestLSECombineSingleShard(CustomTestCase):
|
|
"""N=1 should return input unchanged."""
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
if not torch.cuda.is_available():
|
|
raise unittest.SkipTest("CUDA required")
|
|
cls.device = "cuda"
|
|
|
|
def test_single_shard(self):
|
|
from sglang.kernels.ops.attention.dcp_kernels import dcp_lse_combine_triton
|
|
|
|
N, B, H_local, D = 1, 4, 8, 64
|
|
partial_outputs = torch.randn(
|
|
N, B, H_local, D, device=self.device, dtype=torch.bfloat16
|
|
)
|
|
partial_lses = torch.randn(
|
|
N, B, H_local, device=self.device, dtype=torch.float32
|
|
)
|
|
|
|
triton_result, _ = dcp_lse_combine_triton(
|
|
partial_outputs, partial_lses, is_lse_base_on_e=True
|
|
)
|
|
|
|
torch.testing.assert_close(
|
|
triton_result.float().cpu(),
|
|
partial_outputs.squeeze(0).float().cpu(),
|
|
atol=1e-3,
|
|
rtol=1e-3,
|
|
)
|
|
|
|
|
|
class TestLSECombineReturnLSE(CustomTestCase):
|
|
"""Verify return_lse=True produces valid global LSE."""
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
if not torch.cuda.is_available():
|
|
raise unittest.SkipTest("CUDA required")
|
|
cls.device = "cuda"
|
|
|
|
def test_return_lse(self):
|
|
from sglang.kernels.ops.attention.dcp_kernels import dcp_lse_combine_triton
|
|
|
|
N, B, H_local, D = 2, 4, 8, 64
|
|
partial_outputs = torch.randn(
|
|
N, B, H_local, D, device=self.device, dtype=torch.bfloat16
|
|
)
|
|
partial_lses = torch.randn(
|
|
N, B, H_local, device=self.device, dtype=torch.float32
|
|
)
|
|
|
|
triton_result, triton_lse = dcp_lse_combine_triton(
|
|
partial_outputs, partial_lses, is_lse_base_on_e=True, return_lse=True
|
|
)
|
|
|
|
self.assertIsNotNone(triton_lse)
|
|
self.assertEqual(triton_lse.shape, (B, H_local))
|
|
self.assertFalse(torch.isnan(triton_lse).any())
|
|
|
|
|
|
class TestLSECombineEdgeCases(CustomTestCase):
|
|
"""Test edge cases for LSE combine."""
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
if not torch.cuda.is_available():
|
|
raise unittest.SkipTest("CUDA required")
|
|
cls.device = "cuda"
|
|
|
|
def test_one_shard_dominant(self):
|
|
"""One shard has much larger LSE -- output should be close to that shard."""
|
|
from sglang.kernels.ops.attention.dcp_kernels import (
|
|
_lse_weighted_combine_cpu,
|
|
dcp_lse_combine_triton,
|
|
)
|
|
|
|
N, B, H_local, D = 2, 1, 1, 64
|
|
partial_outputs = torch.randn(
|
|
N, B, H_local, D, device=self.device, dtype=torch.bfloat16
|
|
)
|
|
partial_lses = torch.tensor(
|
|
[[[100.0]], [[-100.0]]], device=self.device, dtype=torch.float32
|
|
)
|
|
|
|
triton_result, _ = dcp_lse_combine_triton(
|
|
partial_outputs, partial_lses, is_lse_base_on_e=True
|
|
)
|
|
cpu_result = _lse_weighted_combine_cpu(
|
|
partial_outputs.cpu(), partial_lses.cpu(), is_lse_base_on_e=True
|
|
)
|
|
|
|
torch.testing.assert_close(
|
|
triton_result.float().cpu(), cpu_result.float(), atol=1e-2, rtol=1e-2
|
|
)
|
|
torch.testing.assert_close(
|
|
triton_result.float().cpu(),
|
|
partial_outputs[0].float().cpu(),
|
|
atol=1e-2,
|
|
rtol=1e-2,
|
|
)
|
|
|
|
def test_equal_lse(self):
|
|
"""Equal LSE across shards -- output should be mean of outputs."""
|
|
from sglang.kernels.ops.attention.dcp_kernels import dcp_lse_combine_triton
|
|
|
|
N, B, H_local, D = 2, 1, 1, 64
|
|
partial_outputs = torch.randn(
|
|
N, B, H_local, D, device=self.device, dtype=torch.bfloat16
|
|
)
|
|
partial_lses = torch.tensor(
|
|
[[[5.0]], [[5.0]]], device=self.device, dtype=torch.float32
|
|
)
|
|
|
|
triton_result, _ = dcp_lse_combine_triton(
|
|
partial_outputs, partial_lses, is_lse_base_on_e=True
|
|
)
|
|
expected = partial_outputs.float().mean(dim=0)
|
|
|
|
torch.testing.assert_close(
|
|
triton_result.float().cpu(), expected.cpu(), atol=1e-2, rtol=1e-2
|
|
)
|
|
|
|
|
|
class TestLSEBaseByBackend(CustomTestCase):
|
|
"""Which attention backends report LSE in natural log."""
|
|
|
|
def test_natural_log_lse_backends(self):
|
|
from sglang.srt.models.deepseek_common.attention_forward_methods.forward_mla import (
|
|
is_mla_dcp_lse_base_on_e,
|
|
)
|
|
|
|
self.assertTrue(is_mla_dcp_lse_base_on_e("flashmla"))
|
|
self.assertTrue(is_mla_dcp_lse_base_on_e("cutedsl_mla"))
|
|
self.assertFalse(is_mla_dcp_lse_base_on_e("flashinfer_mla"))
|
|
self.assertFalse(is_mla_dcp_lse_base_on_e("tokenspeed_mla"))
|
|
self.assertFalse(is_mla_dcp_lse_base_on_e("trtllm_mla"))
|
|
self.assertFalse(is_mla_dcp_lse_base_on_e(None))
|
|
|
|
|
|
class TestDCPA2AReduceWithCUDAGraphBuffers(CustomTestCase):
|
|
"""Test dcp_a2a_lse_reduce with pre-allocated CUDA graph buffers."""
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
if not torch.cuda.is_available():
|
|
raise unittest.SkipTest("CUDA required")
|
|
cls.device = "cuda"
|
|
|
|
def _make_mock_group(self, world_size):
|
|
group = MagicMock()
|
|
group.world_size = world_size
|
|
|
|
def identity_a2a(output, input_):
|
|
output.copy_(input_)
|
|
|
|
group.all_to_all_single = MagicMock(side_effect=identity_a2a)
|
|
return group
|
|
|
|
def _make_cuda_graph_buffers(self, N, max_bs, H_per_rank, D, lpd=2):
|
|
"""Create fused CUDA graph buffers matching dcp_a2a_lse_reduce API."""
|
|
return {
|
|
"send_combined": torch.empty(
|
|
N, max_bs, H_per_rank, D + lpd, dtype=torch.bfloat16, device=self.device
|
|
),
|
|
"recv_combined": torch.empty(
|
|
N, max_bs, H_per_rank, D + lpd, dtype=torch.bfloat16, device=self.device
|
|
),
|
|
"send_lse": torch.empty(
|
|
N, max_bs, H_per_rank, dtype=torch.float32, device=self.device
|
|
),
|
|
"recv_lse": torch.empty(
|
|
N, max_bs, H_per_rank, dtype=torch.float32, device=self.device
|
|
),
|
|
}
|
|
|
|
def test_cuda_graph_buffers_same_as_dynamic(self):
|
|
from sglang.srt.layers.dcp import dcp_a2a_lse_reduce
|
|
|
|
torch.manual_seed(123)
|
|
N, B, H_per_rank, D = 2, 4, 8, 128
|
|
H = H_per_rank * N
|
|
max_bs = 16
|
|
|
|
group = self._make_mock_group(N)
|
|
|
|
attn_out = torch.randn(B, H, D, device=self.device, dtype=torch.bfloat16)
|
|
attn_lse = torch.randn(B, H, device=self.device, dtype=torch.float32)
|
|
|
|
result_dynamic = dcp_a2a_lse_reduce(
|
|
attn_out.clone(), attn_lse.clone(), group, is_lse_base_on_e=True
|
|
)
|
|
|
|
cuda_graph_buffers = self._make_cuda_graph_buffers(N, max_bs, H_per_rank, D)
|
|
|
|
result_graph = dcp_a2a_lse_reduce(
|
|
attn_out.clone(),
|
|
attn_lse.clone(),
|
|
group,
|
|
is_lse_base_on_e=True,
|
|
cuda_graph_buffers=cuda_graph_buffers,
|
|
)
|
|
|
|
torch.testing.assert_close(
|
|
result_graph.float().cpu(),
|
|
result_dynamic.float().cpu(),
|
|
atol=1e-5,
|
|
rtol=1e-5,
|
|
)
|
|
|
|
def test_cuda_graph_buffers_partial_batch(self):
|
|
"""Buffer max_bs > actual B -- should correctly slice."""
|
|
from sglang.srt.layers.dcp import dcp_a2a_lse_reduce
|
|
|
|
torch.manual_seed(789)
|
|
N, B, H_per_rank, D = 2, 3, 8, 128
|
|
H = H_per_rank * N
|
|
max_bs = 32
|
|
|
|
group = self._make_mock_group(N)
|
|
|
|
attn_out = torch.randn(B, H, D, device=self.device, dtype=torch.bfloat16)
|
|
attn_lse = torch.randn(B, H, device=self.device, dtype=torch.float32)
|
|
|
|
cuda_graph_buffers = self._make_cuda_graph_buffers(N, max_bs, H_per_rank, D)
|
|
|
|
result = dcp_a2a_lse_reduce(
|
|
attn_out,
|
|
attn_lse,
|
|
group,
|
|
is_lse_base_on_e=True,
|
|
cuda_graph_buffers=cuda_graph_buffers,
|
|
)
|
|
|
|
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[:, :, :, :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)
|
|
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_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
|
|
|
|
N, B, H_per_rank, D = 2, 4, 8, 64
|
|
H = H_per_rank * N
|
|
max_bs = 16
|
|
|
|
group = self._make_mock_group(N)
|
|
|
|
buffers = self._make_cuda_graph_buffers(N, max_bs, H_per_rank, D)
|
|
send_ptr = buffers["send_combined"].data_ptr()
|
|
recv_ptr = buffers["recv_combined"].data_ptr()
|
|
|
|
attn_out = torch.randn(B, H, D, device=self.device, dtype=torch.bfloat16)
|
|
attn_lse = torch.randn(B, H, device=self.device, dtype=torch.float32)
|
|
|
|
dcp_a2a_lse_reduce(
|
|
attn_out,
|
|
attn_lse,
|
|
group,
|
|
is_lse_base_on_e=True,
|
|
cuda_graph_buffers=buffers,
|
|
)
|
|
|
|
self.assertEqual(buffers["send_combined"].data_ptr(), send_ptr)
|
|
self.assertEqual(buffers["recv_combined"].data_ptr(), recv_ptr)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|