[CI] Derive registered-test kind from the registry call instead of the path (#40294)
This commit is contained in:
Liangsheng Yin
@@ -1,315 +0,0 @@
|
||||
"""
|
||||
End-to-end tests for strict reasoning + constrained decoding.
|
||||
|
||||
Tests that the full pipeline works:
|
||||
- AC-5.1: Strict reasoning + JSON schema constrained generation
|
||||
- AC-5.2: Strict reasoning + tool call parsing (basic validation only)
|
||||
|
||||
These tests launch a real server with a small model and verify
|
||||
the constrained decoding pipeline produces valid output.
|
||||
"""
|
||||
|
||||
import json
|
||||
import unittest
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=96, stage="base-b", runner_config="1-gpu-small")
|
||||
register_amd_ci(est_time=120, suite="stage-b-test-1-gpu-small-amd")
|
||||
|
||||
MODEL = "Qwen/Qwen3-0.6B"
|
||||
BASE_URL = "http://127.0.0.1:39877"
|
||||
API_KEY = "sk-test-1234"
|
||||
|
||||
|
||||
class TestConstrainedReasoningE2E(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = MODEL
|
||||
cls.base_url = BASE_URL
|
||||
cls.api_key = API_KEY
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
api_key=cls.api_key,
|
||||
other_args=[
|
||||
"--reasoning-parser",
|
||||
"qwen3",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def _chat(self, **kwargs):
|
||||
default = {
|
||||
"model": self.model,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "What is 2+2? Answer with just the number.",
|
||||
}
|
||||
],
|
||||
"temperature": 0,
|
||||
"max_tokens": 256,
|
||||
}
|
||||
default.update(kwargs)
|
||||
resp = requests.post(
|
||||
f"{self.base_url}/v1/chat/completions",
|
||||
headers={"Authorization": f"Bearer {self.api_key}"},
|
||||
json=default,
|
||||
timeout=60,
|
||||
)
|
||||
self.assertEqual(resp.status_code, 200, f"Request failed: {resp.text}")
|
||||
return resp.json()
|
||||
|
||||
def test_reasoning_with_json_schema(self):
|
||||
"""AC-5.1: Reasoning + JSON schema produces valid JSON output."""
|
||||
schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"answer": {"type": "integer"},
|
||||
},
|
||||
"required": ["answer"],
|
||||
}
|
||||
data = self._chat(
|
||||
response_format={
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": "answer_schema",
|
||||
"schema": schema,
|
||||
},
|
||||
},
|
||||
chat_template_kwargs={"enable_thinking": True},
|
||||
separate_reasoning=True,
|
||||
)
|
||||
|
||||
choice = data["choices"][0]
|
||||
content = choice["message"]["content"] or ""
|
||||
|
||||
# Content should be valid JSON conforming to schema when non-empty.
|
||||
# With small models + separate_reasoning, content may be empty if the
|
||||
# model puts everything in reasoning_content. That's acceptable.
|
||||
if content.strip():
|
||||
try:
|
||||
parsed = json.loads(content)
|
||||
self.assertIn("answer", parsed)
|
||||
self.assertIsInstance(parsed["answer"], int)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
# Small models may produce imperfect JSON
|
||||
self.assertTrue(
|
||||
content.strip().startswith("{"),
|
||||
f"Expected JSON-like output, got: {content!r}",
|
||||
)
|
||||
|
||||
# Content should NOT contain <think> tags (those go to reasoning_content)
|
||||
self.assertNotIn("<think>", content)
|
||||
|
||||
def test_reasoning_disabled_with_json_schema(self):
|
||||
"""JSON schema still works when reasoning is explicitly disabled."""
|
||||
schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"answer": {"type": "integer"},
|
||||
},
|
||||
"required": ["answer"],
|
||||
}
|
||||
data = self._chat(
|
||||
response_format={
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": "answer_schema",
|
||||
"schema": schema,
|
||||
},
|
||||
},
|
||||
chat_template_kwargs={"enable_thinking": False},
|
||||
)
|
||||
|
||||
choice = data["choices"][0]
|
||||
content = choice["message"]["content"]
|
||||
|
||||
# Should still produce valid JSON
|
||||
parsed = json.loads(content)
|
||||
self.assertIn("answer", parsed)
|
||||
|
||||
def test_reasoning_with_separate_output(self):
|
||||
"""Reasoning content is correctly separated from normal content."""
|
||||
data = self._chat(
|
||||
chat_template_kwargs={"enable_thinking": True},
|
||||
separate_reasoning=True,
|
||||
)
|
||||
|
||||
choice = data["choices"][0]
|
||||
content = choice["message"]["content"]
|
||||
reasoning = choice["message"].get("reasoning_content")
|
||||
|
||||
# Content should not contain think tags
|
||||
self.assertNotIn("<think>", content)
|
||||
self.assertNotIn("</think>", content)
|
||||
|
||||
def test_tool_call_after_reasoning(self):
|
||||
"""AC-5.2: Tool call parsing works with reasoning enabled."""
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get the current weather",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": {"type": "string"},
|
||||
},
|
||||
"required": ["location"],
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
data = self._chat(
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": "What's the weather in Paris?",
|
||||
}
|
||||
],
|
||||
tools=tools,
|
||||
chat_template_kwargs={"enable_thinking": True},
|
||||
separate_reasoning=True,
|
||||
)
|
||||
|
||||
choice = data["choices"][0]
|
||||
# The model may or may not produce tool calls (depends on model capability)
|
||||
# but the response should be well-formed (no crashes)
|
||||
self.assertIn("message", choice)
|
||||
self.assertIn("finish_reason", choice)
|
||||
# finish_reason should be either "stop" or "tool_calls"
|
||||
self.assertIn(choice["finish_reason"], ["stop", "tool_calls", "length"])
|
||||
|
||||
|
||||
class TestStrictThinkingE2E(CustomTestCase):
|
||||
"""E2E tests with --enable-strict-thinking flag.
|
||||
|
||||
Validates that the strict thinking flag is correctly propagated through
|
||||
the full pipeline: server_args -> grammar_backend -> ReasonerGrammarBackend
|
||||
-> token filtering during thinking phase.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = MODEL
|
||||
cls.base_url = "http://127.0.0.1:39878"
|
||||
cls.api_key = API_KEY
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
api_key=cls.api_key,
|
||||
other_args=[
|
||||
"--reasoning-parser",
|
||||
"qwen3",
|
||||
"--enable-strict-thinking",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def _chat(self, **kwargs):
|
||||
default = {
|
||||
"model": self.model,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "What is 2+2? Answer with just the number.",
|
||||
}
|
||||
],
|
||||
"temperature": 0,
|
||||
"max_tokens": 256,
|
||||
}
|
||||
default.update(kwargs)
|
||||
resp = requests.post(
|
||||
f"{self.base_url}/v1/chat/completions",
|
||||
headers={"Authorization": f"Bearer {self.api_key}"},
|
||||
json=default,
|
||||
timeout=60,
|
||||
)
|
||||
self.assertEqual(resp.status_code, 200, f"Request failed: {resp.text}")
|
||||
return resp.json()
|
||||
|
||||
def test_strict_thinking_with_json_schema(self):
|
||||
"""Strict thinking + JSON schema: server starts and produces valid output."""
|
||||
schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"answer": {"type": "integer"},
|
||||
},
|
||||
"required": ["answer"],
|
||||
}
|
||||
data = self._chat(
|
||||
response_format={
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": "answer_schema",
|
||||
"schema": schema,
|
||||
},
|
||||
},
|
||||
chat_template_kwargs={"enable_thinking": True},
|
||||
separate_reasoning=True,
|
||||
)
|
||||
|
||||
choice = data["choices"][0]
|
||||
content = choice["message"]["content"] or ""
|
||||
|
||||
if content.strip():
|
||||
try:
|
||||
parsed = json.loads(content)
|
||||
self.assertIn("answer", parsed)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
self.assertTrue(
|
||||
content.strip().startswith("{"),
|
||||
f"Expected JSON-like output, got: {content!r}",
|
||||
)
|
||||
|
||||
# Think tags must not leak into content
|
||||
self.assertNotIn("<think>", content)
|
||||
|
||||
def test_strict_thinking_disabled_per_request(self):
|
||||
"""When thinking is disabled per-request, strict server still works."""
|
||||
data = self._chat(
|
||||
chat_template_kwargs={"enable_thinking": False},
|
||||
)
|
||||
|
||||
choice = data["choices"][0]
|
||||
self.assertIn("message", choice)
|
||||
self.assertIn("finish_reason", choice)
|
||||
# Should complete normally without errors
|
||||
self.assertIn(choice["finish_reason"], ["stop", "length"])
|
||||
|
||||
def test_strict_thinking_separate_reasoning(self):
|
||||
"""Strict thinking with separate_reasoning produces well-formed output."""
|
||||
data = self._chat(
|
||||
chat_template_kwargs={"enable_thinking": True},
|
||||
separate_reasoning=True,
|
||||
)
|
||||
|
||||
choice = data["choices"][0]
|
||||
content = choice["message"]["content"] or ""
|
||||
|
||||
# Think tags must not leak into content
|
||||
self.assertNotIn("<think>", content)
|
||||
self.assertNotIn("</think>", content)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,515 @@
|
||||
"""Dispatch-contract tests for the unified ``BaseFusedOp`` (RFC #29630, #26426).
|
||||
|
||||
``BaseFusedOp`` replaced ``MultiPlatformOp`` as the single operator
|
||||
abstraction; these tests pin down the parts of that contract that a refactor
|
||||
could silently break:
|
||||
|
||||
- the priority ladder: explicit ``backend=`` > global forced backend > OOT
|
||||
platform override > declared optimized kernel backends > platform-specific
|
||||
forward > native fallback;
|
||||
- the standard ``nn.Module`` behavior (hooks, traversal);
|
||||
- static-dispatch caching and per-call ``backend_eligible`` gating;
|
||||
- the torch.compile enter/leave protocol (idempotency, TopK / FusedMoE
|
||||
special paths);
|
||||
- the deprecated ``MultiPlatformOp`` alias and its OOT plugin surface.
|
||||
|
||||
Platform detection is mocked, so everything here runs on a CPU-only box.
|
||||
"""
|
||||
|
||||
import warnings
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
import sglang.kernels.fused_op as fo
|
||||
from sglang.kernels.fused_op import BaseFusedOp
|
||||
from sglang.kernels.spec import CapabilityRequirement as Cap
|
||||
from sglang.kernels.spec import KernelBackend, PlatformInfo
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=12, suite="base-a-test-cpu")
|
||||
|
||||
_CUDA = PlatformInfo(device_type="cuda", cuda_arch_major=9, cuda_arch_minor=0)
|
||||
_HIP = PlatformInfo(device_type="hip")
|
||||
_CPU = PlatformInfo()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_global_state():
|
||||
saved_oot = {k: dict(v) for k, v in BaseFusedOp._oot_forward_registry.items()}
|
||||
yield
|
||||
fo.set_fused_op_backend(None)
|
||||
fo.disable_fused_op_trace()
|
||||
fo.clear_fused_op_trace()
|
||||
BaseFusedOp._oot_forward_registry.clear()
|
||||
BaseFusedOp._oot_forward_registry.update(saved_oot)
|
||||
|
||||
|
||||
def _mock_platform(monkeypatch, *, key="", info=_CPU, oot_key=None):
|
||||
monkeypatch.setattr(fo, "_platform_key", lambda: key)
|
||||
monkeypatch.setattr(fo, "_platform", lambda: info)
|
||||
monkeypatch.setattr(fo, "_oot_dispatch_key", lambda: oot_key)
|
||||
|
||||
|
||||
class _AllPlatformsOp(BaseFusedOp):
|
||||
"""Marks which path ran by returning its name."""
|
||||
|
||||
op = "test.all_platforms"
|
||||
|
||||
def forward_native(self, x):
|
||||
return "native"
|
||||
|
||||
def forward_cuda(self, x):
|
||||
return "cuda"
|
||||
|
||||
def forward_hip(self, x):
|
||||
return "hip"
|
||||
|
||||
def forward_npu(self, x):
|
||||
return "npu"
|
||||
|
||||
def forward_xpu(self, x):
|
||||
return "xpu"
|
||||
|
||||
def forward_musa(self, x):
|
||||
return "musa"
|
||||
|
||||
def forward_cpu(self, x):
|
||||
return "cpu"
|
||||
|
||||
|
||||
class _CudaOnlyPlatformOp(BaseFusedOp):
|
||||
op = "test.cuda_only_platform"
|
||||
|
||||
def forward_native(self, x):
|
||||
return "native"
|
||||
|
||||
def forward_cuda(self, x):
|
||||
return "cuda"
|
||||
|
||||
|
||||
class _NativeOnlyOp(BaseFusedOp):
|
||||
op = "test.native_only"
|
||||
|
||||
def forward_native(self, x):
|
||||
return "native"
|
||||
|
||||
|
||||
class _BackendAndPlatformOp(BaseFusedOp):
|
||||
"""Declared JIT backend + a CUDA platform forward."""
|
||||
|
||||
op = "test.backend_and_platform"
|
||||
priority = (KernelBackend.JIT, KernelBackend.TORCH)
|
||||
capabilities = {KernelBackend.JIT: frozenset({Cap.CUDA})}
|
||||
|
||||
def forward_native(self, x):
|
||||
return "native"
|
||||
|
||||
def forward_jit(self, x):
|
||||
return "jit"
|
||||
|
||||
def forward_cuda(self, x):
|
||||
return "cuda"
|
||||
|
||||
|
||||
class _UndeclaredBackendOp(BaseFusedOp):
|
||||
"""Overrides forward_aiter but does not declare it in ``capabilities``."""
|
||||
|
||||
op = "test.undeclared_backend"
|
||||
|
||||
def forward_native(self, x):
|
||||
return "native"
|
||||
|
||||
def forward_aiter(self, x):
|
||||
return "aiter"
|
||||
|
||||
|
||||
# --- nn.Module contract -------------------------------------------------------
|
||||
|
||||
|
||||
def test_is_standard_nn_module(monkeypatch):
|
||||
_mock_platform(monkeypatch)
|
||||
op = _NativeOnlyOp()
|
||||
assert isinstance(op, nn.Module)
|
||||
|
||||
parent = nn.Module()
|
||||
parent.act = op
|
||||
assert dict(parent.named_modules())["act"] is op
|
||||
|
||||
seen = []
|
||||
op.register_forward_hook(lambda module, args, output: seen.append(output))
|
||||
assert op(torch.zeros(1)) == "native"
|
||||
assert seen == ["native"] # __call__ goes through nn.Module, hooks fire
|
||||
|
||||
|
||||
# --- platform dispatch + native fallback ---------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"key, expect",
|
||||
[
|
||||
("cuda", "cuda"),
|
||||
("hip", "hip"),
|
||||
("npu", "npu"),
|
||||
("xpu", "xpu"),
|
||||
("musa", "musa"),
|
||||
("cpu", "cpu"),
|
||||
("", "native"),
|
||||
],
|
||||
)
|
||||
def test_platform_forward_dispatch(monkeypatch, key, expect):
|
||||
_mock_platform(monkeypatch, key=key)
|
||||
assert _AllPlatformsOp()(torch.zeros(1)) == expect
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"key, expect",
|
||||
[
|
||||
("hip", "cuda"), # HIP falls back to the CUDA path (hipified kernels)
|
||||
# MUSA has no implicit CUDA fallback: srt kernel imports are gated on
|
||||
# is_cuda(), so silently entering forward_cuda on a MUSA box can
|
||||
# NameError; ops opt in with an explicit forward_musa instead.
|
||||
("musa", "native"),
|
||||
("npu", "native"), # no NPU path -> native
|
||||
("cpu", "native"),
|
||||
("cuda", "cuda"),
|
||||
],
|
||||
)
|
||||
def test_platform_default_chains(monkeypatch, key, expect):
|
||||
_mock_platform(monkeypatch, key=key)
|
||||
assert _CudaOnlyPlatformOp()(torch.zeros(1)) == expect
|
||||
|
||||
|
||||
def test_native_fallback_without_any_override(monkeypatch):
|
||||
_mock_platform(monkeypatch, key="cuda", info=_CUDA)
|
||||
assert _NativeOnlyOp()(torch.zeros(1)) == "native"
|
||||
|
||||
|
||||
# --- optimized-backend selection ------------------------------------------------
|
||||
|
||||
|
||||
def test_declared_backend_beats_platform_forward(monkeypatch):
|
||||
_mock_platform(monkeypatch, key="cuda", info=_CUDA)
|
||||
assert _BackendAndPlatformOp()(torch.zeros(1)) == "jit"
|
||||
|
||||
|
||||
def test_capability_filters_backend_to_platform_forward(monkeypatch):
|
||||
# JIT is declared CUDA-only; on HIP the platform chain (-> forward_cuda) runs.
|
||||
_mock_platform(monkeypatch, key="hip", info=_HIP)
|
||||
assert _BackendAndPlatformOp()(torch.zeros(1)) == "cuda"
|
||||
|
||||
|
||||
def test_undeclared_backend_not_auto_selected(monkeypatch):
|
||||
_mock_platform(monkeypatch, key="", info=_CPU)
|
||||
op = _UndeclaredBackendOp()
|
||||
assert op(torch.zeros(1)) == "native"
|
||||
# ... but stays reachable by explicit request.
|
||||
assert op(torch.zeros(1), backend=KernelBackend.AITER) == "aiter"
|
||||
|
||||
|
||||
def test_priority_order_decides_between_backends(monkeypatch):
|
||||
class _TwoBackends(BaseFusedOp):
|
||||
op = "test.two_backends"
|
||||
priority = (KernelBackend.TRITON, KernelBackend.JIT, KernelBackend.TORCH)
|
||||
capabilities = {
|
||||
KernelBackend.TRITON: frozenset(),
|
||||
KernelBackend.JIT: frozenset(),
|
||||
}
|
||||
|
||||
def forward_native(self, x):
|
||||
return "native"
|
||||
|
||||
def forward_triton(self, x):
|
||||
return "triton"
|
||||
|
||||
def forward_jit(self, x):
|
||||
return "jit"
|
||||
|
||||
class _Flipped(_TwoBackends):
|
||||
priority = (KernelBackend.JIT, KernelBackend.TRITON, KernelBackend.TORCH)
|
||||
|
||||
_mock_platform(monkeypatch, key="cuda", info=_CUDA)
|
||||
assert _TwoBackends()(torch.zeros(1)) == "triton"
|
||||
assert _Flipped()(torch.zeros(1)) == "jit"
|
||||
|
||||
|
||||
def test_explicit_backend_beats_forced_global(monkeypatch):
|
||||
_mock_platform(monkeypatch, key="cuda", info=_CUDA)
|
||||
op = _BackendAndPlatformOp()
|
||||
fo.set_fused_op_backend(KernelBackend.TORCH)
|
||||
assert op(torch.zeros(1)) == "native" # forced global
|
||||
assert op(torch.zeros(1), backend=KernelBackend.JIT) == "jit" # explicit wins
|
||||
|
||||
|
||||
def test_forced_global_falls_back_when_unimplemented(monkeypatch):
|
||||
# The global debug switch must not take down ops that lack the forced
|
||||
# backend (e.g. forcing "torch" on a device-only op like the DSA indexer,
|
||||
# whose forward_native raises NotImplementedError).
|
||||
_mock_platform(monkeypatch, key="cuda", info=_CUDA)
|
||||
|
||||
class _DeviceOnly(BaseFusedOp):
|
||||
op = "test.device_only"
|
||||
|
||||
def forward_native(self, x):
|
||||
raise NotImplementedError
|
||||
|
||||
def forward_cuda(self, x):
|
||||
return "cuda"
|
||||
|
||||
op = _DeviceOnly()
|
||||
fo.set_fused_op_backend(KernelBackend.TORCH)
|
||||
assert op(torch.zeros(1)) == "cuda" # fell back to normal dispatch
|
||||
fo.set_fused_op_backend(KernelBackend.JIT)
|
||||
assert op(torch.zeros(1)) == "cuda" # no jit backend -> fall back too
|
||||
# Explicit per-call selection stays strict.
|
||||
fo.set_fused_op_backend(None)
|
||||
with pytest.raises(NotImplementedError):
|
||||
op(torch.zeros(1), backend=KernelBackend.JIT)
|
||||
|
||||
|
||||
def test_forced_global_beats_platform_and_oot(monkeypatch):
|
||||
_mock_platform(monkeypatch, key="", oot_key="myplat")
|
||||
BaseFusedOp.register_oot_forward(
|
||||
_CudaOnlyPlatformOp, lambda self, x: "oot", "myplat"
|
||||
)
|
||||
op = _CudaOnlyPlatformOp()
|
||||
fo.set_fused_op_backend(KernelBackend.TORCH)
|
||||
assert op(torch.zeros(1)) == "native"
|
||||
fo.set_fused_op_backend(None)
|
||||
assert op(torch.zeros(1)) == "oot"
|
||||
|
||||
|
||||
# --- OOT platform overrides ----------------------------------------------------
|
||||
|
||||
|
||||
def test_oot_registered_forward_wins_over_method(monkeypatch):
|
||||
class _OotOp(BaseFusedOp):
|
||||
op = "test.oot"
|
||||
|
||||
def forward_native(self, x):
|
||||
return "native"
|
||||
|
||||
def forward_myplat(self, x):
|
||||
return "method"
|
||||
|
||||
_mock_platform(monkeypatch, oot_key="myplat")
|
||||
assert _OotOp()(torch.zeros(1)) == "method" # forward_<key> lookup
|
||||
|
||||
BaseFusedOp.register_oot_forward(_OotOp, lambda self, x: "registered", "myplat")
|
||||
assert _OotOp()(torch.zeros(1)) == "registered" # registry beats method
|
||||
|
||||
|
||||
def test_oot_registration_is_exact_type(monkeypatch):
|
||||
_mock_platform(monkeypatch, oot_key="myplat")
|
||||
BaseFusedOp.register_oot_forward(
|
||||
_CudaOnlyPlatformOp, lambda self, x: "oot", "myplat"
|
||||
)
|
||||
|
||||
class _Sub(_CudaOnlyPlatformOp):
|
||||
pass
|
||||
|
||||
assert _CudaOnlyPlatformOp()(torch.zeros(1)) == "oot"
|
||||
# Subclasses do not inherit the registered forward (pre-existing
|
||||
# MultiPlatformOp semantics: lookup is by exact type).
|
||||
assert _Sub()(torch.zeros(1)) == "native"
|
||||
|
||||
|
||||
def test_oot_falls_back_to_native(monkeypatch):
|
||||
_mock_platform(monkeypatch, oot_key="myplat")
|
||||
assert _CudaOnlyPlatformOp()(torch.zeros(1)) == "native"
|
||||
|
||||
|
||||
def test_oot_registered_fn_is_bound(monkeypatch):
|
||||
_mock_platform(monkeypatch, oot_key="myplat")
|
||||
BaseFusedOp.register_oot_forward(
|
||||
_CudaOnlyPlatformOp, lambda self, x: type(self).__name__, "myplat"
|
||||
)
|
||||
assert _CudaOnlyPlatformOp()(torch.zeros(1)) == "_CudaOnlyPlatformOp"
|
||||
|
||||
|
||||
# --- dispatch caching + per-call gates ------------------------------------------
|
||||
|
||||
|
||||
def test_static_dispatch_resolved_once(monkeypatch):
|
||||
_mock_platform(monkeypatch, key="cuda", info=_CUDA)
|
||||
op = _CudaOnlyPlatformOp()
|
||||
calls = []
|
||||
original = op._resolve_forward_method
|
||||
monkeypatch.setattr(
|
||||
op,
|
||||
"_resolve_forward_method",
|
||||
lambda: calls.append(1) or original(),
|
||||
)
|
||||
op(torch.zeros(1))
|
||||
op(torch.zeros(1))
|
||||
assert len(calls) == 1 # hot path must not re-resolve per call
|
||||
|
||||
|
||||
def test_init_preseeded_forward_method_is_kept(monkeypatch):
|
||||
# srt layers pin instance paths in __init__ (e.g. env-gated aiter modes);
|
||||
# lazy resolution must not clobber that.
|
||||
_mock_platform(monkeypatch, key="cuda", info=_CUDA)
|
||||
op = _AllPlatformsOp()
|
||||
op._forward_method = op.forward_xpu
|
||||
assert op(torch.zeros(1)) == "xpu"
|
||||
|
||||
|
||||
def test_backend_eligible_override_gates_per_call(monkeypatch):
|
||||
class _Gated(BaseFusedOp):
|
||||
op = "test.gated"
|
||||
priority = (KernelBackend.JIT, KernelBackend.TORCH)
|
||||
capabilities = {KernelBackend.JIT: frozenset()}
|
||||
|
||||
def forward_native(self, x):
|
||||
return "native"
|
||||
|
||||
def forward_jit(self, x):
|
||||
return "jit"
|
||||
|
||||
def backend_eligible(self, backend, *args, **kwargs):
|
||||
if not super().backend_eligible(backend, *args, **kwargs):
|
||||
return False
|
||||
if backend is KernelBackend.JIT:
|
||||
return args[0].shape[-1] % 2 == 0
|
||||
return True
|
||||
|
||||
_mock_platform(monkeypatch, key="", info=_CPU)
|
||||
op = _Gated()
|
||||
assert op(torch.zeros(4)) == "jit"
|
||||
assert op(torch.zeros(3)) == "native" # same instance, per-call bounce
|
||||
assert op(torch.zeros(8)) == "jit"
|
||||
|
||||
|
||||
# --- torch.compile protocol -----------------------------------------------------
|
||||
|
||||
|
||||
def test_enter_leave_torch_compile_roundtrip(monkeypatch):
|
||||
_mock_platform(monkeypatch, key="cuda", info=_CUDA)
|
||||
op = _CudaOnlyPlatformOp()
|
||||
assert op(torch.zeros(1)) == "cuda"
|
||||
|
||||
op.enter_torch_compile(num_tokens=16)
|
||||
assert op.is_torch_compile
|
||||
assert op(torch.zeros(1)) == "native"
|
||||
|
||||
# Reused-module idempotency: a second enter must not overwrite the saved
|
||||
# original forward, otherwise leave() cannot restore it.
|
||||
op.enter_torch_compile(num_tokens=16)
|
||||
op.leave_torch_compile()
|
||||
assert not op.is_torch_compile
|
||||
assert op(torch.zeros(1)) == "cuda"
|
||||
op.leave_torch_compile() # double leave is a no-op
|
||||
assert op(torch.zeros(1)) == "cuda"
|
||||
|
||||
|
||||
def test_torch_compile_hook_none_keeps_dispatch(monkeypatch):
|
||||
class _KeepOptimized(_CudaOnlyPlatformOp):
|
||||
def _torch_compile_forward(self, num_tokens):
|
||||
return None if num_tokens > 1 else self.forward_native
|
||||
|
||||
_mock_platform(monkeypatch, key="cuda", info=_CUDA)
|
||||
op = _KeepOptimized()
|
||||
op.enter_torch_compile(num_tokens=8)
|
||||
assert op.is_torch_compile
|
||||
assert op(torch.zeros(1)) == "cuda" # dispatch unchanged for bs > 1
|
||||
op.leave_torch_compile()
|
||||
|
||||
op.enter_torch_compile(num_tokens=1)
|
||||
assert op(torch.zeros(1)) == "native"
|
||||
op.leave_torch_compile()
|
||||
|
||||
|
||||
def test_topk_compile_hook_is_bs1_only():
|
||||
from sglang.srt.layers.moe.topk import TopK
|
||||
|
||||
class _Probe:
|
||||
forward_native = "native-sentinel"
|
||||
|
||||
assert TopK._torch_compile_forward(_Probe(), num_tokens=1) == "native-sentinel"
|
||||
assert TopK._torch_compile_forward(_Probe(), num_tokens=2) is None
|
||||
|
||||
|
||||
def test_fused_moe_compile_hook_is_bs1_only():
|
||||
from sglang.srt.layers.moe.fused_moe_native import fused_moe_forward_native
|
||||
from sglang.srt.layers.quantization.unquant import UnquantizedFusedMoEMethod
|
||||
|
||||
probe = object.__new__(UnquantizedFusedMoEMethod)
|
||||
assert (
|
||||
UnquantizedFusedMoEMethod._torch_compile_forward(probe, num_tokens=1)
|
||||
is fused_moe_forward_native
|
||||
)
|
||||
assert UnquantizedFusedMoEMethod._torch_compile_forward(probe, num_tokens=2) is None
|
||||
|
||||
|
||||
# --- tracing --------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_trace_labels_explicit_backend(monkeypatch):
|
||||
_mock_platform(monkeypatch, key="cuda", info=_CUDA)
|
||||
op = _CudaOnlyPlatformOp()
|
||||
fo.enable_fused_op_trace()
|
||||
op(torch.zeros(2, 3))
|
||||
op(torch.zeros(2, 3), backend=KernelBackend.TORCH)
|
||||
_, explicit_rec = fo.get_fused_op_trace()
|
||||
assert explicit_rec.backend == "torch"
|
||||
|
||||
|
||||
# --- deprecated MultiPlatformOp alias --------------------------------------------
|
||||
|
||||
|
||||
def test_deprecated_alias_contract(monkeypatch):
|
||||
from sglang.srt.layers.utils import MultiPlatformOp
|
||||
|
||||
assert issubclass(MultiPlatformOp, BaseFusedOp)
|
||||
|
||||
with warnings.catch_warnings(record=True) as caught:
|
||||
warnings.simplefilter("always")
|
||||
|
||||
class _LegacyOp(MultiPlatformOp):
|
||||
# Old-style subclass: platform forwards only, no forward_native.
|
||||
def forward_cuda(self, x):
|
||||
return "cuda"
|
||||
|
||||
assert any(issubclass(w.category, DeprecationWarning) for w in caught)
|
||||
|
||||
_mock_platform(monkeypatch, key="cuda", info=_CUDA)
|
||||
op = _LegacyOp() # instantiable without forward_native (lenient alias)
|
||||
assert op(torch.zeros(1)) == "cuda"
|
||||
with pytest.raises(NotImplementedError):
|
||||
op.forward_native(torch.zeros(1))
|
||||
|
||||
# register_oot_forward via the alias lands in the shared registry.
|
||||
MultiPlatformOp.register_oot_forward(_LegacyOp, lambda self, x: "oot", "aliasplat")
|
||||
_mock_platform(monkeypatch, key="", oot_key="aliasplat")
|
||||
assert _LegacyOp()(torch.zeros(1)) == "oot"
|
||||
|
||||
|
||||
def test_deprecated_alias_keeps_legacy_platform_defaults(monkeypatch):
|
||||
"""Old MultiPlatformOp defined per-platform default methods (hip/musa ->
|
||||
cuda, npu/xpu/cpu -> native); plugin code may call them directly, and a
|
||||
subclass without forward_cuda must still raise on CUDA like before."""
|
||||
from sglang.srt.layers.utils import MultiPlatformOp
|
||||
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("ignore", DeprecationWarning)
|
||||
|
||||
class _NativeOnlyLegacy(MultiPlatformOp):
|
||||
def forward_native(self, x):
|
||||
return "native"
|
||||
|
||||
op = _NativeOnlyLegacy()
|
||||
assert op.forward_cpu(torch.zeros(1)) == "native"
|
||||
assert op.forward_npu(torch.zeros(1)) == "native"
|
||||
with pytest.raises(NotImplementedError):
|
||||
op.forward_hip(torch.zeros(1)) # chains to the raising forward_cuda
|
||||
|
||||
_mock_platform(monkeypatch, key="cuda", info=_CUDA)
|
||||
with pytest.raises(NotImplementedError):
|
||||
_NativeOnlyLegacy()(torch.zeros(1)) # old CUDA behavior preserved
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -0,0 +1,532 @@
|
||||
"""CPU-only tests for the JIT build pipeline: ninja generation and the cache.
|
||||
|
||||
Everything here runs against synthetic files, so the invariants that matter — a
|
||||
bad recorded dependency list never causes reuse, differing flags never share a
|
||||
directory, a moved clone still hits — are checked without a GPU or a compiler.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import pathlib
|
||||
import sys
|
||||
|
||||
import msgspec
|
||||
import pytest
|
||||
|
||||
from sglang.kernels.jit.utils.compile import cache, ninja
|
||||
from sglang.kernels.jit.utils.compile.paths import KERNEL_PATH
|
||||
from sglang.kernels.jit.utils.compile.spec import BuildSpec
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=8, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _fresh_digests():
|
||||
cache.clear_digest_cache()
|
||||
yield
|
||||
cache.clear_digest_cache()
|
||||
|
||||
|
||||
def _write(path: pathlib.Path, text: str) -> pathlib.Path:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(text)
|
||||
return path
|
||||
|
||||
|
||||
def _spec(**overrides) -> BuildSpec:
|
||||
base = dict(
|
||||
module_args=("activation", "bf16_t"),
|
||||
cpp_files=(),
|
||||
cuda_files=(),
|
||||
cpp_wrappers=(("run", "Kernel::run"),),
|
||||
cuda_wrappers=(),
|
||||
cflags=("-O3",),
|
||||
cuda_cflags=("-O3",),
|
||||
ldflags=(),
|
||||
include_paths=(),
|
||||
header_only=True,
|
||||
)
|
||||
base.update(overrides)
|
||||
return BuildSpec(**base)
|
||||
|
||||
|
||||
def _build_key(**overrides) -> str:
|
||||
spec = _spec(**overrides)
|
||||
return cache.compute_build_key(spec, build_file=ninja.generate(spec))
|
||||
|
||||
|
||||
def _entries(paths) -> list:
|
||||
out = []
|
||||
for path in paths:
|
||||
root, relpath = cache._normalize_path(path)
|
||||
out.append(
|
||||
cache._DepEntry(root=root, relpath=relpath, digest=cache._file_digest(path))
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def _publish_leaf(scope: pathlib.Path, paths, *, module_name="m") -> pathlib.Path:
|
||||
"""Create a leaf the way commit_build would: name derived from its own list."""
|
||||
entries = _entries(paths)
|
||||
leaf = scope / f"{cache._DEPS_KEY_PREFIX}{cache._deps_key(entries)}"
|
||||
leaf.mkdir(parents=True)
|
||||
(leaf / cache._DEPS_FILE).write_bytes(msgspec.json.encode(entries))
|
||||
(leaf / f"{module_name}.so").write_bytes(b"")
|
||||
return leaf
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Anchor roots
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_in_tree_paths_normalize_to_an_anchor():
|
||||
header = KERNEL_PATH / "include" / "sgl_kernel" / "utils.cuh"
|
||||
root, relpath = cache._normalize_path(header)
|
||||
assert root == "kernels"
|
||||
assert relpath == "include/sgl_kernel/utils.cuh"
|
||||
assert cache._resolve_path(root=root, relpath=relpath) == header
|
||||
|
||||
|
||||
def test_unknown_paths_fall_back_to_absolute(tmp_path):
|
||||
root, relpath = cache._normalize_path(tmp_path / "elsewhere.h")
|
||||
assert root == "abs"
|
||||
assert cache._resolve_path(root=root, relpath=relpath) == tmp_path / "elsewhere.h"
|
||||
|
||||
|
||||
def test_unresolvable_anchor_is_a_miss_not_a_crash():
|
||||
assert cache._resolve_path(root="pkg:does-not-exist", relpath="x.h") is None
|
||||
|
||||
|
||||
def test_anchor_roots_are_symlink_resolved(tmp_path, monkeypatch):
|
||||
"""Anchors must be symlink-resolved, since the paths matched against them are.
|
||||
|
||||
`/usr/local/cuda` is a symlink to `/usr/local/cuda-<version>`; an unresolved
|
||||
anchor makes every toolkit header miss it and fall through to `sys`, whose
|
||||
relpath then carries the CUDA version and breaks reuse across upgrades.
|
||||
|
||||
The toolkit stands in for all of them — one comprehension resolves every
|
||||
anchor — and it is faked rather than read off the machine so this still
|
||||
guards on the CPU-only runners, which have no toolkit at all.
|
||||
"""
|
||||
from sglang.kernels.jit.utils.compile import toolchain
|
||||
|
||||
versioned = tmp_path / "cuda-12.9"
|
||||
versioned.mkdir()
|
||||
(tmp_path / "cuda").symlink_to(versioned)
|
||||
monkeypatch.setattr(toolchain, "toolkit_home", lambda: tmp_path / "cuda")
|
||||
|
||||
# `__wrapped__` is the undecorated function: the anchors are memoized for
|
||||
# the process, and the real ones were already computed by an earlier test.
|
||||
roots = dict(cache._anchor_roots.__wrapped__())
|
||||
assert roots["toolkit"] == versioned
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# build_key — what must and must not change it
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"overrides",
|
||||
[
|
||||
{"module_args": ("activation", "fp16_t")},
|
||||
{"cuda_cflags": ("-O3", "--use_fast_math")},
|
||||
{"cflags": ("-O2",)},
|
||||
{"ldflags": ("-lfoo",)},
|
||||
{"cpp_wrappers": (("run", "Other::run"),)},
|
||||
{"include_paths": ("/opt/extra",)},
|
||||
{"header_only": False, "cpp_files": ("/tmp/x.cpp",), "cpp_wrappers": ()},
|
||||
],
|
||||
ids=["args", "cuda_cflags", "cflags", "ldflags", "wrappers", "includes", "mode"],
|
||||
)
|
||||
def test_build_key_separates_every_build_input(overrides):
|
||||
"""Two builds differing in any of these must never share a directory.
|
||||
|
||||
Sharing one would let a lookup select a leaf produced under different flags,
|
||||
which is the failure the whole two-key split exists to prevent.
|
||||
"""
|
||||
assert _build_key(**overrides) != _build_key()
|
||||
|
||||
|
||||
def test_build_key_covers_the_whole_ninja_file():
|
||||
"""Every flag reaching the compiler reaches the key, because the key is
|
||||
taken over the generated build file itself rather than over a hand-listed
|
||||
subset of inputs."""
|
||||
spec = _spec()
|
||||
baseline = cache.compute_build_key(spec, build_file=ninja.generate(spec))
|
||||
tampered = ninja.generate(spec).replace("-O3", "-O0")
|
||||
assert cache.compute_build_key(spec, build_file=tampered) != baseline
|
||||
|
||||
|
||||
def test_build_key_tracks_direct_source_contents(tmp_path):
|
||||
source = _write(tmp_path / "a.cu", "// v1")
|
||||
before = _build_key(cuda_files=(str(source),))
|
||||
source.write_text("// v2")
|
||||
cache.clear_digest_cache()
|
||||
assert _build_key(cuda_files=(str(source),)) != before
|
||||
|
||||
|
||||
def test_no_unordered_container_reaches_the_key(monkeypatch):
|
||||
"""Nothing hashed into a key may iterate in `PYTHONHASHSEED` order.
|
||||
|
||||
A set or dict among the hashed parts would make the same tree key
|
||||
differently in two processes: no error, no wrong result, the cache simply
|
||||
never hits again. Verified end-to-end by running the key computation under
|
||||
several hash seeds; this pins the property cheaply.
|
||||
"""
|
||||
recorded = []
|
||||
original = cache._hash_parts
|
||||
monkeypatch.setattr(
|
||||
cache,
|
||||
"_hash_parts",
|
||||
lambda parts: recorded.append(list(parts)) or original(recorded[-1]),
|
||||
)
|
||||
_build_key()
|
||||
|
||||
def walk(value, path="parts"):
|
||||
assert not isinstance(value, (set, frozenset, dict)), (
|
||||
f"unordered container at {path}: {type(value).__name__}"
|
||||
)
|
||||
if isinstance(value, (list, tuple)):
|
||||
for index, item in enumerate(value):
|
||||
walk(item, f"{path}[{index}]")
|
||||
|
||||
walk(recorded)
|
||||
|
||||
|
||||
def test_build_key_is_independent_of_install_location():
|
||||
"""Absolute paths are anchor-normalized, which is what lets a second clone
|
||||
of the same tree reuse the first clone's builds."""
|
||||
text = ninja.generate(_spec())
|
||||
assert str(KERNEL_PATH) not in cache._normalize_text(text)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# deps_key — a leaf that reproduces its own name
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_leaf_is_found_when_nothing_changed(tmp_path):
|
||||
dep = _write(tmp_path / "dep.h", "// v1")
|
||||
leaf = _publish_leaf(tmp_path, [dep])
|
||||
assert cache.find_prebuilt(scope=tmp_path, module_name="m") == leaf / "m.so"
|
||||
|
||||
|
||||
def test_leaf_is_skipped_when_a_dependency_changed(tmp_path):
|
||||
dep = _write(tmp_path / "dep.h", "// v1")
|
||||
_publish_leaf(tmp_path, [dep])
|
||||
dep.write_text("// v2")
|
||||
cache.clear_digest_cache()
|
||||
assert cache.find_prebuilt(scope=tmp_path, module_name="m") is None
|
||||
|
||||
|
||||
def test_leaf_survives_an_mtime_only_change(tmp_path):
|
||||
dep = _write(tmp_path / "dep.h", "// stable")
|
||||
leaf = _publish_leaf(tmp_path, [dep])
|
||||
os.utime(dep, (0, 0))
|
||||
cache.clear_digest_cache()
|
||||
assert cache.find_prebuilt(scope=tmp_path, module_name="m") == leaf / "m.so"
|
||||
|
||||
|
||||
def test_leaf_is_skipped_when_a_dependency_disappeared(tmp_path):
|
||||
dep = _write(tmp_path / "dep.h", "// here")
|
||||
_publish_leaf(tmp_path, [dep])
|
||||
dep.unlink()
|
||||
cache.clear_digest_cache()
|
||||
assert cache.find_prebuilt(scope=tmp_path, module_name="m") is None
|
||||
|
||||
|
||||
def test_a_leaf_that_does_not_match_its_own_name_is_skipped(tmp_path):
|
||||
"""The recorded list is verified, not trusted.
|
||||
|
||||
This is what replaces a shared manifest plus a format-version check: a
|
||||
truncated, tampered, or foreign list simply fails to reproduce the leaf's
|
||||
own name, so no schema field has to be believed.
|
||||
"""
|
||||
dep = _write(tmp_path / "dep.h", "// v1")
|
||||
leaf = _publish_leaf(tmp_path, [dep])
|
||||
other = _write(tmp_path / "other.h", "// x")
|
||||
(leaf / cache._DEPS_FILE).write_bytes(msgspec.json.encode(_entries([dep, other])))
|
||||
assert cache.find_prebuilt(scope=tmp_path, module_name="m") is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"payload",
|
||||
[b"", b"not json", b'[["kernels", "a.h"]]'],
|
||||
ids=["empty", "garbage", "wrong-shape"],
|
||||
)
|
||||
def test_an_unreadable_dependency_list_is_a_miss(tmp_path, payload):
|
||||
dep = _write(tmp_path / "dep.h", "// v1")
|
||||
leaf = _publish_leaf(tmp_path, [dep])
|
||||
(leaf / cache._DEPS_FILE).write_bytes(payload)
|
||||
assert cache.find_prebuilt(scope=tmp_path, module_name="m") is None
|
||||
|
||||
|
||||
def test_a_foreign_leaf_does_not_block_a_valid_one(tmp_path):
|
||||
"""A leaf naming a file that does not exist here is skipped, not fatal.
|
||||
|
||||
A shared, merged manifest could not do this: one unresolvable entry would
|
||||
make every lookup fail permanently on this machine.
|
||||
"""
|
||||
dep = _write(tmp_path / "dep.h", "// v1")
|
||||
good = _publish_leaf(tmp_path, [dep])
|
||||
|
||||
foreign = tmp_path / f"{cache._DEPS_KEY_PREFIX}{'0' * 16}"
|
||||
foreign.mkdir()
|
||||
(foreign / cache._DEPS_FILE).write_bytes(
|
||||
msgspec.json.encode(
|
||||
[cache._DepEntry(root="sys", relpath="include/c++/99/absent.h", digest="x")]
|
||||
)
|
||||
)
|
||||
(foreign / "m.so").write_bytes(b"")
|
||||
os.utime(foreign, None) # make the foreign leaf the newest
|
||||
|
||||
assert cache.find_prebuilt(scope=tmp_path, module_name="m") == good / "m.so"
|
||||
|
||||
|
||||
def test_a_hit_survives_an_unwritable_cache(tmp_path, monkeypatch):
|
||||
"""Touching the leaf is bookkeeping; it must never turn a hit into a crash.
|
||||
|
||||
The cache root can be a read-only mount, and a prune racing the lookup
|
||||
leaves nothing to touch -- either way `os.utime` raises, and before this it
|
||||
escaped `find_prebuilt` and took `load_jit` down on an otherwise good hit.
|
||||
"""
|
||||
header = _write(tmp_path / "a.h", "x")
|
||||
scope = tmp_path / "scope"
|
||||
_publish_leaf(scope, [header], module_name="m")
|
||||
|
||||
def deny(*args, **kwargs):
|
||||
raise PermissionError("read-only file system")
|
||||
|
||||
monkeypatch.setattr(cache.os, "utime", deny)
|
||||
assert cache.find_prebuilt(scope=scope, module_name="m") is not None
|
||||
|
||||
|
||||
def test_missing_library_is_not_a_hit(tmp_path):
|
||||
dep = _write(tmp_path / "dep.h", "// v1")
|
||||
leaf = _publish_leaf(tmp_path, [dep])
|
||||
(leaf / "m.so").unlink()
|
||||
assert cache.find_prebuilt(scope=tmp_path, module_name="m") is None
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Commit-side guards
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_empty_dependency_scan_is_rejected():
|
||||
"""An empty scan must not be recorded: it would narrow the checked set to
|
||||
nothing, the one shape of bad recorded data that could cause reuse."""
|
||||
assert not cache._covers_direct_sources(entries=[], direct_sources=["/x/a.cuh"])
|
||||
|
||||
|
||||
def test_scan_missing_a_direct_source_is_rejected(tmp_path):
|
||||
other = _write(tmp_path / "other.h", "// x")
|
||||
entries = _entries([other])
|
||||
assert not cache._covers_direct_sources(
|
||||
entries=entries, direct_sources=[str(tmp_path / "a.cuh")]
|
||||
)
|
||||
assert cache._covers_direct_sources(entries=entries, direct_sources=[str(other)])
|
||||
|
||||
|
||||
def test_build_directory_entries_are_dropped(tmp_path):
|
||||
"""The generated translation unit is not a dependency of itself.
|
||||
|
||||
Its path is unstable and its contents are already a function of inputs the
|
||||
build key covers, so recording it would defeat reuse across clones.
|
||||
"""
|
||||
build_dir = (tmp_path / "build").resolve()
|
||||
generated = _write(build_dir / "cuda.cu", "// generated")
|
||||
outside = _write(tmp_path / "real.h", "// real")
|
||||
entries = cache._to_entries(dependencies=[generated, outside], build_dir=build_dir)
|
||||
assert [entry.relpath for entry in entries] == [str(outside)]
|
||||
|
||||
|
||||
def test_publish_loses_the_race_gracefully(tmp_path):
|
||||
"""Two processes building identical content: the loser adopts the winner's leaf."""
|
||||
winner = tmp_path / "leaf"
|
||||
winner.mkdir()
|
||||
(winner / "m.so").write_bytes(b"winner")
|
||||
staging = tmp_path / "staging"
|
||||
staging.mkdir()
|
||||
(staging / "m.so").write_bytes(b"loser")
|
||||
|
||||
assert cache._publish(staging=staging, leaf=winner) == winner
|
||||
assert (winner / "m.so").read_bytes() == b"winner"
|
||||
|
||||
|
||||
def test_build_lock_excludes_a_second_holder(tmp_path):
|
||||
"""One compile per module variant per node, not one per process.
|
||||
|
||||
Every tensor-parallel rank hits the same cold cache at the same instant.
|
||||
Without exclusion each runs a full compile — measured with 8 ranks, all
|
||||
eight compiled; with it, one compiled and seven took the cache. The lock
|
||||
saves duplicated work only; the atomic rename is what makes publication
|
||||
safe, so a missing lock is slow rather than wrong.
|
||||
"""
|
||||
import threading
|
||||
|
||||
from sglang.kernels.jit.utils.compile import loader
|
||||
|
||||
held, release, contender_entered = (threading.Event() for _ in range(3))
|
||||
|
||||
def holder():
|
||||
with loader._build_lock(tmp_path):
|
||||
held.set()
|
||||
release.wait(5)
|
||||
|
||||
def contender():
|
||||
with loader._build_lock(tmp_path):
|
||||
contender_entered.set()
|
||||
|
||||
first = threading.Thread(target=holder)
|
||||
first.start()
|
||||
assert held.wait(5)
|
||||
|
||||
second = threading.Thread(target=contender)
|
||||
second.start()
|
||||
assert not contender_entered.wait(0.3), "entered while the lock was held"
|
||||
|
||||
release.set()
|
||||
first.join(5)
|
||||
second.join(5)
|
||||
assert contender_entered.is_set(), "never entered after the lock was released"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# ninja generation
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_generated_ninja_keeps_depfiles_on_disk():
|
||||
"""`deps = gcc` must not be emitted.
|
||||
|
||||
That setting folds each depfile into ninja's binary log and deletes it,
|
||||
leaving the cache with nothing to record the dependency closure from.
|
||||
"""
|
||||
text = ninja.generate(_spec(cuda_wrappers=(("run", "K::run"),)))
|
||||
assert "deps = gcc" not in text
|
||||
assert "depfile = $out.d" in text
|
||||
|
||||
|
||||
def test_generated_ninja_asks_both_compilers_for_dependencies():
|
||||
"""Both rules must write a depfile on every backend.
|
||||
|
||||
tvm-ffi's HIP branch declared `depfile =` while running a command that never
|
||||
produced one, so ROCm builds silently carried no header dependencies.
|
||||
"""
|
||||
text = ninja.generate(_spec(cuda_wrappers=(("run", "K::run"),)))
|
||||
compile_commands = [
|
||||
line
|
||||
for line in text.splitlines()
|
||||
if line.startswith(" command = ") and ' -c "$in"' in line
|
||||
]
|
||||
assert len(compile_commands) == 2
|
||||
assert all('-MD -MF "$out.d"' in line for line in compile_commands)
|
||||
|
||||
|
||||
def test_pure_cpp_module_does_not_link_the_gpu_runtime():
|
||||
"""A module with no `.cu` sources must not ask the linker for libcudart.
|
||||
|
||||
`ngram_corpus` is five .cpp files and no device code, and it is built on
|
||||
CPU-only CI runners that have no CUDA toolkit — linking it there fails with
|
||||
`cannot find -lcudart`. tvm-ffi keyed the runtime flags off the presence of
|
||||
`.cu` sources for exactly this reason.
|
||||
"""
|
||||
cpu_only = _spec(cpp_files=("/tmp/a.cpp",), cpp_wrappers=(), header_only=False)
|
||||
ldflags = next(
|
||||
line
|
||||
for line in ninja.generate(cpu_only).splitlines()
|
||||
if line.startswith("ldflags = ")
|
||||
)
|
||||
assert "cudart" not in ldflags and "amdhip" not in ldflags
|
||||
|
||||
with_device = _spec(cuda_files=("/tmp/a.cu",), cpp_wrappers=(), header_only=False)
|
||||
ldflags = next(
|
||||
line
|
||||
for line in ninja.generate(with_device).splitlines()
|
||||
if line.startswith("ldflags = ")
|
||||
)
|
||||
assert "cudart" in ldflags or "amdhip" in ldflags
|
||||
|
||||
|
||||
def test_generated_ninja_is_deterministic():
|
||||
spec = _spec(cuda_wrappers=(("run", "K::run"),))
|
||||
assert ninja.generate(spec) == ninja.generate(spec)
|
||||
|
||||
|
||||
def test_header_only_module_compiles_through_a_generated_wrapper(tmp_path):
|
||||
source = str(tmp_path / "kernel.cuh")
|
||||
units = _spec(
|
||||
cuda_files=(source,), cuda_wrappers=(("run", "K::run"),)
|
||||
).translation_units()
|
||||
assert [unit.filename for unit in units] == ["main.cpp", "cuda.cu"]
|
||||
generated = next(unit for unit in units if unit.filename == "cuda.cu").source
|
||||
assert f'#include "{source}"' in generated
|
||||
assert "TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, (K::run));" in generated
|
||||
# The wrapper must include the header defining the macro it uses. Every
|
||||
# kernel in tree happens to drag it in transitively, so dropping this would
|
||||
# only break whichever future kernel does not.
|
||||
assert "#include <tvm/ffi/function.h>" in generated
|
||||
|
||||
|
||||
def test_non_header_only_module_compiles_its_sources_in_place(tmp_path):
|
||||
source = str(tmp_path / "kernel.cu")
|
||||
units = _spec(
|
||||
cuda_files=(source,), cpp_wrappers=(), header_only=False
|
||||
).translation_units()
|
||||
assert [(unit.filename, unit.source, unit.is_cuda) for unit in units] == [
|
||||
(source, None, True)
|
||||
]
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# depfile parsing
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_depfile_parsing_handles_continuations_and_escaped_spaces():
|
||||
text = "cuda_0.o: /a/cuda.cu \\\n /a/with\\ space.cuh \\\n /b/plain.h\n"
|
||||
assert ninja._parse_depfile(text) == [
|
||||
"/a/cuda.cu",
|
||||
"/a/with space.cuh",
|
||||
"/b/plain.h",
|
||||
]
|
||||
|
||||
|
||||
def test_depfile_parsing_ignores_a_target_with_no_prerequisites():
|
||||
assert ninja._parse_depfile("a.o:\n") == []
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Layout
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_layout_is_readable_and_scoped_by_build_key():
|
||||
scope = cache.build_key_dir(
|
||||
module_name="sgl_kernel_jit_activation_bf16_t", build_key="abc123"
|
||||
)
|
||||
assert scope.name == "build-abc123"
|
||||
assert scope.parent.name == "sgl_kernel_jit_activation_bf16_t"
|
||||
assert scope.parent.parent.name == cache._target_tag()
|
||||
|
||||
|
||||
def test_module_name_is_derived_from_the_args():
|
||||
assert _spec().module_name == "sgl_kernel_jit_activation_bf16_t"
|
||||
|
||||
|
||||
def test_relative_sources_resolve_against_csrc():
|
||||
from sglang.kernels.jit.utils.compile.spec import resolve_sources
|
||||
|
||||
assert resolve_sources(["elementwise/activation.cuh"]) == (
|
||||
str(KERNEL_PATH / "csrc" / "elementwise" / "activation.cuh"),
|
||||
)
|
||||
assert resolve_sources(["/usr/include/stdio.h"]) == ("/usr/include/stdio.h",)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -1,172 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
from typing import TYPE_CHECKING, List
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.entrypoints.engine import Engine
|
||||
from sglang.srt.layers.sampler import Sampler, register_sampler_backend
|
||||
from sglang.srt.managers.scheduler import run_scheduler_process
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
from sglang.test.mock_model.utils import MOCK_MODEL_PATH
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
|
||||
from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
|
||||
|
||||
register_cuda_ci(est_time=26, stage="base-b", runner_config="1-gpu-small")
|
||||
register_amd_ci(est_time=120, stage="stage-b", runner_config="1-gpu-small-amd")
|
||||
|
||||
|
||||
CUSTOMIZED_INFO_FIELD = "sampled_token_ids_copy"
|
||||
CUSTOMIZED_INFO_SAMPLER_BACKEND = "customized_info_probe"
|
||||
_INPUT_IDS = [464, 9345, 3958, 1752, 13]
|
||||
_MAX_NEW_TOKENS = 17
|
||||
|
||||
|
||||
class CustomizedInfoSampler(Sampler):
|
||||
"""Sampler probe that mirrors every sampled token into customized_info.
|
||||
|
||||
The scheduler already appends sampled token ids to each request's output_ids.
|
||||
By copying the same values into customized_info at the sampler boundary, the
|
||||
test can assert that customized_info is sliced and accumulated exactly like
|
||||
output_ids throughout the scheduler -> tokenizer manager -> Engine path.
|
||||
"""
|
||||
|
||||
def forward(
|
||||
self,
|
||||
logits_output: LogitsProcessorOutput,
|
||||
sampling_info: SamplingBatchInfo,
|
||||
return_logprob: bool,
|
||||
top_logprobs_nums: List[int],
|
||||
token_ids_logprobs: List[List[int]],
|
||||
positions: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
batch_next_token_ids = super().forward(
|
||||
logits_output,
|
||||
sampling_info,
|
||||
return_logprob,
|
||||
top_logprobs_nums,
|
||||
token_ids_logprobs,
|
||||
positions,
|
||||
)
|
||||
|
||||
if logits_output.customized_info is None:
|
||||
logits_output.customized_info = {}
|
||||
logits_output.customized_info[CUSTOMIZED_INFO_FIELD] = (
|
||||
batch_next_token_ids.detach().cpu().tolist()
|
||||
)
|
||||
return batch_next_token_ids
|
||||
|
||||
|
||||
def install_customized_info_sampler() -> None:
|
||||
# Register before ServerArgs validation in the parent and before sampler
|
||||
# construction in the scheduler subprocess.
|
||||
register_sampler_backend(
|
||||
CUSTOMIZED_INFO_SAMPLER_BACKEND,
|
||||
CustomizedInfoSampler,
|
||||
)
|
||||
|
||||
|
||||
def run_scheduler_process_with_customized_info_sampler(*args, **kwargs):
|
||||
# Engine launches the scheduler in a subprocess. Install the sampler there
|
||||
# too so create_sampler() can resolve CUSTOMIZED_INFO_SAMPLER_BACKEND.
|
||||
install_customized_info_sampler()
|
||||
return run_scheduler_process(*args, **kwargs)
|
||||
|
||||
|
||||
class _CustomizedInfoEngine(Engine):
|
||||
run_scheduler_process_func = staticmethod(
|
||||
run_scheduler_process_with_customized_info_sampler
|
||||
)
|
||||
|
||||
|
||||
class TestCustomizedInfoStreaming(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
install_customized_info_sampler()
|
||||
cls.engine = _CustomizedInfoEngine(
|
||||
model_path=MOCK_MODEL_PATH,
|
||||
load_format="dummy",
|
||||
sampling_backend=CUSTOMIZED_INFO_SAMPLER_BACKEND,
|
||||
incremental_streaming_output=True,
|
||||
skip_tokenizer_init=True,
|
||||
disable_cuda_graph=True,
|
||||
disable_radix_cache=True,
|
||||
random_seed=0,
|
||||
log_level="error",
|
||||
mem_fraction_static=0.5,
|
||||
max_total_tokens=1024,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
cls.engine.shutdown()
|
||||
|
||||
def _sampling_params(self, *, stream_interval: int | None = None) -> dict:
|
||||
sampling_params = {
|
||||
"temperature": 0.0,
|
||||
"max_new_tokens": _MAX_NEW_TOKENS,
|
||||
"ignore_eos": True,
|
||||
}
|
||||
if stream_interval is not None:
|
||||
sampling_params["stream_interval"] = stream_interval
|
||||
return sampling_params
|
||||
|
||||
def _generate(self, *, stream: bool, stream_interval: int | None = None):
|
||||
self.engine.flush_cache()
|
||||
# skip_tokenizer_init keeps this test focused on streaming output
|
||||
# handling; input_ids bypass tokenizer setup while the real Engine,
|
||||
# scheduler, and tokenizer-manager response path still run.
|
||||
return self.engine.generate(
|
||||
input_ids=_INPUT_IDS,
|
||||
sampling_params=self._sampling_params(stream_interval=stream_interval),
|
||||
stream=stream,
|
||||
)
|
||||
|
||||
def _assert_customized_info_matches_output_ids(self, output: dict):
|
||||
# For streaming chunks this should compare per-chunk lists. For the
|
||||
# non-streaming final response it should compare fully accumulated
|
||||
# lists. Either failure means customized_info drifted from output_ids.
|
||||
self.assertIn("output_ids", output)
|
||||
self.assertIn("meta_info", output)
|
||||
self.assertIn(CUSTOMIZED_INFO_FIELD, output["meta_info"])
|
||||
self.assertEqual(
|
||||
output["meta_info"][CUSTOMIZED_INFO_FIELD], output["output_ids"]
|
||||
)
|
||||
|
||||
def test_non_streaming_returns_accumulated_customized_info(self):
|
||||
output = self._generate(stream=False)
|
||||
|
||||
self._assert_customized_info_matches_output_ids(output)
|
||||
self.assertEqual(len(output["output_ids"]), _MAX_NEW_TOKENS)
|
||||
|
||||
def test_incremental_streaming_returns_chunk_customized_info(self):
|
||||
chunks = list(self._generate(stream=True, stream_interval=1))
|
||||
|
||||
self.assertEqual(len(chunks), _MAX_NEW_TOKENS)
|
||||
output_ids = []
|
||||
for chunk in chunks:
|
||||
self._assert_customized_info_matches_output_ids(chunk)
|
||||
output_ids.extend(chunk["output_ids"])
|
||||
self.assertEqual(len(output_ids), _MAX_NEW_TOKENS)
|
||||
|
||||
def test_incremental_streaming_interval_returns_chunk_customized_info(self):
|
||||
chunks = list(self._generate(stream=True, stream_interval=4))
|
||||
|
||||
# stream_interval should coalesce multiple scheduler token events into
|
||||
# at least one multi-token Engine chunk while preserving per-chunk
|
||||
# customized_info alignment.
|
||||
self.assertGreater(len(chunks), 1)
|
||||
self.assertTrue(any(len(chunk["output_ids"]) > 1 for chunk in chunks))
|
||||
output_ids = []
|
||||
for chunk in chunks:
|
||||
self._assert_customized_info_matches_output_ids(chunk)
|
||||
output_ids.extend(chunk["output_ids"])
|
||||
self.assertEqual(len(output_ids), _MAX_NEW_TOKENS)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user