propagate pytest exit code from test __main__ entries (#24487)
This commit is contained in:
Liangsheng Yin
@@ -1,88 +0,0 @@
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.diffusion.cutedsl.norm_tanh_mul_add_norm_scale import (
|
||||
fused_norm_tanh_mul_add,
|
||||
fused_norm_tanh_mul_add_norm_scale,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=45, suite="stage-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=180, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
BSD_CONFIG = [
|
||||
(1, 3648, 3840), # Z-image
|
||||
(1, 4128, 3840), # Z-image
|
||||
(3, 7, 256), # bound
|
||||
(7, 1, 8192), # bound
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("B,S,D", BSD_CONFIG)
|
||||
@pytest.mark.parametrize("norm_type", ["rms", "layer"])
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
|
||||
def test_norm_tanh_mul_add(B: int, S: int, D: int, norm_type: str, dtype: str) -> None:
|
||||
device = "cuda"
|
||||
eps = 1e-5
|
||||
x = torch.randn(B, S, D, device=device, dtype=dtype)
|
||||
weight = torch.randn(D, device=device, dtype=dtype)
|
||||
bias = torch.randn(D, device=device, dtype=dtype) if norm_type == "layer" else None
|
||||
scale = torch.randn(B, 1, D, device=device, dtype=dtype)
|
||||
shift = torch.randn(B, 1, D, device=device, dtype=dtype)
|
||||
|
||||
y = fused_norm_tanh_mul_add(x, weight, bias, scale, shift, norm_type, eps)
|
||||
if norm_type == "rms":
|
||||
normed = torch.rms_norm(x, x.shape[-1:], weight=weight, eps=eps)
|
||||
else:
|
||||
normed = torch.layer_norm(x, x.shape[-1:], weight=weight, bias=bias, eps=eps)
|
||||
ref_y = normed * torch.tanh(scale) + shift
|
||||
# Accuracy check
|
||||
if dtype == "float32":
|
||||
torch.testing.assert_close(y, ref_y, atol=1e-5, rtol=1e-5)
|
||||
else:
|
||||
torch.testing.assert_close(y, ref_y, atol=5e-2, rtol=5e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("B,S,D", BSD_CONFIG)
|
||||
@pytest.mark.parametrize("norm_type", ["rms", "layer"])
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
|
||||
def test_norm_tanh_mul_add_norm_scale(
|
||||
B: int, S: int, D: int, norm_type: str, dtype: str
|
||||
) -> None:
|
||||
device = "cuda"
|
||||
eps = 1e-5
|
||||
x = torch.randn(B, S, D, device=device, dtype=dtype)
|
||||
weight = torch.randn(D, device=device, dtype=dtype)
|
||||
bias = torch.randn(D, device=device, dtype=dtype) if norm_type == "layer" else None
|
||||
scale = torch.randn(B, 1, D, device=device, dtype=dtype)
|
||||
shift = torch.randn(B, 1, D, device=device, dtype=dtype)
|
||||
weight2 = torch.randn(D, device=device, dtype=dtype)
|
||||
bias2 = torch.randn(D, device=device, dtype=dtype) if norm_type == "layer" else None
|
||||
scale2 = torch.randn(B, 1, D, device=device, dtype=dtype)
|
||||
|
||||
y, y2 = fused_norm_tanh_mul_add_norm_scale(
|
||||
x, weight, bias, scale, shift, weight2, bias2, scale2, norm_type, eps
|
||||
)
|
||||
if norm_type == "rms":
|
||||
normed = torch.rms_norm(x, x.shape[-1:], weight=weight, eps=eps)
|
||||
else:
|
||||
normed = torch.layer_norm(x, x.shape[-1:], weight=weight, bias=bias, eps=eps)
|
||||
ref_y = normed * torch.tanh(scale) + shift
|
||||
if norm_type == "rms":
|
||||
normed2 = torch.rms_norm(ref_y, ref_y.shape[-1:], weight=weight2, eps=eps)
|
||||
else:
|
||||
normed2 = torch.layer_norm(
|
||||
ref_y, ref_y.shape[-1:], weight=weight2, bias=bias2, eps=eps
|
||||
)
|
||||
ref_y2 = normed2 * (1 + scale2)
|
||||
# Accuracy check
|
||||
if dtype == "float32":
|
||||
torch.testing.assert_close(y, ref_y, atol=1e-5, rtol=1e-5)
|
||||
torch.testing.assert_close(y2, ref_y2, atol=1e-5, rtol=1e-5)
|
||||
else:
|
||||
torch.testing.assert_close(y, ref_y, atol=5e-2, rtol=5e-2)
|
||||
torch.testing.assert_close(y2, ref_y2, atol=5e-2, rtol=5e-2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__])
|
||||
@@ -1,314 +0,0 @@
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.cast import downcast_fp8
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=15, suite="stage-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
DTYPES = [torch.bfloat16, torch.float16]
|
||||
|
||||
# FP8 E4M3 representable range (matches kFP8E4M3Max in type.cuh)
|
||||
_FP8_E4M3_MAX = 448.0
|
||||
|
||||
|
||||
def _run(input_sl, head, dim, out_sl, dtype):
|
||||
k = torch.randn(input_sl, head, dim, dtype=dtype, device="cuda")
|
||||
v = torch.randn(input_sl, head, dim, dtype=dtype, device="cuda")
|
||||
k_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda")
|
||||
v_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda")
|
||||
k_scale = torch.tensor([1.0], dtype=torch.float32, device="cuda")
|
||||
v_scale = torch.tensor([1.0], dtype=torch.float32, device="cuda")
|
||||
loc = torch.arange(input_sl, dtype=torch.int64, device="cuda")
|
||||
downcast_fp8(k, v, k_out, v_out, k_scale, v_scale, loc)
|
||||
return k_out, v_out
|
||||
|
||||
|
||||
def _ref_fp8(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
|
||||
"""Reference: replicate kernel precision — scale_inv in dtype T, then to fp8.
|
||||
|
||||
Mirrors the kernel logic:
|
||||
scale_inv = cast<T>(1.0f) / cast<T>(scale[0])
|
||||
out[j] = cast<fp8_e4m3_t>(clamp(x[j] * scale_inv))
|
||||
"""
|
||||
dtype = x.dtype
|
||||
scale_inv = x.new_ones(1) / scale[0].to(dtype)
|
||||
x_scaled = (x * scale_inv).clamp(-_FP8_E4M3_MAX, _FP8_E4M3_MAX)
|
||||
return x_scaled.to(torch.float8_e4m3fn).view(torch.uint8)
|
||||
|
||||
|
||||
def _ref_downcast(
|
||||
x: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
loc: torch.Tensor,
|
||||
out_sl: int,
|
||||
mult: int = 1,
|
||||
offset: int = 0,
|
||||
) -> torch.Tensor:
|
||||
"""Scatter _ref_fp8 output to the correct output slots via loc/mult/offset."""
|
||||
head, dim = x.shape[1], x.shape[2]
|
||||
out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device=x.device)
|
||||
fp8 = _ref_fp8(x, scale)
|
||||
for i, dst in enumerate(loc.tolist()):
|
||||
out[dst * mult + offset] = fp8[i]
|
||||
return out
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Existing sanity test
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("input_sl,head,dim,out_sl", [(4, 8, 128, 16)])
|
||||
def test_downcast_fp8(input_sl, head, dim, out_sl, dtype):
|
||||
k = torch.randn(input_sl, head, dim, dtype=dtype, device="cuda")
|
||||
v = torch.randn(input_sl, head, dim, dtype=dtype, device="cuda")
|
||||
k_scale = torch.tensor([1.0], dtype=torch.float32, device="cuda")
|
||||
v_scale = torch.tensor([1.0], dtype=torch.float32, device="cuda")
|
||||
loc = torch.arange(input_sl, dtype=torch.int64, device="cuda")
|
||||
|
||||
k_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda")
|
||||
v_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda")
|
||||
downcast_fp8(k, v, k_out, v_out, k_scale, v_scale, loc)
|
||||
|
||||
# Verify written slots are non-zero (fp8 of random non-zero values)
|
||||
assert k_out[:input_sl].any(), "k_out should have non-zero fp8 values"
|
||||
assert v_out[:input_sl].any(), "v_out should have non-zero fp8 values"
|
||||
# Verify unwritten slots remain zero
|
||||
assert not k_out[input_sl:].any(), "k_out slots beyond input_sl should be zero"
|
||||
assert not v_out[input_sl:].any(), "v_out slots beyond input_sl should be zero"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Numerical correctness: kernel output must match PyTorch fp8 reference.
|
||||
# This verifies that cast<T>(float) and cast<fp8_e4m3_t>(T) produce the
|
||||
# same bit patterns as the removed ConvertFromFloat / ConvertToFP8 structs.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("input_sl,head,dim,out_sl", [(4, 8, 128, 16), (1, 4, 64, 8)])
|
||||
def test_downcast_fp8_matches_reference(input_sl, head, dim, out_sl, dtype):
|
||||
torch.manual_seed(42)
|
||||
k = torch.randn(input_sl, head, dim, dtype=dtype, device="cuda")
|
||||
v = torch.randn(input_sl, head, dim, dtype=dtype, device="cuda")
|
||||
k_scale = torch.tensor([1.0], dtype=torch.float32, device="cuda")
|
||||
v_scale = torch.tensor([1.0], dtype=torch.float32, device="cuda")
|
||||
loc = torch.arange(input_sl, dtype=torch.int64, device="cuda")
|
||||
|
||||
k_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda")
|
||||
v_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda")
|
||||
downcast_fp8(k, v, k_out, v_out, k_scale, v_scale, loc)
|
||||
|
||||
k_ref = _ref_downcast(k, k_scale, loc, out_sl)
|
||||
v_ref = _ref_downcast(v, v_scale, loc, out_sl)
|
||||
|
||||
torch.testing.assert_close(k_out, k_ref, msg="k: kernel vs reference mismatch")
|
||||
torch.testing.assert_close(v_out, v_ref, msg="v: kernel vs reference mismatch")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Scale: a non-unit scale divides the values before fp8 conversion.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("scale_val", [0.5, 2.0, 0.1])
|
||||
def test_downcast_fp8_scale(scale_val, dtype):
|
||||
torch.manual_seed(0)
|
||||
input_sl, head, dim, out_sl = 4, 4, 64, 8
|
||||
|
||||
k = torch.randn(input_sl, head, dim, dtype=dtype, device="cuda")
|
||||
v = torch.randn(input_sl, head, dim, dtype=dtype, device="cuda")
|
||||
k_scale = torch.tensor([scale_val], dtype=torch.float32, device="cuda")
|
||||
v_scale = torch.tensor([scale_val], dtype=torch.float32, device="cuda")
|
||||
loc = torch.arange(input_sl, dtype=torch.int64, device="cuda")
|
||||
|
||||
k_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda")
|
||||
v_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda")
|
||||
downcast_fp8(k, v, k_out, v_out, k_scale, v_scale, loc)
|
||||
|
||||
k_ref = _ref_downcast(k, k_scale, loc, out_sl)
|
||||
v_ref = _ref_downcast(v, v_scale, loc, out_sl)
|
||||
|
||||
torch.testing.assert_close(
|
||||
k_out, k_ref, msg=f"scale={scale_val}: kernel vs reference mismatch"
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
v_out, v_ref, msg=f"scale={scale_val}: kernel vs reference mismatch"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Clamping: values exceeding ±448 must be saturated to fp8 max/min.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
def test_downcast_fp8_clamp(dtype):
|
||||
input_sl, head, dim, out_sl = 2, 1, 8, 4
|
||||
|
||||
# All values well outside fp8 range so clamping is unavoidable.
|
||||
k = torch.full((input_sl, head, dim), 1000.0, dtype=dtype, device="cuda")
|
||||
v = torch.full((input_sl, head, dim), -1000.0, dtype=dtype, device="cuda")
|
||||
scale = torch.tensor([1.0], dtype=torch.float32, device="cuda")
|
||||
loc = torch.arange(input_sl, dtype=torch.int64, device="cuda")
|
||||
|
||||
k_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda")
|
||||
v_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda")
|
||||
downcast_fp8(k, v, k_out, v_out, scale, scale, loc)
|
||||
|
||||
# Reference fp8 max/min byte values (E4M3: 0x7e = 448.0, 0xfe = -448.0)
|
||||
fp8_pos_max = (
|
||||
torch.tensor([_FP8_E4M3_MAX], dtype=dtype, device="cuda")
|
||||
.to(torch.float8_e4m3fn)
|
||||
.view(torch.uint8)
|
||||
.item()
|
||||
)
|
||||
fp8_neg_max = (
|
||||
torch.tensor([-_FP8_E4M3_MAX], dtype=dtype, device="cuda")
|
||||
.to(torch.float8_e4m3fn)
|
||||
.view(torch.uint8)
|
||||
.item()
|
||||
)
|
||||
|
||||
assert (
|
||||
k_out[:input_sl] == fp8_pos_max
|
||||
).all(), "large positive values should clamp to fp8 max"
|
||||
assert (
|
||||
v_out[:input_sl] == fp8_neg_max
|
||||
).all(), "large negative values should clamp to fp8 min"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Scatter: loc controls which output rows receive the converted values.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
def test_downcast_fp8_loc(dtype):
|
||||
torch.manual_seed(7)
|
||||
input_sl, head, dim, out_sl = 3, 2, 32, 10
|
||||
|
||||
k = torch.randn(input_sl, head, dim, dtype=dtype, device="cuda")
|
||||
v = torch.randn(input_sl, head, dim, dtype=dtype, device="cuda")
|
||||
scale = torch.tensor([1.0], dtype=torch.float32, device="cuda")
|
||||
|
||||
# Write to non-contiguous output positions: 0, 5, 9
|
||||
loc = torch.tensor([0, 5, 9], dtype=torch.int64, device="cuda")
|
||||
|
||||
k_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda")
|
||||
v_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda")
|
||||
downcast_fp8(k, v, k_out, v_out, scale, scale, loc)
|
||||
|
||||
k_ref = _ref_downcast(k, scale, loc, out_sl)
|
||||
v_ref = _ref_downcast(v, scale, loc, out_sl)
|
||||
|
||||
torch.testing.assert_close(
|
||||
k_out, k_ref, msg="loc scatter: kernel vs reference mismatch"
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
v_out, v_ref, msg="loc scatter: kernel vs reference mismatch"
|
||||
)
|
||||
|
||||
# Slots not in loc must remain zero
|
||||
written = {0, 5, 9}
|
||||
for i in range(out_sl):
|
||||
if i not in written:
|
||||
assert not k_out[i].any(), f"k_out[{i}] should be zero (not a loc target)"
|
||||
assert not v_out[i].any(), f"v_out[{i}] should be zero (not a loc target)"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# mult/offset: output index = loc[i] * mult + offset
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("mult,offset", [(2, 0), (1, 3), (2, 1)])
|
||||
def test_downcast_fp8_mult_offset(mult, offset, dtype):
|
||||
torch.manual_seed(3)
|
||||
input_sl, head, dim = 2, 2, 32
|
||||
out_sl = input_sl * mult + offset + 4 # ensure output is large enough
|
||||
|
||||
k = torch.randn(input_sl, head, dim, dtype=dtype, device="cuda")
|
||||
v = torch.randn(input_sl, head, dim, dtype=dtype, device="cuda")
|
||||
scale = torch.tensor([1.0], dtype=torch.float32, device="cuda")
|
||||
loc = torch.arange(input_sl, dtype=torch.int64, device="cuda")
|
||||
|
||||
k_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda")
|
||||
v_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda")
|
||||
downcast_fp8(k, v, k_out, v_out, scale, scale, loc, mult=mult, offset=offset)
|
||||
|
||||
k_ref = _ref_downcast(k, scale, loc, out_sl, mult=mult, offset=offset)
|
||||
v_ref = _ref_downcast(v, scale, loc, out_sl, mult=mult, offset=offset)
|
||||
|
||||
torch.testing.assert_close(
|
||||
k_out, k_ref, msg=f"mult={mult},offset={offset}: kernel vs reference mismatch"
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
v_out, v_ref, msg=f"mult={mult},offset={offset}: kernel vs reference mismatch"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# static_cast conversion: verify static_cast<fp8_e4m3_t> matches PyTorch fp8
|
||||
# for a comprehensive sweep including values near and at the fp8 boundary.
|
||||
# This specifically validates that the static_cast fallback (used after
|
||||
# removing explicit __nv_cvt_*raw_to_fp8 from dtype_trait) produces the
|
||||
# same bit patterns as the reference path.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
def test_downcast_fp8_static_cast_boundary(dtype):
|
||||
"""Test conversion accuracy near ±448 fp8 boundary using static_cast path."""
|
||||
torch.manual_seed(0)
|
||||
# Values specifically chosen to stress the static_cast conversion path:
|
||||
# - exactly at ±448 (representable fp8 max)
|
||||
# - just inside the range
|
||||
# - just outside (must saturate)
|
||||
# - zero, small, and mid-range values
|
||||
boundary_vals = [
|
||||
0.0,
|
||||
1.0,
|
||||
-1.0,
|
||||
100.0,
|
||||
-100.0,
|
||||
447.0,
|
||||
-447.0,
|
||||
448.0,
|
||||
-448.0,
|
||||
449.0,
|
||||
-449.0,
|
||||
1000.0,
|
||||
-1000.0,
|
||||
]
|
||||
input_sl = len(boundary_vals)
|
||||
head, dim, out_sl = 1, 8, input_sl
|
||||
|
||||
base = torch.tensor(boundary_vals, dtype=dtype, device="cuda")
|
||||
k = base.unsqueeze(1).unsqueeze(2).expand(input_sl, head, dim).contiguous()
|
||||
v = (-base).unsqueeze(1).unsqueeze(2).expand(input_sl, head, dim).contiguous()
|
||||
scale = torch.tensor([1.0], dtype=torch.float32, device="cuda")
|
||||
loc = torch.arange(input_sl, dtype=torch.int64, device="cuda")
|
||||
|
||||
k_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda")
|
||||
v_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device="cuda")
|
||||
downcast_fp8(k, v, k_out, v_out, scale, scale, loc)
|
||||
|
||||
k_ref = _ref_downcast(k, scale, loc, out_sl)
|
||||
v_ref = _ref_downcast(v, scale, loc, out_sl)
|
||||
|
||||
torch.testing.assert_close(
|
||||
k_out, k_ref, msg="boundary values: k static_cast vs reference mismatch"
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
v_out, v_ref, msg="boundary values: v static_cast vs reference mismatch"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,448 +0,0 @@
|
||||
"""
|
||||
Correctness tests for the fused_qknorm_rope JIT kernel.
|
||||
|
||||
Validates fused_qk_norm_rope against a pure-PyTorch reference and (when
|
||||
available) the sgl_kernel AOT implementation.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.fused_qknorm_rope import fused_qk_norm_rope
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=35, suite="stage-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=256, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
try:
|
||||
from sgl_kernel import fused_qk_norm_rope as fused_qk_norm_rope_aot
|
||||
|
||||
AOT_AVAILABLE = True
|
||||
except ImportError:
|
||||
AOT_AVAILABLE = False
|
||||
|
||||
HEAD_DIMS = [64, 128, 256]
|
||||
NUM_TOKENS = [1, 16, 128]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pure-PyTorch reference
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _compute_inv_freq_yarn(base, rotary_dim, factor, low, high, device):
|
||||
"""Compute YaRN-adjusted inverse frequencies for rotary_dim//2 positions."""
|
||||
half_dims = torch.arange(rotary_dim // 2, dtype=torch.float32, device=device)
|
||||
inv_freq = base ** (-2.0 * half_dims / rotary_dim)
|
||||
|
||||
if factor != 1.0:
|
||||
inv_freq_interp = inv_freq / factor
|
||||
inv_freq_extrap = inv_freq
|
||||
high_adj = high if abs(high - low) > 1e-6 else high + 0.001
|
||||
linear = (half_dims - low) / (high_adj - low)
|
||||
ramp = linear.clamp(0.0, 1.0)
|
||||
extrap_factor = 1.0 - ramp
|
||||
inv_freq = (
|
||||
inv_freq_interp * (1 - extrap_factor) + inv_freq_extrap * extrap_factor
|
||||
)
|
||||
|
||||
return inv_freq
|
||||
|
||||
|
||||
def fused_qk_norm_rope_ref(
|
||||
qkv,
|
||||
num_heads_q,
|
||||
num_heads_k,
|
||||
num_heads_v,
|
||||
head_dim,
|
||||
eps,
|
||||
q_weight,
|
||||
k_weight,
|
||||
base,
|
||||
is_neox,
|
||||
position_ids,
|
||||
factor,
|
||||
low,
|
||||
high,
|
||||
attention_factor,
|
||||
rotary_dim,
|
||||
):
|
||||
"""
|
||||
Pure-PyTorch reference: RMSNorm per head, then RoPE on Q and K.
|
||||
|
||||
Returns a new tensor (same shape as qkv) with the transformation applied.
|
||||
"""
|
||||
num_tokens = qkv.shape[0]
|
||||
total_heads = num_heads_q + num_heads_k + num_heads_v
|
||||
|
||||
qkv_f = qkv.float()
|
||||
qw = q_weight.float()
|
||||
kw = k_weight.float()
|
||||
|
||||
# Reshape to [num_tokens, total_heads, head_dim]
|
||||
qkv_3d = qkv_f.view(num_tokens, total_heads, head_dim)
|
||||
q = qkv_3d[:, :num_heads_q].clone() # [num_tokens, nq, head_dim]
|
||||
k = qkv_3d[:, num_heads_q : num_heads_q + num_heads_k].clone()
|
||||
|
||||
# RMSNorm per head
|
||||
def rms_norm_heads(x, w):
|
||||
# x: [num_tokens, n_heads, head_dim], w: [head_dim]
|
||||
rms = (x**2).mean(-1, keepdim=True)
|
||||
return x * torch.rsqrt(rms + eps) * w
|
||||
|
||||
q = rms_norm_heads(q, qw)
|
||||
k = rms_norm_heads(k, kw)
|
||||
|
||||
# Compute frequencies
|
||||
inv_freq = _compute_inv_freq_yarn(base, rotary_dim, factor, low, high, qkv.device)
|
||||
# theta: [num_tokens, rotary_dim//2]
|
||||
theta = position_ids.float().unsqueeze(1) * inv_freq.unsqueeze(0)
|
||||
cos = torch.cos(theta) # [num_tokens, rotary_dim//2]
|
||||
sin = torch.sin(theta)
|
||||
# Broadcast across heads: [num_tokens, 1, rotary_dim//2]
|
||||
c = cos.unsqueeze(1)
|
||||
s = sin.unsqueeze(1)
|
||||
|
||||
if not is_neox:
|
||||
# Interleave (GPT-J) style: rotate pairs (x[2i], x[2i+1])
|
||||
def apply_interleave(x):
|
||||
# x: [num_tokens, n_heads, head_dim]
|
||||
x_rot = x[:, :, :rotary_dim] # [num_tokens, n_heads, rotary_dim]
|
||||
x_pairs = x_rot.view(num_tokens, -1, rotary_dim // 2, 2)
|
||||
x0, x1 = x_pairs[..., 0], x_pairs[..., 1]
|
||||
x0_new = x0 * c - x1 * s
|
||||
x1_new = x1 * c + x0 * s
|
||||
x_rot_new = torch.stack([x0_new, x1_new], dim=-1).view(
|
||||
num_tokens, -1, rotary_dim
|
||||
)
|
||||
result = x.clone()
|
||||
result[:, :, :rotary_dim] = x_rot_new * attention_factor
|
||||
return result
|
||||
|
||||
q = apply_interleave(q)
|
||||
k = apply_interleave(k)
|
||||
else:
|
||||
# NeoX style: first half × cos − second half × sin (and vice versa)
|
||||
def apply_neox(x):
|
||||
# x: [num_tokens, n_heads, head_dim]
|
||||
x1 = x[:, :, : rotary_dim // 2]
|
||||
x2 = x[:, :, rotary_dim // 2 : rotary_dim]
|
||||
x1_new = x1 * c - x2 * s
|
||||
x2_new = x2 * c + x1 * s
|
||||
result = x.clone()
|
||||
result[:, :, : rotary_dim // 2] = x1_new * attention_factor
|
||||
result[:, :, rotary_dim // 2 : rotary_dim] = x2_new * attention_factor
|
||||
return result
|
||||
|
||||
q = apply_neox(q)
|
||||
k = apply_neox(k)
|
||||
|
||||
# Write back into a copy of the full QKV
|
||||
result_3d = qkv_f.view(num_tokens, total_heads, head_dim).clone()
|
||||
result_3d[:, :num_heads_q] = q
|
||||
result_3d[:, num_heads_q : num_heads_q + num_heads_k] = k
|
||||
return result_3d.view(num_tokens, -1).bfloat16()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: correctness vs PyTorch reference
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("head_dim", HEAD_DIMS)
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("is_neox", [False, True])
|
||||
def test_fused_qknorm_rope_vs_ref(head_dim, num_tokens, is_neox):
|
||||
torch.manual_seed(head_dim * num_tokens + int(is_neox))
|
||||
device = "cuda"
|
||||
num_heads_q, num_heads_k, num_heads_v = 4, 2, 2
|
||||
total_heads = num_heads_q + num_heads_k + num_heads_v
|
||||
rotary_dim = head_dim # full rotary
|
||||
|
||||
qkv = torch.randn(
|
||||
(num_tokens, total_heads * head_dim), dtype=torch.bfloat16, device=device
|
||||
)
|
||||
q_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device)
|
||||
k_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device)
|
||||
position_ids = torch.arange(num_tokens, dtype=torch.int32, device=device)
|
||||
|
||||
eps = 1e-5
|
||||
base = 10000.0
|
||||
factor = 1.0 # no YaRN
|
||||
low, high = 1.0, 32.0
|
||||
attention_factor = 1.0
|
||||
|
||||
ref = fused_qk_norm_rope_ref(
|
||||
qkv,
|
||||
num_heads_q,
|
||||
num_heads_k,
|
||||
num_heads_v,
|
||||
head_dim,
|
||||
eps,
|
||||
q_weight,
|
||||
k_weight,
|
||||
base,
|
||||
is_neox,
|
||||
position_ids,
|
||||
factor,
|
||||
low,
|
||||
high,
|
||||
attention_factor,
|
||||
rotary_dim,
|
||||
)
|
||||
|
||||
qkv_jit = qkv.clone()
|
||||
fused_qk_norm_rope(
|
||||
qkv_jit,
|
||||
num_heads_q,
|
||||
num_heads_k,
|
||||
num_heads_v,
|
||||
head_dim,
|
||||
eps,
|
||||
q_weight,
|
||||
k_weight,
|
||||
base,
|
||||
is_neox,
|
||||
position_ids,
|
||||
factor,
|
||||
low,
|
||||
high,
|
||||
attention_factor,
|
||||
rotary_dim,
|
||||
)
|
||||
|
||||
assert torch.allclose(qkv_jit.float(), ref.float(), atol=5e-3, rtol=1e-2), (
|
||||
f"mismatch: head_dim={head_dim}, num_tokens={num_tokens}, "
|
||||
f"is_neox={is_neox}, "
|
||||
f"max_err={( qkv_jit.float() - ref.float()).abs().max().item():.4e}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("head_dim", HEAD_DIMS)
|
||||
@pytest.mark.parametrize("is_neox", [False, True])
|
||||
def test_fused_qknorm_rope_partial_rotary(head_dim, is_neox):
|
||||
"""Test with rotary_dim < head_dim: non-rotary elements should be RMSNorm-only."""
|
||||
torch.manual_seed(42 + head_dim + int(is_neox))
|
||||
device = "cuda"
|
||||
num_tokens = 16
|
||||
num_heads_q, num_heads_k, num_heads_v = 2, 2, 2
|
||||
total_heads = num_heads_q + num_heads_k + num_heads_v
|
||||
rotary_dim = head_dim // 2 # half of head_dim
|
||||
|
||||
# NeoX requires half_rotary_lanes to be power of 2.
|
||||
# half_rotary_lanes = rotary_dim / (head_dim / 32) / 2 = (head_dim//2) / (head_dim/32) / 2
|
||||
# = 16 / 2 = 8 → power of 2, OK for all supported head_dims.
|
||||
|
||||
qkv = torch.randn(
|
||||
(num_tokens, total_heads * head_dim), dtype=torch.bfloat16, device=device
|
||||
)
|
||||
q_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device)
|
||||
k_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device)
|
||||
position_ids = torch.arange(num_tokens, dtype=torch.int32, device=device)
|
||||
|
||||
ref = fused_qk_norm_rope_ref(
|
||||
qkv,
|
||||
num_heads_q,
|
||||
num_heads_k,
|
||||
num_heads_v,
|
||||
head_dim,
|
||||
1e-5,
|
||||
q_weight,
|
||||
k_weight,
|
||||
10000.0,
|
||||
is_neox,
|
||||
position_ids,
|
||||
1.0,
|
||||
1.0,
|
||||
32.0,
|
||||
1.0,
|
||||
rotary_dim,
|
||||
)
|
||||
|
||||
qkv_jit = qkv.clone()
|
||||
fused_qk_norm_rope(
|
||||
qkv_jit,
|
||||
num_heads_q,
|
||||
num_heads_k,
|
||||
num_heads_v,
|
||||
head_dim,
|
||||
1e-5,
|
||||
q_weight,
|
||||
k_weight,
|
||||
10000.0,
|
||||
is_neox,
|
||||
position_ids,
|
||||
1.0,
|
||||
1.0,
|
||||
32.0,
|
||||
1.0,
|
||||
rotary_dim,
|
||||
)
|
||||
|
||||
assert torch.allclose(qkv_jit.float(), ref.float(), atol=5e-3, rtol=1e-2), (
|
||||
f"partial rotary mismatch: head_dim={head_dim}, is_neox={is_neox}, "
|
||||
f"max_err={(qkv_jit.float() - ref.float()).abs().max().item():.4e}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("head_dim", HEAD_DIMS)
|
||||
def test_fused_qknorm_rope_yarn_scaling(head_dim):
|
||||
"""Test with YaRN scaling (factor != 1.0)."""
|
||||
torch.manual_seed(99 + head_dim)
|
||||
device = "cuda"
|
||||
num_tokens = 32
|
||||
num_heads_q, num_heads_k, num_heads_v = 2, 2, 2
|
||||
total_heads = num_heads_q + num_heads_k + num_heads_v
|
||||
rotary_dim = head_dim
|
||||
|
||||
qkv = torch.randn(
|
||||
(num_tokens, total_heads * head_dim), dtype=torch.bfloat16, device=device
|
||||
)
|
||||
q_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device)
|
||||
k_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device)
|
||||
position_ids = torch.arange(num_tokens, dtype=torch.int32, device=device)
|
||||
|
||||
factor = 2.5
|
||||
low, high = 4.0, 32.0
|
||||
attention_factor = 0.9
|
||||
is_neox = False # test with interleave; NeoX also tested in other tests
|
||||
|
||||
ref = fused_qk_norm_rope_ref(
|
||||
qkv,
|
||||
num_heads_q,
|
||||
num_heads_k,
|
||||
num_heads_v,
|
||||
head_dim,
|
||||
1e-5,
|
||||
q_weight,
|
||||
k_weight,
|
||||
500000.0,
|
||||
is_neox,
|
||||
position_ids,
|
||||
factor,
|
||||
low,
|
||||
high,
|
||||
attention_factor,
|
||||
rotary_dim,
|
||||
)
|
||||
|
||||
qkv_jit = qkv.clone()
|
||||
fused_qk_norm_rope(
|
||||
qkv_jit,
|
||||
num_heads_q,
|
||||
num_heads_k,
|
||||
num_heads_v,
|
||||
head_dim,
|
||||
1e-5,
|
||||
q_weight,
|
||||
k_weight,
|
||||
500000.0,
|
||||
is_neox,
|
||||
position_ids,
|
||||
factor,
|
||||
low,
|
||||
high,
|
||||
attention_factor,
|
||||
rotary_dim,
|
||||
)
|
||||
|
||||
assert torch.allclose(qkv_jit.float(), ref.float(), atol=5e-3, rtol=1e-2), (
|
||||
f"YaRN mismatch: head_dim={head_dim}, "
|
||||
f"max_err={(qkv_jit.float() - ref.float()).abs().max().item():.4e}"
|
||||
)
|
||||
|
||||
|
||||
def test_fused_qknorm_rope_default_rotary_dim():
|
||||
"""rotary_dim=None should default to head_dim."""
|
||||
device = "cuda"
|
||||
num_tokens = 8
|
||||
num_heads_q, num_heads_k, num_heads_v = 2, 2, 2
|
||||
head_dim = 128
|
||||
total_heads = num_heads_q + num_heads_k + num_heads_v
|
||||
|
||||
torch.manual_seed(0)
|
||||
qkv1 = torch.randn(
|
||||
(num_tokens, total_heads * head_dim), dtype=torch.bfloat16, device=device
|
||||
)
|
||||
qkv2 = qkv1.clone()
|
||||
q_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device)
|
||||
k_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device)
|
||||
position_ids = torch.zeros(num_tokens, dtype=torch.int32, device=device)
|
||||
|
||||
common_kwargs = dict(
|
||||
num_heads_q=num_heads_q,
|
||||
num_heads_k=num_heads_k,
|
||||
num_heads_v=num_heads_v,
|
||||
head_dim=head_dim,
|
||||
eps=1e-5,
|
||||
q_weight=q_weight,
|
||||
k_weight=k_weight,
|
||||
base=10000.0,
|
||||
is_neox=False,
|
||||
position_ids=position_ids,
|
||||
factor=1.0,
|
||||
low=1.0,
|
||||
high=32.0,
|
||||
attention_factor=1.0,
|
||||
)
|
||||
|
||||
fused_qk_norm_rope(qkv1, **common_kwargs, rotary_dim=None)
|
||||
fused_qk_norm_rope(qkv2, **common_kwargs, rotary_dim=head_dim)
|
||||
|
||||
assert torch.equal(qkv1, qkv2), "rotary_dim=None must equal rotary_dim=head_dim"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cross-validation against AOT sgl_kernel
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.skipif(not AOT_AVAILABLE, reason="sgl_kernel not available")
|
||||
@pytest.mark.parametrize("head_dim", [64, 128, 256])
|
||||
@pytest.mark.parametrize("is_neox", [False, True])
|
||||
def test_fused_qknorm_rope_vs_aot(head_dim, is_neox):
|
||||
torch.manual_seed(head_dim * 7 + int(is_neox))
|
||||
device = "cuda"
|
||||
num_tokens = 32
|
||||
num_heads_q, num_heads_k, num_heads_v = 4, 2, 2
|
||||
total_heads = num_heads_q + num_heads_k + num_heads_v
|
||||
|
||||
qkv = torch.randn(
|
||||
(num_tokens, total_heads * head_dim), dtype=torch.bfloat16, device=device
|
||||
)
|
||||
q_weight = torch.randn(head_dim, dtype=torch.bfloat16, device=device).abs() + 0.5
|
||||
k_weight = torch.randn(head_dim, dtype=torch.bfloat16, device=device).abs() + 0.5
|
||||
position_ids = torch.arange(num_tokens, dtype=torch.int32, device=device)
|
||||
|
||||
common = dict(
|
||||
num_heads_q=num_heads_q,
|
||||
num_heads_k=num_heads_k,
|
||||
num_heads_v=num_heads_v,
|
||||
head_dim=head_dim,
|
||||
eps=1e-5,
|
||||
q_weight=q_weight,
|
||||
k_weight=k_weight,
|
||||
base=10000.0,
|
||||
is_neox=is_neox,
|
||||
position_ids=position_ids,
|
||||
factor=1.0,
|
||||
low=1.0,
|
||||
high=32.0,
|
||||
attention_factor=1.0,
|
||||
rotary_dim=head_dim,
|
||||
)
|
||||
|
||||
qkv_jit = qkv.clone()
|
||||
fused_qk_norm_rope(qkv_jit, **common)
|
||||
|
||||
qkv_aot = qkv.clone()
|
||||
fused_qk_norm_rope_aot(qkv_aot, **common)
|
||||
|
||||
assert torch.allclose(qkv_jit.float(), qkv_aot.float(), atol=1e-2, rtol=1e-2), (
|
||||
f"JIT vs AOT mismatch: head_dim={head_dim}, is_neox={is_neox}, "
|
||||
f"max_err={(qkv_jit.float() - qkv_aot.float()).abs().max().item():.4e}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
@@ -1,155 +0,0 @@
|
||||
"""Integration test for OpenTelemetry tracing in the diffusion pipeline.
|
||||
|
||||
Spins up a lightweight in-process OTLP collector, launches a diffusion server
|
||||
with ``--enable-trace``, sends an image-generation request with a
|
||||
``traceparent`` header, and asserts that the expected spans
|
||||
(``scheduler_dispatch``, ``gpu_forward``) are exported.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
# Configure OTLP exporter for faster test execution.
|
||||
# Must be set before importing any sglang trace module.
|
||||
os.environ.setdefault("SGLANG_OTLP_EXPORTER_SCHEDULE_DELAY_MILLIS", "50")
|
||||
os.environ.setdefault("SGLANG_OTLP_EXPORTER_MAX_EXPORT_BATCH_SIZE", "4")
|
||||
|
||||
import logging
|
||||
import time
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from sglang.multimodal_gen.test.server.test_server_utils import ServerManager
|
||||
from sglang.multimodal_gen.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
from sglang.test.otel_collector import LightweightOtlpCollector
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Expected diffusion trace span names (from DiffStage in trace_wrapper.py)
|
||||
EXPECTED_DIFF_SPANS = ["scheduler_dispatch", "gpu_forward"]
|
||||
|
||||
COLLECTOR_PORT = 4317
|
||||
SERVER_PORT = 39812
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def tracing_env():
|
||||
"""Start the OTLP collector and diffusion server once for all tests."""
|
||||
collector = LightweightOtlpCollector(port=COLLECTOR_PORT)
|
||||
collector.start()
|
||||
time.sleep(0.3)
|
||||
|
||||
mgr = ServerManager(
|
||||
model=DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
port=SERVER_PORT,
|
||||
extra_args=f"--enable-trace --otlp-traces-endpoint 127.0.0.1:{COLLECTOR_PORT}",
|
||||
)
|
||||
ctx = mgr.start()
|
||||
|
||||
# Clear any warmup spans
|
||||
time.sleep(2)
|
||||
collector.clear()
|
||||
|
||||
yield collector, ctx
|
||||
|
||||
ctx.cleanup()
|
||||
collector.stop()
|
||||
|
||||
|
||||
def _generate_image(headers=None):
|
||||
"""Send a single image-generation request."""
|
||||
resp = requests.post(
|
||||
f"http://127.0.0.1:{SERVER_PORT}/v1/images/generations",
|
||||
json={
|
||||
"model": DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
"prompt": "A white cat",
|
||||
"size": "256x256",
|
||||
"n": 1,
|
||||
},
|
||||
headers=headers or {},
|
||||
timeout=300,
|
||||
)
|
||||
assert resp.status_code == 200, f"Generation failed: {resp.text}"
|
||||
return resp
|
||||
|
||||
|
||||
def _wait_for_spans(collector, required_names=None, min_count=1, timeout=30):
|
||||
"""Wait until collector has the required span names (or at least ``min_count`` spans)."""
|
||||
deadline = time.time() + timeout
|
||||
while time.time() < deadline:
|
||||
if required_names:
|
||||
if all(collector.has_span(n) for n in required_names):
|
||||
return
|
||||
elif collector.count_spans() >= min_count:
|
||||
return
|
||||
time.sleep(0.5)
|
||||
|
||||
|
||||
def test_spans_exported(tracing_env):
|
||||
"""After a generation request the expected diffusion spans appear."""
|
||||
collector, _ = tracing_env
|
||||
collector.clear()
|
||||
|
||||
# W3C Trace Context traceparent header
|
||||
_generate_image(
|
||||
headers={
|
||||
"traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"
|
||||
}
|
||||
)
|
||||
_wait_for_spans(collector, required_names=EXPECTED_DIFF_SPANS)
|
||||
|
||||
span_names = collector.get_span_names()
|
||||
for expected in EXPECTED_DIFF_SPANS:
|
||||
assert (
|
||||
expected in span_names
|
||||
), f"Missing span '{expected}'. Collected: {sorted(span_names)}"
|
||||
|
||||
|
||||
def test_spans_without_traceparent(tracing_env):
|
||||
"""Requests without a traceparent header still produce spans as a new
|
||||
root trace (not linked to any upstream)."""
|
||||
collector, _ = tracing_env
|
||||
collector.clear()
|
||||
|
||||
_generate_image()
|
||||
_wait_for_spans(collector, required_names=EXPECTED_DIFF_SPANS)
|
||||
|
||||
span_names = collector.get_span_names()
|
||||
for expected in EXPECTED_DIFF_SPANS:
|
||||
assert (
|
||||
expected in span_names
|
||||
), f"Missing span '{expected}'. Collected: {sorted(span_names)}"
|
||||
|
||||
|
||||
def test_batch_requests(tracing_env):
|
||||
"""Multiple requests each produce their own set of spans."""
|
||||
collector, _ = tracing_env
|
||||
collector.clear()
|
||||
|
||||
batch_size = 3
|
||||
for i in range(batch_size):
|
||||
# Each request gets a unique trace-id
|
||||
trace_id = f"0af7651916cd43dd8448eb211c8031{i:02x}"
|
||||
_generate_image(headers={"traceparent": f"00-{trace_id}-b7ad6b7169203331-01"})
|
||||
|
||||
# Wait until all scheduler_dispatch spans have arrived (they come from a
|
||||
# separate process so may lag behind gpu_forward).
|
||||
deadline = time.time() + 60
|
||||
while time.time() < deadline:
|
||||
if all(
|
||||
len(collector.get_spans_by_name(n)) >= batch_size
|
||||
for n in EXPECTED_DIFF_SPANS
|
||||
):
|
||||
break
|
||||
time.sleep(0.5)
|
||||
|
||||
for span_name in EXPECTED_DIFF_SPANS:
|
||||
matching = collector.get_spans_by_name(span_name)
|
||||
assert len(matching) >= batch_size, (
|
||||
f"Expected at least {batch_size} '{span_name}' spans, "
|
||||
f"got {len(matching)}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v", "-s"])
|
||||
@@ -1,292 +0,0 @@
|
||||
"""
|
||||
Test that the optimized PatchEmbed (reshape + F.linear) is equivalent
|
||||
to the original Conv3d-based PatchEmbed from upstream/main.
|
||||
|
||||
The opt_krea branch replaces Conv3d forward with manual
|
||||
reshape + permute + F.linear for 5D input. This is valid because
|
||||
Conv3d with stride==kernel_size is a non-overlapping patch extraction
|
||||
followed by linear projection, which is exactly what the manual path does.
|
||||
|
||||
We disable TF32 so cuDNN (Conv3d) and cuBLAS (F.linear) both use
|
||||
full FP32 precision, enabling strict numerical comparison.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class PatchEmbed3D(nn.Module):
|
||||
"""PatchEmbed from upstream/main: uses Conv3d directly."""
|
||||
|
||||
def __init__(
|
||||
self, patch_size, in_chans, embed_dim, flatten=True, bias=True, dtype=None
|
||||
):
|
||||
super().__init__()
|
||||
if isinstance(patch_size, list | tuple):
|
||||
if len(patch_size) == 1:
|
||||
patch_size = (patch_size[0], patch_size[0])
|
||||
else:
|
||||
patch_size = (patch_size, patch_size)
|
||||
self.patch_size = patch_size
|
||||
self.flatten = flatten
|
||||
self.proj = nn.Conv3d(
|
||||
in_chans,
|
||||
embed_dim,
|
||||
kernel_size=patch_size,
|
||||
stride=patch_size,
|
||||
bias=bias,
|
||||
dtype=dtype,
|
||||
)
|
||||
self.norm = nn.Identity()
|
||||
|
||||
def forward(self, x):
|
||||
x = self.proj(x)
|
||||
if self.flatten:
|
||||
x = x.flatten(2).transpose(1, 2)
|
||||
x = self.norm(x)
|
||||
return x
|
||||
|
||||
|
||||
class PatchEmbed(nn.Module):
|
||||
"""PatchEmbed from opt_krea: replaces Conv3d with reshape + F.linear for 5D input."""
|
||||
|
||||
def __init__(
|
||||
self, patch_size, in_chans, embed_dim, flatten=True, bias=True, dtype=None
|
||||
):
|
||||
super().__init__()
|
||||
if isinstance(patch_size, list | tuple):
|
||||
if len(patch_size) == 1:
|
||||
patch_size = (1, patch_size[0], patch_size[0])
|
||||
elif len(patch_size) == 2:
|
||||
patch_size = (1, patch_size[0], patch_size[1])
|
||||
else:
|
||||
patch_size = (1, patch_size, patch_size)
|
||||
self.patch_size = patch_size
|
||||
self.flatten = flatten
|
||||
self.proj = nn.Conv3d(
|
||||
in_chans,
|
||||
embed_dim,
|
||||
kernel_size=patch_size,
|
||||
stride=patch_size,
|
||||
bias=bias,
|
||||
dtype=dtype,
|
||||
)
|
||||
self.norm = nn.Identity()
|
||||
|
||||
def forward(self, x):
|
||||
if x.dim() == 5:
|
||||
B, C, T, H, W = x.shape
|
||||
pt, ph, pw = self.patch_size
|
||||
T_ = T // pt
|
||||
H_ = H // ph
|
||||
W_ = W // pw
|
||||
x = x.reshape(B, C, T_, pt, H_, ph, W_, pw)
|
||||
x = x.permute(0, 2, 4, 6, 1, 3, 5, 7).contiguous()
|
||||
x = x.reshape(B, T_ * H_ * W_, C * pt * ph * pw)
|
||||
w = self.proj.weight.reshape(self.proj.weight.shape[0], -1)
|
||||
x = F.linear(x, w, self.proj.bias)
|
||||
if not self.flatten:
|
||||
x = x.reshape(B, T_, H_, W_, -1).permute(0, 4, 1, 2, 3).contiguous()
|
||||
x = self.norm(x)
|
||||
return x
|
||||
x = self.proj(x)
|
||||
if self.flatten:
|
||||
x = x.flatten(2).transpose(1, 2)
|
||||
x = self.norm(x)
|
||||
return x
|
||||
|
||||
|
||||
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
|
||||
def _copy_weights(src, dst):
|
||||
dst.proj.weight.data.copy_(src.proj.weight.data)
|
||||
if src.proj.bias is not None:
|
||||
dst.proj.bias.data.copy_(src.proj.bias.data)
|
||||
|
||||
|
||||
def _run_equivalence(
|
||||
patch_size,
|
||||
in_chans,
|
||||
embed_dim,
|
||||
flatten,
|
||||
bias,
|
||||
weight_dtype,
|
||||
input_dtype,
|
||||
B,
|
||||
T,
|
||||
H,
|
||||
W,
|
||||
atol,
|
||||
rtol,
|
||||
):
|
||||
"""Helper: build both models with shared weights, run forward, compare.
|
||||
|
||||
Args:
|
||||
weight_dtype: dtype for Conv3d weights (None = FP32).
|
||||
input_dtype: dtype for the input tensor (None = FP32).
|
||||
"""
|
||||
torch.manual_seed(42)
|
||||
main = (
|
||||
PatchEmbed3D(patch_size, in_chans, embed_dim, flatten, bias, dtype=weight_dtype)
|
||||
.to(DEVICE)
|
||||
.eval()
|
||||
)
|
||||
opt = (
|
||||
PatchEmbed(patch_size, in_chans, embed_dim, flatten, bias, dtype=weight_dtype)
|
||||
.to(DEVICE)
|
||||
.eval()
|
||||
)
|
||||
_copy_weights(main, opt)
|
||||
|
||||
x = torch.randn(
|
||||
B, in_chans, T, H, W, device=DEVICE, dtype=input_dtype or torch.float32
|
||||
)
|
||||
with torch.no_grad():
|
||||
out_main = main(x)
|
||||
out_opt = opt(x)
|
||||
|
||||
assert (
|
||||
out_main.shape == out_opt.shape
|
||||
), f"Shape mismatch: {out_main.shape} vs {out_opt.shape}"
|
||||
assert (
|
||||
out_main.dtype == out_opt.dtype
|
||||
), f"Dtype mismatch: {out_main.dtype} vs {out_opt.dtype}"
|
||||
torch.testing.assert_close(out_main, out_opt, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _disable_tf32():
|
||||
prev_cudnn = torch.backends.cudnn.allow_tf32
|
||||
prev_matmul = torch.backends.cuda.matmul.allow_tf32
|
||||
torch.backends.cudnn.allow_tf32 = False
|
||||
torch.backends.cuda.matmul.allow_tf32 = False
|
||||
yield
|
||||
torch.backends.cudnn.allow_tf32 = prev_cudnn
|
||||
torch.backends.cuda.matmul.allow_tf32 = prev_matmul
|
||||
|
||||
|
||||
# ── Wan2.1 / Wan2.2 / CausalWan / Helios ────────────────────────────────────
|
||||
# patch_size=(1,2,2), in_channels=16, embed_dim=5120, flatten=False
|
||||
# Real usage: weight=FP32 (no dtype passed), input=BF16 from VAE latent
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"dtype,atol,rtol",
|
||||
[
|
||||
(None, 1e-4, 1e-4), # weight=FP32, input=FP32
|
||||
(torch.bfloat16, 1e-2, 1e-2), # weight=BF16, input=BF16
|
||||
(torch.float16, 1e-2, 1e-2), # weight=FP16, input=FP16
|
||||
],
|
||||
ids=["fp32", "bf16", "fp16"],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"B,T,H,W",
|
||||
[
|
||||
(1, 21, 60, 104), # 480p typical
|
||||
(2, 9, 40, 64), # smaller resolution, batch=2
|
||||
(1, 33, 90, 160), # 720p longer video
|
||||
],
|
||||
ids=["480p-B1", "small-B2", "720p-B1"],
|
||||
)
|
||||
def test_wan_helios(dtype, atol, rtol, B, T, H, W):
|
||||
_run_equivalence(
|
||||
patch_size=(1, 2, 2),
|
||||
in_chans=16,
|
||||
embed_dim=5120,
|
||||
flatten=False,
|
||||
bias=True,
|
||||
weight_dtype=dtype,
|
||||
input_dtype=dtype,
|
||||
B=B,
|
||||
T=T,
|
||||
H=H,
|
||||
W=W,
|
||||
atol=atol,
|
||||
rtol=rtol,
|
||||
)
|
||||
|
||||
|
||||
# ── HunyuanVideo ─────────────────────────────────────────────────────────────
|
||||
# patch_size=[1,2,2] (list!), in_channels=16, embed_dim=3072, flatten=True
|
||||
# Real usage: dtype passed to PatchEmbed, so weight & input share same dtype
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"dtype,atol,rtol",
|
||||
[
|
||||
(None, 1e-4, 1e-4), # weight=FP32, input=FP32
|
||||
(torch.bfloat16, 1e-2, 1e-2), # weight=BF16, input=BF16
|
||||
(torch.float16, 1e-2, 1e-2), # weight=FP16, input=FP16
|
||||
],
|
||||
ids=["fp32", "bf16", "fp16"],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"B,T,H,W",
|
||||
[
|
||||
(1, 21, 60, 104),
|
||||
(2, 9, 40, 64),
|
||||
],
|
||||
ids=["480p-B1", "small-B2"],
|
||||
)
|
||||
def test_hunyuanvideo(dtype, atol, rtol, B, T, H, W):
|
||||
_run_equivalence(
|
||||
patch_size=[1, 2, 2],
|
||||
in_chans=16,
|
||||
embed_dim=3072,
|
||||
flatten=True,
|
||||
bias=True,
|
||||
weight_dtype=dtype,
|
||||
input_dtype=dtype,
|
||||
B=B,
|
||||
T=T,
|
||||
H=H,
|
||||
W=W,
|
||||
atol=atol,
|
||||
rtol=rtol,
|
||||
)
|
||||
|
||||
|
||||
# ── No-bias variants ─────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_wan_no_bias():
|
||||
_run_equivalence(
|
||||
patch_size=(1, 2, 2),
|
||||
in_chans=16,
|
||||
embed_dim=5120,
|
||||
flatten=False,
|
||||
bias=False,
|
||||
weight_dtype=None,
|
||||
input_dtype=None,
|
||||
B=1,
|
||||
T=21,
|
||||
H=60,
|
||||
W=104,
|
||||
atol=1e-4,
|
||||
rtol=1e-4,
|
||||
)
|
||||
|
||||
|
||||
def test_hunyuanvideo_no_bias():
|
||||
_run_equivalence(
|
||||
patch_size=[1, 2, 2],
|
||||
in_chans=16,
|
||||
embed_dim=3072,
|
||||
flatten=True,
|
||||
bias=False,
|
||||
weight_dtype=None,
|
||||
input_dtype=None,
|
||||
B=1,
|
||||
T=21,
|
||||
H=60,
|
||||
W=104,
|
||||
atol=1e-4,
|
||||
rtol=1e-4,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v", "--tb=short"])
|
||||
@@ -1,788 +0,0 @@
|
||||
# Copyright 2023-2025 SGLang Team
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# ==============================================================================
|
||||
|
||||
import random
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
|
||||
from sglang.srt.layers.moe.moe_runner.runner import MoeRunner
|
||||
from sglang.srt.layers.moe.moe_runner.triton import (
|
||||
TritonMoeQuantInfo,
|
||||
)
|
||||
from sglang.srt.layers.moe.token_dispatcher.standard import StandardDispatchOutput
|
||||
from sglang.srt.layers.moe.topk import StandardTopKOutput
|
||||
from sglang.srt.layers.moe.utils import MoeRunnerBackend
|
||||
from sglang.srt.lora.lora_moe_runners import LoRAInfo
|
||||
from sglang.srt.utils import set_random_seed
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=17, suite="stage-b-test-1-gpu-large")
|
||||
|
||||
|
||||
def generate_request_data(
|
||||
num_tokens: int, num_sequences: int, max_loras: int, device="cuda"
|
||||
):
|
||||
"""
|
||||
Generates segment-based request data instead of token-based data.
|
||||
"""
|
||||
assert num_sequences > 0 and max_loras > 0
|
||||
assert num_tokens >= num_sequences, "num_tokens must be >= num_sequences"
|
||||
|
||||
# 1. Generate random segment lengths
|
||||
remaining = num_tokens
|
||||
seg_lens = []
|
||||
for _ in range(num_sequences - 1):
|
||||
# Ensure at least 1 token per sequence
|
||||
max_len = remaining - (num_sequences - len(seg_lens)) + 1
|
||||
length = random.randint(1, min(max_len, num_tokens // num_sequences * 2))
|
||||
seg_lens.append(length)
|
||||
remaining -= length
|
||||
seg_lens.append(remaining) # Last segment gets the rest
|
||||
|
||||
# 2. Build seg_indptr [0, len1, len1+len2, ...]
|
||||
seg_indptr = torch.cumsum(
|
||||
torch.tensor([0] + seg_lens, dtype=torch.int32, device=device),
|
||||
dim=0,
|
||||
dtype=torch.int32,
|
||||
)
|
||||
|
||||
# 3. Assign one LoRA ID per Request
|
||||
req_to_lora = torch.randint(
|
||||
0, max_loras, (num_sequences,), dtype=torch.int32, device=device
|
||||
)
|
||||
|
||||
# 4. Create dense mapping for the Naive verification function
|
||||
# (Expand req_to_lora based on seg_lens)
|
||||
token_lora_mapping = torch.repeat_interleave(
|
||||
req_to_lora, torch.tensor(seg_lens, device=device)
|
||||
)
|
||||
|
||||
return seg_indptr, req_to_lora, token_lora_mapping
|
||||
|
||||
|
||||
def assign_experts_to_tokens(
|
||||
num_tokens: int, num_experts: int, top_k_num: int, dtype=torch.float32
|
||||
):
|
||||
assert top_k_num <= num_experts, "top_k_num must be <= num_experts"
|
||||
|
||||
expert_indices = torch.empty((num_tokens, top_k_num), dtype=torch.int32)
|
||||
for i in range(num_tokens):
|
||||
selected = torch.randperm(num_experts)[:top_k_num]
|
||||
expert_indices[i] = selected
|
||||
|
||||
expert_weights = torch.rand((num_tokens, top_k_num), dtype=dtype)
|
||||
expert_weights = expert_weights / expert_weights.sum(dim=1, keepdim=True)
|
||||
|
||||
return expert_indices, expert_weights
|
||||
|
||||
|
||||
def sample_data(
|
||||
num_tokens: int,
|
||||
num_sequences: int,
|
||||
max_loras: int,
|
||||
num_experts: int,
|
||||
top_k_num: int,
|
||||
dtype=torch.float32,
|
||||
device="cuda",
|
||||
):
|
||||
topk_ids, topk_weights = assign_experts_to_tokens(
|
||||
num_tokens, num_experts, top_k_num, dtype
|
||||
)
|
||||
seg_indptr, req_to_lora, token_lora_mapping = generate_request_data(
|
||||
num_tokens, num_sequences, max_loras, device
|
||||
)
|
||||
return topk_ids, topk_weights, seg_indptr, req_to_lora, token_lora_mapping
|
||||
|
||||
|
||||
def create_lora_info(
|
||||
seg_indptr,
|
||||
weight_indices,
|
||||
topk_ids,
|
||||
max_loras,
|
||||
num_experts,
|
||||
max_lora_rank,
|
||||
hidden_dim,
|
||||
intermediate_dim,
|
||||
gate_up_dim,
|
||||
dtype,
|
||||
device,
|
||||
lora_use_virtual_experts=False,
|
||||
):
|
||||
# -------------------------------------------------------------------------
|
||||
# 1. Deterministic LoRA A Initialization
|
||||
# -------------------------------------------------------------------------
|
||||
|
||||
val_gate_up_a = 0.1
|
||||
gate_up_lora_a_weights = torch.full(
|
||||
(max_loras, num_experts, max_lora_rank * 2, hidden_dim),
|
||||
val_gate_up_a,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
|
||||
val_down_a = 1.0 / intermediate_dim
|
||||
down_lora_a_weights = torch.full(
|
||||
(max_loras, num_experts, max_lora_rank, intermediate_dim),
|
||||
val_down_a,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# 2. Deterministic LoRA B Initialization
|
||||
# -------------------------------------------------------------------------
|
||||
base_target = 0.05
|
||||
|
||||
gate_up_lora_b_weights = torch.zeros(
|
||||
(max_loras, num_experts, gate_up_dim, max_lora_rank),
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
down_lora_b_weights = torch.zeros(
|
||||
(max_loras, num_experts, hidden_dim, max_lora_rank), dtype=dtype, device=device
|
||||
)
|
||||
|
||||
for i in range(num_experts):
|
||||
expert_multiplier = i + 1
|
||||
divisor = max(1, max_lora_rank)
|
||||
fill_val = (base_target * expert_multiplier) / divisor
|
||||
|
||||
gate_up_lora_b_weights[:, i, :, :] = fill_val
|
||||
down_lora_b_weights[:, i, :, :] = fill_val
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# 3. Setup Metadata
|
||||
# -------------------------------------------------------------------------
|
||||
lora_ranks = torch.full(
|
||||
(max_loras,), max_lora_rank, dtype=torch.int32, device=device
|
||||
)
|
||||
|
||||
# Enable all adapters referenced in weight_indices
|
||||
adapter_enabled = torch.zeros(max_loras + 1, dtype=torch.int32, device=device)
|
||||
adapter_enabled.index_fill_(0, weight_indices.long(), 1)
|
||||
|
||||
return LoRAInfo(
|
||||
gate_up_lora_a_weights=gate_up_lora_a_weights,
|
||||
gate_up_lora_b_weights=gate_up_lora_b_weights,
|
||||
down_lora_a_weights=down_lora_a_weights,
|
||||
down_lora_b_weights=down_lora_b_weights,
|
||||
# UPDATED FIELDS
|
||||
seg_indptr=seg_indptr,
|
||||
req_to_lora=weight_indices,
|
||||
lora_ranks=lora_ranks,
|
||||
adapter_enabled=adapter_enabled,
|
||||
max_lora_rank=max_lora_rank,
|
||||
num_experts=num_experts,
|
||||
lora_use_virtual_experts=lora_use_virtual_experts,
|
||||
)
|
||||
|
||||
|
||||
def torch_naive_moe_with_lora(
|
||||
hidden_states,
|
||||
w13,
|
||||
w2,
|
||||
b13,
|
||||
b2,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
lora_info,
|
||||
token_lora_mapping,
|
||||
):
|
||||
"""
|
||||
Naive implementation. Note: We pass 'token_lora_mapping' explicitly because
|
||||
lora_info no longer contains it, but the naive token-loop logic needs it.
|
||||
"""
|
||||
num_tokens, hidden_dim = hidden_states.shape
|
||||
top_k = topk_ids.shape[1]
|
||||
num_experts = w13.shape[0]
|
||||
|
||||
# Expand hidden states for top-k routing
|
||||
hidden_expanded = (
|
||||
hidden_states.unsqueeze(1).expand(-1, top_k, -1).reshape(-1, hidden_dim)
|
||||
)
|
||||
|
||||
# 1. Gate/Up Projection (Base)
|
||||
gate_up_out = torch.zeros(
|
||||
num_tokens * top_k,
|
||||
w13.shape[1],
|
||||
dtype=hidden_states.dtype,
|
||||
device=hidden_states.device,
|
||||
)
|
||||
|
||||
for expert_id in range(num_experts):
|
||||
mask = (topk_ids == expert_id).flatten()
|
||||
if mask.any():
|
||||
expert_result = hidden_expanded[mask] @ w13[expert_id].T
|
||||
gate_up_out[mask] = expert_result
|
||||
if b13 is not None:
|
||||
gate_up_out[mask] += b13[expert_id]
|
||||
|
||||
gate_up_out = gate_up_out.view(num_tokens, top_k, -1)
|
||||
|
||||
# 1.5. LoRA Gate/Up Delta
|
||||
# gate_up_lora_a is packed as [gate_a; up_a] along rank dim → [2*r, hidden_dim]
|
||||
# gate_up_lora_b is packed as [gate_b; up_b] along output dim → [2*inter, r]
|
||||
# Correct computation splits them: gate uses first r rows of A with first half of B,
|
||||
# up uses last r rows of A with second half of B.
|
||||
if lora_info.max_lora_rank > 0:
|
||||
r = lora_info.max_lora_rank
|
||||
for i in range(num_tokens):
|
||||
for k in range(top_k):
|
||||
expert_id = topk_ids[i, k]
|
||||
lora_id = token_lora_mapping[i]
|
||||
|
||||
if lora_id < len(lora_info.lora_ranks):
|
||||
lora_a = lora_info.gate_up_lora_a_weights[lora_id, expert_id]
|
||||
lora_b = lora_info.gate_up_lora_b_weights[lora_id, expert_id]
|
||||
half = lora_b.shape[0] // 2
|
||||
lora_a_result = lora_a @ hidden_states[i]
|
||||
gate_delta = lora_b[:half, :] @ lora_a_result[:r]
|
||||
up_delta = lora_b[half:, :] @ lora_a_result[r:]
|
||||
gate_up_out[i, k] += torch.cat([gate_delta, up_delta])
|
||||
|
||||
# 2. Activation
|
||||
gate_up_dim = gate_up_out.shape[-1]
|
||||
gate_dim = gate_up_dim // 2
|
||||
gate = gate_up_out[..., :gate_dim]
|
||||
up = gate_up_out[..., gate_dim:]
|
||||
|
||||
silu_gate = torch.nn.functional.silu(gate)
|
||||
intermediate_out = silu_gate * up
|
||||
|
||||
# 3. Down Projection (Base)
|
||||
down_out = torch.zeros(
|
||||
num_tokens,
|
||||
top_k,
|
||||
hidden_dim,
|
||||
dtype=hidden_states.dtype,
|
||||
device=hidden_states.device,
|
||||
)
|
||||
|
||||
for expert_id in range(num_experts):
|
||||
mask = topk_ids == expert_id
|
||||
if mask.any():
|
||||
masked_intermediate = intermediate_out[mask]
|
||||
expert_down_result = masked_intermediate @ w2[expert_id].T
|
||||
down_out[mask] = expert_down_result
|
||||
if b2 is not None:
|
||||
down_out[mask] += b2[expert_id]
|
||||
|
||||
# 3.5. LoRA Down Delta
|
||||
if lora_info.max_lora_rank > 0:
|
||||
for i in range(num_tokens):
|
||||
for k in range(top_k):
|
||||
expert_id = topk_ids[i, k]
|
||||
lora_id = token_lora_mapping[i] # Use explicit mapping
|
||||
|
||||
if lora_id < len(lora_info.lora_ranks):
|
||||
lora_a = lora_info.down_lora_a_weights[lora_id, expert_id]
|
||||
lora_b = lora_info.down_lora_b_weights[lora_id, expert_id]
|
||||
lora_a_result = lora_a @ intermediate_out[i, k]
|
||||
lora_b_result = lora_b @ lora_a_result
|
||||
down_out[i, k] += lora_b_result
|
||||
|
||||
# 4. Final Reduction
|
||||
weighted_out = down_out * topk_weights.unsqueeze(-1)
|
||||
final_out = weighted_out.sum(dim=1)
|
||||
|
||||
return final_out
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [32, 64])
|
||||
@pytest.mark.parametrize("top_k_num", [1, 2])
|
||||
@pytest.mark.parametrize("num_experts", [8, 20])
|
||||
@pytest.mark.parametrize("max_lora_rank", [8, 16])
|
||||
def test_lora_moe_runner_multi_expert(
|
||||
num_tokens, top_k_num, num_experts, max_lora_rank
|
||||
):
|
||||
# Fixed parameters
|
||||
max_loras = 2
|
||||
hidden_dim = 512
|
||||
intermediate_dim = 1024
|
||||
|
||||
dtype = torch.float32
|
||||
device = "cuda:0"
|
||||
seed = 42
|
||||
|
||||
torch.set_default_device(device)
|
||||
set_random_seed(seed)
|
||||
|
||||
num_sequences = 4
|
||||
|
||||
# Generate Data using the new Request-Based generator
|
||||
topk_ids, topk_weights, seg_indptr, req_to_lora, token_lora_mapping = sample_data(
|
||||
num_tokens, num_sequences, max_loras, num_experts, top_k_num, dtype, device
|
||||
)
|
||||
|
||||
gate_up_dim = intermediate_dim * 2
|
||||
|
||||
# Initialize experts
|
||||
w13 = torch.randn(num_experts, gate_up_dim, hidden_dim, dtype=dtype) * 0.1
|
||||
w2 = torch.randn(num_experts, hidden_dim, intermediate_dim, dtype=dtype) * 0.1
|
||||
b13 = torch.randn(num_experts, gate_up_dim, dtype=dtype) * 0.1
|
||||
b2 = torch.randn(num_experts, hidden_dim, dtype=dtype) * 0.1
|
||||
|
||||
hidden_states = torch.randn(num_tokens, hidden_dim, dtype=dtype)
|
||||
|
||||
# Create LoRA Info using the new fields
|
||||
lora_info_delta = create_lora_info(
|
||||
seg_indptr=seg_indptr,
|
||||
weight_indices=req_to_lora,
|
||||
topk_ids=topk_ids,
|
||||
max_loras=max_loras,
|
||||
num_experts=num_experts,
|
||||
max_lora_rank=max_lora_rank,
|
||||
hidden_dim=hidden_dim,
|
||||
intermediate_dim=intermediate_dim,
|
||||
gate_up_dim=gate_up_dim,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
|
||||
lora_info_baseline = create_lora_info(
|
||||
seg_indptr=seg_indptr,
|
||||
weight_indices=req_to_lora,
|
||||
topk_ids=topk_ids,
|
||||
max_loras=max_loras,
|
||||
num_experts=num_experts,
|
||||
max_lora_rank=0, # Set rank to 0 for baseline
|
||||
hidden_dim=hidden_dim,
|
||||
intermediate_dim=intermediate_dim,
|
||||
gate_up_dim=gate_up_dim,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
|
||||
# Sort tokens for the runner
|
||||
topk_ids_flat = topk_ids.flatten()
|
||||
sorted_indices = torch.argsort(topk_ids_flat)
|
||||
sorted_token_ids = sorted_indices // top_k_num
|
||||
expert_ids = topk_ids_flat[sorted_indices]
|
||||
|
||||
num_dispatched = num_tokens * top_k_num
|
||||
num_tokens_post_padded = torch.tensor(
|
||||
[num_dispatched], dtype=torch.int32, device=device
|
||||
)
|
||||
|
||||
quant_info = TritonMoeQuantInfo(
|
||||
w13_weight=w13,
|
||||
w2_weight=w2,
|
||||
b13=b13,
|
||||
b2=b2,
|
||||
)
|
||||
|
||||
config = MoeRunnerConfig(
|
||||
activation="silu",
|
||||
is_gated=True,
|
||||
inplace=False,
|
||||
no_combine=False,
|
||||
gemm1_alpha=None,
|
||||
gemm1_clamp_limit=None,
|
||||
routed_scaling_factor=1.0,
|
||||
apply_router_weight_on_input=False,
|
||||
num_local_experts=num_experts,
|
||||
)
|
||||
|
||||
# Create StandardTopKOutput
|
||||
router_logits = torch.randn(num_tokens, num_experts, dtype=dtype, device=device)
|
||||
topk_output = StandardTopKOutput(
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
router_logits=router_logits,
|
||||
)
|
||||
|
||||
# Create StandardDispatchOutput
|
||||
dispatch_output = StandardDispatchOutput(
|
||||
hidden_states=hidden_states,
|
||||
hidden_states_scale=None,
|
||||
topk_output=topk_output,
|
||||
)
|
||||
|
||||
class MockServerArgs:
|
||||
enable_deterministic_inference = False
|
||||
|
||||
with patch(
|
||||
"sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_config.get_global_server_args",
|
||||
return_value=MockServerArgs(),
|
||||
):
|
||||
runner = MoeRunner(MoeRunnerBackend.TRITON, config, lora_enabled=True)
|
||||
|
||||
# 3. Get outputs for both scenarios
|
||||
output_with_lora = runner.run(
|
||||
dispatch_output, quant_info, lora_info_delta
|
||||
).hidden_states
|
||||
output_baseline = runner.run(
|
||||
dispatch_output, quant_info, lora_info_baseline
|
||||
).hidden_states
|
||||
|
||||
# Run Naive Torch Implementation (Uses dense mapping for verification)
|
||||
torch_output_lora = torch_naive_moe_with_lora(
|
||||
hidden_states,
|
||||
w13,
|
||||
w2,
|
||||
b13,
|
||||
b2,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
lora_info_delta,
|
||||
token_lora_mapping,
|
||||
)
|
||||
|
||||
torch_output_base = torch_naive_moe_with_lora(
|
||||
hidden_states,
|
||||
w13,
|
||||
w2,
|
||||
b13,
|
||||
b2,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
lora_info_baseline,
|
||||
token_lora_mapping,
|
||||
)
|
||||
|
||||
# The actual "Delta" (LoRA effect) for both
|
||||
sglang_delta = output_with_lora - output_baseline
|
||||
torch_delta = torch_output_lora - torch_output_base
|
||||
|
||||
# Larger expert counts accumulate more numerical drift in Triton kernels on GB300
|
||||
tol = 0.15 if num_experts >= 20 else 5e-2
|
||||
torch.testing.assert_close(sglang_delta, torch_delta, atol=tol, rtol=tol)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [32, 64])
|
||||
@pytest.mark.parametrize("top_k_num", [1, 2])
|
||||
@pytest.mark.parametrize("num_experts", [8, 20])
|
||||
@pytest.mark.parametrize("max_lora_rank", [8, 16])
|
||||
def test_lora_moe_runner_virtual_experts(
|
||||
num_tokens, top_k_num, num_experts, max_lora_rank
|
||||
):
|
||||
# Fixed parameters
|
||||
max_loras = 2
|
||||
hidden_dim = 512
|
||||
intermediate_dim = 1024
|
||||
|
||||
dtype = torch.float32
|
||||
device = "cuda:0"
|
||||
seed = 42
|
||||
|
||||
torch.set_default_device(device)
|
||||
set_random_seed(seed)
|
||||
|
||||
num_sequences = 4
|
||||
|
||||
# Generate Data using the new Request-Based generator
|
||||
topk_ids, topk_weights, seg_indptr, req_to_lora, token_lora_mapping = sample_data(
|
||||
num_tokens, num_sequences, max_loras, num_experts, top_k_num, dtype, device
|
||||
)
|
||||
|
||||
gate_up_dim = intermediate_dim * 2
|
||||
|
||||
# Initialize experts
|
||||
w13 = torch.randn(num_experts, gate_up_dim, hidden_dim, dtype=dtype) * 0.1
|
||||
w2 = torch.randn(num_experts, hidden_dim, intermediate_dim, dtype=dtype) * 0.1
|
||||
b13 = torch.randn(num_experts, gate_up_dim, dtype=dtype) * 0.1
|
||||
b2 = torch.randn(num_experts, hidden_dim, dtype=dtype) * 0.1
|
||||
|
||||
hidden_states = torch.randn(num_tokens, hidden_dim, dtype=dtype)
|
||||
|
||||
# Create LoRA Info with virtual experts enabled
|
||||
lora_info_delta = create_lora_info(
|
||||
seg_indptr=seg_indptr,
|
||||
weight_indices=req_to_lora,
|
||||
topk_ids=topk_ids,
|
||||
max_loras=max_loras,
|
||||
num_experts=num_experts,
|
||||
max_lora_rank=max_lora_rank,
|
||||
hidden_dim=hidden_dim,
|
||||
intermediate_dim=intermediate_dim,
|
||||
gate_up_dim=gate_up_dim,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
lora_use_virtual_experts=True,
|
||||
)
|
||||
|
||||
lora_info_baseline = create_lora_info(
|
||||
seg_indptr=seg_indptr,
|
||||
weight_indices=req_to_lora,
|
||||
topk_ids=topk_ids,
|
||||
max_loras=max_loras,
|
||||
num_experts=num_experts,
|
||||
max_lora_rank=0,
|
||||
hidden_dim=hidden_dim,
|
||||
intermediate_dim=intermediate_dim,
|
||||
gate_up_dim=gate_up_dim,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
lora_use_virtual_experts=True,
|
||||
)
|
||||
|
||||
quant_info = TritonMoeQuantInfo(
|
||||
w13_weight=w13,
|
||||
w2_weight=w2,
|
||||
b13=b13,
|
||||
b2=b2,
|
||||
)
|
||||
|
||||
config = MoeRunnerConfig(
|
||||
activation="silu",
|
||||
is_gated=True,
|
||||
inplace=False,
|
||||
no_combine=False,
|
||||
gemm1_alpha=None,
|
||||
gemm1_clamp_limit=None,
|
||||
routed_scaling_factor=1.0,
|
||||
apply_router_weight_on_input=False,
|
||||
num_local_experts=num_experts,
|
||||
)
|
||||
|
||||
router_logits = torch.randn(num_tokens, num_experts, dtype=dtype, device=device)
|
||||
topk_output = StandardTopKOutput(
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
router_logits=router_logits,
|
||||
)
|
||||
|
||||
dispatch_output = StandardDispatchOutput(
|
||||
hidden_states=hidden_states,
|
||||
hidden_states_scale=None,
|
||||
topk_output=topk_output,
|
||||
)
|
||||
|
||||
class MockServerArgs:
|
||||
enable_deterministic_inference = False
|
||||
|
||||
with patch(
|
||||
"sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_config.get_global_server_args",
|
||||
return_value=MockServerArgs(),
|
||||
):
|
||||
runner = MoeRunner(MoeRunnerBackend.TRITON, config, lora_enabled=True)
|
||||
|
||||
output_with_lora = runner.run(
|
||||
dispatch_output, quant_info, lora_info_delta
|
||||
).hidden_states
|
||||
output_baseline = runner.run(
|
||||
dispatch_output, quant_info, lora_info_baseline
|
||||
).hidden_states
|
||||
|
||||
# Run Naive Torch Implementation (Uses dense mapping for verification)
|
||||
torch_output_lora = torch_naive_moe_with_lora(
|
||||
hidden_states,
|
||||
w13,
|
||||
w2,
|
||||
b13,
|
||||
b2,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
lora_info_delta,
|
||||
token_lora_mapping,
|
||||
)
|
||||
|
||||
torch_output_base = torch_naive_moe_with_lora(
|
||||
hidden_states,
|
||||
w13,
|
||||
w2,
|
||||
b13,
|
||||
b2,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
lora_info_baseline,
|
||||
token_lora_mapping,
|
||||
)
|
||||
|
||||
# The actual "Delta" (LoRA effect) for both
|
||||
sglang_delta = output_with_lora - output_baseline
|
||||
torch_delta = torch_output_lora - torch_output_base
|
||||
|
||||
# Larger expert counts accumulate more numerical drift in Triton kernels on GB300
|
||||
tol = 0.15 if num_experts >= 20 else 5e-2
|
||||
torch.testing.assert_close(sglang_delta, torch_delta, atol=tol, rtol=tol)
|
||||
|
||||
|
||||
def _setup_marlin_moe_weights(num_experts, n, k, dtype):
|
||||
"""Quantize float weights into AWQ Marlin format for testing."""
|
||||
from sgl_kernel.scalar_type import scalar_types
|
||||
|
||||
from sglang.test.test_marlin_utils import awq_marlin_quantize
|
||||
|
||||
group_size = 128
|
||||
quant_type = scalar_types.uint4
|
||||
|
||||
w = torch.randn((num_experts, n, k), device="cuda", dtype=dtype) / 20
|
||||
|
||||
w_ref_l, qweight_l, scales_l, zeros_l = [], [], [], []
|
||||
for i in range(num_experts):
|
||||
w_ref, qweight, scales, zeros = awq_marlin_quantize(
|
||||
w[i].transpose(1, 0), quant_type, group_size
|
||||
)
|
||||
w_ref_l.append(w_ref.T)
|
||||
qweight_l.append(qweight)
|
||||
scales_l.append(scales)
|
||||
zeros_l.append(zeros)
|
||||
|
||||
def _stack(tensors):
|
||||
dev = tensors[0].device
|
||||
return torch.stack(tensors, dim=0).to(dev)
|
||||
|
||||
return (
|
||||
_stack(w_ref_l),
|
||||
_stack(qweight_l).contiguous(),
|
||||
_stack(scales_l),
|
||||
_stack(zeros_l),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [32, 64])
|
||||
@pytest.mark.parametrize("top_k_num", [1, 2])
|
||||
@pytest.mark.parametrize("num_experts", [8])
|
||||
@pytest.mark.parametrize("max_lora_rank", [8, 16])
|
||||
def test_lora_moe_runner_marlin(num_tokens, top_k_num, num_experts, max_lora_rank):
|
||||
from sglang.srt.layers.moe.moe_runner.marlin import MarlinMoeQuantInfo
|
||||
|
||||
max_loras = 2
|
||||
hidden_dim = 512
|
||||
intermediate_dim = 1024
|
||||
gate_up_dim = intermediate_dim * 2
|
||||
|
||||
dtype = torch.float16
|
||||
device = "cuda:0"
|
||||
seed = 42
|
||||
|
||||
torch.set_default_device(device)
|
||||
set_random_seed(seed)
|
||||
|
||||
num_sequences = 4
|
||||
|
||||
topk_ids, topk_weights, seg_indptr, req_to_lora, token_lora_mapping = sample_data(
|
||||
num_tokens,
|
||||
num_sequences,
|
||||
max_loras,
|
||||
num_experts,
|
||||
top_k_num,
|
||||
dtype,
|
||||
device,
|
||||
)
|
||||
|
||||
# Quantize base weights to Marlin format
|
||||
_, w13_qweight, w13_scales, w13_qzeros = _setup_marlin_moe_weights(
|
||||
num_experts, gate_up_dim, hidden_dim, dtype
|
||||
)
|
||||
_, w2_qweight, w2_scales, w2_qzeros = _setup_marlin_moe_weights(
|
||||
num_experts, hidden_dim, intermediate_dim, dtype
|
||||
)
|
||||
|
||||
hidden_states = torch.randn(num_tokens, hidden_dim, dtype=dtype, device=device)
|
||||
|
||||
lora_info_delta = create_lora_info(
|
||||
seg_indptr=seg_indptr,
|
||||
weight_indices=req_to_lora,
|
||||
topk_ids=topk_ids,
|
||||
max_loras=max_loras,
|
||||
num_experts=num_experts,
|
||||
max_lora_rank=max_lora_rank,
|
||||
hidden_dim=hidden_dim,
|
||||
intermediate_dim=intermediate_dim,
|
||||
gate_up_dim=gate_up_dim,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
lora_use_virtual_experts=True,
|
||||
)
|
||||
|
||||
lora_info_baseline = create_lora_info(
|
||||
seg_indptr=seg_indptr,
|
||||
weight_indices=req_to_lora,
|
||||
topk_ids=topk_ids,
|
||||
max_loras=max_loras,
|
||||
num_experts=num_experts,
|
||||
max_lora_rank=0,
|
||||
hidden_dim=hidden_dim,
|
||||
intermediate_dim=intermediate_dim,
|
||||
gate_up_dim=gate_up_dim,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
lora_use_virtual_experts=True,
|
||||
)
|
||||
|
||||
quant_info = MarlinMoeQuantInfo(
|
||||
w13_qweight=w13_qweight,
|
||||
w2_qweight=w2_qweight,
|
||||
w13_scales=w13_scales,
|
||||
w2_scales=w2_scales,
|
||||
w13_qzeros=w13_qzeros,
|
||||
w2_qzeros=w2_qzeros,
|
||||
w13_g_idx=None,
|
||||
w2_g_idx=None,
|
||||
w13_g_idx_sort_indices=None,
|
||||
w2_g_idx_sort_indices=None,
|
||||
weight_bits=4,
|
||||
)
|
||||
|
||||
config = MoeRunnerConfig(
|
||||
activation="silu",
|
||||
is_gated=True,
|
||||
inplace=False,
|
||||
no_combine=False,
|
||||
gemm1_alpha=None,
|
||||
gemm1_clamp_limit=None,
|
||||
routed_scaling_factor=1.0,
|
||||
apply_router_weight_on_input=False,
|
||||
num_local_experts=num_experts,
|
||||
)
|
||||
|
||||
router_logits = torch.randn(num_tokens, num_experts, dtype=dtype, device=device)
|
||||
topk_output = StandardTopKOutput(
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
router_logits=router_logits,
|
||||
)
|
||||
dispatch_output = StandardDispatchOutput(
|
||||
hidden_states=hidden_states,
|
||||
hidden_states_scale=None,
|
||||
topk_output=topk_output,
|
||||
)
|
||||
|
||||
class MockServerArgs:
|
||||
enable_deterministic_inference = False
|
||||
|
||||
with patch(
|
||||
"sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_config.get_global_server_args",
|
||||
return_value=MockServerArgs(),
|
||||
):
|
||||
runner = MoeRunner(MoeRunnerBackend.MARLIN, config, lora_enabled=True)
|
||||
output_with_lora = runner.run(
|
||||
dispatch_output, quant_info, lora_info_delta
|
||||
).hidden_states
|
||||
output_baseline = runner.run(
|
||||
dispatch_output, quant_info, lora_info_baseline
|
||||
).hidden_states
|
||||
|
||||
marlin_delta = output_with_lora - output_baseline
|
||||
|
||||
# Verify the LoRA hooks fired and produced a non-trivial delta
|
||||
assert marlin_delta.abs().max().item() > 1e-4, (
|
||||
f"LoRA delta is too small ({marlin_delta.abs().max().item():.6f}), "
|
||||
"hooks may not be firing"
|
||||
)
|
||||
assert torch.isfinite(
|
||||
output_with_lora
|
||||
).all(), "Marlin+LoRA output contains non-finite values"
|
||||
assert torch.isfinite(
|
||||
output_baseline
|
||||
).all(), "Marlin baseline output contains non-finite values"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__])
|
||||
@@ -1,289 +0,0 @@
|
||||
# Copyright 2023-2025 SGLang Team
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# ==============================================================================
|
||||
|
||||
"""
|
||||
Correctness test: Marlin (int4 base + LoRA) vs Triton (dequantized base + LoRA).
|
||||
|
||||
Fake-quantizes random weights to int4/Marlin format and dequantizes them with the
|
||||
same path, then runs both backends through MoeRunner and compares LoRA deltas.
|
||||
"""
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
|
||||
from sglang.srt.layers.moe.moe_runner.marlin import MarlinMoeQuantInfo
|
||||
from sglang.srt.layers.moe.moe_runner.runner import MoeRunner
|
||||
from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo
|
||||
from sglang.srt.layers.moe.token_dispatcher.standard import StandardDispatchOutput
|
||||
from sglang.srt.layers.moe.topk import StandardTopKOutput
|
||||
from sglang.srt.layers.moe.utils import MoeRunnerBackend
|
||||
from sglang.srt.lora.lora_moe_runners import LoRAInfo
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=129, suite="stage-b-test-1-gpu-large")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fake quantization helpers (symmetric int4, matching Marlin's dequant path)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _quantize_per_expert(w_float: torch.Tensor, K: int, group_size: int):
|
||||
"""Quantize [N, K] float weight to int4. Returns (q_int [N,K], scales_bf16 [N,groups])."""
|
||||
N = w_float.shape[0]
|
||||
num_groups = K // group_size
|
||||
|
||||
w_grouped = w_float.reshape(N, num_groups, group_size)
|
||||
scales_fp32 = w_grouped.abs().amax(dim=-1) / 7.0
|
||||
scales_fp32 = scales_fp32.clamp(min=1e-6)
|
||||
scales_bf16 = scales_fp32.to(torch.bfloat16)
|
||||
|
||||
scales_for_quant = scales_bf16.float()
|
||||
q_int = torch.zeros(N, K, dtype=torch.int32, device=w_float.device)
|
||||
for g in range(num_groups):
|
||||
s = scales_for_quant[:, g : g + 1]
|
||||
sl = slice(g * group_size, (g + 1) * group_size)
|
||||
q_int[:, sl] = torch.round(w_float[:, sl] / s).clamp(-8, 7).to(torch.int32) + 8
|
||||
|
||||
return q_int, scales_bf16
|
||||
|
||||
|
||||
def _fake_quantize_to_marlin_int4(weight_bf16: torch.Tensor):
|
||||
"""Fake-quantize [E, N, K] bf16 weight to Marlin int4 format.
|
||||
|
||||
Returns: (qweight, scales, g_idx, g_idx_sort_indices)
|
||||
"""
|
||||
from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack
|
||||
from sglang.srt.layers.quantization.marlin_utils import marlin_permute_scales
|
||||
from sglang.srt.layers.quantization.utils import pack_rows
|
||||
|
||||
E, N, K = weight_bf16.shape
|
||||
num_bits = 4
|
||||
group_size = 128
|
||||
device = weight_bf16.device
|
||||
|
||||
all_qweight, all_scales = [], []
|
||||
for e in range(E):
|
||||
q_int, scales_bf16 = _quantize_per_expert(weight_bf16[e].float(), K, group_size)
|
||||
w_quant_t = q_int.t().contiguous()
|
||||
packed = pack_rows(w_quant_t, num_bits, K, N)
|
||||
perm = torch.arange(K, device=device, dtype=torch.int32)
|
||||
all_qweight.append(gptq_marlin_repack(packed.to(device), perm, K, N, num_bits))
|
||||
all_scales.append(
|
||||
marlin_permute_scales(
|
||||
scales_bf16.t().contiguous().to(device), K, N, group_size
|
||||
)
|
||||
)
|
||||
|
||||
g_idx = (
|
||||
(torch.arange(K, device=device, dtype=torch.int32) // group_size)
|
||||
.unsqueeze(0)
|
||||
.expand(E, -1)
|
||||
.contiguous()
|
||||
)
|
||||
sort_indices = (
|
||||
torch.arange(K, device=device, dtype=torch.int32)
|
||||
.unsqueeze(0)
|
||||
.expand(E, -1)
|
||||
.contiguous()
|
||||
)
|
||||
|
||||
return torch.stack(all_qweight), torch.stack(all_scales), g_idx, sort_indices
|
||||
|
||||
|
||||
def _dequantize_from_marlin_int4(weight_bf16_orig: torch.Tensor, group_size: int = 128):
|
||||
"""Dequantize using the same path as _fake_quantize, so Triton reference matches Marlin."""
|
||||
E, N, K = weight_bf16_orig.shape
|
||||
result = torch.zeros_like(weight_bf16_orig)
|
||||
for e in range(E):
|
||||
q_int, scales_bf16 = _quantize_per_expert(
|
||||
weight_bf16_orig[e].float(), K, group_size
|
||||
)
|
||||
num_groups = K // group_size
|
||||
for g in range(num_groups):
|
||||
sl = slice(g * group_size, (g + 1) * group_size)
|
||||
s = scales_bf16[:, g : g + 1]
|
||||
result[e, :, sl] = (q_int[:, sl] - 8).to(torch.bfloat16) * s
|
||||
return result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [1, 8, 32])
|
||||
@pytest.mark.parametrize("top_k", [2, 8])
|
||||
def test_marlin_vs_triton_lora_correctness(num_tokens, top_k):
|
||||
torch.manual_seed(42)
|
||||
|
||||
device = "cuda"
|
||||
dtype = torch.bfloat16
|
||||
|
||||
hidden_dim = 7168
|
||||
intermediate_dim = 2048
|
||||
gate_up_dim = 2 * intermediate_dim
|
||||
num_experts = 64
|
||||
lora_rank = 32
|
||||
num_loras = 1
|
||||
|
||||
hidden = torch.randn(num_tokens, hidden_dim, dtype=dtype, device=device)
|
||||
topk_weights = torch.randn(
|
||||
num_tokens, top_k, dtype=torch.float32, device=device
|
||||
).softmax(dim=-1)
|
||||
topk_ids = torch.randint(
|
||||
0, num_experts, (num_tokens, top_k), dtype=torch.int32, device=device
|
||||
)
|
||||
|
||||
# Base weights (random bf16)
|
||||
w13_bf16 = (
|
||||
torch.randn(num_experts, gate_up_dim, hidden_dim, dtype=dtype, device=device)
|
||||
* 0.01
|
||||
)
|
||||
w2_bf16 = (
|
||||
torch.randn(
|
||||
num_experts, hidden_dim, intermediate_dim, dtype=dtype, device=device
|
||||
)
|
||||
* 0.01
|
||||
)
|
||||
|
||||
# LoRA weights (shared across both paths)
|
||||
gu_lora_a = (
|
||||
torch.randn(num_loras, 1, lora_rank * 2, hidden_dim, dtype=dtype, device=device)
|
||||
* 0.01
|
||||
)
|
||||
gu_lora_b = (
|
||||
torch.randn(
|
||||
num_loras, num_experts, gate_up_dim, lora_rank, dtype=dtype, device=device
|
||||
)
|
||||
* 0.01
|
||||
)
|
||||
dn_lora_a = (
|
||||
torch.randn(
|
||||
num_loras,
|
||||
num_experts,
|
||||
lora_rank,
|
||||
intermediate_dim,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
* 0.01
|
||||
)
|
||||
dn_lora_b = (
|
||||
torch.randn(num_loras, 1, hidden_dim, lora_rank, dtype=dtype, device=device)
|
||||
* 0.01
|
||||
)
|
||||
|
||||
# Token-to-LoRA mapping: all tokens use adapter 0
|
||||
seg_indptr = torch.tensor([0, num_tokens], dtype=torch.int32, device=device)
|
||||
req_to_lora = torch.tensor([0], dtype=torch.int32, device=device)
|
||||
|
||||
def _make_lora_info(rank):
|
||||
return LoRAInfo(
|
||||
gate_up_lora_a_weights=gu_lora_a if rank > 0 else gu_lora_a[:, :, :0, :],
|
||||
gate_up_lora_b_weights=gu_lora_b if rank > 0 else gu_lora_b[:, :, :, :0],
|
||||
down_lora_a_weights=dn_lora_a if rank > 0 else dn_lora_a[:, :, :0, :],
|
||||
down_lora_b_weights=dn_lora_b if rank > 0 else dn_lora_b[:, :, :, :0],
|
||||
seg_indptr=seg_indptr,
|
||||
req_to_lora=req_to_lora,
|
||||
lora_ranks=torch.full((num_loras,), rank, dtype=torch.int32, device=device),
|
||||
adapter_enabled=torch.ones(num_loras + 1, dtype=torch.int32, device=device),
|
||||
max_lora_rank=rank,
|
||||
num_experts=num_experts,
|
||||
lora_use_virtual_experts=True,
|
||||
experts_shared_outer_loras=True,
|
||||
)
|
||||
|
||||
lora_info = _make_lora_info(lora_rank)
|
||||
lora_baseline = _make_lora_info(0)
|
||||
|
||||
# Quantize for Marlin, dequantize for Triton reference
|
||||
w13_qw, w13_sc, w13_gidx, w13_si = _fake_quantize_to_marlin_int4(w13_bf16)
|
||||
w2_qw, w2_sc, w2_gidx, w2_si = _fake_quantize_to_marlin_int4(w2_bf16)
|
||||
w13_deq = _dequantize_from_marlin_int4(w13_bf16)
|
||||
w2_deq = _dequantize_from_marlin_int4(w2_bf16)
|
||||
|
||||
marlin_qi = MarlinMoeQuantInfo(
|
||||
w13_qweight=w13_qw,
|
||||
w2_qweight=w2_qw,
|
||||
w13_scales=w13_sc,
|
||||
w2_scales=w2_sc,
|
||||
w13_g_idx=w13_gidx,
|
||||
w2_g_idx=w2_gidx,
|
||||
w13_g_idx_sort_indices=w13_si,
|
||||
w2_g_idx_sort_indices=w2_si,
|
||||
weight_bits=4,
|
||||
)
|
||||
triton_qi = TritonMoeQuantInfo(
|
||||
w13_weight=w13_deq, w2_weight=w2_deq, b13=None, b2=None
|
||||
)
|
||||
|
||||
config = MoeRunnerConfig(
|
||||
activation="silu",
|
||||
is_gated=True,
|
||||
inplace=False,
|
||||
no_combine=False,
|
||||
gemm1_alpha=None,
|
||||
gemm1_clamp_limit=None,
|
||||
routed_scaling_factor=1.0,
|
||||
apply_router_weight_on_input=False,
|
||||
num_local_experts=num_experts,
|
||||
)
|
||||
|
||||
router_logits = torch.randn(num_tokens, num_experts, dtype=dtype, device=device)
|
||||
topk_output = StandardTopKOutput(
|
||||
topk_weights=topk_weights, topk_ids=topk_ids, router_logits=router_logits
|
||||
)
|
||||
dispatch_output = StandardDispatchOutput(
|
||||
hidden_states=hidden, hidden_states_scale=None, topk_output=topk_output
|
||||
)
|
||||
|
||||
class MockServerArgs:
|
||||
enable_deterministic_inference = False
|
||||
|
||||
with patch(
|
||||
"sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_config.get_global_server_args",
|
||||
return_value=MockServerArgs(),
|
||||
):
|
||||
marlin_runner = MoeRunner(MoeRunnerBackend.MARLIN, config, lora_enabled=True)
|
||||
triton_runner = MoeRunner(MoeRunnerBackend.TRITON, config, lora_enabled=True)
|
||||
|
||||
marlin_out = marlin_runner.run(
|
||||
dispatch_output, marlin_qi, lora_info
|
||||
).hidden_states
|
||||
marlin_base = marlin_runner.run(
|
||||
dispatch_output, marlin_qi, lora_baseline
|
||||
).hidden_states
|
||||
triton_out = triton_runner.run(
|
||||
dispatch_output, triton_qi, lora_info
|
||||
).hidden_states
|
||||
triton_base = triton_runner.run(
|
||||
dispatch_output, triton_qi, lora_baseline
|
||||
).hidden_states
|
||||
|
||||
marlin_delta = marlin_out - marlin_base
|
||||
triton_delta = triton_out - triton_base
|
||||
|
||||
# Remaining error is from kernel-level accumulation differences
|
||||
# (Marlin fp32 reduce vs Triton bf16 dot), not from quantization mismatch.
|
||||
torch.testing.assert_close(
|
||||
marlin_delta.float(), triton_delta.float(), atol=0.01, rtol=0.05
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__])
|
||||
@@ -1,236 +0,0 @@
|
||||
"""Test that sgemm kernels produce identical results with and without SORTED_BY_ADAPTER."""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=11, suite="stage-b-test-1-gpu-large")
|
||||
|
||||
|
||||
def _make_batch_info(
|
||||
bs: int,
|
||||
weight_indices: list[int],
|
||||
lora_ranks: list[int],
|
||||
scalings: list[float],
|
||||
device: str = "cuda",
|
||||
) -> Any:
|
||||
"""Build a per-sequence LoRABatchInfo (no permutation)."""
|
||||
from sglang.srt.lora.utils import LoRABatchInfo
|
||||
|
||||
seg_lens = torch.ones(bs, dtype=torch.int32, device=device)
|
||||
seg_indptr = torch.zeros(bs + 1, dtype=torch.int32, device=device)
|
||||
seg_indptr[1:] = torch.cumsum(seg_lens, dim=0)
|
||||
return LoRABatchInfo(
|
||||
bs=bs,
|
||||
use_cuda_graph=False,
|
||||
num_segments=bs,
|
||||
seg_lens=seg_lens,
|
||||
seg_indptr=seg_indptr,
|
||||
max_len=1,
|
||||
weight_indices=torch.tensor(weight_indices, dtype=torch.int32, device=device),
|
||||
lora_ranks=torch.tensor(lora_ranks, dtype=torch.int32, device=device),
|
||||
scalings=torch.tensor(scalings, dtype=torch.float, device=device),
|
||||
permutation=None,
|
||||
)
|
||||
|
||||
|
||||
def _make_sorted_batch_info(
|
||||
weight_indices: list[int],
|
||||
lora_ranks: list[int],
|
||||
scalings: list[float],
|
||||
max_loras: int,
|
||||
device: str = "cuda",
|
||||
) -> Any:
|
||||
from sglang.srt.lora.utils import LoRABatchInfo
|
||||
|
||||
"""Build a merged-by-adapter LoRABatchInfo (with permutation)."""
|
||||
wi = torch.tensor(weight_indices, dtype=torch.int32, device=device)
|
||||
bs = wi.shape[0]
|
||||
|
||||
perm = torch.argsort(wi, stable=True).to(torch.int32)
|
||||
sorted_wi = wi[perm]
|
||||
adapter_ids = torch.arange(max_loras, device=device, dtype=torch.int32)
|
||||
seg_starts = torch.searchsorted(sorted_wi, adapter_ids)
|
||||
seg_ends = torch.searchsorted(sorted_wi, adapter_ids, right=True)
|
||||
seg_lens = seg_ends - seg_starts
|
||||
|
||||
seg_indptr = torch.zeros(max_loras + 1, dtype=torch.int32, device=device)
|
||||
seg_indptr[1:] = torch.cumsum(seg_lens, dim=0)
|
||||
|
||||
return LoRABatchInfo(
|
||||
bs=max_loras,
|
||||
use_cuda_graph=False,
|
||||
num_segments=max_loras,
|
||||
seg_lens=seg_lens,
|
||||
seg_indptr=seg_indptr,
|
||||
max_len=bs,
|
||||
weight_indices=adapter_ids,
|
||||
lora_ranks=torch.tensor(lora_ranks, dtype=torch.int32, device=device),
|
||||
scalings=torch.tensor(scalings, dtype=torch.float, device=device),
|
||||
permutation=perm,
|
||||
)
|
||||
|
||||
|
||||
def _check_close(
|
||||
a: torch.Tensor, b: torch.Tensor, name: str, atol: float = 1e-4, rtol: float = 1e-3
|
||||
) -> None:
|
||||
diff = (a - b).abs().max().item()
|
||||
assert torch.allclose(a, b, atol=atol, rtol=rtol), f"{name}: max diff = {diff}"
|
||||
|
||||
|
||||
def test_sgemm_lora_a():
|
||||
from sglang.srt.lora.triton_ops import sgemm_lora_a_fwd
|
||||
|
||||
torch.manual_seed(42)
|
||||
bs, input_dim, rank, num_loras = 8, 256, 16, 3
|
||||
x = torch.randn(bs, input_dim, device="cuda", dtype=torch.bfloat16)
|
||||
weights = torch.randn(
|
||||
num_loras, rank, input_dim, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
wi = [i % num_loras for i in range(bs)]
|
||||
lora_ranks = [rank] * num_loras
|
||||
scalings = [1.0] * num_loras
|
||||
|
||||
bi_plain = _make_batch_info(bs, wi, lora_ranks, scalings)
|
||||
bi_sorted = _make_sorted_batch_info(wi, lora_ranks, scalings, num_loras)
|
||||
|
||||
out_plain = sgemm_lora_a_fwd(x, weights, bi_plain)
|
||||
out_sorted = sgemm_lora_a_fwd(x, weights, bi_sorted)
|
||||
_check_close(out_plain, out_sorted, "sgemm_lora_a")
|
||||
|
||||
|
||||
def test_sgemm_lora_b():
|
||||
from sglang.srt.lora.triton_ops import sgemm_lora_b_fwd
|
||||
|
||||
torch.manual_seed(42)
|
||||
bs, output_dim, rank, num_loras = 8, 256, 16, 3
|
||||
x = torch.randn(bs, rank, device="cuda", dtype=torch.bfloat16)
|
||||
weights = torch.randn(
|
||||
num_loras, output_dim, rank, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
wi = [i % num_loras for i in range(bs)]
|
||||
lora_ranks = [rank] * num_loras
|
||||
scalings = [0.5] * num_loras
|
||||
|
||||
bi_plain = _make_batch_info(bs, wi, lora_ranks, scalings)
|
||||
bi_sorted = _make_sorted_batch_info(wi, lora_ranks, scalings, num_loras)
|
||||
|
||||
base_plain = torch.randn(bs, output_dim, device="cuda", dtype=torch.bfloat16)
|
||||
base_sorted = base_plain.clone()
|
||||
|
||||
out_plain = sgemm_lora_b_fwd(x, weights, bi_plain, base_plain)
|
||||
out_sorted = sgemm_lora_b_fwd(x, weights, bi_sorted, base_sorted)
|
||||
_check_close(out_plain, out_sorted, "sgemm_lora_b")
|
||||
|
||||
|
||||
def test_qkv_lora_b():
|
||||
from sglang.srt.lora.triton_ops import qkv_lora_b_fwd
|
||||
|
||||
torch.manual_seed(42)
|
||||
bs, rank, num_loras = 8, 16, 3
|
||||
n_slices = 3
|
||||
q_dim, kv_dim = 128, 64
|
||||
total_out = q_dim + 2 * kv_dim
|
||||
x = torch.randn(bs, n_slices * rank, device="cuda", dtype=torch.bfloat16)
|
||||
weights = torch.randn(
|
||||
num_loras, total_out, rank, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
output_offset = torch.tensor(
|
||||
[0, q_dim, q_dim + kv_dim, total_out], device="cuda", dtype=torch.int32
|
||||
)
|
||||
wi = [i % num_loras for i in range(bs)]
|
||||
lora_ranks = [rank] * num_loras
|
||||
scalings = [1.0] * num_loras
|
||||
|
||||
bi_plain = _make_batch_info(bs, wi, lora_ranks, scalings)
|
||||
bi_sorted = _make_sorted_batch_info(wi, lora_ranks, scalings, num_loras)
|
||||
|
||||
base_plain = torch.randn(bs, total_out, device="cuda", dtype=torch.bfloat16)
|
||||
base_sorted = base_plain.clone()
|
||||
|
||||
max_qkv_out_dim = max(q_dim, kv_dim)
|
||||
out_plain = qkv_lora_b_fwd(
|
||||
x, weights, bi_plain, output_offset, max_qkv_out_dim, base_plain
|
||||
)
|
||||
out_sorted = qkv_lora_b_fwd(
|
||||
x, weights, bi_sorted, output_offset, max_qkv_out_dim, base_sorted
|
||||
)
|
||||
_check_close(out_plain, out_sorted, "qkv_lora_b")
|
||||
|
||||
|
||||
def test_gate_up_lora_b():
|
||||
from sglang.srt.lora.triton_ops import gate_up_lora_b_fwd
|
||||
|
||||
torch.manual_seed(42)
|
||||
bs, rank, num_loras = 8, 16, 3
|
||||
output_dim = 128
|
||||
x = torch.randn(bs, 2 * rank, device="cuda", dtype=torch.bfloat16)
|
||||
weights = torch.randn(
|
||||
num_loras, 2 * output_dim, rank, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
wi = [i % num_loras for i in range(bs)]
|
||||
lora_ranks = [rank] * num_loras
|
||||
scalings = [1.0] * num_loras
|
||||
|
||||
bi_plain = _make_batch_info(bs, wi, lora_ranks, scalings)
|
||||
bi_sorted = _make_sorted_batch_info(wi, lora_ranks, scalings, num_loras)
|
||||
|
||||
base_plain = torch.randn(bs, 2 * output_dim, device="cuda", dtype=torch.bfloat16)
|
||||
base_sorted = base_plain.clone()
|
||||
|
||||
out_plain = gate_up_lora_b_fwd(x, weights, bi_plain, output_dim, base_plain)
|
||||
out_sorted = gate_up_lora_b_fwd(x, weights, bi_sorted, output_dim, base_sorted)
|
||||
_check_close(out_plain, out_sorted, "gate_up_lora_b")
|
||||
|
||||
|
||||
def test_mixed_ranks():
|
||||
"""Test with different LoRA ranks per adapter."""
|
||||
from sglang.srt.lora.triton_ops import sgemm_lora_a_fwd
|
||||
|
||||
torch.manual_seed(42)
|
||||
bs, input_dim, num_loras = 12, 256, 4
|
||||
max_rank = 32
|
||||
lora_ranks = [8, 16, 32, 16]
|
||||
scalings = [0.25, 0.5, 1.0, 2.0]
|
||||
# Use max_rank for weight shape, kernel handles per-adapter rank
|
||||
weights = torch.randn(
|
||||
num_loras, max_rank, input_dim, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
x = torch.randn(bs, input_dim, device="cuda", dtype=torch.bfloat16)
|
||||
wi = [i % num_loras for i in range(bs)]
|
||||
|
||||
bi_plain = _make_batch_info(bs, wi, lora_ranks, scalings)
|
||||
bi_sorted = _make_sorted_batch_info(wi, lora_ranks, scalings, num_loras)
|
||||
|
||||
out_plain = sgemm_lora_a_fwd(x, weights, bi_plain)
|
||||
out_sorted = sgemm_lora_a_fwd(x, weights, bi_sorted)
|
||||
_check_close(out_plain, out_sorted, "sgemm_lora_a_mixed_ranks")
|
||||
|
||||
|
||||
def test_single_adapter():
|
||||
"""All sequences use the same adapter."""
|
||||
from sglang.srt.lora.triton_ops import sgemm_lora_a_fwd
|
||||
|
||||
torch.manual_seed(42)
|
||||
bs, input_dim, rank, num_loras = 16, 256, 16, 2
|
||||
x = torch.randn(bs, input_dim, device="cuda", dtype=torch.bfloat16)
|
||||
weights = torch.randn(
|
||||
num_loras, rank, input_dim, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
wi = [0] * bs # all adapter 0
|
||||
lora_ranks = [rank, rank]
|
||||
scalings = [1.0, 1.0]
|
||||
|
||||
bi_plain = _make_batch_info(bs, wi, lora_ranks, scalings)
|
||||
bi_sorted = _make_sorted_batch_info(wi, lora_ranks, scalings, num_loras)
|
||||
|
||||
out_plain = sgemm_lora_a_fwd(x, weights, bi_plain)
|
||||
out_sorted = sgemm_lora_a_fwd(x, weights, bi_sorted)
|
||||
_check_close(out_plain, out_sorted, "sgemm_lora_a_single_adapter")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
@@ -0,0 +1,90 @@
|
||||
import ast
|
||||
import pathlib
|
||||
import unittest
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=5, suite="stage-a-test-cpu")
|
||||
|
||||
|
||||
_REPO_ROOT = pathlib.Path(__file__).resolve().parents[3]
|
||||
_SCAN_ROOTS = [_REPO_ROOT / "python", _REPO_ROOT / "test"]
|
||||
|
||||
|
||||
class TestNoBarePytestMain(CustomTestCase):
|
||||
def test_no_bare_pytest_main_in_repo(self):
|
||||
offenders = []
|
||||
for root in _SCAN_ROOTS:
|
||||
if not root.exists():
|
||||
continue
|
||||
for path in root.rglob("*.py"):
|
||||
violation = _find_bare_pytest_main(path)
|
||||
if violation is not None:
|
||||
offenders.append(violation)
|
||||
|
||||
self.assertFalse(
|
||||
offenders,
|
||||
msg=(
|
||||
"Found bare `pytest.main(...)` in __main__ blocks (must be "
|
||||
"wrapped in sys.exit(...) so failing tests propagate the exit "
|
||||
"code to the CI runner):\n " + "\n ".join(offenders)
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _find_bare_pytest_main(path: pathlib.Path):
|
||||
"""Return `<rel_path>:<lineno>` if `path` has a bare pytest.main(...) call
|
||||
inside `if __name__ == "__main__":`, else None."""
|
||||
try:
|
||||
source = path.read_text(encoding="utf-8")
|
||||
except (OSError, UnicodeDecodeError):
|
||||
return None
|
||||
try:
|
||||
tree = ast.parse(source, filename=str(path))
|
||||
except SyntaxError:
|
||||
return None
|
||||
|
||||
for node in ast.walk(tree):
|
||||
if not isinstance(node, ast.If):
|
||||
continue
|
||||
if not _is_main_guard(node.test):
|
||||
continue
|
||||
for stmt in node.body:
|
||||
if _is_bare_pytest_main_call(stmt):
|
||||
rel = path.relative_to(_REPO_ROOT)
|
||||
return f"{rel}:{stmt.lineno}"
|
||||
return None
|
||||
|
||||
|
||||
def _is_main_guard(test: ast.expr) -> bool:
|
||||
"""Match `__name__ == "__main__"` (either side)."""
|
||||
if not isinstance(test, ast.Compare) or len(test.ops) != 1:
|
||||
return False
|
||||
if not isinstance(test.ops[0], ast.Eq):
|
||||
return False
|
||||
sides = [test.left, *test.comparators]
|
||||
has_name = any(isinstance(s, ast.Name) and s.id == "__name__" for s in sides)
|
||||
has_main = any(isinstance(s, ast.Constant) and s.value == "__main__" for s in sides)
|
||||
return has_name and has_main
|
||||
|
||||
|
||||
def _is_bare_pytest_main_call(stmt: ast.stmt) -> bool:
|
||||
"""Match `pytest.main(...)` whose return value is discarded.
|
||||
`sys.exit(pytest.main(...))` and `code = pytest.main(...)` are fine."""
|
||||
if not isinstance(stmt, ast.Expr):
|
||||
return False
|
||||
call = stmt.value
|
||||
if not isinstance(call, ast.Call):
|
||||
return False
|
||||
func = call.func
|
||||
return (
|
||||
isinstance(func, ast.Attribute)
|
||||
and func.attr == "main"
|
||||
and isinstance(func.value, ast.Name)
|
||||
and func.value.id == "pytest"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user