minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

move dead sglang.test files to test/manual (#25316)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-05-14 20:02:44 -07:00
committed by GitHub
parent 8d5b347edd
commit d89b678d69
29 changed files with 0 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/test/ascend/disaggregation_utils.py
-153
View File
@@ -1,153 +0,0 @@
import logging
import os
import time
import warnings
from urllib.parse import urlparse
import requests
from sglang.srt.environ import envs
from sglang.srt.utils import kill_process_tree
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_with_error_check,
)
logger = logging.getLogger(__name__)
class TestDisaggregationBase(CustomTestCase):
@classmethod
def setUpClass(cls):
parsed_url = urlparse(DEFAULT_URL_FOR_TEST)
cls.base_host = parsed_url.hostname
base_port = str(parsed_url.port)
cls.lb_port = base_port
cls.prefill_port = f"{int(base_port) + 100}"
cls.decode_port = f"{int(base_port) + 200}"
cls.prefill_url = f"http://{cls.base_host}:{cls.prefill_port}"
cls.decode_url = f"http://{cls.base_host}:{cls.decode_port}"
cls.lb_url = f"http://{cls.base_host}:{cls.lb_port}"
print(f"{cls.base_host=} {cls.lb_port=} {cls.prefill_port=} {cls.decode_port=}")
cls.process_lb, cls.process_decode, cls.process_prefill = None, None, None
# config transfer backend and rdma devices
cls.transfer_backend = [
"--disaggregation-transfer-backend",
envs.SGLANG_TEST_PD_DISAGG_BACKEND.get(),
]
cls.rdma_devices = [
"--disaggregation-ib-device",
envs.SGLANG_TEST_PD_DISAGG_DEVICES.get(),
]
if cls.rdma_devices[1] is None:
cls.rdma_devices = []
msg = "No RDMA devices specified for disaggregation test, using default settings."
warnings.warn(msg)
@classmethod
def launch_lb(cls):
lb_command = [
"python3",
"-m",
"sglang_router.launch_router",
"--pd-disaggregation",
"--mini-lb",
"--prefill",
cls.prefill_url,
"--decode",
cls.decode_url,
"--host",
cls.base_host,
"--port",
cls.lb_port,
]
print("Starting load balancer:", " ".join(lb_command))
cls.process_lb = popen_with_error_check(lb_command)
cls.wait_server_ready(cls.lb_url + "/health")
@classmethod
def wait_server_ready(cls, url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH):
start_time = time.perf_counter()
while True:
try:
response = requests.get(url)
if response.status_code == 200:
print(f"Server {url} is ready")
return
except Exception:
pass
if time.perf_counter() - start_time > timeout:
raise RuntimeError(f"Server {url} failed to start in {timeout}s")
time.sleep(1)
@classmethod
def tearDownClass(cls):
for process in [cls.process_lb, cls.process_decode, cls.process_prefill]:
if process:
try:
kill_process_tree(process.pid)
except Exception as e:
print(f"Error killing process {process.pid}: {e}")
# wait for 5 seconds
time.sleep(5)
def get_rdma_devices_args():
def _parse_list_env(var_name: str):
val = os.getenv(var_name)
if not val:
return None
items = [x.strip() for x in val.split(",") if x.strip()]
return items or None
def _pick_default_pair(rdma_all_devices):
return [rdma_all_devices[0], rdma_all_devices[len(rdma_all_devices) // 2]]
rdma_all_devices = _parse_list_env("SGLANG_CI_RDMA_ALL_DEVICES") or [
f"mlx5_roce{i}" for i in range(8)
]
logger.info("Resolved rdma_all_devices=%s", rdma_all_devices)
n_rdma = len(rdma_all_devices)
# 1. Get visible GPU indices
cuda_visible_devices = os.getenv("CUDA_VISIBLE_DEVICES")
if not cuda_visible_devices:
warnings.warn("CUDA_VISIBLE_DEVICES is not set. Using default RDMA devices.")
return ",".join(_pick_default_pair(rdma_all_devices))
try:
# Convert to list of integers (handling possible spaces and empty strings)
gpu_indices = [
int(idx.strip()) for idx in cuda_visible_devices.split(",") if idx.strip()
]
if not gpu_indices or len(gpu_indices) > 4:
return ",".join(_pick_default_pair(rdma_all_devices))
except ValueError:
warnings.warn(f"Invalid CUDA_VISIBLE_DEVICES format: {cuda_visible_devices}")
return ",".join(_pick_default_pair(rdma_all_devices))
# 2. Calculate base RDMA index group (each group of 4 GPUs uses consecutive devices)
base_rdma_group = (min(gpu_indices) // 4) * 4
for gpu_idx in gpu_indices:
if not (base_rdma_group <= gpu_idx < base_rdma_group + 4):
warnings.warn(
f"GPU index {gpu_idx} is outside expected group "
f"{base_rdma_group}-{base_rdma_group+3}"
)
# 3. Generate RDMA device names
rdma_devices = []
for gpu_idx in gpu_indices:
nic_index = gpu_idx // (8 // n_rdma)
rdma_devices.append(rdma_all_devices[nic_index])
if not rdma_devices:
return ",".join(_pick_default_pair(rdma_all_devices))
return ",".join(rdma_devices)
python/sglang/test/ascend/test_ascend_vocab_mask.py
-81
View File
@@ -1,81 +0,0 @@
import math
import pytest
import torch
from sglang.srt.constrained import xgrammar_backend as xb
def _pack_mask(allowed_ids, vocab_size, batch_size=1):
nwords = math.ceil(vocab_size / 32)
m = torch.zeros((batch_size, nwords), dtype=torch.int32)
for b in range(batch_size):
for tid in allowed_ids[b]:
m[b, tid // 32] |= 1 << (tid % 32)
return m
def _apply_ref_cpu(logits, vocab_mask):
vocab_size = logits.shape[-1]
token_ids = torch.arange(vocab_size, device="cpu", dtype=torch.int64)
word_idx = token_ids // 32
bit_idx = (token_ids % 32).to(torch.int32)
words = vocab_mask.cpu()[:, word_idx].to(torch.int32)
allowed = ((words >> bit_idx) & 1).bool().to(logits.device)
out = logits.clone()
out.masked_fill_(~allowed, float("-inf"))
return out
@pytest.mark.skipif(
not hasattr(torch, "npu") or not torch.npu.is_available(), reason="NPU required"
)
def test_mask_blocks_disallowed_token_on_npu():
device = "npu:0"
vocab_size = 64
logits = torch.zeros((1, vocab_size), device=device, dtype=torch.float32)
logits[0, 16] = 22.125
logits[0, 5] = 10.0
allowed = [[5, 6, 7, 8]]
vocab_mask = _pack_mask(allowed, vocab_size).to(device=device, dtype=torch.int32)
g = xb.XGrammarGrammar.__new__(xb.XGrammarGrammar)
out = logits.clone()
g.apply_vocab_mask(out, vocab_mask)
assert not torch.isfinite(out[0, 16])
assert int(torch.argmax(out[0]).item()) != 16
@pytest.mark.skipif(
not hasattr(torch, "npu") or not torch.npu.is_available(), reason="NPU required"
)
def test_npu_path_matches_reference_random():
device = "npu:0"
B, V = 4, 257
torch.manual_seed(0)
logits = torch.randn(B, V, device=device, dtype=torch.float32)
allowed = []
for _ in range(B):
ids = torch.randperm(V)[: V // 4].tolist()
allowed.append(ids)
vocab_mask = _pack_mask(allowed, V, B).to(device=device, dtype=torch.int32)
g = xb.XGrammarGrammar.__new__(xb.XGrammarGrammar)
out_npu = logits.clone()
g.apply_vocab_mask(out_npu, vocab_mask)
out_ref = _apply_ref_cpu(logits, vocab_mask)
assert torch.equal(torch.isfinite(out_npu), torch.isfinite(out_ref))
diff = (
torch.nan_to_num(out_npu - out_ref, nan=0.0, posinf=0.0, neginf=0.0)
.abs()
.max()
.item()
)
assert diff < 1e-5
python/sglang/test/attention/__init__.py
View File
python/sglang/test/attention/test_flashattn_backend.py
-559
View File
@@ -1,559 +0,0 @@
import unittest
import torch
from sglang.srt.configs.model_config import AttentionArch
from sglang.srt.layers.attention.flashattention_backend import (
FlashAttentionBackend,
draft_decode_set_expand_metadata,
)
from sglang.srt.layers.attention.torch_native_backend import TorchNativeAttnBackend
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.test.test_utils import CustomTestCase
class MockModelRunner:
def __init__(
self,
page_size=1,
num_heads=2,
head_dim=8,
):
self.device = "cuda"
self.dtype = torch.float16
self.kv_cache_dtype = torch.float16
self.is_hybrid_swa = False
self.attention_chunk_size = None
attention_arch = AttentionArch.MHA
# Max batch size for the test.
max_batch_size = 160
# Total tokens(prefix + extend + decode) in the test should not exceed this length.
max_context_len = 2048
self.model_config = type(
"ModelConfig",
(),
{
"context_len": max_context_len,
"is_multimodal": False,
"attention_arch": attention_arch,
"is_encoder_decoder": False,
"is_local_attention_model": False,
},
)()
self.sliding_window_size = None
self.kv_cache_dtype = (
self.dtype
) # torch dtype, required by FlashAttentionBackend
# server_args is still needed for string-based config (kv_cache_dtype_str)
self.server_args = type(
"ServerArgs",
(),
{
"kv_cache_dtype": "auto", # string version for kv_cache_dtype_str
"speculative_eagle_topk": None,
"speculative_num_draft_tokens": 0,
"enable_deterministic_inference": False,
},
)
self.attn_cp_size = 1
# Create a large enough req_to_token_pool to fit the test usage.
self.req_to_token_pool = type(
"TokenPool",
(),
{
# A typical max_bs * max_context_len for cuda graph decode
"size": max_batch_size,
# Add req_to_token attribute
"req_to_token": torch.zeros(
max_batch_size,
max_context_len,
dtype=torch.int32,
device=self.device,
),
},
)()
self.page_size = page_size
max_total_num_tokens = max_batch_size * max_context_len
self.token_to_kv_pool = MHATokenToKVPool(
size=max_total_num_tokens,
page_size=page_size,
dtype=self.dtype,
head_num=num_heads,
head_dim=head_dim,
layer_num=1, # only consider layer=1 for unit test
device=self.device,
enable_memory_saver=False,
)
@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
class TestFlashAttentionBackend(CustomTestCase):
def setUp(self):
# Test parameters
self.batch_size = 2
self.seq_len = 256
self.num_heads = 2
self.head_dim = 8
self.device = "cuda"
self.dtype = torch.float16
def _init_model_runner(self, page_size=1):
self.model_runner = MockModelRunner(
page_size=page_size,
num_heads=self.num_heads,
head_dim=self.head_dim,
)
self.backend = FlashAttentionBackend(self.model_runner)
self.ref_backend = TorchNativeAttnBackend(self.model_runner)
self.model_runner.model_config.num_attention_heads = self.num_heads
def _mock_write_to_req_to_token_pool(self, batch_size, seq_len, page_size):
# if page_size > 1, the token pool stores the index to the page.
# so we need to multiply the index by page_size.
self.req_to_token = (
torch.arange(0, batch_size, dtype=torch.int32, device=self.device)[:, None]
* seq_len
+ torch.arange(0, seq_len, dtype=torch.int32, device=self.device)[None, :]
+ page_size
)
self.model_runner.req_to_token_pool.req_to_token[:batch_size, :seq_len] = (
self.req_to_token
)
def _create_attention_layer(self):
"""Create attention layer for testing."""
return RadixAttention(
num_heads=self.num_heads,
head_dim=self.head_dim,
scaling=1.0,
num_kv_heads=self.num_heads,
layer_id=0,
)
def _create_qkv_tensors(self, tokens_len):
"""Create q, k, v tensors for testing."""
shape = (tokens_len, self.num_heads, self.head_dim)
return (
torch.randn(shape, dtype=self.dtype, device=self.device),
torch.randn(shape, dtype=self.dtype, device=self.device),
torch.randn(shape, dtype=self.dtype, device=self.device),
)
def _run_reference_forward(
self, mode, q, k, v, layer, forward_batch, expected_shape
):
"""Run reference forward pass using native backend."""
if mode == ForwardMode.EXTEND:
output = self.ref_backend.forward_extend(q, k, v, layer, forward_batch)
else: # ForwardMode.DECODE
output = self.ref_backend.forward_decode(q, k, v, layer, forward_batch)
return output.view(expected_shape)
def _verify_output(self, output, expected_shape, output_ref=None):
"""Verify output tensor shape, dtype, and values."""
self.assertEqual(
output.shape,
expected_shape,
f"Expected shape {expected_shape}, got {output.shape}",
)
self.assertEqual(output.dtype, self.dtype)
self.assertEqual(output.device.type, "cuda")
self.assertEqual(
torch.isnan(output).sum().item(), 0, "Output contains NaN values"
)
if output_ref is not None:
if not torch.allclose(output, output_ref, atol=1e-1, rtol=0.0):
# Check where the values differ beyond the given tolerances
diff_mask = ~torch.isclose(output, output_ref, atol=1e-1, rtol=0.0)
# Find the first index where the difference occurs
if diff_mask.any():
first_mismatch_idx = diff_mask.nonzero()[0]
print(
"First mismatch at index:", tuple(first_mismatch_idx.tolist())
)
print("output:", output[tuple(first_mismatch_idx.tolist())])
print("output_ref:", output_ref[tuple(first_mismatch_idx.tolist())])
raise AssertionError(
"Attention output is not close to the torch native backend output"
)
def _create_forward_batch(
self, mode, q_len=None, prefix_len=0, page_size=1, attn_cp_size=1
):
"""Create a forward batch for testing based on mode and lengths."""
self._init_model_runner(page_size=page_size)
# Default to self.seq_len if not specified
q_len = q_len or self.seq_len
if mode == ForwardMode.EXTEND:
total_len = prefix_len + q_len
out_cache_start = prefix_len * self.batch_size
out_cache_end = total_len * self.batch_size
forward_batch = ForwardBatch(
batch_size=self.batch_size,
input_ids=torch.randint(
0, 100, (self.batch_size, q_len), device=self.device
),
out_cache_loc=torch.arange(
out_cache_start, out_cache_end, device=self.device
),
seq_lens_sum=self.batch_size * total_len,
forward_mode=mode,
req_pool_indices=torch.arange(self.batch_size, device=self.device),
seq_lens=torch.tensor(
[total_len] * self.batch_size, device=self.device
),
seq_lens_cpu=torch.tensor([total_len] * self.batch_size, device="cpu"),
extend_prefix_lens=torch.tensor(
[prefix_len] * self.batch_size, device=self.device
),
extend_prefix_lens_cpu=torch.tensor(
[prefix_len] * self.batch_size, device="cpu"
),
extend_seq_lens=torch.tensor(
[q_len] * self.batch_size, device=self.device
),
extend_seq_lens_cpu=torch.tensor(
[q_len] * self.batch_size, device="cpu"
),
attn_backend=self.backend,
)
if attn_cp_size > 1:
forward_batch.attn_cp_metadata = type(
"AttnCPMetadata",
(),
{
"kv_len_prev_tensor": torch.tensor(
[q_len // 2] * self.batch_size,
dtype=torch.int32,
device=self.device,
),
"kv_len_next_tensor": torch.tensor(
[q_len] * self.batch_size,
dtype=torch.int32,
device=self.device,
),
"actual_seq_q_prev": q_len // 2,
"actual_seq_q_next": q_len // 2,
},
)
else: # ForwardMode.DECODE
decode_len = q_len # Assuming 1 for decode testing
total_len = self.seq_len + decode_len
if mode == ForwardMode.DECODE and page_size > 1:
# Get next page_size multiple of self.seq_len
out_cache_start = (
self.batch_size * self.seq_len // page_size + 1
) * page_size
# out_cache_end is the start of the next block
out_cache_end = out_cache_start + decode_len * page_size
else:
out_cache_start = self.batch_size * self.seq_len
out_cache_end = self.batch_size * total_len
forward_batch = ForwardBatch(
batch_size=self.batch_size,
input_ids=torch.randint(
0, 100, (self.batch_size, decode_len), device=self.device
),
out_cache_loc=torch.tensor(
[out_cache_start, out_cache_end], device=self.device
),
seq_lens_sum=self.batch_size * total_len,
forward_mode=mode,
req_pool_indices=torch.arange(self.batch_size, device=self.device),
seq_lens=torch.tensor(
[total_len] * self.batch_size, device=self.device
),
seq_lens_cpu=torch.tensor([total_len] * self.batch_size, device="cpu"),
attn_backend=self.backend,
)
# Add token pool
forward_batch.req_to_token_pool = self.model_runner.req_to_token_pool
# Write current batch's req_to_token to req_to_token_pool
self._mock_write_to_req_to_token_pool(self.batch_size, total_len, page_size)
# Add kv pool for this forward batch
forward_batch.token_to_kv_pool = self.model_runner.token_to_kv_pool
return forward_batch
def _setup_kv_cache(self, forward_batch, layer, cache_len):
# Create constant values for the prefix cache for easy debugging
cache_k = torch.ones(
self.batch_size * cache_len,
self.num_heads,
self.head_dim,
dtype=self.dtype,
device=self.device,
)
cache_v = (
torch.ones(
self.batch_size * cache_len,
self.num_heads,
self.head_dim,
dtype=self.dtype,
device=self.device,
)
* 2
)
# Set the prefix KV cache
forward_batch.token_to_kv_pool.set_kv_buffer(
layer,
torch.arange(self.batch_size * cache_len, device=self.device),
cache_k,
cache_v,
layer.k_scale,
layer.v_scale,
)
def _run_attention_test(self, mode, q_len, prefix_len=0, page_size=1):
"""
Run an attention test with the specified parameters.
Args:
mode: ForwardMode.EXTEND or ForwardMode.DECODE
q_len: Length of the query sequence. For decode mode, q_len is 1.
prefix_len: Length of the prefix sequence for extend mode
page_size: Page size for the KV cache
"""
layer = self._create_attention_layer()
# Create forward batch and set up
forward_batch = self._create_forward_batch(mode, q_len, prefix_len, page_size)
# Create QKV tensors for the input
q, k, v = self._create_qkv_tensors(self.batch_size * q_len)
# KV cache for prefixed extend is prefix_len
# KV cache for decode is same as seq_len
# No KV cache for extend without prefix
if mode == ForwardMode.EXTEND:
if prefix_len > 0:
self._setup_kv_cache(forward_batch, layer, prefix_len)
else:
self._setup_kv_cache(forward_batch, layer, self.seq_len)
self.backend.init_forward_metadata(forward_batch)
if mode == ForwardMode.EXTEND:
expected_shape = (
self.batch_size * q_len,
self.num_heads * self.head_dim,
)
output = self.backend.forward_extend(q, k, v, layer, forward_batch)
else:
expected_shape = (self.batch_size, self.num_heads * self.head_dim)
output = self.backend.forward_decode(q, k, v, layer, forward_batch)
output_ref = self._run_reference_forward(
mode, q, k, v, layer, forward_batch, expected_shape
)
self._verify_output(output, expected_shape, output_ref)
return output
def _run_attention_cp_test(self, mode, q_len, prefix_len=0, page_size=1):
layer = self._create_attention_layer()
# Create forward batch and set up
forward_batch = self._create_forward_batch(
mode, q_len, prefix_len, page_size, attn_cp_size=2
)
self.backend.attn_cp_size = 2
# Create QKV tensors for the input
q, k, v = self._create_qkv_tensors(self.batch_size * q_len)
# KV cache for prefixed extend is prefix_len
# KV cache for decode is same as seq_len
# No KV cache for extend without prefix
# Setup KV cache for CP testing - need KV cache to have actual values
# For extend with CP, we need KV cache populated so attention has something to attend to
self._setup_kv_cache(forward_batch, layer, q_len)
self.backend.init_forward_metadata(forward_batch)
# if mode == ForwardMode.EXTEND:
expected_shape = (
self.batch_size * q_len,
self.num_heads * self.head_dim,
)
output = self.backend.forward_extend(q, k, v, layer, forward_batch)
# else:
# expected_shape = (self.batch_size, self.num_heads * self.head_dim)
# output = self.backend.forward_decode(q, k, v, layer, forward_batch)
output_ref = self._run_reference_forward(
mode, q, k, v, layer, forward_batch, expected_shape
)
self._verify_output(output, expected_shape, output_ref)
return output
def test_forward_extend_cp(self):
"""Test the standard extend operation with context parallel."""
self._run_attention_cp_test(ForwardMode.EXTEND, q_len=self.seq_len)
# def test_forward_extend_cp_with_prefix(self):
# """Test the standard extend operation with context parallel and prefix."""
# prefix_len = self.seq_len // 2
# extend_len = self.seq_len - prefix_len
# self._run_attention_cp_test(ForwardMode.EXTEND, q_len=extend_len, prefix_len=prefix_len)
# def test_forward_extend(self):
# """Test the standard extend operation."""
# self._run_attention_test(ForwardMode.EXTEND, q_len=self.seq_len)
# def test_forward_decode(self):
# """Test the decode operation with cached tokens."""
# self._run_attention_test(ForwardMode.DECODE, q_len=1)
# def test_forward_extend_with_prefix(self):
# """Test extending from cached prefix tokens."""
# prefix_len = self.seq_len // 2
# extend_len = self.seq_len - prefix_len
# self._run_attention_test(
# ForwardMode.EXTEND, q_len=extend_len, prefix_len=prefix_len
# )
# def test_forward_extend_with_page_size_greater_than_1(self):
# """Test extending from cached prefix tokens with page size greater than 1."""
# self._run_attention_test(ForwardMode.EXTEND, q_len=self.seq_len, page_size=64)
# def test_forward_decode_with_page_size_greater_than_1(self):
# """Test decode operation with page size greater than 1."""
# self._run_attention_test(ForwardMode.DECODE, q_len=1, page_size=64)
class TestUpdateDraftDecodeSetExpandMetadata(CustomTestCase):
"""
All the test cases examples have 1 additional cache location than the decode length.
This is to align with the current allocation logic. It does not affect the correctness.
"""
def test_draft_decode_set_expand_metadata(self):
bs, topk, page_size = 1, 2, 4
cases = [
(
torch.tensor(
[
[23, 24],
[31, 32],
],
dtype=torch.int32,
),
torch.tensor(
[
[5, 6],
[7, 8],
],
dtype=torch.int32,
),
1,
),
# Decode span multiple pages:
# duplicated kv cache: 24, 25, 26
# decode locations: 27, 28, 29, 30, 31, 32
# We need 3 pages in total.
(
torch.tensor(
[
[27, 28, 29, 30, 31, 32],
[35, 36, 37, 38, 39, 40],
],
dtype=torch.int32,
),
torch.tensor(
[
[6, 7, 8, 0, 0, 0],
[8, 9, 10, 0, 0, 0],
],
dtype=torch.int32,
),
5,
),
]
last_page_lens = torch.tensor([3], dtype=torch.int32)
for cache_loc, expected_page_table, decode_length in cases:
cache_seqlens_int32 = torch.zeros(bs * topk, dtype=torch.int32)
page_table = torch.zeros_like(cache_loc, dtype=torch.int32)
draft_decode_set_expand_metadata(
cache_seqlens_int32=cache_seqlens_int32,
page_table=page_table,
last_page_lens=last_page_lens,
decode_length=decode_length,
cache_loc=cache_loc,
topk=topk,
page_size=page_size,
)
expected_cache_seqlens = torch.tensor(
[decode_length + 3, decode_length + 3], dtype=torch.int32
)
self.assertTrue(torch.equal(cache_seqlens_int32, expected_cache_seqlens))
self.assertTrue(torch.equal(page_table, expected_page_table))
def test_update_draft_decode_set_expand_metadata_multi_batch(self):
"""
Ensure expand metadata works when batch size > 1 and last pages differ.
"""
bs, topk, decode_length, page_size = 3, 2, 3, 4
cache_loc = torch.tensor(
[
# First batch: last page duplicate is 1, consecutive pages
[1, 2, 3, 4],
[6, 7, 8, 9],
# Second batch: last page duplicate is 3, non-consecutive pages
[3, 8, 9, 10],
[14, 15, 16, 17],
# Third batch: last page duplicate is 0, consecutive pages
[0, 1, 2, 3],
[4, 5, 6, 7],
],
dtype=torch.int32,
)
cache_seqlens_int32 = torch.zeros(bs * topk, dtype=torch.int32)
last_page_lens = torch.tensor([1, 3, 0], dtype=torch.int32)
page_table = torch.zeros_like(cache_loc, dtype=torch.int32)
draft_decode_set_expand_metadata(
cache_seqlens_int32=cache_seqlens_int32,
page_table=page_table,
last_page_lens=last_page_lens,
decode_length=decode_length,
cache_loc=cache_loc,
topk=topk,
page_size=page_size,
)
expected_cache_seqlens = torch.tensor([4, 4, 6, 6, 3, 3], dtype=torch.int32)
expected_page_table = torch.tensor(
[
[0, 1, 0, 0],
[1, 2, 0, 0],
[0, 2, 0, 0],
[3, 4, 0, 0],
[0, 0, 0, 0],
[1, 0, 0, 0],
],
dtype=torch.int32,
)
self.assertTrue(torch.equal(cache_seqlens_int32, expected_cache_seqlens))
self.assertTrue(torch.equal(page_table, expected_page_table))
if __name__ == "__main__":
unittest.main()
python/sglang/test/attention/test_flashattn_mla_backend.py
-331
View File
@@ -1,331 +0,0 @@
import unittest
import torch
from sglang.srt.configs.model_config import AttentionArch
from sglang.srt.layers.attention.flashattention_backend import FlashAttentionBackend
from sglang.srt.layers.attention.torch_native_backend import TorchNativeAttnBackend
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.test.test_utils import CustomTestCase
class MockModelRunner:
def __init__(
self,
kv_lora_rank,
qk_rope_head_dim,
):
attention_arch = AttentionArch.MLA
self.device = "cuda"
self.dtype = torch.float16
self.is_hybrid_swa = False
context_len = 2048
self.model_config = type(
"ModelConfig",
(),
{
"context_len": context_len,
"attention_arch": attention_arch,
"is_encoder_decoder": False,
"is_local_attention_model": False,
},
)
self.sliding_window_size = None
# Add server_args attribute
self.server_args = type(
"ServerArgs",
(),
{
"kv_cache_dtype": torch.float16,
"speculative_eagle_topk": None,
"speculative_num_draft_tokens": 0,
"enable_deterministic_inference": False,
},
)
self.kv_cache_dtype = self.server_args.kv_cache_dtype
batch_size = 160
# Create a proper req_to_token_pool with the req_to_token attribute
self.req_to_token_pool = type(
"TokenPool",
(),
{
# A typical max_bs * max_context_len for cuda graph decode
"size": batch_size,
# Add req_to_token attribute
"req_to_token": torch.zeros(
batch_size, context_len, dtype=torch.int32, device=self.device
),
},
)
self.page_size = 1
max_total_num_tokens = batch_size * context_len
self.token_to_kv_pool = MLATokenToKVPool(
size=max_total_num_tokens,
page_size=self.page_size,
dtype=self.kv_cache_dtype,
kv_lora_rank=kv_lora_rank,
qk_rope_head_dim=qk_rope_head_dim,
layer_num=1, # only consider layer=1 for unit test
device=self.device,
enable_memory_saver=False,
)
class MockReqToTokenPool:
def __init__(self, batch_size, seq_len, device):
self.req_to_token = (
torch.arange(batch_size * seq_len, device=device)
.reshape(batch_size, seq_len)
.to(torch.int32)
)
@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
class TestFlashAttentionMLABackend(CustomTestCase):
def setUp(self):
# MLA with different V headdim requires Hopper architecture (compute capability >= 9.0)
if torch.cuda.is_available():
compute_capability = torch.cuda.get_device_capability()
if compute_capability[0] < 9:
self.skipTest(
f"MLA requires Hopper GPU (compute capability >= 9.0), "
f"but found compute capability {compute_capability[0]}.{compute_capability[1]}"
)
# Test parameters
self.batch_size = 2
self.seq_len = 360
self.num_heads = 2
self.device = "cuda"
self.dtype = torch.float16
self.kv_lora_rank = 512
self.q_lora_rank = 128
self.qk_rope_head_dim = 64
self.qk_head_dim = self.qk_rope_head_dim + self.kv_lora_rank
# Assume no rope scaling
self.scaling = self.qk_head_dim**-0.5
# Initialize model runner and backend
self._init_model_runner()
self.backend = FlashAttentionBackend(self.model_runner)
self.ref_backend = TorchNativeAttnBackend(self.model_runner)
self.num_local_heads = 2
def _init_model_runner(self):
self.model_runner = MockModelRunner(
kv_lora_rank=self.kv_lora_rank,
qk_rope_head_dim=self.qk_rope_head_dim,
)
def _create_attention_layer(self):
"""Create attention layer for testing."""
self.attn_mqa = RadixAttention(
num_heads=self.num_local_heads,
head_dim=self.kv_lora_rank + self.qk_rope_head_dim,
scaling=self.scaling,
num_kv_heads=1,
layer_id=0,
v_head_dim=self.kv_lora_rank,
prefix="attn_mqa",
)
return self.attn_mqa
def _run_reference_forward(
self, mode, q, k, v, layer, forward_batch, expected_shape
):
"""Run reference forward pass using native backend."""
if mode == ForwardMode.EXTEND:
output = self.ref_backend.forward_extend(q, k, v, layer, forward_batch)
else: # ForwardMode.DECODE
output = self.ref_backend.forward_decode(q, k, v, layer, forward_batch)
return output.view(expected_shape)
def _verify_output(self, output, expected_shape):
"""Verify output tensor shape, dtype, and values."""
self.assertEqual(
output.shape,
expected_shape,
f"Expected shape {expected_shape}, got {output.shape}",
)
self.assertEqual(output.dtype, self.dtype)
self.assertEqual(output.device.type, "cuda")
self.assertEqual(
torch.isnan(output).sum().item(), 0, "Output contains NaN values"
)
def _create_forward_batch(self, mode, q_len=None, prefix_len=0):
"""Create a forward batch for testing based on mode and lengths."""
# Default to self.seq_len if not specified
q_len = q_len or self.seq_len
if mode == ForwardMode.EXTEND:
total_len = prefix_len + q_len
out_cache_start = prefix_len * self.batch_size
out_cache_end = total_len * self.batch_size
forward_batch = ForwardBatch(
batch_size=self.batch_size,
input_ids=torch.randint(
0, 100, (self.batch_size, q_len), device=self.device
),
out_cache_loc=torch.arange(
out_cache_start, out_cache_end, device=self.device
),
seq_lens_sum=self.batch_size * total_len,
forward_mode=mode,
req_pool_indices=torch.arange(self.batch_size, device=self.device),
seq_lens=torch.tensor(
[total_len] * self.batch_size, device=self.device
),
seq_lens_cpu=torch.tensor([total_len] * self.batch_size, device="cpu"),
extend_prefix_lens=torch.tensor(
[prefix_len] * self.batch_size, device=self.device
),
extend_prefix_lens_cpu=torch.tensor(
[prefix_len] * self.batch_size, device="cpu"
),
extend_seq_lens=torch.tensor(
[q_len] * self.batch_size, device=self.device
),
extend_seq_lens_cpu=torch.tensor(
[q_len] * self.batch_size, device="cpu"
),
attn_backend=self.backend,
)
else: # ForwardMode.DECODE
decode_len = q_len # typically 1 for decode mode
total_len = self.seq_len + decode_len
out_cache_start = self.batch_size * self.seq_len
out_cache_end = self.batch_size * total_len
forward_batch = ForwardBatch(
batch_size=self.batch_size,
input_ids=torch.randint(
0, 100, (self.batch_size, decode_len), device=self.device
),
out_cache_loc=torch.arange(
out_cache_start, out_cache_end, device=self.device
),
seq_lens_sum=self.batch_size * total_len,
forward_mode=mode,
req_pool_indices=torch.arange(self.batch_size, device=self.device),
seq_lens=torch.tensor(
[total_len] * self.batch_size, device=self.device
),
seq_lens_cpu=torch.tensor([total_len] * self.batch_size, device="cpu"),
attn_backend=self.backend,
)
# Add token pool from model runner to forward batch
forward_batch.req_to_token_pool = self.model_runner.req_to_token_pool
# Add KV cache from model runner to forward batch
forward_batch.token_to_kv_pool = self.model_runner.token_to_kv_pool
return forward_batch
def _setup_kv_cache(self, forward_batch, layer, cache_len):
"""Set up KV cache with prefix tokens."""
if cache_len <= 0:
return
# For MLA, create separate nope and rope caches
cache_k_nope = torch.ones(
self.batch_size * cache_len,
1, # latent cache has only one head in MQA
self.kv_lora_rank,
dtype=self.dtype,
device=self.device,
)
cache_k_rope = torch.ones(
self.batch_size * cache_len,
1, # latent cache has only one head in MQA
self.qk_rope_head_dim,
dtype=self.dtype,
device=self.device,
)
# Set the prefix KV cache using MLA-specific method
forward_batch.token_to_kv_pool.set_mla_kv_buffer(
layer,
torch.arange(self.batch_size * cache_len, device=self.device),
cache_k_nope,
cache_k_rope,
)
def _run_attention_test(self, mode, q_len, prefix_len=0):
"""
Run an attention test with the specified parameters.
Args:
mode: ForwardMode.EXTEND or ForwardMode.DECODE
q_len: Length of the query sequence. For decode mode, q_len is 1.
prefix_len: Length of the prefix sequence for extend mode
"""
layer = self._create_attention_layer()
# Create forward batch and set up
forward_batch = self._create_forward_batch(mode, q_len, prefix_len)
# Create q, kv_compressed for testing
q_shape = (self.batch_size * q_len, self.num_heads, self.qk_head_dim)
kv_shape = (self.batch_size * q_len, self.qk_head_dim)
q = torch.randn(q_shape, dtype=self.dtype, device=self.device)
kv_compressed = torch.randn(kv_shape, dtype=self.dtype, device=self.device)
# For MLA, split kv_compressed into k_nope and k_rope
# k_nope has dimension kv_lora_rank, k_rope has dimension qk_rope_head_dim
k_nope = kv_compressed[:, : self.kv_lora_rank]
k_rope = kv_compressed[:, self.kv_lora_rank :]
# k_nope needs to be unsqueezed for the num_heads dimension
k = k_nope.unsqueeze(1)
# k_rope also needs to be unsqueezed
k_rope = k_rope.unsqueeze(1)
# v is not used for mqa
v = torch.randn((1), dtype=self.dtype, device=self.device)
self._setup_kv_cache(forward_batch, layer, prefix_len)
self.backend.init_forward_metadata(forward_batch)
expected_shape = (
self.batch_size * q_len,
self.num_heads * self.kv_lora_rank,
)
if mode == ForwardMode.EXTEND:
output = self.backend.forward_extend(
q, k, v, layer, forward_batch, k_rope=k_rope
)
else:
output = self.backend.forward_decode(
q, k, v, layer, forward_batch, k_rope=k_rope
)
self._verify_output(output, expected_shape)
return output
def test_forward_extend(self):
"""Test the standard extend operation."""
self._run_attention_test(ForwardMode.EXTEND, q_len=self.seq_len)
def test_forward_decode(self):
"""Test the decode operation with cached tokens."""
self._run_attention_test(ForwardMode.DECODE, q_len=1)
def test_forward_extend_with_prefix(self):
"""Test extending from cached prefix tokens."""
prefix_len = self.seq_len // 2
extend_len = self.seq_len - prefix_len
self._run_attention_test(
ForwardMode.EXTEND, q_len=extend_len, prefix_len=prefix_len
)
if __name__ == "__main__":
unittest.main()
python/sglang/test/attention/test_prefix_chunk_info.py
-228
View File
@@ -1,228 +0,0 @@
import unittest
import torch
from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.utils.common import get_device
from sglang.test.test_utils import CustomTestCase
TEST_CASES = [
# Sequence with same prefix lens
{
"batch_size": 3,
"prefix_lens": [64, 64, 64],
"max_chunk_capacity": 48,
"prefix_chunk_len": 16,
"num_prefix_chunks": 4,
"prefix_chunk_starts": torch.tensor(
[
[0, 0, 0],
[16, 16, 16],
[32, 32, 32],
[48, 48, 48],
],
dtype=torch.int32,
),
"prefix_chunk_seq_lens": torch.tensor(
[
[16, 16, 16],
[16, 16, 16],
[16, 16, 16],
[16, 16, 16],
],
dtype=torch.int32,
),
},
# Sequence with different prefix lens
{
"batch_size": 4,
"prefix_lens": [16, 32, 48, 64],
"max_chunk_capacity": 64,
"prefix_chunk_len": 16,
"num_prefix_chunks": 4,
"prefix_chunk_starts": torch.tensor(
[
[0, 0, 0, 0],
[16, 16, 16, 16],
[32, 32, 32, 32],
[48, 48, 48, 48],
],
dtype=torch.int32,
),
"prefix_chunk_seq_lens": torch.tensor(
[
[16, 16, 16, 16],
[0, 16, 16, 16],
[0, 0, 16, 16],
[0, 0, 0, 16],
],
dtype=torch.int32,
),
},
# Sequence with irregular shapes
{
"batch_size": 2,
"prefix_lens": [1, 64],
"max_chunk_capacity": 31,
"prefix_chunk_len": 15,
"num_prefix_chunks": 5,
"prefix_chunk_starts": torch.tensor(
[
[0, 0],
[15, 15],
[30, 30],
[45, 45],
[60, 60],
],
dtype=torch.int32,
),
"prefix_chunk_seq_lens": torch.tensor(
[
[1, 15],
[0, 15],
[0, 15],
[0, 15],
[0, 4],
],
dtype=torch.int32,
),
},
]
class MockForwardBatch(ForwardBatch):
def __init__(self, max_chunk_capacity: int, *args, **kwargs):
super().__init__(*args, **kwargs)
self.max_chunk_capacity = max_chunk_capacity
def get_max_chunk_capacity(self):
return self.max_chunk_capacity
class MockReqToTokenPool:
def __init__(self, batch_size, seq_len, device):
self.req_to_token = (
torch.arange(batch_size * seq_len, device=device)
.reshape(batch_size, seq_len)
.to(torch.int32)
)
# Test correctness of triton kernel for computing kv indices
def check_kv_indices(forward_batch):
for i in range(forward_batch.num_prefix_chunks):
computed_kv_indices = forward_batch.prefix_chunk_kv_indices[i]
req_to_token = forward_batch.req_to_token_pool.req_to_token[
: forward_batch.batch_size, :
]
ref_kv_indices = torch.empty(
forward_batch.prefix_chunk_num_tokens[i],
dtype=torch.int32,
device=computed_kv_indices.device,
)
running_ptr = 0
for j in range(forward_batch.batch_size):
seq_start = forward_batch.prefix_chunk_starts[i, j].item()
seq_len = forward_batch.prefix_chunk_seq_lens[i, j].item()
ref_kv_indices[running_ptr : running_ptr + seq_len].copy_(
req_to_token[j, seq_start : seq_start + seq_len]
)
running_ptr += seq_len
assert torch.allclose(computed_kv_indices, ref_kv_indices)
@unittest.skipIf(
not (torch.cuda.is_available() or torch.xpu.is_available()),
"Test requires CUDA or XPU",
)
class TestPrefixChunkInfo(CustomTestCase):
def setUp(self):
# Common test parameters
self.num_local_heads = 128
self.kv_lora_rank = 512
self.qk_rope_head_dim = 64
self.device = get_device()
self.dtype = torch.bfloat16
self.extend_len = 64
self.max_bs = 4
self.max_seq_len = 128
# req_to_token_pool
self.req_to_token_pool = MockReqToTokenPool(
self.max_bs,
self.max_seq_len,
self.device,
)
# token_to_kv_pool
self.token_to_kv_pool = MLATokenToKVPool(
size=self.max_bs * self.max_seq_len,
page_size=1, # only consider page=1 for unit test
dtype=self.dtype,
kv_lora_rank=self.kv_lora_rank,
qk_rope_head_dim=self.qk_rope_head_dim,
layer_num=1, # only consider layer=1 for unit test
device=self.device,
enable_memory_saver=False,
)
def test_prefix_chunk_info(self):
"""Test the standard extend operation."""
for test_case in TEST_CASES:
print(
f"Test case with batch_size={test_case['batch_size']}, prefix_lens={test_case['prefix_lens']}, max_chunk_capacity={test_case['max_chunk_capacity']}"
)
batch_size = test_case["batch_size"]
prefix_lens_cpu = test_case["prefix_lens"]
assert len(prefix_lens_cpu) == batch_size
prefix_lens = torch.tensor(prefix_lens_cpu, device=self.device)
max_chunk_capacity = test_case["max_chunk_capacity"]
seq_lens_cpu = [
self.extend_len + prefix_lens_cpu[i] for i in range(batch_size)
]
seq_lens = torch.tensor(seq_lens_cpu, device=self.device)
# Create forward batch
# input_ids and out_cache_loc are dummy tensors in this test
forward_batch = MockForwardBatch(
max_chunk_capacity=max_chunk_capacity,
batch_size=batch_size,
input_ids=torch.randint(
0, 100, (batch_size, self.extend_len), device=self.device
),
out_cache_loc=torch.arange(
self.max_bs * self.max_seq_len - batch_size * self.extend_len,
self.max_bs * self.max_seq_len,
device=self.device,
),
seq_lens_sum=sum(seq_lens_cpu),
forward_mode=ForwardMode.EXTEND,
req_pool_indices=torch.arange(batch_size, device=self.device),
seq_lens=seq_lens,
seq_lens_cpu=seq_lens_cpu,
extend_prefix_lens=prefix_lens,
extend_prefix_lens_cpu=prefix_lens_cpu,
)
forward_batch.req_to_token_pool = self.req_to_token_pool
forward_batch.token_to_kv_pool = self.token_to_kv_pool
forward_batch.prepare_chunked_prefix_cache_info(self.device)
assert forward_batch.get_max_chunk_capacity() == max_chunk_capacity
assert forward_batch.prefix_chunk_len == test_case["prefix_chunk_len"]
assert forward_batch.num_prefix_chunks == test_case["num_prefix_chunks"]
assert torch.allclose(
forward_batch.prefix_chunk_starts,
test_case["prefix_chunk_starts"].to(self.device),
)
assert torch.allclose(
forward_batch.prefix_chunk_seq_lens,
test_case["prefix_chunk_seq_lens"].to(self.device),
)
check_kv_indices(forward_batch)
if __name__ == "__main__":
unittest.main()
python/sglang/test/attention/test_trtllm_mla_backend.py
-1436
View File
File diff suppressed because it is too large Load Diff
python/sglang/test/bench_mxfp4_sm90_kernels.py
-366
View File
@@ -1,366 +0,0 @@
"""Benchmark MXFP4 MoE kernels on H100/H200: SGLang Marlin vs FlashInfer cutlass.
Compares per-call latency of:
* Marlin path : ``fused_marlin_moe(...)`` after Marlin weight repack
* FlashInfer : ``cutlass_fused_moe(use_w4_group_scaling=True, ...)``
(PR #3084's SM90 mixed-input path)
Both run on the same random MXFP4 weights/scales (semantics differ slightly --
Marlin uses a scalar swiglu clamp + no bias, FlashInfer fuses per-expert
SwiGLU with bias -- so the timing comparison reports kernel cost for
*equivalent compute volume*, not bit-exact numerics).
Run on H100/H200:
cd /sgl-workspace/sglang_dev3 && \\
PYTHONPATH=python:/sgl-workspace/flashinfer FLASHINFER_DISABLE_VERSION_CHECK=1 \\
python python/sglang/test/bench_mxfp4_sm90_kernels.py
"""
from __future__ import annotations
import argparse
from dataclasses import dataclass
from typing import Callable, List, Tuple
import torch
from flashinfer.autotuner import autotune
# ---- FlashInfer ----
from flashinfer.fused_moe import (
cutlass_fused_moe,
interleave_moe_scales_for_sm90_mixed_gemm,
interleave_moe_weights_for_sm90_mixed_gemm,
)
from flashinfer.fused_moe.core import ActivationType
# ---- SGLang Marlin ----
from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack
from sglang.srt.layers.moe.fused_moe_triton.fused_marlin_moe import fused_marlin_moe
from sglang.srt.layers.quantization.marlin_utils import (
marlin_make_workspace,
marlin_permute_scales,
)
from sglang.srt.layers.quantization.marlin_utils_fp4 import mxfp4_marlin_process_scales
GROUP_SIZE = 32
@dataclass
class Shape:
tokens: int
hidden: int
inter: int
num_experts: int
top_k: int
def label(self) -> str:
return (
f"m={self.tokens:>4} h={self.hidden} i={self.inter} "
f"E={self.num_experts} k={self.top_k}"
)
# Sweep tokens at a fixed GPT-OSS-like body (hidden=4096, inter=2048, E=256,
# topk=6 -- matches PR #3084's headline shape so the small-batch numbers stay
# directly comparable). Token range covers decode (4-256) and prefill chunks
# (1024-8192).
_BODY = dict(hidden=4096, inter=2048, num_experts=256, top_k=6)
DEFAULT_SHAPES: List[Shape] = [
Shape(tokens=4, **_BODY),
Shape(tokens=16, **_BODY),
Shape(tokens=64, **_BODY),
Shape(tokens=256, **_BODY),
Shape(tokens=1024, **_BODY),
Shape(tokens=2048, **_BODY),
Shape(tokens=4096, **_BODY),
Shape(tokens=8192, **_BODY),
]
def _make_random_mxfp4(shape: Shape, seed: int = 0):
g = torch.Generator(device="cuda").manual_seed(seed)
e = shape.num_experts
n = shape.inter
k = shape.hidden
w13 = torch.randint(
0, 256, (e, 2 * n, k // 2), dtype=torch.uint8, device="cuda", generator=g
)
w2 = torch.randint(
0, 256, (e, k, n // 2), dtype=torch.uint8, device="cuda", generator=g
)
# Narrow E8M0 band so dequant magnitudes stay sane.
w13_s = torch.randint(
125,
130,
(e, 2 * n, k // GROUP_SIZE),
dtype=torch.uint8,
device="cuda",
generator=g,
)
w2_s = torch.randint(
125,
130,
(e, k, n // GROUP_SIZE),
dtype=torch.uint8,
device="cuda",
generator=g,
)
w13_b = (
torch.randn(e, 2 * n, dtype=torch.float32, device="cuda", generator=g).to(
torch.bfloat16
)
* 0.01
)
w2_b = (
torch.randn(e, k, dtype=torch.float32, device="cuda", generator=g).to(
torch.bfloat16
)
* 0.01
)
return w13, w2, w13_s, w2_s, w13_b, w2_b
def _make_topk(shape: Shape, seed: int = 1):
g = torch.Generator(device="cuda").manual_seed(seed)
logits = torch.randn(
shape.tokens,
shape.num_experts,
dtype=torch.float32,
device="cuda",
generator=g,
)
weights, ids = torch.topk(torch.softmax(logits, dim=-1), shape.top_k, dim=-1)
weights = weights / weights.sum(dim=-1, keepdim=True)
return logits, weights.to(torch.float32), ids.to(torch.int32)
# ---------------------------------------------------------------------------
# FlashInfer cutlass path
# ---------------------------------------------------------------------------
def build_flashinfer_inputs(shape: Shape, w13, w2, w13_s, w2_s, w13_b, w2_b):
w13_il = interleave_moe_weights_for_sm90_mixed_gemm(w13, "fp4")
w2_il = interleave_moe_weights_for_sm90_mixed_gemm(w2, "fp4")
w13_s_il = interleave_moe_scales_for_sm90_mixed_gemm(w13_s, group_size=GROUP_SIZE)
w2_s_il = interleave_moe_scales_for_sm90_mixed_gemm(w2_s, group_size=GROUP_SIZE)
e = shape.num_experts
swiglu_alpha = torch.full((e,), 1.702, dtype=torch.float32, device="cuda")
swiglu_beta = torch.full((e,), 1.0, dtype=torch.float32, device="cuda")
swiglu_limit = torch.full((e,), 7.0, dtype=torch.float32, device="cuda")
return {
"w13": w13_il,
"w2": w2_il,
"quant_scales": [w13_s_il.view(torch.int32), w2_s_il.view(torch.int32)],
"w13_b": w13_b,
"w2_b": w2_b,
"swiglu_alpha": swiglu_alpha,
"swiglu_beta": swiglu_beta,
"swiglu_limit": swiglu_limit,
}
def make_flashinfer_runner(
shape: Shape, prep, x, topk_w, topk_i, autotuned: bool, with_bias: bool = True
):
out = torch.empty(shape.tokens, shape.hidden, dtype=torch.bfloat16, device="cuda")
fc1_b = prep["w13_b"] if with_bias else None
fc2_b = prep["w2_b"] if with_bias else None
def _call():
cutlass_fused_moe(
input=x,
token_selected_experts=topk_i,
token_final_scales=topk_w,
fc1_expert_weights=prep["w13"],
fc2_expert_weights=prep["w2"],
output_dtype=torch.bfloat16,
quant_scales=prep["quant_scales"],
fc1_expert_biases=fc1_b,
fc2_expert_biases=fc2_b,
swiglu_alpha=prep["swiglu_alpha"],
swiglu_beta=prep["swiglu_beta"],
swiglu_limit=prep["swiglu_limit"],
use_w4_group_scaling=True,
activation_type=ActivationType.Swiglu,
output=out,
)
if autotuned:
# Populate FlashInfer's tactic cache once before timing.
with autotune(True):
_call()
return _call
# ---------------------------------------------------------------------------
# SGLang Marlin path
# ---------------------------------------------------------------------------
def build_marlin_inputs(shape: Shape, w13, w2, w13_s, w2_s):
"""Repack MXFP4 weights into Marlin layout. Mirrors
``prepare_moe_mxfp4_layer_for_marlin`` but does not require a layer object."""
e = shape.num_experts
n = shape.inter
k = shape.hidden
device = w13.device
perm = torch.empty(0, dtype=torch.int, device=device)
def _repack(weight, size_n, size_k):
out_list = []
for i in range(e):
qweight = weight[i].view(torch.int32).T.contiguous()
out_list.append(
gptq_marlin_repack(
b_q_weight=qweight,
perm=perm,
size_k=size_k,
size_n=size_n,
num_bits=4,
)
)
return torch.stack(out_list)
def _scales_for(scales, size_n, size_k):
out_list = []
# Reinterpret uint8 E8M0 byte as float8_e8m0fnu, then to bf16 numerical.
scales_bf16 = scales.view(torch.float8_e8m0fnu).to(torch.bfloat16)
for i in range(e):
s = scales_bf16[i].T.contiguous()
ms = marlin_permute_scales(
s=s, size_k=size_k, size_n=size_n, group_size=GROUP_SIZE
)
out_list.append(mxfp4_marlin_process_scales(ms, input_dtype=torch.bfloat16))
return torch.stack(out_list)
w13_marlin = _repack(w13, size_n=2 * n, size_k=k)
w2_marlin = _repack(w2, size_n=k, size_k=n)
w13_s_marlin = _scales_for(w13_s, size_n=2 * n, size_k=k)
w2_s_marlin = _scales_for(w2_s, size_n=k, size_k=n)
workspace = marlin_make_workspace(device, 4)
return {
"w13": w13_marlin,
"w2": w2_marlin,
"w13_s": w13_s_marlin,
"w2_s": w2_s_marlin,
"workspace": workspace,
}
def make_marlin_runner(shape: Shape, prep, x_bf16, router_logits, topk_w, topk_i):
def _call():
fused_marlin_moe(
hidden_states=x_bf16,
w1=prep["w13"],
w2=prep["w2"],
w1_scale=prep["w13_s"],
w2_scale=prep["w2_s"],
gating_output=router_logits,
topk_weights=topk_w,
topk_ids=topk_i,
workspace=prep["workspace"],
num_bits=4,
is_k_full=True,
inplace=False,
clamp_limit=7.0,
)
return _call
# ---------------------------------------------------------------------------
# Timing harness
# ---------------------------------------------------------------------------
def time_call(fn: Callable, warmup: int = 5, iters: int = 30) -> Tuple[float, float]:
"""Returns (median_ms, min_ms) across ``iters`` calls after ``warmup``."""
for _ in range(warmup):
fn()
torch.cuda.synchronize()
starts = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
ends = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
for s, e in zip(starts, ends):
s.record()
fn()
e.record()
torch.cuda.synchronize()
times = sorted(s.elapsed_time(e) for s, e in zip(starts, ends))
return times[len(times) // 2], times[0]
def run_one_shape(shape: Shape, run_marlin: bool):
print(f"\n=== {shape.label()} ===")
w13, w2, w13_s, w2_s, w13_b, w2_b = _make_random_mxfp4(shape, seed=0)
router_logits, topk_w, topk_i = _make_topk(shape, seed=1)
x = (
torch.randn(shape.tokens, shape.hidden, dtype=torch.bfloat16, device="cuda")
* 0.1
)
# FlashInfer cutlass (autotune ON, with bias).
fi_prep = build_flashinfer_inputs(shape, w13, w2, w13_s, w2_s, w13_b, w2_b)
fi_at_call = make_flashinfer_runner(
shape, fi_prep, x, topk_w, topk_i, autotuned=True, with_bias=True
)
fi_at_med, fi_at_min = time_call(fi_at_call)
print(
f" FlashInfer cutlass (autotune): median={fi_at_med:.3f} ms "
f"min={fi_at_min:.3f} ms"
)
# FlashInfer cutlass (autotune ON, no bias) -- isolate bias epilogue cost.
fi_at_nb_call = make_flashinfer_runner(
shape, fi_prep, x, topk_w, topk_i, autotuned=True, with_bias=False
)
fi_at_nb_med, fi_at_nb_min = time_call(fi_at_nb_call)
print(
f" FlashInfer cutlass (AT, no-bias): median={fi_at_nb_med:.3f} ms "
f"min={fi_at_nb_min:.3f} ms "
f"(bias overhead = {fi_at_med - fi_at_nb_med:+.3f} ms / "
f"{(fi_at_med / fi_at_nb_med - 1) * 100:+.1f}%)"
)
fi_med = fi_at_med # alias for downstream speedup print
# Marlin
if run_marlin:
try:
ml_prep = build_marlin_inputs(shape, w13, w2, w13_s, w2_s)
ml_call = make_marlin_runner(
shape, ml_prep, x, router_logits, topk_w, topk_i
)
ml_med, ml_min = time_call(ml_call)
print(
f" SGLang Marlin: median={ml_med:.3f} ms "
f"min={ml_min:.3f} ms"
)
print(f" speedup (Marlin / FI autotune): {ml_med / fi_at_med:.2f}x")
print(f" speedup (Marlin / FI AT no-bias): {ml_med / fi_at_nb_med:.2f}x")
except Exception as exc: # pylint: disable=broad-except
print(f" SGLang Marlin: SKIPPED ({type(exc).__name__}: {exc})")
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--no-marlin", action="store_true", help="Skip Marlin path.")
args = parser.parse_args()
if not torch.cuda.is_available():
raise SystemExit("CUDA required.")
cap = torch.cuda.get_device_capability()
if cap[0] != 9:
print(f"WARNING: device cap {cap} is not SM90; SM90-specific kernel may fail.")
print(f"Device: {torch.cuda.get_device_name()} (cap {cap[0]}.{cap[1]})")
for shape in DEFAULT_SHAPES:
run_one_shape(shape, run_marlin=not args.no_marlin)
if __name__ == "__main__":
main()
python/sglang/test/ci/run_with_retry.py
-53
View File
@@ -1,53 +0,0 @@
"""
Simple wrapper to run a test file with retry logic.
Usage:
python3 -m sglang.test.ci.run_with_retry test_file.py [--max-attempts 2] [--retry-wait 60]
"""
import argparse
import sys
from sglang.test.ci.ci_utils import TestFile, run_unittest_files
def main():
parser = argparse.ArgumentParser(description="Run a test file with retry logic")
parser.add_argument("test_file", help="The test file to run")
parser.add_argument(
"--max-attempts",
type=int,
default=2,
help="Maximum number of attempts (default: 2)",
)
parser.add_argument(
"--retry-wait",
type=int,
default=60,
help="Seconds to wait between retries (default: 60)",
)
parser.add_argument(
"--timeout",
type=int,
default=1200,
help="Timeout per attempt in seconds (default: 1200)",
)
args = parser.parse_args()
# Create a TestFile with a reasonable estimated time
test_file = TestFile(name=args.test_file, estimated_time=args.timeout)
exit_code = run_unittest_files(
files=[test_file],
timeout_per_file=args.timeout,
continue_on_error=False,
enable_retry=True,
max_attempts=args.max_attempts,
retry_wait_seconds=args.retry_wait,
)
sys.exit(exit_code)
if __name__ == "__main__":
main()
python/sglang/test/get_logits_ut.py
-57
View File
@@ -1,57 +0,0 @@
import torch
import torch.nn as nn
class DummyModel(nn.Module):
def __init__(self, d_in=2048, n_heads=128, softmax_scale=0.5):
super().__init__()
self.weights_proj = nn.Linear(d_in, 1024)
self.n_heads = n_heads
self.softmax_scale = softmax_scale
def _get_logits_head_gate_orig(self, x: torch.Tensor, q_scale: torch.Tensor):
weights = self.weights_proj(x)
weights = weights * self.n_heads**-0.5
q_scale = q_scale.unsqueeze(1) # (B,1,1)
weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale
return weights
def _get_logits_head_gate_opt(self, x: torch.Tensor, q_scale: torch.Tensor):
weights = self.weights_proj(x)
q_scale = q_scale.unsqueeze(1) # (B,1,1)
scale_const = self.n_heads**-0.5 * q_scale * self.softmax_scale # (B,1,1)
weights = weights.unsqueeze(-1) * scale_const # (B,1024,1)
return weights
def main():
torch.manual_seed(0)
model = DummyModel(d_in=2048, n_heads=128, softmax_scale=0.5)
x = torch.randn(128, 2048) # batch=128, d_in=2048
q_scale = torch.randn(128, 1)
import time
start = time.time()
for _ in range(1000):
out_orig = model._get_logits_head_gate_orig(x, q_scale)
print("Original version time:", time.time() - start)
start = time.time()
for _ in range(1000):
out_opt = model._get_logits_head_gate_opt(x, q_scale)
print("Optimized version time:", time.time() - start)
print("Difference:", (out_orig - out_opt).abs().max().item())
assert torch.allclose(out_orig, out_opt), "Mismatch between original and optimized"
if __name__ == "__main__":
main()
"""
Original version time: 0.49235057830810547
Optimized version time: 0.4087331295013428
Difference: 1.4901161193847656e-08
"""
python/sglang/test/longbench_v2/__init__.py
-1
View File
@@ -1 +0,0 @@
"""LongBench-v2 auxiliary utilities and validation scripts."""
python/sglang/test/longbench_v2/longbench_v2_evaluation.md
-217
View File
@@ -1,217 +0,0 @@
# LongBench-v2 Evaluation Guide
## Overview
LongBench-v2 is a benchmark designed to assess the ability of Large Language Models (LLMs) to handle long-context problems requiring deep understanding and reasoning across real-world multitasks. This guide explains how to use SGLang's LongBench-v2 evaluation utilities.
## Features
- **Context Length**: 8k to 2M words (majority under 128k)
- **Task Categories**: 6 major categories with 503 challenging multiple-choice questions
- **Difficulty**: Challenging enough that human experts achieve only 53.7% accuracy
- **Format**: All questions are multiple-choice for reliable evaluation
## Task Categories
1. **Single-Document QA**: Question answering within a single long document
2. **Multi-Document QA**: Cross-document reasoning and synthesis
3. **Long In-Context Learning**: Few-shot learning with long examples
4. **Long-Dialogue History**: Understanding long conversation histories
5. **Code Repository Understanding**: Analysis of large codebases
6. **Long Structured Data**: Comprehension of tables, JSON, and structured data
## Quick Start
### Basic Usage
```python
from sglang.test.simple_eval_longbench_v2 import LongBenchV2Eval
from sglang.test.simple_eval_common import ChatCompletionSampler
# Initialize evaluator with HuggingFace dataset
eval_obj = LongBenchV2Eval(
data_source="THUDM/LongBench-v2",
num_examples=10, # Limit for testing
num_threads=4
)
# Create sampler (pointing to your SGLang server)
sampler = ChatCompletionSampler(
base_url="http://localhost:30000/v1",
model="your-model-name"
)
# Run evaluation
result = eval_obj(sampler)
print(f"Overall Score: {result.score:.3f}")
print(f"Metrics: {result.metrics}")
```
### Using the Command Line
```bash
# Basic evaluation
python -m sglang.test.run_eval \
--eval-name longbench_v2 \
--port 30000 \
--num-examples 50
# Evaluate specific categories
python -m sglang.test.run_eval \
--eval-name longbench_v2 \
--categories "single_document_qa,multi_document_qa" \
--port 30000
# Filter by context length
python -m sglang.test.run_eval \
--eval-name longbench_v2 \
--max-context-length 100000 \
--min-context-length 10000 \
--port 30000
```
## Advanced Configuration
### Category-Specific Evaluation
```python
# Evaluate only specific task categories
eval_obj = LongBenchV2Eval(
data_source="THUDM/LongBench-v2",
categories=[
"single_document_qa",
"code_repo_understanding"
]
)
```
### Context Length Filtering
```python
# Focus on medium-length contexts
eval_obj = LongBenchV2Eval(
data_source="THUDM/LongBench-v2",
min_context_length=32000, # characters
max_context_length=128000 # characters
)
```
### Using Local Dataset
```python
# Load from local JSON file
eval_obj = LongBenchV2Eval(
data_source="/path/to/longbench_v2.json",
num_examples=100
)
# Load from local CSV file
eval_obj = LongBenchV2Eval(
data_source="/path/to/longbench_v2.csv"
)
```
## Dataset Format
The expected format for LongBench-v2 examples:
```json
{
"context": "Long context text...",
"question": "Question about the context",
"A": "First choice",
"B": "Second choice",
"C": "Third choice",
"D": "Fourth choice",
"answer": "A",
"category": "single_document_qa"
}
```
Alternative format with choices as list:
```json
{
"context": "Long context text...",
"question": "Question about the context",
"choices": ["First choice", "Second choice", "Third choice", "Fourth choice"],
"answer": "A",
"category": "multi_document_qa"
}
```
## Metrics and Scoring
### Overall Metrics
- **score**: Overall accuracy across all examples
- **chars**: Average response length in characters
### Category-Specific Metrics
Each task category gets its own metric:
- `single_document_qa`: Accuracy on single-document QA tasks
- `multi_document_qa`: Accuracy on multi-document QA tasks
- `long_in_context_learning`: Accuracy on in-context learning tasks
- `long_dialogue_history`: Accuracy on dialogue understanding tasks
- `code_repo_understanding`: Accuracy on code analysis tasks
- `long_structured_data`: Accuracy on structured data tasks
### Context Length Metrics
- `short_context`: Accuracy on contexts < 32k characters
- `medium_context`: Accuracy on contexts 32k-128k characters
- `long_context`: Accuracy on contexts > 128k characters
- `difficulty_easy` / `difficulty_hard`: Accuracy grouped by dataset difficulty labels
## Performance Considerations
### Memory Usage
LongBench-v2 contains very long contexts (up to 2M words). Consider:
1. **GPU Memory**: Ensure your model can handle the context lengths
2. **Batch Size**: Use smaller batch sizes for longer contexts
3. **Parallel Processing**: Adjust `num_threads` based on available resources
### Evaluation Time
- Full evaluation (503 examples) can take several hours
- Use `num_examples` parameter to limit evaluation size during development
- Consider filtering by context length to focus on specific ranges
## Troubleshooting
### Common Issues
1. **Out of Memory**: Reduce context length limits or batch size
2. **Slow Evaluation**: Increase `num_threads` or reduce `num_examples`
3. **Dataset Loading**: Ensure `datasets` library is installed for HuggingFace integration
### Installation Requirements
```bash
pip install datasets # For HuggingFace dataset support
```
## Example Results
Typical performance ranges for different model sizes:
- **Small models (7B)**: 35-45% accuracy
- **Medium models (13-30B)**: 45-55% accuracy
- **Large models (70B+)**: 55-65% accuracy
- **Human experts**: 53.7% accuracy
## Citation
If you use LongBench-v2 in your research, please cite:
```bibtex
@article{bai2024longbench,
title={LongBench v2: Towards Deeper Understanding and Reasoning on Realistic Long-Context Multitasks},
author={Bai, Yushi and Tu, Shangqing and Zhang, Jiajie and Peng, Hao and Wang, Xiaozhi and Lv, Xin and Cao, Shulin and Xu, Jiazheng and Hou, Lei and Dong, Yuxiao and Tang, Jie and Li, Juanzi},
journal={arXiv preprint arXiv:2412.15204},
year={2024}
}
```
python/sglang/test/longbench_v2/test_longbench_v2_eval.py
-238
View File
@@ -1,238 +0,0 @@
"""
Test cases for LongBench-v2 evaluation utility.
"""
import json
import os
import tempfile
from sglang.test.simple_eval_longbench_v2 import (
LongBenchV2Eval,
extract_longbench_v2_answer,
format_longbench_v2_question,
)
def test_format_longbench_v2_question():
"""Test the official LongBench-v2 question formatting."""
sample_row = {
"context": "This is a sample context about environmental issues.",
"question": "What is the main theme?",
"A": "Technology",
"B": "Environment",
"C": "Economics",
"D": "Politics",
"answer": "B",
}
formatted = format_longbench_v2_question(sample_row)
# Verify official template structure
assert "This is a sample context about environmental issues." in formatted
assert (
"What is the correct answer to this question: What is the main theme?"
in formatted
)
assert "(A) Technology" in formatted
assert "(B) Environment" in formatted
assert "(C) Economics" in formatted
assert "(D) Politics" in formatted
assert "The correct answer is" in formatted
print("✓ Question formatting works correctly")
def test_extract_longbench_v2_answer():
"""Test the official LongBench-v2 answer extraction."""
# Test official format: "The correct answer is (A)"
response1 = "After analyzing the context, The correct answer is (B)."
assert extract_longbench_v2_answer(response1) == "B"
# Test alternative format: "The correct answer is A"
response2 = "Based on the evidence, The correct answer is C."
assert extract_longbench_v2_answer(response2) == "C"
# Test with asterisks
response3 = "*The correct answer is (D)*"
assert extract_longbench_v2_answer(response3) == "D"
# Test fallback to standard pattern
response4 = "I think the answer is A."
assert extract_longbench_v2_answer(response4) == "A"
# Test no answer
response5 = "I'm not sure about this."
assert extract_longbench_v2_answer(response5) is None
print("✓ Answer extraction works correctly")
def test_longbench_v2_eval_initialization():
"""Test LongBench-v2 evaluation class initialization."""
# Create a temporary JSON file with sample data
sample_data = [
{
"_id": "test_001",
"domain": "single_document_qa",
"question": "What is X?",
"choice_A": "Option A1",
"choice_B": "Option B1",
"choice_C": "Option C1",
"choice_D": "Option D1",
"answer": "A",
"context": "Context 1",
},
{
"_id": "test_002",
"domain": "multi_document_qa",
"question": "What is Y?",
"A": "Option A2",
"B": "Option B2",
"C": "Option C2",
"D": "Option D2",
"answer": "B",
"context": "Context 2",
},
]
with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
json.dump(sample_data, f)
temp_file = f.name
try:
# Test initialization with new data_source parameter
eval_instance = LongBenchV2Eval(data_source=temp_file, num_examples=1)
assert len(eval_instance.examples) == 1
first_example = eval_instance.examples[0]
assert first_example.get("category") in {
"single_document_qa",
"multi_document_qa",
}
assert first_example.get("A") in {"Option A1", "Option A2"}
print("✓ Evaluation class initialization works correctly")
finally:
os.unlink(temp_file)
def test_category_filtering():
"""Ensure category filtering keeps only requested domains."""
sample_data = [
{
"_id": "test_001",
"domain": "single_document_qa",
"question": "What is X?",
"choice_A": "Option A1",
"choice_B": "Option B1",
"choice_C": "Option C1",
"choice_D": "Option D1",
"answer": "A",
"context": "Context 1",
},
{
"_id": "test_002",
"domain": "multi_document_qa",
"question": "What is Y?",
"choice_A": "Option A2",
"choice_B": "Option B2",
"choice_C": "Option C2",
"choice_D": "Option D2",
"answer": "B",
"context": "Context 2",
},
]
with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
json.dump(sample_data, f)
temp_file = f.name
try:
eval_instance = LongBenchV2Eval(
data_source=temp_file,
categories=["multi_document_qa"],
)
assert len(eval_instance.examples) == 1
assert eval_instance.examples[0]["category"] == "multi_document_qa"
print("✓ Category filtering works correctly")
finally:
os.unlink(temp_file)
def test_difficulty_metrics():
"""Validate that difficulty-specific metrics are recorded."""
sample_data = [
{
"_id": "easy_001",
"domain": "single_document_qa",
"difficulty": "easy",
"question": "Easy question?",
"choice_A": "Correct",
"choice_B": "Wrong",
"choice_C": "Wrong",
"choice_D": "Wrong",
"answer": "A",
"context": "Easy context",
},
{
"_id": "hard_001",
"domain": "single_document_qa",
"difficulty": "hard",
"question": "Hard question?",
"choice_A": "Wrong",
"choice_B": "Correct",
"choice_C": "Wrong",
"choice_D": "Wrong",
"answer": "B",
"context": "Hard context",
},
]
with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
json.dump(sample_data, f)
temp_file = f.name
class FixedSampler: # noqa: D401 - simple helper
"""Mock sampler returning the correct answer based on question text."""
def _pack_message(self, content: str, role: str):
return {"content": content, "role": role}
def __call__(self, messages):
prompt = messages[0]["content"]
if "Easy question" in prompt:
return "The correct answer is (A)"
return "The correct answer is (B)"
try:
eval_instance = LongBenchV2Eval(data_source=temp_file, num_threads=1)
result = eval_instance(FixedSampler())
assert result.metrics.get("difficulty_easy") == 1.0
assert result.metrics.get("difficulty_hard") == 1.0
print("✓ Difficulty metrics recorded correctly")
finally:
os.unlink(temp_file)
def main():
"""Run all tests."""
print("Testing simplified LongBench-v2 evaluation utility...\n")
test_format_longbench_v2_question()
test_extract_longbench_v2_answer()
test_longbench_v2_eval_initialization()
test_category_filtering()
test_difficulty_metrics()
print("\n" + "=" * 50)
print("✅ ALL TESTS PASSED!")
print("The simplified implementation follows SGLang patterns")
print("while maintaining LongBench-v2 compatibility.")
print("=" * 50)
if __name__ == "__main__":
main()
python/sglang/test/longbench_v2/validate_longbench_v2.py
-337
View File
@@ -1,337 +0,0 @@
#!/usr/bin/env python3
"""
Validation script for LongBench-v2 implementation.
This script validates our implementation against official LongBench-v2 format and benchmarks.
"""
import json
import os
import tempfile
from typing import Any, Dict, List
from sglang.test.simple_eval_longbench_v2 import (
LongBenchV2Eval,
extract_longbench_v2_answer,
format_longbench_v2_question,
)
def create_sample_official_data() -> List[Dict[str, Any]]:
"""Create sample data in official LongBench-v2 format for validation."""
return [
{
"_id": "test_001",
"domain": "science",
"sub_domain": "physics",
"difficulty": "hard",
"length": "medium",
"question": "What is the fundamental force responsible for holding atomic nuclei together?",
"choice_A": "Electromagnetic force",
"choice_B": "Strong nuclear force",
"choice_C": "Weak nuclear force",
"choice_D": "Gravitational force",
"answer": "B",
"context": "Nuclear physics studies the components and behavior of atomic nuclei. "
* 100,
},
{
"_id": "test_002",
"domain": "literature",
"sub_domain": "analysis",
"difficulty": "hard",
"length": "long",
"question": "What literary technique is primarily used in the given passage?",
"choice_A": "Metaphor",
"choice_B": "Alliteration",
"choice_C": "Symbolism",
"choice_D": "Irony",
"answer": "C",
"context": "Literary analysis involves examining various techniques authors use to convey meaning. "
* 150,
},
{
"_id": "test_003",
"domain": "code",
"sub_domain": "algorithms",
"difficulty": "easy",
"length": "short",
"question": "What is the time complexity of binary search?",
"choice_A": "O(n)",
"choice_B": "O(log n)",
"choice_C": "O(n²)",
"choice_D": "O(1)",
"answer": "B",
"context": "Binary search is a fundamental algorithm in computer science. "
* 50,
},
]
def create_alternative_format_data() -> List[Dict[str, Any]]:
"""Create sample data in alternative format (choices as list) for validation."""
return [
{
"_id": "alt_001",
"question": "What is 2 + 2?",
"choices": ["3", "4", "5", "6"],
"answer": "B",
"category": "single_document_qa",
"context": "Basic arithmetic operations. " * 30,
},
{
"_id": "alt_002",
"question": "What color is the sky?",
"choices": ["Red", "Blue", "Green", "Yellow"],
"answer": "B",
"category": "multi_document_qa",
"context": "Color perception and atmospheric science. " * 40,
},
]
class MockSampler:
"""Mock sampler for testing that returns predictable responses."""
def __init__(self, responses: Dict[str, str]):
self.responses = responses
self.call_count = 0
def _pack_message(self, content: str, role: str) -> Dict[str, str]:
return {"content": content, "role": role}
def __call__(self, messages: List[Dict[str, str]]) -> str:
"""Return a mock response based on the question content."""
prompt = messages[0]["content"]
self.call_count += 1
if "atomic nuclei" in prompt:
return "The correct answer is (B)"
if "literary technique" in prompt:
return "The correct answer is (C)"
if "binary search" in prompt:
return "The correct answer is (B)"
if "2 + 2" in prompt:
return "The correct answer is (B)"
if "color is the sky" in prompt:
return "The correct answer is (B)"
if "Complex reasoning question" in prompt:
return "The correct answer is (B)"
return "The correct answer is (A)"
def test_format_compatibility() -> None:
"""Test that our implementation handles official LongBench-v2 format correctly."""
print("Testing official format compatibility...")
official_sample = {
"context": "Test context",
"question": "Test question?",
"choice_A": "Option A",
"choice_B": "Option B",
"choice_C": "Option C",
"choice_D": "Option D",
"answer": "A",
}
formatted = format_longbench_v2_question(official_sample)
assert "Test context" in formatted
assert "Test question?" in formatted
assert "(A) Option A" in formatted
assert "(B) Option B" in formatted
assert "The correct answer is" in formatted
print("✓ Official format compatibility verified")
alt_sample = {
"context": "Test context",
"question": "Test question?",
"choices": ["Option A", "Option B", "Option C", "Option D"],
"answer": "A",
}
formatted_alt = format_longbench_v2_question(alt_sample)
assert "Test context" in formatted_alt
assert "(A) Option A" in formatted_alt
print("✓ Alternative format compatibility verified")
def test_answer_extraction() -> None:
"""Test answer extraction with various response formats."""
print("Testing answer extraction...")
test_cases = [
("The correct answer is (B)", "B"),
("The correct answer is C", "C"),
("After analysis, The correct answer is (D)", "D"),
("*The correct answer is (A)*", "A"),
("I think the answer is B", "B"),
("No clear answer here", None),
]
for response, expected in test_cases:
result = extract_longbench_v2_answer(response)
assert (
result == expected
), f"Failed for '{response}': got {result}, expected {expected}"
print("✓ Answer extraction verified")
def test_evaluation_pipeline() -> None:
"""Test the complete evaluation pipeline with mock data."""
print("Testing evaluation pipeline...")
official_data = create_sample_official_data()
with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
json.dump(official_data, f)
temp_file = f.name
try:
eval_obj = LongBenchV2Eval(data_source=temp_file, num_examples=3, num_threads=1)
mock_sampler = MockSampler({})
result = eval_obj(mock_sampler)
assert result.score > 0, "Expected positive score"
assert len(result.convos) == 3, "Expected 3 evaluated conversations"
assert "chars" in result.metrics, "Expected chars metric"
print(f"✓ Evaluation pipeline verified (score: {result.score:.3f})")
finally:
os.unlink(temp_file)
def test_category_filtering() -> None:
"""Test category-based filtering functionality."""
print("Testing category filtering...")
alt_data = create_alternative_format_data()
with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
json.dump(alt_data, f)
temp_file = f.name
try:
eval_obj = LongBenchV2Eval(
data_source=temp_file,
categories=["single_document_qa"],
num_threads=1,
)
assert len(eval_obj.examples) == 1, "Expected 1 example after filtering"
assert eval_obj.examples[0]["category"] == "single_document_qa"
print("✓ Category filtering verified")
finally:
os.unlink(temp_file)
def run_accuracy_benchmark() -> None:
"""Run a small accuracy benchmark to compare with expected performance."""
print("Running accuracy benchmark...")
benchmark_data = [
{
"_id": "bench_001",
"question": "Complex reasoning question",
"choice_A": "Incorrect option 1",
"choice_B": "Correct answer",
"choice_C": "Incorrect option 2",
"choice_D": "Incorrect option 3",
"answer": "B",
"context": "This requires careful analysis. " * 200,
}
] * 10
with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
json.dump(benchmark_data, f)
temp_file = f.name
try:
eval_obj = LongBenchV2Eval(data_source=temp_file, num_threads=1)
perfect_sampler = MockSampler({})
result = eval_obj(perfect_sampler)
print(f"✓ Benchmark completed - Perfect sampler accuracy: {result.score:.3f}")
print(f" Total examples: {len(result.convos)}")
print(f" Average response length: {result.metrics.get('chars', 0):.1f} chars")
assert (
result.score == 1.0
), f"Perfect sampler should get 100% accuracy, got {result.score:.3f}"
finally:
os.unlink(temp_file)
def generate_comparison_report() -> None:
"""Generate a comparison report with official benchmarks."""
print("\n" + "=" * 60)
print("LONGBENCH-V2 IMPLEMENTATION VALIDATION REPORT")
print("=" * 60)
print("\n📊 OFFICIAL BENCHMARK RESULTS (for comparison):")
print(" • Human Experts: 53.7% accuracy (15-min constraint)")
print(" • Best Direct Model: 50.1% accuracy")
print(" • o1-preview (with CoT): 57.7% accuracy")
print(" • Dataset: 503 questions, 8k-2M word contexts")
print("\n✅ IMPLEMENTATION VALIDATION:")
print(" • Format compatibility: VERIFIED")
print(" • Answer extraction: VERIFIED")
print(" • Evaluation pipeline: VERIFIED")
print(" • Category filtering: VERIFIED")
print(" • Perfect sampler benchmark: VERIFIED (100% accuracy)")
print("\n🔍 TECHNICAL VERIFICATION:")
print(" • Handles official choice_A/B/C/D format: ✓")
print(" • Handles alternative choices list format: ✓")
print(" • Official answer extraction patterns: ✓")
print(" • Context length filtering: ✓")
print(" • HuggingFace dataset integration: ✓")
print(" • SGLang evaluation framework compliance: ✓")
print("\n📈 EXPECTED PERFORMANCE RANGE:")
print(" • Small models (7B): 35-45% accuracy")
print(" • Medium models (13-30B): 45-55% accuracy")
print(" • Large models (70B+): 55-65% accuracy")
print(
" • Note: Actual results depend on model capabilities and context length handling"
)
print("\n✨ IMPLEMENTATION HIGHLIGHTS:")
print(" • Follows official LongBench-v2 evaluation methodology")
print(" • Compatible with SGLang's existing evaluation patterns")
print(" • Supports multiple data sources (HF, JSON, CSV)")
print(" • Robust error handling and fallback mechanisms")
print(" • Comprehensive filtering and configuration options")
print("\n" + "=" * 60)
print("VALIDATION COMPLETE - IMPLEMENTATION READY FOR USE")
print("=" * 60)
def main() -> None:
"""Run all validation tests."""
print("🔍 Starting LongBench-v2 Implementation Validation...\n")
try:
test_format_compatibility()
test_answer_extraction()
test_evaluation_pipeline()
test_category_filtering()
run_accuracy_benchmark()
generate_comparison_report()
print("\n🎉 All validation tests passed successfully!")
print("The LongBench-v2 implementation is working correctly and ready for use.")
except Exception as exc: # pragma: no cover - debug helper
print(f"\n❌ Validation failed: {exc}")
raise
if __name__ == "__main__":
main()
python/sglang/test/longbench_v2/validate_longbench_v2_standalone.py
-306
View File
@@ -1,306 +0,0 @@
#!/usr/bin/env python3
"""
Standalone validation script for LongBench-v2 implementation.
Tests core functionality without requiring full SGLang dependencies.
"""
import json
import os
import re
import tempfile
from typing import Any, Dict, List, Optional
ANSWER_PATTERN_MULTICHOICE = r"(?i)(?:the\s+)?(?:correct\s+)?(?:answer\s+)?(?:is\s+)?(?:\(?\s*)?([A-D])(?:\s*\)?)"
def format_longbench_v2_question(row: Dict[str, Any]) -> str:
"""Format a LongBench-v2 question using the official template."""
context = row.get("context", "")
question = row.get("question", "")
if "choices" in row:
choices = row["choices"]
choice_A = choices[0] if len(choices) > 0 else ""
choice_B = choices[1] if len(choices) > 1 else ""
choice_C = choices[2] if len(choices) > 2 else ""
choice_D = choices[3] if len(choices) > 3 else ""
else:
choice_A = row.get("choice_A", row.get("A", ""))
choice_B = row.get("choice_B", row.get("B", ""))
choice_C = row.get("choice_C", row.get("C", ""))
choice_D = row.get("choice_D", row.get("D", ""))
prompt = f"""{context.strip()}
What is the correct answer to this question: {question.strip()}
Choices:
(A) {choice_A.strip()}
(B) {choice_B.strip()}
(C) {choice_C.strip()}
(D) {choice_D.strip()}
The correct answer is"""
return prompt
def extract_longbench_v2_answer(response: str) -> Optional[str]:
"""Extract answer from model response using official LongBench-v2 method."""
response = response.replace("*", "")
match = re.search(r"The correct answer is \(([A-D])\)", response, re.IGNORECASE)
if match:
return match.group(1).upper()
match = re.search(r"The correct answer is ([A-D])", response, re.IGNORECASE)
if match:
return match.group(1).upper()
match = re.search(ANSWER_PATTERN_MULTICHOICE, response)
if match:
return match.group(1).upper()
return None
def create_official_format_samples() -> List[Dict[str, Any]]:
"""Create test samples in official LongBench-v2 format."""
return [
{
"_id": "official_001",
"domain": "science",
"sub_domain": "physics",
"difficulty": "hard",
"length": "medium",
"question": "What force holds atomic nuclei together?",
"choice_A": "Electromagnetic force",
"choice_B": "Strong nuclear force",
"choice_C": "Weak nuclear force",
"choice_D": "Gravitational force",
"answer": "B",
"context": "Nuclear physics studies atomic nuclei behavior." * 50,
},
{
"_id": "official_002",
"domain": "literature",
"sub_domain": "analysis",
"difficulty": "hard",
"length": "long",
"question": "What literary device is primarily demonstrated?",
"choice_A": "Metaphor",
"choice_B": "Alliteration",
"choice_C": "Symbolism",
"choice_D": "Irony",
"answer": "C",
"context": "The recurring image of the white whale represents much more than a literal creature."
* 80,
},
]
def create_alternative_format_samples() -> List[Dict[str, Any]]:
"""Create test samples in alternative format."""
return [
{
"_id": "alt_001",
"question": "What is 2 + 2?",
"choices": ["3", "4", "5", "6"],
"answer": "B",
"category": "single_document_qa",
"context": "Basic arithmetic: Addition is a fundamental mathematical operation."
* 30,
}
]
def test_format_compatibility() -> None:
"""Test format compatibility with both official and alternative formats."""
print("Testing format compatibility...")
official_sample = create_official_format_samples()[0]
formatted = format_longbench_v2_question(official_sample)
assert "Nuclear physics studies" in formatted
assert "(A) Electromagnetic force" in formatted
assert "(B) Strong nuclear force" in formatted
assert "The correct answer is" in formatted
print("✓ Official format (choice_A/B/C/D) working correctly")
alt_sample = create_alternative_format_samples()[0]
formatted_alt = format_longbench_v2_question(alt_sample)
assert "What is 2 + 2?" in formatted_alt
assert "(B) 4" in formatted_alt
print("✓ Alternative format (choices list) working correctly")
def test_answer_extraction() -> None:
"""Test answer extraction patterns."""
print("Testing answer extraction...")
test_cases = [
("The correct answer is (B)", "B"),
("The correct answer is C", "C"),
("After analysis, The correct answer is (D)", "D"),
("*The correct answer is (A)*", "A"),
("I believe the answer is B", "B"),
("Looking at this, A seems correct", "A"),
("The answer should be (C)", "C"),
("No clear pattern here", None),
]
for response, expected in test_cases:
result = extract_longbench_v2_answer(response)
assert (
result == expected
), f"Failed for '{response}': got {result}, expected {expected}"
print("✓ Answer extraction patterns working correctly")
def test_data_loading_simulation() -> None:
"""Simulate data loading and processing."""
print("Testing data loading simulation...")
test_data = create_official_format_samples() + create_alternative_format_samples()
with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
json.dump(test_data, f)
temp_file = f.name
try:
with open(temp_file, "r", encoding="utf-8") as fh:
loaded_data = json.load(fh)
assert len(loaded_data) == 3
assert loaded_data[0]["_id"] == "official_001"
assert "choices" in loaded_data[2]
print("✓ JSON data loading working correctly")
finally:
os.unlink(temp_file)
def run_accuracy_simulation() -> None:
"""Simulate accuracy testing with perfect responses."""
print("Running accuracy simulation...")
samples = create_official_format_samples()
correct_responses = {
"official_001": "The correct answer is (B)",
"official_002": "The correct answer is (C)",
}
total_score = 0
for sample in samples:
formatted = format_longbench_v2_question(sample)
response = correct_responses[sample["_id"]]
extracted = extract_longbench_v2_answer(response)
expected = sample["answer"]
score = 1.0 if extracted == expected else 0.0
total_score += score
print(f" Question {sample['_id']}: {extracted} == {expected} -> {score}")
accuracy = total_score / len(samples)
print(f"✓ Simulation accuracy: {accuracy:.3f} (expected: 1.0)")
assert accuracy == 1.0, "Perfect simulation should achieve 100% accuracy"
def generate_validation_report() -> None:
"""Generate comprehensive validation report."""
print("\n" + "=" * 70)
print("LONGBENCH-V2 IMPLEMENTATION VALIDATION REPORT")
print("=" * 70)
print("\n📚 OFFICIAL LONGBENCH-V2 BENCHMARK:")
print(" • Dataset: 503 multiple-choice questions")
print(" • Context length: 8k to 2M words (majority < 128k)")
print(" • Categories: 6 major task categories")
print(" • Human expert accuracy: 53.7%")
print(" • Best direct model: 50.1% accuracy")
print(" • o1-preview (with CoT): 57.7% accuracy")
print("\n✅ IMPLEMENTATION VERIFICATION:")
print(" • Official format compatibility: VERIFIED")
print(" • Alternative format support: VERIFIED")
print(" • Answer extraction patterns: VERIFIED")
print(" • Data loading mechanisms: VERIFIED")
print(" • Accuracy calculation: VERIFIED")
print("\n🔧 TECHNICAL COMPLIANCE:")
print(" • Official question template: ✓")
print(" • Multiple answer extraction patterns: ✓")
print(" • HuggingFace dataset integration: ✓")
print(" • CSV/JSON file support: ✓")
print(" • Category-based filtering: ✓")
print(" • Context length filtering: ✓")
print("\n📊 EXPECTED PERFORMANCE BENCHMARKS:")
print(" Model Category | Expected Accuracy")
print(" ----------------------- | ----------------")
print(" Small models (7B) | 35-45%")
print(" Medium models (13-30B) | 45-55%")
print(" Large models (70B+) | 55-65%")
print(" Human experts | 53.7%")
print(" Advanced reasoning | 57.7%")
print("\n🏗️ IMPLEMENTATION FEATURES:")
print(" • Multiple data source support (HuggingFace, JSON, CSV)")
print(" • Robust answer extraction with fallback patterns")
print(" • Category-based evaluation filtering")
print(" • Context length range filtering")
print(" • SGLang evaluation framework integration")
print(" • Comprehensive error handling")
print("\n📋 FORMAT COMPATIBILITY:")
print(" • Official format: choice_A, choice_B, choice_C, choice_D")
print(' • Alternative format: choices = ["A", "B", "C", "D"]')
print(' • Answer format: "A", "B", "C", or "D"')
print(" • Context field: Long-form text content")
print("\n🚀 USAGE EXAMPLES:")
print(" # Command line usage:")
print(" python -m sglang.test.run_eval --eval-name longbench_v2 --port 30000")
print(" ")
print(" # Python API usage:")
print(" from sglang.test.simple_eval_longbench_v2 import LongBenchV2Eval")
print(" eval_obj = LongBenchV2Eval(data_source='THUDM/LongBench-v2')")
print(" result = eval_obj(sampler)")
print("\n🎯 ACCURACY COMPARISON GUIDANCE:")
print(" • Run evaluation on a subset for validation")
print(" • Compare results within expected performance ranges")
print(" • Verify answer extraction matches official pattern")
print(" • Confirm handling of long-context inputs")
print("\n" + "=" * 70)
print("VALIDATION STATUS: ✅ PASSED - IMPLEMENTATION READY FOR PRODUCTION")
print("=" * 70)
def main() -> bool:
"""Run complete validation suite."""
print("🔍 LongBench-v2 Implementation Validation Starting...\n")
try:
test_format_compatibility()
test_answer_extraction()
test_data_loading_simulation()
run_accuracy_simulation()
generate_validation_report()
print("\n🎉 All validation tests completed successfully!")
print("Implementation is ready for accuracy comparison testing.")
return True
except Exception as exc: # pragma: no cover - debug helper
print(f"\n❌ Validation failed: {exc}")
raise
if __name__ == "__main__":
success = main()
raise SystemExit(0 if success else 1)
python/sglang/test/speculative/test_spec_utils.py
-348
View File
@@ -1,348 +0,0 @@
import unittest
import numpy as np
import torch
from sglang.srt.mem_cache.memory_pool import copy_all_layer_kv_cache_tiled
from sglang.srt.speculative.spec_utils import assign_draft_cache_locs
from sglang.srt.utils import next_power_of_2
BYTES_PER_TILE = 128
class TestSpecUtils(unittest.TestCase):
def setUp(self):
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.data_ptrs = torch.zeros(2, 1, dtype=torch.uint64, device=self.device)
self.k_cache = [
torch.zeros((100, 1, 1), dtype=torch.float32, device=self.device)
]
self.v_cache = [
torch.zeros((100, 1, 1), dtype=torch.float32, device=self.device)
]
self.k_cache[0][:11, 0, 0] = torch.tensor(
[0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0],
dtype=torch.float32,
device=self.device,
)
self.v_cache[0][:11, 0, 0] = torch.tensor(
[-0.0, -0.1, -0.2, -0.3, -0.4, -0.5, -0.6, -0.7, -0.8, -0.9, -1.0],
dtype=torch.float32,
device=self.device,
)
self.data_ptrs[0, 0] = self.k_cache[0].data_ptr()
self.data_ptrs[1, 0] = self.v_cache[0].data_ptr()
self.data_strides = torch.tensor(
[
np.prod(x.shape[1:]) * x.dtype.itemsize
for x in self.k_cache + self.v_cache
],
device=self.device,
dtype=torch.int64,
)
def test_assign_draft_cache_locs_single_seq(self):
# Testing Setup: req_to_token starting from 4
# 4,5,6,7,{8,9,10}, 8,9,10 is the last partial page, 3 tokens < page_size=4
# next kv cache will be stored starting 11,12,13...
device = self.device
num_seqs = 1
page_size = 4
speculative_num_steps = 5
topk = 8
seq_lens_num = 7
extend_lens_num = 61 # includes the duplicated last page
req_pool_indices = torch.arange(num_seqs, dtype=torch.int32, device=device)
req_to_token = torch.zeros((num_seqs, 100), dtype=torch.int32, device=device)
req_to_token[0, :seq_lens_num] = torch.tensor(
[4, 5, 6, 7, 8, 9, 10], device=device
)
seq_lens = torch.tensor([seq_lens_num], dtype=torch.int32, device=device)
extend_lens = torch.tensor([extend_lens_num], dtype=torch.int32, device=device)
num_new_pages_per_topk = torch.tensor([2], dtype=torch.int32, device=device)
out_cache_loc = torch.arange(11, 11 + extend_lens_num, device=device)
last_page_lens = torch.tensor([3], dtype=torch.int32, device=device)
last_page_lens_cumsum = torch.cumsum(last_page_lens, dim=0)
duplicate_cache_len = last_page_lens.sum().item() * (topk - 1)
target_cache_loc = torch.zeros(
duplicate_cache_len, dtype=torch.int32, device=device
)
source_cache_loc = torch.zeros(
duplicate_cache_len, dtype=torch.int32, device=device
)
assign_draft_cache_locs[(num_seqs,)](
req_pool_indices,
req_to_token,
seq_lens,
extend_lens,
num_new_pages_per_topk,
out_cache_loc,
source_cache_loc,
target_cache_loc,
last_page_lens_cumsum,
duplicate_cache_len,
req_to_token.shape[1],
topk,
speculative_num_steps,
page_size,
next_power_of_2(num_seqs),
next_power_of_2(speculative_num_steps + page_size),
)
out_cache_loc = out_cache_loc[: num_seqs * topk * speculative_num_steps]
expected_source_cache_loc = torch.tensor(
[8, 9, 10] * (topk - 1), device=device, dtype=torch.int32
)
assert torch.allclose(source_cache_loc, expected_source_cache_loc)
copy_all_layer_kv_cache_tiled[(len(self.data_ptrs),)](
self.data_ptrs,
self.data_strides,
target_cache_loc,
source_cache_loc,
len(target_cache_loc),
next_power_of_2(len(target_cache_loc)),
BYTES_PER_TILE,
)
assert torch.allclose(
self.k_cache[0][16:19, 0, 0],
torch.tensor(
[0.8, 0.9, 1.0],
dtype=torch.float32,
device=device,
),
)
assert torch.allclose(
self.v_cache[0][16:19, 0, 0],
torch.tensor(
[-0.8, -0.9, -1.0],
dtype=torch.float32,
device=device,
),
)
def test_assign_draft_cache_locs_multi_seq(self):
device = self.device
num_seqs = 3
page_size = 4
speculative_num_steps = 5
topk = 8
req_pool_indices = torch.arange(num_seqs, dtype=torch.int32, device=device)
req_to_token = torch.zeros((num_seqs, 100), dtype=torch.int32, device=device)
seq_lens = torch.tensor([8, 7, 5], dtype=torch.int32, device=device)
extend_lens = torch.tensor([64, 64, 64], dtype=torch.int32, device=device)
num_new_pages_per_topk = torch.tensor(
[2, 2, 2], dtype=torch.int32, device=device
)
req_to_token = torch.zeros((num_seqs, 100), dtype=torch.int32, device=device)
req_to_token[0, :8] = torch.tensor([4, 5, 6, 7, 8, 9, 10, 11], device=device)
req_to_token[1, :7] = torch.tensor([4, 5, 6, 7, 8, 9, 10], device=device)
req_to_token[2, :5] = torch.tensor([4, 5, 6, 7, 8], device=device)
last_page_lens = torch.tensor([0, 3, 1], dtype=torch.int32, device=device)
last_page_lens_cumsum = torch.cumsum(last_page_lens, dim=0)
duplicate_cache_len = last_page_lens.sum().item() * (topk - 1)
out_cache_loc = torch.arange(
12, 12 + torch.sum(extend_lens), dtype=torch.int32, device=device
)
target_cache_loc = torch.zeros(
duplicate_cache_len, dtype=torch.int32, device=device
)
source_cache_loc = torch.zeros(
duplicate_cache_len, dtype=torch.int32, device=device
)
assign_draft_cache_locs[(num_seqs,)](
req_pool_indices,
req_to_token,
seq_lens,
extend_lens,
num_new_pages_per_topk,
out_cache_loc,
source_cache_loc,
target_cache_loc,
last_page_lens_cumsum,
duplicate_cache_len,
req_to_token.shape[1],
topk,
speculative_num_steps,
page_size,
next_power_of_2(num_seqs),
next_power_of_2(speculative_num_steps + page_size),
)
out_cache_loc = out_cache_loc[: num_seqs * topk * speculative_num_steps]
# fmt: off
expected_out_cache_loc = torch.tensor([
12, 13, 14, 15, 16,
20, 21, 22, 23, 24,
28, 29, 30, 31, 32,
36, 37, 38, 39, 40,
44, 45, 46, 47, 48,
52, 53, 54, 55, 56,
60, 61, 62, 63, 64,
68, 69, 70, 71, 72,
76, 77, 78, 79, 80,
84, 85, 86, 87, 88,
92, 93, 94, 95, 96,
100, 101, 102, 103, 104,
108, 109, 110, 111, 112,
116, 117, 118, 119, 120,
124, 125, 126, 127, 128,
132, 133, 134, 135, 136,
140, 141, 142, 143, 144,
148, 149, 150, 151, 152,
156, 157, 158, 159, 160,
164, 165, 166, 167, 168,
172, 173, 174, 175, 176,
180, 181, 182, 183, 184,
188, 189, 190, 191, 192,
196, 197, 198, 199, 200
], device=device, dtype=torch.int32)
expected_source_cache_loc = torch.tensor([8, 9, 10] * 7 + [8] * 7, device=device, dtype=torch.int32)
expected_target_cache_loc = torch.tensor([
81, 82, 83, 89, 90, 91, 97, 98, 99, 105, 106, 107, 113, 114,
115, 121, 122, 123, 129, 130, 131, 147, 155, 163, 171, 179, 187, 195
], device=device, dtype=torch.int32)
# fmt: on
assert torch.allclose(out_cache_loc, expected_out_cache_loc)
assert torch.allclose(source_cache_loc, expected_source_cache_loc)
assert torch.allclose(target_cache_loc, expected_target_cache_loc)
copy_all_layer_kv_cache_tiled[(len(self.data_ptrs),)](
self.data_ptrs,
self.data_strides,
target_cache_loc,
source_cache_loc,
len(target_cache_loc),
next_power_of_2(len(target_cache_loc)),
BYTES_PER_TILE,
)
assert torch.allclose(
self.k_cache[0][81:84, 0, 0],
torch.tensor(
[0.8, 0.9, 1.0],
dtype=torch.float32,
device=device,
),
)
assert torch.allclose(
self.v_cache[0][81:84, 0, 0],
torch.tensor(
[-0.8, -0.9, -1.0],
dtype=torch.float32,
device=device,
),
)
def test_assign_draft_cache_locs_page_size_1(self):
# Test to make sure page_size=1 not affected
device = self.device
num_seqs = 1
page_size = 1
speculative_num_steps = 5
topk = 8
seq_lens_num = 7
extend_lens_num = topk * speculative_num_steps
req_pool_indices = torch.arange(num_seqs, dtype=torch.int32, device=device)
req_to_token = torch.zeros((num_seqs, 100), dtype=torch.int32, device=device)
req_to_token[0, :seq_lens_num] = torch.tensor(
[4, 5, 6, 7, 8, 9, 10], device=device
)
seq_lens = torch.tensor([seq_lens_num], dtype=torch.int32, device=device)
extend_lens = torch.tensor([extend_lens_num], dtype=torch.int32, device=device)
num_new_pages_per_topk = torch.tensor([2], dtype=torch.int32, device=device)
out_cache_loc = torch.arange(11, 11 + extend_lens_num, device=device)
last_page_lens = torch.tensor([3], dtype=torch.int32, device=device)
duplicate_cache_len = 0
target_cache_loc = None
source_cache_loc = None
last_page_lens_cumsum = None
assign_draft_cache_locs[(num_seqs,)](
req_pool_indices,
req_to_token,
seq_lens,
extend_lens,
num_new_pages_per_topk,
out_cache_loc,
source_cache_loc,
target_cache_loc,
last_page_lens_cumsum,
duplicate_cache_len,
req_to_token.shape[1],
topk,
speculative_num_steps,
page_size,
next_power_of_2(num_seqs),
next_power_of_2(speculative_num_steps + page_size),
)
out_cache_loc = out_cache_loc[: num_seqs * topk * speculative_num_steps]
expected_out_cache_loc = torch.arange(11, 11 + extend_lens_num, device=device)
assert torch.allclose(out_cache_loc, expected_out_cache_loc)
def test_assign_draft_cache_locs_page_size_gt_spec_steps(self):
device = self.device
num_seqs = 1
page_size = 16
speculative_num_steps = 4
topk = 3
seq_lens_num = 12
pool_len = 256
req_pool_indices = torch.arange(num_seqs, dtype=torch.int32, device=device)
req_to_token = torch.zeros(
(num_seqs, pool_len), dtype=torch.int32, device=device
)
req_to_token[0, :seq_lens_num] = torch.arange(
seq_lens_num, dtype=torch.int32, device=device
)
seq_lens = torch.tensor([seq_lens_num], dtype=torch.int32, device=device)
last_page_len = seq_lens_num % page_size
last_page_lens = torch.tensor([last_page_len], dtype=torch.int32, device=device)
last_page_lens_cumsum = torch.cumsum(last_page_lens, dim=0)
num_new_pages_per_topk_val = (
last_page_len + speculative_num_steps + page_size - 1
) // page_size
num_new_pages_per_topk = torch.tensor(
[num_new_pages_per_topk_val], dtype=torch.int32, device=device
)
extend_lens_num = num_new_pages_per_topk_val * page_size * topk
extend_lens = torch.tensor([extend_lens_num], dtype=torch.int32, device=device)
out_cache_loc = torch.arange(
2000, 2000 + extend_lens_num, dtype=torch.int32, device=device
)
duplicate_cache_len = last_page_lens.sum().item() * (topk - 1)
target_cache_loc = torch.zeros(
duplicate_cache_len, dtype=torch.int32, device=device
)
source_cache_loc = torch.zeros(
duplicate_cache_len, dtype=torch.int32, device=device
)
assign_draft_cache_locs[(num_seqs,)](
req_pool_indices,
req_to_token,
seq_lens,
extend_lens,
num_new_pages_per_topk,
out_cache_loc,
source_cache_loc,
target_cache_loc,
last_page_lens_cumsum,
duplicate_cache_len,
req_to_token.shape[1],
topk,
speculative_num_steps,
page_size,
next_power_of_2(num_seqs),
next_power_of_2(speculative_num_steps + page_size),
)
trimmed = out_cache_loc[: num_seqs * topk * speculative_num_steps]
expected = []
for topk_id in range(topk):
start = seq_lens_num + topk_id * num_new_pages_per_topk_val * page_size
expected.append(
req_to_token[0, start : start + speculative_num_steps].clone()
)
expected_out_cache_loc = torch.cat(expected)
assert torch.allclose(trimmed, expected_out_cache_loc)
if __name__ == "__main__":
unittest.main()
python/sglang/test/test_activation.py
-105
View File
@@ -1,105 +0,0 @@
import itertools
import unittest
import torch
from sglang.srt.layers.activation import GeluAndMul, QuickGELU
from sglang.srt.utils import is_hip
from sglang.test.test_utils import CustomTestCase
_is_hip = is_hip()
class TestGeluAndMul(CustomTestCase):
DTYPES = [torch.half, torch.bfloat16]
NUM_TOKENS = [7, 83, 2048]
D = [512, 4096, 5120, 13824]
SEEDS = [0]
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
torch.set_default_device("cuda")
def _run_gelu_and_mul_test(self, num_tokens, d, dtype, seed):
torch.manual_seed(seed)
layer = GeluAndMul().to(dtype=dtype)
x = torch.randn(num_tokens, 2 * d, dtype=dtype)
with torch.inference_mode():
ref_out = layer.forward_native(x)
out = layer.forward_cuda(x)
if dtype == torch.bfloat16:
atol = rtol = 1e-2
else:
atol = rtol = 1e-3
self.assertTrue(torch.allclose(out, ref_out, atol=atol, rtol=rtol))
def test_gelu_and_mul(self):
for params in itertools.product(
self.NUM_TOKENS,
self.D,
self.DTYPES,
self.SEEDS,
):
with self.subTest(
num_tokens=params[0],
d=params[1],
dtype=params[2],
seed=params[3],
):
self._run_gelu_and_mul_test(*params)
class TestQuickGELU(CustomTestCase):
DTYPES = [torch.half, torch.bfloat16]
NUM_TOKENS = [7, 83, 2048] # batch = sequence length
DIMS = [512, 4096, 5120, 13824] # all multiples of 16 bytes
SEEDS = [0]
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
torch.set_default_device("cuda")
def _run_gelu_quick_test(self, n_tok: int, dim: int, dtype: torch.dtype, seed: int):
torch.manual_seed(seed)
layer = QuickGELU().to(dtype=dtype)
x = torch.randn(n_tok, dim, dtype=dtype, device="cuda")
with torch.inference_mode():
ref = layer.forward_native(x) # x * sigmoid(1.702 * x), fp32 math
if _is_hip:
out = layer.forward_hip(x) # 128-bit vectorised kernel from sgl-kernel
else:
out = layer.forward_cuda(x)
tol = 1e-2 if dtype is torch.bfloat16 else 1e-3
self.assertTrue(
torch.allclose(out, ref, atol=tol, rtol=tol),
msg=f"Mismatch @ B={n_tok}, D={dim}, dtype={dtype}",
)
print(f"Match @ B={n_tok}, D={dim}, dtype={dtype}")
def test_quick_gelu(self):
for params in itertools.product(
self.NUM_TOKENS, self.DIMS, self.DTYPES, self.SEEDS
):
with self.subTest(
num_tokens=params[0],
dim=params[1],
dtype=params[2],
seed=params[3],
):
self._run_gelu_quick_test(*params)
if __name__ == "__main__":
unittest.main(verbosity=2)
python/sglang/test/test_block_fp8.py
-760
View File
@@ -1,760 +0,0 @@
import itertools
import unittest
from functools import lru_cache
import torch
from sglang.srt.layers.activation import SiluAndMul
from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import fused_moe
from sglang.srt.layers.moe.topk import TopKConfig, select_experts
from sglang.srt.layers.quantization.fp8_kernel import (
per_tensor_quant_mla_fp8,
per_token_group_quant_fp8,
per_token_group_quant_mla_deep_gemm_masked_fp8,
static_quant_fp8,
w8a8_block_fp8_matmul,
)
from sglang.srt.layers.quantization.fp8_utils import (
input_to_float8,
mxfp8_group_quantize,
triton_mxfp8_blockscaled_linear,
)
from sglang.srt.utils import is_sm100_supported, is_sm120_supported
from sglang.test.test_utils import CustomTestCase
_is_cuda = torch.cuda.is_available() and torch.version.cuda
# For test
@lru_cache(maxsize=1)
def _get_triton_mxfp8_upcast():
try:
from triton_kernels.numerics_details.mxfp import upcast_from_mxfp_torch
except Exception as err:
raise RuntimeError(
"MXFP8 dequantization requires triton_kernels with MXFP8 support."
) from err
return upcast_from_mxfp_torch
# For test
def native_per_token_group_quant_fp8(
x, group_size, eps=1e-10, dtype=torch.float8_e4m3fn
):
"""Function to perform per-token-group quantization on an input tensor `x` using native torch.
It converts the tensor values into float8 values and returns the
quantized tensor along with the scaling factor used for quantization.
Note that only `torch.float8_e4m3fn` is supported for now.
"""
assert (
x.shape[-1] % group_size == 0
), "the last dimension of `x` cannot be divisible by `group_size`"
assert x.is_contiguous(), "`x` is not contiguous"
finfo = torch.finfo(dtype)
fp8_min = finfo.min
fp8_max = finfo.max
x_ = x.reshape(x.numel() // group_size, group_size)
amax = x_.abs().max(dim=-1, keepdim=True)[0].clamp(min=eps).to(torch.float32)
x_s = amax / fp8_max
x_q = (x_ / x_s).clamp(min=fp8_min, max=fp8_max).to(dtype)
x_q = x_q.reshape(x.shape)
x_s = x_s.reshape(x.shape[:-1] + (x.shape[-1] // group_size,))
return x_q, x_s
class TestPerTokenGroupQuantFP8(CustomTestCase):
DTYPES = [torch.half, torch.bfloat16, torch.float32]
NUM_TOKENS = [7, 83, 2048]
D = [512, 4096, 5120, 13824]
GROUP_SIZE = [64, 128, 256, 512]
SEEDS = [0]
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
torch.set_default_device("cuda")
def _per_token_group_quant_fp8(self, num_tokens, d, dtype, group_size, seed):
torch.manual_seed(seed)
x = torch.rand(num_tokens, d, dtype=dtype)
with torch.inference_mode():
ref_out, ref_scale = native_per_token_group_quant_fp8(x, group_size)
out, scale = per_token_group_quant_fp8(x, group_size)
self.assertTrue(
torch.allclose(out.to(torch.float32), ref_out.to(torch.float32), rtol=0.20)
)
self.assertTrue(torch.allclose(scale, ref_scale))
def test_per_token_group_quant_fp8(self):
for params in itertools.product(
self.NUM_TOKENS,
self.D,
self.DTYPES,
self.GROUP_SIZE,
self.SEEDS,
):
with self.subTest(
num_tokens=params[0],
d=params[1],
dtype=params[2],
group_size=params[3],
seed=params[4],
):
self._per_token_group_quant_fp8(*params)
# For test
def native_static_quant_fp8(x, x_s, dtype=torch.float8_e4m3fn):
"""Function to perform static quantization on an input tensor `x` using native torch.
It converts the tensor values into float8 values and returns the
quantized tensor along with the scaling factor used for quantization.
"""
assert x.is_contiguous(), "`x` is not contiguous"
assert x_s.numel() == 1, "only supports per-tensor scale"
finfo = torch.finfo(dtype)
fp8_min = finfo.min
fp8_max = finfo.max
x_ = x.reshape(x.numel() // x.shape[-1], x.shape[-1])
x_s_inv = 1.0 / x_s
x_q = (x_ * x_s_inv).clamp(min=fp8_min, max=fp8_max).to(dtype)
x_q = x_q.reshape(x.shape)
return x_q, x_s
class TestStaticQuantFP8(CustomTestCase):
DTYPES = [torch.half, torch.bfloat16, torch.float32]
NUM_TOKENS = [7, 83, 2048]
D = [512, 4096, 5120, 13824]
SEEDS = [0]
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
torch.set_default_device("cuda")
def _static_quant_fp8(self, num_tokens, d, dtype, seed):
torch.manual_seed(seed)
x = torch.rand(num_tokens, d, dtype=dtype)
fp8_max = torch.finfo(torch.float8_e4m3fn).max
x_s = x.max() / fp8_max
with torch.inference_mode():
ref_out, _ = native_static_quant_fp8(x, x_s)
out, _ = static_quant_fp8(x, x_s, repeat_scale=True)
self.assertTrue(
torch.allclose(out.to(torch.float32), ref_out.to(torch.float32), rtol=0.50)
)
def test_static_quant_fp8(self):
for params in itertools.product(
self.NUM_TOKENS,
self.D,
self.DTYPES,
self.SEEDS,
):
with self.subTest(
num_tokens=params[0],
d=params[1],
dtype=params[2],
seed=params[3],
):
self._static_quant_fp8(*params)
class TestPerTensorQuantMlaFP8(CustomTestCase):
DTYPES = [torch.half, torch.bfloat16, torch.float32]
NUM_TOKENS = [7, 83, 2048]
D = [512, 4096, 5120, 13824]
LAST_D_EXT = [1024, 0]
LAST_D = [512]
SEEDS = [0]
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
torch.set_default_device("cuda")
def _per_tensor_quant_mla_fp8(self, num_tokens, d, last_d_ext, last_d, dtype, seed):
torch.manual_seed(seed)
x = torch.rand(
(num_tokens, d // last_d, last_d + last_d_ext),
dtype=dtype,
)
x_sub, _ = x.split([last_d, last_d_ext], dim=-1)
with torch.inference_mode():
ref_out, ref_s = input_to_float8(x_sub.transpose(0, 1))
out, out_s = per_tensor_quant_mla_fp8(x_sub.transpose(0, 1))
self.assertTrue(out.is_contiguous())
self.assertTrue(
torch.allclose(out.to(torch.float32), ref_out.to(torch.float32), rtol=0.50)
)
self.assertTrue(
torch.allclose(out_s.to(torch.float32), ref_s.to(torch.float32))
)
def test_per_tensor_quant_mla_fp8(self):
for params in itertools.product(
self.NUM_TOKENS,
self.D,
self.LAST_D_EXT,
self.LAST_D,
self.DTYPES,
self.SEEDS,
):
with self.subTest(
num_tokens=params[0],
d=params[1],
last_d_ext=params[2],
last_d=params[3],
dtype=params[4],
seed=params[5],
):
self._per_tensor_quant_mla_fp8(*params)
class TestPerTokenGroupQuantMlaDeepGemmMaskedFP8(CustomTestCase):
DTYPES = [torch.half, torch.bfloat16, torch.float32]
B = [128]
NUM_TOKENS = [7, 83, 2048, 1024 * 16]
D = [512, 128]
GROUP_SIZE = [128]
SEEDS = [0]
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
torch.set_default_device("cuda")
def _per_token_group_quant_mla_deep_gemm_masked_fp8(
self, b, num_tokens, d, dtype, group_size, seed
):
torch.manual_seed(seed)
x = torch.rand(b, num_tokens, d, dtype=dtype)
with torch.inference_mode():
ref_out, ref_scale = native_per_token_group_quant_fp8(x, group_size, 1e-12)
out, scale, _, _, _ = per_token_group_quant_mla_deep_gemm_masked_fp8(
x, group_size
)
out = out[:, :num_tokens, :]
scale = scale[:, :num_tokens, :]
self.assertTrue(
torch.allclose(
out.to(torch.float32), ref_out.to(torch.float32), rtol=0.20, atol=1e-2
)
)
self.assertTrue(torch.allclose(scale, ref_scale))
def test_per_token_group_quant_mla_deep_gemm_masked_fp8(self):
for params in itertools.product(
self.B,
self.NUM_TOKENS,
self.D,
self.DTYPES,
self.GROUP_SIZE,
self.SEEDS,
):
with self.subTest(
b=params[0],
num_tokens=params[1],
d=params[2],
dtype=params[3],
group_size=params[4],
seed=params[5],
):
self._per_token_group_quant_mla_deep_gemm_masked_fp8(*params)
# For test
def native_w8a8_block_fp8_matmul(A, B, As, Bs, block_size, output_dtype=torch.float16):
"""This function performs matrix multiplication with block-wise quantization using native torch.
It takes two input tensors `A` and `B` with scales `As` and `Bs`.
The output is returned in the specified `output_dtype`.
"""
A = A.to(torch.float32)
B = B.to(torch.float32)
assert A.shape[-1] == B.shape[-1]
assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2
assert len(block_size) == 2
block_n, block_k = block_size[0], block_size[1]
assert (A.shape[-1] + block_k - 1) // block_k == As.shape[-1]
assert A.shape[:-1] == As.shape[:-1]
M = A.numel() // A.shape[-1]
N, K = B.shape
origin_C_shape = A.shape[:-1] + (N,)
A = A.reshape(M, A.shape[-1])
As = As.reshape(M, As.shape[-1])
n_tiles = (N + block_n - 1) // block_n
k_tiles = (K + block_k - 1) // block_k
assert n_tiles == Bs.shape[0]
assert k_tiles == Bs.shape[1]
C_shape = (M, N)
C = torch.zeros(C_shape, dtype=torch.float32, device=A.device)
A_tiles = [A[:, i * block_k : min((i + 1) * block_k, K)] for i in range(k_tiles)]
B_tiles = [
[
B[
j * block_n : min((j + 1) * block_n, N),
i * block_k : min((i + 1) * block_k, K),
]
for i in range(k_tiles)
]
for j in range(n_tiles)
]
C_tiles = [C[:, j * block_n : min((j + 1) * block_n, N)] for j in range(n_tiles)]
As_tiles = [As[:, i : i + 1] for i in range(k_tiles)]
for i in range(k_tiles):
for j in range(n_tiles):
a = A_tiles[i]
b = B_tiles[j][i]
c = C_tiles[j]
s = As_tiles[i] * Bs[j][i]
c[:, :] += torch.matmul(a, b.t()) * s
C = C.reshape(origin_C_shape).to(output_dtype)
return C
class TestW8A8BlockFP8Matmul(CustomTestCase):
if not _is_cuda:
OUT_DTYPES = [torch.float32, torch.half, torch.bfloat16]
M = [1, 7, 83, 512, 2048]
NKs = [
(N, K)
for N in [128, 512, 1024, 4096, 7748, 13824]
for K in [256, 4096, 5120, 3884, 13824]
]
# BLOCK_SIZE = [[64, 64], [64, 128], [128, 64], [128, 128]]
BLOCK_SIZE = [[128, 128]]
SEEDS = [0]
else:
# use practical shape in DeepSeek V3 for test
OUT_DTYPES = [torch.bfloat16]
M = [64, 128, 512, 1024, 4096]
NKs = [
(2112, 7168),
(1536, 7168),
(3072, 1536),
(24576, 7168),
(4096, 512),
(7168, 2048),
(4608, 7168),
(512, 7168),
(7168, 2304),
(7168, 512),
]
BLOCK_SIZE = [[128, 128]]
SEEDS = [0]
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
torch.set_default_device("cuda")
def _w8a8_block_fp8_matmul(self, M, NK, block_size, out_dtype, seed):
N, K = NK
torch.manual_seed(seed)
# NOTE(HandH1998): to avoid overflow when out_dtype = torch.half
factor_for_scale = 1e-2
fp8_info = torch.finfo(torch.float8_e4m3fn)
fp8_max, fp8_min = fp8_info.max, fp8_info.min
A_fp32 = (torch.rand(M, K, dtype=torch.float32) - 0.5) * 2 * fp8_max
A_fp8 = A_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
B_fp32 = (torch.rand(N, K, dtype=torch.float32) - 0.5) * 2 * fp8_max
B_fp8 = B_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
block_n, block_k = block_size[0], block_size[1]
n_tiles = (N + block_n - 1) // block_n
k_tiles = (K + block_k - 1) // block_k
As = torch.rand(M, k_tiles, dtype=torch.float32) * factor_for_scale
Bs = torch.rand(n_tiles, k_tiles, dtype=torch.float32) * factor_for_scale
with torch.inference_mode():
ref_out = native_w8a8_block_fp8_matmul(
A_fp8, B_fp8, As, Bs, block_size, out_dtype
)
out = w8a8_block_fp8_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype)
self.assertTrue(
torch.mean(torch.abs(out.to(torch.float32) - ref_out.to(torch.float32)))
/ torch.mean(torch.abs(ref_out.to(torch.float32)))
< 0.001
)
def test_w8a8_block_fp8_matmul(self):
for params in itertools.product(
self.M,
self.NKs,
self.BLOCK_SIZE,
self.OUT_DTYPES,
self.SEEDS,
):
with self.subTest(
M=params[0],
NKs=params[1],
block_size=params[2],
out_dtype=params[3],
seed=params[4],
):
self._w8a8_block_fp8_matmul(*params)
def _mxfp8_group_dequant(q: torch.Tensor, scale_u8: torch.Tensor) -> torch.Tensor:
upcast_from_mxfp_torch = _get_triton_mxfp8_upcast()
return upcast_from_mxfp_torch(q, scale_u8, torch.float32, axis=1)
class TestMXFP8DenseLinear(CustomTestCase):
DTYPES = [torch.bfloat16]
M = [1, 127, 128, 129, 255, 256]
NKs = [
(256, 512),
(384, 1024),
(512, 2048),
(768, 1024),
]
SEEDS = [0]
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
if not (is_sm100_supported() or is_sm120_supported()):
raise unittest.SkipTest("MXFP8 requires Blackwell (SM100/SM120)")
torch.set_default_device("cuda")
def _mxfp8_dense_linear(self, M, NK, dtype, seed):
N, K = NK
torch.manual_seed(seed)
input_fp32 = torch.randn((M, K), dtype=torch.float32) / 4
input_fp16 = input_fp32.to(dtype)
weight_fp32 = torch.randn((N, K), dtype=torch.float32) / 4
weight_q, weight_scale_u8 = mxfp8_group_quantize(weight_fp32)
with torch.inference_mode():
q_input, input_scale_u8 = mxfp8_group_quantize(input_fp16.to(torch.float32))
a_dq = _mxfp8_group_dequant(q_input, input_scale_u8)
b_dq = _mxfp8_group_dequant(weight_q, weight_scale_u8)
ref_out = torch.matmul(a_dq, b_dq.t()).to(dtype)
out = triton_mxfp8_blockscaled_linear(
input=input_fp16,
weight=weight_q,
weight_scale=weight_scale_u8,
)
out_prequant = triton_mxfp8_blockscaled_linear(
input=q_input,
weight=weight_q,
weight_scale=weight_scale_u8,
input_scale=input_scale_u8,
output_dtype=dtype,
)
self.assertTrue(
torch.mean(torch.abs(out.to(torch.float32) - ref_out.to(torch.float32)))
/ torch.mean(torch.abs(ref_out.to(torch.float32)))
< 0.02
)
self.assertTrue(
torch.mean(
torch.abs(out_prequant.to(torch.float32) - ref_out.to(torch.float32))
)
/ torch.mean(torch.abs(ref_out.to(torch.float32)))
< 0.02
)
def test_mxfp8_dense_linear(self):
for params in itertools.product(
self.M,
self.NKs,
self.DTYPES,
self.SEEDS,
):
with self.subTest(
M=params[0],
NKs=params[1],
dtype=params[2],
seed=params[3],
):
self._mxfp8_dense_linear(*params)
# For test
def torch_w8a8_block_fp8_moe(a, w1, w2, w1_s, w2_s, score, topk, block_shape):
"""This function performs fused moe with block-wise quantization using native torch."""
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)
_, block_k = block_shape[0], block_shape[1]
a_q, a_s = native_per_token_group_quant_fp8(a, block_k)
# NOTE(HandH1998): Since "index_cuda" not implemented for 'Float8_e4m3fn', we need to cast `float8`` to `float32``.
a_q = a_q.to(torch.float32)
for i in range(w1.shape[0]):
mask = topk_ids == i
if mask.sum():
inter_out = native_w8a8_block_fp8_matmul(
a_q[mask], w1[i], a_s[mask], w1_s[i], block_shape, output_dtype=a.dtype
)
act_out = SiluAndMul().forward_native(inter_out)
act_out_q, act_out_s = native_per_token_group_quant_fp8(act_out, block_k)
act_out = act_out.to(torch.float32)
out[mask] = native_w8a8_block_fp8_matmul(
act_out_q, w2[i], act_out_s, w2_s[i], block_shape, output_dtype=a.dtype
)
return (
out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype)
).sum(dim=1)
class TestW8A8BlockFP8FusedMoE(CustomTestCase):
DTYPES = [torch.float32, torch.half, torch.bfloat16]
M = [1, 33, 64, 222, 1024 * 128]
N = [128, 1024, 2048]
K = [256, 4096, 5120]
E = [8, 24]
TOP_KS = [2, 6]
BLOCK_SIZE = [[64, 64], [64, 128], [128, 64], [128, 128]]
# BLOCK_SIZE = [[128, 128]]
SEEDS = [0]
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
torch.set_default_device("cuda")
def _w8a8_block_fp8_fused_moe(self, M, N, K, E, topk, block_size, dtype, seed):
torch.manual_seed(seed)
# NOTE(HandH1998): to avoid overflow when out_dtype = torch.half
factor_for_scale = 1e-2
fp8_info = torch.finfo(torch.float8_e4m3fn)
fp8_max, fp8_min = fp8_info.max, fp8_info.min
a = torch.randn((M, K), dtype=dtype) / 10
w1_fp32 = (torch.rand((E, 2 * N, K), dtype=torch.float32) - 0.5) * 2 * fp8_max
w1 = w1_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
w2_fp32 = (torch.rand((E, K, N), dtype=torch.float32) - 0.5) * 2 * fp8_max
w2 = w2_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
block_n, block_k = block_size[0], block_size[1]
n_tiles_w1 = (2 * N + block_n - 1) // block_n
n_tiles_w2 = (K + block_n - 1) // block_n
k_tiles_w1 = (K + block_k - 1) // block_k
k_tiles_w2 = (N + block_k - 1) // block_k
w1_s = (
torch.rand((E, n_tiles_w1, k_tiles_w1), dtype=torch.float32)
* factor_for_scale
)
w2_s = (
torch.rand((E, n_tiles_w2, k_tiles_w2), dtype=torch.float32)
* factor_for_scale
)
score = torch.randn((M, E), dtype=dtype)
with torch.inference_mode():
ref_out = torch_w8a8_block_fp8_moe(
a, w1, w2, w1_s, w2_s, score, topk, block_size
)
topk_output = select_experts(
hidden_states=a,
router_logits=score,
topk_config=TopKConfig(top_k=topk, renormalize=False),
)
out = fused_moe(
a,
w1,
w2,
topk_output,
use_fp8_w8a8=True,
w1_scale=w1_s,
w2_scale=w2_s,
block_shape=block_size,
)
self.assertTrue(
torch.mean(torch.abs(out.to(torch.float32) - ref_out.to(torch.float32)))
/ torch.mean(torch.abs(ref_out.to(torch.float32)))
< 0.02
)
def test_w8a8_block_fp8_fused_moe(self):
for params in itertools.product(
self.M,
self.N,
self.K,
self.E,
self.TOP_KS,
self.BLOCK_SIZE,
self.DTYPES,
self.SEEDS,
):
with self.subTest(
M=params[0],
N=params[1],
K=params[2],
E=params[3],
topk=params[4],
block_size=params[5],
dtype=params[6],
seed=params[7],
):
self._w8a8_block_fp8_fused_moe(*params)
# For test
def torch_w8a8_block_fp8_bmm(a, a_s, w, w_s, block_shape, out_dtype):
"""This function performs bmm with block-wise quantization using native torch."""
B, N, _ = w.shape
_, M, _ = a.shape
out = torch.empty((B, M, N), dtype=out_dtype, device=a.device)
for i in range(B):
out[i] = native_w8a8_block_fp8_matmul(
a[i], w[i], a_s[i], w_s[i], block_shape, output_dtype=out_dtype
)
return out
class TestW8A8BlockFP8BatchedDeepGemm(CustomTestCase):
DTYPES = [torch.bfloat16]
M = [1, 33, 64, 222, 8192]
N = [128, 512]
K = [128, 512]
BATCH = [128]
BLOCK_SIZE = [[128, 128]]
SEEDS = [0]
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
try:
import deep_gemm # noqa: F401
except ImportError:
raise unittest.SkipTest("DeepGEMM is not available")
torch.set_default_device("cuda")
def _w8a8_block_fp8_batched_deep_gemm(self, M, N, K, B, block_size, dtype, seed):
torch.manual_seed(seed)
factor_for_scale = 1e-2
fp8_info = torch.finfo(torch.float8_e4m3fn)
fp8_max, fp8_min = fp8_info.max, fp8_info.min
a_fp32 = torch.randn((B, M, K), dtype=torch.float32) / 10
a = a_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
w_fp32 = (torch.rand((B, N, K), dtype=torch.float32) - 0.5) * 2 * fp8_max
w = w_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
block_n, block_k = block_size[0], block_size[1]
n_tiles_w = (N + block_n - 1) // block_n
k_tiles_w = (K + block_k - 1) // block_k
w_s = (
torch.rand((B, n_tiles_w, k_tiles_w), dtype=torch.float32)
* factor_for_scale
)
a_s = torch.rand((B, M, k_tiles_w), dtype=torch.float32) * factor_for_scale
ae = a.new_empty(B, (M + 255) // 256 * 256, K)
ae_s = a_s.new_empty(B, (M + 255) // 256 * 256, k_tiles_w)
oe = torch.empty((B, (M + 255) // 256 * 256, N), dtype=dtype)
ae[:, :M, :] = a
ae_s[:, :M, :] = a_s
masked_m = torch.full((B,), M, dtype=torch.int)
expected_m = M
lhs = (
ae,
ae_s,
)
rhs = (
w,
w_s,
)
from deep_gemm import fp8_m_grouped_gemm_nt_masked
with torch.inference_mode():
ref_out = torch_w8a8_block_fp8_bmm(a, a_s, w, w_s, block_size, dtype)
fp8_m_grouped_gemm_nt_masked(lhs, rhs, oe, masked_m, expected_m)
out = oe[:, :M, :]
self.assertTrue(
torch.mean(torch.abs(out.to(torch.float32) - ref_out.to(torch.float32)))
/ torch.mean(torch.abs(ref_out.to(torch.float32)))
< 0.0001
)
def test_w8a8_block_fp8_batched_deep_gemm(self):
for params in itertools.product(
self.M,
self.N,
self.K,
self.BATCH,
self.BLOCK_SIZE,
self.DTYPES,
self.SEEDS,
):
with self.subTest(
M=params[0],
N=params[1],
K=params[2],
B=params[3],
block_size=params[4],
dtype=params[5],
seed=params[6],
):
self._w8a8_block_fp8_batched_deep_gemm(*params)
if __name__ == "__main__":
unittest.main(verbosity=2)
python/sglang/test/test_block_fp8_deep_gemm_blackwell.py
-251
View File
@@ -1,251 +0,0 @@
import itertools
import unittest
from typing import List, Tuple
import torch
from deep_gemm import fp8_gemm_nt
from sglang.test.test_utils import CustomTestCase
_is_cuda = torch.cuda.is_available() and torch.version.cuda
# Modify form DeepGEMM Blackwell
def ceil_div(x: int, y: int) -> int:
return (x + y - 1) // y
def align(x: int, y: int) -> int:
return ceil_div(x, y) * y
def per_token_group_quant_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
assert x.dim() == 2 and x.size(1) % 128 == 0
m, n = x.shape
x_view = x.view(m, -1, 128)
x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4)
sf = x_amax / 448.0
return (x_view * (1.0 / sf.unsqueeze(2))).to(torch.float8_e4m3fn).view(m, n), sf
def per_block_quant_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
assert x.dim() == 2
m, n = x.shape
x_padded = torch.zeros(
(align(m, 128), align(n, 128)), dtype=x.dtype, device=x.device
)
x_padded[:m, :n] = x
x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, 128)
x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
sf = x_amax / 448.0
x_scaled = (x_view * (1.0 / sf)).to(torch.float8_e4m3fn)
return x_scaled.view_as(x_padded)[:m, :n].contiguous(), sf.view(
x_view.size(0), x_view.size(2)
)
def ceil_to_ue8m0(x: torch.Tensor):
assert x.view(-1).amax().item() > 0
return torch.pow(2.0, torch.ceil(torch.log2(x.abs())))
def per_token_group_quant_mxfp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
assert x.dim() == 2 and x.size(1) % 128 == 0
m, n = x.shape
x_view = x.view(m, -1, 128)
x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4)
sf = ceil_to_ue8m0(x_amax / 448.0)
return (x_view * (1.0 / sf.unsqueeze(2))).to(torch.float8_e4m3fn).view(m, n), sf
def per_block_quant_mxfp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
assert x.dim() == 2
m, n = x.shape
x_padded = torch.zeros(
(align(m, 128), align(n, 128)), dtype=x.dtype, device=x.device
)
x_padded[:m, :n] = x
x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, 128)
x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
sf = ceil_to_ue8m0(x_amax / 448.0)
x_scaled = (x_view * (1.0 / sf)).to(torch.float8_e4m3fn)
return x_scaled.view_as(x_padded)[:m, :n].contiguous(), sf.view(
x_view.size(0), x_view.size(2)
)
# For test
def native_w8a8_block_fp8_matmul(A, B, As, Bs, block_size, output_dtype=torch.float16):
"""This function performs matrix multiplication with block-wise quantization using native torch.
It takes two input tensors `A` and `B` with scales `As` and `Bs`.
The output is returned in the specified `output_dtype`.
"""
A = A.to(torch.float32)
B = B.to(torch.float32)
assert A.shape[-1] == B.shape[-1]
assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2
assert len(block_size) == 2
block_n, block_k = block_size[0], block_size[1]
assert (A.shape[-1] + block_k - 1) // block_k == As.shape[-1]
assert A.shape[:-1] == As.shape[:-1]
M = A.numel() // A.shape[-1]
N, K = B.shape
origin_C_shape = A.shape[:-1] + (N,)
A = A.reshape(M, A.shape[-1])
As = As.reshape(M, As.shape[-1])
n_tiles = (N + block_n - 1) // block_n
k_tiles = (K + block_k - 1) // block_k
assert n_tiles == Bs.shape[0]
assert k_tiles == Bs.shape[1]
C_shape = (M, N)
C = torch.zeros(C_shape, dtype=torch.float32, device=A.device)
A_tiles = [A[:, i * block_k : min((i + 1) * block_k, K)] for i in range(k_tiles)]
B_tiles = [
[
B[
j * block_n : min((j + 1) * block_n, N),
i * block_k : min((i + 1) * block_k, K),
]
for i in range(k_tiles)
]
for j in range(n_tiles)
]
C_tiles = [C[:, j * block_n : min((j + 1) * block_n, N)] for j in range(n_tiles)]
As_tiles = [As[:, i : i + 1] for i in range(k_tiles)]
for i in range(k_tiles):
for j in range(n_tiles):
a = A_tiles[i]
b = B_tiles[j][i]
c = C_tiles[j]
s = As_tiles[i] * Bs[j][i]
c[:, :] += torch.matmul(a, b.t()) * s
C = C.reshape(origin_C_shape).to(output_dtype)
return C
def block_quant_dequant(
x_q_block: torch.Tensor,
x_s: torch.Tensor,
block_size: List[int],
dtype: torch.dtype,
) -> torch.Tensor:
"""This function converts block-wise quantization to unquantized.
The inputs are block-wise quantization tensor `x_q_block`, block-wise quantization scale
and the block size.
The output is an unquantized tensor with dtype.
"""
block_n, block_k = block_size[0], block_size[1]
n, k = x_q_block.shape
n_tiles = (n + block_n - 1) // block_n
k_tiles = (k + block_k - 1) // block_k
assert n_tiles == x_s.shape[0]
assert k_tiles == x_s.shape[1]
x_dq_block = torch.empty_like(x_q_block, dtype=dtype)
for j in range(n_tiles):
for i in range(k_tiles):
x_q_block_tile = x_q_block[
j * block_n : min((j + 1) * block_n, n),
i * block_k : min((i + 1) * block_k, k),
]
x_dq_block_tile = x_dq_block[
j * block_n : min((j + 1) * block_n, n),
i * block_k : min((i + 1) * block_k, k),
]
x_dq_block_tile[:, :] = x_q_block_tile.to(torch.float32) * x_s[j][i]
return x_dq_block
class TestDeepGemmBlackwell(CustomTestCase):
if not _is_cuda:
OUT_DTYPES = [torch.float32, torch.half, torch.bfloat16]
M = [1, 7, 83, 512, 2048]
NKs = [
(N, K)
for N in [128, 512, 1024, 4096, 7748, 13824]
for K in [256, 4096, 5120, 3884, 13824]
]
# BLOCK_SIZE = [[64, 64], [64, 128], [128, 64], [128, 128]]
BLOCK_SIZE = [[128, 128]]
SEEDS = [0]
else:
# use practical shape in DeepSeek V3 for test
OUT_DTYPES = [torch.bfloat16]
M = [64, 128, 512, 1024, 4096]
NKs = [
(2112, 7168),
(1536, 7168),
# (3072, 1536),
# (24576, 7168),
# (4096, 512),
# (7168, 2048),
# (4608, 7168),
# (512, 7168),
# (7168, 2304),
# (7168, 512),
]
BLOCK_SIZE = [[128, 128]]
SEEDS = [0]
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
torch.set_default_device("cuda")
def _test_deep_gemm_blackwell(self, M, NK, block_size, out_dtype, seed):
N, K = NK
torch.manual_seed(seed)
A = torch.empty((M, K), dtype=torch.bfloat16).normal_(0, 0.2)
B = torch.empty((N, K), dtype=torch.bfloat16).normal_(0, 0.2)
A_q, A_s = per_token_group_quant_fp8(A)
B_q, B_s = per_block_quant_fp8(B)
A_dq = block_quant_dequant(A_q, A_s, [1, block_size[1]], out_dtype)
B_dq = block_quant_dequant(B_q, B_s, block_size, out_dtype)
A_qu = per_token_group_quant_mxfp8(A_dq)
B_qu = per_block_quant_mxfp8(B_dq)
out = None
with torch.inference_mode():
ref_out = native_w8a8_block_fp8_matmul(
A_qu[0], B_qu[0], A_qu[1], B_qu[1], block_size, out_dtype
)
out = torch.empty_like(ref_out)
fp8_gemm_nt(A_qu, B_qu, out)
torch.testing.assert_close(out, ref_out, atol=1e-1, rtol=1e-2)
def test_deep_gemm_blackwell(self):
for params in itertools.product(
self.M,
self.NKs,
self.BLOCK_SIZE,
self.OUT_DTYPES,
self.SEEDS,
):
with self.subTest(
M=params[0],
NKs=params[1],
block_size=params[2],
out_dtype=params[3],
seed=params[4],
):
self._test_deep_gemm_blackwell(*params)
if __name__ == "__main__":
unittest.main(verbosity=2)
python/sglang/test/test_custom_ops.py
-152
View File
@@ -1,152 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Adapted from https://github.com/vllm-project/vllm/blob/8ca7a71df787ad711ad3ac70a5bd2eb2bb398938/tests/quantization/test_fp8.py
import sys
import pytest
import torch
from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz, scaled_fp8_quant
from sglang.srt.utils import is_cuda, is_hip
_is_cuda = is_cuda()
_is_hip = is_hip()
_is_fp8_fnuz = is_fp8_fnuz()
fp8_dtype = torch.float8_e4m3fnuz if _is_fp8_fnuz else torch.float8_e4m3fn
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_scaled_fp8_quant_per_tensor(dtype) -> None:
def quantize_ref_per_tensor(tensor, inv_scale):
# The reference implementation that fully aligns to
# the kernel being tested.
finfo = torch.finfo(fp8_dtype)
scale = inv_scale.reciprocal()
qweight = (tensor.to(torch.float32) * scale).clamp(min=finfo.min, max=finfo.max)
qweight = qweight.to(fp8_dtype)
return qweight
def dequantize_per_tensor(tensor, inv_scale, dtype):
fake_qweight = tensor.to(dtype)
dq_weight = fake_qweight * inv_scale
return dq_weight
# Note that we use a shape % 8 != 0 to cover edge cases,
# because scaled_fp8_quant is vectorized by 8.
x = (torch.randn(size=(11, 11), device="cuda") * 13).to(dtype)
# Test Per Tensor Dynamic quantization
# scale = max(abs(x)) / FP8_E4M3_MAX
y, scale = scaled_fp8_quant(x, None)
ref_y = quantize_ref_per_tensor(x, scale)
torch.testing.assert_close(y, ref_y)
torch.testing.assert_close(
dequantize_per_tensor(y, scale, dtype),
dequantize_per_tensor(ref_y, scale, dtype),
)
# Test Per Tensor Static quantization
y, _ = scaled_fp8_quant(x, scale)
ref_y = quantize_ref_per_tensor(x, scale)
torch.testing.assert_close(y, ref_y)
torch.testing.assert_close(
dequantize_per_tensor(y, scale, dtype),
dequantize_per_tensor(ref_y, scale, dtype),
)
if _is_cuda or _is_hip:
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_scaled_fp8_quant_per_token_dynamic(dtype) -> None:
def quantize_ref_per_token(tensor, inv_scale):
# The reference implementation that fully aligns to
# the kernel being tested.
finfo = torch.finfo(fp8_dtype)
scale = inv_scale.reciprocal()
qweight = (tensor.to(torch.float32) * scale).clamp(
min=finfo.min, max=finfo.max
)
qweight = qweight.to(fp8_dtype)
return qweight
def dequantize_per_token(tensor, inv_scale, dtype):
fake_qweight = tensor.to(dtype)
dq_weight = fake_qweight * inv_scale
return dq_weight
# Note that we use a shape % 8 = 0,
# because per_token_quant_fp8 is vectorized by 8 elements.
x = (torch.randn(size=(11, 16), device="cuda") * 13).to(dtype)
# Test Per Tensor Dynamic quantization
# scale = max(abs(x)) / FP8_E4M3_MAX
y, scale = scaled_fp8_quant(x, None, use_per_token_if_dynamic=True)
ref_y = quantize_ref_per_token(x, scale)
torch.testing.assert_close(y, ref_y)
torch.testing.assert_close(
dequantize_per_token(y, scale, dtype),
dequantize_per_token(ref_y, scale, dtype),
)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_scaled_fp8_quant_with_padding(dtype) -> None:
original_rows = 5
x = (torch.randn(size=(original_rows, 16), device="cuda") * 13).to(dtype)
padding_size = 10
# Test with dynamic quantization
y_dynamic, scale_dynamic = scaled_fp8_quant(
x, None, num_token_padding=padding_size
)
# Verify output shape has the padded size
assert y_dynamic.shape[0] == padding_size
assert y_dynamic.shape[1] == x.shape[1]
# Verify that the actual data in the non-padded region is correctly quantized
y_without_padding, scale_without_padding = scaled_fp8_quant(x, None)
torch.testing.assert_close(y_dynamic[:original_rows], y_without_padding)
# Test with static quantization
# First get a scale
_, scale = scaled_fp8_quant(x, None)
# Then use it for static quantization with padding
y_static, _ = scaled_fp8_quant(x, scale, num_token_padding=padding_size)
# Verify output shape has the padded size
assert y_static.shape[0] == padding_size
assert y_static.shape[1] == x.shape[1]
# Verify that the actual data in the non-padded region is correctly quantized
y_static_without_padding, _ = scaled_fp8_quant(x, scale)
torch.testing.assert_close(y_static[:original_rows], y_static_without_padding)
# Test with per-token dynamic quantization
y_per_token, scale_per_token = scaled_fp8_quant(
x, None, num_token_padding=padding_size, use_per_token_if_dynamic=True
)
# Verify output shape has the padded size
assert y_per_token.shape[0] == padding_size
assert y_per_token.shape[1] == x.shape[1]
# Verify that the actual data in the non-padded region is correctly quantized
y_per_token_without_padding, scale_per_token_without_padding = scaled_fp8_quant(
x, None, use_per_token_if_dynamic=True
)
torch.testing.assert_close(
y_per_token[:original_rows], y_per_token_without_padding
)
torch.testing.assert_close(
scale_per_token[:original_rows], scale_per_token_without_padding
)
if __name__ == "__main__":
# Run the specific test function directly
sys.exit(pytest.main([__file__]))
python/sglang/test/test_cutlass_moe.py
-306
View File
@@ -1,306 +0,0 @@
import argparse
import torch
import triton # Added import
import triton.testing # Added import
from transformers import AutoConfig
from sglang.srt.layers.moe.cutlass_moe import cutlass_fused_experts_fp8
from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import fused_experts
from sglang.srt.layers.moe.topk import StandardTopKOutput
# Copy from: https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/utils.py
def calc_diff(x, y):
x, y = x.double(), y.double()
denominator = (x * x + y * y).sum()
sim = 2 * (x * y).sum() / denominator
return 1 - sim
def get_model_config(tp_size: int):
config = AutoConfig.from_pretrained(
"deepseek-ai/Deepseek-R1", trust_remote_code=True
)
E = config.n_routed_experts
topk = config.num_experts_per_tok
intermediate_size = config.moe_intermediate_size
shard_intermediate_size = 2 * intermediate_size // tp_size
return {
"num_experts": E,
"topk": topk,
"hidden_size": config.hidden_size,
"shard_intermediate_size": shard_intermediate_size,
"dtype": config.dtype,
"block_shape": config.quantization_config["weight_block_size"],
}
def to_fp8(tensor: torch.Tensor) -> torch.Tensor:
"""Converts tensor to FP8 E4M3, scaling values to fit the range."""
finfo = torch.finfo(torch.float8_e4m3fn)
# Calculate max absolute value safely
max_val = torch.max(torch.abs(tensor))
# Avoid division by zero if tensor is all zeros
if max_val == 0:
scale_factor = 1.0
else:
# Scale factor to bring the max value to finfo.max
scale_factor = finfo.max / max_val
# Apply scaling
scaled_tensor = tensor * scale_factor
# Clamp and convert
fp8_tensor = scaled_tensor.clamp(min=finfo.min, max=finfo.max).to(
dtype=torch.float8_e4m3fn
)
return fp8_tensor
def run_test(tp_size, batch_size, model_config, check=False):
print(f"\n--- Batch Size: {batch_size} ---")
torch.set_default_device("cuda")
torch.cuda.manual_seed_all(42) # For reproducible random numbers
E = model_config["num_experts"]
topk = model_config["topk"]
H = model_config["hidden_size"]
I = model_config["shard_intermediate_size"]
block_shape = model_config["block_shape"] # Tuple (BLOCK_N, BLOCK_K)
dtype = model_config["dtype"] # e.g., torch.bfloat16
print(
f"Config: E={E}, topk={topk}, H={H}, I_shard={I}, dtype={dtype}, block_shape={block_shape}"
)
# --- Input Data ---
# Use bf16/fp16 for input activation based on model config
x = torch.randn((batch_size, H), device="cuda", dtype=dtype)
# --- Weights (Generate in higher precision, then convert to FP8) ---
# Generate weights suitable for FP8 conversion (e.g., scaled appropriately)
w1_hp = torch.randn((E, I, H), device="cuda", dtype=torch.float32)
w2_hp = torch.randn((E, H, I // 2), device="cuda", dtype=torch.float32)
w1 = to_fp8(w1_hp)
w2 = to_fp8(w2_hp)
# --- Scales for FP8 Weights ---
block_n, block_k = block_shape
# Calculate number of blocks needed
w1_blocks_dim1 = (I + block_n - 1) // block_n
w1_blocks_dim2 = (H + block_k - 1) // block_k
w2_blocks_dim1 = (H + block_n - 1) // block_n
w2_blocks_dim2 = (I // 2 + block_k - 1) // block_k
# Scales are typically float32 or float16/bfloat16
scale_dtype = torch.float32 # Or dtype if scales match model dtype
w1_scale = torch.full(
(E, w1_blocks_dim1, w1_blocks_dim2), 1, device="cuda", dtype=scale_dtype
) # Avoid zero scales
w2_scale = torch.full(
(E, w2_blocks_dim1, w2_blocks_dim2), 1, device="cuda", dtype=scale_dtype
) # Avoid zero scales
# --- Routing Information ---
topk_weights = torch.softmax(
torch.rand(batch_size, topk, device="cuda", dtype=dtype), dim=-1
)
topk_ids = torch.randint(0, E, (batch_size, topk), dtype=torch.int32, device="cuda")
a1_strides = torch.full((E,), H, dtype=torch.int64, device="cuda")
c1_strides = torch.full((E,), I, dtype=torch.int64, device="cuda")
a2_strides = torch.full((E,), I // 2, dtype=torch.int64, device="cuda")
c2_strides = torch.full((E,), H, dtype=torch.int64, device="cuda")
workspace = torch.empty(
(7182 * 1024), device="cuda", dtype=torch.uint8
) # Allocate sufficient workspace
# Pointer arrays (often filled by the kernel or a prep step, but needed as args)
a_ptrs = torch.empty((E,), dtype=torch.int64, device="cuda")
b_ptrs = torch.empty((E,), dtype=torch.int64, device="cuda")
out_ptrs = torch.empty((E,), dtype=torch.int64, device="cuda")
a_scales_ptrs = torch.empty((E,), dtype=torch.int64, device="cuda")
b_scales_ptrs = torch.empty((E,), dtype=torch.int64, device="cuda")
expert_offsets = torch.empty((E + 1,), dtype=torch.int32, device="cuda")
problem_sizes1 = torch.empty((E, 3), dtype=torch.int32, device="cuda")
problem_sizes2 = torch.empty((E, 3), dtype=torch.int32, device="cuda")
enable_es = (False, False)
if torch.cuda.get_device_name(torch.cuda.current_device()) == "NVIDIA H200":
enable_es = (False, True)
elif torch.cuda.get_device_name(torch.cuda.current_device()) == "NVIDIA H20":
enable_es = (True, True)
# --- Lambdas for Benchmarking ---
cutlass_lambda = lambda: cutlass_fused_experts_fp8(
x,
w1.transpose(1, 2), # Transposed
w2.transpose(1, 2), # Transposed
w1_scale.transpose(1, 2),
w2_scale.transpose(1, 2),
topk_weights,
topk_ids,
a1_strides,
c1_strides,
a2_strides,
c2_strides,
workspace,
a_ptrs,
b_ptrs,
out_ptrs,
a_scales_ptrs,
b_scales_ptrs,
expert_offsets,
problem_sizes1,
problem_sizes2,
enable_es=enable_es,
)
topk_output = StandardTopKOutput(
topk_weights=topk_weights,
topk_ids=topk_ids,
router_logits=torch.randn(
(batch_size, topk), device=topk_weights.device, dtype=dtype
),
)
moe_runner_config = MoeRunnerConfig(
num_experts=E,
top_k=topk,
hidden_size=H,
intermediate_size_per_partition=I,
params_dtype=dtype,
activation="silu",
inplace=False,
)
# Note: Triton expects non-transposed weights
triton_lambda = lambda: fused_experts(
x,
w1,
w2,
topk_output,
moe_runner_config,
use_fp8_w8a8=True,
w1_scale=w1_scale,
w2_scale=w2_scale,
block_shape=block_shape,
)
# --- Warmup ---
print("Warming up...")
for _ in range(10):
_ = cutlass_lambda()
_ = triton_lambda()
torch.cuda.synchronize()
# --- Benchmarking ---
quantiles = [0.5, 0.2, 0.8]
print(f"Benchmarking Cutlass fused_experts...")
cutlass_ms, cutlass_min, cutlass_max = triton.testing.do_bench_cudagraph(
cutlass_lambda, rep=1000, quantiles=quantiles
)
print(f"Benchmarking Triton fused_experts...")
triton_ms, triton_min, triton_max = triton.testing.do_bench_cudagraph(
triton_lambda, rep=1000, quantiles=quantiles
)
print(
f"Cutlass fused_experts time: {cutlass_ms:.3f} ms (median) [{cutlass_min:.3f} - {cutlass_max:.3f}]"
)
print(
f"Triton fused_experts time: {triton_ms:.3f} ms (median) [{triton_min:.3f} - {triton_max:.3f}]"
)
# --- Correctness Check ---
if check:
print("Running correctness check...")
with torch.no_grad():
# Run CUTLASS version (requires transposed weights)
y_cutlass = cutlass_fused_experts_fp8(
x,
w1.transpose(1, 2), # Transposed
w2.transpose(1, 2), # Transposed
w1_scale.transpose(1, 2),
w2_scale.transpose(1, 2),
topk_weights,
topk_ids,
a1_strides,
c1_strides,
a2_strides,
c2_strides,
workspace,
a_ptrs,
b_ptrs,
out_ptrs,
a_scales_ptrs,
b_scales_ptrs,
expert_offsets,
problem_sizes1,
problem_sizes2,
enable_es=enable_es,
)
# Run Triton version (requires original shape weights, use inplace=False)
y_triton = fused_experts(
x,
w1, # Original shape
w2, # Original shape
topk_output,
moe_runner_config,
use_fp8_w8a8=True,
w1_scale=w1_scale,
w2_scale=w2_scale,
block_shape=block_shape,
)
diff = calc_diff(y_cutlass, y_triton)
print(f"Diff: {diff:.6f}")
# Tolerance might need adjustment based on FP8 specifics and kernel differences
# FP8 comparisons often require higher tolerance than FP16/BF16
assert diff < 1e-4, f"Diff too high! {diff}"
print("Correctness check passed.")
def main(tp_size=8, batch_sizes=[1, 4, 8, 16, 32, 64, 128, 256, 512], check=False):
model_config = get_model_config(tp_size)
print("Model Config:", model_config)
for batch_size in batch_sizes:
run_test(tp_size, batch_size, model_config, check)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--tp-size", type=int, default=8, help="Tensor Parallel size")
parser.add_argument(
"--batch-sizes",
type=int,
nargs="+",
default=[
1,
4,
8,
16,
32,
64,
128,
256,
512,
1024,
2048,
4096,
8192,
], # Adjusted default
help="List of batch sizes to test",
)
parser.add_argument("--check", action="store_true", help="Enable check mode")
args = parser.parse_args()
print(f"Running benchmarks with TP size: {args.tp_size}")
print(f"Testing batch sizes: {args.batch_sizes}")
main(tp_size=args.tp_size, batch_sizes=args.batch_sizes, check=args.check)
python/sglang/test/test_cutlass_w16a16_moe.py
-118
View File
@@ -1,118 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
import pytest
import torch
from flashinfer.fused_moe import cutlass_fused_moe as flashinfer_cutlass_fused_moe
from sglang.srt.layers.activation import SiluAndMul
from sglang.srt.layers.moe.topk import TopKConfig, select_experts
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
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 for unquantized weights
def torch_moe_reference(a, w13, w2, score, topk):
B, D = a.shape
set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
# Flip w13 layout
dim = -2
size = w13.size(dim)
assert size % 2 == 0, f"Expected even size in dim {dim}, got {size}"
half = size // 2
# Reorder weight
w1, w3 = w13.split(half, dim=dim)
w13 = torch.cat([w3, w1], dim=dim).contiguous()
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)
for i in range(w13.shape[0]):
mask = topk_ids == i
if mask.sum():
out[mask] = SiluAndMul()(a[mask] @ w13[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)
@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
@pytest.mark.parametrize("e", [40, 64, 256])
@pytest.mark.parametrize("topk", [1, 6, 8])
@torch.inference_mode()
def test_flashinfer_bf16_cutlass_moe(m: int, n: int, k: int, e: int, topk: int):
"""
Test the bf16 cutlass moe API.
Args:
m: number of tokens
n: intermediate size
k: hidden size
e: number of experts
topk: top-k experts per token
"""
torch.manual_seed(7)
dtype = torch.bfloat16
# Create unquantized weights
a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
# w13: fused gate_up projection [num_experts, 2*intermediate, hidden]
# FlashInfer CUTLASS expects [up, gate] layout
w13 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
# w2: down projection [num_experts, hidden, intermediate]
w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
# Generate router scores
score = torch.randn((m, e), device="cuda", dtype=dtype)
# Get topk routing
topk_output = select_experts(
hidden_states=a,
router_logits=score,
topk_config=TopKConfig(top_k=topk, renormalize=False),
)
topk_weights, topk_ids, _ = topk_output
# Test: Call FlashInfer CUTLASS fused_moe (unquantized version)
test_output = flashinfer_cutlass_fused_moe(
input=a,
token_selected_experts=topk_ids,
token_final_scales=topk_weights,
fc1_expert_weights=w13,
fc2_expert_weights=w2,
output_dtype=dtype,
quant_scales=None,
)[0]
# Reference: Torch implementation
torch_output = torch_moe_reference(a, w13, w2, score, topk)
# Compare outputs
torch.testing.assert_close(torch_output, test_output, rtol=1e-2, atol=1e-2)
if __name__ == "__main__":
# Run a simple test case
test_flashinfer_bf16_cutlass_moe(224, 1024, 1024, 8, 2)
python/sglang/test/test_cutlass_w4a8_moe.py
-285
View File
@@ -1,285 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
from typing import Optional
import pytest
import torch
from sglang.srt.layers.moe.cutlass_w4a8_moe import cutlass_w4a8_moe
from sglang.srt.layers.moe.topk import TopKConfig, select_experts
def pack_int4_values_to_int8(int4_values_interleaved: torch.Tensor) -> torch.Tensor:
if int4_values_interleaved.shape[-1] % 2 != 0:
raise ValueError(
"the last dim size of int4_values_interleaved tensor must be even."
)
input_tensor_int8 = int4_values_interleaved.to(torch.int8)
low_nibbles = input_tensor_int8[..., 0::2]
high_nibbles = input_tensor_int8[..., 1::2]
packed_tensor = (high_nibbles << 4) | (low_nibbles & 0x0F)
return packed_tensor.to(torch.int8)
def pack_interleave(num_experts, ref_weight, ref_scale, alignment=4):
n, k = ref_weight.shape[1], ref_weight.shape[2]
weight = pack_int4_values_to_int8(ref_weight.cpu()).cuda()
w_q = weight.view((num_experts, n, k // 2)).view(torch.int8)
w_q = w_q.contiguous()
scale_interleaved = ref_scale.reshape(
ref_scale.shape[0],
ref_scale.shape[1],
(ref_scale.shape[2] // alignment),
alignment,
) # [E, N, K/4, 4]
scale_interleaved = scale_interleaved.permute(0, 2, 1, 3) # [E, K/4, N, 4]
scale_interleaved = scale_interleaved.reshape(
ref_scale.shape[0],
ref_scale.shape[2] // alignment,
ref_scale.shape[1] * alignment,
) # [E, K/4, N*4]
w_scale = scale_interleaved.contiguous()
return w_q, w_scale
@pytest.mark.parametrize("M", [1, 2, 4, 8, 16])
@pytest.mark.parametrize("N", [2048])
@pytest.mark.parametrize("K", [7168])
@pytest.mark.parametrize("E", [256])
@pytest.mark.parametrize("tp_size", [8])
@pytest.mark.parametrize("use_ep_moe", [True, False])
@pytest.mark.parametrize("topk", [8])
@pytest.mark.parametrize("group_size", [128])
@pytest.mark.parametrize("dtype", [torch.bfloat16])
def test_cutlass_w4a8_moe(M, N, K, E, tp_size, use_ep_moe, topk, group_size, dtype):
if use_ep_moe:
local_e = E // tp_size
else: # tp mode
local_e = E
N = N // tp_size
debug = False
if debug:
a = torch.ones((M, K), dtype=dtype, device="cuda") * 0.001
ref_weight_1 = torch.ones((local_e, N * 2, K), dtype=torch.int8, device="cuda")
ref_weight_2 = torch.ones((local_e, K, N), dtype=torch.int8, device="cuda")
a1_scale = torch.ones(1, dtype=torch.float32, device="cuda")
a2_scale = torch.ones(1, dtype=torch.float32, device="cuda")
scale_1 = torch.ones(
(local_e, N * 2, K // group_size), dtype=dtype, device="cuda"
)
scale_2 = torch.ones((local_e, K, N // group_size), dtype=dtype, device="cuda")
else:
a = torch.randn(M, K, dtype=dtype, device="cuda")
ref_weight_1 = torch.randint(
-8, 8, (local_e, N * 2, K), dtype=torch.int8, device="cuda"
)
ref_weight_2 = torch.randint(
-8, 8, (local_e, K, N), dtype=torch.int8, device="cuda"
)
affine_coeff = 0.005
a1_scale = torch.randn(1, dtype=torch.float32, device="cuda")
a2_scale = torch.randn(1, dtype=torch.float32, device="cuda")
scale_1 = (
torch.randn(local_e, N * 2, K // group_size, dtype=dtype, device="cuda")
* affine_coeff
)
scale_2 = (
torch.randn(local_e, K, N // group_size, dtype=dtype, device="cuda")
* affine_coeff
)
w1_q, w1_scale = pack_interleave(local_e, ref_weight_1, scale_1)
if use_ep_moe:
w2_q, w2_scale = pack_interleave(local_e, ref_weight_2, scale_2)
else:
w2_q, w2_scale = pack_interleave(local_e, ref_weight_2, scale_2, 1)
device = "cuda"
a_strides1 = torch.full((local_e, 3), K, device=device, dtype=torch.int64)
c_strides1 = torch.full((local_e, 3), 2 * N, device=device, dtype=torch.int64)
a_strides2 = torch.full((local_e, 3), N, device=device, dtype=torch.int64)
c_strides2 = torch.full((local_e, 3), K, device=device, dtype=torch.int64)
b_strides1 = a_strides1
s_strides13 = c_strides1
b_strides2 = a_strides2
s_strides2 = c_strides2
score = torch.randn((M, E), dtype=dtype, device=device)
topk_output = select_experts(
hidden_states=a,
router_logits=score,
topk_config=TopKConfig(top_k=topk, renormalize=False),
)
topk_weights, topk_ids, _ = topk_output
expert_map = torch.arange(E, dtype=torch.int32, device=device)
expert_map[local_e:] = -1
output = cutlass_moe(
a,
w1_q,
w2_q,
w1_scale,
w2_scale,
topk_weights,
topk_ids,
a_strides1,
b_strides1,
c_strides1,
a_strides2,
b_strides2,
c_strides2,
s_strides13,
s_strides2,
local_e,
a1_scale,
a2_scale,
expert_map,
)
ref_output = ref(
a,
local_e,
topk_weights,
topk_ids,
ref_weight_1,
ref_weight_2,
scale_1,
scale_2,
has_pre_quant=True,
has_alpha=True,
pre_quant_scale_1=a1_scale,
pre_quant_scale_2=a2_scale,
alpha_1=a1_scale,
alpha_2=a2_scale,
)
# compare
torch.cuda.synchronize()
# compare final output
torch.testing.assert_close(output, ref_output, rtol=1e-2, atol=0.1)
print("SUCCESS: Final output tensors are close.")
def cutlass_moe(
a: torch.Tensor,
w1_q: torch.Tensor,
w2_q: torch.Tensor,
w1_scale: torch.Tensor,
w2_scale: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
a_strides1: torch.Tensor,
b_strides1: torch.Tensor,
c_strides1: torch.Tensor,
a_strides2: torch.Tensor,
b_strides2: torch.Tensor,
c_strides2: torch.Tensor,
s_strides13: torch.Tensor,
s_strides2: torch.Tensor,
num_local_experts: int,
a1_scale: Optional[torch.Tensor] = None,
a2_scale: Optional[torch.Tensor] = None,
expert_map: Optional[torch.Tensor] = None,
apply_router_weight_on_input: bool = False,
):
topk_ids = expert_map[topk_ids]
device = a.device
expert_offsets = torch.empty(
(num_local_experts + 1), dtype=torch.int32, device=device
)
problem_sizes1 = torch.empty(
(num_local_experts, 3), dtype=torch.int32, device=device
)
problem_sizes2 = torch.empty(
(num_local_experts, 3), dtype=torch.int32, device=device
)
return cutlass_w4a8_moe(
a,
w1_q,
w2_q,
w1_scale,
w2_scale,
topk_weights,
topk_ids,
a_strides1,
b_strides1,
c_strides1,
a_strides2,
b_strides2,
c_strides2,
s_strides13,
s_strides2,
expert_offsets,
problem_sizes1,
problem_sizes2,
a1_scale,
a2_scale,
apply_router_weight_on_input,
)
def ref(
x: torch.Tensor,
num_experts: int,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
ref_weight_1: torch.Tensor,
ref_weight_2: torch.Tensor,
ref_weight_scale_1: torch.Tensor,
ref_weight_scale_2: torch.Tensor,
has_pre_quant: bool = False,
has_alpha: bool = False,
pre_quant_scale_1: Optional[torch.Tensor] = None,
pre_quant_scale_2: Optional[torch.Tensor] = None,
alpha_1: Optional[torch.Tensor] = None,
alpha_2: Optional[torch.Tensor] = None,
):
results = torch.zeros_like(x)
dtype = x.dtype
for e_idx in range(num_experts):
mask = topk_ids == e_idx
activated_tokens = mask.sum(1).bool()
act = x[activated_tokens, :]
if act.shape[0] == 0:
continue
final_scale = (topk_weights * mask).sum(1)[activated_tokens].unsqueeze(1)
act = (
torch.clamp((act / pre_quant_scale_1.float()), -448.0, 448.0)
.to(torch.float8_e4m3fn)
.to(dtype)
)
w3_w1 = ref_weight_1[e_idx]
ref_w_scale_repeat = (
ref_weight_scale_1[e_idx].repeat_interleave(128, dim=1).to(float)
)
w3_w1 = (w3_w1.to(float) * ref_w_scale_repeat).to(dtype)
fc1 = ((torch.matmul(act, w3_w1.T)) * alpha_1).to(torch.float16)
gate, fc1 = fc1.chunk(2, dim=-1)
fc1 = fc1 * torch.nn.functional.silu(gate)
act = torch.clamp((fc1 / pre_quant_scale_2.float()), -448.0, 448.0).to(
torch.float8_e4m3fn
)
act = act.to(dtype)
w2 = ref_weight_2[e_idx]
ref_w_scale_repeat = (
ref_weight_scale_2[e_idx].repeat_interleave(128, dim=1).to(float)
)
w2 = (w2.to(float) * ref_w_scale_repeat).to(dtype)
fc2 = (torch.matmul(act, w2.T) * alpha_2).to(torch.float16)
results[activated_tokens, :] += (fc2 * final_scale).to(results.dtype)
return results
python/sglang/test/test_dump_metric.py
-132
View File
@@ -1,132 +0,0 @@
"""Unit tests for dump_metric() function."""
import json
import os
import tempfile
import unittest
from pathlib import Path
from sglang.test.test_utils import dump_metric
class TestDumpMetric(unittest.TestCase):
"""Test suite for dump_metric() function."""
_ENV_KEYS_TO_CLEAN = ["SGLANG_TEST_METRICS_OUTPUT", "PYTEST_CURRENT_TEST"]
def setUp(self):
"""Clean up env vars before each test."""
for key in self._ENV_KEYS_TO_CLEAN:
os.environ.pop(key, None)
def tearDown(self):
"""Clean up env vars after each test."""
for key in self._ENV_KEYS_TO_CLEAN:
os.environ.pop(key, None)
def test_writes_valid_jsonl(self):
"""Test that dump_metric writes one valid JSON line when env is set."""
with tempfile.TemporaryDirectory() as tmpdir:
base_path = os.path.join(tmpdir, "metrics")
os.environ["SGLANG_TEST_METRICS_OUTPUT"] = base_path
dump_metric("test_accuracy", 0.95, labels={"model": "llama"})
# Check file exists with PID suffix
pid = os.getpid()
jsonl_path = f"{base_path}.{pid}.jsonl"
self.assertTrue(os.path.exists(jsonl_path))
# Read and validate
with open(jsonl_path, encoding="utf-8") as f:
lines = f.readlines()
self.assertEqual(len(lines), 1)
record = json.loads(lines[0])
# Validate required fields
self.assertIn("filename", record)
self.assertIn("test_case", record)
self.assertEqual(record["metric_name"], "test_accuracy")
self.assertEqual(record["value"], 0.95)
# Validate optional fields
self.assertIn("ts", record)
self.assertIsInstance(record["ts"], (int, float))
self.assertEqual(record["labels"], {"model": "llama"})
def test_no_env_no_file(self):
"""Test that dump_metric doesn't create file when env var not set."""
with tempfile.TemporaryDirectory() as tmpdir:
# Don't set env var
dump_metric("test_metric", 42)
# Verify no files created
files = list(Path(tmpdir).glob("*.jsonl"))
self.assertEqual(len(files), 0)
def test_labels_not_serializable_stringified(self):
"""Test that non-serializable labels are stringified."""
with tempfile.TemporaryDirectory() as tmpdir:
base_path = os.path.join(tmpdir, "metrics")
os.environ["SGLANG_TEST_METRICS_OUTPUT"] = base_path
# Non-serializable label
class NonSerializable:
pass
dump_metric("test_metric", 100, labels={"obj": NonSerializable()})
pid = os.getpid()
jsonl_path = f"{base_path}.{pid}.jsonl"
with open(jsonl_path, encoding="utf-8") as f:
lines = f.readlines()
record = json.loads(lines[0])
self.assertIn("labels", record)
self.assertIsInstance(record["labels"], str)
def test_bool_to_int(self):
"""Test that bool values are converted to int."""
with tempfile.TemporaryDirectory() as tmpdir:
base_path = os.path.join(tmpdir, "metrics")
os.environ["SGLANG_TEST_METRICS_OUTPUT"] = base_path
dump_metric("bool_true", True)
dump_metric("bool_false", False)
pid = os.getpid()
jsonl_path = f"{base_path}.{pid}.jsonl"
with open(jsonl_path, encoding="utf-8") as f:
lines = f.readlines()
self.assertEqual(len(lines), 2)
record1 = json.loads(lines[0])
record2 = json.loads(lines[1])
self.assertEqual(record1["value"], 1) # True -> 1
self.assertEqual(record2["value"], 0) # False -> 0
def test_pytest_current_test_parsing(self):
"""Test PYTEST_CURRENT_TEST parsing for test_case."""
with tempfile.TemporaryDirectory() as tmpdir:
base_path = os.path.join(tmpdir, "metrics")
os.environ["SGLANG_TEST_METRICS_OUTPUT"] = base_path
os.environ["PYTEST_CURRENT_TEST"] = (
"test/srt/test_example.py::TestClass::test_method (call)"
)
dump_metric("pytest_metric", 123)
pid = os.getpid()
jsonl_path = f"{base_path}.{pid}.jsonl"
with open(jsonl_path, encoding="utf-8") as f:
lines = f.readlines()
record = json.loads(lines[0])
# Only assert test_case parsing, not filename
self.assertEqual(record["test_case"], "TestClass.test_method")
if __name__ == "__main__":
unittest.main()
python/sglang/test/test_dynamic_grad_mode.py
-58
View File
@@ -1,58 +0,0 @@
import unittest
import torch
from sglang.srt.utils import DynamicGradMode
from sglang.test.test_utils import CustomTestCase
class TestDynamicGradMode(CustomTestCase):
def test_inference(self):
# Test inference_mode
DynamicGradMode.set_inference_mode(True)
@DynamicGradMode()
def create_tensor_x():
return torch.empty(0)
X = create_tensor_x()
self.assertTrue(not X.requires_grad and X.is_inference())
def test_no_grad(self):
# Test no_grad
DynamicGradMode.set_inference_mode(False)
@DynamicGradMode()
def create_tensor_y():
return torch.empty(0)
Y = create_tensor_y()
self.assertTrue(not Y.requires_grad and not Y.is_inference())
def test_nested_inference(self):
# Test no_grad nested inference_mode, inference_mode should has higher priority
DynamicGradMode.set_inference_mode(False)
@DynamicGradMode()
def create_tensor_z():
with torch.inference_mode():
return torch.empty(0)
Z = create_tensor_z()
self.assertTrue(not Z.requires_grad and Z.is_inference())
def test_nested_no_grad(self):
# Test inference_mode nested no_grad, inference_mode should has higher priority
DynamicGradMode.set_inference_mode(True)
@DynamicGradMode()
def create_tensor_w():
with torch.no_grad():
return torch.empty(0)
W = create_tensor_w()
self.assertTrue(not W.requires_grad and W.is_inference())
if __name__ == "__main__":
unittest.main(verbosity=2)
python/sglang/test/test_flashinfer_dispatcher.py
-322
View File
@@ -1,322 +0,0 @@
import unittest
import torch
from sglang.srt.distributed import init_distributed_environment
from sglang.srt.distributed.parallel_state import (
get_tp_group,
initialize_model_parallel,
)
from sglang.srt.layers.dp_attention import set_dp_buffer_len
from sglang.srt.layers.moe.token_dispatcher.flashinfer import FlashinferDispatcher
from sglang.srt.layers.moe.utils import initialize_moe_config
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
from sglang.test.test_utils import CustomTestCase
class TestFlashinferDispatcher(CustomTestCase):
@classmethod
def setUpClass(cls):
server_args = ServerArgs(model_path="dummy")
server_args.moe_runner_backend = "flashinfer_cutlass"
server_args.moe_a2a_backend = "flashinfer"
set_global_server_args_for_scheduler(server_args)
initialize_moe_config(server_args)
init_distributed_environment(
world_size=-1, # Auto-detect from environment
rank=-1, # Auto-detect from environment
local_rank=-1, # Auto-detect from environment
backend="nccl",
)
world_size = torch.distributed.get_world_size()
rank = torch.distributed.get_rank()
device = torch.device(f"cuda:{rank % torch.cuda.device_count()}")
torch.cuda.set_device(device)
initialize_model_parallel(
tensor_model_parallel_size=world_size, expert_model_parallel_size=world_size
)
@classmethod
def tearDownClass(cls):
# Clean up distributed environment
if torch.distributed.is_initialized():
torch.distributed.destroy_process_group()
def create_dispatcher(
self, router_topk=2, num_experts=8, num_local_experts=4, hidden_size=128
):
"""Helper to create dispatcher instance"""
return FlashinferDispatcher(
group=get_tp_group().device_group,
router_topk=router_topk,
num_experts=num_experts,
num_local_experts=num_local_experts,
hidden_size=hidden_size,
params_dtype=torch.bfloat16,
)
def test_dispatch_basic(self):
"""Test basic dispatch functionality"""
num_tokens = 16
hidden_size = 128
router_topk = 1 # Single expert per token for simplicity
world_size = torch.distributed.get_world_size()
rank = torch.distributed.get_rank()
num_experts = world_size
num_local_experts = 1 # One expert per rank
set_dp_buffer_len(
global_dp_buffer_len=num_tokens * world_size,
local_dp_buffer_len=num_tokens,
dp_max_padding=True,
global_num_tokens=None,
)
# Create tokens with rank number
hidden_states = torch.full(
(num_tokens, hidden_size), 100.0 + rank, dtype=torch.bfloat16, device="cuda"
)
# Route all tokens from rank i to expert (i+1) % world_size
target_rank = (rank + 1) % world_size
target_expert = target_rank # Since we have 1 expert per rank
topk_ids = torch.full(
(num_tokens, router_topk), target_expert, dtype=torch.int32, device="cuda"
)
topk_weights = torch.ones(
(num_tokens, router_topk), dtype=torch.float32, device="cuda"
)
from sglang.srt.layers.moe.topk import StandardTopKOutput
topk_output = StandardTopKOutput(
topk_weights=topk_weights, topk_ids=topk_ids, router_logits=None
)
torch.distributed.barrier()
dispatcher = self.create_dispatcher(
router_topk=router_topk,
num_experts=num_experts,
num_local_experts=num_local_experts,
hidden_size=hidden_size,
)
dispatcher.set_quant_config({"input_global_scale": None})
dispatch_output = dispatcher.dispatch(hidden_states, topk_output)
received_hidden_states = dispatch_output.hidden_states
self.assertEqual(dispatch_output.hidden_states_scale, None)
# Expected: we should receive tokens from rank (rank - 1) % world_size
expected_source_rank = (rank - 1 + world_size) % world_size
# Verify we received the right number of tokens
self.assertEqual(
received_hidden_states.shape[0],
num_tokens * world_size,
f"Should receive {num_tokens * world_size} tokens",
)
# Verify tokens came from the expected source
self.assertTrue(
torch.all(
received_hidden_states[
expected_source_rank
* num_tokens : (expected_source_rank + 1)
* num_tokens
]
== 100.0 + expected_source_rank
)
)
self.assertTrue(
torch.all(
received_hidden_states[: expected_source_rank * num_tokens] == 0.0
)
)
self.assertTrue(
torch.all(
received_hidden_states[(expected_source_rank + 1) * num_tokens :] == 0.0
)
)
def test_dispatch_with_empty_tokens(self):
"""Test dispatch when there are no tokens (edge case)"""
# This tests the dummy token handling
num_tokens = 16
hidden_size = 1
router_topk = 1 # Single expert per token for simplicity
world_size = torch.distributed.get_world_size()
rank = torch.distributed.get_rank()
num_experts = world_size
num_local_experts = 1 # One expert per rank
set_dp_buffer_len(
global_dp_buffer_len=num_tokens * world_size,
local_dp_buffer_len=num_tokens,
dp_max_padding=False,
global_num_tokens=[16, 0, 16, 16],
)
# Route all tokens from rank i to expert (i+1) % world_size
target_rank = (rank + 1) % world_size
target_expert = target_rank # Since we have 1 expert per rank
# Create tokens with rank number, rank 1 has no tokens
if rank == 1:
hidden_states = torch.empty(
0, hidden_size, dtype=torch.bfloat16, device="cuda"
)
topk_ids = torch.empty(0, router_topk, dtype=torch.int32, device="cuda")
topk_weights = torch.empty(
0, router_topk, dtype=torch.float32, device="cuda"
)
else:
hidden_states = torch.full(
(num_tokens, hidden_size),
100.0 + rank,
dtype=torch.bfloat16,
device="cuda",
)
topk_ids = torch.full(
(num_tokens, router_topk),
target_expert,
dtype=torch.int32,
device="cuda",
)
topk_weights = torch.ones(
(num_tokens, router_topk), dtype=torch.float32, device="cuda"
)
from sglang.srt.layers.moe.topk import StandardTopKOutput
topk_output = StandardTopKOutput(
topk_weights=topk_weights, topk_ids=topk_ids, router_logits=None
)
dispatcher = self.create_dispatcher(
router_topk=router_topk,
num_experts=num_experts,
num_local_experts=num_local_experts,
hidden_size=hidden_size,
)
dispatcher.set_quant_config({"input_global_scale": None})
dispatch_output = dispatcher.dispatch(hidden_states, topk_output)
received_hidden_states = dispatch_output.hidden_states
# Expected: we should receive tokens from rank (rank - 1) % world_size
expected_source_rank = (rank - 1 + world_size) % world_size
# Verify we received the right number of tokens
self.assertEqual(
received_hidden_states.shape[0],
num_tokens * world_size,
f"Should receive {num_tokens * world_size} tokens",
)
# Verify tokens came from the expected source
if rank == 2:
# Rank 2 should receive no tokens since rank 1 was empty
self.assertTrue(
torch.all(received_hidden_states == 0.0),
"Rank should receive no tokens",
)
else:
self.assertTrue(
torch.all(
received_hidden_states[
expected_source_rank
* num_tokens : (expected_source_rank + 1)
* num_tokens
]
== 100.0 + expected_source_rank
),
"Rank {rank} should receive tokens from the expected source {expected_source_rank}",
)
self.assertTrue(
torch.all(
received_hidden_states[: expected_source_rank * num_tokens] == 0.0
),
"Rank should receive no tokens from previous ranks",
)
self.assertTrue(
torch.all(
received_hidden_states[(expected_source_rank + 1) * num_tokens :]
== 0.0
),
"Rank should receive no tokens from next ranks",
)
def test_dispatch_with_fp4_quantization(self):
"""Test dispatch with FP4 quantization enabled"""
num_tokens = 128
hidden_size = 128
router_topk = 1 # Single expert per token for simplicity
world_size = torch.distributed.get_world_size()
rank = torch.distributed.get_rank()
num_experts = world_size
num_local_experts = 1 # One expert per rank
set_dp_buffer_len(
global_dp_buffer_len=num_tokens * world_size,
local_dp_buffer_len=num_tokens,
dp_max_padding=True,
global_num_tokens=None,
)
# Create tokens with random values
hidden_states = torch.randn(
(num_tokens, hidden_size), dtype=torch.bfloat16, device="cuda"
)
# Route all tokens from rank i to expert (i+1) % world_size
target_rank = (rank + 1) % world_size
target_expert = target_rank # Since we have 1 expert per rank
topk_ids = torch.full(
(num_tokens, router_topk), target_expert, dtype=torch.int32, device="cuda"
)
topk_weights = torch.ones(
(num_tokens, router_topk), dtype=torch.float32, device="cuda"
)
from sglang.srt.layers.moe.topk import StandardTopKOutput
topk_output = StandardTopKOutput(
topk_weights=topk_weights, topk_ids=topk_ids, router_logits=None
)
dispatcher = self.create_dispatcher(
router_topk=router_topk,
num_experts=num_experts,
num_local_experts=num_local_experts,
hidden_size=hidden_size,
)
# Set input global scale to enable FP4 quantization
input_global_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
dispatcher.set_quant_config({"input_global_scale": input_global_scale})
dispatch_output = dispatcher.dispatch(hidden_states, topk_output)
self.assertEqual(
dispatch_output.hidden_states.shape,
(num_tokens * world_size, hidden_size // 2),
)
self.assertEqual(dispatch_output.hidden_states.dtype, torch.uint8)
self.assertNotEqual(dispatch_output.hidden_states_scale, None)
self.assertEqual(
dispatch_output.hidden_states_scale.numel(),
num_tokens * world_size * (hidden_size // 16),
)
self.assertEqual(dispatch_output.hidden_states_scale.dtype, torch.uint8)
if __name__ == "__main__":
"""
Usage
torchrun --nproc_per_node=4 test_flashinfer_dispatcher.py
"""
unittest.main()
python/sglang/test/test_kvfp4_quant_dequant.py
-116
View File
@@ -1,116 +0,0 @@
#!/usr/bin/env python3
import time
import numpy as np
import pytest
import torch
from sglang.srt.layers.quantization.kvfp4_tensor import KVFP4QuantizeUtil
def calculate_accuracy_metrics(
original: torch.Tensor, reconstructed: torch.Tensor
) -> dict[str, float]:
"""Calculate accuracy metrics between original and reconstructed tensors."""
mse = torch.mean((original - reconstructed) ** 2).item()
mae = torch.mean(torch.abs(original - reconstructed)).item()
# PSNR calculation
max_val = torch.max(torch.abs(original)).item()
psnr = 20 * np.log10(max_val / np.sqrt(mse)) if mse > 0 else float("inf")
# Relative error
rel_error = torch.mean(
torch.abs(original - reconstructed) / (torch.abs(original) + 1e-8)
).item()
return {"MSE": mse, "MAE": mae, "PSNR": psnr, "Relative Error": rel_error}
def run_benchmark(m, n, k, num_runs=100) -> dict[str, dict[str, float]]:
"""Run FP8 vs KVFP4 quantization benchmark and return metrics."""
tensor_bf16 = torch.randn(m, n, k, dtype=torch.bfloat16, device="cuda")
# --- FP8 ---
for _ in range(3): # warmup
_ = tensor_bf16 * 2
torch.cuda.synchronize()
start = time.time()
for _ in range(num_runs):
tensor_fp8 = tensor_bf16.to(torch.float8_e4m3fn)
torch.cuda.synchronize()
fp8_quant_time = (time.time() - start) / num_runs
start = time.time()
for _ in range(num_runs):
tensor_fp8_dequant = tensor_fp8.to(torch.bfloat16)
torch.cuda.synchronize()
fp8_dequant_time = (time.time() - start) / num_runs
fp8_metrics = calculate_accuracy_metrics(tensor_bf16, tensor_fp8_dequant)
# --- KVFP4 ---
tensor_fp4, scale_factors = KVFP4QuantizeUtil.batched_quantize(tensor_bf16)
_ = KVFP4QuantizeUtil.batched_dequantize(tensor_fp4, scale_factors)
start = time.time()
for _ in range(num_runs):
tensor_fp4, scale_factors = KVFP4QuantizeUtil.batched_quantize(tensor_bf16)
torch.cuda.synchronize()
fp4_quant_time = (time.time() - start) / num_runs
start = time.time()
for _ in range(num_runs):
tensor_fp4_dequant = KVFP4QuantizeUtil.batched_dequantize(
tensor_fp4, scale_factors
)
torch.cuda.synchronize()
fp4_dequant_time = (time.time() - start) / num_runs
fp4_metrics = calculate_accuracy_metrics(tensor_bf16, tensor_fp4_dequant)
return {
"fp8": {
"quant_time": fp8_quant_time,
"dequant_time": fp8_dequant_time,
**fp8_metrics,
},
"fp4": {
"quant_time": fp4_quant_time,
"dequant_time": fp4_dequant_time,
**fp4_metrics,
},
}
# default tensor shapes (m, n, k)
# [M, 1, 576]: DeepSeekR1-FP4 MLA
# [M, 8, 64]: gpt-oss-20b MHA
MNK_FACTORS = [
(64, 1, 576),
(512, 1, 576),
(1024, 1, 576),
(4096, 1, 576),
(2868672, 1, 576),
(64, 8, 64),
(512, 8, 64),
(1024, 8, 64),
(4096, 8, 64),
(2868672, 8, 64),
]
@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
def test_kvfp4_quant_dequant(m, n, k):
"""Benchmark FP8 vs KVFP4 for predefined tensor shapes."""
print(f"\n=== Running benchmark for tensor shape: [{m}, {n}, {k}] ===")
results = run_benchmark(m, n, k)
print("FP8:", results["fp8"])
print("FP4:", results["fp4"])
# Basic assertions to make sure metrics are reasonable
assert results["fp4"]["MSE"] < 1.0
assert results["fp8"]["MSE"] < 1.0
python/sglang/test/test_layernorm.py
-185
View File
@@ -1,185 +0,0 @@
import itertools
import unittest
import torch
from sglang.srt.layers.layernorm import GemmaRMSNorm, LayerNorm, RMSNorm
from sglang.test.test_utils import CustomTestCase
class TestRMSNorm(CustomTestCase):
DTYPES = [torch.half, torch.bfloat16]
NUM_TOKENS = [7, 83, 4096]
HIDDEN_SIZES = [768, 769, 770, 771, 5120, 5124, 5125, 5126, 8192, 8199]
ADD_RESIDUAL = [False, True]
SEEDS = [0]
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
torch.set_default_device("cuda")
def _run_rms_norm_test(self, num_tokens, hidden_size, add_residual, dtype, seed):
torch.manual_seed(seed)
layer = RMSNorm(hidden_size).to(dtype=dtype)
layer.weight.data.normal_(mean=1.0, std=0.1)
scale = 1 / (2 * hidden_size)
x = torch.randn(num_tokens, hidden_size, dtype=dtype) * scale
residual = torch.randn_like(x) * scale if add_residual else None
with torch.inference_mode():
ref_out = layer.forward_native(x, residual)
out = layer(x, residual)
if add_residual:
self.assertTrue(torch.allclose(out[0], ref_out[0], atol=1e-2, rtol=1e-2))
self.assertTrue(torch.allclose(out[1], ref_out[1], atol=1e-2, rtol=1e-2))
else:
self.assertTrue(torch.allclose(out, ref_out, atol=1e-2, rtol=1e-2))
def test_rms_norm(self):
for params in itertools.product(
self.NUM_TOKENS,
self.HIDDEN_SIZES,
self.ADD_RESIDUAL,
self.DTYPES,
self.SEEDS,
):
with self.subTest(
num_tokens=params[0],
hidden_size=params[1],
add_residual=params[2],
dtype=params[3],
seed=params[4],
):
self._run_rms_norm_test(*params)
class TestGemmaRMSNorm(CustomTestCase):
DTYPES = [torch.half, torch.bfloat16]
NUM_TOKENS = [7, 83, 4096]
HIDDEN_SIZES = [768, 769, 770, 771, 5120, 5124, 5125, 5126, 8192, 8199]
ADD_RESIDUAL = [False, True]
SEEDS = [0]
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
torch.set_default_device("cuda")
def _run_gemma_rms_norm_test(
self, num_tokens, hidden_size, add_residual, dtype, seed
):
torch.manual_seed(seed)
layer = GemmaRMSNorm(hidden_size).to(dtype=dtype)
layer.weight.data.normal_(mean=1.0, std=0.1)
scale = 1 / (2 * hidden_size)
x = torch.randn(num_tokens, hidden_size, dtype=dtype) * scale
residual = torch.randn_like(x) * scale if add_residual else None
with torch.inference_mode():
ref_out = layer.forward_native(x, residual)
out = layer(x, residual)
if add_residual:
self.assertTrue(torch.allclose(out[0], ref_out[0], atol=1e-3, rtol=1e-3))
self.assertTrue(torch.allclose(out[1], ref_out[1], atol=1e-3, rtol=1e-3))
else:
self.assertTrue(torch.allclose(out, ref_out, atol=1e-3, rtol=1e-3))
def test_gemma_rms_norm(self):
for params in itertools.product(
self.NUM_TOKENS,
self.HIDDEN_SIZES,
self.ADD_RESIDUAL,
self.DTYPES,
self.SEEDS,
):
with self.subTest(
num_tokens=params[0],
hidden_size=params[1],
add_residual=params[2],
dtype=params[3],
seed=params[4],
):
self._run_gemma_rms_norm_test(*params)
class TestLayerNorm(CustomTestCase):
DTYPES = [torch.half, torch.bfloat16]
PARAM_DTYPES = [torch.bfloat16, torch.float32]
NUM_TOKENS = [7, 83, 1024]
HIDDEN_SIZES = [128, 512, 1536, 5120, 5124, 5125, 5126, 7168]
USE_AFFINE = [False, True]
USE_BIAS = [False, True]
SEEDS = [0]
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
torch.set_default_device("cuda")
def _run_layer_norm_test(
self, num_tokens, hidden_size, use_affine, use_bias, dtype, seed, param_dtype
):
torch.manual_seed(seed)
layer = LayerNorm(
hidden_size, elementwise_affine=use_affine, bias=use_bias, dtype=param_dtype
)
if use_affine:
layer.weight.data.normal_(mean=1.0, std=0.1)
if use_bias:
layer.bias.data.normal_(mean=0.0, std=0.1)
scale = 1 / (2 * hidden_size)
x = torch.randn(num_tokens, hidden_size, dtype=dtype) * scale
with torch.inference_mode():
ref_out = layer.forward_native(x)
out = layer(x)
self.assertTrue(torch.allclose(out, ref_out, atol=1e-2, rtol=1e-3))
if (
use_affine
and use_bias
and not (dtype == torch.bfloat16 and param_dtype == torch.float32)
):
layer.dtype = torch.float32
layer.weight.data = layer.weight.data.to(torch.float32)
layer.bias.data = layer.bias.data.to(torch.float32)
with torch.inference_mode():
cuda_out = layer(x.to(torch.bfloat16)).to(x.dtype)
self.assertTrue(torch.allclose(cuda_out, ref_out, atol=2e-2, rtol=1e-3))
def test_layer_norm(self):
for params in itertools.product(
self.NUM_TOKENS,
self.HIDDEN_SIZES,
self.USE_AFFINE,
self.USE_BIAS,
self.DTYPES,
self.SEEDS,
self.PARAM_DTYPES,
):
with self.subTest(
num_tokens=params[0],
hidden_size=params[1],
use_affine=params[2],
use_bias=params[3],
dtype=params[4],
seed=params[5],
param_dtype=params[6],
):
self._run_layer_norm_test(*params)
if __name__ == "__main__":
unittest.main(verbosity=2)
python/sglang/test/test_mm_utils.py
-50
View File
@@ -1,50 +0,0 @@
import unittest
from unittest.mock import Mock, patch
import torch
from sglang.srt.managers import mm_utils, schedule_batch
from sglang.srt.managers.schedule_batch import (
Modality,
MultimodalDataItem,
MultimodalInputs,
)
def _make_proxy_with_reconstruct_result(tensor: torch.Tensor):
proxy = mm_utils.CudaIpcTensorTransportProxy.__new__(
mm_utils.CudaIpcTensorTransportProxy
)
proxy.reconstruct_on_target_device = Mock(return_value=tensor)
return proxy
class TestMultimodalInputsFromDict(unittest.TestCase):
def test_materialize_proxy(self):
feature_tensor = torch.tensor([[7.0], [8.0]], dtype=torch.float32)
proxy_feature = _make_proxy_with_reconstruct_result(feature_tensor)
mm_item = MultimodalDataItem(
modality=Modality.IMAGE,
offsets=[(0, 1), (1, 2)],
feature=proxy_feature,
model_specific_data={"image_grid_thw": [[1, 1, 1], [1, 1, 1]]},
)
with (
patch.object(schedule_batch.torch.cuda, "is_available", return_value=True),
patch.object(schedule_batch.torch.cuda, "current_device", return_value=0),
patch.object(
schedule_batch.envs.SGLANG_MM_BUFFER_SIZE_MB, "get", return_value=0
),
):
mm_inputs = MultimodalInputs.from_dict({"mm_items": [mm_item]})
# Splitting happens at the processor layer, not in from_dict.
# from_dict just reconstructs and passes through.
self.assertEqual(len(mm_inputs.mm_items), 1)
self.assertTrue(torch.equal(mm_inputs.mm_items[0].feature, feature_tensor))
proxy_feature.reconstruct_on_target_device.assert_called_once_with(0)
if __name__ == "__main__":
unittest.main(verbosity=2)