[CI] Move nightly tests to test/nightly/ (#13683)
This commit is contained in:
alisonshao
@@ -1,258 +0,0 @@
|
||||
# Adapted from https://github.com/thinking-machines-lab/batch_invariant_ops/blob/main/test_batch_invariance.py
|
||||
import math
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.batch_invariant_ops import batch_invariant_ops
|
||||
from sglang.srt.batch_invariant_ops.batch_invariant_ops import set_batch_invariant_mode
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
device_type = getattr(torch.accelerator.current_accelerator(), "type", "cpu")
|
||||
torch.set_default_device(device_type)
|
||||
|
||||
# Just to get the logging out of the way
|
||||
with set_batch_invariant_mode(True):
|
||||
pass
|
||||
|
||||
|
||||
class TestBatchInvariantOps(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
batch_invariant_ops._ENABLE_MM_COMPARISON_TEST = True
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
batch_invariant_ops._ENABLE_MM_COMPARISON_TEST = False
|
||||
|
||||
def _test_batch_invariance(self, M, K, N, dtype):
|
||||
"""
|
||||
Test that matrix operations produce identical results for:
|
||||
- Method 1: Matrix-vector multiplication (batch size 1)
|
||||
- Method 2: Matrix-matrix multiplication, then slice (full batch)
|
||||
"""
|
||||
a = torch.linspace(-100, 100, M * K, dtype=dtype).reshape(M, K)
|
||||
|
||||
# Create non-contiguous tensor
|
||||
b = torch.linspace(-100, 100, K * N, dtype=dtype).reshape(N, K)
|
||||
b = b.transpose(0, 1)
|
||||
|
||||
# Method 1: Matrix-vector multiplication (batch size 1)
|
||||
out1 = torch.mm(a[:1], b)
|
||||
|
||||
# Method 2: Matrix-matrix multiplication, then slice (full batch)
|
||||
out2_pre = torch.mm(a, b)
|
||||
out2 = out2_pre[:1]
|
||||
|
||||
# Check if results are identical
|
||||
diff = (out1 - out2).abs().max()
|
||||
return diff.item()
|
||||
|
||||
def _run_multiple_iterations(self, iters, M, K, N, dtype):
|
||||
"""Run multiple iterations and collect diff statistics"""
|
||||
difflist = []
|
||||
for _ in range(iters):
|
||||
diff = self._test_batch_invariance(M, K, N, dtype)
|
||||
difflist.append(diff)
|
||||
return difflist
|
||||
|
||||
def _assert_batch_invariant_results(self, difflist, dtype, test_name):
|
||||
"""
|
||||
Assert that in batch-invariant mode:
|
||||
1. All diffs must not be NaN
|
||||
2. All diffs must be exactly 0
|
||||
3. Max, min, and diff of diffs must all be 0
|
||||
"""
|
||||
max_diff = max(difflist)
|
||||
min_diff = min(difflist)
|
||||
diff_range = max_diff - min_diff
|
||||
|
||||
# Check for NaN values
|
||||
self.assertFalse(
|
||||
math.isnan(max_diff), f"{test_name}: max_diff is NaN for {dtype}"
|
||||
)
|
||||
self.assertFalse(
|
||||
math.isnan(min_diff), f"{test_name}: min_diff is NaN for {dtype}"
|
||||
)
|
||||
self.assertFalse(
|
||||
math.isnan(diff_range), f"{test_name}: diff_range is NaN for {dtype}"
|
||||
)
|
||||
|
||||
# Check that all diffs are exactly 0
|
||||
self.assertEqual(
|
||||
max_diff,
|
||||
0.0,
|
||||
f"{test_name}: max_diff must be 0 in batch-invariant mode, got {max_diff} for {dtype}",
|
||||
)
|
||||
self.assertEqual(
|
||||
min_diff,
|
||||
0.0,
|
||||
f"{test_name}: min_diff must be 0 in batch-invariant mode, got {min_diff} for {dtype}",
|
||||
)
|
||||
self.assertEqual(
|
||||
diff_range,
|
||||
0.0,
|
||||
f"{test_name}: diff_range must be 0 in batch-invariant mode, got {diff_range} for {dtype}",
|
||||
)
|
||||
|
||||
def test_small_matrices(self):
|
||||
"""Test batch invariance with small matrix sizes"""
|
||||
test_cases = [
|
||||
("Small-1", 8, 64, 128),
|
||||
("Small-2", 16, 128, 256),
|
||||
("Small-3", 4, 32, 64),
|
||||
]
|
||||
|
||||
for name, M, K, N in test_cases:
|
||||
with self.subTest(name=name, M=M, K=K, N=N):
|
||||
for dtype in [torch.float32, torch.bfloat16]:
|
||||
with self.subTest(dtype=dtype):
|
||||
# Run with batch-invariant mode
|
||||
with set_batch_invariant_mode(True):
|
||||
difflist = self._run_multiple_iterations(
|
||||
iters=5, M=M, K=K, N=N, dtype=dtype
|
||||
)
|
||||
self._assert_batch_invariant_results(difflist, dtype, name)
|
||||
|
||||
def test_medium_matrices(self):
|
||||
"""Test batch invariance with medium matrix sizes"""
|
||||
test_cases = [
|
||||
("Medium-1", 32, 128, 1024),
|
||||
("Medium-2", 64, 512, 2048),
|
||||
("Medium-3", 24, 192, 768),
|
||||
]
|
||||
|
||||
for name, M, K, N in test_cases:
|
||||
with self.subTest(name=name, M=M, K=K, N=N):
|
||||
for dtype in [torch.float32, torch.bfloat16]:
|
||||
with self.subTest(dtype=dtype):
|
||||
# Run with batch-invariant mode
|
||||
with set_batch_invariant_mode(True):
|
||||
difflist = self._run_multiple_iterations(
|
||||
iters=5, M=M, K=K, N=N, dtype=dtype
|
||||
)
|
||||
self._assert_batch_invariant_results(difflist, dtype, name)
|
||||
|
||||
def test_large_matrices(self):
|
||||
"""Test batch invariance with large matrix sizes"""
|
||||
test_cases = [
|
||||
("Large-1", 128, 1024, 4096),
|
||||
("Large-2", 256, 2048, 8192),
|
||||
("Large-3", 96, 768, 3072),
|
||||
]
|
||||
|
||||
for name, M, K, N in test_cases:
|
||||
with self.subTest(name=name, M=M, K=K, N=N):
|
||||
for dtype in [torch.float32, torch.bfloat16]:
|
||||
with self.subTest(dtype=dtype):
|
||||
# Run with batch-invariant mode
|
||||
with set_batch_invariant_mode(True):
|
||||
difflist = self._run_multiple_iterations(
|
||||
iters=5, M=M, K=K, N=N, dtype=dtype
|
||||
)
|
||||
self._assert_batch_invariant_results(difflist, dtype, name)
|
||||
|
||||
def test_without_batch_invariant_mode(self):
|
||||
"""
|
||||
Test that without batch-invariant mode, results may differ.
|
||||
This test demonstrates the difference batch-invariant mode makes.
|
||||
"""
|
||||
M, K, N = 32, 128, 1024
|
||||
dtype = torch.float32
|
||||
|
||||
# Run without batch-invariant mode
|
||||
with set_batch_invariant_mode(False):
|
||||
difflist = self._run_multiple_iterations(
|
||||
iters=5, M=M, K=K, N=N, dtype=dtype
|
||||
)
|
||||
print(f"Without batch-invariant mode, we get diffs: {difflist}")
|
||||
|
||||
def _test_bmm_batch_invariance(self, B, M, K, N, dtype):
|
||||
"""
|
||||
Test that BMM operations produce identical results for:
|
||||
- Method 1: BMM with subset of batches
|
||||
- Method 2: BMM with all batches, then slice
|
||||
"""
|
||||
a = torch.linspace(-100, 100, B * M * K, dtype=dtype).reshape(B, M, K)
|
||||
b = torch.linspace(-100, 100, B * K * N, dtype=dtype).reshape(B, K, N)
|
||||
|
||||
# Method 1: BMM with subset (first 2 batches)
|
||||
subset_size = min(2, B)
|
||||
out1 = torch.bmm(a[:subset_size], b[:subset_size])
|
||||
|
||||
# Method 2: BMM with all batches, then slice
|
||||
out2_pre = torch.bmm(a, b)
|
||||
out2 = out2_pre[:subset_size]
|
||||
|
||||
# Check if results are identical
|
||||
diff = (out1 - out2).abs().max()
|
||||
return diff.item()
|
||||
|
||||
def _run_bmm_multiple_iterations(self, iters, B, M, K, N, dtype):
|
||||
"""Run multiple BMM iterations and collect diff statistics"""
|
||||
difflist = []
|
||||
for _ in range(iters):
|
||||
diff = self._test_bmm_batch_invariance(B, M, K, N, dtype)
|
||||
difflist.append(diff)
|
||||
return difflist
|
||||
|
||||
def test_bmm_small_matrices(self):
|
||||
"""Test BMM batch invariance with small matrix sizes"""
|
||||
test_cases = [
|
||||
("BMM-Small-1", 4, 8, 64, 128),
|
||||
("BMM-Small-2", 8, 16, 128, 256),
|
||||
("BMM-Small-3", 6, 4, 32, 64),
|
||||
]
|
||||
|
||||
for name, B, M, K, N in test_cases:
|
||||
with self.subTest(name=name, B=B, M=M, K=K, N=N):
|
||||
for dtype in [torch.float32, torch.bfloat16]:
|
||||
with self.subTest(dtype=dtype):
|
||||
# Run with batch-invariant mode
|
||||
with set_batch_invariant_mode(True):
|
||||
difflist = self._run_bmm_multiple_iterations(
|
||||
iters=5, B=B, M=M, K=K, N=N, dtype=dtype
|
||||
)
|
||||
self._assert_batch_invariant_results(difflist, dtype, name)
|
||||
|
||||
def test_bmm_medium_matrices(self):
|
||||
"""Test BMM batch invariance with medium matrix sizes"""
|
||||
test_cases = [
|
||||
("BMM-Medium-1", 8, 32, 128, 1024),
|
||||
("BMM-Medium-2", 16, 64, 512, 2048),
|
||||
("BMM-Medium-3", 12, 24, 192, 768),
|
||||
]
|
||||
|
||||
for name, B, M, K, N in test_cases:
|
||||
with self.subTest(name=name, B=B, M=M, K=K, N=N):
|
||||
for dtype in [torch.float32, torch.bfloat16]:
|
||||
with self.subTest(dtype=dtype):
|
||||
# Run with batch-invariant mode
|
||||
with set_batch_invariant_mode(True):
|
||||
difflist = self._run_bmm_multiple_iterations(
|
||||
iters=5, B=B, M=M, K=K, N=N, dtype=dtype
|
||||
)
|
||||
self._assert_batch_invariant_results(difflist, dtype, name)
|
||||
|
||||
def test_bmm_large_matrices(self):
|
||||
"""Test BMM batch invariance with large matrix sizes"""
|
||||
test_cases = [
|
||||
("BMM-Large-1", 16, 128, 1024, 4096),
|
||||
("BMM-Large-2", 32, 256, 2048, 8192),
|
||||
("BMM-Large-3", 24, 96, 768, 3072),
|
||||
]
|
||||
|
||||
for name, B, M, K, N in test_cases:
|
||||
with self.subTest(name=name, B=B, M=M, K=K, N=N):
|
||||
for dtype in [torch.float32, torch.bfloat16]:
|
||||
with self.subTest(dtype=dtype):
|
||||
# Run with batch-invariant mode
|
||||
with set_batch_invariant_mode(True):
|
||||
difflist = self._run_bmm_multiple_iterations(
|
||||
iters=5, B=B, M=M, K=K, N=N, dtype=dtype
|
||||
)
|
||||
self._assert_batch_invariant_results(difflist, dtype, name)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,612 +0,0 @@
|
||||
import unittest
|
||||
from typing import Optional
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers import dp_attention as _dp_attn
|
||||
|
||||
# Patch DP-attention globals before importing backends
|
||||
_dp_attn.get_attention_tp_size = lambda: 1 # TP size = 1 for unit test
|
||||
|
||||
from sglang.srt.configs.model_config import AttentionArch
|
||||
from sglang.srt.layers.attention.nsa.nsa_indexer import (
|
||||
BaseIndexerMetadata,
|
||||
Indexer,
|
||||
rotate_activation,
|
||||
)
|
||||
from sglang.srt.layers.attention.nsa_backend import NativeSparseAttnBackend
|
||||
from sglang.srt.layers.layernorm import LayerNorm
|
||||
from sglang.srt.layers.linear import LinearBase
|
||||
from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
|
||||
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
# Global configuration for all indexer tests
|
||||
DEFAULT_CONFIG = {
|
||||
"device": "cuda",
|
||||
"dtype": torch.bfloat16,
|
||||
"kv_cache_dtype": torch.float8_e4m3fn,
|
||||
"context_len": 2048,
|
||||
"max_bs": 64,
|
||||
"hidden_size": 5120,
|
||||
"index_n_heads": 1,
|
||||
"index_head_dim": 128,
|
||||
"rope_head_dim": 64,
|
||||
"index_topk": 64,
|
||||
"q_lora_rank": 1536,
|
||||
"kv_lora_rank": 512,
|
||||
"qk_rope_head_dim": 64,
|
||||
"max_position_embeddings": 163840,
|
||||
"rope_theta": 10000.0,
|
||||
"layer_id": 0,
|
||||
"page_size": 64,
|
||||
}
|
||||
|
||||
|
||||
class MockIndexerMetadata(BaseIndexerMetadata):
|
||||
"""Mock implementation of BaseIndexerMetadata for testing."""
|
||||
|
||||
def __init__(self, batch_size, seq_lens, page_table=None):
|
||||
self.batch_size = batch_size
|
||||
self.seq_lens = seq_lens
|
||||
self.page_table = page_table
|
||||
self.device = "cuda"
|
||||
|
||||
def get_seqlens_int32(self) -> torch.Tensor:
|
||||
"""Return: (batch_size,) int32 tensor"""
|
||||
return torch.tensor(self.seq_lens, dtype=torch.int32, device=self.device)
|
||||
|
||||
def get_page_table_64(self) -> torch.Tensor:
|
||||
"""Return: (batch_size, num_blocks) int32, page table with page size 64."""
|
||||
if self.page_table is not None:
|
||||
return self.page_table
|
||||
# Create a simple page table for testing
|
||||
max_seq_len = max(self.seq_lens)
|
||||
num_blocks = (max_seq_len + 63) // 64 # Round up to page size 64
|
||||
page_table = torch.zeros(
|
||||
(self.batch_size, num_blocks), dtype=torch.int32, device=self.device
|
||||
)
|
||||
for i in range(self.batch_size):
|
||||
# Simple linear mapping: block i maps to page i
|
||||
num_blocks_needed = (self.seq_lens[i] + 63) // 64
|
||||
page_table[i, :num_blocks_needed] = torch.arange(
|
||||
num_blocks_needed, device=self.device
|
||||
)
|
||||
return page_table
|
||||
|
||||
def get_seqlens_expanded(self) -> torch.Tensor:
|
||||
"""Return: (sum_extend_seq_len,) int32 tensor"""
|
||||
# For extend mode, each new token attends to progressively more tokens
|
||||
# For a sequence being extended from position 0 to seq_len, token i attends to i+1 tokens
|
||||
result = []
|
||||
for seq_len in self.seq_lens:
|
||||
result.extend(range(1, seq_len + 1))
|
||||
return torch.tensor(result, dtype=torch.int32, device=self.device)
|
||||
|
||||
def topk_transform(
|
||||
self,
|
||||
logits: torch.Tensor,
|
||||
topk: int,
|
||||
ks: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Perform topk selection on the logits.
|
||||
For testing, just return the topk indices.
|
||||
"""
|
||||
return torch.topk(logits, k=topk, dim=-1).indices
|
||||
|
||||
|
||||
class MockModelRunner:
|
||||
def __init__(self, config=None):
|
||||
self.device = "cuda"
|
||||
self.config = {**DEFAULT_CONFIG, **(config or {})}
|
||||
self.dtype = self.config["dtype"]
|
||||
self.kv_cache_dtype = self.config["kv_cache_dtype"]
|
||||
self.is_hybrid = False
|
||||
|
||||
# Model configuration
|
||||
attention_arch = AttentionArch.MLA
|
||||
max_context_len = self.config["context_len"]
|
||||
max_batch_size = self.config["max_bs"]
|
||||
|
||||
# Create mock hf_config for NSA - instantiate it as an object, not a type
|
||||
hf_config = type(
|
||||
"HfConfig",
|
||||
(),
|
||||
{
|
||||
"architectures": ["DeepseekV3ForCausalLM"],
|
||||
"index_topk": self.config["index_topk"],
|
||||
"index_head_dim": self.config["index_head_dim"],
|
||||
"index_n_heads": self.config["index_n_heads"],
|
||||
},
|
||||
)()
|
||||
|
||||
self.model_config = type(
|
||||
"ModelConfig",
|
||||
(),
|
||||
{
|
||||
"context_len": max_context_len,
|
||||
"is_multimodal": False,
|
||||
"attention_arch": attention_arch,
|
||||
"num_attention_heads": 128,
|
||||
"kv_lora_rank": self.config["kv_lora_rank"],
|
||||
"qk_rope_head_dim": self.config["qk_rope_head_dim"],
|
||||
"hf_config": hf_config,
|
||||
},
|
||||
)()
|
||||
|
||||
self.sliding_window_size = None
|
||||
self.page_size = self.config["page_size"]
|
||||
|
||||
# Create req_to_token_pool
|
||||
self.req_to_token_pool = type(
|
||||
"TokenPool",
|
||||
(),
|
||||
{
|
||||
"size": max_batch_size,
|
||||
"req_to_token": torch.zeros(
|
||||
max_batch_size,
|
||||
max_context_len,
|
||||
dtype=torch.int32,
|
||||
device=self.device,
|
||||
),
|
||||
},
|
||||
)()
|
||||
|
||||
# Create NSATokenToKVPool
|
||||
max_total_num_tokens = max_batch_size * max_context_len
|
||||
self.token_to_kv_pool = NSATokenToKVPool(
|
||||
size=max_total_num_tokens,
|
||||
page_size=self.config["page_size"],
|
||||
dtype=self.config["kv_cache_dtype"],
|
||||
kv_lora_rank=self.config["kv_lora_rank"],
|
||||
qk_rope_head_dim=self.config["qk_rope_head_dim"],
|
||||
layer_num=1,
|
||||
device=self.device,
|
||||
index_head_dim=self.config["index_head_dim"],
|
||||
enable_memory_saver=False,
|
||||
)
|
||||
|
||||
# Required by backend with NSA-specific attributes
|
||||
self.server_args = type(
|
||||
"ServerArgs",
|
||||
(),
|
||||
{
|
||||
"kv_cache_dtype": "auto",
|
||||
"speculative_eagle_topk": None,
|
||||
"speculative_num_draft_tokens": 0,
|
||||
"enable_deterministic_inference": False,
|
||||
"nsa_prefill_backend": "flashmla_sparse",
|
||||
"nsa_decode_backend": "fa3",
|
||||
},
|
||||
)()
|
||||
|
||||
|
||||
@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
|
||||
class TestNSAIndexer(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""Set up global server args for testing."""
|
||||
server_args = ServerArgs(model_path="dummy")
|
||||
server_args.enable_dp_attention = False
|
||||
server_args.nsa_prefill_backend = "flashmla_sparse"
|
||||
server_args.nsa_decode_backend = "flashmla_sparse"
|
||||
set_global_server_args_for_scheduler(server_args)
|
||||
|
||||
# Check GPU capability for FP8
|
||||
if torch.cuda.is_available():
|
||||
compute_capability = torch.cuda.get_device_capability()
|
||||
cls.supports_fp8 = compute_capability[0] >= 9 # Hopper or newer
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""Clean up after all tests."""
|
||||
pass
|
||||
|
||||
def setUp(self):
|
||||
# Test parameters
|
||||
self.batch_size = 2
|
||||
self.seq_len = 128
|
||||
self.config = DEFAULT_CONFIG.copy()
|
||||
self.device = "cuda"
|
||||
self.dtype = torch.bfloat16
|
||||
|
||||
def _init_model_runner(self, config_override=None):
|
||||
"""Initialize model runner with optional config override."""
|
||||
config = self.config.copy()
|
||||
if config_override:
|
||||
config.update(config_override)
|
||||
self.model_runner = MockModelRunner(config)
|
||||
self.backend = NativeSparseAttnBackend(self.model_runner)
|
||||
|
||||
def _create_indexer(self, **kwargs):
|
||||
"""Create an Indexer instance with default parameters."""
|
||||
params = {
|
||||
"hidden_size": self.config["hidden_size"],
|
||||
"index_n_heads": self.config["index_n_heads"],
|
||||
"index_head_dim": self.config["index_head_dim"],
|
||||
"rope_head_dim": self.config["rope_head_dim"],
|
||||
"index_topk": self.config["index_topk"],
|
||||
"q_lora_rank": self.config["q_lora_rank"],
|
||||
"max_position_embeddings": self.config["max_position_embeddings"],
|
||||
"rope_theta": self.config["rope_theta"],
|
||||
"layer_id": self.config["layer_id"],
|
||||
"scale_fmt": "ue8m0",
|
||||
"block_size": 128,
|
||||
"quant_config": None, # No quantization for testing
|
||||
}
|
||||
params.update(kwargs)
|
||||
|
||||
torch.set_default_dtype(self.dtype)
|
||||
indexer = Indexer(**params)
|
||||
# Move indexer to CUDA device
|
||||
indexer = indexer.to(device=self.device)
|
||||
|
||||
# Convert linear layer weights to bfloat16 (but preserve LayerNorm's float32)
|
||||
# Need to recursively convert LinearBase submodules (like ReplicatedLinear)
|
||||
for name, module in indexer.named_modules():
|
||||
# Check for LinearBase (parent of ReplicatedLinear) but exclude LayerNorm
|
||||
if isinstance(module, LinearBase) and not isinstance(module, LayerNorm):
|
||||
module.to(dtype=self.dtype)
|
||||
|
||||
return indexer
|
||||
|
||||
def _create_forward_batch(
|
||||
self, mode, batch_size=None, seq_len=None, extend_len=None
|
||||
):
|
||||
"""Create a forward batch for testing."""
|
||||
batch_size = batch_size or self.batch_size
|
||||
seq_len = seq_len or self.seq_len
|
||||
|
||||
if mode == ForwardMode.EXTEND:
|
||||
q_len = extend_len or seq_len
|
||||
total_len = seq_len
|
||||
|
||||
forward_batch = ForwardBatch(
|
||||
batch_size=batch_size,
|
||||
input_ids=torch.randint(
|
||||
0, 100, (batch_size, q_len), device=self.device
|
||||
),
|
||||
out_cache_loc=torch.arange(
|
||||
batch_size * (total_len - q_len),
|
||||
batch_size * total_len,
|
||||
device=self.device,
|
||||
),
|
||||
seq_lens_sum=batch_size * total_len,
|
||||
forward_mode=mode,
|
||||
req_pool_indices=torch.arange(batch_size, device=self.device),
|
||||
seq_lens=torch.tensor([total_len] * batch_size, device=self.device),
|
||||
seq_lens_cpu=torch.tensor([total_len] * batch_size, device="cpu"),
|
||||
extend_prefix_lens=torch.tensor(
|
||||
[total_len - q_len] * batch_size, device=self.device
|
||||
),
|
||||
extend_prefix_lens_cpu=torch.tensor(
|
||||
[total_len - q_len] * batch_size, device="cpu"
|
||||
),
|
||||
extend_seq_lens=torch.tensor([q_len] * batch_size, device=self.device),
|
||||
extend_seq_lens_cpu=torch.tensor([q_len] * batch_size, device="cpu"),
|
||||
attn_backend=self.backend,
|
||||
)
|
||||
else: # ForwardMode.DECODE
|
||||
decode_len = 1
|
||||
total_len = seq_len + decode_len
|
||||
|
||||
forward_batch = ForwardBatch(
|
||||
batch_size=batch_size,
|
||||
input_ids=torch.randint(
|
||||
0, 100, (batch_size, decode_len), device=self.device
|
||||
),
|
||||
out_cache_loc=torch.arange(
|
||||
batch_size * seq_len, batch_size * total_len, device=self.device
|
||||
),
|
||||
seq_lens_sum=batch_size * total_len,
|
||||
forward_mode=mode,
|
||||
req_pool_indices=torch.arange(batch_size, device=self.device),
|
||||
seq_lens=torch.tensor([total_len] * batch_size, device=self.device),
|
||||
seq_lens_cpu=torch.tensor([total_len] * batch_size, device="cpu"),
|
||||
attn_backend=self.backend,
|
||||
)
|
||||
|
||||
# Add token pools
|
||||
forward_batch.req_to_token_pool = self.model_runner.req_to_token_pool
|
||||
forward_batch.token_to_kv_pool = self.model_runner.token_to_kv_pool
|
||||
|
||||
# Mock write to req_to_token_pool
|
||||
page_size = self.model_runner.page_size
|
||||
for i in range(batch_size):
|
||||
seq_length = total_len
|
||||
for j in range(seq_length):
|
||||
self.model_runner.req_to_token_pool.req_to_token[i, j] = (
|
||||
i * seq_length + j + page_size
|
||||
)
|
||||
|
||||
return forward_batch
|
||||
|
||||
def _verify_topk_output(self, topk_indices, batch_size, q_len, topk):
|
||||
"""Verify the topk indices output shape and basic properties."""
|
||||
self.assertIsNotNone(topk_indices)
|
||||
self.assertEqual(topk_indices.device.type, "cuda")
|
||||
|
||||
# Check shape - should be (total_q_len, topk_padded)
|
||||
# where topk_padded is aligned to 2048
|
||||
self.assertEqual(len(topk_indices.shape), 2)
|
||||
self.assertEqual(topk_indices.shape[0], batch_size * q_len)
|
||||
|
||||
# Check that topk is padded to at least topk
|
||||
self.assertGreaterEqual(topk_indices.shape[1], topk)
|
||||
|
||||
# Check for padding values (-1)
|
||||
has_padding = (topk_indices == -1).any()
|
||||
self.assertTrue(
|
||||
has_padding or topk_indices.shape[1] == topk,
|
||||
"Output should have padding or exact topk size",
|
||||
)
|
||||
|
||||
@patch("sglang.srt.layers.attention.nsa.nsa_indexer.deep_gemm")
|
||||
def test_indexer_basic_creation(self, mock_deep_gemm):
|
||||
"""Test basic indexer creation and initialization."""
|
||||
mock_deep_gemm.get_num_sms.return_value = 132
|
||||
|
||||
indexer = self._create_indexer()
|
||||
|
||||
self.assertEqual(indexer.hidden_size, self.config["hidden_size"])
|
||||
self.assertEqual(indexer.n_heads, self.config["index_n_heads"])
|
||||
self.assertEqual(indexer.head_dim, self.config["index_head_dim"])
|
||||
self.assertEqual(indexer.rope_head_dim, self.config["rope_head_dim"])
|
||||
self.assertEqual(indexer.index_topk, self.config["index_topk"])
|
||||
self.assertEqual(indexer.layer_id, self.config["layer_id"])
|
||||
|
||||
@patch("sglang.srt.layers.attention.nsa.nsa_indexer.deep_gemm")
|
||||
@patch("sglang.srt.layers.attention.nsa.triton_kernel.act_quant")
|
||||
def test_forward_extend_mode(self, mock_act_quant, mock_deep_gemm):
|
||||
"""Test indexer forward pass in extend mode."""
|
||||
if not self.supports_fp8:
|
||||
self.skipTest("FP8 requires Hopper GPU or newer")
|
||||
|
||||
# Setup mocks
|
||||
mock_deep_gemm.get_num_sms.return_value = 132
|
||||
mock_deep_gemm.get_paged_mqa_logits_metadata.return_value = MagicMock()
|
||||
|
||||
def mock_quant(x, *args, **kwargs):
|
||||
# Return FP8 tensor and scale
|
||||
return x.to(torch.float8_e4m3fn), torch.ones(
|
||||
x.shape[0], dtype=torch.float32, device=x.device
|
||||
)
|
||||
|
||||
mock_act_quant.side_effect = mock_quant
|
||||
|
||||
# Mock deep_gemm.fp8_mqa_logits to return logits (ragged path)
|
||||
def mock_mqa_logits(q, kv, weights, ks, ke, *args, **kwargs):
|
||||
# q shape: (sum_extend_seq_len, ...), return logits for each query token
|
||||
num_queries = q.shape[0]
|
||||
# kv is a tuple (k_fp8, k_scale), get total number of keys from k_fp8
|
||||
k_fp8, k_scale = kv
|
||||
max_kv_len = k_fp8.shape[0] # Total keys across all batches (k_offset)
|
||||
return torch.randn(
|
||||
num_queries, max_kv_len, dtype=torch.float32, device="cuda"
|
||||
)
|
||||
|
||||
mock_deep_gemm.fp8_mqa_logits.side_effect = mock_mqa_logits
|
||||
|
||||
# Also mock the paged version for completeness
|
||||
def mock_paged_mqa_logits(q, kv, weights, *args, **kwargs):
|
||||
batch_size = q.shape[0]
|
||||
seq_len = 128
|
||||
return torch.randn(batch_size, seq_len, dtype=torch.float32, device="cuda")
|
||||
|
||||
mock_deep_gemm.fp8_paged_mqa_logits.side_effect = mock_paged_mqa_logits
|
||||
|
||||
self._init_model_runner()
|
||||
|
||||
indexer = self._create_indexer()
|
||||
forward_batch = self._create_forward_batch(ForwardMode.EXTEND)
|
||||
|
||||
# Create input tensors
|
||||
total_tokens = self.batch_size * self.seq_len
|
||||
hidden_states = torch.randn(
|
||||
total_tokens,
|
||||
self.config["hidden_size"],
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
)
|
||||
q_lora = torch.randn(
|
||||
total_tokens,
|
||||
self.config["q_lora_rank"],
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
)
|
||||
positions = torch.arange(total_tokens, device=self.device)
|
||||
|
||||
# Run forward pass
|
||||
with patch.object(
|
||||
self.backend,
|
||||
"get_indexer_metadata",
|
||||
return_value=MockIndexerMetadata(
|
||||
self.batch_size, [self.seq_len] * self.batch_size
|
||||
),
|
||||
):
|
||||
topk_indices = indexer(
|
||||
x=hidden_states,
|
||||
q_lora=q_lora,
|
||||
positions=positions,
|
||||
forward_batch=forward_batch,
|
||||
layer_id=self.config["layer_id"],
|
||||
)
|
||||
|
||||
# Verify output
|
||||
self._verify_topk_output(
|
||||
topk_indices, self.batch_size, self.seq_len, self.config["index_topk"]
|
||||
)
|
||||
|
||||
@patch("sglang.srt.layers.attention.nsa.nsa_indexer.deep_gemm")
|
||||
@patch("sglang.srt.layers.attention.nsa.triton_kernel.act_quant")
|
||||
def test_forward_decode_mode(self, mock_act_quant, mock_deep_gemm):
|
||||
"""Test indexer forward pass in decode mode."""
|
||||
if not self.supports_fp8:
|
||||
self.skipTest("FP8 requires Hopper GPU or newer")
|
||||
|
||||
# Setup mocks
|
||||
mock_deep_gemm.get_num_sms.return_value = 132
|
||||
mock_deep_gemm.get_paged_mqa_logits_metadata.return_value = MagicMock()
|
||||
|
||||
def mock_quant(x, *args, **kwargs):
|
||||
return x.to(torch.float8_e4m3fn), torch.ones(
|
||||
x.shape[0], dtype=torch.float32, device=x.device
|
||||
)
|
||||
|
||||
mock_act_quant.side_effect = mock_quant
|
||||
|
||||
def mock_paged_mqa_logits(q, kv, weights, *args, **kwargs):
|
||||
batch_size = q.shape[0]
|
||||
seq_len = 128
|
||||
return torch.randn(batch_size, seq_len, dtype=torch.float32, device="cuda")
|
||||
|
||||
mock_deep_gemm.fp8_paged_mqa_logits.side_effect = mock_paged_mqa_logits
|
||||
|
||||
self._init_model_runner()
|
||||
|
||||
indexer = self._create_indexer()
|
||||
forward_batch = self._create_forward_batch(ForwardMode.DECODE)
|
||||
|
||||
# Create input tensors for decode (batch_size tokens only)
|
||||
hidden_states = torch.randn(
|
||||
self.batch_size,
|
||||
self.config["hidden_size"],
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
)
|
||||
q_lora = torch.randn(
|
||||
self.batch_size,
|
||||
self.config["q_lora_rank"],
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
)
|
||||
positions = torch.arange(self.batch_size, device=self.device)
|
||||
|
||||
# Run forward pass
|
||||
with patch.object(
|
||||
self.backend,
|
||||
"get_indexer_metadata",
|
||||
return_value=MockIndexerMetadata(
|
||||
self.batch_size, [self.seq_len + 1] * self.batch_size
|
||||
),
|
||||
):
|
||||
topk_indices = indexer(
|
||||
x=hidden_states,
|
||||
q_lora=q_lora,
|
||||
positions=positions,
|
||||
forward_batch=forward_batch,
|
||||
layer_id=self.config["layer_id"],
|
||||
)
|
||||
|
||||
# Verify output - decode mode has q_len=1
|
||||
self._verify_topk_output(
|
||||
topk_indices, self.batch_size, 1, self.config["index_topk"]
|
||||
)
|
||||
|
||||
def test_rotate_activation(self):
|
||||
"""Test the Hadamard transform (rotate_activation) function."""
|
||||
# Test with power-of-2 hidden size
|
||||
hidden_size = 128
|
||||
x = torch.randn(16, hidden_size, dtype=torch.bfloat16, device=self.device)
|
||||
|
||||
try:
|
||||
output = rotate_activation(x)
|
||||
self.assertEqual(output.shape, x.shape)
|
||||
self.assertEqual(output.dtype, torch.bfloat16)
|
||||
except ImportError:
|
||||
self.skipTest("sgl_kernel not available for hadamard_transform")
|
||||
|
||||
def test_rotate_activation_invalid_size(self):
|
||||
"""Test that rotate_activation fails with non-power-of-2 size."""
|
||||
# Test with non-power-of-2 hidden size
|
||||
hidden_size = 129 # Not a power of 2
|
||||
x = torch.randn(16, hidden_size, dtype=torch.bfloat16, device=self.device)
|
||||
|
||||
with self.assertRaises(AssertionError):
|
||||
rotate_activation(x)
|
||||
|
||||
def test_indexer_metadata_interface(self):
|
||||
"""Test the BaseIndexerMetadata interface implementation."""
|
||||
batch_size = 4
|
||||
seq_lens = [64, 128, 96, 112]
|
||||
|
||||
metadata = MockIndexerMetadata(batch_size, seq_lens)
|
||||
|
||||
# Test get_seqlens_int32
|
||||
seqlens = metadata.get_seqlens_int32()
|
||||
self.assertEqual(seqlens.shape, (batch_size,))
|
||||
self.assertEqual(seqlens.dtype, torch.int32)
|
||||
self.assertTrue(torch.all(seqlens == torch.tensor(seq_lens, device="cuda")))
|
||||
|
||||
# Test get_page_table_64
|
||||
page_table = metadata.get_page_table_64()
|
||||
self.assertEqual(len(page_table.shape), 2)
|
||||
self.assertEqual(page_table.shape[0], batch_size)
|
||||
self.assertEqual(page_table.dtype, torch.int32)
|
||||
|
||||
# Test topk_transform
|
||||
logits = torch.randn(batch_size, 128, device="cuda")
|
||||
topk = 64
|
||||
topk_indices = metadata.topk_transform(logits, topk)
|
||||
self.assertEqual(topk_indices.shape, (batch_size, topk))
|
||||
|
||||
# TODO: enable this test after indexer accuracy aligned
|
||||
# @patch("sglang.srt.layers.attention.nsa.nsa_indexer.deep_gemm")
|
||||
# def test_indexer_with_different_topk(self, mock_deep_gemm):
|
||||
# """Test indexer with different topk values."""
|
||||
# mock_deep_gemm.get_num_sms.return_value = 132
|
||||
|
||||
# for topk in [32, 64, 128]:
|
||||
# with self.subTest(topk=topk):
|
||||
# indexer = self._create_indexer(index_topk=topk)
|
||||
# self.assertEqual(indexer.index_topk, topk)
|
||||
|
||||
@patch("sglang.srt.layers.attention.nsa.nsa_indexer.deep_gemm")
|
||||
def test_indexer_with_fused_wk(self, mock_deep_gemm):
|
||||
"""Test indexer creation with fused wk and weights projection."""
|
||||
mock_deep_gemm.get_num_sms.return_value = 132
|
||||
|
||||
# Note: fuse_wk_and_weights_proj feature is not currently implemented
|
||||
# This test verifies basic indexer creation still works
|
||||
indexer = self._create_indexer()
|
||||
self.assertIsNotNone(indexer)
|
||||
|
||||
@patch("sglang.srt.layers.attention.nsa.nsa_indexer.deep_gemm")
|
||||
def test_indexer_with_alt_stream(self, mock_deep_gemm):
|
||||
"""Test indexer creation with alternative CUDA stream."""
|
||||
mock_deep_gemm.get_num_sms.return_value = 132
|
||||
|
||||
alt_stream = torch.cuda.Stream()
|
||||
indexer = self._create_indexer(alt_stream=alt_stream)
|
||||
self.assertEqual(indexer.alt_stream, alt_stream)
|
||||
|
||||
def test_shape_sanity_checks(self):
|
||||
"""Test various shape combinations for consistency."""
|
||||
test_configs = [
|
||||
{"batch_size": 1, "seq_len": 64},
|
||||
{"batch_size": 4, "seq_len": 128},
|
||||
{"batch_size": 8, "seq_len": 256},
|
||||
]
|
||||
|
||||
for config in test_configs:
|
||||
with self.subTest(**config):
|
||||
batch_size = config["batch_size"]
|
||||
seq_len = config["seq_len"]
|
||||
|
||||
# Test metadata shapes
|
||||
metadata = MockIndexerMetadata(batch_size, [seq_len] * batch_size)
|
||||
|
||||
seqlens = metadata.get_seqlens_int32()
|
||||
self.assertEqual(seqlens.shape, (batch_size,))
|
||||
|
||||
page_table = metadata.get_page_table_64()
|
||||
expected_blocks = (seq_len + 63) // 64
|
||||
self.assertEqual(page_table.shape[0], batch_size)
|
||||
self.assertGreaterEqual(page_table.shape[1], expected_blocks)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,190 +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 eviction policies.
|
||||
Tests LRU and FIFO eviction behavior.
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from sglang.srt.lora.eviction_policy import get_eviction_policy
|
||||
|
||||
|
||||
class TestLoRAEvictionPolicy(unittest.TestCase):
|
||||
"""Unit tests for LoRA eviction policies."""
|
||||
|
||||
def _test_eviction_policy(
|
||||
self, policy_name, access_sequence, candidates, expected_victim
|
||||
):
|
||||
"""
|
||||
Helper to test eviction policy with given access pattern.
|
||||
|
||||
Args:
|
||||
policy_name: Name of eviction policy ("lru" or "fifo")
|
||||
access_sequence: List of adapter IDs in access order
|
||||
candidates: Set of adapter IDs that can be evicted
|
||||
expected_victim: Expected adapter ID to be evicted
|
||||
"""
|
||||
policy = get_eviction_policy(policy_name)
|
||||
|
||||
# Simulate access pattern
|
||||
for adapter_id in access_sequence:
|
||||
policy.mark_used(adapter_id)
|
||||
|
||||
# Select victim from candidates
|
||||
victim = policy.select_victim(candidates)
|
||||
self.assertEqual(
|
||||
victim,
|
||||
expected_victim,
|
||||
f"{policy_name.upper()}: Expected {expected_victim}, got {victim}",
|
||||
)
|
||||
|
||||
def test_lru_basic(self):
|
||||
"""Test LRU selects least recently used adapter."""
|
||||
self._test_eviction_policy(
|
||||
"lru",
|
||||
access_sequence=["lora1", "lora2", "lora3", "lora4"],
|
||||
candidates={"lora1", "lora2", "lora3", "lora4"},
|
||||
expected_victim="lora1",
|
||||
)
|
||||
|
||||
def test_lru_with_reuse(self):
|
||||
"""Test LRU updates order on reuse."""
|
||||
self._test_eviction_policy(
|
||||
"lru",
|
||||
access_sequence=["lora1", "lora2", "lora3", "lora4", "lora1"],
|
||||
candidates={"lora1", "lora2", "lora3", "lora4"},
|
||||
expected_victim="lora2",
|
||||
)
|
||||
|
||||
def test_lru_multiple_reuse(self):
|
||||
"""Test LRU with multiple reuses."""
|
||||
self._test_eviction_policy(
|
||||
"lru",
|
||||
access_sequence=["lora1", "lora2", "lora3", "lora1", "lora2"],
|
||||
candidates={"lora1", "lora2", "lora3"},
|
||||
expected_victim="lora3",
|
||||
)
|
||||
|
||||
def test_lru_with_subset_candidates(self):
|
||||
"""Test LRU with subset of candidates."""
|
||||
self._test_eviction_policy(
|
||||
"lru",
|
||||
access_sequence=["lora1", "lora2", "lora3", "lora4"],
|
||||
candidates={"lora2", "lora3", "lora4"},
|
||||
expected_victim="lora2",
|
||||
)
|
||||
|
||||
def test_lru_base_model_priority(self):
|
||||
"""Test LRU prioritizes base model for eviction."""
|
||||
self._test_eviction_policy(
|
||||
"lru",
|
||||
access_sequence=["lora1", "lora2", "lora3"],
|
||||
candidates={None, "lora1", "lora2", "lora3"},
|
||||
expected_victim=None,
|
||||
)
|
||||
|
||||
def test_fifo_basic(self):
|
||||
"""Test FIFO selects first inserted adapter."""
|
||||
self._test_eviction_policy(
|
||||
"fifo",
|
||||
access_sequence=["lora1", "lora2", "lora3", "lora4"],
|
||||
candidates={"lora1", "lora2", "lora3", "lora4"},
|
||||
expected_victim="lora1",
|
||||
)
|
||||
|
||||
def test_fifo_ignores_reuse(self):
|
||||
"""Test FIFO ignores reuse."""
|
||||
self._test_eviction_policy(
|
||||
"fifo",
|
||||
access_sequence=[
|
||||
"lora1",
|
||||
"lora2",
|
||||
"lora3",
|
||||
"lora4",
|
||||
"lora4",
|
||||
"lora3",
|
||||
"lora2",
|
||||
"lora1",
|
||||
],
|
||||
candidates={"lora1", "lora2", "lora3", "lora4"},
|
||||
expected_victim="lora1",
|
||||
)
|
||||
|
||||
def test_fifo_with_subset_candidates(self):
|
||||
"""Test FIFO with subset of candidates."""
|
||||
self._test_eviction_policy(
|
||||
"fifo",
|
||||
access_sequence=["lora1", "lora2", "lora3", "lora4"],
|
||||
candidates={"lora2", "lora3", "lora4"},
|
||||
expected_victim="lora2",
|
||||
)
|
||||
|
||||
def test_fifo_base_model_priority(self):
|
||||
"""Test FIFO prioritizes base model for eviction."""
|
||||
self._test_eviction_policy(
|
||||
"fifo",
|
||||
access_sequence=["lora1", "lora2", "lora3"],
|
||||
candidates={None, "lora1", "lora2", "lora3"},
|
||||
expected_victim=None,
|
||||
)
|
||||
|
||||
def test_policy_remove(self):
|
||||
"""Test that remove() correctly updates internal state."""
|
||||
lru = get_eviction_policy("lru")
|
||||
lru.mark_used("lora1")
|
||||
lru.mark_used("lora2")
|
||||
lru.mark_used("lora3")
|
||||
|
||||
# Remove lora1, so lora2 becomes LRU
|
||||
lru.remove("lora1")
|
||||
victim = lru.select_victim({"lora1", "lora2", "lora3"})
|
||||
self.assertEqual(victim, "lora2")
|
||||
|
||||
def test_eviction_policy_factory(self):
|
||||
"""Test eviction policy factory function."""
|
||||
# Test valid policies
|
||||
lru = get_eviction_policy("lru")
|
||||
fifo = get_eviction_policy("fifo")
|
||||
|
||||
self.assertIsNotNone(lru)
|
||||
self.assertIsNotNone(fifo)
|
||||
|
||||
# Test invalid policy
|
||||
with self.assertRaises(ValueError):
|
||||
get_eviction_policy("invalid_policy")
|
||||
|
||||
def test_lru_vs_fifo_behavior(self):
|
||||
"""Test that LRU and FIFO behave differently."""
|
||||
access_sequence = ["lora1", "lora2", "lora3", "lora1"]
|
||||
candidates = {"lora1", "lora2", "lora3"}
|
||||
|
||||
lru = get_eviction_policy("lru")
|
||||
for adapter_id in access_sequence:
|
||||
lru.mark_used(adapter_id)
|
||||
lru_victim = lru.select_victim(candidates)
|
||||
|
||||
fifo = get_eviction_policy("fifo")
|
||||
for adapter_id in access_sequence:
|
||||
fifo.mark_used(adapter_id)
|
||||
fifo_victim = fifo.select_victim(candidates)
|
||||
|
||||
self.assertNotEqual(lru_victim, fifo_victim)
|
||||
self.assertEqual(lru_victim, "lora2")
|
||||
self.assertEqual(fifo_victim, "lora1")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main(verbosity=2)
|
||||
@@ -1,327 +0,0 @@
|
||||
"""
|
||||
Unit tests for OpenAI-compatible LoRA API support.
|
||||
|
||||
Tests the model parameter parsing and LoRA adapter resolution logic
|
||||
that enables OpenAI-compatible LoRA adapter selection.
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
|
||||
|
||||
class MockTokenizerManager:
|
||||
"""Mock TokenizerManager for testing."""
|
||||
|
||||
def __init__(self, enable_lora=False):
|
||||
self.server_args = MagicMock(spec=ServerArgs)
|
||||
self.server_args.enable_lora = enable_lora
|
||||
self.server_args.tokenizer_metrics_allowed_custom_labels = None
|
||||
|
||||
|
||||
class ConcreteServingBase(OpenAIServingBase):
|
||||
"""Concrete implementation for testing abstract base class."""
|
||||
|
||||
def _request_id_prefix(self) -> str:
|
||||
return "test-"
|
||||
|
||||
def _convert_to_internal_request(self, request, raw_request=None):
|
||||
pass
|
||||
|
||||
def _validate_request(self, request):
|
||||
pass
|
||||
|
||||
|
||||
class TestParseModelParameter(unittest.TestCase):
|
||||
"""Test _parse_model_parameter method."""
|
||||
|
||||
def setUp(self):
|
||||
self.tokenizer_manager = MockTokenizerManager(enable_lora=True)
|
||||
self.serving = ConcreteServingBase(self.tokenizer_manager)
|
||||
|
||||
def test_model_without_adapter(self):
|
||||
"""Test parsing model without adapter returns None for adapter."""
|
||||
base_model, adapter = self.serving._parse_model_parameter("llama-3.1-8B")
|
||||
self.assertEqual(base_model, "llama-3.1-8B")
|
||||
self.assertIsNone(adapter)
|
||||
|
||||
def test_model_with_adapter(self):
|
||||
"""Test parsing model with adapter extracts both parts."""
|
||||
base_model, adapter = self.serving._parse_model_parameter(
|
||||
"llama-3.1-8B:sql-expert"
|
||||
)
|
||||
self.assertEqual(base_model, "llama-3.1-8B")
|
||||
self.assertEqual(adapter, "sql-expert")
|
||||
|
||||
def test_model_with_path_and_adapter(self):
|
||||
"""Test parsing model path with slashes and adapter."""
|
||||
base_model, adapter = self.serving._parse_model_parameter(
|
||||
"meta-llama/Llama-3.1-8B-Instruct:adapter-name"
|
||||
)
|
||||
self.assertEqual(base_model, "meta-llama/Llama-3.1-8B-Instruct")
|
||||
self.assertEqual(adapter, "adapter-name")
|
||||
|
||||
def test_model_with_multiple_colons(self):
|
||||
"""Test that only first colon is used for splitting."""
|
||||
base_model, adapter = self.serving._parse_model_parameter("model:adapter:extra")
|
||||
self.assertEqual(base_model, "model")
|
||||
self.assertEqual(adapter, "adapter:extra")
|
||||
|
||||
def test_model_with_whitespace(self):
|
||||
"""Test that whitespace is stripped from both parts."""
|
||||
base_model, adapter = self.serving._parse_model_parameter(
|
||||
" model-name : adapter-name "
|
||||
)
|
||||
self.assertEqual(base_model, "model-name")
|
||||
self.assertEqual(adapter, "adapter-name")
|
||||
|
||||
def test_model_with_empty_adapter(self):
|
||||
"""Test model ending with colon returns None for adapter."""
|
||||
base_model, adapter = self.serving._parse_model_parameter("model-name:")
|
||||
self.assertEqual(base_model, "model-name")
|
||||
self.assertIsNone(adapter)
|
||||
|
||||
def test_model_with_only_spaces_after_colon(self):
|
||||
"""Test model with only whitespace after colon returns None for adapter."""
|
||||
base_model, adapter = self.serving._parse_model_parameter("model-name: ")
|
||||
self.assertEqual(base_model, "model-name")
|
||||
self.assertIsNone(adapter)
|
||||
|
||||
|
||||
class TestResolveLoraPath(unittest.TestCase):
|
||||
"""Test _resolve_lora_path method."""
|
||||
|
||||
def setUp(self):
|
||||
self.tokenizer_manager = MockTokenizerManager(enable_lora=True)
|
||||
self.serving = ConcreteServingBase(self.tokenizer_manager)
|
||||
|
||||
def test_no_adapter_specified(self):
|
||||
"""Test when neither model nor explicit lora_path has adapter."""
|
||||
result = self.serving._resolve_lora_path("model-name", None)
|
||||
self.assertIsNone(result)
|
||||
|
||||
def test_adapter_in_model_only(self):
|
||||
"""Test adapter from model parameter when no explicit path."""
|
||||
result = self.serving._resolve_lora_path("model:sql-expert", None)
|
||||
self.assertEqual(result, "sql-expert")
|
||||
|
||||
def test_adapter_in_explicit_only(self):
|
||||
"""Test adapter from explicit lora_path when not in model."""
|
||||
result = self.serving._resolve_lora_path("model-name", "python-expert")
|
||||
self.assertEqual(result, "python-expert")
|
||||
|
||||
def test_model_parameter_takes_precedence(self):
|
||||
"""Test model parameter adapter takes precedence over explicit."""
|
||||
result = self.serving._resolve_lora_path("model:sql-expert", "python-expert")
|
||||
self.assertEqual(result, "sql-expert")
|
||||
|
||||
def test_with_list_explicit_lora_path(self):
|
||||
"""Test that explicit list is returned when no model adapter."""
|
||||
explicit = ["adapter1", "adapter2", None]
|
||||
result = self.serving._resolve_lora_path("model-name", explicit)
|
||||
self.assertEqual(result, explicit)
|
||||
|
||||
def test_model_adapter_overrides_list(self):
|
||||
"""Test model adapter overrides even when explicit is a list."""
|
||||
result = self.serving._resolve_lora_path(
|
||||
"model:sql-expert", ["adapter1", "adapter2"]
|
||||
)
|
||||
self.assertEqual(result, "sql-expert")
|
||||
|
||||
def test_complex_model_name_with_adapter(self):
|
||||
"""Test resolution with complex model name."""
|
||||
result = self.serving._resolve_lora_path(
|
||||
"org/model-v2.1:adapter-name", "other-adapter"
|
||||
)
|
||||
self.assertEqual(result, "adapter-name")
|
||||
|
||||
|
||||
class TestValidateLoraEnabled(unittest.TestCase):
|
||||
"""Test _validate_lora_enabled method."""
|
||||
|
||||
def test_validation_passes_when_lora_enabled(self):
|
||||
"""Test validation passes when LoRA is enabled."""
|
||||
tokenizer_manager = MockTokenizerManager(enable_lora=True)
|
||||
serving = ConcreteServingBase(tokenizer_manager)
|
||||
|
||||
# Should not raise
|
||||
try:
|
||||
serving._validate_lora_enabled("sql-expert")
|
||||
except ValueError:
|
||||
self.fail("_validate_lora_enabled raised ValueError unexpectedly")
|
||||
|
||||
def test_validation_fails_when_lora_disabled(self):
|
||||
"""Test validation fails with helpful message when LoRA is disabled."""
|
||||
tokenizer_manager = MockTokenizerManager(enable_lora=False)
|
||||
serving = ConcreteServingBase(tokenizer_manager)
|
||||
|
||||
with self.assertRaises(ValueError) as context:
|
||||
serving._validate_lora_enabled("sql-expert")
|
||||
|
||||
error_message = str(context.exception)
|
||||
self.assertIn("sql-expert", error_message)
|
||||
self.assertIn("--enable-lora", error_message)
|
||||
self.assertIn("not enabled", error_message)
|
||||
|
||||
def test_validation_error_mentions_adapter_name(self):
|
||||
"""Test that error message includes the requested adapter name."""
|
||||
tokenizer_manager = MockTokenizerManager(enable_lora=False)
|
||||
serving = ConcreteServingBase(tokenizer_manager)
|
||||
|
||||
with self.assertRaises(ValueError) as context:
|
||||
serving._validate_lora_enabled("my-custom-adapter")
|
||||
|
||||
self.assertIn("my-custom-adapter", str(context.exception))
|
||||
|
||||
|
||||
class TestIntegrationScenarios(unittest.TestCase):
|
||||
"""Integration tests for common usage scenarios."""
|
||||
|
||||
def setUp(self):
|
||||
self.tokenizer_manager = MockTokenizerManager(enable_lora=True)
|
||||
self.serving = ConcreteServingBase(self.tokenizer_manager)
|
||||
|
||||
def test_openai_compatible_usage(self):
|
||||
"""Test typical OpenAI-compatible usage pattern."""
|
||||
# User specifies adapter in model parameter
|
||||
model = "meta-llama/Llama-3.1-8B:sql-expert"
|
||||
explicit_lora = None
|
||||
|
||||
lora_path = self.serving._resolve_lora_path(model, explicit_lora)
|
||||
self.assertEqual(lora_path, "sql-expert")
|
||||
|
||||
# Validation should pass
|
||||
self.serving._validate_lora_enabled(lora_path)
|
||||
|
||||
def test_backward_compatible_usage(self):
|
||||
"""Test backward-compatible usage with explicit lora_path."""
|
||||
model = "meta-llama/Llama-3.1-8B"
|
||||
explicit_lora = "sql-expert"
|
||||
|
||||
lora_path = self.serving._resolve_lora_path(model, explicit_lora)
|
||||
self.assertEqual(lora_path, "sql-expert")
|
||||
|
||||
# Validation should pass
|
||||
self.serving._validate_lora_enabled(lora_path)
|
||||
|
||||
def test_base_model_usage(self):
|
||||
"""Test using base model without any adapter."""
|
||||
model = "meta-llama/Llama-3.1-8B"
|
||||
explicit_lora = None
|
||||
|
||||
lora_path = self.serving._resolve_lora_path(model, explicit_lora)
|
||||
self.assertIsNone(lora_path)
|
||||
|
||||
# No validation needed when no adapter
|
||||
|
||||
def test_batch_request_scenario(self):
|
||||
"""Test batch request with list of adapters."""
|
||||
model = "meta-llama/Llama-3.1-8B" # No adapter in model
|
||||
explicit_lora = ["sql-expert", "python-expert", None]
|
||||
|
||||
lora_path = self.serving._resolve_lora_path(model, explicit_lora)
|
||||
self.assertEqual(lora_path, explicit_lora)
|
||||
|
||||
# Validate first adapter in list
|
||||
if isinstance(lora_path, list) and lora_path[0]:
|
||||
self.serving._validate_lora_enabled(lora_path[0])
|
||||
|
||||
def test_adapter_in_model_overrides_batch_list(self):
|
||||
"""Test that adapter in model parameter overrides batch list."""
|
||||
model = "meta-llama/Llama-3.1-8B:preferred-adapter"
|
||||
explicit_lora = ["adapter1", "adapter2"]
|
||||
|
||||
lora_path = self.serving._resolve_lora_path(model, explicit_lora)
|
||||
self.assertEqual(lora_path, "preferred-adapter")
|
||||
|
||||
def test_error_when_lora_not_enabled(self):
|
||||
"""Test comprehensive error flow when LoRA is not enabled."""
|
||||
# Setup server without LoRA enabled
|
||||
tokenizer_manager = MockTokenizerManager(enable_lora=False)
|
||||
serving = ConcreteServingBase(tokenizer_manager)
|
||||
|
||||
# User tries to use adapter
|
||||
model = "meta-llama/Llama-3.1-8B:sql-expert"
|
||||
lora_path = serving._resolve_lora_path(model, None)
|
||||
|
||||
# Should get helpful error
|
||||
with self.assertRaises(ValueError) as context:
|
||||
serving._validate_lora_enabled(lora_path)
|
||||
|
||||
error = str(context.exception)
|
||||
self.assertIn("--enable-lora", error)
|
||||
self.assertIn("sql-expert", error)
|
||||
|
||||
|
||||
class TestEdgeCases(unittest.TestCase):
|
||||
"""Test edge cases and error conditions."""
|
||||
|
||||
def setUp(self):
|
||||
self.tokenizer_manager = MockTokenizerManager(enable_lora=True)
|
||||
self.serving = ConcreteServingBase(self.tokenizer_manager)
|
||||
|
||||
def test_empty_string_model(self):
|
||||
"""Test handling of empty string model."""
|
||||
base, adapter = self.serving._parse_model_parameter("")
|
||||
self.assertEqual(base, "")
|
||||
self.assertIsNone(adapter)
|
||||
|
||||
def test_only_colon(self):
|
||||
"""Test model parameter that is just a colon."""
|
||||
base, adapter = self.serving._parse_model_parameter(":")
|
||||
self.assertEqual(base, "")
|
||||
self.assertIsNone(adapter)
|
||||
|
||||
def test_empty_list_lora_path(self):
|
||||
"""Test validation with empty list doesn't crash."""
|
||||
lora_path = self.serving._resolve_lora_path("model-name", [])
|
||||
# Empty list is falsy, so validation won't be called
|
||||
self.assertEqual(lora_path, [])
|
||||
|
||||
def test_list_with_none_first(self):
|
||||
"""Test validation finds first non-None adapter in list."""
|
||||
lora_path = self.serving._resolve_lora_path("model-name", [None, "adapter2"])
|
||||
self.assertEqual(lora_path, [None, "adapter2"])
|
||||
# In actual usage, validation would find "adapter2"
|
||||
|
||||
def test_list_all_none(self):
|
||||
"""Test validation with list of all None values."""
|
||||
lora_path = self.serving._resolve_lora_path("model-name", [None, None])
|
||||
self.assertEqual(lora_path, [None, None])
|
||||
# In actual usage, no validation would occur (no non-None adapters)
|
||||
|
||||
def test_unicode_in_adapter_name(self):
|
||||
"""Test Unicode characters in adapter name."""
|
||||
base, adapter = self.serving._parse_model_parameter("model:adapter-名前")
|
||||
self.assertEqual(base, "model")
|
||||
self.assertEqual(adapter, "adapter-名前")
|
||||
|
||||
def test_special_characters_in_adapter(self):
|
||||
"""Test special characters in adapter name."""
|
||||
base, adapter = self.serving._parse_model_parameter("model:adapter_v2.1-final")
|
||||
self.assertEqual(base, "model")
|
||||
self.assertEqual(adapter, "adapter_v2.1-final")
|
||||
|
||||
def test_none_as_explicit_lora_path(self):
|
||||
"""Test None as explicit lora_path is handled correctly."""
|
||||
result = self.serving._resolve_lora_path("model:adapter", None)
|
||||
self.assertEqual(result, "adapter")
|
||||
|
||||
def test_empty_string_as_explicit_lora_path(self):
|
||||
"""Test empty string as explicit lora_path."""
|
||||
result = self.serving._resolve_lora_path("model-name", "")
|
||||
self.assertEqual(result, "")
|
||||
|
||||
def test_validation_with_empty_adapter_name(self):
|
||||
"""Test validation with empty adapter name still raises error."""
|
||||
tokenizer_manager = MockTokenizerManager(enable_lora=False)
|
||||
serving = ConcreteServingBase(tokenizer_manager)
|
||||
|
||||
with self.assertRaises(ValueError):
|
||||
serving._validate_lora_enabled("")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,51 +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.
|
||||
# ==============================================================================
|
||||
|
||||
import multiprocessing as mp
|
||||
import unittest
|
||||
|
||||
from utils import LoRAAdaptor, LoRAModelCase, run_lora_multiple_batch_on_model_cases
|
||||
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
LORA_MODELS_QWEN3 = [
|
||||
LoRAModelCase(
|
||||
base="Qwen/Qwen3-4B",
|
||||
adaptors=[
|
||||
LoRAAdaptor(
|
||||
name="nissenj/Qwen3-4B-lora-v2",
|
||||
prefill_tolerance=3e-1,
|
||||
),
|
||||
LoRAAdaptor(
|
||||
name="y9760210/Qwen3-4B-lora_model",
|
||||
prefill_tolerance=3e-1,
|
||||
),
|
||||
],
|
||||
max_loras_per_batch=2,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class TestLoRAQwen3(CustomTestCase):
|
||||
def test_ci_lora_models(self):
|
||||
run_lora_multiple_batch_on_model_cases(LORA_MODELS_QWEN3)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
mp.set_start_method("spawn")
|
||||
except RuntimeError:
|
||||
pass
|
||||
|
||||
unittest.main(warnings="ignore")
|
||||
@@ -1,78 +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 utils import CI_MULTI_LORA_MODELS, run_lora_test_one_by_one
|
||||
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
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")
|
||||
@@ -1,305 +0,0 @@
|
||||
"""Utilities for running nightly performance benchmarks with profiling."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
from sglang.bench_one_batch_server import BenchmarkResult, generate_markdown_report
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
is_in_ci,
|
||||
popen_launch_server,
|
||||
write_github_step_summary,
|
||||
)
|
||||
|
||||
|
||||
class NightlyBenchmarkRunner:
|
||||
"""Helper class for running nightly performance benchmarks with profiling.
|
||||
|
||||
This class encapsulates common patterns used across nightly performance tests,
|
||||
including profile directory management, benchmark command construction,
|
||||
result parsing, and report generation.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
profile_dir: str,
|
||||
test_name: str,
|
||||
base_url: str,
|
||||
gpu_config: str = None,
|
||||
):
|
||||
"""Initialize the benchmark runner.
|
||||
|
||||
Args:
|
||||
profile_dir: Directory to store performance profiles
|
||||
test_name: Name of the test (used for reporting)
|
||||
base_url: Base URL for the server
|
||||
gpu_config: Optional GPU configuration string (e.g., "2-gpu-h100", "8-gpu-b200")
|
||||
"""
|
||||
self.profile_dir = profile_dir
|
||||
self.test_name = test_name
|
||||
self.base_url = base_url
|
||||
self.gpu_config = gpu_config or os.environ.get("GPU_CONFIG", "")
|
||||
|
||||
# Include GPU config in report header if available
|
||||
header = f"## {test_name}"
|
||||
if self.gpu_config:
|
||||
header += f" ({self.gpu_config})"
|
||||
header += "\n"
|
||||
self.full_report = header + BenchmarkResult.help_str()
|
||||
|
||||
def setup_profile_directory(self) -> None:
|
||||
"""Create the profile directory if it doesn't exist."""
|
||||
os.makedirs(self.profile_dir, exist_ok=True)
|
||||
|
||||
def generate_profile_filename(
|
||||
self, model_path: str, variant: str = ""
|
||||
) -> Tuple[str, str]:
|
||||
"""Generate unique profile filename and path for the model.
|
||||
|
||||
Args:
|
||||
model_path: Path to the model (e.g., "deepseek-ai/DeepSeek-V3.1")
|
||||
variant: Optional variant suffix (e.g., "basic", "mtp", "nsa")
|
||||
|
||||
Returns:
|
||||
Tuple of (profile_path_prefix, json_output_file)
|
||||
"""
|
||||
timestamp = int(time.time())
|
||||
model_safe_name = model_path.replace("/", "_")
|
||||
|
||||
# Build filename with optional variant
|
||||
if variant:
|
||||
profile_filename = f"{model_safe_name}_{variant}_{timestamp}"
|
||||
json_filename = f"results_{model_safe_name}_{variant}_{timestamp}.json"
|
||||
else:
|
||||
profile_filename = f"{model_safe_name}_{timestamp}"
|
||||
json_filename = f"results_{model_safe_name}_{timestamp}.json"
|
||||
|
||||
profile_path_prefix = os.path.join(self.profile_dir, profile_filename)
|
||||
|
||||
return profile_path_prefix, json_filename
|
||||
|
||||
def build_benchmark_command(
|
||||
self,
|
||||
model_path: str,
|
||||
batch_sizes: List[int],
|
||||
input_lens: Tuple[int, ...],
|
||||
output_lens: Tuple[int, ...],
|
||||
profile_path_prefix: str,
|
||||
json_output_file: str,
|
||||
extra_args: Optional[List[str]] = None,
|
||||
) -> List[str]:
|
||||
"""Build the benchmark command with all required arguments.
|
||||
|
||||
Args:
|
||||
model_path: Path to the model
|
||||
batch_sizes: List of batch sizes to test
|
||||
input_lens: Tuple of input lengths to test
|
||||
output_lens: Tuple of output lengths to test
|
||||
profile_path_prefix: Prefix for profile output files
|
||||
json_output_file: Path to JSON output file
|
||||
extra_args: Optional extra arguments to append to command
|
||||
|
||||
Returns:
|
||||
List of command arguments ready for subprocess.run()
|
||||
"""
|
||||
command = [
|
||||
"python3",
|
||||
"-m",
|
||||
"sglang.bench_one_batch_server",
|
||||
"--model",
|
||||
model_path,
|
||||
"--base-url",
|
||||
self.base_url,
|
||||
"--batch-size",
|
||||
*[str(x) for x in batch_sizes],
|
||||
"--input-len",
|
||||
*[str(x) for x in input_lens],
|
||||
"--output-len",
|
||||
*[str(x) for x in output_lens],
|
||||
"--show-report",
|
||||
"--profile",
|
||||
"--profile-by-stage",
|
||||
"--profile-filename-prefix",
|
||||
profile_path_prefix,
|
||||
f"--output-path={json_output_file}",
|
||||
"--no-append-to-github-summary",
|
||||
]
|
||||
|
||||
if extra_args:
|
||||
command.extend(extra_args)
|
||||
|
||||
return command
|
||||
|
||||
def run_benchmark_command(
|
||||
self, command: List[str], model_description: str = ""
|
||||
) -> Tuple[subprocess.CompletedProcess, bool]:
|
||||
"""Execute the benchmark command and return the result.
|
||||
|
||||
Args:
|
||||
command: Command to execute
|
||||
model_description: Description for logging (e.g., "model_name (variant)")
|
||||
|
||||
Returns:
|
||||
Tuple of (CompletedProcess, success_bool)
|
||||
"""
|
||||
print(f"Running command: {' '.join(command)}")
|
||||
result = subprocess.run(command, capture_output=True, text=True)
|
||||
|
||||
if result.returncode != 0:
|
||||
desc = model_description or "benchmark"
|
||||
print(f"Error running benchmark for {desc}:")
|
||||
print(result.stderr)
|
||||
return result, False
|
||||
|
||||
return result, True
|
||||
|
||||
def load_benchmark_results(
|
||||
self, json_output_file: str, model_description: str = ""
|
||||
) -> Tuple[List[BenchmarkResult], bool]:
|
||||
"""Load and parse benchmark results from JSON file.
|
||||
|
||||
Args:
|
||||
json_output_file: Path to JSON output file
|
||||
model_description: Description for logging
|
||||
|
||||
Returns:
|
||||
Tuple of (list of BenchmarkResult objects, success_bool)
|
||||
"""
|
||||
benchmark_results = []
|
||||
|
||||
if not os.path.exists(json_output_file):
|
||||
desc = model_description or "model"
|
||||
print(f"Warning: JSON output file {json_output_file} not found for {desc}")
|
||||
return benchmark_results, False
|
||||
|
||||
try:
|
||||
with open(json_output_file, "r") as f:
|
||||
json_data = json.load(f)
|
||||
|
||||
# Convert JSON data to BenchmarkResult objects
|
||||
for data in json_data:
|
||||
benchmark_result = BenchmarkResult(**data)
|
||||
benchmark_results.append(benchmark_result)
|
||||
|
||||
print(
|
||||
f"Loaded {len(benchmark_results)} benchmark results from {json_output_file}"
|
||||
)
|
||||
|
||||
# Clean up JSON file
|
||||
os.remove(json_output_file)
|
||||
|
||||
return benchmark_results, True
|
||||
|
||||
except Exception as e:
|
||||
desc = model_description or "model"
|
||||
print(f"Error loading benchmark results for {desc}: {e}")
|
||||
# Try to clean up the file anyway
|
||||
if os.path.exists(json_output_file):
|
||||
os.remove(json_output_file)
|
||||
return benchmark_results, False
|
||||
|
||||
def run_benchmark_for_model(
|
||||
self,
|
||||
model_path: str,
|
||||
batch_sizes: List[int],
|
||||
input_lens: Tuple[int, ...],
|
||||
output_lens: Tuple[int, ...],
|
||||
other_args: Optional[List[str]] = None,
|
||||
variant: str = "",
|
||||
extra_bench_args: Optional[List[str]] = None,
|
||||
) -> Tuple[List[BenchmarkResult], bool]:
|
||||
"""Run a complete benchmark for a single model with server management.
|
||||
|
||||
This method handles:
|
||||
- Server launch and cleanup
|
||||
- Profile filename generation
|
||||
- Benchmark command construction and execution
|
||||
- Result loading and parsing
|
||||
|
||||
Args:
|
||||
model_path: Path to the model
|
||||
batch_sizes: List of batch sizes to test
|
||||
input_lens: Tuple of input lengths
|
||||
output_lens: Tuple of output lengths
|
||||
other_args: Arguments to pass to server launch
|
||||
variant: Optional variant suffix (e.g., "basic", "mtp")
|
||||
extra_bench_args: Extra arguments for the benchmark command
|
||||
|
||||
Returns:
|
||||
Tuple of (list of BenchmarkResult objects, success_bool)
|
||||
"""
|
||||
benchmark_results = []
|
||||
model_description = f"{model_path}" + (f" ({variant})" if variant else "")
|
||||
|
||||
# Launch server
|
||||
process = popen_launch_server(
|
||||
model=model_path,
|
||||
base_url=self.base_url,
|
||||
other_args=other_args or [],
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
)
|
||||
|
||||
try:
|
||||
# Generate filenames
|
||||
profile_path_prefix, json_output_file = self.generate_profile_filename(
|
||||
model_path, variant
|
||||
)
|
||||
|
||||
# Build and run benchmark command
|
||||
# Prepare extra args with run_name if variant is specified
|
||||
bench_args = list(extra_bench_args) if extra_bench_args else []
|
||||
if variant:
|
||||
bench_args.extend(["--run-name", variant])
|
||||
|
||||
command = self.build_benchmark_command(
|
||||
model_path,
|
||||
batch_sizes,
|
||||
input_lens,
|
||||
output_lens,
|
||||
profile_path_prefix,
|
||||
json_output_file,
|
||||
extra_args=bench_args,
|
||||
)
|
||||
|
||||
result, cmd_success = self.run_benchmark_command(command, model_description)
|
||||
|
||||
if not cmd_success:
|
||||
return benchmark_results, False
|
||||
|
||||
# Load results
|
||||
benchmark_results, load_success = self.load_benchmark_results(
|
||||
json_output_file, model_description
|
||||
)
|
||||
|
||||
return benchmark_results, load_success
|
||||
|
||||
finally:
|
||||
# Always clean up server process
|
||||
kill_process_tree(process.pid)
|
||||
|
||||
def add_report(self, results: List[BenchmarkResult]) -> None:
|
||||
"""Add benchmark results to the full report.
|
||||
|
||||
Args:
|
||||
results: List of BenchmarkResult objects to add to report
|
||||
"""
|
||||
if results:
|
||||
report_part = generate_markdown_report(self.profile_dir, results)
|
||||
self.full_report += report_part + "\n"
|
||||
|
||||
def write_final_report(self) -> None:
|
||||
"""Write the final report to GitHub summary if in CI."""
|
||||
if is_in_ci():
|
||||
write_github_step_summary(self.full_report)
|
||||
|
||||
def get_full_report(self) -> str:
|
||||
"""Get the accumulated full report.
|
||||
|
||||
Returns:
|
||||
The full markdown report as a string
|
||||
"""
|
||||
return self.full_report
|
||||
@@ -1,270 +0,0 @@
|
||||
import argparse
|
||||
import glob
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import subprocess
|
||||
import sys
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
is_in_ci,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
MODELS = [
|
||||
SimpleNamespace(model="Qwen/Qwen2.5-VL-72B-Instruct", mmmu_accuracy=0.55),
|
||||
]
|
||||
|
||||
|
||||
# Set default mem_fraction_static to 0.8
|
||||
DEFAULT_MEM_FRACTION_STATIC = 0.8
|
||||
|
||||
|
||||
class TestVLMEncoderDP(CustomTestCase):
|
||||
parsed_args = None # Class variable to store args
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
# Removed argument parsing from here
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.api_key = "sk-123456"
|
||||
cls.time_out = DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
|
||||
|
||||
if cls.parsed_args is None:
|
||||
cls.parsed_args = SimpleNamespace(
|
||||
mem_fraction_static=DEFAULT_MEM_FRACTION_STATIC
|
||||
)
|
||||
|
||||
# Set OpenAI API key and base URL environment variables. Needed for lmm-evals to work.
|
||||
os.environ["OPENAI_API_KEY"] = cls.api_key
|
||||
os.environ["OPENAI_API_BASE"] = f"{cls.base_url}/v1"
|
||||
|
||||
def run_mmmu_eval(
|
||||
self,
|
||||
model_version: str,
|
||||
output_path: str,
|
||||
*,
|
||||
env: dict | None = None,
|
||||
):
|
||||
"""
|
||||
Evaluate a VLM on the MMMU validation set with lmms‑eval.
|
||||
Only `model_version` (checkpoint) and `chat_template` vary;
|
||||
We are focusing only on the validation set due to resource constraints.
|
||||
"""
|
||||
# -------- fixed settings --------
|
||||
model = "openai_compatible"
|
||||
tp = 1
|
||||
tasks = "mmmu_val"
|
||||
batch_size = 32
|
||||
log_suffix = "openai_compatible"
|
||||
os.makedirs(output_path, exist_ok=True)
|
||||
|
||||
# -------- compose --model_args --------
|
||||
model_args = f'model_version="{model_version}",' f"tp={tp}"
|
||||
|
||||
# -------- build command list --------
|
||||
cmd = [
|
||||
"python3",
|
||||
"-m",
|
||||
"lmms_eval",
|
||||
"--model",
|
||||
model,
|
||||
"--model_args",
|
||||
model_args,
|
||||
"--tasks",
|
||||
tasks,
|
||||
"--batch_size",
|
||||
str(batch_size),
|
||||
"--log_samples",
|
||||
"--log_samples_suffix",
|
||||
log_suffix,
|
||||
"--output_path",
|
||||
str(output_path),
|
||||
]
|
||||
|
||||
subprocess.run(
|
||||
cmd,
|
||||
check=True,
|
||||
timeout=3600,
|
||||
)
|
||||
|
||||
def _run_vlm_mmmu_test(
|
||||
self,
|
||||
model,
|
||||
output_path,
|
||||
test_name="",
|
||||
custom_env=None,
|
||||
log_level="info",
|
||||
capture_output=False,
|
||||
):
|
||||
"""
|
||||
Common method to run VLM MMMU benchmark test.
|
||||
|
||||
Args:
|
||||
model: Model to test
|
||||
output_path: Path for output logs
|
||||
test_name: Optional test name for logging
|
||||
custom_env: Optional custom environment variables
|
||||
log_level: Log level for server (default: "info")
|
||||
capture_output: Whether to capture server stdout/stderr
|
||||
"""
|
||||
print(f"\nTesting model: {model.model}{test_name}")
|
||||
|
||||
process = None
|
||||
mmmu_accuracy = 0 # Initialize to handle potential exceptions
|
||||
server_output = ""
|
||||
|
||||
try:
|
||||
# Prepare environment variables
|
||||
process_env = os.environ.copy()
|
||||
if custom_env:
|
||||
process_env.update(custom_env)
|
||||
# if test vlm with cuda_ipc feature, open this env_var
|
||||
process_env["SGLANG_USE_CUDA_IPC_TRANSPORT"] = "1"
|
||||
|
||||
# Prepare stdout/stderr redirection if needed
|
||||
stdout_file = None
|
||||
stderr_file = None
|
||||
if capture_output:
|
||||
stdout_file = open("/tmp/server_stdout.log", "w")
|
||||
stderr_file = open("/tmp/server_stderr.log", "w")
|
||||
|
||||
# Launch server for testing
|
||||
process = popen_launch_server(
|
||||
model.model,
|
||||
base_url=self.base_url,
|
||||
timeout=self.time_out,
|
||||
api_key=self.api_key,
|
||||
other_args=[
|
||||
"--trust-remote-code",
|
||||
"--cuda-graph-max-bs",
|
||||
"32",
|
||||
"--mm-enable-dp-encoder",
|
||||
"--tp=4",
|
||||
"--mem-fraction-static",
|
||||
str(self.parsed_args.mem_fraction_static), # Use class variable
|
||||
"--log-level",
|
||||
log_level,
|
||||
],
|
||||
env=process_env,
|
||||
return_stdout_stderr=(
|
||||
(stdout_file, stderr_file) if capture_output else None
|
||||
),
|
||||
)
|
||||
|
||||
# Run evaluation
|
||||
self.run_mmmu_eval(model.model, output_path)
|
||||
|
||||
# Get the result file
|
||||
# Search recursively for JSON result files (lmms-eval v0.4.1+ creates subdirectories)
|
||||
result_files = glob.glob(f"{output_path}/**/*.json", recursive=True)
|
||||
if not result_files:
|
||||
result_files = glob.glob(f"{output_path}/*.json")
|
||||
|
||||
if not result_files:
|
||||
raise FileNotFoundError(f"No JSON result files found in {output_path}")
|
||||
|
||||
result_file_path = result_files[0]
|
||||
|
||||
with open(result_file_path, "r") as f:
|
||||
result = json.load(f)
|
||||
print(f"Result{test_name}\n: {result}")
|
||||
|
||||
# Process the result
|
||||
mmmu_accuracy = result["results"]["mmmu_val"]["mmmu_acc,none"]
|
||||
print(
|
||||
f"Model {model.model} achieved accuracy{test_name}: {mmmu_accuracy:.4f}"
|
||||
)
|
||||
|
||||
# Capture server output if requested
|
||||
if capture_output and process:
|
||||
server_output = self._read_output_from_files()
|
||||
|
||||
# Assert performance meets expected threshold
|
||||
self.assertGreaterEqual(
|
||||
mmmu_accuracy,
|
||||
model.mmmu_accuracy,
|
||||
f"Model {model.model} accuracy ({mmmu_accuracy:.4f}) below expected threshold ({model.mmmu_accuracy:.4f}){test_name}",
|
||||
)
|
||||
|
||||
return server_output
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error testing {model.model}{test_name}: {e}")
|
||||
self.fail(f"Test failed for {model.model}{test_name}: {e}")
|
||||
|
||||
finally:
|
||||
# Ensure process cleanup happens regardless of success/failure
|
||||
if process is not None and process.poll() is None:
|
||||
print(f"Cleaning up process {process.pid}")
|
||||
try:
|
||||
kill_process_tree(process.pid)
|
||||
except Exception as e:
|
||||
print(f"Error killing process: {e}")
|
||||
|
||||
# clean up temporary files
|
||||
if capture_output:
|
||||
if stdout_file:
|
||||
stdout_file.close()
|
||||
if stderr_file:
|
||||
stderr_file.close()
|
||||
for filename in ["/tmp/server_stdout.log", "/tmp/server_stderr.log"]:
|
||||
try:
|
||||
if os.path.exists(filename):
|
||||
os.remove(filename)
|
||||
except Exception as e:
|
||||
print(f"Error removing {filename}: {e}")
|
||||
|
||||
def _read_output_from_files(self):
|
||||
output_lines = []
|
||||
|
||||
log_files = [
|
||||
("/tmp/server_stdout.log", "[STDOUT]"),
|
||||
("/tmp/server_stderr.log", "[STDERR]"),
|
||||
]
|
||||
for filename, tag in log_files:
|
||||
try:
|
||||
if os.path.exists(filename):
|
||||
with open(filename, "r") as f:
|
||||
for line in f:
|
||||
output_lines.append(f"{tag} {line.rstrip()}")
|
||||
except Exception as e:
|
||||
print(f"Error reading {tag.lower()} file: {e}")
|
||||
|
||||
return "\n".join(output_lines)
|
||||
|
||||
def test_vlm_mmmu_benchmark(self):
|
||||
"""Test VLM models against MMMU benchmark."""
|
||||
models_to_test = MODELS
|
||||
|
||||
if is_in_ci():
|
||||
models_to_test = [random.choice(MODELS)]
|
||||
|
||||
for model in models_to_test:
|
||||
self._run_vlm_mmmu_test(model, "./logs")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Define and parse arguments here, before unittest.main
|
||||
parser = argparse.ArgumentParser(description="Test VLM models")
|
||||
parser.add_argument(
|
||||
"--mem-fraction-static",
|
||||
type=float,
|
||||
help="Static memory fraction for the model",
|
||||
default=DEFAULT_MEM_FRACTION_STATIC,
|
||||
)
|
||||
|
||||
# Parse args intended for unittest
|
||||
args = parser.parse_args()
|
||||
|
||||
# Store the parsed args object on the class
|
||||
TestVLMEncoderDP.parsed_args = args
|
||||
|
||||
# Pass args to unittest
|
||||
unittest.main(argv=[sys.argv[0]])
|
||||
@@ -1,62 +0,0 @@
|
||||
import os
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.few_shot_gsm8k import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestFlashinferTrtllmGenAttnBackend(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = "Qwen/Qwen3-Next-80B-A3B-Instruct"
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
env={**os.environ, "SGLANG_ENABLE_JIT_DEEPGEMM": "False"},
|
||||
other_args=[
|
||||
"--attention-backend",
|
||||
"trtllm_mha",
|
||||
"--cuda-graph-max-bs",
|
||||
"512",
|
||||
"--tp-size",
|
||||
"4",
|
||||
"--ep-size",
|
||||
"4",
|
||||
"--mem-fraction-static",
|
||||
"0.7",
|
||||
"--mamba-ssm-dtype",
|
||||
"bfloat16",
|
||||
"--disable-radix-cache",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
self.assertGreater(metrics["accuracy"], 0.93)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,65 +0,0 @@
|
||||
import os
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.few_shot_gsm8k import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestFlashinferTrtllmGenMoeBackend(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8"
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
env={**os.environ, "SGLANG_ENABLE_JIT_DEEPGEMM": "False"},
|
||||
other_args=[
|
||||
"--attention-backend",
|
||||
"triton",
|
||||
"--moe-runner-backend",
|
||||
"flashinfer_trtllm",
|
||||
"--cuda-graph-max-bs",
|
||||
"512",
|
||||
"--tp-size",
|
||||
"4",
|
||||
"--ep-size",
|
||||
"4",
|
||||
"--mem-fraction-static",
|
||||
"0.7",
|
||||
"--mamba-ssm-dtype",
|
||||
"bfloat16",
|
||||
"--quantization",
|
||||
"fp8",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
self.assertGreater(metrics["accuracy"], 0.93)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,58 +0,0 @@
|
||||
import unittest
|
||||
|
||||
from nightly_utils import NightlyBenchmarkRunner
|
||||
|
||||
from sglang.test.test_utils import DEFAULT_URL_FOR_TEST
|
||||
|
||||
PROFILE_DIR = "performance_profiles_gpt_oss_4gpu"
|
||||
|
||||
|
||||
class TestNightlyGptOss4GpuPerformance(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.models = [
|
||||
(
|
||||
"openai/gpt-oss-120b",
|
||||
[
|
||||
"--tp",
|
||||
"4",
|
||||
"--cuda-graph-max-bs",
|
||||
"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(PROFILE_DIR, cls.__name__, cls.base_url)
|
||||
cls.runner.setup_profile_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()
|
||||
@@ -1,278 +0,0 @@
|
||||
"""
|
||||
End-to-end tests for OpenAI-compatible LoRA adapter usage.
|
||||
|
||||
Tests the model:adapter syntax and backward compatibility with explicit lora_path.
|
||||
|
||||
Usage:
|
||||
python3 -m unittest openai_server.features.test_lora_openai_compatible.TestLoRAOpenAICompatible.test_model_adapter_syntax
|
||||
python3 -m unittest openai_server.features.test_lora_openai_compatible.TestLoRAOpenAICompatible.test_explicit_lora_path
|
||||
python3 -m unittest openai_server.features.test_lora_openai_compatible.TestLoRAOpenAICompatible.test_priority_model_over_explicit
|
||||
python3 -m unittest openai_server.features.test_lora_openai_compatible.TestLoRAOpenAICompatible.test_base_model_no_adapter
|
||||
python3 -m unittest openai_server.features.test_lora_openai_compatible.TestLoRAOpenAICompatible.test_completions_api_with_adapter
|
||||
python3 -m unittest openai_server.features.test_lora_openai_compatible.TestLoRAOpenAICompatible.test_streaming_with_adapter
|
||||
python3 -m unittest openai_server.features.test_lora_openai_compatible.TestLoRADisabledError.test_lora_disabled_error
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
import openai
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
def get_real_lora_adapter() -> str:
|
||||
"""Use a real LoRA adapter from Hugging Face."""
|
||||
return "codelion/Llama-3.2-1B-Instruct-tool-calling-lora"
|
||||
|
||||
|
||||
def setup_class(cls, enable_lora=True):
|
||||
"""Setup test class with LoRA-enabled server."""
|
||||
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
|
||||
# Use real LoRA adapter
|
||||
cls.lora_adapter_path = get_real_lora_adapter()
|
||||
|
||||
other_args = [
|
||||
"--max-running-requests",
|
||||
"10",
|
||||
"--disable-radix-cache", # Disable cache for cleaner tests
|
||||
]
|
||||
|
||||
if enable_lora:
|
||||
other_args.extend(
|
||||
[
|
||||
"--enable-lora",
|
||||
"--lora-paths",
|
||||
f"tool_calling={cls.lora_adapter_path}",
|
||||
]
|
||||
)
|
||||
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=other_args,
|
||||
)
|
||||
cls.client = openai.Client(api_key="EMPTY", base_url=f"{cls.base_url}/v1")
|
||||
|
||||
|
||||
class TestLoRAOpenAICompatible(CustomTestCase):
|
||||
"""Test OpenAI-compatible LoRA adapter usage."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
setup_class(cls, enable_lora=True)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_model_adapter_syntax(self):
|
||||
"""Test the new model:adapter syntax works correctly."""
|
||||
response = self.client.chat.completions.create(
|
||||
# ← New OpenAI-compatible syntax
|
||||
model=f"{self.model}:tool_calling",
|
||||
messages=[{"role": "user", "content": "What tools do you have available?"}],
|
||||
max_tokens=50,
|
||||
temperature=0,
|
||||
)
|
||||
|
||||
self.assertIsNotNone(response.choices[0].message.content)
|
||||
self.assertGreater(len(response.choices[0].message.content), 0)
|
||||
print(f"Model adapter syntax response: {response.choices[0].message.content}")
|
||||
|
||||
def test_explicit_lora_path(self):
|
||||
"""Test backward compatibility with explicit lora_path via extra_body."""
|
||||
response = self.client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[{"role": "user", "content": "What tools do you have available?"}],
|
||||
# ← Legacy explicit method
|
||||
extra_body={"lora_path": "tool_calling"},
|
||||
max_tokens=50,
|
||||
temperature=0,
|
||||
)
|
||||
|
||||
self.assertIsNotNone(response.choices[0].message.content)
|
||||
self.assertGreater(len(response.choices[0].message.content), 0)
|
||||
print(f"Explicit lora_path response: {response.choices[0].message.content}")
|
||||
|
||||
def test_priority_model_over_explicit(self):
|
||||
"""Test that model:adapter syntax takes precedence over explicit lora_path."""
|
||||
# This test verifies the priority logic in _resolve_lora_path
|
||||
response = self.client.chat.completions.create(
|
||||
# ← Model specifies tool_calling adapter
|
||||
model=f"{self.model}:tool_calling",
|
||||
messages=[{"role": "user", "content": "What tools do you have available?"}],
|
||||
# ← Both specify same adapter
|
||||
extra_body={"lora_path": "tool_calling"},
|
||||
max_tokens=50,
|
||||
temperature=0,
|
||||
)
|
||||
|
||||
# Should use tool_calling adapter (model parameter takes precedence)
|
||||
self.assertIsNotNone(response.choices[0].message.content)
|
||||
self.assertGreater(len(response.choices[0].message.content), 0)
|
||||
print(f"Priority test response: {response.choices[0].message.content}")
|
||||
|
||||
def test_base_model_no_adapter(self):
|
||||
"""Test using base model without any adapter."""
|
||||
response = self.client.chat.completions.create(
|
||||
model=self.model, # ← No adapter specified
|
||||
messages=[{"role": "user", "content": "Hello, how are you?"}],
|
||||
max_tokens=30,
|
||||
temperature=0,
|
||||
)
|
||||
|
||||
self.assertIsNotNone(response.choices[0].message.content)
|
||||
self.assertGreater(len(response.choices[0].message.content), 0)
|
||||
print(f"Base model response: {response.choices[0].message.content}")
|
||||
|
||||
def test_completions_api_with_adapter(self):
|
||||
"""Test completions API with LoRA adapter."""
|
||||
response = self.client.completions.create(
|
||||
model=f"{self.model}:tool_calling", # ← Using model:adapter syntax
|
||||
prompt="What tools do you have available?",
|
||||
max_tokens=50,
|
||||
temperature=0,
|
||||
)
|
||||
|
||||
self.assertIsNotNone(response.choices[0].text)
|
||||
self.assertGreater(len(response.choices[0].text), 0)
|
||||
print(f"Completions API response: {response.choices[0].text}")
|
||||
|
||||
def test_streaming_with_adapter(self):
|
||||
"""Test streaming with LoRA adapter."""
|
||||
stream = self.client.chat.completions.create(
|
||||
model=f"{self.model}:tool_calling",
|
||||
messages=[{"role": "user", "content": "What tools do you have available?"}],
|
||||
max_tokens=50,
|
||||
temperature=0,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
collected_content = ""
|
||||
for chunk in stream:
|
||||
if chunk.choices[0].delta.content:
|
||||
collected_content += chunk.choices[0].delta.content
|
||||
|
||||
self.assertGreater(len(collected_content), 0)
|
||||
print(f"Streaming response: {collected_content}")
|
||||
|
||||
def test_multiple_adapters(self):
|
||||
"""Test using different adapters in sequence."""
|
||||
# Test tool_calling adapter
|
||||
tool_response = self.client.chat.completions.create(
|
||||
model=f"{self.model}:tool_calling",
|
||||
messages=[{"role": "user", "content": "What tools do you have available?"}],
|
||||
max_tokens=30,
|
||||
temperature=0,
|
||||
)
|
||||
|
||||
# Test base model without adapter
|
||||
base_response = self.client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[{"role": "user", "content": "Hello, how are you?"}],
|
||||
max_tokens=30,
|
||||
temperature=0,
|
||||
)
|
||||
|
||||
self.assertIsNotNone(tool_response.choices[0].message.content)
|
||||
self.assertIsNotNone(base_response.choices[0].message.content)
|
||||
print(
|
||||
f"Tool calling adapter response: {tool_response.choices[0].message.content}"
|
||||
)
|
||||
print(f"Base model response: {base_response.choices[0].message.content}")
|
||||
|
||||
|
||||
class TestLoRADisabledError(CustomTestCase):
|
||||
"""Test error handling when LoRA is disabled."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
setup_class(cls, enable_lora=False) # ← LoRA disabled
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_lora_disabled_error(self):
|
||||
"""Test that using LoRA adapter when LoRA is disabled raises appropriate error."""
|
||||
with self.assertRaises(openai.APIError) as context:
|
||||
self.client.chat.completions.create(
|
||||
model=f"{self.model}:tool_calling", # ← Trying to use adapter
|
||||
messages=[
|
||||
{"role": "user", "content": "What tools do you have available?"}
|
||||
],
|
||||
max_tokens=50,
|
||||
)
|
||||
|
||||
# Verify the error message contains helpful guidance
|
||||
error_message = str(context.exception)
|
||||
self.assertIn("LoRA", error_message)
|
||||
self.assertIn("not enabled", error_message)
|
||||
print(f"Expected error message: {error_message}")
|
||||
|
||||
|
||||
class TestLoRAEdgeCases(CustomTestCase):
|
||||
"""Test edge cases for LoRA adapter usage."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
setup_class(cls, enable_lora=True)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_model_with_colon_no_adapter(self):
|
||||
"""Test model parameter ending with colon (empty adapter)."""
|
||||
response = self.client.chat.completions.create(
|
||||
model=f"{self.model}:", # ← Model ends with colon
|
||||
messages=[{"role": "user", "content": "Hello!"}],
|
||||
max_tokens=30,
|
||||
temperature=0,
|
||||
)
|
||||
|
||||
# Should work as base model (no adapter)
|
||||
self.assertIsNotNone(response.choices[0].message.content)
|
||||
print(f"Model with colon response: {response.choices[0].message.content}")
|
||||
|
||||
def test_explicit_lora_path_none(self):
|
||||
"""Test explicit lora_path set to None."""
|
||||
response = self.client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[{"role": "user", "content": "Hello!"}],
|
||||
extra_body={"lora_path": None}, # ← Explicitly None
|
||||
max_tokens=30,
|
||||
temperature=0,
|
||||
)
|
||||
|
||||
# Should work as base model
|
||||
self.assertIsNotNone(response.choices[0].message.content)
|
||||
print(
|
||||
f"Explicit None lora_path response: {response.choices[0].message.content}"
|
||||
)
|
||||
|
||||
def test_invalid_adapter_name(self):
|
||||
"""Test using non-existent adapter name."""
|
||||
with self.assertRaises(openai.APIError) as context:
|
||||
self.client.chat.completions.create(
|
||||
model=f"{self.model}:nonexistent", # ← Non-existent adapter
|
||||
messages=[{"role": "user", "content": "Hello!"}],
|
||||
max_tokens=30,
|
||||
)
|
||||
|
||||
error_message = str(context.exception)
|
||||
print(f"Invalid adapter error: {error_message}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
+1
-31
@@ -198,37 +198,7 @@ suites = {
|
||||
TestFile("test_quantization.py", 185),
|
||||
TestFile("test_gguf.py", 96),
|
||||
],
|
||||
# If the test cases take too long, considering adding them to nightly tests instead of per-commit tests
|
||||
"nightly-1-gpu": [
|
||||
TestFile("layers/attention/nsa/test_nsa_indexer.py", 2),
|
||||
TestFile("lora/test_lora_qwen3.py", 97),
|
||||
TestFile("lora/test_lora_radix_cache.py", 200),
|
||||
TestFile("lora/test_lora_eviction_policy.py", 200),
|
||||
TestFile("lora/test_lora_openai_api.py", 30),
|
||||
TestFile("openai_server/features/test_lora_openai_compatible.py", 150),
|
||||
TestFile("batch_invariant/test_batch_invariant_ops.py", 10),
|
||||
TestFile("test_cpp_radix_cache.py", 60),
|
||||
TestFile("test_deepseek_v3_deterministic.py", 240),
|
||||
],
|
||||
"nightly-4-gpu-b200": [
|
||||
TestFile("nightly/test_flashinfer_trtllm_gen_moe_backend.py", 300),
|
||||
TestFile("nightly/test_gpt_oss_4gpu_perf.py", 600),
|
||||
TestFile("nightly/test_flashinfer_trtllm_gen_attn_backend.py", 300),
|
||||
TestFile("test_deepseek_v3_fp4_cutlass_moe.py", 900),
|
||||
TestFile("test_fp4_moe.py", 300),
|
||||
],
|
||||
"nightly-8-gpu-b200": [
|
||||
TestFile("test_deepseek_r1_fp8_trtllm_backend.py", 3600),
|
||||
],
|
||||
"nightly-4-gpu": [
|
||||
TestFile("nightly/test_encoder_dp.py", 500),
|
||||
TestFile("test_qwen3_next_deterministic.py", 200),
|
||||
],
|
||||
"nightly-8-gpu": [],
|
||||
"nightly-8-gpu-h200": [
|
||||
TestFile("test_deepseek_v32_nsabackend.py", 600),
|
||||
],
|
||||
"nightly-8-gpu-h20": [],
|
||||
# Nightly test suites have been moved to test/run_suite_nightly.py
|
||||
"__not_in_ci__": [
|
||||
TestFile("test_bench_one_batch.py"),
|
||||
TestFile("test_bench_serving.py"),
|
||||
|
||||
@@ -1,47 +0,0 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestCppRadixCache(CustomTestCase):
|
||||
@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)
|
||||
|
||||
def test_mmlu(self):
|
||||
args = SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=self.model,
|
||||
eval_name="mmlu",
|
||||
num_examples=64,
|
||||
num_threads=32,
|
||||
)
|
||||
|
||||
metrics = run_eval(args)
|
||||
print(metrics)
|
||||
self.assertGreaterEqual(metrics["score"], 0.65)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,88 +0,0 @@
|
||||
import os
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
try_cached_model,
|
||||
)
|
||||
|
||||
FULL_DEEPSEEK_V3_MODEL_PATH = "deepseek-ai/DeepSeek-V3-0324"
|
||||
|
||||
|
||||
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",
|
||||
"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",
|
||||
"--moe-runner-backend",
|
||||
"flashinfer_trtllm",
|
||||
"--enable-symm-mem",
|
||||
]
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=other_args,
|
||||
env={
|
||||
**os.environ,
|
||||
"SGLANG_ENABLE_FLASHINFER_FP8_GEMM": "1",
|
||||
},
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=512,
|
||||
parallel=512,
|
||||
max_new_tokens=512,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
print(f"Eval accuracy of GSM8K: {metrics=}")
|
||||
|
||||
self.assertGreater(metrics["accuracy"], 0.92)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,221 +0,0 @@
|
||||
import os
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
is_in_ci,
|
||||
popen_launch_server,
|
||||
write_github_step_summary,
|
||||
)
|
||||
|
||||
DEEPSEEK_V32_MODEL_PATH = "deepseek-ai/DeepSeek-V3.2-Exp"
|
||||
|
||||
# Global list to collect results
|
||||
TEST_RESULTS = []
|
||||
|
||||
|
||||
class TestDeepseekV32NasBackend_flashmla(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEEPSEEK_V32_MODEL_PATH
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
other_args = [
|
||||
"--trust-remote-code",
|
||||
"--attention-backend",
|
||||
"nsa",
|
||||
"--nsa-prefill-backend",
|
||||
"flashmla_sparse",
|
||||
"--nsa-decode-backend",
|
||||
"flashmla_kv",
|
||||
"--tp",
|
||||
"8",
|
||||
"--dp",
|
||||
"8",
|
||||
"--enable-dp-attention",
|
||||
]
|
||||
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,
|
||||
): # Append an "a" to make this test run first (alphabetically) to warm up the server
|
||||
args = SimpleNamespace(
|
||||
num_shots=20,
|
||||
data_path=None,
|
||||
num_questions=1400,
|
||||
parallel=1400,
|
||||
max_new_tokens=512,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
print(f"{metrics=}")
|
||||
|
||||
if is_in_ci():
|
||||
TEST_RESULTS.append(
|
||||
{
|
||||
"variant": "flashmla",
|
||||
"prefill_backend": "flashmla_sparse",
|
||||
"decode_backend": "flashmla_kv",
|
||||
"kv_cache": "fp16",
|
||||
"accuracy": metrics["accuracy"],
|
||||
}
|
||||
)
|
||||
self.assertGreater(metrics["accuracy"], 0.935)
|
||||
|
||||
|
||||
class TestDeepseekV32NasBackend_fa3(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEEPSEEK_V32_MODEL_PATH
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
other_args = [
|
||||
"--trust-remote-code",
|
||||
"--attention-backend",
|
||||
"nsa",
|
||||
"--nsa-prefill-backend",
|
||||
"fa3",
|
||||
"--nsa-decode-backend",
|
||||
"fa3",
|
||||
"--tp",
|
||||
"8",
|
||||
"--dp",
|
||||
"8",
|
||||
"--enable-dp-attention",
|
||||
]
|
||||
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,
|
||||
): # Append an "a" to make this test run first (alphabetically) to warm up the server
|
||||
args = SimpleNamespace(
|
||||
num_shots=20,
|
||||
data_path=None,
|
||||
num_questions=1400,
|
||||
parallel=1400,
|
||||
max_new_tokens=512,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
print(f"{metrics=}")
|
||||
|
||||
if is_in_ci():
|
||||
TEST_RESULTS.append(
|
||||
{
|
||||
"variant": "fa3",
|
||||
"prefill_backend": "fa3",
|
||||
"decode_backend": "fa3",
|
||||
"kv_cache": "fp16",
|
||||
"accuracy": metrics["accuracy"],
|
||||
}
|
||||
)
|
||||
self.assertGreater(metrics["accuracy"], 0.935)
|
||||
|
||||
|
||||
class TestDeepseekV32NasBackend_fp8kvcache(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEEPSEEK_V32_MODEL_PATH
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
other_args = [
|
||||
"--trust-remote-code",
|
||||
"--attention-backend",
|
||||
"nsa",
|
||||
"--kv-cache-dtype",
|
||||
"fp8_e4m3",
|
||||
"--tp",
|
||||
"8",
|
||||
"--dp",
|
||||
"8",
|
||||
"--enable-dp-attention",
|
||||
]
|
||||
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,
|
||||
): # Append an "a" to make this test run first (alphabetically) to warm up the server
|
||||
args = SimpleNamespace(
|
||||
num_shots=20,
|
||||
data_path=None,
|
||||
num_questions=1400,
|
||||
parallel=1400,
|
||||
max_new_tokens=512,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
print(f"{metrics=}")
|
||||
|
||||
if is_in_ci():
|
||||
TEST_RESULTS.append(
|
||||
{
|
||||
"variant": "fp8kvcache",
|
||||
"prefill_backend": "default",
|
||||
"decode_backend": "default",
|
||||
"kv_cache": "fp8_e4m3",
|
||||
"accuracy": metrics["accuracy"],
|
||||
}
|
||||
)
|
||||
|
||||
# Write the summary table after all tests complete
|
||||
_write_summary_table()
|
||||
self.assertGreater(metrics["accuracy"], 0.935)
|
||||
|
||||
|
||||
def _write_summary_table():
|
||||
"""Write a markdown table with all test results."""
|
||||
if not TEST_RESULTS:
|
||||
return
|
||||
|
||||
gpu_config = os.getenv("GPU_CONFIG", "8-gpu-h200")
|
||||
|
||||
# Build table header
|
||||
summary = f"### {DEEPSEEK_V32_MODEL_PATH} GSM8K Accuracy [{gpu_config}]\n\n"
|
||||
summary += "| Variant | Prefill Backend | Decode Backend | KV Cache | Accuracy |\n"
|
||||
summary += "|---------|-----------------|----------------|----------|----------|\n"
|
||||
|
||||
# Add each result as a row
|
||||
for result in TEST_RESULTS:
|
||||
summary += (
|
||||
f"| {result['variant']} | {result['prefill_backend']} | "
|
||||
f"{result['decode_backend']} | {result['kv_cache']} | "
|
||||
f"{result['accuracy']:.3f} |\n"
|
||||
)
|
||||
|
||||
write_github_step_summary(summary)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,54 +0,0 @@
|
||||
"""
|
||||
Usage:
|
||||
cd test/srt
|
||||
python3 -m unittest test_deepseek_v3_deterministic.TestFa3Deterministic
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from sglang.test.test_deterministic_utils import (
|
||||
COMMON_SERVER_ARGS,
|
||||
TestDeterministicBase,
|
||||
)
|
||||
|
||||
DEEPSEEK_MODEL = "lmsys/sglang-ci-dsv3-test"
|
||||
|
||||
|
||||
class TestFa3Deterministic(TestDeterministicBase):
|
||||
@classmethod
|
||||
def get_model(cls):
|
||||
return DEEPSEEK_MODEL
|
||||
|
||||
# Test with fa3 attention backend
|
||||
@classmethod
|
||||
def get_server_args(cls):
|
||||
args = COMMON_SERVER_ARGS
|
||||
args.extend(
|
||||
[
|
||||
"--attention-backend",
|
||||
"fa3",
|
||||
]
|
||||
)
|
||||
return args
|
||||
|
||||
|
||||
class TestTritonDeterministic(TestDeterministicBase):
|
||||
@classmethod
|
||||
def get_model(cls):
|
||||
return DEEPSEEK_MODEL
|
||||
|
||||
# Test with triton attention backend
|
||||
@classmethod
|
||||
def get_server_args(cls):
|
||||
args = COMMON_SERVER_ARGS
|
||||
args.extend(
|
||||
[
|
||||
"--attention-backend",
|
||||
"triton",
|
||||
]
|
||||
)
|
||||
return args
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,72 +0,0 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
is_in_ci,
|
||||
popen_launch_server,
|
||||
write_github_step_summary,
|
||||
)
|
||||
|
||||
FULL_DEEPSEEK_V3_FP4_MODEL_PATH = "nvidia/DeepSeek-V3-0324-FP4"
|
||||
SERVER_LAUNCH_TIMEOUT = 1000
|
||||
|
||||
|
||||
class TestDeepseekV3FP4CutlassMoE(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = FULL_DEEPSEEK_V3_FP4_MODEL_PATH
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
other_args = [
|
||||
"--tp",
|
||||
"4",
|
||||
"--ep",
|
||||
"4",
|
||||
"--attention-backend",
|
||||
"trtllm_mla",
|
||||
"--moe-runner-backend",
|
||||
"flashinfer_cutlass",
|
||||
"--quantization",
|
||||
"modelopt_fp4",
|
||||
"--model-loader-extra-config",
|
||||
'{"enable_multithread_load": true}',
|
||||
]
|
||||
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_a_gsm8k(
|
||||
self,
|
||||
): # Append an "a" to make this test run first (alphabetically) to warm up the server
|
||||
args = SimpleNamespace(
|
||||
num_shots=8,
|
||||
data_path=None,
|
||||
num_questions=1319,
|
||||
parallel=1319,
|
||||
max_new_tokens=512,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
print(f"{metrics=}")
|
||||
|
||||
if is_in_ci():
|
||||
write_github_step_summary(
|
||||
f"### test_gsm8k (deepseek-v3-fp4-cutlass-moe)\n"
|
||||
f'{metrics["accuracy"]=:.3f}\n'
|
||||
)
|
||||
self.assertGreater(metrics["accuracy"], 0.935)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,454 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from typing import Callable
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from flashinfer import fp4_quantize, scaled_fp4_grouped_quantize
|
||||
from flashinfer.fused_moe import cutlass_fused_moe as flashinfer_cutlass_fused_moe
|
||||
from sgl_kernel import scaled_fp4_quant, silu_and_mul
|
||||
from torch.nn import functional as F
|
||||
|
||||
from sglang.srt.layers.moe.cutlass_moe import cutlass_moe_fp4
|
||||
from sglang.srt.layers.moe.cutlass_moe_params import CutlassMoEParams, CutlassMoEType
|
||||
from sglang.srt.layers.moe.topk import TopKConfig, select_experts
|
||||
|
||||
if torch.cuda.get_device_capability() < (10, 0):
|
||||
pytest.skip(
|
||||
reason="Nvfp4 Requires compute capability of 10 or above.",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
kE2M1ToFloat = torch.tensor(
|
||||
[0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32
|
||||
)
|
||||
|
||||
FLOAT8_E4M3_MAX = 448.0
|
||||
FLOAT4_E2M1_MAX = 6.0
|
||||
|
||||
|
||||
def convert_swizzled_to_linear(a_sf_swizzled: torch.Tensor, m, k, block_size):
|
||||
m_tiles = (m + 128 - 1) // 128
|
||||
f = block_size * 4
|
||||
k_tiles = (k + f - 1) // f
|
||||
tmp = torch.reshape(a_sf_swizzled, (1, m_tiles, k_tiles, 32, 4, 4))
|
||||
tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5))
|
||||
out = tmp.reshape(m_tiles * 128, k_tiles * f // block_size)
|
||||
return out[0:m, 0:k]
|
||||
|
||||
|
||||
def dequantize_nvfp4_to_dtype(
|
||||
tensor_fp4, tensor_sf, global_scale, dtype, device, block_size=16
|
||||
):
|
||||
"""Dequantize the fp4 tensor back to high precision."""
|
||||
# Two fp4 values are packed into one uint8.
|
||||
assert tensor_fp4.dtype == torch.uint8
|
||||
m, packed_k = tensor_fp4.shape
|
||||
k = packed_k * 2
|
||||
tensor_f32 = break_fp4_bytes(tensor_fp4, dtype)
|
||||
tensor_f32 = tensor_f32.reshape(m, k // block_size, block_size)
|
||||
tensor_sf = tensor_sf.view(torch.float8_e4m3fn)
|
||||
tensor_sf = convert_swizzled_to_linear(tensor_sf, m, k, block_size)
|
||||
tensor_sf_dtype = tensor_sf.to(torch.float32) / global_scale
|
||||
|
||||
# scale the tensor
|
||||
out = (tensor_f32 * tensor_sf_dtype.unsqueeze(-1)).reshape(m, k)
|
||||
return out.to(dtype=dtype)
|
||||
|
||||
|
||||
def break_fp4_bytes(a, dtype):
|
||||
assert a.dtype == torch.uint8
|
||||
m, n = a.shape
|
||||
|
||||
# Vectorized nibble processing
|
||||
a_flat = a.flatten()
|
||||
high = (a_flat & 0xF0) >> 4 # Upper nibbles
|
||||
low = a_flat & 0x0F # Lower nibbles
|
||||
|
||||
# Combine nibbles for batch processing
|
||||
combined = torch.stack((low, high), dim=1).flatten()
|
||||
|
||||
# Vectorized sign and magnitude extraction
|
||||
signs = (combined & 0x08).to(torch.bool) # Sign bits
|
||||
abs_vals = (combined & 0x07).to(torch.long) # Magnitude indices
|
||||
|
||||
# Device-aware lookup and sign application
|
||||
kE2M1 = kE2M1ToFloat.to(device=a.device)
|
||||
values = kE2M1[abs_vals] * torch.where(signs, -1.0, 1.0)
|
||||
|
||||
# Reshape to final form
|
||||
return values.reshape(m, n * 2).to(dtype=dtype)
|
||||
|
||||
|
||||
def compute_routing(router_logits: torch.Tensor, top_k: int):
|
||||
routing_weights = torch.softmax(router_logits, dim=1, dtype=torch.float)
|
||||
routing_weights, selected_experts = torch.topk(routing_weights, top_k, dim=-1)
|
||||
routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
|
||||
routing_weights = routing_weights.float()
|
||||
return routing_weights, selected_experts
|
||||
|
||||
|
||||
def prepare_inputs(
|
||||
hidden_states: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
num_experts: int,
|
||||
topk: int,
|
||||
):
|
||||
routing_weights, topk_idx = compute_routing(router_logits, topk)
|
||||
|
||||
masked_m = []
|
||||
for i in range(num_experts):
|
||||
mask = topk_idx.view(-1) == i
|
||||
masked_m.append(mask.sum())
|
||||
|
||||
masked_m = torch.tensor(masked_m, dtype=torch.int32)
|
||||
hidden_states_3d = torch.empty(
|
||||
(num_experts, max(masked_m), hidden_states.shape[1]), dtype=hidden_states.dtype
|
||||
)
|
||||
for i in range(num_experts):
|
||||
hidden_states_3d[i, : masked_m[i], :] = hidden_states[topk_idx.view(-1) == i]
|
||||
|
||||
return hidden_states_3d, masked_m, topk_idx, routing_weights
|
||||
|
||||
|
||||
MNK_FACTORS = [
|
||||
(2, 1024, 1024),
|
||||
(2, 1024, 1536),
|
||||
(2, 3072, 1024),
|
||||
(2, 3072, 1536),
|
||||
(64, 1024, 1024),
|
||||
(64, 1024, 1536),
|
||||
(64, 3072, 1024),
|
||||
(64, 2048, 1024),
|
||||
(224, 1024, 1024),
|
||||
(224, 1024, 1536),
|
||||
]
|
||||
|
||||
|
||||
# Reference implementation of torch_moe
|
||||
def torch_moe(a, w1, w2, score, topk, expert_map):
|
||||
B, D = a.shape
|
||||
a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
|
||||
out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
|
||||
score = torch.softmax(score, dim=-1, dtype=torch.float32)
|
||||
topk_weight, topk_ids = torch.topk(score, topk)
|
||||
topk_weight = topk_weight.view(-1)
|
||||
topk_ids = topk_ids.view(-1)
|
||||
if expert_map is not None:
|
||||
topk_ids = expert_map[topk_ids]
|
||||
for i in range(w1.shape[0]):
|
||||
mask = topk_ids == i
|
||||
if mask.sum():
|
||||
out[mask] = silu_and_mul(a[mask] @ w1[i].transpose(0, 1)) @ w2[i].transpose(
|
||||
0, 1
|
||||
)
|
||||
return (
|
||||
out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype)
|
||||
).sum(dim=1)
|
||||
|
||||
|
||||
def torch_moe_nvfp4(a, w1, w2, topk, topk_weight, topk_ids):
|
||||
B, D = a.shape
|
||||
a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
|
||||
out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
|
||||
|
||||
topk_weight = topk_weight.view(-1)
|
||||
topk_ids = topk_ids.view(-1)
|
||||
|
||||
for i in range(w1.shape[0]):
|
||||
mask = topk_ids == i
|
||||
if mask.sum():
|
||||
m = w1[i].shape[0]
|
||||
assert m % 2 == 0
|
||||
# Note: w1 and w3 are swapped!
|
||||
w3_expert, w1_expert = w1[i][m // 2 :, :], w1[i][: m // 2, :]
|
||||
inter = F.silu(a[mask] @ w1_expert.t()) * (a[mask] @ w3_expert.t())
|
||||
inter_gs = torch.tensor(1.0).cuda()
|
||||
inter_q, inter_blockscale = fp4_quantize(inter, inter_gs)
|
||||
inter = dequantize_nvfp4_to_dtype(
|
||||
inter_q,
|
||||
inter_blockscale,
|
||||
inter_gs,
|
||||
dtype=inter.dtype,
|
||||
device=inter.device,
|
||||
block_size=16,
|
||||
).cuda()
|
||||
out[mask] = inter @ w2[i].transpose(0, 1)
|
||||
return (
|
||||
out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype)
|
||||
).sum(dim=1)
|
||||
|
||||
|
||||
def flashinfer_cutedsl_grouped_gemm_nt_masked(
|
||||
hidden_states: torch.Tensor, # 3d
|
||||
input_global_scale: torch.Tensor, # (l,)
|
||||
weights: torch.Tensor,
|
||||
w_global_scale: torch.Tensor, # (l,)
|
||||
masked_m: torch.Tensor,
|
||||
):
|
||||
from flashinfer.cute_dsl.blockscaled_gemm import grouped_gemm_nt_masked
|
||||
|
||||
# hidden_states: [l, m, k]
|
||||
# weights: [l, n, k]
|
||||
aq, aq_sf = scaled_fp4_grouped_quantize(
|
||||
hidden_states,
|
||||
masked_m.to(hidden_states.device),
|
||||
input_global_scale,
|
||||
)
|
||||
num_experts, n, k = weights.shape
|
||||
bq, bq_sf = scaled_fp4_grouped_quantize(
|
||||
weights,
|
||||
torch.ones(num_experts, device=weights.device, dtype=torch.int32) * n,
|
||||
w_global_scale,
|
||||
)
|
||||
|
||||
out = torch.zeros(
|
||||
(num_experts, max(masked_m), n), dtype=weights.dtype, device=aq.device
|
||||
)
|
||||
out = out.permute(1, 2, 0) # requirement of kernel
|
||||
sf_vec_size = 16
|
||||
ab_dtype = "float4_e2m1fn"
|
||||
sf_dtype = "float8_e4m3fn"
|
||||
c_dtype = "bfloat16"
|
||||
alpha = 1.0 / (input_global_scale * w_global_scale).to(out.dtype).view(
|
||||
1, 1, num_experts
|
||||
)
|
||||
|
||||
def get_cute_dtype(input: torch.Tensor) -> str:
|
||||
if input.dtype == torch.bfloat16:
|
||||
return "bfloat16"
|
||||
elif input.dtype == torch.float16:
|
||||
return "float16"
|
||||
elif input.dtype == torch.float32:
|
||||
return "float32"
|
||||
else:
|
||||
raise ValueError(f"Unsupported cute dtype {input.dtype}")
|
||||
|
||||
grouped_gemm_nt_masked(
|
||||
(aq, aq_sf),
|
||||
(bq, bq_sf),
|
||||
out,
|
||||
masked_m.to(aq.device),
|
||||
ab_dtype=ab_dtype,
|
||||
sf_dtype=sf_dtype,
|
||||
c_dtype=c_dtype,
|
||||
sf_vec_size=sf_vec_size,
|
||||
alpha=alpha,
|
||||
alpha_dtype=get_cute_dtype(alpha),
|
||||
)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def check_moe(
|
||||
m: int,
|
||||
n: int,
|
||||
k: int,
|
||||
e: int,
|
||||
topk: int,
|
||||
dtype: torch.dtype,
|
||||
moe_impl: Callable,
|
||||
flip_w13: bool,
|
||||
):
|
||||
torch.manual_seed(7)
|
||||
a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
|
||||
w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
|
||||
quant_blocksize = 16
|
||||
round_up = lambda x, y: (x + y - 1) // y * y
|
||||
sf_w1_2n = round_up(2 * n, 128)
|
||||
sf_w1_k = round_up(k // quant_blocksize, 4)
|
||||
w1_blockscale = torch.empty(
|
||||
(e, sf_w1_2n, sf_w1_k), device="cuda", dtype=torch.float8_e4m3fn
|
||||
)
|
||||
|
||||
w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
|
||||
sf_w2_k = round_up(k, 128)
|
||||
sf_w2_n = round_up(n // quant_blocksize, 4)
|
||||
w2_blockscale = torch.empty(
|
||||
(e, sf_w2_k, sf_w2_n), device="cuda", dtype=torch.float8_e4m3fn
|
||||
)
|
||||
|
||||
w1_q = torch.empty((e, 2 * n, k // 2), device="cuda", dtype=torch.uint8)
|
||||
w2_q = torch.empty((e, k, n // 2), device="cuda", dtype=torch.uint8)
|
||||
w1_gs = torch.empty((e,), device="cuda", dtype=torch.float32)
|
||||
w2_gs = torch.empty((e,), device="cuda", dtype=torch.float32)
|
||||
|
||||
for expert in range(e):
|
||||
w1_amax = torch.abs(w1).max().to(torch.float32)
|
||||
w2_amax = torch.abs(w2).max().to(torch.float32)
|
||||
w1_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w1_amax
|
||||
w2_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w2_amax
|
||||
|
||||
w1_q[expert], w1_blockscale[expert] = scaled_fp4_quant(
|
||||
w1[expert], w1_gs[expert]
|
||||
)
|
||||
|
||||
w2_q[expert], w2_blockscale[expert] = scaled_fp4_quant(
|
||||
w2[expert], w2_gs[expert]
|
||||
)
|
||||
|
||||
score = torch.randn((m, e), device="cuda", dtype=dtype)
|
||||
|
||||
topk_output = select_experts(
|
||||
hidden_states=a,
|
||||
router_logits=score,
|
||||
topk_config=TopKConfig(top_k=topk, renormalize=False),
|
||||
)
|
||||
topk_weights, topk_ids, _ = topk_output
|
||||
|
||||
a1_gs = torch.ones((e,), device="cuda", dtype=torch.float32)
|
||||
a2_gs = torch.ones((e,), device="cuda", dtype=torch.float32)
|
||||
test_output = moe_impl(
|
||||
a=a,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
w1_q=w1_q,
|
||||
w2_q=w2_q,
|
||||
a1_gs=a1_gs,
|
||||
w1_blockscale=w1_blockscale,
|
||||
w1_alphas=(1 / w1_gs),
|
||||
a2_gs=a2_gs,
|
||||
w2_blockscale=w2_blockscale,
|
||||
w2_alphas=(1 / w2_gs),
|
||||
)
|
||||
|
||||
# Reference check:
|
||||
a_global_scale = (
|
||||
(FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a.flatten(), dim=-1)
|
||||
).to(torch.float32)
|
||||
a_fp4, a_scale_interleaved = scaled_fp4_quant(a, a_global_scale)
|
||||
_, m_k = a_fp4.shape
|
||||
a_in_dtype = dequantize_nvfp4_to_dtype(
|
||||
a_fp4,
|
||||
a_scale_interleaved,
|
||||
a_global_scale,
|
||||
dtype=a.dtype,
|
||||
device=a.device,
|
||||
block_size=quant_blocksize,
|
||||
)
|
||||
|
||||
w1_d = torch.empty((e, 2 * n, k), device="cuda", dtype=dtype)
|
||||
w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype)
|
||||
|
||||
for idx in range(0, e):
|
||||
w1_d[idx] = dequantize_nvfp4_to_dtype(
|
||||
w1_q[idx],
|
||||
w1_blockscale[idx],
|
||||
w1_gs[idx],
|
||||
dtype=w1.dtype,
|
||||
device=w1.device,
|
||||
block_size=quant_blocksize,
|
||||
)
|
||||
w2_d[idx] = dequantize_nvfp4_to_dtype(
|
||||
w2_q[idx],
|
||||
w2_blockscale[idx],
|
||||
w2_gs[idx],
|
||||
dtype=w2.dtype,
|
||||
device=w2.device,
|
||||
block_size=quant_blocksize,
|
||||
)
|
||||
|
||||
if flip_w13:
|
||||
dim = -2
|
||||
size = w1_d.size(dim)
|
||||
assert size % 2 == 0, f"Expected even size in dim {dim}, got {size}"
|
||||
half = size // 2
|
||||
# Reorder weight
|
||||
w1, w3 = w1_d.split(half, dim=dim)
|
||||
w1_d = torch.cat([w3, w1], dim=dim).contiguous()
|
||||
|
||||
torch_output = torch_moe(a_in_dtype, w1_d, w2_d, score, topk, None)
|
||||
|
||||
torch.testing.assert_close(torch_output, test_output, atol=1e-1, rtol=1e-1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
|
||||
@pytest.mark.parametrize("e", [40, 64, 256])
|
||||
@pytest.mark.parametrize("topk", [1, 6, 8])
|
||||
@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16])
|
||||
@torch.inference_mode()
|
||||
def test_cutlass_fp4_moe_no_graph(
|
||||
m: int, n: int, k: int, e: int, topk: int, dtype: torch.dtype
|
||||
):
|
||||
def cutlass_moe_impl(
|
||||
a,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
w1_q,
|
||||
w2_q,
|
||||
a1_gs,
|
||||
w1_blockscale,
|
||||
w1_alphas,
|
||||
a2_gs,
|
||||
w2_blockscale,
|
||||
w2_alphas,
|
||||
):
|
||||
params = CutlassMoEParams(
|
||||
CutlassMoEType.BlockscaledFP4,
|
||||
device=a.device,
|
||||
num_experts=e,
|
||||
intermediate_size_per_partition=n, # n
|
||||
hidden_size=k,
|
||||
) # k
|
||||
return cutlass_moe_fp4(
|
||||
a=a,
|
||||
a1_gscale=a1_gs,
|
||||
w1_fp4=w1_q,
|
||||
w1_blockscale=w1_blockscale,
|
||||
w1_alphas=w1_alphas,
|
||||
a2_gscale=a2_gs,
|
||||
w2_fp4=w2_q,
|
||||
w2_blockscale=w2_blockscale,
|
||||
w2_alphas=w2_alphas,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
params=params,
|
||||
apply_router_weight_on_input=False,
|
||||
)
|
||||
|
||||
check_moe(m, n, k, e, topk, dtype, cutlass_moe_impl, flip_w13=False)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
|
||||
@pytest.mark.parametrize("e", [40, 64, 256])
|
||||
@pytest.mark.parametrize("topk", [1, 6, 8])
|
||||
@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16])
|
||||
@torch.inference_mode()
|
||||
def test_flashinfer_fp4_moe_no_graph(
|
||||
m: int, n: int, k: int, e: int, topk: int, dtype: torch.dtype
|
||||
):
|
||||
def flashinfer_moe_impl(
|
||||
a,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
w1_q,
|
||||
w2_q,
|
||||
a1_gs,
|
||||
w1_blockscale,
|
||||
w1_alphas,
|
||||
a2_gs,
|
||||
w2_blockscale,
|
||||
w2_alphas,
|
||||
):
|
||||
return flashinfer_cutlass_fused_moe(
|
||||
a,
|
||||
topk_ids.to(torch.int),
|
||||
topk_weights,
|
||||
w1_q.view(torch.long),
|
||||
w2_q.view(torch.long),
|
||||
a.dtype,
|
||||
quant_scales=[
|
||||
a1_gs,
|
||||
w1_blockscale.view(torch.int32),
|
||||
w1_alphas,
|
||||
a2_gs,
|
||||
w2_blockscale.view(torch.int32),
|
||||
w2_alphas,
|
||||
],
|
||||
)[0]
|
||||
|
||||
check_moe(m, n, k, e, topk, dtype, flashinfer_moe_impl, flip_w13=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_cutlass_fp4_moe_no_graph(224, 1024, 1024, 256, 8, torch.half)
|
||||
test_flashinfer_fp4_moe_no_graph(224, 1024, 1024, 256, 8, torch.half)
|
||||
@@ -1,44 +0,0 @@
|
||||
"""
|
||||
Usage:
|
||||
cd test/srt
|
||||
python3 -m unittest test_qwen3_next_deterministic.TestFlashInferDeterministic
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from sglang.test.test_deterministic_utils import (
|
||||
COMMON_SERVER_ARGS,
|
||||
TestDeterministicBase,
|
||||
)
|
||||
|
||||
QWEN3_NEXT = "Qwen/Qwen3-Next-80B-A3B-Instruct"
|
||||
|
||||
|
||||
class TestFlashInferDeterministic(TestDeterministicBase):
|
||||
@classmethod
|
||||
def get_model(cls):
|
||||
return QWEN3_NEXT
|
||||
|
||||
# Test with flashinfer attention backend
|
||||
@classmethod
|
||||
def get_server_args(cls):
|
||||
args = COMMON_SERVER_ARGS
|
||||
args.extend(["--attention-backend", "flashinfer", "--tp", "4"])
|
||||
return args
|
||||
|
||||
|
||||
class TestTritonDeterministic(TestDeterministicBase):
|
||||
@classmethod
|
||||
def get_model(cls):
|
||||
return QWEN3_NEXT
|
||||
|
||||
# Test with triton attention backend
|
||||
@classmethod
|
||||
def get_server_args(cls):
|
||||
args = COMMON_SERVER_ARGS
|
||||
args.extend(["--attention-backend", "triton", "--tp", "4"])
|
||||
return args
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user