[CI] Move existing unit tests into unit directory (#20631)
This commit is contained in:
Ke Bao
@@ -0,0 +1,367 @@
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=10, suite="stage-b-test-small-1-gpu")
|
||||
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from sglang.srt.layers.conv import Conv2dLayer, Conv3dLayer
|
||||
|
||||
|
||||
def _copy_weights(src, dst_nn):
|
||||
"""Copy weights from Conv*dLayer to nn.Conv*d for comparison."""
|
||||
with torch.no_grad():
|
||||
dst_nn.weight.copy_(src.weight)
|
||||
if src.bias is not None:
|
||||
dst_nn.bias.copy_(src.bias)
|
||||
|
||||
|
||||
class TestConv2dLayer(unittest.TestCase):
|
||||
|
||||
def test_basic_patch_embedding(self):
|
||||
layer = Conv2dLayer(3, 768, kernel_size=14, stride=14, bias=False)
|
||||
ref = nn.Conv2d(3, 768, kernel_size=14, stride=14, bias=False)
|
||||
self.assertFalse(layer.enable_linear)
|
||||
_copy_weights(layer, ref)
|
||||
x = torch.randn(2, 3, 224, 224)
|
||||
with torch.no_grad():
|
||||
torch.testing.assert_close(
|
||||
layer.forward_native(x), ref(x), rtol=1e-4, atol=1e-4
|
||||
)
|
||||
|
||||
def test_enable_linear(self):
|
||||
layer = Conv2dLayer(
|
||||
3, 768, kernel_size=14, stride=14, bias=True, disable_linear=False
|
||||
)
|
||||
ref = nn.Conv2d(3, 768, kernel_size=14, stride=14, bias=True)
|
||||
self.assertTrue(layer.enable_linear)
|
||||
_copy_weights(layer, ref)
|
||||
x = torch.randn(1, 3, 224, 224)
|
||||
with torch.no_grad():
|
||||
torch.testing.assert_close(
|
||||
layer.forward_native(x), ref(x), rtol=1e-4, atol=1e-4
|
||||
)
|
||||
|
||||
def test_padding_valid(self):
|
||||
layer = Conv2dLayer(3, 768, kernel_size=14, stride=14, padding="valid")
|
||||
self.assertFalse(layer.enable_linear)
|
||||
self.assertEqual(layer.padding, (0, 0))
|
||||
|
||||
def test_padding_same_disables_linear(self):
|
||||
layer = Conv2dLayer(3, 64, kernel_size=3, stride=1, padding="same")
|
||||
self.assertFalse(layer.enable_linear)
|
||||
|
||||
def test_non_matching_stride_disables_linear(self):
|
||||
layer = Conv2dLayer(3, 64, kernel_size=3, stride=1, padding=1)
|
||||
self.assertFalse(layer.enable_linear)
|
||||
|
||||
def test_groups_disable_linear(self):
|
||||
layer = Conv2dLayer(4, 8, kernel_size=2, stride=2, groups=2)
|
||||
self.assertFalse(layer.enable_linear)
|
||||
|
||||
def test_default_disables_linear(self):
|
||||
layer = Conv2dLayer(3, 768, kernel_size=14, stride=14)
|
||||
self.assertFalse(layer.enable_linear)
|
||||
|
||||
def test_dilation_disables_linear(self):
|
||||
layer = Conv2dLayer(3, 64, kernel_size=3, stride=3, dilation=2)
|
||||
self.assertFalse(layer.enable_linear)
|
||||
|
||||
def test_padding_mode_reflect(self):
|
||||
layer = Conv2dLayer(
|
||||
3, 64, kernel_size=3, stride=1, padding=1, padding_mode="reflect", bias=True
|
||||
)
|
||||
ref = nn.Conv2d(
|
||||
3, 64, kernel_size=3, stride=1, padding=1, padding_mode="reflect", bias=True
|
||||
)
|
||||
self.assertFalse(layer.enable_linear)
|
||||
_copy_weights(layer, ref)
|
||||
x = torch.randn(1, 3, 16, 16)
|
||||
with torch.no_grad():
|
||||
torch.testing.assert_close(
|
||||
layer.forward_native(x), ref(x), rtol=1e-4, atol=1e-4
|
||||
)
|
||||
|
||||
def test_conv_path_with_padding(self):
|
||||
layer = Conv2dLayer(3, 64, kernel_size=3, stride=1, padding=1, bias=True)
|
||||
ref = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=True)
|
||||
_copy_weights(layer, ref)
|
||||
x = torch.randn(1, 3, 32, 32)
|
||||
with torch.no_grad():
|
||||
torch.testing.assert_close(
|
||||
layer.forward_native(x), ref(x), rtol=1e-4, atol=1e-4
|
||||
)
|
||||
|
||||
def test_mulmat_matches_conv(self):
|
||||
layer = Conv2dLayer(
|
||||
3, 768, kernel_size=14, stride=14, bias=True, disable_linear=False
|
||||
)
|
||||
self.assertTrue(layer.enable_linear)
|
||||
x = torch.randn(2, 3, 224, 224)
|
||||
with torch.no_grad():
|
||||
torch.testing.assert_close(
|
||||
layer._forward_mulmat(x),
|
||||
layer._forward_conv(x),
|
||||
rtol=1e-4,
|
||||
atol=1e-4,
|
||||
)
|
||||
|
||||
def test_forward_cuda_uses_mulmat_when_enabled(self):
|
||||
layer = Conv2dLayer(
|
||||
3, 64, kernel_size=4, stride=4, bias=False, disable_linear=False
|
||||
)
|
||||
self.assertTrue(layer.enable_linear)
|
||||
x = torch.randn(1, 3, 16, 16)
|
||||
with torch.no_grad():
|
||||
torch.testing.assert_close(layer.forward_cuda(x), layer._forward_mulmat(x))
|
||||
|
||||
def test_forward_cuda_uses_conv_when_not_eligible(self):
|
||||
layer = Conv2dLayer(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
|
||||
self.assertFalse(layer.enable_linear)
|
||||
x = torch.randn(1, 3, 16, 16)
|
||||
with torch.no_grad():
|
||||
torch.testing.assert_close(layer.forward_cuda(x), layer._forward_conv(x))
|
||||
|
||||
def test_tuple_kernel_size(self):
|
||||
layer = Conv2dLayer(
|
||||
3,
|
||||
768,
|
||||
kernel_size=(14, 14),
|
||||
stride=(14, 14),
|
||||
bias=False,
|
||||
disable_linear=False,
|
||||
)
|
||||
self.assertTrue(layer.enable_linear)
|
||||
ref = nn.Conv2d(3, 768, kernel_size=(14, 14), stride=(14, 14), bias=False)
|
||||
_copy_weights(layer, ref)
|
||||
x = torch.randn(1, 3, 224, 224)
|
||||
with torch.no_grad():
|
||||
torch.testing.assert_close(
|
||||
layer.forward_native(x), ref(x), rtol=1e-4, atol=1e-4
|
||||
)
|
||||
|
||||
def test_output_shape(self):
|
||||
layer = Conv2dLayer(3, 768, kernel_size=16, stride=16, bias=False)
|
||||
x = torch.randn(4, 3, 224, 224)
|
||||
out = layer.forward_native(x)
|
||||
self.assertEqual(out.shape, (4, 768, 14, 14))
|
||||
|
||||
def test_no_bias_parameter(self):
|
||||
layer = Conv2dLayer(3, 64, kernel_size=4, stride=4, bias=False)
|
||||
self.assertIsNone(layer.bias)
|
||||
|
||||
|
||||
class TestConvValidation(unittest.TestCase):
|
||||
|
||||
def test_in_channels_not_divisible_by_groups(self):
|
||||
with self.assertRaises(ValueError):
|
||||
Conv2dLayer(3, 64, kernel_size=3, stride=1, groups=2)
|
||||
|
||||
def test_out_channels_not_divisible_by_groups(self):
|
||||
with self.assertRaises(ValueError):
|
||||
Conv2dLayer(4, 6, kernel_size=3, stride=1, groups=4)
|
||||
|
||||
def test_invalid_padding_string(self):
|
||||
with self.assertRaises(ValueError):
|
||||
Conv2dLayer(3, 64, kernel_size=3, stride=1, padding="full")
|
||||
|
||||
def test_padding_same_with_stride(self):
|
||||
with self.assertRaises(ValueError):
|
||||
Conv2dLayer(3, 64, kernel_size=3, stride=2, padding="same")
|
||||
|
||||
def test_padding_same_with_non_zeros_padding_mode(self):
|
||||
layer = Conv2dLayer(
|
||||
3,
|
||||
64,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding="same",
|
||||
padding_mode="reflect",
|
||||
bias=True,
|
||||
)
|
||||
ref = nn.Conv2d(
|
||||
3,
|
||||
64,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding="same",
|
||||
padding_mode="reflect",
|
||||
bias=True,
|
||||
)
|
||||
self.assertFalse(layer.enable_linear)
|
||||
_copy_weights(layer, ref)
|
||||
x = torch.randn(1, 3, 16, 16)
|
||||
with torch.no_grad():
|
||||
torch.testing.assert_close(
|
||||
layer.forward_native(x), ref(x), rtol=1e-4, atol=1e-4
|
||||
)
|
||||
|
||||
def test_invalid_padding_mode(self):
|
||||
with self.assertRaises(ValueError):
|
||||
Conv3dLayer(3, 64, kernel_size=3, stride=1, padding_mode="invalid")
|
||||
|
||||
def test_conv3d_in_channels_not_divisible_by_groups(self):
|
||||
with self.assertRaises(ValueError):
|
||||
Conv3dLayer(3, 64, kernel_size=3, stride=1, groups=2)
|
||||
|
||||
|
||||
class TestConv3dLayer(unittest.TestCase):
|
||||
|
||||
def test_basic_temporal_patch_embedding(self):
|
||||
layer = Conv3dLayer(
|
||||
3, 1152, kernel_size=[2, 14, 14], stride=[2, 14, 14], bias=False
|
||||
)
|
||||
ref = nn.Conv3d(
|
||||
3, 1152, kernel_size=[2, 14, 14], stride=[2, 14, 14], bias=False
|
||||
)
|
||||
self.assertTrue(layer.enable_linear)
|
||||
_copy_weights(layer, ref)
|
||||
x = torch.randn(1, 3, 2, 14, 14)
|
||||
with torch.no_grad():
|
||||
torch.testing.assert_close(
|
||||
layer.forward_native(x), ref(x), rtol=1e-4, atol=1e-4
|
||||
)
|
||||
|
||||
def test_with_bias(self):
|
||||
layer = Conv3dLayer(
|
||||
3, 1536, kernel_size=[2, 14, 14], stride=[2, 14, 14], bias=True
|
||||
)
|
||||
ref = nn.Conv3d(3, 1536, kernel_size=[2, 14, 14], stride=[2, 14, 14], bias=True)
|
||||
self.assertTrue(layer.enable_linear)
|
||||
_copy_weights(layer, ref)
|
||||
x = torch.randn(4, 3, 2, 14, 14)
|
||||
with torch.no_grad():
|
||||
torch.testing.assert_close(
|
||||
layer.forward_native(x), ref(x), rtol=1e-4, atol=1e-4
|
||||
)
|
||||
|
||||
def test_mulmat_matches_conv(self):
|
||||
layer = Conv3dLayer(
|
||||
3, 1152, kernel_size=[2, 14, 14], stride=[2, 14, 14], bias=True
|
||||
)
|
||||
self.assertTrue(layer.enable_linear)
|
||||
x = torch.randn(2, 3, 2, 14, 14)
|
||||
with torch.no_grad():
|
||||
torch.testing.assert_close(
|
||||
layer._forward_mulmat(x),
|
||||
layer._forward_conv(x),
|
||||
rtol=1e-4,
|
||||
atol=1e-4,
|
||||
)
|
||||
|
||||
def test_non_matching_stride_disables_linear(self):
|
||||
layer = Conv3dLayer(3, 64, kernel_size=3, stride=1, padding=1)
|
||||
self.assertFalse(layer.enable_linear)
|
||||
|
||||
def test_dilation_disables_linear(self):
|
||||
layer = Conv3dLayer(3, 64, kernel_size=3, stride=3, dilation=2)
|
||||
self.assertFalse(layer.enable_linear)
|
||||
|
||||
def test_disable_linear(self):
|
||||
layer = Conv3dLayer(
|
||||
3,
|
||||
1152,
|
||||
kernel_size=[2, 14, 14],
|
||||
stride=[2, 14, 14],
|
||||
bias=False,
|
||||
disable_linear=True,
|
||||
)
|
||||
self.assertFalse(layer.enable_linear)
|
||||
ref = nn.Conv3d(
|
||||
3, 1152, kernel_size=[2, 14, 14], stride=[2, 14, 14], bias=False
|
||||
)
|
||||
_copy_weights(layer, ref)
|
||||
x = torch.randn(1, 3, 2, 14, 14)
|
||||
with torch.no_grad():
|
||||
torch.testing.assert_close(
|
||||
layer.forward_native(x), ref(x), rtol=1e-4, atol=1e-4
|
||||
)
|
||||
|
||||
def test_conv_path_with_padding(self):
|
||||
layer = Conv3dLayer(3, 64, kernel_size=3, stride=1, padding=1, bias=True)
|
||||
ref = nn.Conv3d(3, 64, kernel_size=3, stride=1, padding=1, bias=True)
|
||||
_copy_weights(layer, ref)
|
||||
x = torch.randn(1, 3, 4, 8, 8)
|
||||
with torch.no_grad():
|
||||
torch.testing.assert_close(
|
||||
layer.forward_native(x), ref(x), rtol=1e-4, atol=1e-4
|
||||
)
|
||||
|
||||
def test_output_shape(self):
|
||||
layer = Conv3dLayer(
|
||||
3, 1152, kernel_size=[2, 14, 14], stride=[2, 14, 14], bias=False
|
||||
)
|
||||
x = torch.randn(1, 3, 2, 14, 14)
|
||||
out = layer.forward_native(x)
|
||||
self.assertEqual(out.shape, (1, 1152, 1, 1, 1))
|
||||
|
||||
def test_batch_processing(self):
|
||||
layer = Conv3dLayer(
|
||||
3, 1536, kernel_size=[2, 14, 14], stride=[2, 14, 14], bias=True
|
||||
)
|
||||
ref = nn.Conv3d(3, 1536, kernel_size=[2, 14, 14], stride=[2, 14, 14], bias=True)
|
||||
_copy_weights(layer, ref)
|
||||
x = torch.randn(8, 3, 2, 14, 14)
|
||||
with torch.no_grad():
|
||||
torch.testing.assert_close(
|
||||
layer.forward_native(x), ref(x), rtol=1e-4, atol=1e-4
|
||||
)
|
||||
|
||||
def test_forward_native_uses_mulmat_when_eligible(self):
|
||||
layer = Conv3dLayer(3, 128, kernel_size=[2, 4, 4], stride=[2, 4, 4], bias=True)
|
||||
self.assertTrue(layer.enable_linear)
|
||||
x = torch.randn(1, 3, 2, 4, 4)
|
||||
with torch.no_grad():
|
||||
torch.testing.assert_close(
|
||||
layer.forward_native(x), layer._forward_mulmat(x)
|
||||
)
|
||||
|
||||
def test_padding_valid(self):
|
||||
layer = Conv3dLayer(
|
||||
3, 64, kernel_size=[2, 4, 4], stride=[2, 4, 4], padding="valid"
|
||||
)
|
||||
self.assertTrue(layer.enable_linear)
|
||||
self.assertEqual(layer.padding, (0, 0, 0))
|
||||
|
||||
def test_weight_shape(self):
|
||||
layer = Conv3dLayer(
|
||||
3, 1152, kernel_size=[2, 14, 14], stride=[2, 14, 14], bias=False
|
||||
)
|
||||
self.assertEqual(layer.weight.shape, (1152, 3, 2, 14, 14))
|
||||
|
||||
def test_glm4v_workflow(self):
|
||||
"""GLM4V-style: 2D input -> reshape to 5D -> Conv3dLayer -> flatten."""
|
||||
in_channels, temporal_patch_size, patch_size = 3, 2, 14
|
||||
hidden_size = 1536
|
||||
layer = Conv3dLayer(
|
||||
in_channels,
|
||||
hidden_size,
|
||||
kernel_size=[temporal_patch_size, patch_size, patch_size],
|
||||
stride=[temporal_patch_size, patch_size, patch_size],
|
||||
bias=True,
|
||||
)
|
||||
ref = nn.Conv3d(
|
||||
in_channels,
|
||||
hidden_size,
|
||||
kernel_size=[temporal_patch_size, patch_size, patch_size],
|
||||
stride=[temporal_patch_size, patch_size, patch_size],
|
||||
bias=True,
|
||||
)
|
||||
_copy_weights(layer, ref)
|
||||
num_patches = 4
|
||||
flat_dim = in_channels * temporal_patch_size * patch_size * patch_size
|
||||
x_2d = torch.randn(num_patches, flat_dim)
|
||||
x_5d = x_2d.view(-1, in_channels, temporal_patch_size, patch_size, patch_size)
|
||||
with torch.no_grad():
|
||||
torch.testing.assert_close(
|
||||
layer.forward_native(x_5d).view(-1, hidden_size),
|
||||
ref(x_5d).view(-1, hidden_size),
|
||||
rtol=1e-4,
|
||||
atol=1e-4,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,358 @@
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=10, suite="stage-b-test-small-1-gpu")
|
||||
|
||||
import os
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
try:
|
||||
from sglang.srt.layers.attention.mamba.mamba_state_scatter_triton import (
|
||||
fused_mamba_state_scatter_with_mask,
|
||||
)
|
||||
|
||||
_FUSED_IMPORT_ERROR = None
|
||||
except Exception as e: # pragma: no cover
|
||||
fused_mamba_state_scatter_with_mask = None
|
||||
_FUSED_IMPORT_ERROR = e
|
||||
|
||||
|
||||
def _dtype_from_str(name: str) -> torch.dtype:
|
||||
mapping = {
|
||||
"bfloat16": torch.bfloat16,
|
||||
"float16": torch.float16,
|
||||
"float32": torch.float32,
|
||||
}
|
||||
if name not in mapping:
|
||||
raise ValueError(
|
||||
f"Unsupported dtype string {name!r}. Supported: {sorted(mapping.keys())}"
|
||||
)
|
||||
return mapping[name]
|
||||
|
||||
|
||||
def _ref_scatter(dst, src, dst_indices, src_indices, step_indices):
|
||||
"""Reference implementation using PyTorch advanced indexing."""
|
||||
# dst: [L, C, E]
|
||||
# src: [L, S, D, E]
|
||||
dst[:, dst_indices] = src[:, src_indices, step_indices].to(dst.dtype, copy=False)
|
||||
|
||||
|
||||
def _ref_update_like(
|
||||
ssm_states,
|
||||
intermediate_ssm,
|
||||
conv_states,
|
||||
intermediate_conv,
|
||||
*,
|
||||
state_indices_tensor,
|
||||
accepted_steps,
|
||||
mamba_track_indices=None,
|
||||
mamba_steps_to_track=None,
|
||||
):
|
||||
"""Reference implementation using PyTorch advanced indexing for correctness verification."""
|
||||
request_number = accepted_steps.shape[0]
|
||||
intermediate_state_indices = torch.arange(
|
||||
request_number, dtype=torch.int32, device=accepted_steps.device
|
||||
)
|
||||
|
||||
valid_mask = accepted_steps >= 0
|
||||
dst_state_indices = state_indices_tensor[valid_mask].to(torch.int64)
|
||||
src_state_indices = intermediate_state_indices[valid_mask].to(torch.int64)
|
||||
last_steps = accepted_steps[valid_mask].to(torch.int64)
|
||||
|
||||
# Only scatter if there are valid indices (but don't early return -
|
||||
# mamba_track_indices processing is independent)
|
||||
if dst_state_indices.numel() > 0:
|
||||
_ref_scatter(
|
||||
ssm_states,
|
||||
intermediate_ssm,
|
||||
dst_state_indices,
|
||||
src_state_indices,
|
||||
last_steps,
|
||||
)
|
||||
_ref_scatter(
|
||||
conv_states,
|
||||
intermediate_conv,
|
||||
dst_state_indices,
|
||||
src_state_indices,
|
||||
last_steps,
|
||||
)
|
||||
|
||||
if mamba_track_indices is not None:
|
||||
assert mamba_steps_to_track is not None
|
||||
track_mask = mamba_steps_to_track >= 0
|
||||
if not track_mask.any():
|
||||
return
|
||||
dst_track_indices = mamba_track_indices[track_mask].to(torch.int64)
|
||||
src_track_indices = intermediate_state_indices[track_mask].to(torch.int64)
|
||||
track_steps = mamba_steps_to_track[track_mask].to(torch.int64)
|
||||
|
||||
_ref_scatter(
|
||||
ssm_states,
|
||||
intermediate_ssm,
|
||||
dst_track_indices,
|
||||
src_track_indices,
|
||||
track_steps,
|
||||
)
|
||||
_ref_scatter(
|
||||
conv_states,
|
||||
intermediate_conv,
|
||||
dst_track_indices,
|
||||
src_track_indices,
|
||||
track_steps,
|
||||
)
|
||||
|
||||
|
||||
def _fused_update_like(
|
||||
ssm_states,
|
||||
intermediate_ssm,
|
||||
conv_states,
|
||||
intermediate_conv,
|
||||
*,
|
||||
state_indices_tensor,
|
||||
accepted_steps,
|
||||
mamba_track_indices=None,
|
||||
mamba_steps_to_track=None,
|
||||
):
|
||||
"""Matches the fully fused logic that avoids index_select and nonzero calls."""
|
||||
# Use fully fused kernel that handles masking internally
|
||||
fused_mamba_state_scatter_with_mask(
|
||||
ssm_states,
|
||||
intermediate_ssm,
|
||||
state_indices_tensor,
|
||||
accepted_steps,
|
||||
)
|
||||
fused_mamba_state_scatter_with_mask(
|
||||
conv_states,
|
||||
intermediate_conv,
|
||||
state_indices_tensor,
|
||||
accepted_steps,
|
||||
)
|
||||
|
||||
if mamba_track_indices is not None:
|
||||
assert mamba_steps_to_track is not None
|
||||
fused_mamba_state_scatter_with_mask(
|
||||
ssm_states,
|
||||
intermediate_ssm,
|
||||
mamba_track_indices,
|
||||
mamba_steps_to_track,
|
||||
)
|
||||
fused_mamba_state_scatter_with_mask(
|
||||
conv_states,
|
||||
intermediate_conv,
|
||||
mamba_track_indices,
|
||||
mamba_steps_to_track,
|
||||
)
|
||||
|
||||
|
||||
def _time_cuda_ms(fn, iters=50, warmup=10):
|
||||
"""Measure average CUDA time (ms) using CUDA events."""
|
||||
for _ in range(warmup):
|
||||
fn()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
for _ in range(iters):
|
||||
fn()
|
||||
end.record()
|
||||
torch.cuda.synchronize()
|
||||
return start.elapsed_time(end) / iters
|
||||
|
||||
|
||||
class TestMambaStateScatterCorrectness(unittest.TestCase):
|
||||
@unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for this test.")
|
||||
def test_fused_matches_reference(self):
|
||||
"""Test that fused_mamba_state_scatter_with_mask matches the reference."""
|
||||
if fused_mamba_state_scatter_with_mask is None:
|
||||
self.skipTest(
|
||||
f"fused_mamba_state_scatter_with_mask import failed: {_FUSED_IMPORT_ERROR}"
|
||||
)
|
||||
|
||||
torch.manual_seed(42)
|
||||
device = torch.device("cuda")
|
||||
|
||||
# Keep sizes moderate so this test is quick.
|
||||
L = 8
|
||||
B = 32
|
||||
C = 49
|
||||
D = 5
|
||||
ssm_elems = 1024
|
||||
conv_elems = 512
|
||||
|
||||
ssm_states0 = torch.randn(
|
||||
(L, C, ssm_elems), device=device, dtype=torch.bfloat16
|
||||
)
|
||||
conv_states0 = torch.randn(
|
||||
(L, C, conv_elems), device=device, dtype=torch.bfloat16
|
||||
)
|
||||
intermediate_ssm = torch.randn(
|
||||
(L, B, D, ssm_elems), device=device, dtype=torch.bfloat16
|
||||
)
|
||||
intermediate_conv = torch.randn(
|
||||
(L, B, D, conv_elems), device=device, dtype=torch.bfloat16
|
||||
)
|
||||
|
||||
# unique cache lines (no duplicates) to avoid nondeterministic write order
|
||||
state_indices_tensor = torch.randperm(C, device=device, dtype=torch.int64)[
|
||||
:B
|
||||
].to(torch.int32)
|
||||
|
||||
accepted_steps = torch.randint(0, D, (B,), device=device, dtype=torch.int64)
|
||||
# set ~10% invalid
|
||||
invalid = torch.rand((B,), device=device) < 0.1
|
||||
accepted_steps[invalid] = -1
|
||||
|
||||
# Optional track update
|
||||
mamba_track_indices = torch.randperm(C, device=device, dtype=torch.int64)[:B]
|
||||
mamba_steps_to_track = torch.randint(
|
||||
0, D, (B,), device=device, dtype=torch.int64
|
||||
)
|
||||
track_invalid = torch.rand((B,), device=device) < 0.7
|
||||
mamba_steps_to_track[track_invalid] = -1
|
||||
|
||||
ssm_ref = ssm_states0.clone()
|
||||
conv_ref = conv_states0.clone()
|
||||
ssm_fused = ssm_states0.clone()
|
||||
conv_fused = conv_states0.clone()
|
||||
|
||||
_ref_update_like(
|
||||
ssm_ref,
|
||||
intermediate_ssm,
|
||||
conv_ref,
|
||||
intermediate_conv,
|
||||
state_indices_tensor=state_indices_tensor,
|
||||
accepted_steps=accepted_steps,
|
||||
mamba_track_indices=mamba_track_indices,
|
||||
mamba_steps_to_track=mamba_steps_to_track,
|
||||
)
|
||||
_fused_update_like(
|
||||
ssm_fused,
|
||||
intermediate_ssm,
|
||||
conv_fused,
|
||||
intermediate_conv,
|
||||
state_indices_tensor=state_indices_tensor,
|
||||
accepted_steps=accepted_steps,
|
||||
mamba_track_indices=mamba_track_indices,
|
||||
mamba_steps_to_track=mamba_steps_to_track,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(ssm_fused, ssm_ref)
|
||||
torch.testing.assert_close(conv_fused, conv_ref)
|
||||
|
||||
|
||||
class TestMambaStateScatterPerf(unittest.TestCase):
|
||||
@unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for this test.")
|
||||
def test_perf_report_old_vs_fused(self):
|
||||
"""Optional microbenchmark comparing baseline vs fused kernel.
|
||||
|
||||
Enable with: SGLANG_RUN_MAMBA_SCATTER_PERF_TEST=1
|
||||
"""
|
||||
if os.environ.get("SGLANG_RUN_MAMBA_SCATTER_PERF_TEST", "0") != "1":
|
||||
self.skipTest("Set SGLANG_RUN_MAMBA_SCATTER_PERF_TEST=1 to run perf test.")
|
||||
if fused_mamba_state_scatter_with_mask is None:
|
||||
self.skipTest(
|
||||
f"fused_mamba_state_scatter_with_mask import failed: {_FUSED_IMPORT_ERROR}"
|
||||
)
|
||||
|
||||
torch.manual_seed(0)
|
||||
device = torch.device("cuda")
|
||||
|
||||
# Parameterize sizes via env vars so we can match a real model more closely.
|
||||
L = int(os.environ.get("SGLANG_MAMBA_SCATTER_LAYERS", "32"))
|
||||
B = int(os.environ.get("SGLANG_MAMBA_SCATTER_BATCH", "48"))
|
||||
C = int(os.environ.get("SGLANG_MAMBA_SCATTER_CACHE", "49"))
|
||||
D = int(os.environ.get("SGLANG_MAMBA_SCATTER_DRAFT_TOKENS", "5"))
|
||||
ssm_elems = int(os.environ.get("SGLANG_MAMBA_SCATTER_SSM_ELEMS", "4096"))
|
||||
conv_elems = int(os.environ.get("SGLANG_MAMBA_SCATTER_CONV_ELEMS", "512"))
|
||||
invalid_ratio = float(
|
||||
os.environ.get("SGLANG_MAMBA_SCATTER_INVALID_RATIO", "0.0")
|
||||
)
|
||||
track_ratio = float(os.environ.get("SGLANG_MAMBA_SCATTER_TRACK_RATIO", "0.0"))
|
||||
ssm_dtype = _dtype_from_str(
|
||||
os.environ.get("SGLANG_MAMBA_SCATTER_SSM_DTYPE", "bfloat16")
|
||||
)
|
||||
conv_dtype = _dtype_from_str(
|
||||
os.environ.get("SGLANG_MAMBA_SCATTER_CONV_DTYPE", "bfloat16")
|
||||
)
|
||||
|
||||
# Use zeros for dst so each iteration overwrites the same memory.
|
||||
ssm_states = torch.zeros((L, C, ssm_elems), device=device, dtype=ssm_dtype)
|
||||
conv_states = torch.zeros((L, C, conv_elems), device=device, dtype=conv_dtype)
|
||||
intermediate_ssm = torch.randn(
|
||||
(L, B, D, ssm_elems), device=device, dtype=ssm_dtype
|
||||
)
|
||||
intermediate_conv = torch.randn(
|
||||
(L, B, D, conv_elems), device=device, dtype=conv_dtype
|
||||
)
|
||||
|
||||
state_indices_tensor = torch.randperm(C, device=device, dtype=torch.int64)[
|
||||
:B
|
||||
].to(torch.int32)
|
||||
accepted_steps = torch.randint(0, D, (B,), device=device, dtype=torch.int64)
|
||||
if invalid_ratio > 0:
|
||||
invalid = torch.rand((B,), device=device) < invalid_ratio
|
||||
accepted_steps[invalid] = -1
|
||||
|
||||
mamba_track_indices = None
|
||||
mamba_steps_to_track = None
|
||||
if track_ratio > 0:
|
||||
mamba_track_indices = torch.randperm(C, device=device, dtype=torch.int64)[
|
||||
:B
|
||||
]
|
||||
mamba_steps_to_track = torch.randint(
|
||||
0, D, (B,), device=device, dtype=torch.int64
|
||||
)
|
||||
track_invalid = torch.rand((B,), device=device) >= track_ratio
|
||||
mamba_steps_to_track[track_invalid] = -1
|
||||
|
||||
def ref_fn():
|
||||
_ref_update_like(
|
||||
ssm_states,
|
||||
intermediate_ssm,
|
||||
conv_states,
|
||||
intermediate_conv,
|
||||
state_indices_tensor=state_indices_tensor,
|
||||
accepted_steps=accepted_steps,
|
||||
mamba_track_indices=mamba_track_indices,
|
||||
mamba_steps_to_track=mamba_steps_to_track,
|
||||
)
|
||||
|
||||
def fused_fn():
|
||||
_fused_update_like(
|
||||
ssm_states,
|
||||
intermediate_ssm,
|
||||
conv_states,
|
||||
intermediate_conv,
|
||||
state_indices_tensor=state_indices_tensor,
|
||||
accepted_steps=accepted_steps,
|
||||
mamba_track_indices=mamba_track_indices,
|
||||
mamba_steps_to_track=mamba_steps_to_track,
|
||||
)
|
||||
|
||||
# Warm up JIT compilation for triton kernels (and caches for torch indexing)
|
||||
ref_fn()
|
||||
fused_fn()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
ref_ms = _time_cuda_ms(ref_fn)
|
||||
fused_ms = _time_cuda_ms(fused_fn)
|
||||
|
||||
num_valid = int((accepted_steps >= 0).sum().item())
|
||||
ratio = fused_ms / ref_ms if ref_ms > 0 else float("inf")
|
||||
speedup = ref_ms / fused_ms if fused_ms > 0 else float("inf")
|
||||
|
||||
# Print a concise report
|
||||
print(
|
||||
"\n[MambaStateScatterPerf]\n"
|
||||
f" shapes: L={L} B={B} C={C} D={D} ssm_elems={ssm_elems} conv_elems={conv_elems}\n"
|
||||
f" dtypes: ssm={ssm_dtype} conv={conv_dtype}\n"
|
||||
f" valid: {num_valid}/{B} invalid_ratio={invalid_ratio} track_ratio={track_ratio}\n"
|
||||
f" ref_total_ms (baseline): {ref_ms:.4f}\n"
|
||||
f" fused_total_ms: {fused_ms:.4f} (ratio={ratio:.3f}x, speedup={speedup:.2f}x)\n"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__": # pragma: no cover
|
||||
unittest.main()
|
||||
+2
-1
@@ -20,7 +20,8 @@ class TestSLRUAccuracy(unittest.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
"""Setup minimal memory pools for testing"""
|
||||
device = "cpu" # Using CPU for testing simplicity
|
||||
torch.set_default_device(None)
|
||||
device = "cpu"
|
||||
dtype = torch.float16
|
||||
|
||||
# Create smaller KV cache to ensure evictions occur
|
||||
@@ -0,0 +1,84 @@
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=5, suite="stage-a-cpu-only")
|
||||
|
||||
import unittest
|
||||
|
||||
from sglang.srt.utils.gauge_histogram import BucketLabels
|
||||
|
||||
|
||||
class TestBucketLabels(unittest.TestCase):
|
||||
"""Test BucketLabels with hardcoded expected values."""
|
||||
|
||||
def test_labels_basic(self):
|
||||
buckets = BucketLabels([10, 30, 60])
|
||||
self.assertEqual(
|
||||
list(buckets),
|
||||
[("0", "10"), ("10", "30"), ("30", "60"), ("60", "+Inf")],
|
||||
)
|
||||
|
||||
def test_labels_single_bound(self):
|
||||
buckets = BucketLabels([100])
|
||||
self.assertEqual(list(buckets), [("0", "100"), ("100", "+Inf")])
|
||||
|
||||
def test_labels_many_bounds(self):
|
||||
buckets = BucketLabels([1, 2, 5, 10])
|
||||
self.assertEqual(
|
||||
list(buckets),
|
||||
[("0", "1"), ("1", "2"), ("2", "5"), ("5", "10"), ("10", "+Inf")],
|
||||
)
|
||||
|
||||
def test_len(self):
|
||||
buckets = BucketLabels([10, 30, 60])
|
||||
self.assertEqual(len(buckets), 4)
|
||||
|
||||
|
||||
class TestBucketLabelsCounts(unittest.TestCase):
|
||||
"""Test BucketLabels.compute_bucket_counts with hardcoded expected values."""
|
||||
|
||||
def test_empty_observations(self):
|
||||
buckets = BucketLabels([10, 30, 60])
|
||||
self.assertEqual(buckets.compute_bucket_counts([]), [0, 0, 0, 0])
|
||||
|
||||
def test_single_value_first_bucket(self):
|
||||
# bounds: [10, 30, 60] -> buckets: (0,10], (10,30], (30,60], (60,+Inf]
|
||||
buckets = BucketLabels([10, 30, 60])
|
||||
self.assertEqual(buckets.compute_bucket_counts([5]), [1, 0, 0, 0])
|
||||
|
||||
def test_single_value_last_bucket(self):
|
||||
buckets = BucketLabels([10, 30, 60])
|
||||
self.assertEqual(buckets.compute_bucket_counts([100]), [0, 0, 0, 1])
|
||||
|
||||
def test_exact_boundary_values(self):
|
||||
# Values at exact boundaries: 10 -> (0,10], 30 -> (10,30], 60 -> (30,60]
|
||||
buckets = BucketLabels([10, 30, 60])
|
||||
self.assertEqual(buckets.compute_bucket_counts([10, 30, 60]), [1, 1, 1, 0])
|
||||
|
||||
def test_just_above_boundary(self):
|
||||
# 11 -> (10,30], 31 -> (30,60], 61 -> (60,+Inf]
|
||||
buckets = BucketLabels([10, 30, 60])
|
||||
self.assertEqual(buckets.compute_bucket_counts([11, 31, 61]), [0, 1, 1, 1])
|
||||
|
||||
def test_multiple_values_same_bucket(self):
|
||||
buckets = BucketLabels([10, 30, 60])
|
||||
self.assertEqual(buckets.compute_bucket_counts([1, 2, 3, 4, 5]), [5, 0, 0, 0])
|
||||
|
||||
def test_all_overflow(self):
|
||||
buckets = BucketLabels([10, 30, 60])
|
||||
self.assertEqual(buckets.compute_bucket_counts([100, 200, 300]), [0, 0, 0, 3])
|
||||
|
||||
def test_distribution(self):
|
||||
# 5 (<=10), 10 (<=10), 15 (<=30), 40 (<=60), 100 (+Inf)
|
||||
buckets = BucketLabels([10, 30, 60])
|
||||
self.assertEqual(
|
||||
buckets.compute_bucket_counts([5, 10, 15, 40, 100]), [2, 1, 1, 1]
|
||||
)
|
||||
|
||||
def test_float_values(self):
|
||||
# 9.9 -> (0,10], 10.1 -> (10,30], 30.5 -> (30,60]
|
||||
buckets = BucketLabels([10, 30, 60])
|
||||
self.assertEqual(buckets.compute_bucket_counts([9.9, 10.1, 30.5]), [1, 1, 1, 0])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user