From f6a6f5bf1e18155acbed3e72d37f46e0cd2010ca Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sat, 8 Aug 2026 01:42:46 -0700 Subject: [PATCH] [CI] Trim redundant nightly test registrations (#34070) Co-authored-by: Baizhou Zhang --- .../test_deepseek_r1_fp8_trtllm_backend.py | 90 -- .../backends/test_qwen3_fp4_trtllm_gen_moe.py | 69 - .../test_bench_serving_functionality.py | 191 --- .../piecewise/test_pcg_glm52_fp4.py | 71 - .../piecewise/test_pcg_glm52_fp8_tp8.py | 75 - .../test_disaggregation_dwdp_mimo.py | 122 -- test/registered/gb300/test_kimi_k25.py | 63 - test/registered/gb300/test_qwen35_nvfp4.py | 83 -- .../lora/test_chunked_sgmv_backend.py | 1261 ----------------- .../lora/test_embedding_lora_support.py | 232 --- test/registered/lora/test_lora_openai_api.py | 7 +- test/registered/lora/test_lora_radix_cache.py | 82 -- .../registered/lora/test_lora_tied_lm_head.py | 224 --- .../models_e2e/test_qwen3_next_models_fp4.py | 34 - .../perf/test_dpsk_v3_fp4_4gpu_perf.py | 77 - .../registered/perf/test_gpt_oss_4gpu_perf.py | 60 - .../test_fa_skip_kv_cache_piecewise_nan.py | 162 --- .../quant/test_kimi_k25_nvfp4_eagle.py | 68 - .../quant/test_kimi_k26_nvfp4_dflash.py | 70 - .../radix_cache/test_cpp_radix_cache.py | 41 - ...est_eagle_infer_beta_dp_attention_large.py | 100 -- .../utils/test_model_file_verifier.py | 3 +- test/registered/utils/test_request_logger.py | 309 ---- .../utils/test_scheduler_status_logger.py | 79 -- 24 files changed, 2 insertions(+), 3571 deletions(-) delete mode 100644 test/registered/backends/test_deepseek_r1_fp8_trtllm_backend.py delete mode 100644 test/registered/backends/test_qwen3_fp4_trtllm_gen_moe.py delete mode 100644 test/registered/bench_fn/test_bench_serving_functionality.py delete mode 100644 test/registered/cuda_graph/piecewise/test_pcg_glm52_fp4.py delete mode 100644 test/registered/cuda_graph/piecewise/test_pcg_glm52_fp8_tp8.py delete mode 100644 test/registered/disaggregation/test_disaggregation_dwdp_mimo.py delete mode 100644 test/registered/gb300/test_kimi_k25.py delete mode 100644 test/registered/gb300/test_qwen35_nvfp4.py delete mode 100644 test/registered/lora/test_chunked_sgmv_backend.py delete mode 100644 test/registered/lora/test_embedding_lora_support.py delete mode 100644 test/registered/lora/test_lora_radix_cache.py delete mode 100644 test/registered/lora/test_lora_tied_lm_head.py delete mode 100644 test/registered/models_e2e/test_qwen3_next_models_fp4.py delete mode 100644 test/registered/perf/test_dpsk_v3_fp4_4gpu_perf.py delete mode 100644 test/registered/perf/test_gpt_oss_4gpu_perf.py delete mode 100644 test/registered/prefill_only/test_fa_skip_kv_cache_piecewise_nan.py delete mode 100644 test/registered/quant/test_kimi_k25_nvfp4_eagle.py delete mode 100644 test/registered/quant/test_kimi_k26_nvfp4_dflash.py delete mode 100644 test/registered/radix_cache/test_cpp_radix_cache.py delete mode 100644 test/registered/spec/eagle/test_eagle_infer_beta_dp_attention_large.py delete mode 100644 test/registered/utils/test_request_logger.py delete mode 100644 test/registered/utils/test_scheduler_status_logger.py diff --git a/test/registered/backends/test_deepseek_r1_fp8_trtllm_backend.py b/test/registered/backends/test_deepseek_r1_fp8_trtllm_backend.py deleted file mode 100644 index 385caaa6d..000000000 --- a/test/registered/backends/test_deepseek_r1_fp8_trtllm_backend.py +++ /dev/null @@ -1,90 +0,0 @@ -import unittest -from types import SimpleNamespace - -from sglang.srt.utils import kill_process_tree -from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.run_eval import run_eval -from sglang.test.test_utils import ( - DEFAULT_URL_FOR_TEST, - CustomTestCase, - popen_launch_server, - try_cached_model, -) - -register_cuda_ci(est_time=3600, suite="nightly-8-gpu-b200", nightly=True) - -FULL_DEEPSEEK_V3_MODEL_PATH = "deepseek-ai/DeepSeek-V3-0324" -SERVER_LAUNCH_TIMEOUT = 1000 - - -class TestDeepseekR1Fp8Flashinfer(CustomTestCase): - @classmethod - def setUpClass(cls): - cls.model = try_cached_model(FULL_DEEPSEEK_V3_MODEL_PATH) - cls.base_url = DEFAULT_URL_FOR_TEST - other_args = [ - "--trust-remote-code", - "--disable-radix-cache", - "--max-running-requests", - "512", - "--chunked-prefill-size", - "8192", - "--mem-fraction-static", - "0.9", - "--cuda-graph-max-bs-decode", - "128", - "--max-prefill-tokens", - "8192", - "--kv-cache-dtype", - "fp8_e4m3", - "--quantization", - "fp8", - "--tensor-parallel-size", - "8", - "--data-parallel-size", - "1", - "--expert-parallel-size", - "1", - "--scheduler-recv-interval", - "10", - "--stream-interval", - "10", - "--attention-backend", - "trtllm_mla", - "--fp8-gemm-backend", - "flashinfer_trtllm", - "--moe-runner-backend", - "flashinfer_trtllm", - "--enable-symm-mem", - "--model-loader-extra-config", - '{"enable_multithread_load": true,"num_threads": 64}', - ] - cls.process = popen_launch_server( - cls.model, - cls.base_url, - timeout=SERVER_LAUNCH_TIMEOUT, - other_args=other_args, - ) - - @classmethod - def tearDownClass(cls): - kill_process_tree(cls.process.pid) - - def test_gsm8k(self): - args = SimpleNamespace( - base_url=self.base_url, - model=self.model, - eval_name="gsm8k", - api="completion", - max_tokens=512, - num_examples=512, - num_threads=512, - ) - metrics = run_eval(args) - print(f"Eval accuracy of GSM8K: {metrics=}") - - self.assertGreater(metrics["score"], 0.92) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/backends/test_qwen3_fp4_trtllm_gen_moe.py b/test/registered/backends/test_qwen3_fp4_trtllm_gen_moe.py deleted file mode 100644 index 4011d3b2a..000000000 --- a/test/registered/backends/test_qwen3_fp4_trtllm_gen_moe.py +++ /dev/null @@ -1,69 +0,0 @@ -import unittest -from types import SimpleNamespace - -from sglang.srt.utils import get_device_sm, kill_process_tree -from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.run_eval import run_eval -from sglang.test.test_utils import ( - DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - DEFAULT_URL_FOR_TEST, - CustomTestCase, - popen_launch_server, -) - -# modelopt_fp4 requires SM 100+ (Blackwell) -register_cuda_ci(est_time=300, suite="nightly-1-gpu", nightly=True) - - -@unittest.skipIf( - get_device_sm() < 100, "Test requires CUDA SM 100 or higher (Blackwell)" -) -class TestFlashinferTrtllmGenMoeBackend(CustomTestCase): - @classmethod - def setUpClass(cls): - cls.model = "nvidia/Qwen3-30B-A3B-NVFP4" - cls.base_url = DEFAULT_URL_FOR_TEST - cls.process = popen_launch_server( - cls.model, - cls.base_url, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=[ - "--moe-runner-backend", - "flashinfer_trtllm", - "--quantization", - "modelopt_fp4", - "--trust-remote-code", - "--disable-radix-cache", - "--max-running-requests", - "1024", - "--chunked-prefill-size", - "16384", - "--mem-fraction-static", - "0.89", - "--max-prefill-tokens", - "16384", - ], - ) - - @classmethod - def tearDownClass(cls): - kill_process_tree(cls.process.pid) - - def test_gsm8k(self): - args = SimpleNamespace( - base_url=self.base_url, - model=self.model, - eval_name="gsm8k", - api="completion", - max_tokens=512, - num_examples=1319, - num_threads=1319, - num_shots=8, - ) - metrics = run_eval(args) - print(f"{metrics=}") - self.assertGreater(metrics["score"], 0.88) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/bench_fn/test_bench_serving_functionality.py b/test/registered/bench_fn/test_bench_serving_functionality.py deleted file mode 100644 index 4f0126a52..000000000 --- a/test/registered/bench_fn/test_bench_serving_functionality.py +++ /dev/null @@ -1,191 +0,0 @@ -import json -import tempfile -import threading -import time -import unittest -from http.server import BaseHTTPRequestHandler, HTTPServer -from pathlib import Path - -from sglang.benchmark.serving import run_benchmark -from sglang.benchmark.utils import parse_custom_headers -from sglang.srt.constants import HEALTH_CHECK_RID_PREFIX -from sglang.srt.utils import kill_process_tree -from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci -from sglang.test.test_utils import ( - DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - DEFAULT_URL_FOR_TEST, - CustomTestCase, - get_benchmark_args, - popen_launch_server, -) - -register_cuda_ci(est_time=300, suite="nightly-1-gpu", nightly=True) -register_amd_ci(est_time=300, suite="nightly-amd-1-gpu", nightly=True) - -MODEL = "Qwen/Qwen3-0.6B" -NUM_CONVERSATIONS, NUM_TURNS = 4, 3 - - -class TestBenchServingFunctionality(CustomTestCase): - def test_gsp_multi_turn(self): - with tempfile.TemporaryDirectory() as temp_dir: - process = popen_launch_server( - MODEL, - DEFAULT_URL_FOR_TEST, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=[ - "--mem-fraction-static", - "0.7", - "--log-requests", - "--log-requests-level", - "3", - "--log-requests-format", - "json", - "--log-requests-target", - "stdout", - temp_dir, - ], - ) - try: - args = get_benchmark_args( - base_url=DEFAULT_URL_FOR_TEST, - backend="sglang-oai-chat", - tokenizer=MODEL, - dataset_name="generated-shared-prefix", - num_prompts=NUM_CONVERSATIONS, - request_rate=float("inf"), - gsp_num_groups=2, - gsp_prompts_per_group=2, - gsp_system_prompt_len=64, - gsp_question_len=16, - gsp_output_len=16, - gsp_num_turns=NUM_TURNS, - ) - args.warmup_requests = 0 - res = run_benchmark(args) - self.assertEqual(res["completed"], NUM_CONVERSATIONS * NUM_TURNS) - - time.sleep(1) - logs = "".join(f.read_text() for f in Path(temp_dir).glob("*.log")) - self._verify_multi_turn_logs(logs) - finally: - kill_process_tree(process.pid) - - def _verify_multi_turn_logs(self, content: str): - reqs = [] - for line in content.splitlines(): - idx = line.find("{") - if idx == -1: - continue - try: - obj = json.loads(line[idx:]) - except json.JSONDecodeError: - continue - if obj.get("event") != "request.finished": - continue - text = obj.get("obj", {}).get("text") - rid = obj.get("rid", "") - if text and not rid.startswith(HEALTH_CHECK_RID_PREFIX): - reqs.append(text) - - self.assertGreaterEqual(len(reqs), NUM_CONVERSATIONS * NUM_TURNS) - - # Verify prefix relationships - reqs_sorted = sorted(reqs, key=len) - prefix_count = 0 - for i, text in enumerate(reqs_sorted): - for j in range(i + 1, len(reqs_sorted)): - if reqs_sorted[j].startswith(text): - prefix_count += 1 - break - - expected = NUM_CONVERSATIONS * (NUM_TURNS - 1) - self.assertGreaterEqual( - prefix_count, expected, f"Expected at least {expected} prefix pairs" - ) - - -class TestBenchServingCustomHeaders(CustomTestCase): - def test_parse_custom_headers(self): - headers = parse_custom_headers(["MyHeader=MY_VALUE", "Another=value=hello"]) - self.assertEqual(headers, {"MyHeader": "MY_VALUE", "Another": "value=hello"}) - - headers = parse_custom_headers(["InvalidNoEquals"]) - self.assertEqual(headers, {}) - - headers = parse_custom_headers(["=NoKey"]) - self.assertEqual(headers, {}) - - # TODO: Using well-implemented mock server, e.g. the on in sgl-router - def test_custom_headers_sent_to_server(self): - import queue - - received_requests = queue.Queue() - - class HeaderEchoHandler(BaseHTTPRequestHandler): - def _handle(self): - received_requests.put( - { - "method": self.command, - "path": self.path, - "headers": dict(self.headers), - } - ) - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.end_headers() - if self.path == "/v1/models": - self.wfile.write(json.dumps({"data": [{"id": "gpt2"}]}).encode()) - elif self.path == "/generate": - self.wfile.write( - json.dumps( - {"text": "ok", "meta_info": {"completion_tokens": 1}} - ).encode() - ) - else: - self.wfile.write(json.dumps({}).encode()) - - do_GET = do_POST = _handle - - server = HTTPServer(("127.0.0.1", 0), HeaderEchoHandler) - port = server.server_address[1] - server_thread = threading.Thread(target=server.serve_forever) - server_thread.daemon = True - server_thread.start() - - try: - args = get_benchmark_args( - base_url=f"http://127.0.0.1:{port}", - backend="sglang", - dataset_name="random", - tokenizer="gpt2", - num_prompts=1, - random_input_len=8, - random_output_len=8, - header=["X-Custom-Test=TestValue123", "X-Another=AnotherVal"], - ) - args.warmup_requests = 0 - args.disable_tqdm = True - run_benchmark(args) - except Exception: - pass - finally: - server.shutdown() - - all_reqs = [] - while not received_requests.empty(): - all_reqs.append(received_requests.get_nowait()) - - generate_reqs = [r for r in all_reqs if r["path"] == "/generate"] - self.assertGreater( - len(generate_reqs), - 0, - f"No /generate request. All: {[r['path'] for r in all_reqs]}", - ) - headers = generate_reqs[0]["headers"] - self.assertEqual(headers.get("X-Custom-Test"), "TestValue123") - self.assertEqual(headers.get("X-Another"), "AnotherVal") - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/cuda_graph/piecewise/test_pcg_glm52_fp4.py b/test/registered/cuda_graph/piecewise/test_pcg_glm52_fp4.py deleted file mode 100644 index 41a978113..000000000 --- a/test/registered/cuda_graph/piecewise/test_pcg_glm52_fp4.py +++ /dev/null @@ -1,71 +0,0 @@ -import unittest -from types import SimpleNamespace - -from sglang.srt.utils import kill_process_tree -from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.run_eval import run_eval -from sglang.test.test_utils import ( - DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - DEFAULT_URL_FOR_TEST, - CustomTestCase, - popen_launch_server, -) - -register_cuda_ci(est_time=900, suite="nightly-4-gpu-b200", nightly=True) - -GLM52_FP4_MODEL = "nvidia/GLM-5.2-NVFP4" - - -class TestPCGGlm52Fp4(CustomTestCase): - """PCG prefill on GLM-5.2-NVFP4 (DSA model, TP=4, B200). - - GLM-5.2 uses GlmMoeDsaForCausalLM (DSA attention). This test verifies that - piecewise CUDA graph works correctly after the DSA indexer was updated to - cache k_fp8/k_scale for PCG-compatible prefill. - """ - - @classmethod - def setUpClass(cls): - cls.model = GLM52_FP4_MODEL - cls.base_url = DEFAULT_URL_FOR_TEST - cls.process = popen_launch_server( - cls.model, - cls.base_url, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=[ - "--tp-size", - "4", - "--trust-remote-code", - "--reasoning-parser", - "glm45", - "--tool-call-parser", - "glm47", - "--quantization", - "modelopt_fp4", - "--disable-flashinfer-autotune", - "--cuda-graph-backend-prefill=tc_piecewise", - "--model-loader-extra-config", - '{"enable_multithread_load": true, "num_threads": 64}', - ], - ) - - @classmethod - def tearDownClass(cls): - kill_process_tree(cls.process.pid) - - def test_gsm8k(self): - args = SimpleNamespace( - base_url=self.base_url, - model=self.model, - eval_name="gsm8k", - num_examples=200, - num_threads=200, - max_tokens=4096, - ) - metrics = run_eval(args) - print(f"{metrics=}") - self.assertGreater(metrics["score"], 0.92) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/cuda_graph/piecewise/test_pcg_glm52_fp8_tp8.py b/test/registered/cuda_graph/piecewise/test_pcg_glm52_fp8_tp8.py deleted file mode 100644 index 3c26bb842..000000000 --- a/test/registered/cuda_graph/piecewise/test_pcg_glm52_fp8_tp8.py +++ /dev/null @@ -1,75 +0,0 @@ -import unittest -from types import SimpleNamespace - -from sglang.srt.utils import kill_process_tree -from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.run_eval import run_eval -from sglang.test.test_utils import ( - DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - DEFAULT_URL_FOR_TEST, - CustomTestCase, - popen_launch_server, -) - -register_cuda_ci(est_time=900, suite="nightly-8-gpu-h200", nightly=True) - -GLM52_FP8_MODEL = "zai-org/GLM-5.2-FP8" - - -class TestBCGGlm52Fp8TP8(CustomTestCase): - """Breakable CUDA graph prefill on GLM-5.2-FP8 (DSA model, TP=8, H200).""" - - @classmethod - def setUpClass(cls): - cls.model = GLM52_FP8_MODEL - cls.base_url = DEFAULT_URL_FOR_TEST - cls.process = popen_launch_server( - cls.model, - cls.base_url, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=[ - "--tp-size", - "8", - "--trust-remote-code", - "--reasoning-parser", - "glm45", - "--tool-call-parser", - "glm47", - "--mem-fraction-static", - "0.8", - "--disable-flashinfer-autotune", - "--cuda-graph-backend-prefill=breakable", - # Small chunks => many prefill iterations, each <= the 2048 - # capture max, so every prefill batch replays the BCG graph and - # exercises the DSA split-op / dual-stream / MLA-fusion paths. - "--chunked-prefill-size", - "512", - "--model-loader-extra-config", - '{"enable_multithread_load": true, "num_threads": 64}', - ], - env={ - "SGLANG_ENABLE_PCG_DSV2_DUAL_STREAM": "1", - }, - ) - - @classmethod - def tearDownClass(cls): - if hasattr(cls, "process") and cls.process: - kill_process_tree(cls.process.pid) - - def test_gsm8k(self): - args = SimpleNamespace( - base_url=self.base_url, - model=self.model, - eval_name="gsm8k", - num_examples=200, - num_threads=200, - max_tokens=4096, - ) - metrics = run_eval(args) - print(f"{metrics=}") - self.assertGreater(metrics["score"], 0.92) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/disaggregation/test_disaggregation_dwdp_mimo.py b/test/registered/disaggregation/test_disaggregation_dwdp_mimo.py deleted file mode 100644 index 9b8ee5ea5..000000000 --- a/test/registered/disaggregation/test_disaggregation_dwdp_mimo.py +++ /dev/null @@ -1,122 +0,0 @@ -import unittest -from types import SimpleNamespace - -from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.run_eval import run_eval -from sglang.test.server_fixtures.disaggregation_fixture import ( - PDDisaggregationServerBase, -) -from sglang.test.test_utils import ( - DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - is_in_ci, - popen_launch_pd_server, -) - -register_cuda_ci(est_time=900, suite="nightly-8-gpu-b200", nightly=True) - -MIMO_V2_MODEL_PATH = "XiaomiMiMo/MiMo-V2.5" -GSM8K_BASELINE_ACCURACY = 0.93 - - -@unittest.skipIf(is_in_ci(), "Temporarily disable the flaky test.") -class TestDisaggregationDWDPMiMo(PDDisaggregationServerBase): - """PD disagg with DWDP prefill (4 GPUs) and DP-attention decode (4 GPUs).""" - - NUM_PREFILL_GPUS = 4 - NUM_DECODE_GPUS = 4 - - @classmethod - def setUpClass(cls): - super().setUpClass() - cls.model = MIMO_V2_MODEL_PATH - - cls.start_prefill() - cls.start_decode() - - cls.wait_server_ready(cls.prefill_url + "/health", process=cls.process_prefill) - cls.wait_server_ready(cls.decode_url + "/health", process=cls.process_decode) - - cls.launch_lb() - - @classmethod - def start_prefill(cls): - prefill_args = [ - "--trust-remote-code", - "--disaggregation-mode", - "prefill", - "--disaggregation-bootstrap-port", - cls.bootstrap_port, - "--tp", - str(cls.NUM_PREFILL_GPUS), - "--dwdp-size", - str(cls.NUM_PREFILL_GPUS), - "--mm-enable-dp-encoder", - "--attention-backend", - "fa4", - "--mem-fraction-static", - "0.78", - ] - prefill_args += cls.transfer_backend + cls.rdma_devices - cls.process_prefill = popen_launch_pd_server( - cls.model, - cls.prefill_url, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=prefill_args, - ) - - @classmethod - def start_decode(cls): - decode_args = [ - "--trust-remote-code", - "--disaggregation-mode", - "decode", - "--disaggregation-bootstrap-port", - cls.bootstrap_port, - "--tp", - str(cls.NUM_DECODE_GPUS), - "--dp", - str(cls.NUM_DECODE_GPUS), - "--enable-dp-attention", - "--moe-dense-tp-size", - "1", - "--ep-size", - str(cls.NUM_DECODE_GPUS), - "--attention-backend", - "fa4", - "--mem-fraction-static", - "0.78", - "--base-gpu-id", - str(cls.NUM_PREFILL_GPUS), - ] - decode_args += cls.transfer_backend + cls.rdma_devices - cls.process_decode = popen_launch_pd_server( - cls.model, - cls.decode_url, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=decode_args, - ) - - def test_gsm8k(self): - metrics = run_eval( - SimpleNamespace( - base_url=self.base_url, - model=self.model, - eval_name="gsm8k", - api="chat", - num_shots=5, - num_examples=200, - max_tokens=4096, - num_threads=8, - repeat=1, - temperature=0.0, - top_p=1.0, - host="http://127.0.0.1", - port=int(self.base_url.split(":")[-1]), - ) - ) - print(f"{metrics=}") - self.assertGreaterEqual(metrics["score"], GSM8K_BASELINE_ACCURACY) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/gb300/test_kimi_k25.py b/test/registered/gb300/test_kimi_k25.py deleted file mode 100644 index 8937263d4..000000000 --- a/test/registered/gb300/test_kimi_k25.py +++ /dev/null @@ -1,63 +0,0 @@ -import unittest - -from sglang.test.accuracy_test_runner import AccuracyTestParams -from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.performance_test_runner import PerformanceTestParams -from sglang.test.run_combined_tests import run_combined_tests -from sglang.test.test_utils import ModelLaunchSettings - -register_cuda_ci( - est_time=7200, - suite="nightly-4-gpu-gb300-kimi-k25", - nightly=True, - disabled="not needed", -) - -MODEL_PATH = "moonshotai/Kimi-K2.5" - -COMMON_ARGS = [ - "--trust-remote-code", - "--reasoning-parser=kimi_k2", - "--tool-call-parser=kimi_k2", - "--mem-fraction-static=0.8", - "--enable-multimodal", - "--enable-metrics", -] - - -class TestKimiK25(unittest.TestCase): - """Kimi-K2.5 (native INT4) on GB300 (4x GB300 NVL4, tp=4). - - No EAGLE/MTP support for Kimi-K2.5 — only TP and TP+DP+DPA variants. - """ - - def test_kimi_k25(self): - variants = [ - ModelLaunchSettings( - MODEL_PATH, - tp_size=4, - extra_args=COMMON_ARGS, - variant="TP4", - ), - ModelLaunchSettings( - MODEL_PATH, - tp_size=4, - extra_args=COMMON_ARGS + ["--dp-size=4", "--enable-dp-attention"], - variant="TP4+DP4+DPA", - ), - ] - - run_combined_tests( - models=variants, - test_name="Kimi-K2.5", - accuracy_params=AccuracyTestParams( - dataset="mmmu-pro", baseline_accuracy=0.69, repeat=1, max_tokens=32768 - ), - performance_params=PerformanceTestParams( - result_dir="performance_results_gb300", - ), - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/gb300/test_qwen35_nvfp4.py b/test/registered/gb300/test_qwen35_nvfp4.py deleted file mode 100644 index e8d47f42a..000000000 --- a/test/registered/gb300/test_qwen35_nvfp4.py +++ /dev/null @@ -1,83 +0,0 @@ -import unittest - -from sglang.test.accuracy_test_runner import AccuracyTestParams -from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.performance_test_runner import PerformanceTestParams -from sglang.test.run_combined_tests import run_combined_tests -from sglang.test.test_utils import ModelLaunchSettings - -register_cuda_ci( - est_time=7200, - suite="nightly-4-gpu-gb300-qwen35-nvfp4", - nightly=True, - disabled="not needed", -) - -MODEL_PATH = "nvidia/Qwen3.5-397B-A17B-NVFP4" - -COMMON_ARGS = [ - "--trust-remote-code", - "--reasoning-parser=qwen3", - "--tool-call-parser=qwen3_coder", - "--quantization=modelopt_fp4", - "--fp4-gemm-backend=flashinfer_cutlass", - "--moe-runner-backend=flashinfer_trtllm", - "--kv-cache-dtype=fp8_e4m3", - "--enable-flashinfer-allreduce-fusion", - "--attention-backend=trtllm_mha", - "--mem-fraction-static=0.8", - "--enable-multimodal", - "--enable-metrics", -] - -MTP_ARGS = [ - "--speculative-algorithm=EAGLE", - "--speculative-num-steps=3", - "--speculative-eagle-topk=1", - "--speculative-num-draft-tokens=4", - "--mamba-scheduler-strategy=extra_buffer", - "--page-size=64", -] - - -class TestQwen35Nvfp4(unittest.TestCase): - """Qwen3.5-397B NVFP4 on GB300 (4x GB300 NVL4, tp=4).""" - - def test_qwen35_nvfp4(self): - variants = [ - ModelLaunchSettings( - MODEL_PATH, - tp_size=4, - extra_args=COMMON_ARGS, - variant="TP4", - ), - ModelLaunchSettings( - MODEL_PATH, - tp_size=4, - extra_args=COMMON_ARGS + ["--dp-size=4", "--enable-dp-attention"], - variant="TP4+DP4+DPA", - ), - ModelLaunchSettings( - MODEL_PATH, - tp_size=4, - extra_args=COMMON_ARGS - + ["--dp-size=4", "--enable-dp-attention"] - + MTP_ARGS, - variant="TP4+DP4+DPA+MTP", - ), - ] - - run_combined_tests( - models=variants, - test_name="Qwen3.5-397B-NVFP4", - accuracy_params=AccuracyTestParams( - dataset="mmmu-pro", baseline_accuracy=0.76, repeat=1, max_tokens=32768 - ), - performance_params=PerformanceTestParams( - result_dir="performance_results_gb300", - ), - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/lora/test_chunked_sgmv_backend.py b/test/registered/lora/test_chunked_sgmv_backend.py deleted file mode 100644 index dbbba922b..000000000 --- a/test/registered/lora/test_chunked_sgmv_backend.py +++ /dev/null @@ -1,1261 +0,0 @@ -import random -import unittest -from enum import Enum -from typing import List, Optional, Tuple - -import torch - -from sglang.kernels.ops.gemm.chunked_embedding_lora_a import ( - chunked_embedding_lora_a_forward, -) -from sglang.kernels.ops.gemm.chunked_sgmv_expand import ( - _chunked_lora_expand_kernel, - chunked_sgmv_lora_expand_forward, -) -from sglang.kernels.ops.gemm.chunked_sgmv_shrink import ( - _chunked_lora_shrink_kernel, - chunked_sgmv_lora_shrink_forward, -) -from sglang.kernels.ops.gemm.kv_b_lora_absorbed import ( - step_a_q_fwd, - step_a_v_fwd, - step_b_q_fwd, - step_b_v_fwd, -) -from sglang.srt.layers.logits_processor import LogitsMetadata, LogitsProcessor -from sglang.srt.lora.backend.chunked_backend import ChunkedSgmvLoRABackend -from sglang.srt.lora.utils import LoRABatchInfo, get_lm_head_pruned_lens -from sglang.srt.model_executor.forward_batch_info import ForwardMode -from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.lora_utils import ( - reference_embedding_lora_a_shrink, - reference_sgmv_expand, - reference_sgmv_shrink, -) - -CHUNK_SIZE = 16 - -register_cuda_ci(est_time=60, suite="nightly-1-gpu", nightly=True) - - -def reset_kernel_cache(): - _chunked_lora_shrink_kernel._clear_cache() - _chunked_lora_expand_kernel._clear_cache() - - -class BatchComposition(Enum): - UNIFORM = "uniform" - MIXED = "mixed" - SKEWED = "skewed" - NONE = "_NO_LORA_" - - -class BatchMode(Enum): - PREFILL = "prefill" - DECODE = "decode" - TARGET_VERIFY = "verify" - - -class TestChunkedSGMV(unittest.TestCase): - - # Test configuration constants - RTOL = 1e-3 - ATOL = 1e-3 - DEFAULT_BATCH_SIZE = 8 - - def _compare_shrink_outputs( - self, - chunked_output: torch.Tensor, - reference_output: torch.Tensor, - seq_lengths: List[int], - lora_assignments: List[int], - batch_info: LoRABatchInfo, - num_slices: int, - test_name: str, - ): - """ - Compare only the valid portions of shrink outputs. - - The chunked SGMV shrink kernel only guarantees correctness for - output[seq_start:seq_end, :rank * num_slices] for each sequence. - """ - lora_ranks = batch_info.lora_ranks.cpu().numpy() - - token_offset = 0 - for seq_idx, (lora_idx, seq_len) in enumerate( - zip(lora_assignments, seq_lengths) - ): - if seq_len == 0: - continue - - rank = lora_ranks[lora_idx] - - if rank > 0: - # Only compare the valid columns for this sequence - valid_cols = num_slices * rank - - chunked_seq = chunked_output[ - token_offset : token_offset + seq_len, :valid_cols - ] - reference_seq = reference_output[ - token_offset : token_offset + seq_len, :valid_cols - ] - - torch.testing.assert_close( - chunked_seq, - reference_seq, - rtol=self.RTOL, - atol=self.ATOL, - msg=f"Shrink operation failed for {test_name}, sequence {seq_idx} ({lora_idx})", - ) - - token_offset += seq_len - - def setUp(self): - """Set up common test parameters""" - torch.manual_seed(42) - random.seed(42) - - self.device = torch.device("cuda") - self.dtype = torch.float16 - self.input_dim = 2560 # Hidden dimension - self.max_seq_len = 1024 - self.vocab_size = 32000 # Vocabulary size for embedding tests - - # LoRA configurations: name -> (rank, output_q, output_k, output_v) - self.lora_configs = { - "lora_A": (8, 4096, 1024, 1024), - "lora_B": (16, 4096, 1024, 1024), - "lora_C": (32, 4096, 1024, 1024), - "_NO_LORA_": (0, 4096, 1024, 1024), - } - - # QKV slice offsets: 4096 (Q) + 1024 (K) + 1024 (V) = 6144 total - self.slice_offsets = torch.tensor( - [0, 4096, 5120, 6144], dtype=torch.int32, device=self.device - ) - self.max_slice_size = 4096 - - def generate_sequence_lengths( - self, - batch_size: int, - batch_mode: BatchMode = BatchMode.PREFILL, - min_len: int = 1, - max_len: int = None, - ) -> List[int]: - """Generate sequence lengths for a batch based on mode""" - if batch_mode == BatchMode.DECODE: - return [1] * batch_size - else: - if max_len is None: - max_len = self.max_seq_len - return [random.randint(min_len, max_len) for _ in range(batch_size)] - - def create_lora_weights( - self, lora_name: str, include_missing_k: bool = False - ) -> Tuple[torch.Tensor, torch.Tensor]: - """Create LoRA A and B weights for given configuration""" - rank, out_q, out_k, out_v = self.lora_configs[lora_name] - - if rank == 0: - lora_a = torch.empty( - 0, self.input_dim, dtype=self.dtype, device=self.device - ) - lora_b = torch.empty( - out_q + out_k + out_v, 0, dtype=self.dtype, device=self.device - ) - return lora_a, lora_b - - # Create LoRA A weights (3 slices for QKV) - lora_a = torch.randn( - 3 * rank, self.input_dim, dtype=self.dtype, device=self.device - ) - - if include_missing_k: - lora_a[rank : 2 * rank, :] = 0.0 - - # Create LoRA B weights (stacked Q, K, V) - total_output_dim = out_q + out_k + out_v - lora_b = torch.randn( - total_output_dim, rank, dtype=self.dtype, device=self.device - ) - - if include_missing_k: - lora_b[out_q : out_q + out_k, :] = 0.0 - - return lora_a, lora_b - - def create_batch_info( - self, - lora_names: List[str], - seq_lengths: List[int], - lora_assignments: List[Optional[int]], - batch_mode: BatchMode = BatchMode.PREFILL, - ) -> LoRABatchInfo: - """Create LoRABatchInfo using the same logic as chunked backend""" - lora_ranks = [self.lora_configs[name][0] for name in lora_names] - - def create_mock_batch(): - # Create a minimal mock ForwardBatch for the test - class MockForwardBatch: - def __init__(self, batch_size, seq_lengths, device): - self.batch_size = batch_size - self.extend_seq_lens = torch.tensor( - seq_lengths, dtype=torch.int32, device=device - ) - self.extend_seq_lens_cpu = seq_lengths - self.forward_mode = MockForwardMode() - - class MockForwardMode: - def is_extend(self): - return batch_mode == BatchMode.PREFILL - - def is_decode(self): - return batch_mode == BatchMode.DECODE - - def is_target_verify(self): - return batch_mode == BatchMode.TARGET_VERIFY - - def is_prefill(self): - return self.is_extend() - - return MockForwardBatch(len(seq_lengths), seq_lengths, self.device) - - mock_batch = create_mock_batch() - - # Use the same functions as chunked backend - permutation, weights_reordered = ChunkedSgmvLoRABackend._get_permutation( - lora_assignments, mock_batch - ) - - # Create a minimal backend instance to access _get_segments_info - mock_server_args = type( - "ServerArgs", (object,), {"max_lora_chunk_size": "MOCK_NEVER_USED"} - ) - mock_backend = ChunkedSgmvLoRABackend( - max_loras_per_batch=8, device=self.device, server_args=mock_server_args - ) - weight_indices_list, seg_indptr = mock_backend._get_segments_info( - weights_reordered, - chunk_size=CHUNK_SIZE, - ) - - scalings = [1.0] * len(lora_names) - seg_indptr_tensor = seg_indptr.to(self.device) - weight_indices_tensor = weight_indices_list.to(self.device) - lora_ranks_tensor = ( - torch.tensor(lora_ranks, dtype=torch.int32, device=self.device) - if lora_ranks - else torch.empty(0, dtype=torch.int32, device=self.device) - ) - scalings_tensor = ( - torch.tensor(scalings, dtype=torch.float32, device=self.device) - if scalings - else torch.empty(0, dtype=torch.float32, device=self.device) - ) - permutation_tensor = permutation.to( - self.device, dtype=torch.int32 - ) # Convert to int32 for LoRABatchInfo - seq_lens_tensor = torch.tensor( - seq_lengths, dtype=torch.int32, device=self.device - ) - - return LoRABatchInfo( - use_cuda_graph=False, - bs=len(seq_lengths), - num_segments=len(weight_indices_list), # Number of segments, not sequences! - seg_indptr=seg_indptr_tensor, - weight_indices=weight_indices_tensor, - lora_ranks=lora_ranks_tensor, - scalings=scalings_tensor, - seg_lens=seq_lens_tensor, # Original sequence lengths for reference - max_len=CHUNK_SIZE, - permutation=permutation_tensor, # Token reordering permutation - ) - - def stack_lora_weights( - self, weight_list: List[torch.Tensor], is_lora_a: bool - ) -> torch.Tensor: - """Stack LoRA weights from different adapters into a single tensor""" - if not weight_list: - return torch.empty(0, 0, 0, dtype=self.dtype, device=self.device) - - first_non_empty = next((w for w in weight_list if w.numel() > 0), None) - if first_non_empty is None: - return torch.empty( - len(weight_list), 0, 0, dtype=self.dtype, device=self.device - ) - if is_lora_a: - # LoRA A: (slice_num * rank, input_dim) -> (num_loras, slice_num * max_rank, input_dim) - max_rank = max(w.shape[0] // 3 if w.numel() > 0 else 0 for w in weight_list) - final_shape = (len(weight_list), 3 * max_rank, self.input_dim) - else: - # LoRA B: (output_dim, rank) -> (num_loras, output_dim, max_rank) - max_rank = max(w.shape[1] if w.numel() > 0 else 0 for w in weight_list) - output_dim = first_non_empty.shape[0] - final_shape = (len(weight_list), output_dim, max_rank) - - stacked = torch.zeros(final_shape, dtype=self.dtype, device=self.device) - - for i, weight in enumerate(weight_list): - if weight.numel() > 0: - if is_lora_a: - stacked[i, : weight.shape[0], :] = weight - else: - stacked[i, :, : weight.shape[1]] = weight - - return stacked - - def create_embedding_lora_a_weights(self, lora_ranks: torch.Tensor) -> torch.Tensor: - """Create LoRA A weights for embedding lookup. - - Args: - lora_ranks: Tensor of ranks for each LoRA adapter - - Returns: - Tensor of shape (num_loras, max_rank, vocab_size) - """ - lora_ranks_cpu = lora_ranks.cpu().numpy() - num_loras = len(lora_ranks_cpu) - max_rank = int(lora_ranks_cpu.max()) if num_loras > 0 else 0 - - if max_rank == 0: - return torch.empty( - num_loras, 0, self.vocab_size, dtype=self.dtype, device=self.device - ) - - weights = torch.zeros( - num_loras, max_rank, self.vocab_size, dtype=self.dtype, device=self.device - ) - - for i, rank in enumerate(lora_ranks_cpu): - if rank > 0: - weights[i, :rank, :] = torch.randn( - rank, self.vocab_size, dtype=self.dtype, device=self.device - ) - - return weights - - def create_test_input_ids(self, total_tokens: int) -> torch.Tensor: - """Create random token IDs for embedding test.""" - return torch.randint( - 0, self.vocab_size, (total_tokens,), dtype=torch.int64, device=self.device - ) - - def create_test_batch( - self, - batch_composition: BatchComposition, - batch_size: int, - batch_mode: BatchMode = BatchMode.PREFILL, - include_missing_k: bool = False, - ) -> Tuple[ - torch.Tensor, - List[Tuple[torch.Tensor, torch.Tensor]], - LoRABatchInfo, - List[int], - List[str], - ]: - """Create test batch with specified composition and mode""" - - # Reset kernel cache to avoid cross-test contamination - reset_kernel_cache() - - seq_lengths = self.generate_sequence_lengths( - batch_size, batch_mode, 1, self.max_seq_len - ) - if batch_composition == BatchComposition.UNIFORM: - lora_names = ["lora_A"] - lora_assignments = [lora_names.index("lora_A")] * batch_size - elif batch_composition == BatchComposition.MIXED: - lora_names = ["lora_A", "lora_B", "lora_C", None] - lora_assignments = [(i % len(lora_names)) for i in range(batch_size)] - elif batch_composition == BatchComposition.SKEWED: - lora_names = ["lora_A", "lora_B"] - num_minority = max(1, batch_size // 8) - lora_assignments = [lora_names.index("lora_A")] * num_minority + [ - lora_names.index("lora_B") - ] * (batch_size - num_minority) - random.shuffle(lora_assignments) - elif batch_composition == BatchComposition.NONE: - lora_names = [None] - lora_assignments = [0] * batch_size - else: - raise ValueError(f"Unknown batch composition: {batch_composition}") - - total_seq_len = sum(seq_lengths) - x = torch.randn( - total_seq_len, self.input_dim, dtype=self.dtype, device=self.device - ) - - normalized_lora_names = [ - "_NO_LORA_" if name is None else name for name in lora_names - ] - weights = [] - for lora_name in normalized_lora_names: - weights.append(self.create_lora_weights(lora_name, include_missing_k)) - - batch_info = self.create_batch_info( - normalized_lora_names, seq_lengths, lora_assignments, batch_mode - ) - - return x, weights, batch_info, seq_lengths, lora_assignments - - def run_test_comparison( - self, - x: torch.Tensor, - weights: List[Tuple[torch.Tensor, torch.Tensor]], - batch_info: LoRABatchInfo, - seq_lengths: List[int], - lora_assignments: List[int], - test_name: str, - ): - """Run comparison between chunked and reference implementations""" - if not weights: # Handle case with no LoRA weights - return - - lora_assignments_tensor = torch.tensor( - lora_assignments, dtype=torch.int32, device="cpu" - ) - seq_lengths_tensor = torch.tensor(seq_lengths, dtype=torch.int32, device="cpu") - lora_ranks_tensor = batch_info.lora_ranks.detach().cpu() - scalings_tensor = batch_info.scalings.detach().cpu() - - # Stack LoRA A weights - lora_a_weights = [weight[0] for weight in weights] - stacked_lora_a = self.stack_lora_weights(lora_a_weights, is_lora_a=True) - - # Stack LoRA B weights - lora_b_weights = [weight[1] for weight in weights] - stacked_lora_b = self.stack_lora_weights(lora_b_weights, is_lora_a=False) - - # Test shrink operation - chunked_shrink = chunked_sgmv_lora_shrink_forward( - x, stacked_lora_a, batch_info, num_slices=3 - ) - reference_shrink = reference_sgmv_shrink( - x, - stacked_lora_a, - lora_assignments_tensor, - seq_lengths_tensor, - lora_ranks_tensor, - scalings_tensor, - num_slices=3, - ) - - # Only compare valid portions of shrink output (first rank * num_slices columns per sequence) - self._compare_shrink_outputs( - chunked_shrink, - reference_shrink, - seq_lengths, - lora_assignments, - batch_info, - num_slices=3, - test_name=test_name, - ) - - # Test expand operation - chunked_expand = chunked_sgmv_lora_expand_forward( - reference_shrink, - stacked_lora_b, - batch_info, - self.slice_offsets, - self.max_slice_size, - base_output=None, - ) - reference_expand = reference_sgmv_expand( - reference_shrink, - stacked_lora_b, - lora_assignments_tensor, - seq_lengths_tensor, - lora_ranks_tensor, - self.slice_offsets, - ) - - torch.testing.assert_close( - chunked_expand, - reference_expand, - rtol=self.RTOL, - atol=self.ATOL, - msg=f"Expand operation failed for {test_name}", - ) - - # === Basic Operations Tests === - - def test_shrink_basic(self): - """Test basic shrink operation against PyTorch reference""" - for batch_size in [1, 2, 16, 64]: - with self.subTest(batch_size=batch_size): - x, weights, batch_info, seq_lengths, lora_assignments = ( - self.create_test_batch(BatchComposition.UNIFORM, batch_size) - ) - - lora_assignments_tensor = torch.tensor( - lora_assignments, dtype=torch.int32, device="cpu" - ) - seq_lengths_tensor = torch.tensor( - seq_lengths, dtype=torch.int32, device="cpu" - ) - lora_ranks_tensor = batch_info.lora_ranks.detach().cpu() - scalings_tensor = batch_info.scalings.detach().cpu() - - lora_a_weights = [weight[0] for weight in weights] - stacked_lora_a = self.stack_lora_weights(lora_a_weights, is_lora_a=True) - - chunked_shrink = chunked_sgmv_lora_shrink_forward( - x, stacked_lora_a, batch_info, num_slices=3 - ) - reference_shrink = reference_sgmv_shrink( - x, - stacked_lora_a, - lora_assignments_tensor, - seq_lengths_tensor, - lora_ranks_tensor, - scalings_tensor, - num_slices=3, - ) - - torch.testing.assert_close( - chunked_shrink, reference_shrink, rtol=self.RTOL, atol=self.ATOL - ) - - # Test chunked embedding LoRA A forward - # Create embedding-specific LoRA A weights with shape (num_loras, rank, vocab_size) - embedding_lora_a = self.create_embedding_lora_a_weights( - batch_info.lora_ranks - ) - - # Create input_ids (token indices) instead of hidden states - total_tokens = x.shape[0] - input_ids = self.create_test_input_ids(total_tokens) - - chunked_shrink_embeddings = chunked_embedding_lora_a_forward( - input_ids, embedding_lora_a, batch_info, self.vocab_size - ) - - reference_shrink_embeddings = reference_embedding_lora_a_shrink( - input_ids, - embedding_lora_a, - lora_assignments_tensor, - seq_lengths_tensor, - lora_ranks_tensor, - scalings_tensor, - self.vocab_size, - ) - torch.testing.assert_close( - chunked_shrink_embeddings, - reference_shrink_embeddings, - rtol=self.RTOL, - atol=self.ATOL, - msg=f"Shrink test embedding loRA A operation failed for batch_size={batch_size}", - ) - - def test_expand_basic(self): - """Test basic expand operation against PyTorch reference""" - for batch_size in [1, 2, 16, 64]: - with self.subTest(batch_size=batch_size): - x, weights, batch_info, seq_lengths, lora_assignments = ( - self.create_test_batch(BatchComposition.UNIFORM, batch_size) - ) - - lora_assignments_tensor = torch.tensor( - lora_assignments, dtype=torch.int32, device="cpu" - ) - seq_lengths_tensor = torch.tensor( - seq_lengths, dtype=torch.int32, device="cpu" - ) - lora_ranks_tensor = batch_info.lora_ranks.detach().cpu() - scalings_tensor = batch_info.scalings.detach().cpu() - - lora_a_weights = [weight[0] for weight in weights] - stacked_lora_a = self.stack_lora_weights(lora_a_weights, is_lora_a=True) - - intermediate = reference_sgmv_shrink( - x, - stacked_lora_a, - lora_assignments_tensor, - seq_lengths_tensor, - lora_ranks_tensor, - scalings_tensor, - num_slices=3, - ) - - lora_b_weights = [weight[1] for weight in weights] - stacked_lora_b = self.stack_lora_weights( - lora_b_weights, is_lora_a=False - ) - - chunked_expand = chunked_sgmv_lora_expand_forward( - intermediate, - stacked_lora_b, - batch_info, - self.slice_offsets, - self.max_slice_size, - base_output=None, - ) - reference_expand = reference_sgmv_expand( - intermediate, - stacked_lora_b, - lora_assignments_tensor, - seq_lengths_tensor, - lora_ranks_tensor, - self.slice_offsets, - ) - - torch.testing.assert_close( - chunked_expand, reference_expand, rtol=self.RTOL, atol=self.ATOL - ) - - # === QKV Operations Test === - - def test_qkv_missing_projections(self): - """Test QKV operations with missing k_proj (Qwen3 scenario)""" - for batch_size in [1, 2, 16, 64]: - with self.subTest(batch_size=batch_size): - x, weights, batch_info, seq_lengths, lora_assignments = ( - self.create_test_batch( - BatchComposition.MIXED, batch_size, include_missing_k=True - ) - ) - self.run_test_comparison( - x, - weights, - batch_info, - seq_lengths, - lora_assignments, - f"QKV missing k_proj batch_size={batch_size}", - ) - - def test_4_slice_gdn_qkvz(self): - """Test 4-slice shrink+expand operations (GDN in_proj_qkvz).""" - num_slices = 4 - # GDN-style: 4 slices with different sizes [2048, 2048, 4096, 4096] - slice_offsets = torch.tensor( - [0, 2048, 4096, 8192, 12288], dtype=torch.int32, device=self.device - ) - total_out = 12288 - max_slice_size = 4096 - - for batch_size in [1, 2, 16]: - with self.subTest(batch_size=batch_size): - # Build batch (reuse the 3-slice helper just for x, batch_info, etc.) - x, _, batch_info, seq_lengths, lora_assignments = ( - self.create_test_batch(BatchComposition.MIXED, batch_size) - ) - - # Build 4-slice LoRA weights and stack them manually - lora_names = [n for n in self.lora_configs if n != "_NO_LORA_"] - max_rank = max(self.lora_configs[n][0] for n in lora_names) - stacked_a = torch.zeros( - len(lora_names), - num_slices * max_rank, - self.input_dim, - dtype=self.dtype, - device=self.device, - ) - stacked_b = torch.zeros( - len(lora_names), - total_out, - max_rank, - dtype=self.dtype, - device=self.device, - ) - for i, name in enumerate(lora_names): - rank = self.lora_configs[name][0] - if rank > 0: - stacked_a[i, : num_slices * rank, :] = torch.randn( - num_slices * rank, - self.input_dim, - dtype=self.dtype, - device=self.device, - ) - stacked_b[i, :, :rank] = torch.randn( - total_out, rank, dtype=self.dtype, device=self.device - ) - - lora_assignments_tensor = torch.tensor( - lora_assignments, dtype=torch.int32, device="cpu" - ) - seq_lengths_tensor = torch.tensor( - seq_lengths, dtype=torch.int32, device="cpu" - ) - lora_ranks_tensor = batch_info.lora_ranks.detach().cpu() - scalings_tensor = batch_info.scalings.detach().cpu() - - # Shrink - chunked_shrink = chunked_sgmv_lora_shrink_forward( - x, stacked_a, batch_info, num_slices=num_slices - ) - reference_shrink = reference_sgmv_shrink( - x, - stacked_a, - lora_assignments_tensor, - seq_lengths_tensor, - lora_ranks_tensor, - scalings_tensor, - num_slices=num_slices, - ) - self._compare_shrink_outputs( - chunked_shrink, - reference_shrink, - seq_lengths, - lora_assignments, - batch_info, - num_slices=num_slices, - test_name=f"4-slice shrink bs={batch_size}", - ) - - # Expand - chunked_expand = chunked_sgmv_lora_expand_forward( - reference_shrink, - stacked_b, - batch_info, - slice_offsets, - max_slice_size, - base_output=None, - ) - reference_expand = reference_sgmv_expand( - reference_shrink, - stacked_b, - lora_assignments_tensor, - seq_lengths_tensor, - lora_ranks_tensor, - slice_offsets, - ) - torch.testing.assert_close( - chunked_expand, - reference_expand, - rtol=self.RTOL, - atol=self.ATOL, - msg=f"4-slice expand failed bs={batch_size}", - ) - - # === Batch Composition Tests === - - def test_uniform_lora_batch(self): - """All sequences use same LoRA, random sequence lengths""" - for batch_size in [1, 2, 16, 64]: - with self.subTest(batch_size=batch_size): - x, weights, batch_info, seq_lengths, lora_assignments = ( - self.create_test_batch(BatchComposition.UNIFORM, batch_size) - ) - self.run_test_comparison( - x, - weights, - batch_info, - seq_lengths, - lora_assignments, - f"uniform batch_size={batch_size}", - ) - - def test_evenly_mixed_lora_batch(self): - """Sequences evenly distributed across LoRAs, random lengths""" - for batch_size in [1, 2, 16, 64]: - with self.subTest(batch_size=batch_size): - x, weights, batch_info, seq_lengths, lora_assignments = ( - self.create_test_batch(BatchComposition.MIXED, batch_size) - ) - self.run_test_comparison( - x, - weights, - batch_info, - seq_lengths, - lora_assignments, - f"mixed batch_size={batch_size}", - ) - - def test_highly_skewed_lora_batch(self): - """Highly uneven LoRA distribution, random lengths""" - for batch_size in [1, 2, 16, 64]: - with self.subTest(batch_size=batch_size): - x, weights, batch_info, seq_lengths, lora_assignments = ( - self.create_test_batch(BatchComposition.SKEWED, batch_size) - ) - self.run_test_comparison( - x, - weights, - batch_info, - seq_lengths, - lora_assignments, - f"skewed batch_size={batch_size}", - ) - - # === Decode Mode Tests === - - def test_decode_uniform_lora_batch(self): - """Decode mode: All sequences use same LoRA, all length 1""" - for batch_size in [1, 2, 16, 64]: - with self.subTest(batch_size=batch_size): - x, weights, batch_info, seq_lengths, lora_assignments = ( - self.create_test_batch( - BatchComposition.UNIFORM, batch_size, BatchMode.DECODE - ) - ) - self.run_test_comparison( - x, - weights, - batch_info, - seq_lengths, - lora_assignments, - f"decode uniform batch_size={batch_size}", - ) - - def test_decode_mixed_lora_batch(self): - """Decode mode: Sequences distributed across LoRAs, all length 1""" - for batch_size in [1, 2, 16, 64]: - with self.subTest(batch_size=batch_size): - x, weights, batch_info, seq_lengths, lora_assignments = ( - self.create_test_batch( - BatchComposition.MIXED, batch_size, BatchMode.DECODE - ) - ) - self.run_test_comparison( - x, - weights, - batch_info, - seq_lengths, - lora_assignments, - f"decode mixed batch_size={batch_size}", - ) - - def test_decode_skewed_lora_batch(self): - """Decode mode: Highly uneven LoRA distribution, all length 1""" - for batch_size in [1, 2, 16, 64]: - with self.subTest(batch_size=batch_size): - x, weights, batch_info, seq_lengths, lora_assignments = ( - self.create_test_batch( - BatchComposition.SKEWED, batch_size, BatchMode.DECODE - ) - ) - self.run_test_comparison( - x, - weights, - batch_info, - seq_lengths, - lora_assignments, - f"decode skewed batch_size={batch_size}", - ) - - def _make_cuda_graph_batch_info(self, bs: int, num_loras: int) -> LoRABatchInfo: - return LoRABatchInfo( - use_cuda_graph=True, - bs=bs, - num_segments=None, - max_len=CHUNK_SIZE, - seg_lens=None, - seg_indptr=torch.zeros(bs + 1, dtype=torch.int32, device=self.device), - weight_indices=torch.zeros(bs, dtype=torch.int32, device=self.device), - lora_ranks=torch.zeros(num_loras, dtype=torch.int32, device=self.device), - scalings=torch.ones(num_loras, dtype=torch.float, device=self.device), - permutation=torch.arange(bs, dtype=torch.int32, device=self.device), - ) - - def _set_cuda_graph_segment_state( - self, - batch_info: LoRABatchInfo, - lora_ranks: List[int], - weight_indices: List[int], - seg_indptr: List[int], - ): - num_segments = len(weight_indices) - total_tokens = seg_indptr[-1] - device = batch_info.weight_indices.device - - batch_info.lora_ranks.zero_() - batch_info.lora_ranks[: len(lora_ranks)].copy_( - torch.tensor(lora_ranks, dtype=torch.int32, device=device) - ) - batch_info.weight_indices.zero_() - batch_info.weight_indices[:num_segments].copy_( - torch.tensor(weight_indices, dtype=torch.int32, device=device) - ) - batch_info.seg_indptr.fill_(total_tokens) - batch_info.seg_indptr[: num_segments + 1].copy_( - torch.tensor(seg_indptr, dtype=torch.int32, device=device) - ) - batch_info.num_segments = num_segments - - def _set_cuda_graph_capture_state(self, batch_info: LoRABatchInfo, bs: int): - self._set_cuda_graph_segment_state( - batch_info=batch_info, - lora_ranks=[0] * batch_info.lora_ranks.shape[0], - weight_indices=[0], - seg_indptr=[0, bs], - ) - - def _set_cuda_graph_replay_state( - self, batch_info: LoRABatchInfo, max_rank: int, bs: int - ): - self._set_cuda_graph_segment_state( - batch_info=batch_info, - lora_ranks=[max_rank] * batch_info.lora_ranks.shape[0], - weight_indices=[1, 2, 3, 4], - seg_indptr=[0, 2, 4, 6, bs], - ) - - @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required") - def test_cuda_graph_shrink_replay_with_more_segments_than_capture(self): - """CUDA graph replay must honor updated shrink segment metadata.""" - reset_kernel_cache() - - bs = 8 - num_loras = 5 - max_rank = 8 - input_dim = 64 - num_slices = 1 - - x = torch.randn(bs, input_dim, dtype=self.dtype, device=self.device) - weights = torch.randn( - num_loras, max_rank, input_dim, dtype=self.dtype, device=self.device - ) - batch_info = self._make_cuda_graph_batch_info(bs, num_loras) - - self._set_cuda_graph_replay_state(batch_info, max_rank, bs) - expected = chunked_sgmv_lora_shrink_forward( - x, weights, batch_info, num_slices=num_slices - ).clone() - torch.cuda.synchronize() - - self._set_cuda_graph_capture_state(batch_info, bs) - warmup_stream = torch.cuda.Stream() - warmup_stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(warmup_stream): - for _ in range(3): - chunked_sgmv_lora_shrink_forward( - x, weights, batch_info, num_slices=num_slices - ) - torch.cuda.current_stream().wait_stream(warmup_stream) - torch.cuda.synchronize() - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - captured_output = chunked_sgmv_lora_shrink_forward( - x, weights, batch_info, num_slices=num_slices - ) - - self._set_cuda_graph_replay_state(batch_info, max_rank, bs) - captured_output.zero_() - graph.replay() - torch.cuda.synchronize() - - torch.testing.assert_close( - captured_output[:, :max_rank], - expected[:, :max_rank], - rtol=self.RTOL, - atol=self.ATOL, - ) - - @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required") - def test_cuda_graph_expand_replay_with_more_segments_than_capture(self): - """CUDA graph replay must honor updated expand segment metadata.""" - reset_kernel_cache() - - bs = 8 - num_loras = 5 - max_rank = 8 - output_dim = 32 - slice_offsets = torch.tensor( - [0, output_dim], dtype=torch.int32, device=self.device - ) - - x = torch.randn(bs, max_rank, dtype=self.dtype, device=self.device) - weights = torch.randn( - num_loras, output_dim, max_rank, dtype=self.dtype, device=self.device - ) - base_output = torch.randn(bs, output_dim, dtype=self.dtype, device=self.device) - graph_base_output = base_output.clone() - batch_info = self._make_cuda_graph_batch_info(bs, num_loras) - - self._set_cuda_graph_replay_state(batch_info, max_rank, bs) - expected = chunked_sgmv_lora_expand_forward( - x, - weights, - batch_info, - slice_offsets, - output_dim, - base_output=base_output.clone(), - ).clone() - torch.cuda.synchronize() - - self._set_cuda_graph_capture_state(batch_info, bs) - warmup_stream = torch.cuda.Stream() - warmup_stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(warmup_stream): - for _ in range(3): - graph_base_output.copy_(base_output) - chunked_sgmv_lora_expand_forward( - x, - weights, - batch_info, - slice_offsets, - output_dim, - base_output=graph_base_output, - ) - torch.cuda.current_stream().wait_stream(warmup_stream) - torch.cuda.synchronize() - - graph_base_output.copy_(base_output) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - captured_output = chunked_sgmv_lora_expand_forward( - x, - weights, - batch_info, - slice_offsets, - output_dim, - base_output=graph_base_output, - ) - - self._set_cuda_graph_replay_state(batch_info, max_rank, bs) - graph_base_output.copy_(base_output) - graph.replay() - torch.cuda.synchronize() - - torch.testing.assert_close( - captured_output, - expected, - rtol=self.RTOL, - atol=self.ATOL, - ) - - @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required") - def test_prepare_lora_batch_cuda_graph_zero_length_tail(self): - """prepare_lora_batch must neutralize stale CUDA graph tail segments.""" - - class MockForwardBatch: - def __init__(self, batch_size): - self.batch_size = batch_size - self.forward_mode = ForwardMode.DECODE - - mock_server_args = type( - "ServerArgs", (object,), {"max_lora_chunk_size": CHUNK_SIZE} - ) - backend = ChunkedSgmvLoRABackend( - max_loras_per_batch=5, device=self.device, server_args=mock_server_args - ) - backend.init_cuda_graph_batch_info(max_bs_in_cuda_graph=8, num_tokens_per_req=1) - - lora_ranks = [8] * 5 - scalings = [1.0] * 5 - backend.prepare_lora_batch( - forward_batch=MockForwardBatch(8), - weight_indices=[0, 1, 2, 3, 4, 0, 1, 2], - lora_ranks=lora_ranks, - scalings=scalings, - use_cuda_graph=True, - ) - backend.prepare_lora_batch( - forward_batch=MockForwardBatch(2), - weight_indices=[0, 0], - lora_ranks=lora_ranks, - scalings=scalings, - use_cuda_graph=True, - ) - torch.cuda.synchronize() - - batch_info = backend.batch_info - self.assertEqual(batch_info.num_segments, 1) - torch.testing.assert_close( - batch_info.weight_indices.cpu(), - torch.tensor([0] * 8, dtype=torch.int32), - ) - torch.testing.assert_close( - batch_info.seg_indptr.cpu(), - torch.tensor([0, 2, 2, 2, 2, 2, 2, 2, 2], dtype=torch.int32), - ) - - @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required") - def test_kv_b_cuda_graph_replay_with_more_segments_than_capture(self): - """Absorbed MLA kv_b LoRA kernels must replay dynamic segment metadata.""" - bs = 8 - num_loras = 5 - max_rank = 8 - num_heads = 2 - qk_nope_head_dim = 16 - v_head_dim = 16 - kv_lora_rank = 32 - full_k_per_head = qk_nope_head_dim + v_head_dim - - q_nope = torch.randn( - bs, num_heads, qk_nope_head_dim, dtype=self.dtype, device=self.device - ) - attn_output = torch.randn( - bs, num_heads, kv_lora_rank, dtype=self.dtype, device=self.device - ) - a_buf = torch.randn( - num_loras, max_rank, kv_lora_rank, dtype=self.dtype, device=self.device - ) - b_buf = torch.randn( - num_loras, - num_heads * full_k_per_head, - max_rank, - dtype=self.dtype, - device=self.device, - ) - base_q = torch.randn( - bs, num_heads, kv_lora_rank, dtype=self.dtype, device=self.device - ) - base_v = torch.randn( - bs, num_heads, v_head_dim, dtype=self.dtype, device=self.device - ) - graph_base_q = base_q.clone() - graph_base_v = base_v.clone() - batch_info = self._make_cuda_graph_batch_info(bs, num_loras) - - def run_kv_b(base_q_out, base_v_out): - q_lora_a = step_a_q_fwd(q_nope, b_buf, batch_info, full_k_per_head) - q_out = step_b_q_fwd(q_lora_a, a_buf, batch_info, base_q_out) - v_lora_a = step_a_v_fwd(attn_output, a_buf, batch_info) - v_out = step_b_v_fwd( - v_lora_a, - b_buf, - batch_info, - base_v_out, - qk_nope_head_dim, - v_head_dim, - ) - return q_out, v_out - - self._set_cuda_graph_replay_state(batch_info, max_rank, bs) - expected_q, expected_v = run_kv_b(base_q.clone(), base_v.clone()) - expected_q = expected_q.clone() - expected_v = expected_v.clone() - torch.cuda.synchronize() - - self._set_cuda_graph_capture_state(batch_info, bs) - warmup_stream = torch.cuda.Stream() - warmup_stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(warmup_stream): - for _ in range(3): - graph_base_q.copy_(base_q) - graph_base_v.copy_(base_v) - run_kv_b(graph_base_q, graph_base_v) - torch.cuda.current_stream().wait_stream(warmup_stream) - torch.cuda.synchronize() - - graph_base_q.copy_(base_q) - graph_base_v.copy_(base_v) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - captured_q, captured_v = run_kv_b(graph_base_q, graph_base_v) - - self._set_cuda_graph_replay_state(batch_info, max_rank, bs) - graph_base_q.copy_(base_q) - graph_base_v.copy_(base_v) - graph.replay() - torch.cuda.synchronize() - - torch.testing.assert_close( - captured_q, - expected_q, - rtol=self.RTOL, - atol=self.ATOL, - ) - torch.testing.assert_close( - captured_v, - expected_v, - rtol=self.RTOL, - atol=self.ATOL, - ) - - -class TestLmHeadPruningConsistency(unittest.TestCase): - """Verify get_lm_head_pruned_lens (LoRA) stays consistent with - LogitsProcessor._get_pruned_states (logits_processor). - - If this test fails, it likely means one side was changed without - updating the other. See cross-references in both functions. - """ - - def _make_mock_forward_batch( - self, - forward_mode, - extend_seq_lens_cpu, - return_logprob=False, - logprob_start_lens_cpu=None, - ): - class MockForwardBatch: - pass - - batch = MockForwardBatch() - batch.forward_mode = forward_mode - batch.batch_size = len(extend_seq_lens_cpu) - batch.return_logprob = return_logprob - batch.extend_seq_lens_cpu = extend_seq_lens_cpu - batch.extend_logprob_start_lens_cpu = logprob_start_lens_cpu - return batch - - def _count_pruned_states_tokens( - self, - forward_mode, - extend_seq_lens_cpu, - return_logprob=False, - logprob_start_lens_cpu=None, - ): - """Call _get_pruned_states and return the number of output tokens.""" - total_tokens = sum(extend_seq_lens_cpu) - hidden_states = torch.zeros(total_tokens, 4) - - logits_meta = LogitsMetadata( - forward_mode=forward_mode, - extend_return_logprob=return_logprob, - extend_seq_lens=torch.tensor(extend_seq_lens_cpu, dtype=torch.int64), - extend_seq_lens_cpu=extend_seq_lens_cpu, - extend_logprob_start_lens_cpu=logprob_start_lens_cpu, - ) - - # _get_pruned_states does not use self, so pass None - result = LogitsProcessor._get_pruned_states( - None, hidden_states, None, None, logits_meta - ) - pruned_states = result[0] - return pruned_states.shape[0] - - def _assert_consistency( - self, - forward_mode, - extend_seq_lens_cpu, - return_logprob=False, - logprob_start_lens_cpu=None, - ): - mock_batch = self._make_mock_forward_batch( - forward_mode, - extend_seq_lens_cpu, - return_logprob, - logprob_start_lens_cpu, - ) - pruned_lens = get_lm_head_pruned_lens(mock_batch) - - actual_count = self._count_pruned_states_tokens( - forward_mode, - extend_seq_lens_cpu, - return_logprob, - logprob_start_lens_cpu, - ) - - if pruned_lens is None: - expected_count = sum(extend_seq_lens_cpu) - else: - expected_count = sum(pruned_lens) - - self.assertEqual( - expected_count, - actual_count, - f"get_lm_head_pruned_lens expects {expected_count} tokens, " - f"but _get_pruned_states produces {actual_count}. " - f"These functions must stay in sync — see their cross-reference comments.", - ) - - def test_extend_no_logprob(self): - self._assert_consistency(ForwardMode.EXTEND, [4, 5, 6]) - - def test_extend_with_logprob(self): - self._assert_consistency( - ForwardMode.EXTEND, - [4, 5, 6], - return_logprob=True, - logprob_start_lens_cpu=[0, 5, 3], - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/lora/test_embedding_lora_support.py b/test/registered/lora/test_embedding_lora_support.py deleted file mode 100644 index 32f397c3a..000000000 --- a/test/registered/lora/test_embedding_lora_support.py +++ /dev/null @@ -1,232 +0,0 @@ -# Copyright 2023-2024 SGLang Team -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== -""" -Unit tests for LoRA support in embedding models. - -Validates that EmbeddingReqInput correctly handles LoRA fields through -normalization, batching, and request splitting. -""" - -import multiprocessing as mp -import unittest - -import numpy as np -import torch - -from sglang.srt.entrypoints.openai.protocol import EmbeddingRequest -from sglang.srt.managers.io_struct import EmbeddingReqInput, TokenizedEmbeddingReqInput -from sglang.srt.sampling.sampling_params import SamplingParams -from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.runners import SRTRunner -from sglang.test.test_utils import DEFAULT_PORT_FOR_SRT_TEST_RUNNER, CustomTestCase - -# Test configuration (same model/LoRA as test_lora_hf_sgl_logprob_diff.py) -MODEL_PATH = "meta-llama/Llama-2-7b-hf" -LORA_PATH = "yushengsu/sglang_lora_logprob_diff_without_tuning" -LORA_BACKEND = "triton" -SIMILARITY_THRESHOLD = 0.9999 - -register_cuda_ci( - est_time=150, - suite="nightly-1-gpu", -) - - -class TestEmbeddingLoraSupport(unittest.TestCase): - """Test LoRA support in embedding request structures.""" - - def test_engine_encode_validates_enable_lora(self): - """Test Engine.encode() validates enable_lora before processing lora_path.""" - # Use a simple non-gated model for this validation test - with SRTRunner( - MODEL_PATH, - torch_dtype=torch.float16, - model_type="embedding", - port=DEFAULT_PORT_FOR_SRT_TEST_RUNNER, - ) as runner: - # Should raise ValueError because enable_lora was not set for the server - with self.assertRaises(ValueError) as context: - runner.engine.encode(prompt="Test", lora_path="fake-adapter") - - error_msg = str(context.exception) - self.assertIn("not enabled", error_msg.lower()) - self.assertIn("--enable-lora", error_msg) - self.assertIn("fake-adapter", error_msg) - - def test_embedding_lora_fields(self): - """Test LoRA fields exist and work correctly across all embedding structures.""" - # EmbeddingReqInput: fields exist, normalization expands single to batch, indexing works - req = EmbeddingReqInput( - text=["Hello", "World"], lora_path="my-adapter", lora_id=["id1", "id2"] - ) - self.assertIsNotNone(req.lora_path) - req.normalize_batch_and_arguments() - self.assertEqual(req.lora_path, ["my-adapter", "my-adapter"]) - self.assertEqual(req[0].lora_path, "my-adapter") - self.assertEqual(req[1].lora_id, "id2") - - # EmbeddingReqInput: mismatched list length raises error - req = EmbeddingReqInput(text=["Hello", "World", "Test"], lora_path=["adapter1"]) - with self.assertRaises(ValueError): - req.normalize_batch_and_arguments() - - # TokenizedEmbeddingReqInput and EmbeddingRequest have lora fields - tokenized = TokenizedEmbeddingReqInput( - input_text="Hello", - input_ids=[1, 2, 3], - image_inputs={}, - token_type_ids=[], - sampling_params=SamplingParams(), - lora_id="my-lora-id", - ) - self.assertEqual(tokenized.lora_id, "my-lora-id") - self.assertEqual( - EmbeddingRequest( - input="Hello", model="test", lora_path="adapter" - ).lora_path, - "adapter", - ) - - -class TestEmbeddingLoraHFComparison(CustomTestCase): - """Compare HF+LoRA vs SGLang+LoRA embedding outputs.""" - - @classmethod - def get_hf_embedding_with_lora(cls, model_path, lora_path, texts, torch_dtype): - """Get embeddings from HuggingFace model with LoRA adapter.""" - from peft import PeftModel - from transformers import AutoModelForCausalLM, AutoTokenizer - - # Load base model as CausalLM to match adapter's expected structure - base_model = AutoModelForCausalLM.from_pretrained( - model_path, - torch_dtype=torch_dtype, - trust_remote_code=True, - ).cuda() - - # Load LoRA adapter - model = PeftModel.from_pretrained(base_model, lora_path) - model.eval() - - tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) - if tokenizer.pad_token is None: - tokenizer.pad_token = tokenizer.eos_token - - with torch.no_grad(): - inputs = tokenizer( - texts, padding=True, truncation=True, return_tensors="pt" - ).to("cuda") - - # Access the inner model (CausalLM wraps the base model) - outputs = model.model(**inputs, output_hidden_states=True) - hidden_states = outputs.hidden_states[-1] - - # Last token pooling with L2 normalization (matching SGLang) - attention_mask = inputs["attention_mask"] - last_token_indices = attention_mask.sum(dim=1) - 1 - batch_size = hidden_states.shape[0] - embeddings = hidden_states[ - torch.arange(batch_size, device="cuda"), last_token_indices - ] - embeddings = embeddings / embeddings.norm(dim=1, keepdim=True) - - # Cleanup - del model, base_model - torch.cuda.empty_cache() - - return embeddings.cpu().numpy() - - @classmethod - def get_sglang_embedding_with_lora(cls, model_path, lora_path, texts, torch_dtype): - """Get embeddings from SGLang with LoRA adapter.""" - with SRTRunner( - model_path, - torch_dtype=torch_dtype, - model_type="embedding", - lora_paths=[lora_path], - lora_backend=LORA_BACKEND, - port=DEFAULT_PORT_FOR_SRT_TEST_RUNNER, - trust_remote_code=True, - mem_fraction_static=0.88, - ) as runner: - # Call engine.encode directly with lora_path - response = runner.engine.encode(prompt=texts, lora_path=lora_path) - if isinstance(response, list): - embeddings = [r["embedding"] for r in response] - else: - embeddings = [response["embedding"]] - - return np.array(embeddings) - - @staticmethod - def cosine_similarity(a, b): - """Compute cosine similarity between vectors.""" - return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)) - - def test_embedding_lora_hf_sglang_similarity(self): - """Test that HF+LoRA and SGLang+LoRA produce similar embeddings.""" - test_texts = [ - "Hello world", - "This is a test sentence for embedding comparison", - ] - - print(f"\nModel: {MODEL_PATH}") - print(f"LoRA: {LORA_PATH}") - - # Get SGLang embeddings first (before HF loads model into GPU) - # This order matches test_lora_hf_sgl_logprob_diff.py and avoids OOM - print("\nGetting SGLang embeddings...") - sglang_embeddings = self.get_sglang_embedding_with_lora( - MODEL_PATH, LORA_PATH, test_texts, torch.float16 - ) - - # Clear GPU memory - torch.cuda.empty_cache() - - # Get HF embeddings - print("Getting HF embeddings...") - hf_embeddings = self.get_hf_embedding_with_lora( - MODEL_PATH, LORA_PATH, test_texts, torch.float16 - ) - - # Compare embeddings - print("\nHF vs SGLang LoRA Embedding Comparison:") - similarities = [] - for i, (hf_emb, sgl_emb) in enumerate(zip(hf_embeddings, sglang_embeddings)): - sim = self.cosine_similarity(hf_emb, sgl_emb) - similarities.append(sim) - print(f" Text {i}: cosine similarity = {sim:.6f}") - self.assertGreater( - sim, - SIMILARITY_THRESHOLD, - f"Text {i} similarity {sim:.6f} below threshold {SIMILARITY_THRESHOLD}", - ) - - avg_similarity = np.mean(similarities) - print(f" Average similarity: {avg_similarity:.6f}") - print(f" Threshold: {SIMILARITY_THRESHOLD}") - - self.assertGreater( - avg_similarity, - SIMILARITY_THRESHOLD, - f"Average similarity {avg_similarity:.4f} below threshold {SIMILARITY_THRESHOLD}", - ) - - -if __name__ == "__main__": - try: - mp.set_start_method("spawn") - except RuntimeError: - pass - unittest.main() diff --git a/test/registered/lora/test_lora_openai_api.py b/test/registered/lora/test_lora_openai_api.py index 53f3fd397..f71acec43 100644 --- a/test/registered/lora/test_lora_openai_api.py +++ b/test/registered/lora/test_lora_openai_api.py @@ -10,13 +10,8 @@ from unittest.mock import MagicMock from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase from sglang.srt.server_args import ServerArgs -from sglang.test.ci.ci_register import ( - register_amd_ci, - register_cpu_ci, - register_cuda_ci, -) +from sglang.test.ci.ci_register import register_amd_ci, register_cpu_ci -register_cuda_ci(est_time=30, suite="nightly-1-gpu", nightly=True) register_amd_ci(est_time=30, suite="nightly-amd-1-gpu", nightly=True) register_cpu_ci(est_time=8, suite="base-c-test-cpu") diff --git a/test/registered/lora/test_lora_radix_cache.py b/test/registered/lora/test_lora_radix_cache.py deleted file mode 100644 index d59572ed8..000000000 --- a/test/registered/lora/test_lora_radix_cache.py +++ /dev/null @@ -1,82 +0,0 @@ -# Copyright 2023-2024 SGLang Team -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== - -import multiprocessing as mp -import unittest - -import torch - -from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci -from sglang.test.lora_utils import CI_MULTI_LORA_MODELS, run_lora_test_one_by_one -from sglang.test.test_utils import CustomTestCase - -register_cuda_ci(est_time=200, suite="nightly-1-gpu", nightly=True) -register_amd_ci(est_time=200, suite="nightly-amd-1-gpu", nightly=True) - -PROMPTS = [ - "AI is a field of computer science focused on", - """ - ### Instruction: - Tell me about llamas and alpacas - ### Response: - Llamas are large, long-necked animals with a woolly coat. They have two toes on each foot instead of three like other camelids. - ### Question: - What do you know about llamas? - ### Answer: - """, -] - - -class TestLoRARadixCache(CustomTestCase): - - def test_lora_radix_cache(self): - # Here we need a model case with multiple adaptors for testing correctness of radix cache - model_case = CI_MULTI_LORA_MODELS[0] - - torch_dtype = torch.float16 - max_new_tokens = 32 - batch_prompts = ( - PROMPTS - if not model_case.skip_long_prompt - else [p for p in PROMPTS if len(p) < 1000] - ) - - # Test lora with radix cache - run_lora_test_one_by_one( - batch_prompts, - model_case, - torch_dtype, - max_new_tokens=max_new_tokens, - disable_radix_cache=False, - test_tag="lora-with-radix-cache", - ) - - # Test lora without radix cache - run_lora_test_one_by_one( - batch_prompts, - model_case, - torch_dtype, - max_new_tokens=max_new_tokens, - disable_radix_cache=True, - test_tag="lora-without-radix-cache", - ) - - -if __name__ == "__main__": - try: - mp.set_start_method("spawn") - except RuntimeError: - pass - - unittest.main(warnings="ignore") diff --git a/test/registered/lora/test_lora_tied_lm_head.py b/test/registered/lora/test_lora_tied_lm_head.py deleted file mode 100644 index 4070a5321..000000000 --- a/test/registered/lora/test_lora_tied_lm_head.py +++ /dev/null @@ -1,224 +0,0 @@ -# Copyright 2023-2025 SGLang Team -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== - -""" -Test LoRA on models with tied lm_head (tie_word_embeddings=True). - -When tie_word_embeddings=True, lm_head shares the same weight tensor as -embed_tokens. PyTorch's named_modules() deduplicates by object identity, -so lm_head won't appear as a separate module. This test validates that -SGLang correctly handles this case by untying lm_head before LoRA wrapping. - -The test: -1. Programmatically creates a LoRA adapter with lm_head in target_modules - using PEFT on a model with tie_word_embeddings=True (Qwen/Qwen2.5-0.5B). -2. Compares logprobs between HuggingFace+PEFT and SGLang to ensure numerical - consistency. This implicitly verifies no NaN values are produced and that - LoRA is actually being applied (since HF+PEFT is the trusted reference). -""" - -import multiprocessing as mp -import os -import shutil -import tempfile -import unittest - -import torch - -try: - from peft import LoraConfig, get_peft_model -except ImportError: - import subprocess - - subprocess.check_call(["pip", "install", "peft", "--no-deps"]) - from peft import LoraConfig, get_peft_model - -from transformers import AutoModelForCausalLM - -from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.runners import HFRunner, SRTRunner -from sglang.test.test_utils import DEFAULT_PORT_FOR_SRT_TEST_RUNNER, CustomTestCase - -register_cuda_ci(est_time=120, suite="nightly-1-gpu", nightly=True) - -# Use a small model with tie_word_embeddings=True -BASE_MODEL = "Qwen/Qwen2.5-0.5B" - -TEST_PROMPTS = [ - "AI is a field of computer science focused on", - "The capital of France is", -] - -MAX_NEW_TOKENS = 16 -LOGPROB_THRESHOLD = 2e-1 - - -def create_lora_adapter_with_lm_head(base_model_name: str, output_dir: str): - """ - Programmatically create a LoRA adapter that targets lm_head, - using a model with tie_word_embeddings=True. - - The adapter uses randomly initialized LoRA weights (no training). - This is sufficient to test that: - - SGLang can load the adapter without errors - - lm_head LoRA is applied (output differs from base model) - - Logprobs match between HF and SGLang - """ - model = AutoModelForCausalLM.from_pretrained( - base_model_name, - torch_dtype=torch.float16, - device_map="cpu", - ) - - # Verify the model actually has tied embeddings - assert ( - model.config.tie_word_embeddings - ), f"Expected tie_word_embeddings=True for {base_model_name}" - - # Only target lm_head to isolate the test to the tied-embedding scenario. - lora_config = LoraConfig( - r=8, - lora_alpha=16, - target_modules=["lm_head"], - lora_dropout=0, - bias="none", - task_type="CAUSAL_LM", - ) - - peft_model = get_peft_model(model, lora_config) - - # PEFT initializes lora_B to zeros by default, which makes the adapter - # produce identical output to the base model. Initialize lora_B with - # non-zero random weights so the adapter has a visible effect. - with torch.no_grad(): - for name, param in peft_model.named_parameters(): - if "lora_B" in name: - torch.nn.init.normal_(param, mean=0.0, std=0.02) - - peft_model.save_pretrained(output_dir) - - # Verify the saved adapter contains lm_head keys - from safetensors import safe_open - - safetensors_path = os.path.join(output_dir, "adapter_model.safetensors") - f = safe_open(safetensors_path, framework="pt") - lm_head_keys = [k for k in f.keys() if "lm_head" in k] - assert ( - len(lm_head_keys) > 0 - ), f"Expected lm_head LoRA weights in adapter, got keys: {sorted(f.keys())}" - - print(f"Created LoRA adapter at {output_dir}") - print(f" lm_head keys: {lm_head_keys}") - - # Clean up the model to free memory - del peft_model, model - torch.cuda.empty_cache() - - -class TestLoRATiedLMHead(CustomTestCase): - """ - Test that LoRA works correctly on models with tied lm_head. - """ - - _adapter_dir = None - - @classmethod - def setUpClass(cls): - """Create a temporary LoRA adapter with lm_head targeting.""" - super().setUpClass() - cls._adapter_dir = tempfile.mkdtemp(prefix="sglang_test_lora_tied_lm_head_") - create_lora_adapter_with_lm_head(BASE_MODEL, cls._adapter_dir) - - @classmethod - def tearDownClass(cls): - """Clean up the temporary adapter directory.""" - if cls._adapter_dir and os.path.exists(cls._adapter_dir): - shutil.rmtree(cls._adapter_dir) - super().tearDownClass() - - def test_tied_lm_head_lora_hf_sgl_logprob_match(self): - """ - Compare logprobs between HuggingFace+PEFT and SGLang+LoRA - for a tied lm_head adapter, ensuring numerical consistency. - """ - prompts = TEST_PROMPTS[:2] - - # Run SGLang with LoRA - with SRTRunner( - BASE_MODEL, - torch_dtype=torch.float16, - model_type="generation", - lora_paths=[self._adapter_dir], - max_loras_per_batch=1, - lora_backend="triton", - lora_target_modules=["lm_head"], - disable_cuda_graph=True, - disable_radix_cache=True, - mem_fraction_static=0.80, - port=DEFAULT_PORT_FOR_SRT_TEST_RUNNER, - ) as srt_runner: - srt_outputs = srt_runner.forward( - prompts, - max_new_tokens=MAX_NEW_TOKENS, - lora_paths=[self._adapter_dir] * len(prompts), - ) - - torch.cuda.empty_cache() - - # Run HuggingFace with LoRA (via PEFT) - with HFRunner( - BASE_MODEL, - torch_dtype=torch.float16, - model_type="generation", - ) as hf_runner: - hf_outputs = hf_runner.forward( - prompts, - max_new_tokens=MAX_NEW_TOKENS, - lora_paths=[self._adapter_dir] * len(prompts), - ) - - # Compare prefill logprobs - for i in range(len(prompts)): - srt_logprobs = torch.tensor(srt_outputs.top_input_logprobs[i]) - hf_logprobs = torch.tensor(hf_outputs.top_input_logprobs[i]) - max_diff = torch.max(torch.abs(srt_logprobs - hf_logprobs)).item() - print(f"Prompt {i} prefill logprob max_diff (SGLang vs HF): {max_diff:.6e}") - self.assertLess( - max_diff, - LOGPROB_THRESHOLD, - f"Prompt {i}: prefill logprob diff {max_diff:.6e} " - f"exceeds threshold {LOGPROB_THRESHOLD:.0e}", - ) - - # Compare decode logprobs - for i in range(len(prompts)): - srt_logprobs = torch.tensor(srt_outputs.top_output_logprobs[i]) - hf_logprobs = torch.tensor(hf_outputs.top_output_logprobs[i]) - max_diff = torch.max(torch.abs(srt_logprobs - hf_logprobs)).item() - print(f"Prompt {i} decode logprob max_diff (SGLang vs HF): {max_diff:.6e}") - self.assertLess( - max_diff, - LOGPROB_THRESHOLD, - f"Prompt {i}: decode logprob diff {max_diff:.6e} " - f"exceeds threshold {LOGPROB_THRESHOLD:.0e}", - ) - - -if __name__ == "__main__": - try: - mp.set_start_method("spawn") - except RuntimeError: - pass - - unittest.main(warnings="ignore") diff --git a/test/registered/models_e2e/test_qwen3_next_models_fp4.py b/test/registered/models_e2e/test_qwen3_next_models_fp4.py deleted file mode 100644 index 2c5d6bdf2..000000000 --- a/test/registered/models_e2e/test_qwen3_next_models_fp4.py +++ /dev/null @@ -1,34 +0,0 @@ -import unittest - -from sglang.srt.utils import get_device_sm -from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.kits.eval_accuracy_kit import GSM8KMixin -from sglang.test.server_fixtures.default_fixture import DefaultServerBase - -register_cuda_ci(est_time=500, suite="nightly-4-gpu-b200", nightly=True) - -QWEN3_NEXT_MODEL_FP4 = "nvidia/Qwen3-Next-80B-A3B-Instruct-NVFP4" - - -@unittest.skipIf( - get_device_sm() < 100, "Test requires CUDA SM 100 or higher (Blackwell)" -) -class TestQwen3NextFp4(GSM8KMixin, DefaultServerBase): - model = QWEN3_NEXT_MODEL_FP4 - gsm8k_accuracy_thres = 0.93 - other_args = [ - "--tp-size", - "4", - "--chunked-prefill-size", - "2048", - "--quantization", - "modelopt_fp4", - "--mamba-scheduler-strategy", - "extra_buffer", - "--mamba-track-interval", - "128", - ] - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/perf/test_dpsk_v3_fp4_4gpu_perf.py b/test/registered/perf/test_dpsk_v3_fp4_4gpu_perf.py deleted file mode 100644 index f9aec8612..000000000 --- a/test/registered/perf/test_dpsk_v3_fp4_4gpu_perf.py +++ /dev/null @@ -1,77 +0,0 @@ -import unittest - -from sglang.test.accuracy_test_runner import AccuracyTestParams -from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.performance_test_runner import PerformanceTestParams -from sglang.test.run_combined_tests import run_combined_tests -from sglang.test.test_utils import ModelLaunchSettings - -# Runs on B200 via nightly-4-gpu-b200 suite -register_cuda_ci(est_time=2000, suite="nightly-4-gpu-b200", nightly=True) - -FULL_DEEPSEEK_V3_FP4_MODEL_PATH = "nvidia/DeepSeek-V3-0324-FP4" - - -class TestDeepseekR1FP4Unified(unittest.TestCase): - """Unified test class for DeepSeek-V3-0324-FP4 performance and accuracy. - - Two variants: - - basic: Standard TP=4 - - mtp: TP=4 + EAGLE speculative decoding - - Each variant runs BOTH: - - Performance test (using NightlyBenchmarkRunner) - - Accuracy test (using run_eval with mgsm_en) - """ - - def test_deepseek_r1_fp4_all_variants(self): - """Run performance and accuracy for all DeepSeek-R1-0528-NVFP4-v2 variants.""" - # Define base arguments shared by most variants - base_args = [ - "--tp=4", - "--trust-remote-code", - "--model-loader-extra-config", - '{"enable_multithread_load": true}', - ] - mtp_args = [ - "--speculative-algorithm=EAGLE", - "--speculative-num-steps=3", - "--speculative-eagle-topk=1", - "--speculative-num-draft-tokens=4", - "--mem-frac=0.7", - ] - - variants = [ - # Variant: "basic" - Standard TP=4 - ModelLaunchSettings( - FULL_DEEPSEEK_V3_FP4_MODEL_PATH, - tp_size=4, - extra_args=base_args, - variant="TP4", - ), - # Variant: "mtp" - TP=4 + EAGLE speculative decoding - ModelLaunchSettings( - FULL_DEEPSEEK_V3_FP4_MODEL_PATH, - tp_size=4, - extra_args=base_args + mtp_args, - variant="TP4+MTP", - ), - ] - - run_combined_tests( - models=variants, - test_name="DeepSeek-V3-0324-FP4 Unified", - accuracy_params=AccuracyTestParams( - dataset="gsm8k", - baseline_accuracy=0.935, - num_examples=200, - api="completion", - ), - performance_params=PerformanceTestParams( - result_dir="performance_results_deepseek_v3_fp4", - ), - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/perf/test_gpt_oss_4gpu_perf.py b/test/registered/perf/test_gpt_oss_4gpu_perf.py deleted file mode 100644 index d4a048fd6..000000000 --- a/test/registered/perf/test_gpt_oss_4gpu_perf.py +++ /dev/null @@ -1,60 +0,0 @@ -import unittest - -from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.nightly_utils import NightlyBenchmarkRunner -from sglang.test.test_utils import DEFAULT_URL_FOR_TEST - -register_cuda_ci(est_time=600, suite="nightly-4-gpu-b200", nightly=True) - -RESULT_DIR = "performance_results_gpt_oss_4gpu" - - -class TestNightlyGptOss4GpuPerformance(unittest.TestCase): - @classmethod - def setUpClass(cls): - cls.models = [ - ( - "openai/gpt-oss-120b", - [ - "--tp", - "4", - "--cuda-graph-max-bs-decode", - "200", - "--mem-fraction-static", - "0.93", - ], - ), - ] - cls.base_url = DEFAULT_URL_FOR_TEST - cls.batch_sizes = [1, 1, 8, 16, 64] - cls.input_lens = (4096,) - cls.output_lens = (512,) - cls.runner = NightlyBenchmarkRunner(RESULT_DIR, cls.__name__, cls.base_url) - cls.runner.setup_result_directory() - - def test_bench_one_batch(self): - all_model_succeed = True - - for model_path, other_args in self.models: - with self.subTest(model=model_path): - results, success, _ = self.runner.run_benchmark_for_model( - model_path=model_path, - batch_sizes=self.batch_sizes, - input_lens=self.input_lens, - output_lens=self.output_lens, - other_args=other_args, - ) - - if not success: - all_model_succeed = False - - self.runner.add_report(results) - - self.runner.write_final_report() - - if not all_model_succeed: - raise AssertionError("Some models failed the perf tests.") - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/prefill_only/test_fa_skip_kv_cache_piecewise_nan.py b/test/registered/prefill_only/test_fa_skip_kv_cache_piecewise_nan.py deleted file mode 100644 index e51027b5b..000000000 --- a/test/registered/prefill_only/test_fa_skip_kv_cache_piecewise_nan.py +++ /dev/null @@ -1,162 +0,0 @@ -# Copyright 2023-2024 SGLang Team -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== -"""Regression test for the fa_skip_kv_cache embedding fast path under piecewise -CUDA graph. - -PR #21971 added an embedding fast path (`fa_skip_kv_cache`) that serves attention -with `flash_attn_varlen_func` on raw K/V. Under a piecewise CUDA graph the model -forward runs at a padded token-bucket size, so `q` has more rows than -`cu_seqlens_q` covers. `flash_attn_varlen_func` requires -`q.shape[0] == cu_seqlens_q[-1]`; when that is violated the boundary query block -corrupts the **last real token's** output. Because embedding models use LAST-token -pooling, that corrupted row IS the returned embedding -> ~40% of *short* inputs -came back fully NaN (long inputs, which fill the bucket, were unaffected). - -This test feeds a spread of short inputs through `fa3 + piecewise + fa_skip_kv_cache` -and asserts no embedding contains NaN, and that the embeddings match the -non-piecewise path. -""" - -import os -import unittest - -import torch - -from sglang import Engine -from sglang.srt.utils import get_device_sm -from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.test_utils import CustomTestCase - -# Route to the nightly 1-GPU suite, which runs on the H100 pool (1-gpu-h100, SM90). -# FA3 + the piecewise embedding path this regression covers only runs on -# Ampere/Ada/Hopper (SM 80-90), so the test must land on the H100 pool to actually -# execute (on the RTX 5090 SM120/Blackwell 1-gpu-small runner it would skip 100%). -register_cuda_ci(est_time=600, suite="nightly-1-gpu", nightly=True) - -# Lowest/highest CUDA SM that supports the FA3 + piecewise embedding path. FA3 is -# unavailable on Blackwell (sm100 B200 / sm120 consumer e.g. RTX 5090); the gate is -# applied at RUNTIME (see setUp) so the SM is read after CUDA is initialized on the -# actual runner, never at import/collection time. -_FA3_SM_MIN, _FA3_SM_MAX = 80, 90 - -# Overridable so the test can run against a locally-mounted model in dev. -MODEL_PATH = os.environ.get("SGLANG_TEST_EMB_MODEL", "Qwen/Qwen3-Embedding-0.6B") - -_WORDS = [ - "the", - "quick", - "brown", - "fox", - "jumps", - "over", - "lazy", - "dog", - "embedding", - "vector", - "token", - "sample", -] - - -def _short_prompts(): - """A spread of short inputs (~1..150 tokens). - - The bug only triggers when a prefill is PADDED up to a piecewise bucket, i.e. - for token counts that are not exactly a capture size. Using many lengths - guarantees several land just below a bucket boundary (80/96/112/128 ...). - """ - return [" ".join(_WORDS[i % len(_WORDS)] for i in range(n)) for n in range(1, 150)] - - -def _embed(prompts, **engine_kwargs): - # fa_skip_kv_cache is enabled by: is_embedding + chunked_prefill_size == -1 - # + disable_radix_cache (+ a non-MLA model + the FA3 backend). - engine = Engine( - model_path=MODEL_PATH, - is_embedding=True, - attention_backend="fa3", - chunked_prefill_size=-1, - disable_radix_cache=True, - **engine_kwargs, - ) - try: - # Encode one request per forward (batch size 1). The bug corrupts the - # last real token, which sits exactly at the real/pad boundary; when many - # requests are batched into one forward only the tail request hits the - # boundary, which hides the per-request failure rate. - embs = [] - for prompt in prompts: - out = engine.encode(prompt) - emb = out["embedding"] if isinstance(out, dict) else out[0]["embedding"] - embs.append(torch.tensor(emb, dtype=torch.float32)) - return embs - finally: - engine.shutdown() - - -# Enables the piecewise CUDA graph for prefill the way production does. After the -# cuda-graph refactor (#23906) the piecewise config lives in cuda_graph_config; the -# convenience kwargs below fold into cuda_graph_config[prefill]: -# - cuda_graph_backend_prefill="tc_piecewise" -> prefill.backend (also the default) -# - cuda_graph_max_bs_prefill=32768 -> prefill.max_bs (for tc_piecewise -# prefill, max_bs/bs carries the captured TOKEN count -- the old -# piecewise_cuda_graph_max_tokens) -# - cuda_graph_tc_compiler="inductor" -> prefill.tc_compiler -_PIECEWISE_KWARGS = dict( - cuda_graph_backend_prefill="tc_piecewise", - cuda_graph_max_bs_prefill=32768, - cuda_graph_tc_compiler="inductor", -) - - -class TestFaSkipKvCachePiecewiseNoNaN(CustomTestCase): - def setUp(self): - # Gate at runtime: read the SM after CUDA is initialized on the runner. If - # the hardware can't run FA3 (e.g. SM120 RTX 5090 / SM100 B200), skip -- - # a skip is NOT a CI failure, it just records the test as inapplicable here. - sm = get_device_sm() - if not (_FA3_SM_MIN <= sm <= _FA3_SM_MAX): - self.skipTest( - f"fa3 + piecewise embedding repro requires CUDA SM " - f"{_FA3_SM_MIN}-{_FA3_SM_MAX} (Ampere/Ada/Hopper); got SM {sm}" - ) - - def test_no_nan_with_piecewise(self): - prompts = _short_prompts() - embs = _embed(prompts, **_PIECEWISE_KWARGS) - nan_idx = [i for i, e in enumerate(embs) if torch.isnan(e).any()] - self.assertEqual( - nan_idx, - [], - f"{len(nan_idx)}/{len(embs)} short-input embeddings contain NaN under " - f"fa_skip_kv_cache + piecewise CUDA graph (e.g. prompt indices {nan_idx[:10]})", - ) - - def test_matches_non_piecewise(self): - prompts = _short_prompts() - with_pcg = _embed(prompts, **_PIECEWISE_KWARGS) - without_pcg = _embed(prompts, disable_prefill_cuda_graph=True) - for i, (a, b) in enumerate(zip(with_pcg, without_pcg)): - self.assertFalse( - torch.isnan(a).any(), - f"prompt {i}: NaN embedding with piecewise CUDA graph", - ) - cos = torch.nn.functional.cosine_similarity(a, b, dim=0).item() - self.assertGreater( - cos, 0.99, f"prompt {i}: cosine {cos:.4f} < 0.99 vs non-piecewise" - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/quant/test_kimi_k25_nvfp4_eagle.py b/test/registered/quant/test_kimi_k25_nvfp4_eagle.py deleted file mode 100644 index c5073abb8..000000000 --- a/test/registered/quant/test_kimi_k25_nvfp4_eagle.py +++ /dev/null @@ -1,68 +0,0 @@ -import unittest - -from sglang.test.accuracy_test_runner import AccuracyTestParams -from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.performance_test_runner import PerformanceTestParams -from sglang.test.run_combined_tests import run_combined_tests -from sglang.test.test_utils import ModelLaunchSettings - -# Kimi-K2.5 NVFP4 + EAGLE3 (MLA draft) speculative decoding on 4x B200, tp=4. -register_cuda_ci(est_time=3000, suite="nightly-4-gpu-b200", nightly=True) - -MODEL_PATH = "nvidia/Kimi-K2.5-NVFP4" -DRAFT_MODEL_PATH = "lightseekorg/kimi-k2.5-eagle3-mla" - -EXTRA_ARGS = [ - "--trust-remote-code", - "--attention-backend=tokenspeed_mla", - "--moe-runner-backend=flashinfer_trtllm", - "--quantization=modelopt_fp4", - "--kv-cache-dtype=fp8_e4m3", - "--mem-fraction-static=0.85", - "--max-running-requests=16", - "--speculative-algorithm=EAGLE3", - f"--speculative-draft-model-path={DRAFT_MODEL_PATH}", - "--speculative-num-steps=3", - "--speculative-eagle-topk=1", - "--speculative-num-draft-tokens=4", - "--speculative-draft-model-quantization=unquant", -] - - -class TestKimiK25Nvfp4Eagle(unittest.TestCase): - """Kimi-K2.5 NVFP4 with EAGLE3 speculative decoding on 4x B200 (tp=4). - - Runs both an accuracy test (gsm8k) and a performance test (bs=1/8/16), - and gates the speculative-decoding accept length. - """ - - def test_kimi_k25_nvfp4_eagle(self): - variants = [ - ModelLaunchSettings( - MODEL_PATH, - tp_size=4, - extra_args=EXTRA_ARGS, - variant="TP4+EAGLE3", - ), - ] - - run_combined_tests( - models=variants, - test_name="Kimi-K2.5-NVFP4 EAGLE3", - # Thresholds from a measured tp=4 run: gsm8k 0.945, perf accept ~3.0-3.3. - accuracy_params=AccuracyTestParams( - dataset="gsm8k", - baseline_accuracy=0.92, - num_examples=200, - api="completion", - ), - performance_params=PerformanceTestParams( - batch_sizes=[1, 8, 16], - spec_accept_length_threshold=2.8, - result_dir="performance_results_kimi_k25_nvfp4_eagle", - ), - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/quant/test_kimi_k26_nvfp4_dflash.py b/test/registered/quant/test_kimi_k26_nvfp4_dflash.py deleted file mode 100644 index 47722386f..000000000 --- a/test/registered/quant/test_kimi_k26_nvfp4_dflash.py +++ /dev/null @@ -1,70 +0,0 @@ -import unittest - -from sglang.test.accuracy_test_runner import AccuracyTestParams -from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.performance_test_runner import PerformanceTestParams -from sglang.test.run_combined_tests import run_combined_tests -from sglang.test.test_utils import ModelLaunchSettings - -# Kimi-K2.6 NVFP4 (pure-MLA target, fp8 KV) + DFlash speculative decoding on 8x B200, tp=8. -register_cuda_ci(est_time=3600, suite="nightly-8-gpu-b200", nightly=True) - -MODEL_PATH = "nvidia/Kimi-K2.6-NVFP4" -DRAFT_MODEL_PATH = "nvidia/Kimi-K2.6-DFlash" - -# trtllm_mla verify only; cuteDSL fold verify depends on the flashinfer version. -EXTRA_ARGS = [ - "--trust-remote-code", - "--quantization=modelopt_fp4", - "--moe-runner-backend=flashinfer_trtllm", - "--fp4-gemm-backend=flashinfer_cutlass", - "--attention-backend=trtllm_mla", - "--kv-cache-dtype=fp8_e4m3", - "--mem-fraction-static=0.85", - "--max-running-requests=16", - "--speculative-algorithm=DFLASH", - f"--speculative-draft-model-path={DRAFT_MODEL_PATH}", - "--speculative-num-draft-tokens=8", - "--speculative-draft-attention-backend=fa4", - "--speculative-draft-model-quantization=unquant", - "--speculative-draft-window-size=4096", -] - - -class TestKimiK26Nvfp4Dflash(unittest.TestCase): - """Kimi-K2.6 NVFP4 (pure-MLA, fp8 KV) with DFlash speculative decoding on 8x B200 (tp=8). - - Runs both an accuracy test (gsm8k) and a performance test (bs=1/8/16), and gates the - speculative-decoding accept length. Guards the pure-MLA fp8-KV DFlash path. - """ - - def test_kimi_k26_nvfp4_dflash(self): - variants = [ - ModelLaunchSettings( - MODEL_PATH, - tp_size=8, - extra_args=EXTRA_ARGS, - variant="TP8+DFLASH", - ), - ] - - run_combined_tests( - models=variants, - test_name="Kimi-K2.6-NVFP4 DFlash", - # Thresholds from a measured tp=8 run: gsm8k 0.936 (full set), accept length ~2.66. - accuracy_params=AccuracyTestParams( - dataset="gsm8k", - baseline_accuracy=0.92, - num_examples=200, - api="completion", - ), - performance_params=PerformanceTestParams( - batch_sizes=[1, 8, 16], - spec_accept_length_threshold=2.0, - result_dir="performance_results_kimi_k26_nvfp4_dflash", - ), - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/radix_cache/test_cpp_radix_cache.py b/test/registered/radix_cache/test_cpp_radix_cache.py deleted file mode 100644 index c8fad0637..000000000 --- a/test/registered/radix_cache/test_cpp_radix_cache.py +++ /dev/null @@ -1,41 +0,0 @@ -import unittest - -from sglang.srt.environ import envs -from sglang.srt.utils import kill_process_tree -from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.kits.eval_accuracy_kit import MMLUMixin -from sglang.test.test_utils import ( - DEFAULT_MODEL_NAME_FOR_TEST, - DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - DEFAULT_URL_FOR_TEST, - CustomTestCase, - popen_launch_server, -) - -# Note: AMD registration removed - test_cpp_radix_cache fails on AMD due to C++ radix tree issues -register_cuda_ci(est_time=60, suite="nightly-1-gpu", nightly=True) - - -class TestCppRadixCache(CustomTestCase, MMLUMixin): - mmlu_score_threshold = 0.65 - mmlu_num_examples = 64 - mmlu_num_threads = 32 - - @classmethod - def setUpClass(cls): - envs.SGLANG_EXPERIMENTAL_CPP_RADIX_TREE.set(True) - cls.model = DEFAULT_MODEL_NAME_FOR_TEST - cls.base_url = DEFAULT_URL_FOR_TEST - cls.process = popen_launch_server( - cls.model, - cls.base_url, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - ) - - @classmethod - def tearDownClass(cls): - kill_process_tree(cls.process.pid) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/spec/eagle/test_eagle_infer_beta_dp_attention_large.py b/test/registered/spec/eagle/test_eagle_infer_beta_dp_attention_large.py deleted file mode 100644 index 6f7511244..000000000 --- a/test/registered/spec/eagle/test_eagle_infer_beta_dp_attention_large.py +++ /dev/null @@ -1,100 +0,0 @@ -import unittest -from types import SimpleNamespace - -import requests - -from sglang.srt.utils import kill_process_tree -from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.run_eval import run_eval -from sglang.test.test_utils import ( - DEFAULT_DEEPSEEK_NVFP4_MODEL_FOR_TEST, - DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - DEFAULT_URL_FOR_TEST, - CustomTestCase, - is_in_ci, - popen_launch_server, - write_github_step_summary, -) - -# 16 GPU test (4 TP x 4 DP), runs on 2x 8-GPU B200 nodes -register_cuda_ci(est_time=600, suite="nightly-8-gpu-b200", nightly=True) - - -def test_gsm8k(base_url: str, model: str): - requests.get(base_url + "/flush_cache") - - args = SimpleNamespace( - base_url=base_url, - model=model, - eval_name="gsm8k", - api="completion", - max_tokens=512, - num_examples=200, - num_threads=128, - ) - metrics = run_eval(args) - server_info = requests.get(base_url + "/server_info").json() - avg_spec_accept_length = server_info["internal_states"][0]["avg_spec_accept_length"] - - print(f"{metrics=}") - print(f"{avg_spec_accept_length=}") - return metrics, avg_spec_accept_length - - -class TestEagleDPAttnServerLarge(CustomTestCase): - # FIXME: move this large mode test into nightly tests - @classmethod - def setUpClass(cls): - cls.model = DEFAULT_DEEPSEEK_NVFP4_MODEL_FOR_TEST - cls.base_url = DEFAULT_URL_FOR_TEST - other_args = [ - "--tp-size", - "4", - "--dp-size", - "4", - "--enable-dp-attention", - "--attention-backend", - "trtllm_mla", - "--moe-runner-backend", - "flashinfer_trtllm", - "--quantization", - "modelopt_fp4", - "--speculative-algorithm", - "EAGLE", - "--speculative-num-steps", - "3", - "--speculative-eagle-topk", - "1", - "--speculative-num-draft-tokens", - "4", - "--kv-cache-dtype", - "fp8_e4m3", - "--model-loader-extra-config", - '{"enable_multithread_load": true,"num_threads": 64}', - ] - cls.process = popen_launch_server( - cls.model, - cls.base_url, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=other_args, - ) - - @classmethod - def tearDownClass(cls): - kill_process_tree(cls.process.pid) - - def test_a_gsm8k(self): - metrics, avg_spec_accept_length = test_gsm8k(self.base_url, self.model) - - self.assertGreater(metrics["score"], 0.94) - self.assertGreater(avg_spec_accept_length, 2.7) - if is_in_ci(): - write_github_step_summary( - f"### test_gsm8k (deepseek-v3-fp4 mtp)\n" - f'{metrics["score"]=:.3f}\n' - f"{avg_spec_accept_length=:.2f}\n" - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/utils/test_model_file_verifier.py b/test/registered/utils/test_model_file_verifier.py index 7cb87cc24..b83763ef3 100644 --- a/test/registered/utils/test_model_file_verifier.py +++ b/test/registered/utils/test_model_file_verifier.py @@ -20,7 +20,7 @@ from sglang.srt.utils.model_file_verifier import ( generate_checksums, verify, ) -from sglang.test.ci.ci_register import register_cpu_ci, register_cuda_ci +from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import ( DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, DEFAULT_URL_FOR_TEST, @@ -28,7 +28,6 @@ from sglang.test.test_utils import ( ) # Note: AMD registration removed - test_model_file_verifier fails on AMD -register_cuda_ci(est_time=120, suite="nightly-1-gpu", nightly=True) register_cpu_ci(est_time=540, suite="base-c-test-cpu") MODEL_NAME = "Qwen/Qwen3-0.6B" diff --git a/test/registered/utils/test_request_logger.py b/test/registered/utils/test_request_logger.py deleted file mode 100644 index ae6c7e352..000000000 --- a/test/registered/utils/test_request_logger.py +++ /dev/null @@ -1,309 +0,0 @@ -import io -import json -import os -import tempfile -import time -import unittest -from pathlib import Path - -import requests - -from sglang.srt.constants import HEALTH_CHECK_RID_PREFIX -from sglang.srt.utils import kill_process_tree -from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci -from sglang.test.test_utils import ( - DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - DEFAULT_URL_FOR_TEST, - CustomTestCase, - popen_launch_server, -) - -register_cuda_ci(est_time=120, suite="nightly-1-gpu", nightly=True) -register_amd_ci(est_time=120, suite="nightly-amd-1-gpu", nightly=True) - -TEST_ROUTING_KEY = "test-routing-key-12345" -TEST_CUSTOM_HEADER_NAME = "X-Test-Header" -TEST_CUSTOM_HEADER_VALUE = "test-header-value-67890" -TEST_MODEL_NAME = "Qwen/Qwen3-0.6B" - - -class BaseTestRequestLogger: - log_requests_format = None - env_vars: dict[str, str] = {} # Env vars to set before server launch - request_headers: dict[str, str] = {"X-SMG-Routing-Key": TEST_ROUTING_KEY} - - @classmethod - def setUpClass(cls): - cls._temp_dir_obj = tempfile.TemporaryDirectory() - cls.temp_dir = cls._temp_dir_obj.name - cls.stdout = io.StringIO() - cls.stderr = io.StringIO() - other_args = [ - "--log-requests", - "--log-requests-level", - "2", - "--log-requests-format", - cls.log_requests_format, - "--skip-server-warmup", - "--log-requests-target", - "stdout", - cls.temp_dir, - ] - # Set env vars and save old values for restoration - cls._old_env_vars = {} - for key, value in cls.env_vars.items(): - cls._old_env_vars[key] = os.environ.get(key) - os.environ[key] = value - - cls.process = popen_launch_server( - TEST_MODEL_NAME, - DEFAULT_URL_FOR_TEST, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=other_args, - return_stdout_stderr=(cls.stdout, cls.stderr), - ) - - @classmethod - def tearDownClass(cls): - kill_process_tree(cls.process.pid) - cls.stdout.close() - cls.stderr.close() - cls._temp_dir_obj.cleanup() - # Restore env vars - for key, old_value in cls._old_env_vars.items(): - if old_value is None: - os.environ.pop(key, None) - else: - os.environ[key] = old_value - - def _verify_logs(self, content: str, source_name: str): - raise NotImplementedError - - def _verify_openai_logs(self, content: str, source_name: str): - raise NotImplementedError - - def _wait_until_verified( - self, - verify_fn, - get_content_fn, - source_name: str, - timeout: float = 10.0, - interval: float = 0.1, - ): - deadline = time.time() + timeout - last_error = None - - while time.time() < deadline: - content = get_content_fn() - try: - verify_fn(content, source_name) - return - except AssertionError as err: - last_error = err - time.sleep(interval) - - if last_error is not None: - raise last_error - - def test_logging(self): - response = requests.post( - DEFAULT_URL_FOR_TEST + "/generate", - json={ - "text": "Hello", - "sampling_params": {"max_new_tokens": 8, "temperature": 0}, - }, - headers=self.request_headers, - timeout=30, - ) - self.assertEqual(response.status_code, 200) - self._wait_until_verified( - self._verify_logs, - lambda: self.stdout.getvalue() + self.stderr.getvalue(), - "stdout", - ) - self._wait_until_verified( - self._verify_logs, - lambda: "".join(f.read_text() for f in Path(self.temp_dir).glob("*.log")), - "log files", - ) - - log_files = list(Path(self.temp_dir).glob("*.log")) - self.assertGreater(len(log_files), 0, "No log files found in temp directory") - - def test_openai_chat_logging(self): - response = requests.post( - DEFAULT_URL_FOR_TEST + "/v1/chat/completions", - json={ - "model": TEST_MODEL_NAME, - "messages": [{"role": "user", "content": "hello request logger"}], - "max_tokens": 8, - "temperature": 0, - }, - headers=self.request_headers, - timeout=30, - ) - self.assertEqual(response.status_code, 200) - self._wait_until_verified( - self._verify_openai_logs, - lambda: self.stdout.getvalue() + self.stderr.getvalue(), - "stdout", - ) - self._wait_until_verified( - self._verify_openai_logs, - lambda: "".join(f.read_text() for f in Path(self.temp_dir).glob("*.log")), - "log files", - ) - - log_files = list(Path(self.temp_dir).glob("*.log")) - self.assertGreater(len(log_files), 0, "No log files found in temp directory") - - -class TestRequestLoggerText(BaseTestRequestLogger, CustomTestCase): - log_requests_format = "text" - - def _verify_logs(self, content: str, source_name: str): - self.assertIn("Receive:", content, f"'Receive:' not found in {source_name}") - self.assertIn("Finish:", content, f"'Finish:' not found in {source_name}") - self.assertIn( - TEST_ROUTING_KEY, content, f"Routing key not found in {source_name}" - ) - self.assertIn( - "x-smg-routing-key", content, f"Header name not found in {source_name}" - ) - - def _verify_openai_logs(self, content: str, source_name: str): - self.assertIn( - "Receive OpenAI:", content, f"OpenAI receive log not found in {source_name}" - ) - self.assertIn("'messages':", content, f"Messages not found in {source_name}") - self.assertIn( - "hello request logger", - content, - f"OpenAI user prompt not found in {source_name}", - ) - - -class TestRequestLoggerJson(BaseTestRequestLogger, CustomTestCase): - log_requests_format = "json" - - def _verify_logs(self, content: str, source_name: str): - received_found = False - finished_found = False - for line in content.splitlines(): - idx = line.find("{") - if idx == -1: - continue - try: - data = json.loads(line[idx:]) - except json.JSONDecodeError: - continue - - rid = data.get("rid", "") - if rid.startswith(HEALTH_CHECK_RID_PREFIX): - continue - - if data.get("event") == "request.received": - self.assertIn("rid", data) - self.assertIn("obj", data) - self.assertEqual( - data.get("headers", {}).get("x-smg-routing-key"), TEST_ROUTING_KEY - ) - received_found = True - elif data.get("event") == "request.finished": - self.assertIn("rid", data) - self.assertIn("obj", data) - self.assertIn("out", data) - self.assertEqual( - data.get("headers", {}).get("x-smg-routing-key"), TEST_ROUTING_KEY - ) - finished_found = True - - self.assertTrue( - received_found, f"request.received event not found in {source_name}" - ) - self.assertTrue( - finished_found, f"request.finished event not found in {source_name}" - ) - - def _verify_openai_logs(self, content: str, source_name: str): - openai_received_found = False - for line in content.splitlines(): - idx = line.find("{") - if idx == -1: - continue - try: - data = json.loads(line[idx:]) - except json.JSONDecodeError: - continue - if data.get("event") != "request.received.openai": - continue - - obj = data.get("obj", {}) - self.assertEqual(obj.get("model"), TEST_MODEL_NAME) - self.assertIsInstance(obj.get("messages"), list) - self.assertGreater(len(obj.get("messages")), 0) - self.assertEqual(obj["messages"][0].get("content"), "hello request logger") - self.assertEqual( - data.get("headers", {}).get("x-smg-routing-key"), TEST_ROUTING_KEY - ) - openai_received_found = True - break - - self.assertTrue( - openai_received_found, - f"request.received.openai event not found in {source_name}", - ) - - -class TestCustomHeaderViaEnvVar(BaseTestRequestLogger, CustomTestCase): - """Test that custom headers can be added via SGLANG_LOG_REQUEST_HEADERS env var.""" - - log_requests_format = "text" - env_vars = {"SGLANG_LOG_REQUEST_HEADERS": TEST_CUSTOM_HEADER_NAME} - request_headers = { - "X-SMG-Routing-Key": TEST_ROUTING_KEY, - TEST_CUSTOM_HEADER_NAME: TEST_CUSTOM_HEADER_VALUE, - } - - def _verify_logs(self, content: str, source_name: str): - # Verify custom header is logged - self.assertIn( - TEST_CUSTOM_HEADER_NAME.lower(), - content, - f"Custom header name not found in {source_name}", - ) - self.assertIn( - TEST_CUSTOM_HEADER_VALUE, - content, - f"Custom header value not found in {source_name}", - ) - # Verify default header is still logged (env var appends, not replaces) - self.assertIn( - "x-smg-routing-key", - content, - f"Default header should still be in whitelist in {source_name}", - ) - self.assertIn( - TEST_ROUTING_KEY, - content, - f"Default header value not found in {source_name}", - ) - - def _verify_openai_logs(self, content: str, source_name: str): - self.assertIn( - "Receive OpenAI:", content, f"OpenAI receive log not found in {source_name}" - ) - self.assertIn( - TEST_CUSTOM_HEADER_NAME.lower(), - content, - f"Custom header name not found in {source_name}", - ) - self.assertIn( - TEST_CUSTOM_HEADER_VALUE, - content, - f"Custom header value not found in {source_name}", - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/utils/test_scheduler_status_logger.py b/test/registered/utils/test_scheduler_status_logger.py deleted file mode 100644 index d579a77e8..000000000 --- a/test/registered/utils/test_scheduler_status_logger.py +++ /dev/null @@ -1,79 +0,0 @@ -import json -import os -import shutil -import tempfile -import time -import unittest -from pathlib import Path - -import requests - -from sglang.srt.utils import kill_process_tree -from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci -from sglang.test.test_utils import ( - DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - DEFAULT_URL_FOR_TEST, - CustomTestCase, - popen_launch_server, -) - -register_cuda_ci(est_time=120, suite="nightly-1-gpu", nightly=True) -register_amd_ci(est_time=120, suite="nightly-amd-1-gpu", nightly=True) - - -class TestSchedulerStatusLogger(CustomTestCase): - @classmethod - def setUpClass(cls): - cls.temp_dir = tempfile.mkdtemp() - cls.addClassCleanup(shutil.rmtree, cls.temp_dir) - env = os.environ.copy() - env["SGLANG_LOG_SCHEDULER_STATUS_TARGET"] = cls.temp_dir - env["SGLANG_LOG_SCHEDULER_STATUS_INTERVAL"] = "1" - cls.process = popen_launch_server( - "Qwen/Qwen3-0.6B", - DEFAULT_URL_FOR_TEST, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=["--skip-server-warmup", "--enable-metrics"], - env=env, - ) - cls.addClassCleanup(kill_process_tree, cls.process.pid) - - def test_scheduler_status_dump(self): - response = requests.post( - DEFAULT_URL_FOR_TEST + "/generate", - json={ - "text": "Hello", - "sampling_params": {"max_new_tokens": 8, "temperature": 0}, - }, - timeout=30, - ) - self.assertEqual(response.status_code, 200) - - time.sleep(2) - - events = list(_find_log_events(self.temp_dir, "scheduler.status")) - print(f"{events=}") - self.assertGreater(len(events), 0, "scheduler.status event not found") - data = events[0] - for field in ["timestamp", "rank", "running_rids", "queued_rids"]: - self.assertIn(field, data) - self.assertIsInstance(data["running_rids"], list) - self.assertIsInstance(data["queued_rids"], list) - - -def _find_log_events(log_dir: str, event_name: str): - for f in Path(log_dir).glob("*.log"): - for line in f.read_text().splitlines(): - idx = line.find("{") - if idx == -1: - continue - try: - data = json.loads(line[idx:]) - except json.JSONDecodeError: - continue - if data.get("event") == event_name: - yield data - - -if __name__ == "__main__": - unittest.main()