chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

ci: migrate Debug Utils, Ops, and Rotary Embedding tests to test/registered/ (#16422)

Browse Source
This commit is contained in:
Alison Shao
2026-01-06 14:35:57 -08:00
committed by GitHub
parent 6beb50d612
commit 2e0527dd74
4 changed files with 17 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/srt/debug_utils/test_tensor_dump_forward_hook.py
-95
View File
@@ -1,95 +0,0 @@
import unittest
import torch
from torch import nn
from sglang.srt.debug_utils.tensor_dump_forward_hook import (
register_forward_hook_for_model,
)
from sglang.srt.distributed.parallel_state import (
init_distributed_environment,
initialize_model_parallel,
)
from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.linear import LinearBase
from sglang.srt.models.qwen2 import Qwen2MLP
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
from sglang.srt.utils import add_prefix
TEST_HIDDEN_SIZE = 32
class SimpleModel(nn.Module):
def __init__(self) -> None:
super().__init__()
self.hidden_size = TEST_HIDDEN_SIZE
self.rms_norm_eps = 1e-5
self.mlp = Qwen2MLP(
hidden_size=self.hidden_size,
intermediate_size=self.hidden_size,
hidden_act="silu",
quant_config=None,
prefix=add_prefix("mlp", ""),
)
self.layernorm = RMSNorm(self.hidden_size, eps=self.rms_norm_eps)
@torch.no_grad()
def forward(
self,
hidden_states: torch.Tensor,
) -> torch.Tensor:
hidden_states = self.layernorm(hidden_states)
hidden_states = self.mlp(hidden_states)
return hidden_states
class MockCausalLM(nn.Module):
def __init__(self) -> None:
super().__init__()
self.model = SimpleModel()
@torch.no_grad()
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
return self.model(hidden_states)
def init_weights(module):
if isinstance(module, LinearBase):
torch.nn.init.uniform_(module.weight)
if module.bias is not None:
torch.nn.init.zeros_(module.bias)
elif isinstance(module, RMSNorm):
torch.nn.init.ones_(module.weight)
def test_model_forward_dump(tmp_path):
set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
init_distributed_environment(
backend="nccl",
world_size=1,
rank=0,
local_rank=0,
distributed_init_method="tcp://127.0.0.1:2646",
)
initialize_model_parallel()
model = MockCausalLM()
model.apply(init_weights)
model = model.cuda().bfloat16()
dumper = register_forward_hook_for_model(
model, tmp_path / "sglang_dump", [0], 0, 0, 0
)
dir_path = dumper.get_dump_dir()
inp = torch.randn(4, TEST_HIDDEN_SIZE, dtype=torch.bfloat16) * 0.01
result = model(inp.cuda())
data = torch.load(f"{dir_path}/Pass00000.pt")
assert "model.layernorm" in data
assert "model.mlp.down_proj" in data
assert torch.allclose(
data["model.mlp.down_proj"], result.cpu(), rtol=1e-5, atol=1e-5
)
if __name__ == "__main__":
unittest.main()
test/srt/ops/test_repeat_interleave.py
-141
View File
@@ -1,141 +0,0 @@
import time
from typing import Tuple
import numpy as np
import pytest
import torch
from sglang.srt.models.utils import compute_cu_seqlens_from_grid_numpy as cpu_numpy_impl
def torch_ref_impl(grid_thw: torch.Tensor) -> torch.Tensor:
"""
Pure PyTorch implementation of cu_seqlens computation.
Assumes grid_thw is already on the correct device (CPU here).
Shape: [T, 3], columns: [repeat_count, H, W]
"""
cu_seqlens = torch.repeat_interleave(
grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]
).cumsum(dim=0)
cu_seqlens = torch.cat(
[
torch.zeros(1, dtype=torch.int32, device=cu_seqlens.device),
cu_seqlens.to(torch.int32),
]
)
return cu_seqlens
def benchmark_once(fn, grid_thw, iters: int = 1000):
"""
Run a function `fn` on the same input `grid_thw` for `iters` times
and measure total elapsed time.
"""
start = time.perf_counter()
for _ in range(iters):
out = fn(grid_thw)
end = time.perf_counter()
return (end - start), out
# (T, repeat_min, repeat_max)
GRID_TEST_CONFIGS: list[Tuple[int, int, int]] = [
(16, 1, 4), # small T, small repeat counts
(128, 0, 4), # allow repeat=0 to test edge cases
(512, 1, 8),
(1024, 1, 16),
]
NUM_CASES_PER_CONFIG = 10
def _generate_random_grid(T: int, repeat_min: int, repeat_max: int) -> torch.Tensor:
"""
grid_thw: [T, 3]
col0: repeat count
col1, col2: arbitrary positive integers (here 1..16)
"""
repeats = torch.randint(repeat_min, repeat_max + 1, (T, 1), dtype=torch.int32)
th = torch.randint(1, 17, (T, 1), dtype=torch.int32)
tw = torch.randint(1, 17, (T, 1), dtype=torch.int32)
grid_thw = torch.cat([repeats, th, tw], dim=1)
return grid_thw
class TestRepeatInterleave:
@classmethod
def setup_class(cls):
torch.set_num_threads(1)
def setup_method(self, method):
torch.manual_seed(0)
np.random.seed(0)
@pytest.mark.parametrize(
"T,repeat_min,repeat_max",
GRID_TEST_CONFIGS,
)
@pytest.mark.parametrize("case_idx", range(NUM_CASES_PER_CONFIG))
def test_cpu_correctness_random_cases(
self,
T: int,
repeat_min: int,
repeat_max: int,
case_idx: int,
):
torch.manual_seed(case_idx)
np.random.seed(case_idx)
grid_thw = _generate_random_grid(T, repeat_min, repeat_max)
grid_clone = grid_thw.clone()
out_torch = torch_ref_impl(grid_thw)
out_numpy = cpu_numpy_impl(grid_thw)
assert torch.equal(grid_thw, grid_clone), "Function modified input grid_thw!"
assert (
out_torch.shape == out_numpy.shape
), f"Shape mismatch: torch={out_torch.shape}, numpy={out_numpy.shape}"
assert (
out_torch.dtype == torch.int32
), f"Unexpected torch dtype: {out_torch.dtype}"
assert (
out_numpy.dtype == torch.int32
), f"Unexpected numpy impl dtype: {out_numpy.dtype}"
if not torch.equal(out_torch.cpu(), out_numpy.cpu()):
diff_idx = (out_torch.cpu() != out_numpy.cpu()).nonzero(as_tuple=False)
idx0 = diff_idx[0].item()
pytest.fail(
f"Value mismatch, T={T}, case_idx={case_idx}, first differing index={idx0}, "
f"torch={out_torch[idx0].item()}, "
f"numpy={out_numpy[idx0].item()}"
)
def test_zero_repeat_edge_case(self):
T = 4
grid_thw = torch.tensor(
[
[0, 4, 4],
[1, 2, 3], # 6
[2, 1, 5], # 5, 5
[0, 7, 7], # 0
],
dtype=torch.int32,
)
grid_clone = grid_thw.clone()
out_torch = torch_ref_impl(grid_thw)
out_numpy = cpu_numpy_impl(grid_thw)
assert torch.equal(
grid_thw, grid_clone
), "Function modified input grid_thw with zero repeats!"
assert torch.equal(
out_torch.cpu(), out_numpy.cpu()
), f"Zero-repeat case mismatch: torch={out_torch}, numpy={out_numpy}"
test/srt/rotary_embedding/test_mrope.py
-143
View File
@@ -1,143 +0,0 @@
from typing import NamedTuple
import pytest
import torch
from packaging.version import Version
from transformers import AutoConfig
from transformers import __version__ as TRANSFORMERS_VERSION
from sglang.srt.layers.rotary_embedding import get_rope
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
from sglang.srt.utils import (
cpu_has_amx_support,
is_cpu,
is_cuda,
is_hip,
is_npu,
is_xpu,
)
_is_cuda = is_cuda()
_is_hip = is_hip()
_is_cpu = is_cpu()
_is_cpu_amx_available = cpu_has_amx_support()
_is_npu = is_npu()
_is_xpu = is_xpu()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def generate_test_data(
num_tokens: int,
num_q_heads: int,
num_kv_heads: int,
head_size: int,
max_position_embeddings: int,
dtype: torch.dtype,
device: torch.device,
):
"""Generate test data for given configuration."""
torch.manual_seed(42)
# Create 2D positions (3, num_tokens) for multimodal case
positions = torch.randint(
0, max_position_embeddings // 4, (3, num_tokens), device=device
)
# Create query and key tensors
query = torch.randn(num_tokens, num_q_heads * head_size, dtype=dtype, device=device)
key = torch.randn(num_tokens, num_kv_heads * head_size, dtype=dtype, device=device)
return positions, query, key
class MRoPETestInfo(NamedTuple):
model_name: str
atol: float = 1e-2
rtol: float = 1.6e-2
marks: list[pytest.MarkDecorator] = []
TRANSFORMERS_BASE_VERSION = Version(TRANSFORMERS_VERSION).base_version
MODELS_TO_TEST = [
MRoPETestInfo(model_name="Qwen/Qwen2-VL-7B-Instruct"),
MRoPETestInfo(model_name="Qwen/Qwen2-VL-72B-Instruct"),
MRoPETestInfo(model_name="Qwen/Qwen2.5-VL-72B-Instruct"),
]
num_tokens_list = [11, 8192]
@pytest.mark.skipif(not (_is_cuda or _is_hip), reason="Skipping CUDA/ROCm only tests.")
@pytest.mark.parametrize(
"model_info, model_name",
[
pytest.param(test_config, test_config.model_name, marks=test_config.marks)
for test_config in MODELS_TO_TEST
],
)
@pytest.mark.parametrize("tp_size", [1, 2])
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
@pytest.mark.parametrize("num_tokens", num_tokens_list)
def test_mrope(
model_name: str,
model_info: MRoPETestInfo,
tp_size: int,
dtype: torch.dtype,
num_tokens: int,
):
set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
atol = model_info.atol
rtol = model_info.rtol
config = AutoConfig.from_pretrained(model_name)
config = config.get_text_config()
# get the model config
total_num_kv_heads = config.num_key_value_heads
total_num_heads = config.num_attention_heads
num_heads = total_num_heads // tp_size
num_kv_heads = max(1, total_num_kv_heads // tp_size)
head_dim = (
config.head_dim
if hasattr(config, "head_dim")
else config.hidden_size // total_num_heads
)
is_neox_style = True
rope_theta = config.rope_theta
max_position = config.max_position_embeddings
partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0)
rotary_dim = int(head_dim * partial_rotary_factor)
mrope_helper_class = get_rope(
head_size=head_dim,
rotary_dim=rotary_dim,
max_position=max_position,
base=rope_theta,
is_neox_style=is_neox_style,
rope_scaling=config.rope_scaling,
dtype=dtype,
).to(device=device)
# create q k v input tensors
# create rotary pos emb input tensors
positions, query, key = generate_test_data(
num_tokens, num_heads, num_kv_heads, head_dim, max_position, dtype, device
)
query_native, key_native = mrope_helper_class._forward_native(
positions,
query.clone(),
key.clone(),
)
query_cuda, key_cuda = mrope_helper_class.forward(
positions,
query.clone(),
key.clone(),
)
torch.testing.assert_close(query_native, query_cuda, atol=atol, rtol=rtol)
torch.testing.assert_close(key_native, key_cuda, atol=atol, rtol=rtol)
test/srt/run_suite.py
-6
View File
@@ -9,8 +9,6 @@ from sglang.test.ci.ci_utils import TestFile, run_unittest_files
# NOTE: please sort the test cases alphabetically by the test file name
suites = {
"per-commit-1-gpu": [
TestFile("debug_utils/test_tensor_dump_forward_hook.py", 9),
TestFile("ops/test_repeat_interleave.py", 60),
TestFile("test_abort.py", 131),
TestFile("test_chunked_prefill.py", 312),
TestFile("test_deterministic.py", 228),
@@ -18,7 +16,6 @@ suites = {
TestFile("test_evs.py", 20),
TestFile("test_external_models.py", 30),
TestFile("test_fp8_utils.py", 9),
TestFile("rotary_embedding/test_mrope.py", 10),
TestFile("test_gpt_oss_1gpu.py", 402),
TestFile("test_hidden_states.py", 55),
TestFile("test_input_embeddings.py", 38),
@@ -130,7 +127,6 @@ suites = {
# NOTE: please sort the test cases alphabetically by the test file name
suite_amd = {
"per-commit-amd": [
TestFile("debug_utils/test_tensor_dump_forward_hook.py", 15),
# TestFile("hicache/test_hicache.py", 116), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/12575
# TestFile("hicache/test_hicache_mla.py", 127), # Disabled temporarily, # Temporarily disabled, see https://github.com/sgl-project/sglang/issues/12574
# TestFile("hicache/test_hicache_storage.py", 127), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/12575
@@ -138,8 +134,6 @@ suite_amd = {
# TestFile("lora/test_lora_backend.py", 99), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/13107
# TestFile("lora/test_lora_cuda_graph.py", 250), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/13107
# TestFile("lora/test_lora_qwen3.py", 97), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/13107
TestFile("ops/test_repeat_interleave.py", 75),
TestFile("rotary_embedding/test_mrope.py", 15),
TestFile("test_abort.py", 51),
TestFile("test_bench_typebaseddispatcher.py", 10),
TestFile("test_chunked_prefill.py", 312),