[misc] CI hygiene: enforce __main__ entry, drop silent-skipped tests, fix rerun-test protoc (#23305)
This commit is contained in:
Liangsheng Yin
@@ -138,6 +138,14 @@ jobs:
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v5
|
||||
|
||||
# Needed by setuptools-rust to build the bundled native gRPC extension
|
||||
# (rust/sglang-grpc) when installing the main `sglang` wheel from source.
|
||||
- name: Install protoc
|
||||
run: sudo bash scripts/ci/utils/install_protoc.sh
|
||||
|
||||
- name: Install Rust toolchain
|
||||
run: bash scripts/ci/utils/install_rustup.sh
|
||||
|
||||
- name: Install dependencies
|
||||
timeout-minutes: 20
|
||||
env:
|
||||
|
||||
@@ -1,13 +0,0 @@
|
||||
import pytest
|
||||
|
||||
from sglang.jit_kernel.utils import _REGISTERED_DEPENDENCIES
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=30, suite="stage-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=30, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name", _REGISTERED_DEPENDENCIES.keys())
|
||||
def test_availability(name: str) -> None:
|
||||
# NOTE: the path resolution should not fail
|
||||
_REGISTERED_DEPENDENCIES[name]()
|
||||
@@ -2,7 +2,7 @@ import ast
|
||||
import warnings
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum, auto
|
||||
from typing import List, Optional
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
__all__ = [
|
||||
"HWBackend",
|
||||
@@ -87,6 +87,7 @@ class RegistryVisitor(ast.NodeVisitor):
|
||||
def __init__(self, filename: str):
|
||||
self.filename = filename
|
||||
self.registries: list[CIRegistry] = []
|
||||
self.has_main_entry: bool = False
|
||||
|
||||
def _constant_value(self, node: ast.AST) -> object:
|
||||
if isinstance(node, ast.Constant):
|
||||
@@ -180,25 +181,53 @@ class RegistryVisitor(ast.NodeVisitor):
|
||||
disabled=disabled,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _is_main_block_with_call(stmt: ast.If) -> bool:
|
||||
"""True iff `stmt` is `if __name__ == "__main__":` with a body that
|
||||
contains at least one call (i.e. actually runs something, not just
|
||||
`pass`). This is what makes `python3 file.py` execute tests."""
|
||||
test = stmt.test
|
||||
if not isinstance(test, ast.Compare):
|
||||
return False
|
||||
if not (isinstance(test.left, ast.Name) and test.left.id == "__name__"):
|
||||
return False
|
||||
if len(test.ops) != 1 or not isinstance(test.ops[0], ast.Eq):
|
||||
return False
|
||||
if len(test.comparators) != 1:
|
||||
return False
|
||||
rhs = test.comparators[0]
|
||||
if not (isinstance(rhs, ast.Constant) and rhs.value == "__main__"):
|
||||
return False
|
||||
for child in ast.walk(ast.Module(body=stmt.body, type_ignores=[])):
|
||||
if isinstance(child, ast.Call):
|
||||
return True
|
||||
return False
|
||||
|
||||
def visit_Module(self, node):
|
||||
for stmt in node.body:
|
||||
if not isinstance(stmt, ast.Expr) or not isinstance(stmt.value, ast.Call):
|
||||
continue
|
||||
|
||||
cr = self._collect_ci_registry(stmt.value)
|
||||
if cr is not None:
|
||||
self.registries.append(cr)
|
||||
if isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Call):
|
||||
cr = self._collect_ci_registry(stmt.value)
|
||||
if cr is not None:
|
||||
self.registries.append(cr)
|
||||
elif isinstance(stmt, ast.If) and self._is_main_block_with_call(stmt):
|
||||
self.has_main_entry = True
|
||||
|
||||
self.generic_visit(node)
|
||||
|
||||
|
||||
def ut_parse_one_file(filename: str) -> List[CIRegistry]:
|
||||
def ut_parse_one_file(filename: str) -> Tuple[List[CIRegistry], bool]:
|
||||
"""Parse a test file and return (registries, has_main_entry).
|
||||
|
||||
`has_main_entry` is True iff the file has `if __name__ == "__main__":`
|
||||
with a call in its body -- required for `python3 file.py` to actually
|
||||
run tests (the CI runner's invocation pattern).
|
||||
"""
|
||||
with open(filename, "r") as f:
|
||||
file_content = f.read()
|
||||
tree = ast.parse(file_content, filename=filename)
|
||||
visitor = RegistryVisitor(filename=filename)
|
||||
visitor.visit(tree)
|
||||
return visitor.registries
|
||||
return visitor.registries, visitor.has_main_entry
|
||||
|
||||
|
||||
def auto_partition(files: List[CIRegistry], rank: int, size: int) -> List[CIRegistry]:
|
||||
@@ -230,7 +259,7 @@ def auto_partition(files: List[CIRegistry], rank: int, size: int) -> List[CIRegi
|
||||
def collect_tests(files: list[str], sanity_check: bool = True) -> List[CIRegistry]:
|
||||
ci_tests = []
|
||||
for file in files:
|
||||
registries = ut_parse_one_file(file)
|
||||
registries, has_main_entry = ut_parse_one_file(file)
|
||||
if len(registries) == 0:
|
||||
msg = f"No CI registry found in {file}"
|
||||
if sanity_check:
|
||||
@@ -239,6 +268,20 @@ def collect_tests(files: list[str], sanity_check: bool = True) -> List[CIRegistr
|
||||
warnings.warn(msg)
|
||||
continue
|
||||
|
||||
# Every file with at least one enabled registry must have an
|
||||
# executable `if __name__ == "__main__":` block; otherwise
|
||||
# `python3 file.py -f` (how run_unittest_files invokes tests)
|
||||
# silently exits and the file shows green without running.
|
||||
has_enabled = any(r.disabled is None for r in registries)
|
||||
if sanity_check and has_enabled and not has_main_entry:
|
||||
raise ValueError(
|
||||
f'{file}: missing `if __name__ == "__main__":` entry. '
|
||||
f"Pytest-style tests in this file will silently skip under "
|
||||
f"`python3 file.py -f`. Add `unittest.main()` (for "
|
||||
f"unittest.TestCase) or `sys.exit(pytest.main([__file__, "
|
||||
f'"-v"]))` (for pytest-style).'
|
||||
)
|
||||
|
||||
ci_tests.extend(registries)
|
||||
|
||||
return ci_tests
|
||||
|
||||
@@ -34,7 +34,7 @@ def main() -> int:
|
||||
errors = []
|
||||
for f in files:
|
||||
try:
|
||||
registries = ci_register.ut_parse_one_file(f)
|
||||
registries, _has_main_entry = ci_register.ut_parse_one_file(f)
|
||||
if len(registries) == 0:
|
||||
errors.append(f)
|
||||
except Exception:
|
||||
|
||||
@@ -35,7 +35,7 @@ def collect_all_tests(registered_dir: str) -> list[CIRegistry]:
|
||||
|
||||
for file in sorted(files):
|
||||
try:
|
||||
registries = ut_parse_one_file(file)
|
||||
registries, _ = ut_parse_one_file(file)
|
||||
all_tests.extend(registries)
|
||||
except Exception as e:
|
||||
print(f"Warning: Failed to parse {file}: {e}", file=sys.stderr)
|
||||
|
||||
@@ -1,150 +0,0 @@
|
||||
"""AMD nightly test for Z-Image-Turbo diffusion model (text-to-image)."""
|
||||
|
||||
import io
|
||||
import logging
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from sglang.multimodal_gen.test.server.test_server_common import ( # noqa: F401
|
||||
DiffusionServerBase,
|
||||
diffusion_server,
|
||||
)
|
||||
from sglang.multimodal_gen.test.server.test_server_utils import (
|
||||
ServerContext,
|
||||
get_generate_fn,
|
||||
)
|
||||
from sglang.multimodal_gen.test.server.testcase_configs import (
|
||||
DiffusionSamplingParams,
|
||||
DiffusionServerArgs,
|
||||
DiffusionTestCase,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_amd_ci
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
register_amd_ci(est_time=1800, suite="nightly-amd-1-gpu-zimage-turbo", nightly=True)
|
||||
|
||||
AMD_ZIMAGE_CASES = [
|
||||
DiffusionTestCase(
|
||||
"zimage_image_t2i",
|
||||
DiffusionServerArgs(model_path="Tongyi-MAI/Z-Image-Turbo", modality="image"),
|
||||
DiffusionSamplingParams(
|
||||
prompt="Doraemon is eating dorayaki",
|
||||
output_size="1024x1024",
|
||||
),
|
||||
),
|
||||
]
|
||||
|
||||
CLIP_SCORE_THRESHOLD = 0.20
|
||||
|
||||
|
||||
ARTIFACT_DIR = os.environ.get(
|
||||
"SGLANG_DIFFUSION_ARTIFACT_DIR", "/tmp/diffusion-artifacts"
|
||||
)
|
||||
|
||||
|
||||
def _save_image_and_write_summary(
|
||||
case_id: str, prompt: str, image_bytes: bytes, clip_score: float | None = None
|
||||
):
|
||||
"""Save generated image to artifact dir and write summary."""
|
||||
ext = "jpg" if image_bytes[:2] == b"\xff\xd8" else "png"
|
||||
os.makedirs(ARTIFACT_DIR, exist_ok=True)
|
||||
img_path = os.path.join(ARTIFACT_DIR, f"{case_id}.{ext}")
|
||||
with open(img_path, "wb") as f:
|
||||
f.write(image_bytes)
|
||||
logger.info("Saved image artifact: %s (%d bytes)", img_path, len(image_bytes))
|
||||
|
||||
summary_file = os.environ.get("GITHUB_STEP_SUMMARY")
|
||||
if not summary_file:
|
||||
return
|
||||
|
||||
clip_line = ""
|
||||
if clip_score is not None:
|
||||
status = "PASS" if clip_score >= CLIP_SCORE_THRESHOLD else "FAIL"
|
||||
clip_line = f"| CLIP Score | {clip_score:.4f} ({status}, threshold: {CLIP_SCORE_THRESHOLD}) |\n"
|
||||
|
||||
md = (
|
||||
f"### Z-Image-Turbo — `{case_id}`\n\n"
|
||||
f"| | |\n|---|---|\n"
|
||||
f"| Prompt | {prompt} |\n"
|
||||
f"| Size | {len(image_bytes):,} bytes |\n"
|
||||
f"{clip_line}"
|
||||
f"| Artifact | `{case_id}.{ext}` (download from Artifacts section above) |\n\n"
|
||||
)
|
||||
|
||||
with open(summary_file, "a") as f:
|
||||
f.write(md)
|
||||
|
||||
|
||||
def _compute_clip_score(image_bytes: bytes, prompt: str) -> float | None:
|
||||
"""Compute CLIP cosine similarity between the image and prompt."""
|
||||
try:
|
||||
import torch
|
||||
from PIL import Image
|
||||
from transformers import CLIPModel, CLIPProcessor
|
||||
|
||||
model_name = "openai/clip-vit-base-patch32"
|
||||
processor = CLIPProcessor.from_pretrained(model_name)
|
||||
model = CLIPModel.from_pretrained(model_name)
|
||||
model.eval()
|
||||
|
||||
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
|
||||
inputs = processor(text=[prompt], images=image, return_tensors="pt")
|
||||
|
||||
with torch.no_grad():
|
||||
outputs = model(**inputs)
|
||||
score = outputs.logits_per_image.item() / 100.0
|
||||
|
||||
logger.info("CLIP score for '%s': %.4f", prompt, score)
|
||||
return score
|
||||
except Exception as e:
|
||||
logger.warning("CLIP score computation failed: %s", e)
|
||||
return None
|
||||
|
||||
|
||||
class TestZImageTurboAMD(DiffusionServerBase):
|
||||
"""AMD nightly test for Z-Image-Turbo text-to-image generation."""
|
||||
|
||||
@classmethod
|
||||
def teardown_class(cls):
|
||||
try:
|
||||
super().teardown_class()
|
||||
except AttributeError:
|
||||
pass
|
||||
|
||||
@pytest.fixture(params=AMD_ZIMAGE_CASES, ids=lambda c: c.id)
|
||||
def case(self, request) -> DiffusionTestCase:
|
||||
return request.param
|
||||
|
||||
def test_diffusion_generation(
|
||||
self,
|
||||
case: DiffusionTestCase,
|
||||
diffusion_server: ServerContext,
|
||||
):
|
||||
generate_fn = get_generate_fn(
|
||||
model_path=case.server_args.model_path,
|
||||
modality=case.server_args.modality,
|
||||
sampling_params=case.sampling_params,
|
||||
)
|
||||
|
||||
perf_record, content = self.run_and_collect(
|
||||
diffusion_server, case.id, generate_fn
|
||||
)
|
||||
|
||||
self._validate_and_record(case, perf_record)
|
||||
self._test_v1_models_endpoint(diffusion_server, case)
|
||||
|
||||
prompt = case.sampling_params.prompt or ""
|
||||
clip_score = _compute_clip_score(content, prompt)
|
||||
|
||||
if clip_score is not None:
|
||||
logger.info(
|
||||
"CLIP score: %.4f (threshold: %.2f)", clip_score, CLIP_SCORE_THRESHOLD
|
||||
)
|
||||
assert clip_score >= CLIP_SCORE_THRESHOLD, (
|
||||
f"CLIP score {clip_score:.4f} below threshold {CLIP_SCORE_THRESHOLD} "
|
||||
f"for prompt '{prompt}'"
|
||||
)
|
||||
|
||||
_save_image_and_write_summary(case.id, prompt, content, clip_score)
|
||||
@@ -42,3 +42,9 @@ class TestExpandPreset:
|
||||
"""Unknown preset name raises ValueError."""
|
||||
with pytest.raises(ValueError, match="Unknown value for --preset"):
|
||||
expand_preset(["--preset", "nonexistent"], presets=PRESETS)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
|
||||
@@ -254,3 +254,9 @@ class TestCodePatcher:
|
||||
raise RuntimeError("test error")
|
||||
|
||||
assert obj.greet("world") == "hello world"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
|
||||
@@ -55,3 +55,11 @@ class TestDumperApplySourcePatches:
|
||||
cls.greet.__code__ = original_code
|
||||
|
||||
assert obj.greet("world") == "hello world"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
|
||||
@@ -290,3 +290,9 @@ class TestApplyEdits:
|
||||
]
|
||||
result = apply_edits(source=source, edits=edits)
|
||||
assert result == ("def foo():\n" " x = 1\n" " y = 20\n" " return x\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
|
||||
@@ -1,6 +1,3 @@
|
||||
from argparse import Namespace
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
@@ -9,9 +6,7 @@ from sglang.srt.debug_utils.dump_comparator import (
|
||||
_calc_rel_diff,
|
||||
_compute_smaller_dtype,
|
||||
_try_unify_shape,
|
||||
main,
|
||||
)
|
||||
from sglang.srt.debug_utils.dumper import DumperConfig, _Dumper
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=30, suite="stage-a-test-cpu", nightly=True)
|
||||
@@ -59,101 +54,7 @@ class TestComputeSmallerDtype:
|
||||
assert _compute_smaller_dtype(torch.float32, torch.float32) is None
|
||||
|
||||
|
||||
# ----------------------------- Integration tests -----------------------------
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
|
||||
def _make_dumper(directory: Path) -> _Dumper:
|
||||
return _Dumper(
|
||||
config=DumperConfig(
|
||||
enable=True,
|
||||
dir=str(directory),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _create_dumps(
|
||||
tmp_path: Path,
|
||||
tensor_names: list[str],
|
||||
*,
|
||||
baseline_names: list[str] | None = None,
|
||||
) -> tuple[Path, Path]:
|
||||
if baseline_names is None:
|
||||
baseline_names = tensor_names
|
||||
|
||||
d_baseline: Path = tmp_path / "baseline"
|
||||
d_target: Path = tmp_path / "target"
|
||||
d_baseline.mkdir()
|
||||
d_target.mkdir()
|
||||
|
||||
torch.manual_seed(42)
|
||||
baseline_tensor: torch.Tensor = torch.randn(10, 10)
|
||||
target_tensor: torch.Tensor = baseline_tensor + torch.randn(10, 10) * 0.01
|
||||
|
||||
exp_paths: list[Path] = []
|
||||
for d, names, tensor in [
|
||||
(d_baseline, baseline_names, baseline_tensor),
|
||||
(d_target, tensor_names, target_tensor),
|
||||
]:
|
||||
dumper: _Dumper = _make_dumper(d)
|
||||
for name in names:
|
||||
dumper.dump(name, tensor)
|
||||
dumper.step()
|
||||
exp_paths.append(d / dumper._config.exp_name)
|
||||
|
||||
return exp_paths[0], exp_paths[1]
|
||||
|
||||
|
||||
def _make_args(
|
||||
baseline_path: Path,
|
||||
target_path: Path,
|
||||
*,
|
||||
filter_pattern: str | None = None,
|
||||
) -> Namespace:
|
||||
return Namespace(
|
||||
baseline_path=str(baseline_path),
|
||||
target_path=str(target_path),
|
||||
start_step=0,
|
||||
end_step=1000000,
|
||||
diff_threshold=1e-3,
|
||||
filter=filter_pattern,
|
||||
)
|
||||
|
||||
|
||||
class TestMainBasic:
|
||||
def test_matching_tensors(
|
||||
self, tmp_path: Path, capsys: pytest.CaptureFixture
|
||||
) -> None:
|
||||
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
|
||||
args: Namespace = _make_args(baseline_path, target_path)
|
||||
|
||||
main(args)
|
||||
|
||||
captured: str = capsys.readouterr().out
|
||||
assert "✅" in captured
|
||||
|
||||
def test_with_filter(self, tmp_path: Path, capsys: pytest.CaptureFixture) -> None:
|
||||
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
|
||||
args: Namespace = _make_args(
|
||||
baseline_path, target_path, filter_pattern="tensor_a"
|
||||
)
|
||||
|
||||
main(args)
|
||||
|
||||
captured: str = capsys.readouterr().out
|
||||
assert "tensor_a" in captured
|
||||
assert "Check:" in captured
|
||||
|
||||
def test_no_match_skips(
|
||||
self, tmp_path: Path, capsys: pytest.CaptureFixture
|
||||
) -> None:
|
||||
baseline_path, target_path = _create_dumps(
|
||||
tmp_path,
|
||||
["only_in_target"],
|
||||
baseline_names=["only_in_baseline"],
|
||||
)
|
||||
args: Namespace = _make_args(baseline_path, target_path)
|
||||
|
||||
main(args)
|
||||
|
||||
captured: str = capsys.readouterr().out
|
||||
assert "Skip" in captured
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
|
||||
@@ -1,148 +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
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
|
||||
# Ops - Repeat Interleave tests (1-GPU)
|
||||
|
||||
|
||||
register_cuda_ci(est_time=7, suite="stage-b-test-1-gpu-small")
|
||||
register_amd_ci(est_time=75, suite="stage-b-test-1-gpu-small-amd")
|
||||
|
||||
|
||||
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}"
|
||||
@@ -1,310 +0,0 @@
|
||||
"""
|
||||
Usage:
|
||||
python3 -m unittest test_bnb.TestVisionModel.test_vlm
|
||||
python3 -m unittest test_bnb.TestLanguageModel.test_mmlu
|
||||
"""
|
||||
|
||||
import multiprocessing as mp
|
||||
import random
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from types import SimpleNamespace
|
||||
|
||||
import openai
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
is_in_ci,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=6, suite="stage-b-test-1-gpu-small")
|
||||
|
||||
VISION_MODELS = [
|
||||
"unsloth/Qwen2.5-VL-7B-Instruct-bnb-4bit",
|
||||
"unsloth/Qwen2-VL-7B-Instruct-bnb-4bit",
|
||||
"unsloth/Llama-3.2-11B-Vision-Instruct-bnb-4bit",
|
||||
"unsloth/Llama-3.2-11B-Vision-bnb-4bit",
|
||||
"unsloth/gemma-3-4b-it-bnb-4bit",
|
||||
"unsloth/gemma-3-4b-it-unsloth-bnb-4bit",
|
||||
]
|
||||
LANGUAGE_MODELS = [
|
||||
"unsloth/Qwen2.5-7B-Instruct-bnb-4bit",
|
||||
"unsloth/Qwen2-7B-Instruct-bnb-4bit",
|
||||
"unsloth/Llama-3.2-3B-Instruct-bnb-4bit",
|
||||
"unsloth/gemma-3-1b-it-bnb-4bit",
|
||||
]
|
||||
|
||||
# image
|
||||
IMAGE_MAN_IRONING_URL = "https://raw.githubusercontent.com/sgl-project/sgl-test-files/refs/heads/main/images/man_ironing_on_back_of_suv.png"
|
||||
IMAGE_SGL_LOGO_URL = "https://raw.githubusercontent.com/sgl-project/sgl-test-files/refs/heads/main/images/sgl_logo.png"
|
||||
|
||||
# video
|
||||
VIDEO_JOBS_URL = "https://raw.githubusercontent.com/sgl-project/sgl-test-files/refs/heads/main/videos/jobs_presenting_ipod.mp4"
|
||||
|
||||
# audio
|
||||
AUDIO_TRUMP_SPEECH_URL = "https://raw.githubusercontent.com/sgl-project/sgl-test-files/refs/heads/main/audios/Trump_WEF_2018_10s.mp3"
|
||||
AUDIO_BIRD_SONG_URL = "https://raw.githubusercontent.com/sgl-project/sgl-test-files/refs/heads/main/audios/bird_song.mp3"
|
||||
|
||||
|
||||
def popen_launch_server_wrapper(base_url, model, other_args):
|
||||
process = popen_launch_server(
|
||||
model,
|
||||
base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=other_args,
|
||||
)
|
||||
return process
|
||||
|
||||
|
||||
class TestVisionModel(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
mp.set_start_method("spawn", force=True)
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.base_url += "/v1"
|
||||
cls.api_key = "sk-123456"
|
||||
|
||||
def _run_single_image_chat_completion(self):
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model="default",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": IMAGE_MAN_IRONING_URL},
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Describe this image in a very short sentence.",
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
temperature=0,
|
||||
)
|
||||
|
||||
assert response.choices[0].message.role == "assistant"
|
||||
text = response.choices[0].message.content
|
||||
assert isinstance(text, str)
|
||||
# `driver` is for gemma-3-it
|
||||
assert "man" in text or "person" or "driver" in text, text
|
||||
assert "cab" in text or "taxi" in text or "SUV" in text, text
|
||||
# MiniCPMO fails to recognize `iron`, but `hanging`
|
||||
assert "iron" in text or "hang" in text, text
|
||||
assert response.id
|
||||
assert response.created
|
||||
assert response.usage.prompt_tokens > 0
|
||||
assert response.usage.completion_tokens > 0
|
||||
assert response.usage.total_tokens > 0
|
||||
|
||||
def _run_multi_turn_chat_completion(self):
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model="default",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": IMAGE_MAN_IRONING_URL},
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Describe this image in a very short sentence.",
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "There is a man at the back of a yellow cab ironing his clothes.",
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Repeat your previous answer."}
|
||||
],
|
||||
},
|
||||
],
|
||||
temperature=0,
|
||||
)
|
||||
|
||||
assert response.choices[0].message.role == "assistant"
|
||||
text = response.choices[0].message.content
|
||||
assert isinstance(text, str)
|
||||
assert "man" in text or "cab" in text, text
|
||||
assert response.id
|
||||
assert response.created
|
||||
assert response.usage.prompt_tokens > 0
|
||||
assert response.usage.completion_tokens > 0
|
||||
assert response.usage.total_tokens > 0
|
||||
|
||||
def _run_multi_images_chat_completion(self):
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
response = client.chat.completions.create(
|
||||
model="default",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": IMAGE_MAN_IRONING_URL},
|
||||
"modalities": "multi-images",
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": IMAGE_SGL_LOGO_URL},
|
||||
"modalities": "multi-images",
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "I have two very different images. They are not related at all. "
|
||||
"Please describe the first image in one sentence, and then describe the second image in another sentence.",
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
temperature=0,
|
||||
)
|
||||
|
||||
assert response.choices[0].message.role == "assistant"
|
||||
text = response.choices[0].message.content
|
||||
assert isinstance(text, str)
|
||||
print("-" * 30)
|
||||
print(f"Multi images response:\n{text}")
|
||||
print("-" * 30)
|
||||
assert "man" in text or "cab" in text or "SUV" in text or "taxi" in text, text
|
||||
assert "logo" in text or '"S"' in text or "SG" in text, text
|
||||
assert response.id
|
||||
assert response.created
|
||||
assert response.usage.prompt_tokens > 0
|
||||
assert response.usage.completion_tokens > 0
|
||||
assert response.usage.total_tokens > 0
|
||||
|
||||
def run_decode_with_image(self, image_id):
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
content = []
|
||||
if image_id == 0:
|
||||
content.append(
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": IMAGE_MAN_IRONING_URL},
|
||||
}
|
||||
)
|
||||
elif image_id == 1:
|
||||
content.append(
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": IMAGE_SGL_LOGO_URL},
|
||||
}
|
||||
)
|
||||
else:
|
||||
pass
|
||||
|
||||
content.append(
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Describe this image in a very short sentence.",
|
||||
}
|
||||
)
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model="default",
|
||||
messages=[
|
||||
{"role": "user", "content": content},
|
||||
],
|
||||
temperature=0,
|
||||
)
|
||||
|
||||
assert response.choices[0].message.role == "assistant"
|
||||
text = response.choices[0].message.content
|
||||
assert isinstance(text, str)
|
||||
|
||||
def _run_test_mixed_batch(self):
|
||||
image_ids = [0, 1, 2] * 4
|
||||
with ThreadPoolExecutor(4) as executor:
|
||||
list(executor.map(self.run_decode_with_image, image_ids))
|
||||
|
||||
def test_vlm(self):
|
||||
models_to_test = VISION_MODELS
|
||||
|
||||
if is_in_ci():
|
||||
models_to_test = [random.choice(VISION_MODELS)]
|
||||
|
||||
for model in models_to_test:
|
||||
with self.subTest(model=model):
|
||||
other_args = [
|
||||
"--mem-fraction-static",
|
||||
"0.6",
|
||||
"--load-format",
|
||||
"bitsandbytes",
|
||||
"--enable-multimodal",
|
||||
]
|
||||
try:
|
||||
process = popen_launch_server_wrapper(
|
||||
DEFAULT_URL_FOR_TEST, model, other_args
|
||||
)
|
||||
self._run_test_mixed_batch()
|
||||
self._run_multi_images_chat_completion()
|
||||
self._run_multi_turn_chat_completion()
|
||||
self._run_single_image_chat_completion()
|
||||
finally:
|
||||
kill_process_tree(process.pid)
|
||||
|
||||
|
||||
class TestLanguageModel(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
mp.set_start_method("spawn", force=True)
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
# cls.base_url += "/v1"
|
||||
cls.api_key = "sk-123456"
|
||||
|
||||
def test_mmlu(self):
|
||||
models_to_test = LANGUAGE_MODELS
|
||||
|
||||
if is_in_ci():
|
||||
models_to_test = [random.choice(LANGUAGE_MODELS)]
|
||||
|
||||
for model in models_to_test:
|
||||
with self.subTest(model=model):
|
||||
other_args = [
|
||||
"--mem-fraction-static",
|
||||
"0.6",
|
||||
"--load-format",
|
||||
"bitsandbytes",
|
||||
]
|
||||
try:
|
||||
process = popen_launch_server_wrapper(
|
||||
DEFAULT_URL_FOR_TEST, model, other_args
|
||||
)
|
||||
args = SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=model,
|
||||
eval_name="mmlu",
|
||||
num_examples=32,
|
||||
num_threads=16,
|
||||
)
|
||||
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
self.assertGreater(metrics["score"], 0.3)
|
||||
finally:
|
||||
kill_process_tree(process.pid)
|
||||
@@ -1,59 +0,0 @@
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_amd_ci
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
register_amd_ci(est_time=313, suite="stage-b-test-1-gpu-small-amd")
|
||||
|
||||
|
||||
class TestMixtralAccuracy(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = "mistralai/Mixtral-8x7B-Instruct-v0.1"
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
|
||||
other_args = [
|
||||
"--tp",
|
||||
"2",
|
||||
"--mem-fraction-static",
|
||||
"0.9",
|
||||
"--context-length",
|
||||
"38768",
|
||||
"--quantization",
|
||||
"quark_int4fp8_moe",
|
||||
# The default aiter attention backend raises segmentation faults and other errors - as quark_int4fp8_moe is not related to attention, let's just use triton here.
|
||||
"--attention-backend",
|
||||
"triton",
|
||||
]
|
||||
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=45 * 60,
|
||||
other_args=other_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=self.model,
|
||||
eval_name="gsm8k",
|
||||
api="completion",
|
||||
max_tokens=512,
|
||||
num_examples=1400,
|
||||
num_threads=128,
|
||||
num_shots=8,
|
||||
)
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
self.assertGreater(metrics["score"], 0.56)
|
||||
@@ -1,174 +0,0 @@
|
||||
# Rotary Embedding - MRoPE tests (1-GPU)
|
||||
|
||||
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,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=7, suite="stage-b-test-1-gpu-large")
|
||||
register_amd_ci(est_time=15, suite="stage-b-test-1-gpu-small-amd")
|
||||
|
||||
_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]
|
||||
|
||||
|
||||
def create_yarn_rope_scaling(original_config, scaling_factor=2.0):
|
||||
yarn_config = {
|
||||
"rope_type": "yarn",
|
||||
"factor": scaling_factor,
|
||||
"original_max_position_embeddings": original_config.max_position_embeddings,
|
||||
}
|
||||
if hasattr(original_config, "rope_scaling") and original_config.rope_scaling:
|
||||
if "mrope_section" in original_config.rope_scaling:
|
||||
yarn_config["mrope_section"] = original_config.rope_scaling["mrope_section"]
|
||||
if "mrope_interleaved" in original_config.rope_scaling:
|
||||
yarn_config["mrope_interleaved"] = original_config.rope_scaling[
|
||||
"mrope_interleaved"
|
||||
]
|
||||
return yarn_config
|
||||
|
||||
|
||||
@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)
|
||||
@pytest.mark.parametrize(
|
||||
"rope_scaling_type", ["default", "yarn"], ids=["mrope_default", "mrope_yarn"]
|
||||
)
|
||||
def test_mrope(
|
||||
model_name: str,
|
||||
model_info: MRoPETestInfo,
|
||||
tp_size: int,
|
||||
dtype: torch.dtype,
|
||||
num_tokens: int,
|
||||
rope_scaling_type: str,
|
||||
):
|
||||
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)
|
||||
|
||||
if rope_scaling_type == "yarn":
|
||||
rope_scaling_config = create_yarn_rope_scaling(config, scaling_factor=2.0)
|
||||
else:
|
||||
rope_scaling_config = config.rope_scaling
|
||||
|
||||
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=rope_scaling_config,
|
||||
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)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -4,7 +4,6 @@ import torch
|
||||
|
||||
from sglang.srt.managers.schedule_batch import FINISH_ABORT
|
||||
from sglang.srt.mem_cache.base_prefix_cache import MatchResult
|
||||
from sglang.srt.mem_cache.common import release_kv_cache
|
||||
from sglang.srt.session.streaming_session import SessionSlot, StreamingSession
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
@@ -76,6 +75,7 @@ class _FakeReq:
|
||||
self.pop_overallocated_calls = 0
|
||||
self.to_finish = None
|
||||
self.finished_reason = None
|
||||
self.finished_len = None
|
||||
|
||||
def pop_committed_kv_cache(self):
|
||||
assert not self.kv_committed_freed
|
||||
@@ -89,36 +89,6 @@ class _FakeReq:
|
||||
return self.kv_committed_len, self.kv_allocated_len
|
||||
|
||||
|
||||
def test_streaming_release_kv_cache_defers_tail_free(monkeypatch):
|
||||
"""Spec tail is NOT trimmed in cache_finished_req; it is deferred to
|
||||
match_prefix's orphan tail free on the next turn. cache_finished_req
|
||||
only sets bookkeeping flags and saves the slot as-is."""
|
||||
page_size = 16
|
||||
req_to_token = torch.arange(128, dtype=torch.int32).reshape(1, 128)
|
||||
req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[])
|
||||
allocator = _FakeAllocator()
|
||||
tree_cache = StreamingSession(
|
||||
_FakeInnerCache(req_to_token_pool, allocator, page_size)
|
||||
)
|
||||
req = _FakeReq("session-a", req_pool_idx=0, committed=17, allocated=40)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"sglang.srt.mem_cache.common.get_global_server_args",
|
||||
lambda: SimpleNamespace(page_size=page_size, speculative_algorithm="eagle"),
|
||||
)
|
||||
|
||||
release_kv_cache(req, tree_cache)
|
||||
|
||||
slot = tree_cache.slots["session-a"]
|
||||
assert req.kv_committed_freed is True
|
||||
assert req.kv_overallocated_freed is True
|
||||
assert req.req_pool_idx is None
|
||||
# Slot keeps the full allocation — tail free is deferred to match_prefix.
|
||||
assert slot.kv_committed_len == 17
|
||||
assert slot.kv_allocated_len == 40
|
||||
assert len(allocator.freed) == 0
|
||||
|
||||
|
||||
def test_preabort_detaches_session_and_preserves_slot():
|
||||
"""Pre-aborted req (to_finish set before match_prefix) is detached from
|
||||
the session: session=None, abort_req() called. Slot stays intact."""
|
||||
@@ -271,3 +241,11 @@ def test_trim_overshoot_postcondition():
|
||||
# Tail [38, 44) freed by _free_kv_aligned.
|
||||
assert len(allocator.freed) == 1
|
||||
assert allocator.freed[0].tolist() == list(range(38, 44))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
|
||||
@@ -1,156 +0,0 @@
|
||||
import argparse
|
||||
import json
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from sglang.srt.model_executor.hook_manager import register_forward_hooks
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=6, suite="stage-b-test-1-gpu-small")
|
||||
register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd")
|
||||
|
||||
HOOK_CALLS = []
|
||||
|
||||
|
||||
def dummy_hook_factory(config):
|
||||
"""Factory that returns a forward hook capturing a tag from config."""
|
||||
tag = config.get("tag", "default")
|
||||
|
||||
def hook(module, inputs, output):
|
||||
HOOK_CALLS.append(
|
||||
{
|
||||
"module_type": type(module).__name__,
|
||||
"tag": tag,
|
||||
"shape": tuple(output.shape),
|
||||
}
|
||||
)
|
||||
return output
|
||||
|
||||
return hook
|
||||
|
||||
|
||||
class TinyModel(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.inner = nn.Sequential(
|
||||
nn.Linear(4, 2),
|
||||
nn.ReLU(),
|
||||
)
|
||||
self.outer = nn.Sequential(
|
||||
nn.Linear(4, 4),
|
||||
nn.ReLU(),
|
||||
self.inner,
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.outer(x)
|
||||
|
||||
|
||||
class TestAttachHooks(CustomTestCase):
|
||||
"""Tests for register_forward_hooks / resolve_callable integration."""
|
||||
|
||||
def setUp(self):
|
||||
HOOK_CALLS.clear()
|
||||
|
||||
def test_hook_is_attached(self):
|
||||
"""Hook from a factory string is registered and fired."""
|
||||
hook_specs = [
|
||||
{
|
||||
"target_modules": ["outer.0", "outer.1"],
|
||||
"hook_factory": "test_model_hooks:dummy_hook_factory",
|
||||
"config": {"tag": "forward-ok"},
|
||||
},
|
||||
{
|
||||
"target_modules": ["inner.*"],
|
||||
"hook_factory": "test_model_hooks:dummy_hook_factory",
|
||||
"config": {"tag": "forward-ok"},
|
||||
},
|
||||
]
|
||||
|
||||
model = TinyModel()
|
||||
register_forward_hooks(model, hook_specs)
|
||||
|
||||
x = torch.randn(3, 4)
|
||||
_ = model(x)
|
||||
|
||||
self.assertEqual(
|
||||
len(HOOK_CALLS),
|
||||
4,
|
||||
"Forward hook was not called correct number of times",
|
||||
)
|
||||
tags = {call["tag"] for call in HOOK_CALLS}
|
||||
self.assertIn("forward-ok", tags)
|
||||
|
||||
def test_no_matching_modules_does_not_crash(self):
|
||||
"""Hook spec with no matching modules should not crash."""
|
||||
model = TinyModel()
|
||||
hook_specs = [
|
||||
{
|
||||
"name": "no_match",
|
||||
"target_modules": ["does_not_exist.*"],
|
||||
"hook_factory": "test_model_hooks:dummy_hook_factory",
|
||||
"config": {"tag": "unused"},
|
||||
}
|
||||
]
|
||||
|
||||
register_forward_hooks(model, hook_specs)
|
||||
|
||||
x = torch.randn(3, 4)
|
||||
_ = model(x)
|
||||
|
||||
# No hooks should have fired
|
||||
self.assertEqual(len(HOOK_CALLS), 0)
|
||||
|
||||
def test_cli_hooks_reach_model(self):
|
||||
"""
|
||||
Ensure that when hooks are provided via CLI, they are parsed into
|
||||
ServerArgs, passed to register_forward_hooks, and actually
|
||||
run during a forward pass.
|
||||
"""
|
||||
parser = argparse.ArgumentParser()
|
||||
ServerArgs.add_cli_args(parser)
|
||||
|
||||
hooks_spec = [
|
||||
{
|
||||
"name": "outer_and_inner_from_cli",
|
||||
"target_modules": ["outer.0", "outer.1", "inner.*"],
|
||||
"hook_factory": "test_model_hooks:dummy_hook_factory",
|
||||
"config": {"tag": "cli-hook"},
|
||||
}
|
||||
]
|
||||
|
||||
cli_args = [
|
||||
"--model-path",
|
||||
"Qwen/Qwen2-7B-Instruct", # Dummy value; not used in this test
|
||||
"--forward-hooks",
|
||||
json.dumps(hooks_spec),
|
||||
]
|
||||
|
||||
args = parser.parse_args(cli_args)
|
||||
server_args = ServerArgs.from_cli_args(args)
|
||||
|
||||
self.assertEqual(server_args.forward_hooks, hooks_spec)
|
||||
|
||||
model = TinyModel()
|
||||
register_forward_hooks(model, server_args.forward_hooks)
|
||||
|
||||
x = torch.randn(3, 4)
|
||||
_ = model(x)
|
||||
|
||||
# We expect hooks on outer.0, outer.1, inner.0, inner.1 => 4 calls
|
||||
self.assertEqual(
|
||||
len(HOOK_CALLS),
|
||||
4,
|
||||
"CLI-configured hooks did not fire expected number of times",
|
||||
)
|
||||
|
||||
tags = {call["tag"] for call in HOOK_CALLS}
|
||||
self.assertEqual(tags, {"cli-hook"})
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pass
|
||||
# unittest.main()
|
||||
@@ -1,166 +0,0 @@
|
||||
import os
|
||||
import statistics
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from sglang.srt.models.glm4v import Glm4vVisionPatchEmbed
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=8, suite="stage-b-test-1-gpu-large")
|
||||
|
||||
PATCH_SIZE = 14
|
||||
TEMPORAL_PATCH_SIZE = 2
|
||||
IN_CHANNELS = 3
|
||||
HIDDEN_SIZE = 1536
|
||||
FLAT_DIM = IN_CHANNELS * TEMPORAL_PATCH_SIZE * PATCH_SIZE * PATCH_SIZE
|
||||
|
||||
|
||||
class ReferenceConv3dPatchEmbed(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
patch_size=PATCH_SIZE,
|
||||
temporal_patch_size=TEMPORAL_PATCH_SIZE,
|
||||
in_channels=IN_CHANNELS,
|
||||
hidden_size=HIDDEN_SIZE,
|
||||
):
|
||||
super().__init__()
|
||||
self.patch_size = patch_size
|
||||
self.temporal_patch_size = temporal_patch_size
|
||||
self.in_channels = in_channels
|
||||
self.hidden_size = hidden_size
|
||||
|
||||
kernel_size = (temporal_patch_size, patch_size, patch_size)
|
||||
self.proj = nn.Conv3d(
|
||||
in_channels,
|
||||
hidden_size,
|
||||
kernel_size=kernel_size,
|
||||
stride=kernel_size,
|
||||
bias=True,
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = x.view(
|
||||
-1,
|
||||
self.in_channels,
|
||||
self.temporal_patch_size,
|
||||
self.patch_size,
|
||||
self.patch_size,
|
||||
)
|
||||
x = self.proj(x).view(-1, self.hidden_size)
|
||||
return x
|
||||
|
||||
|
||||
def _build_modules(device: str, dtype: torch.dtype):
|
||||
conv_mod = ReferenceConv3dPatchEmbed().to(device=device, dtype=dtype).eval()
|
||||
linear_mod = (
|
||||
Glm4vVisionPatchEmbed(
|
||||
patch_size=PATCH_SIZE,
|
||||
temporal_patch_size=TEMPORAL_PATCH_SIZE,
|
||||
in_channels=IN_CHANNELS,
|
||||
hidden_size=HIDDEN_SIZE,
|
||||
)
|
||||
.to(device=device, dtype=dtype)
|
||||
.eval()
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
linear_mod.proj.weight.copy_(conv_mod.proj.weight)
|
||||
linear_mod.proj.bias.copy_(conv_mod.proj.bias)
|
||||
|
||||
linear_mod.copy_conv3d_weight_to_linear()
|
||||
|
||||
return conv_mod, linear_mod
|
||||
|
||||
|
||||
def _benchmark_cuda_module(
|
||||
module: nn.Module,
|
||||
x: torch.Tensor,
|
||||
warmup: int = 50,
|
||||
inner_iters: int = 200,
|
||||
repeats: int = 10,
|
||||
) -> float:
|
||||
assert x.is_cuda
|
||||
module.eval()
|
||||
|
||||
with torch.inference_mode():
|
||||
for _ in range(warmup):
|
||||
module(x)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
samples = []
|
||||
for _ in range(repeats):
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
|
||||
start.record()
|
||||
for _ in range(inner_iters):
|
||||
module(x)
|
||||
end.record()
|
||||
|
||||
torch.cuda.synchronize()
|
||||
samples.append(start.elapsed_time(end) / inner_iters)
|
||||
|
||||
return statistics.median(samples)
|
||||
|
||||
|
||||
def test_patch_embed_linear_matches_conv3d():
|
||||
torch.manual_seed(0)
|
||||
|
||||
device = "cpu"
|
||||
dtype = torch.float32
|
||||
|
||||
conv_mod, linear_mod = _build_modules(device=device, dtype=dtype)
|
||||
|
||||
x = torch.randn(512, FLAT_DIM, device=device, dtype=dtype)
|
||||
|
||||
with torch.inference_mode():
|
||||
y_conv = conv_mod(x)
|
||||
y_linear = linear_mod(x)
|
||||
|
||||
torch.testing.assert_close(
|
||||
y_conv,
|
||||
y_linear,
|
||||
rtol=1e-5,
|
||||
atol=1e-5,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not torch.cuda.is_available(), reason="CUDA is required for perf benchmark"
|
||||
)
|
||||
def test_patch_embed_linear_conv3d():
|
||||
torch.manual_seed(0)
|
||||
torch.backends.cudnn.benchmark = True
|
||||
|
||||
device = "cuda"
|
||||
dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
|
||||
|
||||
conv_mod, linear_mod = _build_modules(device=device, dtype=dtype)
|
||||
|
||||
num_patches = int(os.getenv("GLM4V_NUM_PATCHES", "4096"))
|
||||
warmup = int(os.getenv("GLM4V_WARMUP", "50"))
|
||||
inner_iters = int(os.getenv("GLM4V_INNER_ITERS", "200"))
|
||||
repeats = int(os.getenv("GLM4V_REPEATS", "10"))
|
||||
|
||||
x = torch.randn(num_patches, FLAT_DIM, device=device, dtype=dtype).contiguous()
|
||||
|
||||
conv_ms = _benchmark_cuda_module(
|
||||
conv_mod, x, warmup=warmup, inner_iters=inner_iters, repeats=repeats
|
||||
)
|
||||
linear_ms = _benchmark_cuda_module(
|
||||
linear_mod, x, warmup=warmup, inner_iters=inner_iters, repeats=repeats
|
||||
)
|
||||
|
||||
speedup = conv_ms / linear_ms
|
||||
print(
|
||||
f"\n[patch_embed perf] conv3d={conv_ms:.4f} ms | "
|
||||
f"linear={linear_ms:.4f} ms | speedup={speedup:.3f}x"
|
||||
)
|
||||
|
||||
min_speedup = float(os.getenv("GLM4V_MIN_SPEEDUP", "1.00"))
|
||||
assert speedup >= min_speedup, (
|
||||
f"Expected speedup >= {min_speedup:.3f}x, but got {speedup:.3f}x "
|
||||
f"(conv3d={conv_ms:.4f} ms, linear={linear_ms:.4f} ms)"
|
||||
)
|
||||
Reference in New Issue
Block a user