151 lines
4.8 KiB
Python
151 lines
4.8 KiB
Python
import itertools
|
||
import unittest
|
||
|
||
import torch
|
||
import torch.nn.functional as F
|
||
|
||
from sglang.test.ci.ci_register import register_amd_ci
|
||
from sglang.test.test_utils import CustomTestCase
|
||
|
||
register_amd_ci(est_time=10, suite="stage-a-test-1-amd")
|
||
|
||
|
||
def _fp8_available() -> bool:
|
||
# requirement：1) GPU；2) ROCm；3) torch support float8_e4m3fn
|
||
if not torch.cuda.is_available():
|
||
return False
|
||
if getattr(torch.version, "hip", None) is None:
|
||
return False
|
||
return hasattr(torch, "float8_e4m3fn")
|
||
|
||
|
||
def _rmsnorm(x, weight, eps=1e-6):
|
||
# row-wise RMSNorm
|
||
row_norm = (x * x).sum(dim=-1)
|
||
norm = torch.rsqrt(row_norm / x.shape[1] + eps)
|
||
return x * norm[:, None] * weight[None, :]
|
||
|
||
|
||
def _per_token_fp8_group_quant(x, dtype_quant, group_size=128):
|
||
"""per token、group-size quant, return (quantized, scale)。"""
|
||
DTYPE_MAX = torch.finfo(dtype_quant).max
|
||
M, N = x.shape
|
||
|
||
pad = (group_size - (N % group_size)) % group_size
|
||
if pad:
|
||
x_reshape = F.pad(x, (0, pad, 0, 0), "constant", 0)
|
||
else:
|
||
x_reshape = x
|
||
|
||
G = (N + group_size - 1) // group_size
|
||
x_reshape = x_reshape.view(M, G, group_size).to(torch.float32)
|
||
x_max = torch.max(torch.abs(x_reshape), dim=-1, keepdim=True)[0].clamp_min_(1e-10)
|
||
x_scale = x_max / DTYPE_MAX
|
||
inv = 1.0 / x_scale
|
||
|
||
x_q = torch.clamp(x_reshape * inv, -DTYPE_MAX, DTYPE_MAX).to(dtype_quant)
|
||
x_q = x_q.view(M, G * group_size)
|
||
if pad:
|
||
x_q = x_q[:, :N]
|
||
x_scale = x_scale.squeeze(-1) # [M, G]
|
||
return x_q, x_scale
|
||
|
||
|
||
def _upcast_fp8_group(x_q, x_s, out_dtype=torch.float32, group_size=128):
|
||
"""unqaunt"""
|
||
M, N = x_q.shape
|
||
G = (N + group_size - 1) // group_size
|
||
pad = (group_size - (N % group_size)) % group_size
|
||
|
||
if pad:
|
||
x_q = F.pad(x_q, (0, pad, 0, 0), "constant", 0)
|
||
|
||
x_q = x_q.view(M, G, group_size).to(torch.float32)
|
||
x = x_q * x_s.view(M, G, 1)
|
||
x = x.view(M, G * group_size)[:, :N]
|
||
return x.to(out_dtype)
|
||
|
||
|
||
class TestFusedRMSFP8GroupQuant(CustomTestCase):
|
||
#
|
||
DTYPES = [torch.bfloat16, torch.float16]
|
||
# (M, N1, N2)
|
||
SHAPES = [(32, 128, 7168), (128, 7168, 7168)]
|
||
GROUP_SIZE = [128]
|
||
SEEDS = [0]
|
||
|
||
@classmethod
|
||
def setUpClass(cls):
|
||
if not _fp8_available():
|
||
raise unittest.SkipTest("Skip: ROCm/FP8 is not available")
|
||
torch.set_default_device("cuda")
|
||
|
||
def _run_ref(self, x1, w1, eps1, x2, w2, eps2, res1, dtype_quant, group_size):
|
||
s = x1 + (res1 if res1 is not None else 0)
|
||
y1 = _rmsnorm(s, w1, eps1)
|
||
y2 = _rmsnorm(x2, w2, eps2) if x2 is not None else None
|
||
y1_q, y1_s = _per_token_fp8_group_quant(y1, dtype_quant, group_size)
|
||
return (
|
||
(y1_q, y1_s),
|
||
y1.to(x1.dtype),
|
||
(y2.to(x1.dtype) if y2 is not None else None),
|
||
(s.to(x1.dtype) if res1 is not None else None),
|
||
)
|
||
|
||
def _case(self, M, N1, N2, group_size, dtype, seed):
|
||
torch.manual_seed(seed)
|
||
fp8 = torch.float8_e4m3fn
|
||
device = "cuda"
|
||
|
||
x1 = torch.randn(M, N1, dtype=dtype, device=device) / 10
|
||
x2 = torch.randn(M, N2, dtype=dtype, device=device) / 10
|
||
w1 = torch.ones(N1, dtype=torch.float32, device=device)
|
||
w2 = torch.ones(N2, dtype=torch.float32, device=device)
|
||
res1 = torch.randn(M, N1, dtype=dtype, device=device) / 10
|
||
|
||
# ref
|
||
(y1_q_ref, y1_s_ref), y1_ref, y2_ref, s_ref = self._run_ref(
|
||
x1, w1, 1e-6, x2, w2, 1e-6, res1, fp8, group_size
|
||
)
|
||
|
||
# be tested：aiter fused op
|
||
from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant
|
||
|
||
(y1_q, y1_s), y1, y2, s = fused_rms_fp8_group_quant(
|
||
x1,
|
||
w1,
|
||
1e-6,
|
||
inp2=x2,
|
||
inp2_weight=w2,
|
||
inp2_epsilon=1e-6,
|
||
group_size=group_size,
|
||
dtype_quant=fp8,
|
||
res1=res1,
|
||
output_unquantized_inp1=True, # get unqaunt y1
|
||
)
|
||
|
||
torch.testing.assert_close(y1, y1_ref, atol=0.1, rtol=0.1)
|
||
torch.testing.assert_close(y2, y2_ref, atol=0.1, rtol=0.1)
|
||
torch.testing.assert_close(s, s_ref, atol=0.1, rtol=0.1)
|
||
|
||
# check unquant
|
||
y1_up_ref = _upcast_fp8_group(
|
||
y1_q_ref, y1_s_ref, out_dtype=torch.float32, group_size=group_size
|
||
)
|
||
y1_up = _upcast_fp8_group(
|
||
y1_q, y1_s, out_dtype=torch.float32, group_size=group_size
|
||
)
|
||
torch.testing.assert_close(y1_up, y1_up_ref, atol=0.1, rtol=0.1)
|
||
|
||
def test_fused_rms_fp8_group_quant(self):
|
||
for params in itertools.product(
|
||
self.SHAPES, self.GROUP_SIZE, self.DTYPES, self.SEEDS
|
||
):
|
||
(M, N1, N2), g, dtype, seed = params
|
||
with self.subTest(M=M, N1=N1, N2=N2, group_size=g, dtype=dtype, seed=seed):
|
||
self._case(M, N1, N2, g, dtype, seed)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main(verbosity=2)
|