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"])
|
||||
Reference in New Issue
Block a user