[Test] Consolidate test cleanup and CI taxonomy (net -11.4K lines) (#37436)
Co-authored-by: Mick Qian <mickqian@users.noreply.github.com>
This commit is contained in:
Xiaoyu Zhang
co-authored by
Mick Qian
Mick Qian
parent
6a1ff90f2d
commit
4d23a4fa6d
+37
-4
@@ -72,9 +72,28 @@ Parameters: `est_time` (seconds), `stage` + `runner_config` (target stage and ru
|
||||
|
||||
Keep `est_time`, `stage`, `runner_config` as **literal values** — `run_suite.py` collects them by AST parsing.
|
||||
|
||||
JIT kernel correctness tests and benchmarks live under `test/registered/jit/`, same as other registered tests (their helpers stay alongside the kernel source under `python/sglang/kernels/jit/` and are imported by absolute path):
|
||||
- Correctness tests: `test/registered/jit/test_*.py` → `base-b-kernel-unit-test-1-gpu-large`
|
||||
- Benchmarks: `test/registered/jit/benchmark/bench_*.py` → `base-b-kernel-benchmark-test-1-gpu-large`
|
||||
New and renamed tests use this layout:
|
||||
|
||||
```text
|
||||
test/registered/<kind>/<subsystem>/test_*.py
|
||||
```
|
||||
|
||||
`<kind>` is one of `unit`, `kernel`, `e2e`, `accuracy`, `perf`, or `stress`.
|
||||
Hardware is expressed by one or more `register_*_ci` calls, never by creating a
|
||||
new top-level hardware directory. The admission checker applies the layout and
|
||||
kind/suite contract incrementally while legacy paths are migrated.
|
||||
|
||||
Diffusion workflows also enter through `test/run_suite.py`; registered bridge
|
||||
files preserve their case-level pytest partitioning until the remaining
|
||||
diffusion cases are moved out of the package test-support tree.
|
||||
|
||||
New JIT kernel correctness tests and benchmarks live under
|
||||
`test/registered/kernel/jit/`; legacy `test/registered/jit/` files are migrated
|
||||
incrementally. Helpers stay alongside the kernel source under
|
||||
`python/sglang/kernels/jit/` and are imported by absolute path:
|
||||
|
||||
- Correctness tests: `test/registered/kernel/jit/test_*.py` → `base-b-kernel-unit-test-1-gpu-large`
|
||||
- Benchmarks: `test/registered/kernel/jit/benchmark/bench_*.py` → `base-b-kernel-benchmark-test-1-gpu-large`
|
||||
|
||||
## Choosing a Suite
|
||||
|
||||
@@ -94,6 +113,20 @@ Use the lightest suite that meets your test's needs. Full suite tables are in th
|
||||
|
||||
See the [write-sglang-test skill](../.claude/skills/write-sglang-test/SKILL.md) for templates, fixtures, model selection, and a complete checklist.
|
||||
|
||||
Before adding a registered test, identify the production change that would make
|
||||
it fail. Prefer extending an existing fixture/server launch over adding another
|
||||
file. The incremental admission check applies these ratchets to new or modified
|
||||
registered tests:
|
||||
|
||||
- Temporary `disabled=` registrations and unconditional skips must reference an
|
||||
issue and include `until YYYY-MM-DD`; expired entries fail lint.
|
||||
- A file registered on CUDA plus another accelerator must place a nearby
|
||||
`backend-specific:` comment above the extra registration and name the path or
|
||||
failure mode that only that backend can catch.
|
||||
- Default PR registrations are limited to 1,200 estimated weighted accelerator-seconds
|
||||
per backend (`est_time * GPU count`). Move larger matrices to extra/nightly,
|
||||
or document a nearby `ci-cost-override:` rationale.
|
||||
|
||||
## Multi-Hardware Backends
|
||||
|
||||
This README mostly describes the NVIDIA GPU CI pipeline. Other hardware backends (AMD, NPU) follow the same practices and use the multi-backend registry system. A scheduled job summarizes test coverage across all backends; [here is an example run](https://github.com/sgl-project/sglang/actions/runs/23424304300).
|
||||
@@ -111,4 +144,4 @@ This README mostly describes the NVIDIA GPU CI pipeline. Other hardware backends
|
||||
|
||||
### Adding New Models to Nightly CI
|
||||
- **Text models**: Extend the [global model list variables](https://github.com/sgl-project/sglang/blob/85c1f7937781199203b38bb46325a2840f353a04/python/sglang/test/test_utils.py#L104) in `test_utils.py`.
|
||||
- **VLMs**: Extend the `MODEL_THRESHOLDS` dictionary in `test/registered/eval/test_vlms_mmmu_eval.py`.
|
||||
- **VLMs**: Extend the `MODEL_THRESHOLDS` dictionary in `test/registered/accuracy/models/test_vlms_mmmu_eval.py`.
|
||||
|
||||
@@ -1,290 +0,0 @@
|
||||
"""AMD GROK GSM8K Completion Evaluation Test (8-GPU)
|
||||
|
||||
Tests GROK models (Grok-1 FP8, Grok-1 INT4, Grok-2) using
|
||||
few-shot completion benchmark on MI300X.
|
||||
|
||||
Registry: nightly-amd-8-gpu-grok suite
|
||||
"""
|
||||
|
||||
import ast
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
import unittest
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_amd_ci
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
is_in_ci,
|
||||
popen_launch_server,
|
||||
write_github_step_summary,
|
||||
)
|
||||
from sglang.utils import download_and_cache_file, read_jsonl
|
||||
|
||||
# DISABLED: Split into individual files for each model variant
|
||||
# See: test_grok1_fp8_eval_amd.py, test_grok1_int4_eval_amd.py, test_grok2_eval_amd.py
|
||||
register_amd_ci(
|
||||
est_time=2700,
|
||||
suite="nightly-amd-8-gpu-grok",
|
||||
nightly=True,
|
||||
disabled="Split into test_grok1_fp8_eval_amd.py, test_grok1_int4_eval_amd.py, test_grok2_eval_amd.py",
|
||||
)
|
||||
|
||||
INVALID = -9999999
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelConfig:
|
||||
"""Configuration for a model to test."""
|
||||
|
||||
model_path: str
|
||||
tp_size: int = 8
|
||||
accuracy_threshold: float = 0.50
|
||||
other_args: Optional[List[str]] = None
|
||||
env_vars: Optional[dict] = None
|
||||
tokenizer_path: Optional[str] = None
|
||||
timeout: Optional[int] = None
|
||||
|
||||
def __post_init__(self):
|
||||
if self.other_args is None:
|
||||
self.other_args = []
|
||||
if self.env_vars is None:
|
||||
self.env_vars = {}
|
||||
|
||||
|
||||
# GROK models for MI300X
|
||||
GROK_MODELS = [
|
||||
# GROK1-FP8
|
||||
ModelConfig(
|
||||
model_path="lmzheng/grok-1",
|
||||
tp_size=8,
|
||||
accuracy_threshold=0.80,
|
||||
timeout=3600,
|
||||
tokenizer_path="Xenova/grok-1-tokenizer",
|
||||
other_args=[
|
||||
"--quantization",
|
||||
"fp8",
|
||||
"--attention-backend",
|
||||
"aiter",
|
||||
"--mem-fraction-static",
|
||||
"0.85",
|
||||
"--trust-remote-code",
|
||||
],
|
||||
env_vars={
|
||||
"RCCL_MSCCL_ENABLE": "0",
|
||||
"SGLANG_USE_AITER": "1",
|
||||
"SGLANG_INT4_WEIGHT": "0",
|
||||
},
|
||||
),
|
||||
# GROK1-INT4
|
||||
ModelConfig(
|
||||
model_path="amd/grok-1-W4A8KV8",
|
||||
tp_size=8,
|
||||
accuracy_threshold=0.80,
|
||||
timeout=3600,
|
||||
tokenizer_path="Xenova/grok-1-tokenizer",
|
||||
other_args=[
|
||||
"--quantization",
|
||||
"fp8",
|
||||
"--attention-backend",
|
||||
"aiter",
|
||||
"--mem-fraction-static",
|
||||
"0.85",
|
||||
"--trust-remote-code",
|
||||
],
|
||||
env_vars={
|
||||
"RCCL_MSCCL_ENABLE": "0",
|
||||
"SGLANG_USE_AITER": "1",
|
||||
"SGLANG_INT4_WEIGHT": "1",
|
||||
},
|
||||
),
|
||||
# GROK2
|
||||
ModelConfig(
|
||||
model_path="xai-org/grok-2",
|
||||
tp_size=8,
|
||||
accuracy_threshold=0.915,
|
||||
timeout=3600,
|
||||
tokenizer_path="alvarobartt/grok-2-tokenizer",
|
||||
other_args=[
|
||||
"--quantization",
|
||||
"fp8",
|
||||
"--attention-backend",
|
||||
"aiter",
|
||||
"--mem-fraction-static",
|
||||
"0.85",
|
||||
"--trust-remote-code",
|
||||
],
|
||||
env_vars={
|
||||
"RCCL_MSCCL_ENABLE": "0",
|
||||
"SGLANG_USE_AITER": "1",
|
||||
"SGLANG_INT4_WEIGHT": "0",
|
||||
},
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def get_one_example(lines, i, include_answer):
|
||||
"""Format a single GSM8K example."""
|
||||
ret = "Question: " + lines[i]["question"] + "\nAnswer:"
|
||||
if include_answer:
|
||||
ret += " " + lines[i]["answer"]
|
||||
return ret
|
||||
|
||||
|
||||
def get_few_shot_examples(lines, k):
|
||||
"""Get k few-shot examples for prompting."""
|
||||
ret = ""
|
||||
for i in range(k):
|
||||
ret += get_one_example(lines, i, True) + "\n\n"
|
||||
return ret
|
||||
|
||||
|
||||
def get_answer_value(answer_str):
|
||||
"""Extract numerical answer from response."""
|
||||
answer_str = answer_str.replace(",", "")
|
||||
numbers = re.findall(r"\d+", answer_str)
|
||||
if len(numbers) < 1:
|
||||
return INVALID
|
||||
try:
|
||||
return ast.literal_eval(numbers[-1])
|
||||
except SyntaxError:
|
||||
return INVALID
|
||||
|
||||
|
||||
def run_gsm8k_benchmark(
|
||||
base_url: str,
|
||||
num_questions: int = 200,
|
||||
num_shots: int = 5,
|
||||
parallel: int = 64,
|
||||
) -> Tuple[float, float, float]:
|
||||
"""Run GSM8K few-shot completion benchmark."""
|
||||
import sglang as sgl
|
||||
from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint
|
||||
|
||||
url = "https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl"
|
||||
data_path = download_and_cache_file(url)
|
||||
lines = list(read_jsonl(data_path))
|
||||
|
||||
few_shot_examples = get_few_shot_examples(lines, num_shots)
|
||||
|
||||
questions = []
|
||||
labels = []
|
||||
for i in range(len(lines[:num_questions])):
|
||||
questions.append(get_one_example(lines, i, False))
|
||||
labels.append(get_answer_value(lines[i]["answer"]))
|
||||
assert all(l != INVALID for l in labels)
|
||||
arguments = [{"question": q} for q in questions]
|
||||
|
||||
@sgl.function
|
||||
def few_shot_gsm8k(s, question):
|
||||
s += few_shot_examples + question
|
||||
s += sgl.gen(
|
||||
"answer", max_tokens=512, stop=["Question", "Assistant:", "<|separator|>"]
|
||||
)
|
||||
|
||||
backend = RuntimeEndpoint(base_url)
|
||||
sgl.set_default_backend(backend)
|
||||
|
||||
tic = time.perf_counter()
|
||||
states = few_shot_gsm8k.run_batch(
|
||||
arguments, temperature=0, num_threads=parallel, progress_bar=True
|
||||
)
|
||||
latency = time.perf_counter() - tic
|
||||
|
||||
preds = [get_answer_value(states[i]["answer"]) for i in range(len(states))]
|
||||
acc = np.mean(np.array(preds) == np.array(labels))
|
||||
invalid = np.mean(np.array(preds) == INVALID)
|
||||
|
||||
return float(acc), float(invalid), float(latency)
|
||||
|
||||
|
||||
class TestGrokEvalAMD(unittest.TestCase):
|
||||
"""GROK GSM8K Completion Evaluation Test for AMD MI300X."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.models = GROK_MODELS
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.num_questions = int(os.environ.get("GSM8K_NUM_QUESTIONS", "200"))
|
||||
|
||||
def test_grok_accuracy(self):
|
||||
"""Test GROK models with GSM8K completion benchmark."""
|
||||
all_results = []
|
||||
summary = "### GROK Models (MI300X)\n\n"
|
||||
summary += "| Model | TP | Accuracy | Threshold | Status |\n"
|
||||
summary += "| ----- | -- | -------- | --------- | ------ |\n"
|
||||
|
||||
for config in self.models:
|
||||
with self.subTest(model=config.model_path):
|
||||
print(f"\n{'=' * 60}")
|
||||
print(f"Testing: {config.model_path}")
|
||||
print(f"{'=' * 60}")
|
||||
|
||||
env = os.environ.copy()
|
||||
for key, value in config.env_vars.items():
|
||||
env[key] = value
|
||||
|
||||
other_args = list(config.other_args)
|
||||
other_args.extend(["--tp", str(config.tp_size)])
|
||||
if config.tokenizer_path:
|
||||
other_args.extend(["--tokenizer-path", config.tokenizer_path])
|
||||
timeout = config.timeout or DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
|
||||
|
||||
try:
|
||||
process = popen_launch_server(
|
||||
model=config.model_path,
|
||||
base_url=self.base_url,
|
||||
timeout=timeout,
|
||||
other_args=other_args,
|
||||
env=env,
|
||||
)
|
||||
|
||||
try:
|
||||
acc, invalid, latency = run_gsm8k_benchmark(
|
||||
self.base_url, num_questions=self.num_questions
|
||||
)
|
||||
passed = acc >= config.accuracy_threshold
|
||||
status = "✅ PASS" if passed else "❌ FAIL"
|
||||
print(
|
||||
f" accuracy={acc:.3f} threshold={config.accuracy_threshold} {status}"
|
||||
)
|
||||
|
||||
all_results.append(
|
||||
{
|
||||
"model": config.model_path,
|
||||
"accuracy": acc,
|
||||
"passed": passed,
|
||||
}
|
||||
)
|
||||
summary += f"| {config.model_path} | {config.tp_size} | {acc:.3f} | {config.accuracy_threshold} | {status} |\n"
|
||||
|
||||
finally:
|
||||
kill_process_tree(process.pid)
|
||||
|
||||
except Exception as e:
|
||||
summary += f"| {config.model_path} | {config.tp_size} | N/A | {config.accuracy_threshold} | ❌ ERROR |\n"
|
||||
all_results.append(
|
||||
{
|
||||
"model": config.model_path,
|
||||
"accuracy": None,
|
||||
"passed": False,
|
||||
"error": str(e),
|
||||
}
|
||||
)
|
||||
|
||||
if is_in_ci():
|
||||
write_github_step_summary(summary)
|
||||
|
||||
failed = [r for r in all_results if not r["passed"]]
|
||||
if failed:
|
||||
raise AssertionError(f"Failed models: {[r['model'] for r in failed]}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -23,7 +23,7 @@ which is below the >=3.5.0 the aiter gluon DSA kernels need, so it logs
|
||||
what loses the accuracy. Re-add a 7.0 job once its image ships Triton >=3.5.0,
|
||||
or once the legacy DSA fallback is fixed on gfx950.
|
||||
|
||||
The eval matches the CUDA GLM-5.2-FP8 nightly (`test/registered/8-gpu-models/
|
||||
The eval matches the CUDA GLM-5.2-FP8 nightly (`test/registered/e2e/models_large/
|
||||
test_glm52_fp8.py`): same dataset and same 0.92 baseline, so a red run here
|
||||
means AMD diverged from CUDA rather than the harness diverging.
|
||||
|
||||
|
||||
@@ -14,14 +14,13 @@ import unittest
|
||||
import psutil
|
||||
|
||||
import sglang as sgl
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
CustomTestCase,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=38, stage="base-b", runner_config="1-gpu-small")
|
||||
register_amd_ci(est_time=77, suite="stage-b-test-1-gpu-small-amd")
|
||||
register_cuda_ci(est_time=77, stage="base-b", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
class TestEngineChildPids(CustomTestCase):
|
||||
|
||||
@@ -4,7 +4,7 @@ import re
|
||||
import unittest
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
@@ -17,8 +17,7 @@ from sglang.test.test_utils import (
|
||||
send_generate_requests,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=65, stage="base-b", runner_config="1-gpu-small")
|
||||
register_amd_ci(est_time=70, suite="stage-b-test-1-gpu-small-amd")
|
||||
register_cuda_ci(est_time=53, stage="base-b", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
class TestMaxQueuedRequests(CustomTestCase):
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
Balanced recipe (TP=4, DeepEP, EAGLE) plus --attn-cp-size=4 with the
|
||||
DSA prefill-CP interleave strategy. Split out of
|
||||
models_e2e/test_deepseek_v4_flash_fp4_b200.py so the `cp` group covers
|
||||
e2e/models/test_deepseek_v4_flash_fp4_b200.py so the `cp` group covers
|
||||
all context-parallel tests.
|
||||
|
||||
Registry: extra-b-test-4-gpu-b200 (label-gated extra CI, 4x B200)
|
||||
|
||||
@@ -1,71 +0,0 @@
|
||||
"""EAGLE spec-decoding core on CPU: the standard config (topk=1, page_size=1)
|
||||
on the synchronous (non-overlap) path. topk > 1 tree drafting is covered in
|
||||
test_spec_eagle_topk_cpu.py (split to stay under the per-file CI timeout).
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.kits.matched_stop_kit import MatchedStopMixin
|
||||
from sglang.test.kits.spec_server_kits import (
|
||||
SpecAccuracyKit,
|
||||
SpecCorrectnessKit,
|
||||
SpecFeatureKit,
|
||||
SpecLogprobKit,
|
||||
SpecPenaltyKit,
|
||||
)
|
||||
from sglang.test.server_fixtures.spec_eagle_fixture import EagleLlama2Base
|
||||
|
||||
# Measured 780s all-green on a 40-core GNR socket (1 launch + 18 methods).
|
||||
register_cpu_ci(
|
||||
est_time=800,
|
||||
suite="stage-a-test-cpu-intel",
|
||||
disabled="EagleLlama2Base needs gated meta-llama/Llama-2-7b-chat-hf",
|
||||
)
|
||||
|
||||
_KITS = (
|
||||
SpecCorrectnessKit,
|
||||
SpecAccuracyKit,
|
||||
SpecLogprobKit,
|
||||
SpecPenaltyKit,
|
||||
SpecFeatureKit,
|
||||
MatchedStopMixin,
|
||||
)
|
||||
|
||||
|
||||
class _Core(EagleLlama2Base):
|
||||
"""EAGLE (Llama-2) preset on CPU."""
|
||||
|
||||
attention_backend = "intel_amx"
|
||||
disable_overlap = True
|
||||
mem_fraction_static = 0.3
|
||||
# CPU decode is compute-bound; a wider batch buys nothing here.
|
||||
max_running_requests = 8
|
||||
gsm8k_num_examples = 64
|
||||
env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
|
||||
|
||||
|
||||
class TestEagleLlama2NoOverlap(_Core, *_KITS):
|
||||
"""Spec v1 (overlap scheduler off) -- the only mode reachable on CPU."""
|
||||
|
||||
# Standard chain config (topk=1, page_size=1), same shape as the CUDA core.
|
||||
spec_steps = 5
|
||||
spec_topk = 1
|
||||
spec_tokens = 6
|
||||
# EAGLE/Llama-2 topk=1 accepts modestly; tune against CI if needed.
|
||||
acc_length_thres = 1.6
|
||||
batch_accept_len_thres = 1.3
|
||||
gsm8k_accept_len_thres = 1.3
|
||||
|
||||
@unittest.skip(
|
||||
"constrained decoding on CPU needs a vocab-mask CPU branch in the "
|
||||
"xgrammar backend (upstream gap, not spec-specific); the other grammar "
|
||||
"backends lack the rollback spec verification requires"
|
||||
)
|
||||
def test_constrained_decoding(self):
|
||||
pass
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,29 +0,0 @@
|
||||
import unittest
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.kits.spec_server_kits import SpecParityKit
|
||||
from sglang.test.server_fixtures.spec_eagle_fixture import Eagle3Base
|
||||
|
||||
# Estimated: 2 sequential 8B server launches + one 4-prompt greedy method
|
||||
# (CUDA sibling: 360); tune from CI TIMINGS once it has run there.
|
||||
register_cpu_ci(
|
||||
est_time=480,
|
||||
suite="stage-a-test-cpu-intel",
|
||||
disabled="EAGLE3 numerical parity mismatches on CPU intel_amx",
|
||||
)
|
||||
|
||||
|
||||
class TestEagle3ParityCPU(SpecParityKit, Eagle3Base):
|
||||
"""EAGLE3 spec (intel_amx) greedy output == non-spec reference."""
|
||||
|
||||
attention_backend = "intel_amx"
|
||||
disable_overlap = True
|
||||
mem_fraction_static = 0.3
|
||||
# CPU decode is compute-bound; a wider batch buys nothing here.
|
||||
max_running_requests = 8
|
||||
env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,68 +0,0 @@
|
||||
"""EAGLE topk > 1 tree drafting on CPU (Llama-2 topk=4, synchronous path).
|
||||
|
||||
Split from test_spec_eagle_cpu.py, mirroring the CUDA test_spec_eagle.py /
|
||||
test_spec_eagle_topk.py layout, so each file stays under the per-file CI
|
||||
timeout.
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.kits.spec_server_kits import (
|
||||
SpecAccuracyKit,
|
||||
SpecCorrectnessKit,
|
||||
SpecFeatureKit,
|
||||
SpecLogprobKit,
|
||||
SpecPenaltyKit,
|
||||
)
|
||||
from sglang.test.server_fixtures.spec_eagle_fixture import EagleLlama2Base
|
||||
|
||||
# Measured 830s all-green on a 40-core GNR socket (1 launch + 14 methods).
|
||||
register_cpu_ci(
|
||||
est_time=850,
|
||||
suite="stage-a-test-cpu-intel",
|
||||
disabled="EagleLlama2Base needs gated meta-llama/Llama-2-7b-chat-hf",
|
||||
)
|
||||
|
||||
|
||||
class _Core(EagleLlama2Base):
|
||||
"""EAGLE (Llama-2) preset on CPU."""
|
||||
|
||||
attention_backend = "intel_amx"
|
||||
disable_overlap = True
|
||||
mem_fraction_static = 0.3
|
||||
# CPU decode is compute-bound; a wider batch buys nothing here.
|
||||
max_running_requests = 8
|
||||
gsm8k_num_examples = 64
|
||||
env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
|
||||
|
||||
|
||||
class TestEagleLlama2Topk4(
|
||||
_Core,
|
||||
SpecCorrectnessKit,
|
||||
SpecAccuracyKit,
|
||||
SpecLogprobKit,
|
||||
SpecPenaltyKit,
|
||||
SpecFeatureKit,
|
||||
):
|
||||
"""EAGLE/Llama-2 topk=4 tree coverage (kits listed in bases)."""
|
||||
|
||||
spec_steps = 3
|
||||
spec_topk = 4
|
||||
spec_tokens = 8
|
||||
acc_length_thres = 2.4
|
||||
batch_accept_len_thres = 1.6
|
||||
gsm8k_accept_len_thres = 2.0
|
||||
|
||||
@unittest.skip(
|
||||
"constrained decoding on CPU needs a vocab-mask CPU branch in the "
|
||||
"xgrammar backend (upstream gap, not spec-specific); the other grammar "
|
||||
"backends lack the rollback spec verification requires"
|
||||
)
|
||||
def test_constrained_decoding(self):
|
||||
pass
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -802,269 +802,6 @@ class TestReduceSum:
|
||||
without_dim_names(unsharder_result.tensors[0]), without_dim_names(expected)
|
||||
)
|
||||
|
||||
def test_recompute_pseudo_mismatch(self) -> None:
|
||||
"""_verify_replicated_group returns failed check for RECOMPUTE_PSEUDO axis mismatch."""
|
||||
tensor_a = torch.ones(4)
|
||||
tensor_b = torch.ones(4) + 0.1
|
||||
|
||||
checks: list[ReplicatedCheckResult] = _verify_replicated_group(
|
||||
[tensor_a, tensor_b],
|
||||
axis=ParallelAxis.RECOMPUTE_PSEUDO,
|
||||
group_index=0,
|
||||
)
|
||||
assert len(checks) == 1
|
||||
assert checks[0].axis == "recompute_pseudo"
|
||||
assert checks[0].group_index == 0
|
||||
assert checks[0].compared_index == 1
|
||||
assert checks[0].baseline_index == 0
|
||||
assert not checks[0].passed
|
||||
assert checks[0].diff.max_abs_diff == pytest.approx(0.1, abs=1e-5)
|
||||
|
||||
|
||||
class TestThdCpConcat:
|
||||
def test_single_seq(self) -> None:
|
||||
"""Single seq THD unshard: 2 ranks → per-seq concat."""
|
||||
rank0 = apply_dim_names(torch.tensor([1, 2, 3]), ["t"])
|
||||
rank1 = apply_dim_names(torch.tensor([4, 5, 6]), ["t"])
|
||||
|
||||
plan = UnsharderPlan(
|
||||
axis=ParallelAxis.CP,
|
||||
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3]),
|
||||
groups=[[0, 1]],
|
||||
)
|
||||
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1])
|
||||
|
||||
assert len(unsharder_result.tensors) == 1
|
||||
expected = torch.tensor([1, 2, 3, 4, 5, 6])
|
||||
assert torch.equal(without_dim_names(unsharder_result.tensors[0]), expected)
|
||||
|
||||
def test_multi_seq(self) -> None:
|
||||
"""Multi-seq THD unshard: 2 ranks, seq_lens=[50, 32, 46]."""
|
||||
# rank0: [seqA_r0(50) | seqB_r0(32) | pad_r0(46)]
|
||||
# rank1: [seqA_r1(50) | seqB_r1(32) | pad_r1(46)]
|
||||
seq_a_r0 = torch.arange(0, 50)
|
||||
seq_b_r0 = torch.arange(100, 132)
|
||||
pad_r0 = torch.full((46,), -1)
|
||||
rank0 = apply_dim_names(torch.cat([seq_a_r0, seq_b_r0, pad_r0]), ["t"])
|
||||
|
||||
seq_a_r1 = torch.arange(50, 100)
|
||||
seq_b_r1 = torch.arange(132, 164)
|
||||
pad_r1 = torch.full((46,), -2)
|
||||
rank1 = apply_dim_names(torch.cat([seq_a_r1, seq_b_r1, pad_r1]), ["t"])
|
||||
|
||||
plan = UnsharderPlan(
|
||||
axis=ParallelAxis.CP,
|
||||
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[50, 32, 46]),
|
||||
groups=[[0, 1]],
|
||||
)
|
||||
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1])
|
||||
|
||||
assert len(unsharder_result.tensors) == 1
|
||||
unsharded: torch.Tensor = without_dim_names(unsharder_result.tensors[0])
|
||||
|
||||
# seqA: r0(50) + r1(50) = 100 tokens, values 0..99
|
||||
assert torch.equal(unsharded[:100], torch.cat([seq_a_r0, seq_a_r1]))
|
||||
# seqB: r0(32) + r1(32) = 64 tokens
|
||||
assert torch.equal(unsharded[100:164], torch.cat([seq_b_r0, seq_b_r1]))
|
||||
# pad: r0(46) + r1(46) = 92 tokens
|
||||
assert torch.equal(unsharded[164:256], torch.cat([pad_r0, pad_r1]))
|
||||
|
||||
def test_with_hidden_dim(self) -> None:
|
||||
"""THD unshard with trailing hidden dim: shape [T, H]."""
|
||||
torch.manual_seed(42)
|
||||
hidden: int = 4
|
||||
# rank0: [seqA_r0(3, 4) | seqB_r0(2, 4)]
|
||||
# rank1: [seqA_r1(3, 4) | seqB_r1(2, 4)]
|
||||
seq_a_r0 = torch.randn(3, hidden)
|
||||
seq_b_r0 = torch.randn(2, hidden)
|
||||
rank0 = apply_dim_names(torch.cat([seq_a_r0, seq_b_r0]), ["t", "h"])
|
||||
|
||||
seq_a_r1 = torch.randn(3, hidden)
|
||||
seq_b_r1 = torch.randn(2, hidden)
|
||||
rank1 = apply_dim_names(torch.cat([seq_a_r1, seq_b_r1]), ["t", "h"])
|
||||
|
||||
plan = UnsharderPlan(
|
||||
axis=ParallelAxis.CP,
|
||||
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3, 2]),
|
||||
groups=[[0, 1]],
|
||||
)
|
||||
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1])
|
||||
|
||||
assert len(unsharder_result.tensors) == 1
|
||||
unsharded: torch.Tensor = without_dim_names(unsharder_result.tensors[0])
|
||||
|
||||
assert unsharded.shape == (10, hidden)
|
||||
assert torch.equal(unsharded[:6], torch.cat([seq_a_r0, seq_a_r1]))
|
||||
assert torch.equal(unsharded[6:10], torch.cat([seq_b_r0, seq_b_r1]))
|
||||
|
||||
def test_with_leading_batch_dim(self) -> None:
|
||||
"""THD unshard with leading batch dim: shape [B, T, H], t is dim=1."""
|
||||
torch.manual_seed(42)
|
||||
batch: int = 2
|
||||
hidden: int = 4
|
||||
# rank0: [seqA_r0(3) | seqB_r0(2)] per batch item
|
||||
# rank1: [seqA_r1(3) | seqB_r1(2)] per batch item
|
||||
seq_a_r0 = torch.randn(batch, 3, hidden)
|
||||
seq_b_r0 = torch.randn(batch, 2, hidden)
|
||||
rank0 = apply_dim_names(torch.cat([seq_a_r0, seq_b_r0], dim=1), ["b", "t", "h"])
|
||||
|
||||
seq_a_r1 = torch.randn(batch, 3, hidden)
|
||||
seq_b_r1 = torch.randn(batch, 2, hidden)
|
||||
rank1 = apply_dim_names(torch.cat([seq_a_r1, seq_b_r1], dim=1), ["b", "t", "h"])
|
||||
|
||||
plan = UnsharderPlan(
|
||||
axis=ParallelAxis.CP,
|
||||
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3, 2]),
|
||||
groups=[[0, 1]],
|
||||
)
|
||||
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1])
|
||||
|
||||
assert len(unsharder_result.tensors) == 1
|
||||
unsharded: torch.Tensor = without_dim_names(unsharder_result.tensors[0])
|
||||
|
||||
assert unsharded.shape == (batch, 10, hidden)
|
||||
# seqA: r0(3) + r1(3) = 6 tokens per batch
|
||||
assert torch.equal(unsharded[:, :6, :], torch.cat([seq_a_r0, seq_a_r1], dim=1))
|
||||
# seqB: r0(2) + r1(2) = 4 tokens per batch
|
||||
assert torch.equal(
|
||||
unsharded[:, 6:10, :], torch.cat([seq_b_r0, seq_b_r1], dim=1)
|
||||
)
|
||||
|
||||
|
||||
class TestReduceSum:
|
||||
def test_basic_tp2_reduce(self) -> None:
|
||||
"""2 partial tensors sum to full tensor."""
|
||||
torch.manual_seed(42)
|
||||
full_tensor = torch.randn(4, 8)
|
||||
part_a = full_tensor * 0.6
|
||||
part_b = full_tensor * 0.4
|
||||
|
||||
dim_specs = parse_dims("h[tp:partial] d").dims
|
||||
parallel_infos = [
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
|
||||
]
|
||||
plans = compute_unsharder_plan(dim_specs, parallel_infos)
|
||||
assert len(plans) == 1
|
||||
assert isinstance(plans[0].params, ReduceSumParams)
|
||||
|
||||
named_parts: list[torch.Tensor] = _name_tensors([part_a, part_b], dim_specs)
|
||||
unsharder_result: UnsharderResult = execute_unsharder_plan(
|
||||
plans[0], named_parts
|
||||
)
|
||||
|
||||
assert len(unsharder_result.tensors) == 1
|
||||
assert torch.allclose(
|
||||
without_dim_names(unsharder_result.tensors[0]), full_tensor
|
||||
)
|
||||
|
||||
def test_tp4_reduce(self) -> None:
|
||||
"""4 partial tensors sum to full tensor."""
|
||||
torch.manual_seed(42)
|
||||
full_tensor = torch.randn(4, 8)
|
||||
parts: list[torch.Tensor] = [full_tensor * 0.25 for _ in range(4)]
|
||||
|
||||
dim_specs = parse_dims("h[tp:partial] d").dims
|
||||
parallel_infos = [
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4)
|
||||
]
|
||||
plans = compute_unsharder_plan(dim_specs, parallel_infos)
|
||||
assert len(plans) == 1
|
||||
|
||||
named_parts: list[torch.Tensor] = _name_tensors(parts, dim_specs)
|
||||
unsharder_result: UnsharderResult = execute_unsharder_plan(
|
||||
plans[0], named_parts
|
||||
)
|
||||
|
||||
assert len(unsharder_result.tensors) == 1
|
||||
assert torch.allclose(
|
||||
without_dim_names(unsharder_result.tensors[0]), full_tensor
|
||||
)
|
||||
|
||||
def test_multi_axis_concat_then_reduce(self) -> None:
|
||||
"""CP concat + TP reduce end-to-end."""
|
||||
torch.manual_seed(42)
|
||||
full_tensor = torch.randn(4, 8, 16)
|
||||
|
||||
cp_chunks = list(full_tensor.chunk(2, dim=1))
|
||||
# Each CP chunk is held as partial sums across TP ranks
|
||||
tensors: list[torch.Tensor] = []
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
|
||||
for cp_rank in range(2):
|
||||
for tp_rank in range(2):
|
||||
tensors.append(cp_chunks[cp_rank] * 0.5)
|
||||
parallel_infos.append(
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=2),
|
||||
}
|
||||
)
|
||||
|
||||
dim_specs = parse_dims("b s[cp] h[tp:partial]").dims
|
||||
plans = compute_unsharder_plan(dim_specs, parallel_infos)
|
||||
assert len(plans) == 2
|
||||
|
||||
current: list[torch.Tensor] = _name_tensors(tensors, dim_specs)
|
||||
for plan in plans:
|
||||
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, current)
|
||||
current = unsharder_result.tensors
|
||||
|
||||
assert len(current) == 1
|
||||
assert torch.allclose(without_dim_names(current[0]), full_tensor)
|
||||
|
||||
def test_reduce_scrambled_ranks(self) -> None:
|
||||
"""Scrambled rank order — sum is commutative so result is the same."""
|
||||
torch.manual_seed(42)
|
||||
full_tensor = torch.randn(4, 8)
|
||||
parts: list[torch.Tensor] = [
|
||||
full_tensor * 0.1,
|
||||
full_tensor * 0.2,
|
||||
full_tensor * 0.3,
|
||||
full_tensor * 0.4,
|
||||
]
|
||||
|
||||
parallel_infos = [
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=2, axis_size=4)},
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=4)},
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=3, axis_size=4)},
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=4)},
|
||||
]
|
||||
dim_specs = parse_dims("h[tp:partial] d").dims
|
||||
plans = compute_unsharder_plan(dim_specs, parallel_infos)
|
||||
|
||||
named_parts: list[torch.Tensor] = _name_tensors(parts, dim_specs)
|
||||
unsharder_result: UnsharderResult = execute_unsharder_plan(
|
||||
plans[0], named_parts
|
||||
)
|
||||
|
||||
assert len(unsharder_result.tensors) == 1
|
||||
assert torch.allclose(
|
||||
without_dim_names(unsharder_result.tensors[0]), full_tensor
|
||||
)
|
||||
|
||||
def test_reduce_preserves_named_dims(self) -> None:
|
||||
"""Named tensor dimensions are preserved through reduce_sum."""
|
||||
dim_specs = parse_dims("h[tp:partial] d").dims
|
||||
part_a = apply_dim_names(torch.randn(4, 8), ["h", "d"])
|
||||
part_b = apply_dim_names(torch.randn(4, 8), ["h", "d"])
|
||||
|
||||
plan = UnsharderPlan(
|
||||
axis=ParallelAxis.TP,
|
||||
params=ReduceSumParams(),
|
||||
groups=[[0, 1]],
|
||||
)
|
||||
unsharder_result: UnsharderResult = execute_unsharder_plan(
|
||||
plan, [part_a, part_b]
|
||||
)
|
||||
|
||||
assert len(unsharder_result.tensors) == 1
|
||||
assert get_dim_names(unsharder_result.tensors[0]) == ("h", "d")
|
||||
expected = apply_dim_names(
|
||||
without_dim_names(part_a) + without_dim_names(part_b), ["h", "d"]
|
||||
)
|
||||
assert torch.allclose(
|
||||
without_dim_names(unsharder_result.tensors[0]), without_dim_names(expected)
|
||||
)
|
||||
|
||||
|
||||
class TestFusedDimExecutor:
|
||||
def test_fused_tp2_concat(self) -> None:
|
||||
|
||||
@@ -19,79 +19,6 @@ from sglang.test.ci.ci_register import register_cpu_ci
|
||||
register_cpu_ci(est_time=20, stage="weekly", runner_config="cpu")
|
||||
|
||||
|
||||
class TestComputeTensorInfo:
|
||||
def test_basic_tensor_returns_correct_shape_and_dtype(self) -> None:
|
||||
tensor = torch.randn(2, 3)
|
||||
info = compute_tensor_info(tensor)
|
||||
assert info.shape == [2, 3]
|
||||
assert info.dtype == "torch.float32"
|
||||
assert info.stats.mean == pytest.approx(tensor.float().mean().item(), abs=1e-4)
|
||||
|
||||
def test_include_sample_false_returns_none_sample(self) -> None:
|
||||
tensor = torch.randn(2, 3)
|
||||
info = compute_tensor_info(tensor, include_sample=False)
|
||||
assert info.sample is None
|
||||
|
||||
def test_include_sample_true_returns_string_sample(self) -> None:
|
||||
tensor = torch.randn(2, 3)
|
||||
info = compute_tensor_info(tensor, include_sample=True)
|
||||
assert info.sample is not None
|
||||
assert isinstance(info.sample, str)
|
||||
|
||||
def test_empty_tensor_stats_are_zero(self) -> None:
|
||||
tensor = torch.tensor([])
|
||||
info = compute_tensor_info(tensor)
|
||||
assert info.stats.mean == 0.0
|
||||
assert info.stats.std == 0.0
|
||||
assert info.shape == [0]
|
||||
|
||||
def test_integer_tensor_converted_to_float_for_stats(self) -> None:
|
||||
"""Integer tensors should be cast to float internally for stats computation."""
|
||||
tensor = torch.tensor([1, 2, 3, 4], dtype=torch.int32)
|
||||
info = compute_tensor_info(tensor)
|
||||
assert info.dtype == "torch.int32"
|
||||
assert info.stats.mean == pytest.approx(2.5, abs=1e-4)
|
||||
assert info.stats.min == pytest.approx(1.0, abs=1e-4)
|
||||
assert info.stats.max == pytest.approx(4.0, abs=1e-4)
|
||||
|
||||
def test_bfloat16_tensor_shape_and_stats(self) -> None:
|
||||
"""bfloat16 tensors produce correct shape and dtype string."""
|
||||
tensor = torch.ones(3, 4, dtype=torch.bfloat16)
|
||||
info = compute_tensor_info(tensor)
|
||||
assert info.shape == [3, 4]
|
||||
assert info.dtype == "torch.bfloat16"
|
||||
assert info.stats.mean == pytest.approx(1.0, abs=1e-2)
|
||||
|
||||
def test_multidimensional_shape(self) -> None:
|
||||
"""Shape is preserved for high-rank tensors."""
|
||||
tensor = torch.randn(2, 3, 4, 5)
|
||||
info = compute_tensor_info(tensor)
|
||||
assert info.shape == [2, 3, 4, 5]
|
||||
|
||||
def test_scalar_tensor(self) -> None:
|
||||
"""Scalar (0-dim) tensor produces empty shape list."""
|
||||
tensor = torch.tensor(3.14)
|
||||
info = compute_tensor_info(tensor)
|
||||
assert info.shape == []
|
||||
assert info.stats.mean == pytest.approx(3.14, abs=1e-4)
|
||||
assert info.stats.min == pytest.approx(3.14, abs=1e-4)
|
||||
assert info.stats.max == pytest.approx(3.14, abs=1e-4)
|
||||
|
||||
def test_include_sample_true_contains_tensor_representation(self) -> None:
|
||||
"""Sample string should contain some recognizable tensor content."""
|
||||
tensor = torch.tensor([1.0, 2.0])
|
||||
info = compute_tensor_info(tensor, include_sample=True)
|
||||
assert info.sample is not None
|
||||
assert "1." in info.sample or "2." in info.sample
|
||||
|
||||
def test_percentiles_present_for_small_tensor(self) -> None:
|
||||
"""Small tensors (< threshold) should have percentile data."""
|
||||
tensor = torch.randn(100)
|
||||
info = compute_tensor_info(tensor)
|
||||
assert len(info.stats.percentiles) > 0
|
||||
assert 50 in info.stats.percentiles
|
||||
|
||||
|
||||
class TestComputeTensorInfo:
|
||||
def test_basic_tensor_returns_correct_shape_and_dtype(self) -> None:
|
||||
tensor = torch.randn(2, 3)
|
||||
|
||||
@@ -389,47 +389,6 @@ class TestTorchSave:
|
||||
assert "skip the tensor" in captured.out
|
||||
|
||||
|
||||
class TestLog:
|
||||
def test_log_format(self):
|
||||
with _capture_stdout() as captured:
|
||||
_log("hello")
|
||||
out = captured.getvalue()
|
||||
assert "hello" in out, out
|
||||
assert "[Dumper, rank=" in out, out
|
||||
assert ", t=" in out, out
|
||||
|
||||
|
||||
class TestCompareTensorsQuick:
|
||||
def test_identical(self):
|
||||
a = torch.tensor([1.0, 2.0, 3.0])
|
||||
s = _compare_tensors_quick(a, a.clone())
|
||||
assert "rel_diff=0" in s, s
|
||||
assert "max_abs=0" in s, s
|
||||
|
||||
def test_diverged(self):
|
||||
a = torch.tensor([1.0, 2.0, 3.0])
|
||||
b = torch.tensor([1.0, 2.0, 4.0]) # last element differs by 1
|
||||
s = _compare_tensors_quick(a, b)
|
||||
assert "max_abs=1" in s, s
|
||||
assert "rel_diff=" in s, s
|
||||
|
||||
def test_shape_mismatch(self):
|
||||
s = _compare_tensors_quick(torch.zeros(3), torch.zeros(4))
|
||||
assert "shape mismatch" in s, s
|
||||
|
||||
def test_dtype_unified(self):
|
||||
s = _compare_tensors_quick(
|
||||
torch.zeros(3, dtype=torch.float32),
|
||||
torch.zeros(3, dtype=torch.float64),
|
||||
)
|
||||
assert "rel_diff=" in s, s
|
||||
assert "max_abs=" in s, s
|
||||
|
||||
def test_empty(self):
|
||||
s = _compare_tensors_quick(torch.zeros(0), torch.zeros(0))
|
||||
assert s == "empty"
|
||||
|
||||
|
||||
class TestCollectiveTimeout:
|
||||
def test_watchdog_fires_on_timeout(self):
|
||||
block_event = threading.Event()
|
||||
|
||||
@@ -1,4 +1,6 @@
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
@@ -12,20 +14,11 @@ from sglang.srt.layers.linear import LinearBase
|
||||
from sglang.srt.models.qwen2 import Qwen2MLP
|
||||
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
|
||||
from sglang.srt.utils import add_prefix, get_device
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.layer_ut_utils import init_single_process_dist
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=9,
|
||||
stage="base-b",
|
||||
runner_config="1-gpu-small",
|
||||
disabled="Test uses pytest-style function without TestCase class - see #17145",
|
||||
)
|
||||
register_amd_ci(
|
||||
est_time=15,
|
||||
suite="stage-b-test-1-gpu-small-amd",
|
||||
disabled="Test uses pytest-style function without TestCase class - see #17145",
|
||||
)
|
||||
register_cuda_ci(est_time=9, stage="base-b", runner_config="1-gpu-small")
|
||||
|
||||
TEST_HIDDEN_SIZE = 32
|
||||
|
||||
@@ -73,26 +66,29 @@ def init_weights(module):
|
||||
torch.nn.init.ones_(module.weight)
|
||||
|
||||
|
||||
def test_model_forward_dump(tmp_path):
|
||||
set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
|
||||
device = get_device()
|
||||
init_single_process_dist(backend=get_default_distributed_backend(device))
|
||||
model = MockCausalLM()
|
||||
model.apply(init_weights)
|
||||
model = model.to(device=device, dtype=torch.bfloat16)
|
||||
dumper = register_forward_hook_for_model(
|
||||
model, tmp_path / "sglang_dump", [0], 0, 0, 0
|
||||
)
|
||||
class TestTensorDumpForwardHook(CustomTestCase):
|
||||
def test_model_forward_dump(self):
|
||||
set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
|
||||
device = get_device()
|
||||
init_single_process_dist(backend=get_default_distributed_backend(device))
|
||||
model = MockCausalLM()
|
||||
model.apply(init_weights)
|
||||
model = model.to(device=device, dtype=torch.bfloat16)
|
||||
|
||||
dir_path = dumper.get_dump_dir()
|
||||
inp = torch.randn(4, TEST_HIDDEN_SIZE, dtype=torch.bfloat16) * 0.01
|
||||
result = model(inp.to(device))
|
||||
data = torch.load(f"{dir_path}/Pass00000.pt")
|
||||
assert "model.layernorm" in data
|
||||
assert "model.mlp.down_proj" in data
|
||||
assert torch.allclose(
|
||||
data["model.mlp.down_proj"], result.cpu(), rtol=1e-5, atol=1e-5
|
||||
)
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
dumper = register_forward_hook_for_model(
|
||||
model, Path(temp_dir) / "sglang_dump", [0], 0, 0, 0
|
||||
)
|
||||
dir_path = dumper.get_dump_dir()
|
||||
inp = torch.randn(4, TEST_HIDDEN_SIZE, dtype=torch.bfloat16) * 0.01
|
||||
result = model(inp.to(device))
|
||||
data = torch.load(f"{dir_path}/Pass00000.pt")
|
||||
|
||||
self.assertIn("model.layernorm", data)
|
||||
self.assertIn("model.mlp.down_proj", data)
|
||||
torch.testing.assert_close(
|
||||
data["model.mlp.down_proj"], result.cpu(), rtol=1e-5, atol=1e-5
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -0,0 +1,13 @@
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.ci.diffusion_suite_bridge import run_diffusion_suite
|
||||
|
||||
# ci-cost-override: compatibility bridge preserves the existing diffusion lane.
|
||||
register_cuda_ci(
|
||||
est_time=14400,
|
||||
stage="base-b",
|
||||
runner_config="diffusion-1-gpu-h100",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_diffusion_suite("1-gpu")
|
||||
@@ -0,0 +1,13 @@
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.ci.diffusion_suite_bridge import run_diffusion_suite
|
||||
|
||||
# ci-cost-override: compatibility bridge preserves the existing diffusion lane.
|
||||
register_cuda_ci(
|
||||
est_time=7200,
|
||||
stage="base-b",
|
||||
runner_config="diffusion-1-gpu-5090",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_diffusion_suite("1-gpu-5090")
|
||||
@@ -0,0 +1,13 @@
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.ci.diffusion_suite_bridge import run_diffusion_suite
|
||||
|
||||
# ci-cost-override: compatibility bridge preserves the existing diffusion lane.
|
||||
register_cuda_ci(
|
||||
est_time=14400,
|
||||
stage="base-b",
|
||||
runner_config="diffusion-1-gpu-b200",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_diffusion_suite("1-gpu-b200")
|
||||
@@ -0,0 +1,13 @@
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.ci.diffusion_suite_bridge import run_diffusion_suite
|
||||
|
||||
# ci-cost-override: compatibility bridge preserves the existing diffusion lane.
|
||||
register_cuda_ci(
|
||||
est_time=14400,
|
||||
stage="base-b",
|
||||
runner_config="diffusion-2-gpu-h100",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_diffusion_suite("2-gpu")
|
||||
@@ -0,0 +1,13 @@
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.ci.diffusion_suite_bridge import run_diffusion_suite
|
||||
|
||||
# ci-cost-override: compatibility bridge preserves the existing diffusion lane.
|
||||
register_cuda_ci(
|
||||
est_time=3600,
|
||||
stage="base-b",
|
||||
runner_config="diffusion-bcg-1-gpu-h100",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_diffusion_suite("bcg-diffusion")
|
||||
@@ -0,0 +1,13 @@
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.ci.diffusion_suite_bridge import run_diffusion_suite
|
||||
|
||||
# ci-cost-override: compatibility bridge preserves the existing diffusion lane.
|
||||
register_cuda_ci(
|
||||
est_time=14400,
|
||||
stage="base-b",
|
||||
runner_config="diffusion-component-2-gpu-h100",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_diffusion_suite("component-accuracy")
|
||||
@@ -0,0 +1,13 @@
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.ci.diffusion_suite_bridge import run_diffusion_suite
|
||||
|
||||
# ci-cost-override: compatibility bridge preserves the existing diffusion lane.
|
||||
register_cuda_ci(
|
||||
est_time=3600,
|
||||
stage="base-b",
|
||||
runner_config="diffusion-unit-1-gpu-h100",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_diffusion_suite("unit")
|
||||
+1
-1
@@ -80,7 +80,7 @@ class TestInklingNVFP4Nightly(unittest.TestCase):
|
||||
|
||||
class TestInklingSmallCacheConsistencyNightly(unittest.TestCase):
|
||||
"""Bitwise version of the per-commit check in
|
||||
``test/registered/models_e2e/test_inkling.py``, on the real checkpoint and
|
||||
``test/registered/e2e/models/test_inkling.py``, on the real checkpoint and
|
||||
at a batch shape the tiny checkpoint never reaches."""
|
||||
|
||||
@unittest.skipIf(not is_blackwell_system(), "NVFP4 requires Blackwell")
|
||||
@@ -1,228 +0,0 @@
|
||||
"""Guards that keep the ``diffusion`` package's import surface from eroding.
|
||||
|
||||
The reorganization only stays useful if two invariants hold:
|
||||
|
||||
1. runtime code imports from ``sglang.kernels.ops.diffusion`` and not from a
|
||||
submodule, so the internal layout can move without touching call sites;
|
||||
2. the facade's ``_EXPORTS`` table and the registry's ``_SPECS`` table both
|
||||
point at symbols that actually exist.
|
||||
|
||||
Neither is checkable by the type system, and both fail silently -- a stale
|
||||
``_EXPORTS`` entry only raises when some model happens to call that kernel, on
|
||||
a GPU, at serving time. These are pure-CPU tests: they read the tables and
|
||||
resolve them with ``importlib``/``ast`` without importing torch backends.
|
||||
"""
|
||||
|
||||
import ast
|
||||
import functools
|
||||
import importlib
|
||||
import pathlib
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
from sglang.kernels.ops.diffusion import _EXPORTS, _SPECS
|
||||
from sglang.kernels.registry import registry
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=16, suite="base-a-test-cpu")
|
||||
|
||||
PACKAGE = "sglang.kernels.ops.diffusion"
|
||||
_PACKAGE_DIR = pathlib.Path(importlib.import_module(PACKAGE).__file__ or "").parent
|
||||
_REPO_ROOT = _PACKAGE_DIR.parents[4] # <repo>/python/sglang/kernels/ops/diffusion
|
||||
|
||||
# Backend-specific test files may name a leaf module on purpose; everything
|
||||
# else -- all runtime code -- must go through the facade.
|
||||
_DEEP_IMPORT_ALLOWLIST = {
|
||||
"python/sglang/multimodal_gen/test/unit/test_latent_upsampler_group_norm_silu.py",
|
||||
"test/registered/kernels/ops/diffusion/test_model_fast_paths.py",
|
||||
"test/registered/kernels/ops/diffusion/test_sites.py",
|
||||
# This test exercises the pure-Torch fallback implementation directly.
|
||||
"test/registered/unit/utils/test_diffusion_torch_fallback.py",
|
||||
}
|
||||
|
||||
|
||||
def _module_defines(module_path: str) -> set[str]:
|
||||
"""Top-level names bound by a submodule, without importing it.
|
||||
|
||||
Importing would pull in Triton / CuTe-DSL / FlyDSL, none of which are
|
||||
installed on the CPU CI lane -- so this reads the source instead.
|
||||
"""
|
||||
if module_path.startswith("sglang."):
|
||||
spec = importlib.util.find_spec(module_path)
|
||||
assert spec is not None and spec.origin is not None, module_path
|
||||
path = pathlib.Path(spec.origin)
|
||||
else:
|
||||
path = _PACKAGE_DIR / (module_path.replace(".", "/") + ".py")
|
||||
if not path.exists():
|
||||
path = _PACKAGE_DIR / module_path.replace(".", "/") / "__init__.py"
|
||||
assert path.exists(), f"{PACKAGE}.{module_path} does not exist"
|
||||
|
||||
names: set[str] = set()
|
||||
for node in ast.parse(path.read_text(encoding="utf-8")).body:
|
||||
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
|
||||
names.add(node.name)
|
||||
elif isinstance(node, ast.Assign):
|
||||
names.update(t.id for t in node.targets if isinstance(t, ast.Name))
|
||||
elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name):
|
||||
names.add(node.target.id)
|
||||
elif isinstance(node, (ast.Import, ast.ImportFrom)):
|
||||
names.update((a.asname or a.name).split(".")[0] for a in node.names)
|
||||
elif isinstance(node, (ast.If, ast.Try)):
|
||||
# Platform-conditional rebinds (``x = select_impl(...)``) and
|
||||
# guarded defs still bind a public name.
|
||||
for inner in ast.walk(node):
|
||||
if isinstance(inner, (ast.FunctionDef, ast.ClassDef)):
|
||||
names.add(inner.name)
|
||||
elif isinstance(inner, ast.Assign):
|
||||
names.update(t.id for t in inner.targets if isinstance(t, ast.Name))
|
||||
return names
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=None)
|
||||
def _scan_root(root: str) -> tuple[frozenset[str], tuple[str, ...]]:
|
||||
unexported: set[str] = set()
|
||||
offenders: list[str] = []
|
||||
root_dir = _REPO_ROOT / root
|
||||
if not root_dir.exists():
|
||||
return frozenset(), ()
|
||||
|
||||
for path in root_dir.rglob("*.py"):
|
||||
rel = path.relative_to(_REPO_ROOT).as_posix()
|
||||
if rel.startswith(
|
||||
(
|
||||
"python/sglang/kernels/ops/diffusion/",
|
||||
"python/sglang/kernels/kda_kernels/",
|
||||
)
|
||||
):
|
||||
continue
|
||||
try:
|
||||
source = path.read_text(encoding="utf-8")
|
||||
except UnicodeDecodeError:
|
||||
continue
|
||||
if PACKAGE not in source:
|
||||
continue
|
||||
try:
|
||||
tree = ast.parse(source)
|
||||
except SyntaxError:
|
||||
continue
|
||||
allowlisted = rel in _DEEP_IMPORT_ALLOWLIST
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(node, ast.ImportFrom):
|
||||
if node.module == PACKAGE:
|
||||
unexported.update(
|
||||
a.name
|
||||
for a in node.names
|
||||
if a.name not in _EXPORTS and not a.name.startswith("_")
|
||||
)
|
||||
elif (
|
||||
not allowlisted
|
||||
and node.module
|
||||
and node.module.startswith(f"{PACKAGE}.")
|
||||
):
|
||||
offenders.append(f"{rel}:{node.lineno} imports {node.module}")
|
||||
elif isinstance(node, ast.Import) and not allowlisted:
|
||||
offenders.extend(
|
||||
f"{rel}:{node.lineno} imports {a.name}"
|
||||
for a in node.names
|
||||
if a.name.startswith(f"{PACKAGE}.")
|
||||
)
|
||||
return frozenset(unexported), tuple(offenders)
|
||||
|
||||
|
||||
def test_every_export_resolves_to_a_real_symbol():
|
||||
missing = [
|
||||
f"{symbol} -> {module}"
|
||||
for symbol, module in sorted(_EXPORTS.items())
|
||||
if symbol not in _module_defines(module)
|
||||
]
|
||||
assert not missing, f"stale _EXPORTS entries: {missing}"
|
||||
|
||||
|
||||
def test_every_symbol_imported_from_the_facade_is_exported():
|
||||
"""The reverse of the check above, and the one that actually bites.
|
||||
|
||||
A missing ``_EXPORTS`` entry raises ``ImportError`` at module import, so a
|
||||
module-level ``from ...diffusion import x`` fails loudly. A *function-local*
|
||||
one -- the pattern used for optional backends -- fails only when that test
|
||||
or code path runs, on the platform that has the backend. Enumerating the
|
||||
call sites catches it here instead.
|
||||
"""
|
||||
unexported: set[str] = set()
|
||||
for root in ("python/sglang", "test", "benchmark"):
|
||||
unexported.update(_scan_root(root)[0])
|
||||
assert not unexported, f"imported but not in _EXPORTS: {sorted(unexported)}"
|
||||
|
||||
|
||||
def test_every_registered_spec_target_resolves():
|
||||
missing = []
|
||||
for _op, _backend, target, _caps, _description in _SPECS:
|
||||
module, _, attr = target.partition(":")
|
||||
if attr not in _module_defines(module):
|
||||
missing.append(target)
|
||||
assert not missing, f"stale _SPECS targets: {missing}"
|
||||
|
||||
|
||||
def test_registry_holds_the_diffusion_ops():
|
||||
# Registration happens at package import, is metadata-only, and is what
|
||||
# ``select_kernel`` / the tracing tools read.
|
||||
registered = {op for op in registry.ops() if op.startswith("diffusion.")}
|
||||
assert {op for op, *_ in _SPECS} <= registered
|
||||
|
||||
|
||||
def test_facade_rejects_unknown_attributes():
|
||||
module = sys.modules[PACKAGE]
|
||||
with pytest.raises(AttributeError):
|
||||
module.definitely_not_a_kernel
|
||||
assert set(module.__all__) == set(_EXPORTS)
|
||||
assert set(_EXPORTS) <= set(dir(module))
|
||||
|
||||
|
||||
def test_importing_the_package_does_not_import_any_leaf_module():
|
||||
"""The reason ``__getattr__`` is lazy rather than a block of re-exports.
|
||||
|
||||
The backends have disjoint, heavy, mutually-exclusive dependencies --
|
||||
Triton (CUDA/ROCm), CUTLASS/CuTe-DSL, and FlyDSL (gfx950). If
|
||||
``_EXPORTS`` ever degrades into eager ``from .norm.x import y`` lines, all
|
||||
of them become import-time requirements on every platform, which is how a
|
||||
CPU-only or Apple install starts failing at ``import sglang``.
|
||||
|
||||
Asserted on this package's own leaf modules rather than on ``triton`` in
|
||||
``sys.modules``: sibling operator groups import Triton for their own
|
||||
reasons, so a global check would not isolate this package's behavior.
|
||||
Run in a fresh interpreter because this process has already resolved
|
||||
exports through the facade.
|
||||
"""
|
||||
code = (
|
||||
"import importlib, sys\n"
|
||||
f"importlib.import_module('{PACKAGE}')\n"
|
||||
f"prefix = '{PACKAGE}.'\n"
|
||||
"leaves = [m for m in sys.modules if m.startswith(prefix)"
|
||||
" and not m.endswith('__init__')]\n"
|
||||
"print(','.join(sorted(m for m in leaves if '.' in m[len(prefix):]"
|
||||
" or sys.modules[m].__file__ and not sys.modules[m].__file__"
|
||||
".endswith('__init__.py'))))\n"
|
||||
)
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", code], capture_output=True, text=True, timeout=600
|
||||
)
|
||||
assert result.returncode == 0, result.stderr
|
||||
leaked = [m for m in result.stdout.strip().split(",") if m]
|
||||
assert not leaked, f"importing {PACKAGE} eagerly imported: {leaked}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("root", ["python/sglang", "test", "benchmark"])
|
||||
def test_runtime_code_imports_only_through_the_facade(root):
|
||||
if not (_REPO_ROOT / root).exists(): # source checkouts only
|
||||
pytest.skip(f"{root} not present in this install")
|
||||
|
||||
offenders = _scan_root(root)[1]
|
||||
assert not offenders, (
|
||||
"import from sglang.kernels.ops.diffusion instead of a submodule:\n "
|
||||
+ "\n ".join(offenders)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -1,7 +1,6 @@
|
||||
"""GPU-free import / registry / selector tests for ``sglang.kernels`` (RFC #29630)."""
|
||||
|
||||
import importlib
|
||||
import importlib.util
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
@@ -9,127 +8,19 @@ import pytest
|
||||
|
||||
import sglang.kernels as K
|
||||
import sglang.kernels.fused_op as fo
|
||||
import sglang.kernels.ops # noqa: F401 -- populate the registry
|
||||
import sglang.kernels.selector as sel
|
||||
from sglang.kernels import DeviceType, KernelBackend, PlatformInfo
|
||||
from sglang.kernels import KernelBackend, PlatformInfo
|
||||
from sglang.kernels.spec import CapabilityRequirement as Cap
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=24, suite="base-a-test-cpu")
|
||||
|
||||
GROUPS = K.ops.__all__
|
||||
|
||||
# Representative ops checked as a subset (the registry holds many more).
|
||||
EXPECTED = {
|
||||
"activation.silu_and_mul": {"aot", "jit", "aiter", "torch", "torch_compile"},
|
||||
"activation.relu2": {"jit", "torch", "torch_compile"},
|
||||
"layernorm.rmsnorm": {"aot", "jit", "aiter", "torch_npu", "torch", "torch_compile"},
|
||||
"layernorm.gemma_rmsnorm": {"aot", "jit", "torch_npu", "torch", "torch_compile"},
|
||||
"gemm.fp8_scaled_mm": {"aot", "torch", "torch_compile"},
|
||||
"moe.moe_align_block_size": {"aot", "jit"},
|
||||
"quantization.nvfp4_gemm_swiglu_nvfp4_quant": {"cute_dsl"},
|
||||
"kvcache.reshape_and_cache_flash": {"triton"},
|
||||
"diffusion.apply_group_norm_silu": {"triton"},
|
||||
"diffusion.norm_scale_shift": {"KDA", "cute_dsl", "flydsl"},
|
||||
"diffusion.scale_residual_norm_scale_shift": {
|
||||
"KDA",
|
||||
"triton",
|
||||
"cute_dsl",
|
||||
"flydsl",
|
||||
},
|
||||
"diffusion.residual_gate_add": {"KDA"},
|
||||
"diffusion.ltx2_qknorm_split_rope": {"KDA"},
|
||||
"diffusion.causal_conv3d_cat_pad": {"KDA", "triton"},
|
||||
"diffusion.flux2_layernorm_modulate_fp8_quant": {"KDA"},
|
||||
"diffusion.flux2_qkv_epilogue": {"KDA"},
|
||||
"diffusion.flux2_token_cat_fp8": {"KDA"},
|
||||
"gemm.qwen3x_nvfp4": {"KDA"},
|
||||
"gemm.sm120_fp8_linear": {"KDA"},
|
||||
}
|
||||
|
||||
_CPU = PlatformInfo(device_type="cpu")
|
||||
_SM90 = PlatformInfo(device_type="cuda", cuda_arch_major=9, cuda_arch_minor=0)
|
||||
_SM100 = PlatformInfo(device_type="cuda", cuda_arch_major=10, cuda_arch_minor=0)
|
||||
_HIP = PlatformInfo(device_type="hip")
|
||||
|
||||
|
||||
def test_top_level_exports():
|
||||
for name in (
|
||||
"KernelSpec",
|
||||
"KernelBackend",
|
||||
"FormatSignature",
|
||||
"CapabilityRequirement",
|
||||
"PlatformInfo",
|
||||
"registry",
|
||||
"get_kernel",
|
||||
"select_kernel",
|
||||
):
|
||||
assert hasattr(K, name), name
|
||||
|
||||
|
||||
@pytest.mark.parametrize("group", GROUPS)
|
||||
def test_group_importable(group):
|
||||
assert importlib.import_module(f"sglang.kernels.ops.{group}") is not None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("op, backends", list(EXPECTED.items()))
|
||||
def test_registry_backends(op, backends):
|
||||
assert {s.backend.value for s in K.registry.get(op)} == backends
|
||||
|
||||
|
||||
def test_specs_well_formed():
|
||||
for spec in K.registry.all_specs():
|
||||
assert spec.op == f"{spec.group}.{spec.name}"
|
||||
mod, sep, attr = spec.target.partition(":")
|
||||
assert sep == ":" and mod and attr, spec.target
|
||||
|
||||
|
||||
def test_internal_registry_target_modules_exist():
|
||||
for spec in K.registry.all_specs():
|
||||
module, _, _ = spec.target.partition(":")
|
||||
if module.startswith("sglang.kernels."):
|
||||
assert importlib.util.find_spec(module) is not None, spec.target
|
||||
|
||||
|
||||
def test_sparse_linear_attention_registry_targets_forward_kernel():
|
||||
spec = K.registry.get_backend(
|
||||
"diffusion.sparse_linear_attn_fwd", KernelBackend.TRITON
|
||||
)
|
||||
assert spec.target.endswith(":_attn_fwd")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"op, target_suffix",
|
||||
(
|
||||
("diffusion.norm_scale_shift", ":kda_norm_scale_shift"),
|
||||
(
|
||||
"diffusion.scale_residual_norm_scale_shift",
|
||||
":kda_scale_residual_norm_scale_shift",
|
||||
),
|
||||
("diffusion.residual_gate_add", ":residual_gate_add"),
|
||||
(
|
||||
"diffusion.ltx2_qknorm_split_rope",
|
||||
":ltx2_qknorm_split_rope_cuda",
|
||||
),
|
||||
(
|
||||
"diffusion.causal_conv3d_cat_pad",
|
||||
":fused_causal_conv3d_cat_pad_cuda",
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_merged_diffusion_kda_provenance_backend(op, target_suffix):
|
||||
spec = K.registry.get_backend(op, KernelBackend.KDA)
|
||||
assert spec.target.endswith(target_suffix)
|
||||
|
||||
|
||||
def test_kda_backend_implementations_live_in_kda_home():
|
||||
specs = [
|
||||
spec for spec in K.registry.all_specs() if spec.backend is KernelBackend.KDA
|
||||
]
|
||||
assert specs
|
||||
assert all(spec.target.startswith("sglang.kernels.kda_kernels.") for spec in specs)
|
||||
|
||||
|
||||
def test_single_backend_resolves_without_backend():
|
||||
assert (
|
||||
K.select_kernel("kvcache.reshape_and_cache_flash").backend
|
||||
@@ -193,15 +84,6 @@ def test_layernorm_default_backend(monkeypatch, op_attr, device, expect):
|
||||
assert getattr(ln, op_attr).auto_selected_backend().value == expect
|
||||
|
||||
|
||||
def test_per_op_backend_subset():
|
||||
# silu_and_mul ships an aiter (HIP) kernel; the gelu siblings deliberately
|
||||
# do not -- ROCm coverage is a per-(op, backend) subset.
|
||||
from sglang.kernels.ops.activation import _GELU_AND_MUL, _SILU_AND_MUL
|
||||
|
||||
assert KernelBackend.AITER in _SILU_AND_MUL.available_backends()
|
||||
assert KernelBackend.AITER not in _GELU_AND_MUL.available_backends()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"req, plat, ok",
|
||||
[
|
||||
@@ -227,20 +109,6 @@ def test_capabilities_or_semantics():
|
||||
assert K.capabilities_satisfied(Cap.CUDA, _SM90) # single tolerated
|
||||
|
||||
|
||||
def test_capability_shortcuts():
|
||||
assert Cap.CUDA == Cap(device=DeviceType.CUDA)
|
||||
assert Cap.HIP == Cap(device=DeviceType.HIP)
|
||||
assert Cap.NPU == Cap(device=DeviceType.NPU)
|
||||
assert {Cap.CUDA, Cap.HIP} == {Cap.HIP, Cap.CUDA}
|
||||
assert Cap.cuda(min_sm=(10, 0)) == Cap(
|
||||
device=DeviceType.CUDA, min_cuda_arch=(10, 0)
|
||||
)
|
||||
|
||||
|
||||
def test_platform_detect_does_not_raise():
|
||||
assert PlatformInfo.detect().device_type in ("cpu", "cuda", "hip", "npu")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"relative_path",
|
||||
(
|
||||
|
||||
@@ -1,239 +0,0 @@
|
||||
"""CPU-only structural checks for the unified kernel tree."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
import importlib.util
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
import sglang.kernels as kernels
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=13, suite="base-a-test-cpu")
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
KERNELS_ROOT = REPO_ROOT / "python" / "sglang" / "kernels"
|
||||
OPS_ROOT = KERNELS_ROOT / "ops"
|
||||
JIT_CSRC_ROOT = KERNELS_ROOT / "jit" / "csrc"
|
||||
AOT_ROOT = KERNELS_ROOT / "aot"
|
||||
|
||||
|
||||
def _directory_names(root: Path) -> set[str]:
|
||||
return {
|
||||
path.name
|
||||
for path in root.iterdir()
|
||||
if path.is_dir()
|
||||
and not path.name.startswith((".", "__"))
|
||||
and any(path.rglob("*.py"))
|
||||
}
|
||||
|
||||
|
||||
def _target_names(target: ast.expr) -> set[str]:
|
||||
if isinstance(target, ast.Name):
|
||||
return {target.id}
|
||||
if isinstance(target, (ast.List, ast.Tuple)):
|
||||
return {name for element in target.elts for name in _target_names(element)}
|
||||
return set()
|
||||
|
||||
|
||||
def _bound_names(statements: list[ast.stmt]) -> set[str]:
|
||||
"""Collect names a module can bind without importing it."""
|
||||
names: set[str] = set()
|
||||
for statement in statements:
|
||||
if isinstance(statement, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
|
||||
names.add(statement.name)
|
||||
elif isinstance(statement, ast.Assign):
|
||||
for target in statement.targets:
|
||||
names.update(_target_names(target))
|
||||
elif isinstance(statement, (ast.AnnAssign, ast.AugAssign)):
|
||||
names.update(_target_names(statement.target))
|
||||
elif isinstance(statement, (ast.Import, ast.ImportFrom)):
|
||||
for alias in statement.names:
|
||||
names.add(alias.asname or alias.name.split(".", 1)[0])
|
||||
elif isinstance(statement, (ast.For, ast.AsyncFor)):
|
||||
names.update(_target_names(statement.target))
|
||||
names.update(_bound_names(statement.body))
|
||||
names.update(_bound_names(statement.orelse))
|
||||
elif isinstance(statement, ast.If):
|
||||
names.update(_bound_names(statement.body))
|
||||
names.update(_bound_names(statement.orelse))
|
||||
elif isinstance(statement, (ast.With, ast.AsyncWith)):
|
||||
names.update(_bound_names(statement.body))
|
||||
elif isinstance(statement, ast.Try):
|
||||
names.update(_bound_names(statement.body))
|
||||
names.update(_bound_names(statement.orelse))
|
||||
names.update(_bound_names(statement.finalbody))
|
||||
for handler in statement.handlers:
|
||||
names.update(_bound_names(handler.body))
|
||||
elif isinstance(statement, ast.Match):
|
||||
for case in statement.cases:
|
||||
names.update(_bound_names(case.body))
|
||||
return names
|
||||
|
||||
|
||||
def _module_string_constants(tree: ast.Module) -> dict[str, str]:
|
||||
constants: dict[str, str] = {}
|
||||
for statement in tree.body:
|
||||
if not isinstance(statement, (ast.Assign, ast.AnnAssign)):
|
||||
continue
|
||||
value = statement.value
|
||||
if not isinstance(value, ast.Constant) or not isinstance(value.value, str):
|
||||
continue
|
||||
targets = (
|
||||
statement.targets
|
||||
if isinstance(statement, ast.Assign)
|
||||
else [statement.target]
|
||||
)
|
||||
for target in targets:
|
||||
for name in _target_names(target):
|
||||
constants[name] = value.value
|
||||
return constants
|
||||
|
||||
|
||||
def _source_patterns(expression: ast.expr, constants: dict[str, str]) -> list[str]:
|
||||
if isinstance(expression, (ast.List, ast.Tuple)):
|
||||
return [
|
||||
pattern
|
||||
for element in expression.elts
|
||||
for pattern in _source_patterns(element, constants)
|
||||
]
|
||||
if isinstance(expression, ast.Constant) and isinstance(expression.value, str):
|
||||
return [expression.value]
|
||||
if isinstance(expression, ast.Name) and expression.id in constants:
|
||||
return [constants[expression.id]]
|
||||
if isinstance(expression, ast.JoinedStr):
|
||||
parts = []
|
||||
for value in expression.values:
|
||||
if isinstance(value, ast.Constant):
|
||||
parts.append(str(value.value))
|
||||
elif isinstance(value, ast.FormattedValue):
|
||||
parts.append("*")
|
||||
else:
|
||||
raise AssertionError(f"Unsupported f-string segment: {ast.dump(value)}")
|
||||
return ["".join(parts)]
|
||||
raise AssertionError(
|
||||
f"Unsupported JIT source declaration: {ast.unparse(expression)}"
|
||||
)
|
||||
|
||||
|
||||
def test_declared_operator_groups_match_packages():
|
||||
assert set(kernels.ops.__all__) == _directory_names(OPS_ROOT)
|
||||
|
||||
|
||||
def test_registered_kernel_test_groups_are_known():
|
||||
declared_groups = set(kernels.ops.__all__)
|
||||
registered_root = REPO_ROOT / "test" / "registered" / "kernels"
|
||||
for kind in ("ops", "benchmark"):
|
||||
unknown = _directory_names(registered_root / kind) - declared_groups
|
||||
assert not unknown, (
|
||||
f"Unknown {kind} kernel group directories: {sorted(unknown)}"
|
||||
)
|
||||
|
||||
|
||||
def test_internal_registry_target_attributes_are_declared():
|
||||
missing = []
|
||||
for spec in kernels.registry.all_specs():
|
||||
module_name, _, attribute_path = spec.target.partition(":")
|
||||
if not module_name.startswith("sglang.kernels."):
|
||||
continue
|
||||
module_spec = importlib.util.find_spec(module_name)
|
||||
if (
|
||||
module_spec is None
|
||||
or module_spec.origin is None
|
||||
or not module_spec.origin.endswith(".py")
|
||||
):
|
||||
continue
|
||||
tree = ast.parse(Path(module_spec.origin).read_text())
|
||||
root_attribute = attribute_path.split(".", 1)[0]
|
||||
if root_attribute not in _bound_names(tree.body):
|
||||
missing.append(spec.target)
|
||||
assert not missing, f"KernelSpec targets missing attributes: {missing}"
|
||||
|
||||
|
||||
# `load_jit` takes in-tree names and absolute paths on the same keyword, so this
|
||||
# check can only reach the declarations spelled out in the source. A module that
|
||||
# assembles its file list at runtime from a package outside `jit/csrc` has no
|
||||
# in-tree name to verify and belongs here; there is none at the moment.
|
||||
_RUNTIME_JIT_SOURCE_MODULES: set[str] = set()
|
||||
|
||||
|
||||
def test_jit_source_declarations_exist():
|
||||
missing = []
|
||||
unsupported = []
|
||||
for python_file in OPS_ROOT.rglob("*.py"):
|
||||
if python_file.relative_to(OPS_ROOT).as_posix() in _RUNTIME_JIT_SOURCE_MODULES:
|
||||
continue
|
||||
tree = ast.parse(python_file.read_text())
|
||||
constants = _module_string_constants(tree)
|
||||
for call in (node for node in ast.walk(tree) if isinstance(node, ast.Call)):
|
||||
function_name = (
|
||||
call.func.id
|
||||
if isinstance(call.func, ast.Name)
|
||||
else call.func.attr
|
||||
if isinstance(call.func, ast.Attribute)
|
||||
else None
|
||||
)
|
||||
if function_name != "load_jit":
|
||||
continue
|
||||
for keyword in call.keywords:
|
||||
if keyword.arg not in {"cpp_files", "cuda_files"}:
|
||||
continue
|
||||
try:
|
||||
patterns = _source_patterns(keyword.value, constants)
|
||||
except AssertionError as exc:
|
||||
unsupported.append(f"{python_file.relative_to(REPO_ROOT)}: {exc}")
|
||||
continue
|
||||
for pattern in patterns:
|
||||
matches = list(JIT_CSRC_ROOT.glob(pattern))
|
||||
if not matches:
|
||||
missing.append(
|
||||
f"{python_file.relative_to(REPO_ROOT)} -> {pattern}"
|
||||
)
|
||||
assert not unsupported, "Unsupported JIT source declarations:\n" + "\n".join(
|
||||
unsupported
|
||||
)
|
||||
assert not missing, "Missing JIT sources:\n" + "\n".join(missing)
|
||||
|
||||
|
||||
def test_aot_compilation_units_are_accounted_for():
|
||||
manifests = [
|
||||
AOT_ROOT / "CMakeLists.txt",
|
||||
AOT_ROOT / "setup_metal.py",
|
||||
AOT_ROOT / "setup_musa.py",
|
||||
AOT_ROOT / "setup_rocm.py",
|
||||
AOT_ROOT / "csrc" / "cpu" / "CMakeLists.txt",
|
||||
*sorted((AOT_ROOT / "cmake").rglob("*.cmake")),
|
||||
]
|
||||
manifest_text = "\n".join(path.read_text() for path in manifests)
|
||||
source_text = {
|
||||
path: path.read_text(errors="ignore")
|
||||
for path in (AOT_ROOT / "csrc").rglob("*")
|
||||
if path.is_file()
|
||||
}
|
||||
compilation_suffixes = {".cc", ".cpp", ".cu", ".hip", ".metal", ".mu"}
|
||||
missing = []
|
||||
for source in source_text:
|
||||
if source.suffix not in compilation_suffixes:
|
||||
continue
|
||||
if AOT_ROOT / "csrc" / "cpu" in source.parents:
|
||||
# The CPU build intentionally uses file(GLOB_RECURSE ... *.cpp).
|
||||
continue
|
||||
relative_path = source.relative_to(AOT_ROOT).as_posix()
|
||||
if relative_path in manifest_text:
|
||||
continue
|
||||
if any(
|
||||
source.name in text
|
||||
for other_source, text in source_text.items()
|
||||
if other_source != source
|
||||
):
|
||||
# Some CUDA translation units are included by another source.
|
||||
continue
|
||||
missing.append(relative_path)
|
||||
assert not missing, f"AOT compilation units missing from build manifests: {missing}"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -42,7 +42,7 @@ import unittest
|
||||
import requests
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci
|
||||
from sglang.test.ci.ci_register import register_mlx_ci
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
@@ -54,7 +54,6 @@ from sglang.test.test_utils import (
|
||||
# Registered on the CPU suite but skipped wherever mlx is absent; runs for real
|
||||
# only on Apple Silicon. Also registered under stage-b-e2e-mlx, which the
|
||||
# macOS CI lane (pr-test-mlx.yml) only dispatches via a gated workflow_dispatch.
|
||||
register_cpu_ci(est_time=11, suite="base-a-test-cpu")
|
||||
register_mlx_ci(est_time=1, suite="stage-b-e2e-mlx")
|
||||
|
||||
_HAS_MLX = (
|
||||
|
||||
@@ -5,7 +5,7 @@ import unittest
|
||||
import requests
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci
|
||||
from sglang.test.ci.ci_register import register_mlx_ci
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
@@ -17,7 +17,6 @@ from sglang.test.test_utils import (
|
||||
# Registered on the CPU suite but skipped wherever mlx is absent; runs for real
|
||||
# only on Apple Silicon. Also registered under stage-b-e2e-mlx, which the
|
||||
# macOS CI lane (pr-test-mlx.yml) only dispatches via a gated workflow_dispatch.
|
||||
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
|
||||
register_mlx_ci(est_time=1, suite="stage-b-e2e-mlx")
|
||||
|
||||
_HAS_MLX = importlib.util.find_spec("mlx") is not None
|
||||
|
||||
@@ -5,7 +5,7 @@ import unittest
|
||||
import requests
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci
|
||||
from sglang.test.ci.ci_register import register_mlx_ci
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
@@ -17,7 +17,6 @@ from sglang.test.test_utils import (
|
||||
# Registered on the CPU suite but skipped wherever mlx is absent; runs for real
|
||||
# only on Apple Silicon. Also registered under stage-b-e2e-mlx, which the
|
||||
# macOS CI lane (pr-test-mlx.yml) only dispatches via a gated workflow_dispatch.
|
||||
register_cpu_ci(est_time=11, suite="base-a-test-cpu")
|
||||
register_mlx_ci(est_time=1, suite="stage-b-e2e-mlx")
|
||||
|
||||
_HAS_MLX = importlib.util.find_spec("mlx") is not None
|
||||
|
||||
@@ -1,41 +0,0 @@
|
||||
import unittest
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase, is_in_ci, run_bench_one_batch
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=120,
|
||||
stage="base-b",
|
||||
runner_config="2-gpu-large",
|
||||
disabled="Temporarily disabled",
|
||||
)
|
||||
|
||||
|
||||
class TestDummyGrok1(CustomTestCase):
|
||||
def test_dummy_grok_1(self):
|
||||
_, output_throughput, _ = run_bench_one_batch(
|
||||
None,
|
||||
[
|
||||
"--model",
|
||||
"/dummy-grok",
|
||||
"--tokenizer-path",
|
||||
"Xenova/grok-1-tokenizer",
|
||||
"--batch-size",
|
||||
"2",
|
||||
"--tp",
|
||||
"2",
|
||||
"--quantization",
|
||||
"fp8",
|
||||
"--load-format",
|
||||
"dummy",
|
||||
"--json-model-override-args",
|
||||
'{"num_hidden_layers": 2}',
|
||||
],
|
||||
)
|
||||
|
||||
if is_in_ci():
|
||||
self.assertGreater(output_throughput, 0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,34 +0,0 @@
|
||||
import unittest
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
|
||||
from sglang.test.kits.mmmu_vlm_kit import MMMUMixin
|
||||
from sglang.test.server_fixtures.default_fixture import DefaultServerBase
|
||||
from sglang.test.server_fixtures.mmmu_fixture import MMMUServerBase
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=200,
|
||||
stage="base-b",
|
||||
runner_config="1-gpu-small",
|
||||
disabled="Temporarily disabled",
|
||||
)
|
||||
|
||||
MODEL = "mistralai/Ministral-3-3B-Instruct-2512"
|
||||
|
||||
|
||||
class TestMinistral3TextOnly(GSM8KMixin, DefaultServerBase):
|
||||
gsm8k_accuracy_thres = 0.6
|
||||
model = MODEL
|
||||
other_args = ["--trust-remote-code"]
|
||||
|
||||
|
||||
class TestMinistral3MMMU(MMMUMixin, MMMUServerBase):
|
||||
accuracy = 0.3
|
||||
model = MODEL
|
||||
other_args = ["--trust-remote-code"]
|
||||
mmmu_args = ["--limit=0.1"]
|
||||
"""`--limit=0.1`: 10 percent of each task - this is fine for testing since the nominal result isn't interesting - this run is just to prevent relative regressions."""
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,940 +0,0 @@
|
||||
import json
|
||||
import unittest
|
||||
|
||||
import openai
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.srt.utils.hf_transformers_utils import get_tokenizer
|
||||
from sglang.test.ascend.test_ascend_utils import LLAMA_3_2_1B_INSTRUCT_WEIGHTS_PATH
|
||||
from sglang.test.ci.ci_register import register_npu_ci
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
register_npu_ci(
|
||||
est_time=400,
|
||||
suite="full-1-npu-a3",
|
||||
nightly=True,
|
||||
)
|
||||
|
||||
|
||||
class TestOpenAIServerFunctionCalling(CustomTestCase):
|
||||
"""Testcase：Verify the correctness of full-scenario OpenAI-style function calling with llama3 parser for Llama-3.2-1B-Instruct model.
|
||||
Cover: Single/multi-turn calls, streaming/non-streaming returns, multi-parameter verification of tool_choice, and JSON parsing validity of function parameters.
|
||||
|
||||
[Test Category] Interface
|
||||
[Test Target] /v1/chat/completions
|
||||
"""
|
||||
|
||||
# NOTE: this system_message is for Llama3.2 system prompt. Without this,
|
||||
# sometimes Llama3.2 gives a different tool call format such as:
|
||||
# '<|python_tag|>{"type": "function", "function": "add", "parameters": {"a": "3", "b": "5"}}'
|
||||
SYSTEM_MESSAGE = (
|
||||
"You are a helpful assistant with tool calling capabilities. "
|
||||
"Only reply with a tool call if the function exists in the library provided by the user. "
|
||||
"If it doesn't exist, just reply directly in natural language. "
|
||||
"When you receive a tool call response, use the output to format an answer to the original user question. "
|
||||
"You have access to the following functions. "
|
||||
"To call a function, please respond with JSON for a function call. "
|
||||
'Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}. '
|
||||
"Do not use variables.\n\n"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
# Replace with the model name needed for testing
|
||||
cls.model = LLAMA_3_2_1B_INSTRUCT_WEIGHTS_PATH
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.api_key = "sk-123456"
|
||||
|
||||
# Start the local OpenAI Server. If necessary, you can add other parameters such as --enable-tools.
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
api_key=cls.api_key,
|
||||
other_args=[
|
||||
# If your server needs extra parameters to test function calling, please add them here.
|
||||
"--attention-backend",
|
||||
"ascend",
|
||||
"--disable-cuda-graph",
|
||||
"--tool-call-parser",
|
||||
"llama3",
|
||||
],
|
||||
)
|
||||
cls.base_url += "/v1"
|
||||
cls.tokenizer = get_tokenizer(cls.model)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_function_calling_format(self):
|
||||
"""
|
||||
Test: Whether the function call format returned by the AI is correct.
|
||||
When returning a tool call, message.content should be None, and tool_calls should be a list.
|
||||
"""
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "add",
|
||||
"description": "Compute the sum of two numbers",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"a": {
|
||||
"type": "integer",
|
||||
"description": "A number",
|
||||
},
|
||||
"b": {
|
||||
"type": "integer",
|
||||
"description": "A number",
|
||||
},
|
||||
},
|
||||
"required": ["a", "b"],
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": self.SYSTEM_MESSAGE},
|
||||
{"role": "user", "content": "Compute (3+5)"},
|
||||
]
|
||||
response = client.chat.completions.create(
|
||||
model=self.model,
|
||||
max_tokens=2048,
|
||||
messages=messages,
|
||||
temperature=0.8,
|
||||
top_p=0.8,
|
||||
stream=False,
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
tool_calls = response.choices[0].message.tool_calls
|
||||
|
||||
assert isinstance(tool_calls, list) and len(tool_calls) > 0, (
|
||||
"tool_calls should be a non-empty list"
|
||||
)
|
||||
|
||||
function_name = tool_calls[0].function.name
|
||||
assert function_name == "add", "Function name should be 'add'"
|
||||
|
||||
# This unit test is too difficult for default model. Mark it as optional unit tests so it won't trigger unless specified.
|
||||
def _test_function_calling_multiturn(self):
|
||||
"""
|
||||
Test: Whether the function call format returned by the AI is correct.
|
||||
When returning a tool call, message.content should be None, and tool_calls should be a list.
|
||||
"""
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "add",
|
||||
"description": "Compute the sum of two numbers",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"a": {
|
||||
"type": "integer",
|
||||
"description": "A number",
|
||||
},
|
||||
"b": {
|
||||
"type": "integer",
|
||||
"description": "A number",
|
||||
},
|
||||
},
|
||||
"required": ["a", "b"],
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
messages = [{"role": "user", "content": "Compute (3+5)"}]
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model=self.model,
|
||||
max_tokens=2048,
|
||||
messages=messages,
|
||||
temperature=0.8,
|
||||
top_p=0.8,
|
||||
stream=False,
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
tool_call = response.choices[0].message.tool_calls[0]
|
||||
function_name = tool_call.function.name
|
||||
assert function_name == "add", "Function name should be 'add'"
|
||||
function_arguments = json.loads(tool_call.function.arguments)
|
||||
assert function_arguments in [
|
||||
{"a": 3, "b": 5},
|
||||
{"a": "3", "b": "5"},
|
||||
], f"Unexpected function arguments: {function_arguments}"
|
||||
|
||||
messages.append(response.choices[0].message)
|
||||
messages.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": tool_call.id,
|
||||
"content": "8",
|
||||
"name": function_name,
|
||||
}
|
||||
)
|
||||
|
||||
final_response = client.chat.completions.create(
|
||||
model=self.model,
|
||||
max_tokens=2048,
|
||||
messages=messages,
|
||||
temperature=0.8,
|
||||
top_p=0.8,
|
||||
stream=False,
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
assert "8" in final_response.choices[0].message.content, (
|
||||
"tool_call response should have the sum 8 in the content"
|
||||
)
|
||||
|
||||
def test_function_calling_streaming_simple(self):
|
||||
"""
|
||||
Test: Whether the function name can be correctly recognized in streaming mode.
|
||||
- Expect a function call to be found, and the function name to be correct.
|
||||
- Verify that streaming mode returns at least multiple chunks.
|
||||
"""
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_current_weather",
|
||||
"description": "Get the current weather in a given location",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {
|
||||
"type": "string",
|
||||
"description": "The city to find the weather for",
|
||||
},
|
||||
"unit": {
|
||||
"type": "string",
|
||||
"description": "Weather unit (celsius or fahrenheit)",
|
||||
"enum": ["celsius", "fahrenheit"],
|
||||
},
|
||||
},
|
||||
"required": ["city", "unit"],
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": self.SYSTEM_MESSAGE},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "What is the temperature in Paris in celsius??",
|
||||
},
|
||||
]
|
||||
|
||||
response_stream = client.chat.completions.create(
|
||||
model=self.model,
|
||||
max_tokens=2048,
|
||||
messages=messages,
|
||||
temperature=0.8,
|
||||
top_p=0.8,
|
||||
stream=True,
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
chunks = list(response_stream)
|
||||
self.assertTrue(len(chunks) > 0, "Streaming should return at least one chunk")
|
||||
|
||||
found_function_name = False
|
||||
for chunk in chunks:
|
||||
choice = chunk.choices[0]
|
||||
# Check whether the current chunk contains tool_calls
|
||||
if choice.delta.tool_calls:
|
||||
tool_call = choice.delta.tool_calls[0]
|
||||
if tool_call.function.name:
|
||||
self.assertEqual(
|
||||
tool_call.function.name,
|
||||
"get_current_weather",
|
||||
"Function name should be 'get_current_weather'",
|
||||
)
|
||||
found_function_name = True
|
||||
break
|
||||
|
||||
self.assertTrue(
|
||||
found_function_name,
|
||||
"Target function name 'get_current_weather' was not found in the streaming chunks",
|
||||
)
|
||||
|
||||
finish_reason = chunks[-1].choices[0].finish_reason
|
||||
self.assertEqual(
|
||||
finish_reason,
|
||||
"tool_calls",
|
||||
"Final response of function calling should have finish_reason 'tool_calls'",
|
||||
)
|
||||
|
||||
def test_function_calling_streaming_args_parsing(self):
|
||||
"""
|
||||
Test: Whether the function call arguments returned in streaming mode can be correctly concatenated into valid JSON.
|
||||
- The user request requires multiple parameters.
|
||||
- AI may return the arguments in chunks that need to be concatenated.
|
||||
"""
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "add",
|
||||
"description": "Compute the sum of two integers",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"a": {
|
||||
"type": "integer",
|
||||
"description": "First integer",
|
||||
},
|
||||
"b": {
|
||||
"type": "integer",
|
||||
"description": "Second integer",
|
||||
},
|
||||
},
|
||||
"required": ["a", "b"],
|
||||
},
|
||||
"strict": True, # Llama-3.2-1B is flaky in tool call. It won't always respond with parameters unless we set strict.
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": self.SYSTEM_MESSAGE},
|
||||
{"role": "user", "content": "Please sum 5 and 7, just call the function."},
|
||||
]
|
||||
|
||||
response_stream = client.chat.completions.create(
|
||||
model=self.model,
|
||||
max_tokens=2048,
|
||||
messages=messages,
|
||||
temperature=0.9,
|
||||
top_p=0.9,
|
||||
stream=True,
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
argument_fragments = []
|
||||
chunks = list(response_stream)
|
||||
function_name = None
|
||||
for chunk in chunks:
|
||||
choice = chunk.choices[0]
|
||||
if choice.delta.tool_calls:
|
||||
tool_call = choice.delta.tool_calls[0]
|
||||
# Record the function name on first occurrence
|
||||
function_name = tool_call.function.name or function_name
|
||||
# In case of multiple chunks, JSON fragments may need to be concatenated
|
||||
if tool_call.function.arguments is not None:
|
||||
argument_fragments.append(tool_call.function.arguments)
|
||||
|
||||
self.assertEqual(function_name, "add", "Function name should be 'add'")
|
||||
joined_args = "".join(argument_fragments)
|
||||
self.assertTrue(
|
||||
len(joined_args) > 0,
|
||||
"No parameter fragments were returned in the function call",
|
||||
)
|
||||
|
||||
finish_reason = chunks[-1].choices[0].finish_reason
|
||||
self.assertEqual(
|
||||
finish_reason,
|
||||
"tool_calls",
|
||||
"Final response of function calling should have finish_reason 'tool_calls'",
|
||||
)
|
||||
|
||||
# Check whether the concatenated JSON is valid
|
||||
try:
|
||||
args_obj = json.loads(joined_args)
|
||||
except json.JSONDecodeError:
|
||||
self.fail(
|
||||
"The concatenated tool call arguments are not valid JSON, parsing failed"
|
||||
)
|
||||
|
||||
self.assertIn("a", args_obj, "Missing parameter 'a'")
|
||||
self.assertIn("b", args_obj, "Missing parameter 'b'")
|
||||
self.assertEqual(str(args_obj["a"]), "5", "Parameter a should be 5")
|
||||
self.assertEqual(str(args_obj["b"]), "7", "Parameter b should be 7")
|
||||
|
||||
def test_function_call_strict(self):
|
||||
"""
|
||||
Test: Whether the strict mode of function calling works as expected.
|
||||
- When strict mode is enabled, the AI should not return a function call if the function name is not recognized.
|
||||
"""
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "sub",
|
||||
"description": "Compute the difference of two integers",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"int_a": {
|
||||
"type": "integer",
|
||||
"description": "First integer",
|
||||
},
|
||||
"int_b": {
|
||||
"type": "integer",
|
||||
"description": "Second integer",
|
||||
},
|
||||
},
|
||||
"required": ["int_a", "int_b"],
|
||||
},
|
||||
"strict": True,
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
messages = [
|
||||
{"role": "user", "content": "Please compute 5 - 7, using your tool."}
|
||||
]
|
||||
response = client.chat.completions.create(
|
||||
model=self.model,
|
||||
max_tokens=2048,
|
||||
messages=messages,
|
||||
temperature=0.8,
|
||||
top_p=0.8,
|
||||
stream=False,
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
tool_calls = response.choices[0].message.tool_calls
|
||||
function_name = tool_calls[0].function.name
|
||||
arguments = tool_calls[0].function.arguments
|
||||
args_obj = json.loads(arguments)
|
||||
|
||||
self.assertEqual(function_name, "sub", "Function name should be 'sub'")
|
||||
self.assertEqual(str(args_obj["int_a"]), "5", "Parameter int_a should be 5")
|
||||
self.assertEqual(str(args_obj["int_b"]), "7", "Parameter int_b should be 7")
|
||||
|
||||
def test_function_call_required(self):
|
||||
"""
|
||||
Test: Whether tool_choice: "required" works as expected.
|
||||
- When tool_choice == "required", the model MUST return one or more tool_calls.
|
||||
- The model may choose ANY of the provided tools; we only verify that
|
||||
a tool call exists and the selected name is among the candidates.
|
||||
"""
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "sub",
|
||||
"description": "Compute the difference of two integers",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"int_a": {
|
||||
"type": "integer",
|
||||
"description": "First integer",
|
||||
},
|
||||
"int_b": {
|
||||
"type": "integer",
|
||||
"description": "Second integer",
|
||||
},
|
||||
},
|
||||
"required": ["int_a", "int_b"],
|
||||
},
|
||||
"strict": True,
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "use this to get latest weather information for a city given its name",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {
|
||||
"type": "string",
|
||||
"description": "name of the city to get weather for",
|
||||
}
|
||||
},
|
||||
"required": ["city"],
|
||||
},
|
||||
"strict": True,
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
valid_tool_names = {t["function"]["name"] for t in tools}
|
||||
|
||||
messages = [{"role": "user", "content": "Tell me about Paris"}]
|
||||
response = client.chat.completions.create(
|
||||
model=self.model,
|
||||
max_tokens=2048,
|
||||
messages=messages,
|
||||
temperature=0,
|
||||
stream=False,
|
||||
tools=tools,
|
||||
tool_choice="required",
|
||||
)
|
||||
|
||||
tool_calls = response.choices[0].message.tool_calls
|
||||
self.assertIsNotNone(
|
||||
tool_calls, "tool_choice='required' must produce tool_calls"
|
||||
)
|
||||
self.assertGreater(len(tool_calls), 0, "tool_calls list should be non-empty")
|
||||
|
||||
function_name = tool_calls[0].function.name
|
||||
self.assertIn(
|
||||
function_name,
|
||||
valid_tool_names,
|
||||
f"Function name '{function_name}' is not among the provided tools: {valid_tool_names}",
|
||||
)
|
||||
|
||||
# Verify the arguments are parseable JSON
|
||||
arguments = tool_calls[0].function.arguments
|
||||
args_obj = json.loads(arguments)
|
||||
self.assertIsInstance(
|
||||
args_obj, dict, "Function arguments should be a JSON object"
|
||||
)
|
||||
|
||||
def test_function_call_specific(self):
|
||||
"""
|
||||
Test: Whether tool_choice: ToolChoice works as expected
|
||||
- When tool_choice is a specific ToolChoice, the model should return one or more tool_calls.
|
||||
"""
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "sub",
|
||||
"description": "Compute the difference of two integers",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"int_a": {
|
||||
"type": "integer",
|
||||
"description": "First integer",
|
||||
},
|
||||
"int_b": {
|
||||
"type": "integer",
|
||||
"description": "Second integer",
|
||||
},
|
||||
},
|
||||
"required": ["int_a", "int_b"],
|
||||
},
|
||||
"strict": True,
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "use this to get latest weather information for a city given its name",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {
|
||||
"type": "string",
|
||||
"description": "name of the city to get weather for",
|
||||
}
|
||||
},
|
||||
"required": ["city"],
|
||||
},
|
||||
"strict": True,
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
messages = [{"role": "user", "content": "What is the capital of France?"}]
|
||||
response = client.chat.completions.create(
|
||||
model=self.model,
|
||||
max_tokens=2048,
|
||||
messages=messages,
|
||||
temperature=0.8,
|
||||
top_p=0.8,
|
||||
stream=False,
|
||||
tools=tools,
|
||||
tool_choice={"type": "function", "function": {"name": "get_weather"}},
|
||||
)
|
||||
|
||||
tool_calls = response.choices[0].message.tool_calls
|
||||
self.assertIsNotNone(tool_calls, "No tool_calls in the response")
|
||||
function_name = tool_calls[0].function.name
|
||||
arguments = tool_calls[0].function.arguments
|
||||
args_obj = json.loads(arguments)
|
||||
|
||||
self.assertEqual(
|
||||
function_name, "get_weather", "Function name should be 'get_weather'"
|
||||
)
|
||||
self.assertIn("city", args_obj, "Function arguments should have 'city'")
|
||||
|
||||
def test_streaming_multiple_choices_finish_reason(self):
|
||||
"""
|
||||
Test: Verify that each choice gets its own finish_reason chunk in streaming mode with n > 1.
|
||||
This tests the fix for the bug where only the last index got a finish_reason chunk.
|
||||
"""
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_current_weather",
|
||||
"description": "Get the current weather in a given location",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": {
|
||||
"type": "string",
|
||||
"description": "The city and state, e.g. San Francisco, CA",
|
||||
},
|
||||
"unit": {
|
||||
"type": "string",
|
||||
"enum": ["celsius", "fahrenheit"],
|
||||
},
|
||||
},
|
||||
"required": ["location"],
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
messages = [
|
||||
{"role": "user", "content": "What is the weather like in Los Angeles?"}
|
||||
]
|
||||
|
||||
# Request with n=2 to get multiple choices
|
||||
response_stream = client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=messages,
|
||||
max_tokens=2048,
|
||||
temperature=0.8,
|
||||
stream=True,
|
||||
tools=tools,
|
||||
tool_choice="required", # Force tool calls
|
||||
n=2, # Multiple choices
|
||||
)
|
||||
|
||||
chunks = list(response_stream)
|
||||
|
||||
# Track finish_reason chunks for each index
|
||||
finish_reason_chunks = {}
|
||||
for chunk in chunks:
|
||||
if chunk.choices:
|
||||
for choice in chunk.choices:
|
||||
if choice.finish_reason is not None:
|
||||
index = choice.index
|
||||
if index not in finish_reason_chunks:
|
||||
finish_reason_chunks[index] = []
|
||||
finish_reason_chunks[index].append(choice.finish_reason)
|
||||
|
||||
# Verify we got finish_reason chunks for both indices
|
||||
self.assertEqual(
|
||||
len(finish_reason_chunks),
|
||||
2,
|
||||
f"Expected finish_reason chunks for 2 indices, got {len(finish_reason_chunks)}",
|
||||
)
|
||||
|
||||
# Verify both index 0 and 1 have finish_reason
|
||||
self.assertIn(
|
||||
0, finish_reason_chunks, "Missing finish_reason chunk for index 0"
|
||||
)
|
||||
self.assertIn(
|
||||
1, finish_reason_chunks, "Missing finish_reason chunk for index 1"
|
||||
)
|
||||
|
||||
# Verify the finish_reason is "tool_calls" since we forced tool calls
|
||||
for index, reasons in finish_reason_chunks.items():
|
||||
self.assertEqual(
|
||||
reasons[-1], # Last finish_reason for this index
|
||||
"tool_calls",
|
||||
f"Expected finish_reason 'tool_calls' for index {index}, got {reasons[-1]}",
|
||||
)
|
||||
|
||||
def test_function_calling_streaming_no_tool_call(self):
|
||||
"""
|
||||
Test: Whether the finish_reason is stop in streaming mode when no tool call is given.
|
||||
- Expect no function call to be found.
|
||||
- Verify that finish_reason is stop
|
||||
"""
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_current_weather",
|
||||
"description": "Get the current weather in a given location",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {
|
||||
"type": "string",
|
||||
"description": "The city to find the weather for",
|
||||
},
|
||||
"unit": {
|
||||
"type": "string",
|
||||
"description": "Weather unit (celsius or fahrenheit)",
|
||||
"enum": ["celsius", "fahrenheit"],
|
||||
},
|
||||
},
|
||||
"required": ["city", "unit"],
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
messages = [{"role": "user", "content": "Who are you?"}]
|
||||
|
||||
response_stream = client.chat.completions.create(
|
||||
model=self.model,
|
||||
max_tokens=2048,
|
||||
messages=messages,
|
||||
temperature=0.8,
|
||||
top_p=0.8,
|
||||
stream=True,
|
||||
tools=tools,
|
||||
tool_choice="none",
|
||||
)
|
||||
|
||||
chunks = list(response_stream)
|
||||
self.assertTrue(len(chunks) > 0, "Streaming should return at least one chunk")
|
||||
|
||||
found_tool_call = False
|
||||
for chunk in chunks:
|
||||
choice = chunk.choices[0]
|
||||
# Check whether the current chunk contains tool_calls
|
||||
found_tool_call = choice.delta.tool_calls is not None
|
||||
|
||||
self.assertFalse(
|
||||
found_tool_call,
|
||||
"Shouldn't have any tool_call in the streaming chunks",
|
||||
)
|
||||
|
||||
finish_reason = chunks[-1].choices[0].finish_reason
|
||||
self.assertEqual(
|
||||
finish_reason,
|
||||
"stop",
|
||||
"Final response of no function calling should have finish_reason 'stop'",
|
||||
)
|
||||
|
||||
def test_streaming_multiple_choices_without_tools(self):
|
||||
"""
|
||||
Test: Verify that each choice gets its own finish_reason chunk without tool calls.
|
||||
This tests the fix for regular content streaming with multiple choices.
|
||||
"""
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
|
||||
messages = [{"role": "user", "content": "Say hello in one word."}]
|
||||
|
||||
# Request with n=2 to get multiple choices, no tools
|
||||
response_stream = client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=messages,
|
||||
temperature=0.8,
|
||||
stream=True,
|
||||
max_tokens=10, # Keep it short
|
||||
n=2, # Multiple choices
|
||||
)
|
||||
|
||||
chunks = list(response_stream)
|
||||
|
||||
# Track finish_reason chunks for each index
|
||||
finish_reason_chunks = {}
|
||||
for chunk in chunks:
|
||||
if chunk.choices:
|
||||
for choice in chunk.choices:
|
||||
if choice.finish_reason is not None:
|
||||
index = choice.index
|
||||
if index not in finish_reason_chunks:
|
||||
finish_reason_chunks[index] = []
|
||||
finish_reason_chunks[index].append(choice.finish_reason)
|
||||
|
||||
# Verify we got finish_reason chunks for both indices
|
||||
self.assertEqual(
|
||||
len(finish_reason_chunks),
|
||||
2,
|
||||
f"Expected finish_reason chunks for 2 indices, got {len(finish_reason_chunks)}",
|
||||
)
|
||||
|
||||
# Verify both index 0 and 1 have finish_reason
|
||||
self.assertIn(
|
||||
0, finish_reason_chunks, "Missing finish_reason chunk for index 0"
|
||||
)
|
||||
self.assertIn(
|
||||
1, finish_reason_chunks, "Missing finish_reason chunk for index 1"
|
||||
)
|
||||
|
||||
# Verify the finish_reason is "stop" (regular completion)
|
||||
for index, reasons in finish_reason_chunks.items():
|
||||
self.assertIn(
|
||||
reasons[-1],
|
||||
["stop", "length"], # Could be either depending on how model responds
|
||||
f"Expected finish_reason 'stop' or 'length' for index {index}, got {reasons[-1]}",
|
||||
)
|
||||
|
||||
|
||||
class TestOpenAIPythonicFunctionCalling(CustomTestCase):
|
||||
"""Testcase：Verify the functionality of Python-style list-format function calling with pythonic parser for Llama-3.2-1B-Instruct model on Ascend NPU backend.
|
||||
Cover: Explicit format prompt verification, streaming call index integrity, and return validity of parallel tool calls.
|
||||
|
||||
[Test Category] Interface
|
||||
[Test Target] /v1/chat/completions
|
||||
"""
|
||||
|
||||
PYTHONIC_TOOLS = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get the current weather for a given location.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": {
|
||||
"type": "string",
|
||||
"description": "The name of the city or location.",
|
||||
}
|
||||
},
|
||||
"required": ["location"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_tourist_attractions",
|
||||
"description": "Get a list of top tourist attractions for a given city.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {
|
||||
"type": "string",
|
||||
"description": "The name of the city to find attractions for.",
|
||||
}
|
||||
},
|
||||
"required": ["city"],
|
||||
},
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
PYTHONIC_MESSAGES = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"You are a travel assistant. "
|
||||
"When asked to call functions, ALWAYS respond ONLY with a python list of function calls, "
|
||||
"using this format: [func_name1(param1=value1, param2=value2), func_name2(param=value)]. "
|
||||
"Do NOT use JSON, do NOT use variables, do NOT use any other format. "
|
||||
"Here is an example:\n"
|
||||
'[get_weather(location="Paris"), get_tourist_attractions(city="Paris")]'
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": (
|
||||
"I'm planning a trip to Tokyo next week. What's the weather like and what are some top tourist attractions? "
|
||||
"Propose parallel tool calls at once, using the python list of function calls format as shown above."
|
||||
),
|
||||
},
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = LLAMA_3_2_1B_INSTRUCT_WEIGHTS_PATH
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.api_key = "sk-123456"
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
api_key=cls.api_key,
|
||||
other_args=[
|
||||
"--attention-backend",
|
||||
"ascend",
|
||||
"--disable-cuda-graph",
|
||||
"--tool-call-parser",
|
||||
"pythonic",
|
||||
],
|
||||
)
|
||||
cls.base_url += "/v1"
|
||||
cls.tokenizer = get_tokenizer(cls.model)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_pythonic_tool_call_prompt(self):
|
||||
"""
|
||||
Test: Explicit prompt for pythonic tool call format without chat template.
|
||||
"""
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
response = client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=self.PYTHONIC_MESSAGES,
|
||||
tools=self.PYTHONIC_TOOLS,
|
||||
temperature=0.1,
|
||||
stream=False,
|
||||
)
|
||||
tool_calls = response.choices[0].message.tool_calls
|
||||
self.assertIsInstance(tool_calls, list, "No tool_calls found")
|
||||
self.assertGreaterEqual(len(tool_calls), 1)
|
||||
names = [tc.function.name for tc in tool_calls]
|
||||
self.assertTrue(
|
||||
"get_weather" in names or "get_tourist_attractions" in names,
|
||||
f"Function name '{names}' should container either 'get_weather' or 'get_tourist_attractions'",
|
||||
)
|
||||
|
||||
def test_pythonic_tool_call_streaming(self):
|
||||
"""
|
||||
Test: Streaming pythonic tool call format; assert tool_call index is present.
|
||||
"""
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
response_stream = client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=self.PYTHONIC_MESSAGES,
|
||||
tools=self.PYTHONIC_TOOLS,
|
||||
temperature=0.1,
|
||||
stream=True,
|
||||
)
|
||||
found_tool_calls = False
|
||||
found_index = False
|
||||
found_names = set()
|
||||
for chunk in response_stream:
|
||||
choice = chunk.choices[0]
|
||||
if getattr(choice.delta, "tool_calls", None):
|
||||
found_tool_calls = True
|
||||
tool_call = choice.delta.tool_calls[0]
|
||||
if hasattr(tool_call, "index") or (
|
||||
isinstance(tool_call, dict) and "index" in tool_call
|
||||
):
|
||||
found_index = True
|
||||
found_names.add(str(tool_call.function.name))
|
||||
|
||||
self.assertTrue(found_tool_calls, "No tool_calls found in streaming response")
|
||||
self.assertTrue(found_index, "No index field found in any streamed tool_call")
|
||||
self.assertTrue(
|
||||
"get_weather" in found_names or "get_tourist_attractions" in found_names,
|
||||
f"Function name '{found_names}' should container either 'get_weather' or 'get_tourist_attractions'",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -9,7 +9,6 @@ from prometheus_client.samples import Sample
|
||||
from sglang.srt.observability.metrics_collector import QueueCount
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import (
|
||||
register_amd_ci,
|
||||
register_cpu_ci,
|
||||
register_cuda_ci,
|
||||
)
|
||||
@@ -25,7 +24,6 @@ register_cuda_ci(
|
||||
stage="base-b",
|
||||
runner_config="1-gpu-small",
|
||||
)
|
||||
register_amd_ci(est_time=60, suite="stage-b-test-1-gpu-small-amd")
|
||||
register_cpu_ci(est_time=49, suite="stage-b-test-cpu-intel")
|
||||
|
||||
_MODEL_NAME = "Qwen/Qwen3-0.6B"
|
||||
|
||||
@@ -34,7 +34,7 @@ from sglang.srt.observability.trace import (
|
||||
)
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.srt.utils.network import get_zmq_socket
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
@@ -46,8 +46,7 @@ from sglang.test.test_utils import (
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# CI registration
|
||||
register_cuda_ci(est_time=172, stage="extra-a", runner_config="1-gpu-small")
|
||||
register_amd_ci(est_time=113, suite="stage-b-test-1-gpu-small-amd")
|
||||
register_cuda_ci(est_time=113, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
# ============================================================================
|
||||
|
||||
@@ -11,7 +11,7 @@ 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.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
@@ -27,8 +27,7 @@ try:
|
||||
except ImportError:
|
||||
_HAS_GRANIAN = False
|
||||
|
||||
register_cuda_ci(est_time=108, stage="base-b", runner_config="1-gpu-small")
|
||||
register_amd_ci(est_time=150, suite="stage-b-test-1-gpu-small-amd")
|
||||
register_cuda_ci(est_time=150, stage="base-b", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
@unittest.skipUnless(_HAS_GRANIAN, "granian not installed (pip install sglang[http2])")
|
||||
|
||||
@@ -17,7 +17,6 @@ import requests
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import (
|
||||
register_amd_ci,
|
||||
register_cpu_ci,
|
||||
register_cuda_ci,
|
||||
)
|
||||
@@ -29,8 +28,7 @@ from sglang.test.test_utils import (
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=51, stage="base-b", runner_config="1-gpu-large")
|
||||
register_amd_ci(est_time=140, suite="stage-b-test-1-gpu-small-amd")
|
||||
register_cuda_ci(est_time=50, stage="base-b", runner_config="1-gpu-large")
|
||||
register_cpu_ci(est_time=54, suite="stage-b-test-cpu-intel")
|
||||
|
||||
# System message to guide Llama3.2 to produce proper tool call format
|
||||
|
||||
@@ -3,9 +3,9 @@ import unittest
|
||||
|
||||
import openai
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.srt.utils import is_npu, kill_process_tree
|
||||
from sglang.srt.utils.hf_transformers_utils import get_tokenizer
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
from sglang.test.ci.ci_register import register_cuda_ci, register_npu_ci
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
@@ -15,8 +15,27 @@ from sglang.test.test_utils import (
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=210, stage="base-b", runner_config="1-gpu-large")
|
||||
register_amd_ci(est_time=73, suite="stage-b-test-1-gpu-small-amd")
|
||||
register_cuda_ci(est_time=100, stage="base-b", runner_config="1-gpu-large")
|
||||
# Backend-specific: Ascend uses a local model mirror and its native
|
||||
# attention backend, while sharing the protocol assertions below.
|
||||
register_npu_ci(est_time=400, suite="full-1-npu-a3", nightly=True)
|
||||
|
||||
|
||||
def _model_path():
|
||||
if is_npu():
|
||||
from sglang.test.ascend.test_ascend_utils import (
|
||||
LLAMA_3_2_1B_INSTRUCT_WEIGHTS_PATH,
|
||||
)
|
||||
|
||||
return LLAMA_3_2_1B_INSTRUCT_WEIGHTS_PATH
|
||||
return DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
|
||||
|
||||
def _server_args(parser):
|
||||
args = ["--tool-call-parser", parser]
|
||||
if is_npu():
|
||||
args[:0] = ["--attention-backend", "ascend", "--disable-cuda-graph"]
|
||||
return args
|
||||
|
||||
|
||||
class TestOpenAIServerFunctionCalling(CustomTestCase):
|
||||
@@ -36,8 +55,7 @@ class TestOpenAIServerFunctionCalling(CustomTestCase):
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
# Replace with the model name needed for testing; if not required, reuse DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
cls.model = _model_path()
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.api_key = "sk-123456"
|
||||
|
||||
@@ -47,11 +65,7 @@ class TestOpenAIServerFunctionCalling(CustomTestCase):
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
api_key=cls.api_key,
|
||||
other_args=[
|
||||
# If your server needs extra parameters to test function calling, please add them here.
|
||||
"--tool-call-parser",
|
||||
"llama3",
|
||||
],
|
||||
other_args=_server_args("llama3"),
|
||||
)
|
||||
cls.base_url += "/v1"
|
||||
cls.tokenizer = get_tokenizer(cls.model)
|
||||
@@ -97,7 +111,7 @@ class TestOpenAIServerFunctionCalling(CustomTestCase):
|
||||
{"role": "system", "content": self.SYSTEM_MESSAGE},
|
||||
{"role": "user", "content": "Compute (3+5)"},
|
||||
]
|
||||
response = client.chat.completions.create(
|
||||
request = dict(
|
||||
model=self.model,
|
||||
max_tokens=2048,
|
||||
messages=messages,
|
||||
@@ -105,8 +119,12 @@ class TestOpenAIServerFunctionCalling(CustomTestCase):
|
||||
top_p=0.8,
|
||||
stream=False,
|
||||
tools=tools,
|
||||
tool_choice="required",
|
||||
)
|
||||
# Ascend keeps the historical auto-choice coverage; CUDA forces the
|
||||
# call so this assertion never depends on a stochastic model decision.
|
||||
if not is_npu():
|
||||
request["tool_choice"] = "required"
|
||||
response = client.chat.completions.create(**request)
|
||||
|
||||
tool_calls = response.choices[0].message.tool_calls
|
||||
|
||||
@@ -843,7 +861,7 @@ class TestOpenAIPythonicFunctionCalling(CustomTestCase):
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
cls.model = _model_path()
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.api_key = "sk-123456"
|
||||
cls.process = popen_launch_server(
|
||||
@@ -851,10 +869,7 @@ class TestOpenAIPythonicFunctionCalling(CustomTestCase):
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
api_key=cls.api_key,
|
||||
other_args=[
|
||||
"--tool-call-parser",
|
||||
"pythonic",
|
||||
],
|
||||
other_args=_server_args("pythonic"),
|
||||
)
|
||||
cls.base_url += "/v1"
|
||||
cls.tokenizer = get_tokenizer(cls.model)
|
||||
@@ -922,6 +937,7 @@ class TestOpenAIPythonicFunctionCalling(CustomTestCase):
|
||||
is_rust_server_built(),
|
||||
"embedded rust server extension not built",
|
||||
)
|
||||
@unittest.skipIf(is_npu(), "the embedded Rust server is not an Ascend path")
|
||||
class TestOpenAIFunctionCallingWithRust(TestOpenAIServerFunctionCalling):
|
||||
"""Run the registered unary/streaming function-call suite through Rust."""
|
||||
|
||||
@@ -946,6 +962,7 @@ class TestOpenAIFunctionCallingWithRust(TestOpenAIServerFunctionCalling):
|
||||
is_rust_server_built(),
|
||||
"embedded rust server extension not built",
|
||||
)
|
||||
@unittest.skipIf(is_npu(), "the embedded Rust server is not an Ascend path")
|
||||
class TestOpenAIPythonicFunctionCallingWithRust(TestOpenAIPythonicFunctionCalling):
|
||||
"""Run Pythonic unary/streaming tool calls through Rust."""
|
||||
|
||||
|
||||
@@ -12,7 +12,6 @@ import openai
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.srt.utils.hf_transformers_utils import get_tokenizer
|
||||
from sglang.test.ci.ci_register import (
|
||||
register_amd_ci,
|
||||
register_cpu_ci,
|
||||
register_cuda_ci,
|
||||
)
|
||||
@@ -26,8 +25,7 @@ from sglang.test.test_utils import (
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=59, stage="base-b", runner_config="1-gpu-large")
|
||||
register_amd_ci(est_time=41, suite="stage-b-test-1-gpu-small-amd")
|
||||
register_cuda_ci(est_time=58, stage="base-b", runner_config="1-gpu-large")
|
||||
register_cpu_ci(est_time=101, suite="stage-b-test-cpu-intel")
|
||||
|
||||
|
||||
|
||||
@@ -2,7 +2,6 @@ import unittest
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import (
|
||||
register_amd_ci,
|
||||
register_cpu_ci,
|
||||
register_cuda_ci,
|
||||
)
|
||||
@@ -14,8 +13,7 @@ from sglang.test.test_utils import (
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=63, stage="base-b", runner_config="1-gpu-small")
|
||||
register_amd_ci(est_time=60, suite="stage-b-test-1-gpu-small-amd")
|
||||
register_cuda_ci(est_time=52, stage="base-b", runner_config="1-gpu-small")
|
||||
register_cpu_ci(est_time=83, suite="stage-b-test-cpu-intel")
|
||||
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ import openai
|
||||
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.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
@@ -13,8 +13,7 @@ from sglang.test.test_utils import (
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=50, stage="base-b", runner_config="1-gpu-large")
|
||||
register_amd_ci(est_time=31, suite="stage-b-test-1-gpu-small-amd")
|
||||
register_cuda_ci(est_time=49, stage="base-b", runner_config="1-gpu-large")
|
||||
|
||||
|
||||
class TestRequestLengthValidation(CustomTestCase):
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user