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