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

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

Browse Source
This commit is contained in:
Liangsheng Yin
2026-05-14 20:02:44 -07:00
committed by GitHub
parent 8d5b347edd
commit d89b678d69
29 changed files with 0 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/manual/quant/test_block_fp8.py
+760
View File
@@ -0,0 +1,760 @@
import itertools
import unittest
from functools import lru_cache
import torch
from sglang.srt.layers.activation import SiluAndMul
from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import fused_moe
from sglang.srt.layers.moe.topk import TopKConfig, select_experts
from sglang.srt.layers.quantization.fp8_kernel import (
per_tensor_quant_mla_fp8,
per_token_group_quant_fp8,
per_token_group_quant_mla_deep_gemm_masked_fp8,
static_quant_fp8,
w8a8_block_fp8_matmul,
)
from sglang.srt.layers.quantization.fp8_utils import (
input_to_float8,
mxfp8_group_quantize,
triton_mxfp8_blockscaled_linear,
)
from sglang.srt.utils import is_sm100_supported, is_sm120_supported
from sglang.test.test_utils import CustomTestCase
_is_cuda = torch.cuda.is_available() and torch.version.cuda
# For test
@lru_cache(maxsize=1)
def _get_triton_mxfp8_upcast():
try:
from triton_kernels.numerics_details.mxfp import upcast_from_mxfp_torch
except Exception as err:
raise RuntimeError(
"MXFP8 dequantization requires triton_kernels with MXFP8 support."
) from err
return upcast_from_mxfp_torch
# For test
def native_per_token_group_quant_fp8(
x, group_size, eps=1e-10, dtype=torch.float8_e4m3fn
):
"""Function to perform per-token-group quantization on an input tensor `x` using native torch.
It converts the tensor values into float8 values and returns the
quantized tensor along with the scaling factor used for quantization.
Note that only `torch.float8_e4m3fn` is supported for now.
"""
assert (
x.shape[-1] % group_size == 0
), "the last dimension of `x` cannot be divisible by `group_size`"
assert x.is_contiguous(), "`x` is not contiguous"
finfo = torch.finfo(dtype)
fp8_min = finfo.min
fp8_max = finfo.max
x_ = x.reshape(x.numel() // group_size, group_size)
amax = x_.abs().max(dim=-1, keepdim=True)[0].clamp(min=eps).to(torch.float32)
x_s = amax / fp8_max
x_q = (x_ / x_s).clamp(min=fp8_min, max=fp8_max).to(dtype)
x_q = x_q.reshape(x.shape)
x_s = x_s.reshape(x.shape[:-1] + (x.shape[-1] // group_size,))
return x_q, x_s
class TestPerTokenGroupQuantFP8(CustomTestCase):
DTYPES = [torch.half, torch.bfloat16, torch.float32]
NUM_TOKENS = [7, 83, 2048]
D = [512, 4096, 5120, 13824]
GROUP_SIZE = [64, 128, 256, 512]
SEEDS = [0]
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
torch.set_default_device("cuda")
def _per_token_group_quant_fp8(self, num_tokens, d, dtype, group_size, seed):
torch.manual_seed(seed)
x = torch.rand(num_tokens, d, dtype=dtype)
with torch.inference_mode():
ref_out, ref_scale = native_per_token_group_quant_fp8(x, group_size)
out, scale = per_token_group_quant_fp8(x, group_size)
self.assertTrue(
torch.allclose(out.to(torch.float32), ref_out.to(torch.float32), rtol=0.20)
)
self.assertTrue(torch.allclose(scale, ref_scale))
def test_per_token_group_quant_fp8(self):
for params in itertools.product(
self.NUM_TOKENS,
self.D,
self.DTYPES,
self.GROUP_SIZE,
self.SEEDS,
):
with self.subTest(
num_tokens=params[0],
d=params[1],
dtype=params[2],
group_size=params[3],
seed=params[4],
):
self._per_token_group_quant_fp8(*params)
# For test
def native_static_quant_fp8(x, x_s, dtype=torch.float8_e4m3fn):
"""Function to perform static quantization on an input tensor `x` using native torch.
It converts the tensor values into float8 values and returns the
quantized tensor along with the scaling factor used for quantization.
"""
assert x.is_contiguous(), "`x` is not contiguous"
assert x_s.numel() == 1, "only supports per-tensor scale"
finfo = torch.finfo(dtype)
fp8_min = finfo.min
fp8_max = finfo.max
x_ = x.reshape(x.numel() // x.shape[-1], x.shape[-1])
x_s_inv = 1.0 / x_s
x_q = (x_ * x_s_inv).clamp(min=fp8_min, max=fp8_max).to(dtype)
x_q = x_q.reshape(x.shape)
return x_q, x_s
class TestStaticQuantFP8(CustomTestCase):
DTYPES = [torch.half, torch.bfloat16, torch.float32]
NUM_TOKENS = [7, 83, 2048]
D = [512, 4096, 5120, 13824]
SEEDS = [0]
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
torch.set_default_device("cuda")
def _static_quant_fp8(self, num_tokens, d, dtype, seed):
torch.manual_seed(seed)
x = torch.rand(num_tokens, d, dtype=dtype)
fp8_max = torch.finfo(torch.float8_e4m3fn).max
x_s = x.max() / fp8_max
with torch.inference_mode():
ref_out, _ = native_static_quant_fp8(x, x_s)
out, _ = static_quant_fp8(x, x_s, repeat_scale=True)
self.assertTrue(
torch.allclose(out.to(torch.float32), ref_out.to(torch.float32), rtol=0.50)
)
def test_static_quant_fp8(self):
for params in itertools.product(
self.NUM_TOKENS,
self.D,
self.DTYPES,
self.SEEDS,
):
with self.subTest(
num_tokens=params[0],
d=params[1],
dtype=params[2],
seed=params[3],
):
self._static_quant_fp8(*params)
class TestPerTensorQuantMlaFP8(CustomTestCase):
DTYPES = [torch.half, torch.bfloat16, torch.float32]
NUM_TOKENS = [7, 83, 2048]
D = [512, 4096, 5120, 13824]
LAST_D_EXT = [1024, 0]
LAST_D = [512]
SEEDS = [0]
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
torch.set_default_device("cuda")
def _per_tensor_quant_mla_fp8(self, num_tokens, d, last_d_ext, last_d, dtype, seed):
torch.manual_seed(seed)
x = torch.rand(
(num_tokens, d // last_d, last_d + last_d_ext),
dtype=dtype,
)
x_sub, _ = x.split([last_d, last_d_ext], dim=-1)
with torch.inference_mode():
ref_out, ref_s = input_to_float8(x_sub.transpose(0, 1))
out, out_s = per_tensor_quant_mla_fp8(x_sub.transpose(0, 1))
self.assertTrue(out.is_contiguous())
self.assertTrue(
torch.allclose(out.to(torch.float32), ref_out.to(torch.float32), rtol=0.50)
)
self.assertTrue(
torch.allclose(out_s.to(torch.float32), ref_s.to(torch.float32))
)
def test_per_tensor_quant_mla_fp8(self):
for params in itertools.product(
self.NUM_TOKENS,
self.D,
self.LAST_D_EXT,
self.LAST_D,
self.DTYPES,
self.SEEDS,
):
with self.subTest(
num_tokens=params[0],
d=params[1],
last_d_ext=params[2],
last_d=params[3],
dtype=params[4],
seed=params[5],
):
self._per_tensor_quant_mla_fp8(*params)
class TestPerTokenGroupQuantMlaDeepGemmMaskedFP8(CustomTestCase):
DTYPES = [torch.half, torch.bfloat16, torch.float32]
B = [128]
NUM_TOKENS = [7, 83, 2048, 1024 * 16]
D = [512, 128]
GROUP_SIZE = [128]
SEEDS = [0]
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
torch.set_default_device("cuda")
def _per_token_group_quant_mla_deep_gemm_masked_fp8(
self, b, num_tokens, d, dtype, group_size, seed
):
torch.manual_seed(seed)
x = torch.rand(b, num_tokens, d, dtype=dtype)
with torch.inference_mode():
ref_out, ref_scale = native_per_token_group_quant_fp8(x, group_size, 1e-12)
out, scale, _, _, _ = per_token_group_quant_mla_deep_gemm_masked_fp8(
x, group_size
)
out = out[:, :num_tokens, :]
scale = scale[:, :num_tokens, :]
self.assertTrue(
torch.allclose(
out.to(torch.float32), ref_out.to(torch.float32), rtol=0.20, atol=1e-2
)
)
self.assertTrue(torch.allclose(scale, ref_scale))
def test_per_token_group_quant_mla_deep_gemm_masked_fp8(self):
for params in itertools.product(
self.B,
self.NUM_TOKENS,
self.D,
self.DTYPES,
self.GROUP_SIZE,
self.SEEDS,
):
with self.subTest(
b=params[0],
num_tokens=params[1],
d=params[2],
dtype=params[3],
group_size=params[4],
seed=params[5],
):
self._per_token_group_quant_mla_deep_gemm_masked_fp8(*params)
# For test
def native_w8a8_block_fp8_matmul(A, B, As, Bs, block_size, output_dtype=torch.float16):
"""This function performs matrix multiplication with block-wise quantization using native torch.
It takes two input tensors `A` and `B` with scales `As` and `Bs`.
The output is returned in the specified `output_dtype`.
"""
A = A.to(torch.float32)
B = B.to(torch.float32)
assert A.shape[-1] == B.shape[-1]
assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2
assert len(block_size) == 2
block_n, block_k = block_size[0], block_size[1]
assert (A.shape[-1] + block_k - 1) // block_k == As.shape[-1]
assert A.shape[:-1] == As.shape[:-1]
M = A.numel() // A.shape[-1]
N, K = B.shape
origin_C_shape = A.shape[:-1] + (N,)
A = A.reshape(M, A.shape[-1])
As = As.reshape(M, As.shape[-1])
n_tiles = (N + block_n - 1) // block_n
k_tiles = (K + block_k - 1) // block_k
assert n_tiles == Bs.shape[0]
assert k_tiles == Bs.shape[1]
C_shape = (M, N)
C = torch.zeros(C_shape, dtype=torch.float32, device=A.device)
A_tiles = [A[:, i * block_k : min((i + 1) * block_k, K)] for i in range(k_tiles)]
B_tiles = [
[
B[
j * block_n : min((j + 1) * block_n, N),
i * block_k : min((i + 1) * block_k, K),
]
for i in range(k_tiles)
]
for j in range(n_tiles)
]
C_tiles = [C[:, j * block_n : min((j + 1) * block_n, N)] for j in range(n_tiles)]
As_tiles = [As[:, i : i + 1] for i in range(k_tiles)]
for i in range(k_tiles):
for j in range(n_tiles):
a = A_tiles[i]
b = B_tiles[j][i]
c = C_tiles[j]
s = As_tiles[i] * Bs[j][i]
c[:, :] += torch.matmul(a, b.t()) * s
C = C.reshape(origin_C_shape).to(output_dtype)
return C
class TestW8A8BlockFP8Matmul(CustomTestCase):
if not _is_cuda:
OUT_DTYPES = [torch.float32, torch.half, torch.bfloat16]
M = [1, 7, 83, 512, 2048]
NKs = [
(N, K)
for N in [128, 512, 1024, 4096, 7748, 13824]
for K in [256, 4096, 5120, 3884, 13824]
]
# BLOCK_SIZE = [[64, 64], [64, 128], [128, 64], [128, 128]]
BLOCK_SIZE = [[128, 128]]
SEEDS = [0]
else:
# use practical shape in DeepSeek V3 for test
OUT_DTYPES = [torch.bfloat16]
M = [64, 128, 512, 1024, 4096]
NKs = [
(2112, 7168),
(1536, 7168),
(3072, 1536),
(24576, 7168),
(4096, 512),
(7168, 2048),
(4608, 7168),
(512, 7168),
(7168, 2304),
(7168, 512),
]
BLOCK_SIZE = [[128, 128]]
SEEDS = [0]
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
torch.set_default_device("cuda")
def _w8a8_block_fp8_matmul(self, M, NK, block_size, out_dtype, seed):
N, K = NK
torch.manual_seed(seed)
# NOTE(HandH1998): to avoid overflow when out_dtype = torch.half
factor_for_scale = 1e-2
fp8_info = torch.finfo(torch.float8_e4m3fn)
fp8_max, fp8_min = fp8_info.max, fp8_info.min
A_fp32 = (torch.rand(M, K, dtype=torch.float32) - 0.5) * 2 * fp8_max
A_fp8 = A_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
B_fp32 = (torch.rand(N, K, dtype=torch.float32) - 0.5) * 2 * fp8_max
B_fp8 = B_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
block_n, block_k = block_size[0], block_size[1]
n_tiles = (N + block_n - 1) // block_n
k_tiles = (K + block_k - 1) // block_k
As = torch.rand(M, k_tiles, dtype=torch.float32) * factor_for_scale
Bs = torch.rand(n_tiles, k_tiles, dtype=torch.float32) * factor_for_scale
with torch.inference_mode():
ref_out = native_w8a8_block_fp8_matmul(
A_fp8, B_fp8, As, Bs, block_size, out_dtype
)
out = w8a8_block_fp8_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype)
self.assertTrue(
torch.mean(torch.abs(out.to(torch.float32) - ref_out.to(torch.float32)))
/ torch.mean(torch.abs(ref_out.to(torch.float32)))
< 0.001
)
def test_w8a8_block_fp8_matmul(self):
for params in itertools.product(
self.M,
self.NKs,
self.BLOCK_SIZE,
self.OUT_DTYPES,
self.SEEDS,
):
with self.subTest(
M=params[0],
NKs=params[1],
block_size=params[2],
out_dtype=params[3],
seed=params[4],
):
self._w8a8_block_fp8_matmul(*params)
def _mxfp8_group_dequant(q: torch.Tensor, scale_u8: torch.Tensor) -> torch.Tensor:
upcast_from_mxfp_torch = _get_triton_mxfp8_upcast()
return upcast_from_mxfp_torch(q, scale_u8, torch.float32, axis=1)
class TestMXFP8DenseLinear(CustomTestCase):
DTYPES = [torch.bfloat16]
M = [1, 127, 128, 129, 255, 256]
NKs = [
(256, 512),
(384, 1024),
(512, 2048),
(768, 1024),
]
SEEDS = [0]
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
if not (is_sm100_supported() or is_sm120_supported()):
raise unittest.SkipTest("MXFP8 requires Blackwell (SM100/SM120)")
torch.set_default_device("cuda")
def _mxfp8_dense_linear(self, M, NK, dtype, seed):
N, K = NK
torch.manual_seed(seed)
input_fp32 = torch.randn((M, K), dtype=torch.float32) / 4
input_fp16 = input_fp32.to(dtype)
weight_fp32 = torch.randn((N, K), dtype=torch.float32) / 4
weight_q, weight_scale_u8 = mxfp8_group_quantize(weight_fp32)
with torch.inference_mode():
q_input, input_scale_u8 = mxfp8_group_quantize(input_fp16.to(torch.float32))
a_dq = _mxfp8_group_dequant(q_input, input_scale_u8)
b_dq = _mxfp8_group_dequant(weight_q, weight_scale_u8)
ref_out = torch.matmul(a_dq, b_dq.t()).to(dtype)
out = triton_mxfp8_blockscaled_linear(
input=input_fp16,
weight=weight_q,
weight_scale=weight_scale_u8,
)
out_prequant = triton_mxfp8_blockscaled_linear(
input=q_input,
weight=weight_q,
weight_scale=weight_scale_u8,
input_scale=input_scale_u8,
output_dtype=dtype,
)
self.assertTrue(
torch.mean(torch.abs(out.to(torch.float32) - ref_out.to(torch.float32)))
/ torch.mean(torch.abs(ref_out.to(torch.float32)))
< 0.02
)
self.assertTrue(
torch.mean(
torch.abs(out_prequant.to(torch.float32) - ref_out.to(torch.float32))
)
/ torch.mean(torch.abs(ref_out.to(torch.float32)))
< 0.02
)
def test_mxfp8_dense_linear(self):
for params in itertools.product(
self.M,
self.NKs,
self.DTYPES,
self.SEEDS,
):
with self.subTest(
M=params[0],
NKs=params[1],
dtype=params[2],
seed=params[3],
):
self._mxfp8_dense_linear(*params)
# For test
def torch_w8a8_block_fp8_moe(a, w1, w2, w1_s, w2_s, score, topk, block_shape):
"""This function performs fused moe with block-wise quantization using native torch."""
B, D = a.shape
a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
score = torch.softmax(score, dim=-1, dtype=torch.float32)
topk_weight, topk_ids = torch.topk(score, topk)
topk_weight = topk_weight.view(-1)
topk_ids = topk_ids.view(-1)
_, block_k = block_shape[0], block_shape[1]
a_q, a_s = native_per_token_group_quant_fp8(a, block_k)
# NOTE(HandH1998): Since "index_cuda" not implemented for 'Float8_e4m3fn', we need to cast `float8`` to `float32``.
a_q = a_q.to(torch.float32)
for i in range(w1.shape[0]):
mask = topk_ids == i
if mask.sum():
inter_out = native_w8a8_block_fp8_matmul(
a_q[mask], w1[i], a_s[mask], w1_s[i], block_shape, output_dtype=a.dtype
)
act_out = SiluAndMul().forward_native(inter_out)
act_out_q, act_out_s = native_per_token_group_quant_fp8(act_out, block_k)
act_out = act_out.to(torch.float32)
out[mask] = native_w8a8_block_fp8_matmul(
act_out_q, w2[i], act_out_s, w2_s[i], block_shape, output_dtype=a.dtype
)
return (
out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype)
).sum(dim=1)
class TestW8A8BlockFP8FusedMoE(CustomTestCase):
DTYPES = [torch.float32, torch.half, torch.bfloat16]
M = [1, 33, 64, 222, 1024 * 128]
N = [128, 1024, 2048]
K = [256, 4096, 5120]
E = [8, 24]
TOP_KS = [2, 6]
BLOCK_SIZE = [[64, 64], [64, 128], [128, 64], [128, 128]]
# BLOCK_SIZE = [[128, 128]]
SEEDS = [0]
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
torch.set_default_device("cuda")
def _w8a8_block_fp8_fused_moe(self, M, N, K, E, topk, block_size, dtype, seed):
torch.manual_seed(seed)
# NOTE(HandH1998): to avoid overflow when out_dtype = torch.half
factor_for_scale = 1e-2
fp8_info = torch.finfo(torch.float8_e4m3fn)
fp8_max, fp8_min = fp8_info.max, fp8_info.min
a = torch.randn((M, K), dtype=dtype) / 10
w1_fp32 = (torch.rand((E, 2 * N, K), dtype=torch.float32) - 0.5) * 2 * fp8_max
w1 = w1_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
w2_fp32 = (torch.rand((E, K, N), dtype=torch.float32) - 0.5) * 2 * fp8_max
w2 = w2_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
block_n, block_k = block_size[0], block_size[1]
n_tiles_w1 = (2 * N + block_n - 1) // block_n
n_tiles_w2 = (K + block_n - 1) // block_n
k_tiles_w1 = (K + block_k - 1) // block_k
k_tiles_w2 = (N + block_k - 1) // block_k
w1_s = (
torch.rand((E, n_tiles_w1, k_tiles_w1), dtype=torch.float32)
* factor_for_scale
)
w2_s = (
torch.rand((E, n_tiles_w2, k_tiles_w2), dtype=torch.float32)
* factor_for_scale
)
score = torch.randn((M, E), dtype=dtype)
with torch.inference_mode():
ref_out = torch_w8a8_block_fp8_moe(
a, w1, w2, w1_s, w2_s, score, topk, block_size
)
topk_output = select_experts(
hidden_states=a,
router_logits=score,
topk_config=TopKConfig(top_k=topk, renormalize=False),
)
out = fused_moe(
a,
w1,
w2,
topk_output,
use_fp8_w8a8=True,
w1_scale=w1_s,
w2_scale=w2_s,
block_shape=block_size,
)
self.assertTrue(
torch.mean(torch.abs(out.to(torch.float32) - ref_out.to(torch.float32)))
/ torch.mean(torch.abs(ref_out.to(torch.float32)))
< 0.02
)
def test_w8a8_block_fp8_fused_moe(self):
for params in itertools.product(
self.M,
self.N,
self.K,
self.E,
self.TOP_KS,
self.BLOCK_SIZE,
self.DTYPES,
self.SEEDS,
):
with self.subTest(
M=params[0],
N=params[1],
K=params[2],
E=params[3],
topk=params[4],
block_size=params[5],
dtype=params[6],
seed=params[7],
):
self._w8a8_block_fp8_fused_moe(*params)
# For test
def torch_w8a8_block_fp8_bmm(a, a_s, w, w_s, block_shape, out_dtype):
"""This function performs bmm with block-wise quantization using native torch."""
B, N, _ = w.shape
_, M, _ = a.shape
out = torch.empty((B, M, N), dtype=out_dtype, device=a.device)
for i in range(B):
out[i] = native_w8a8_block_fp8_matmul(
a[i], w[i], a_s[i], w_s[i], block_shape, output_dtype=out_dtype
)
return out
class TestW8A8BlockFP8BatchedDeepGemm(CustomTestCase):
DTYPES = [torch.bfloat16]
M = [1, 33, 64, 222, 8192]
N = [128, 512]
K = [128, 512]
BATCH = [128]
BLOCK_SIZE = [[128, 128]]
SEEDS = [0]
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
try:
import deep_gemm # noqa: F401
except ImportError:
raise unittest.SkipTest("DeepGEMM is not available")
torch.set_default_device("cuda")
def _w8a8_block_fp8_batched_deep_gemm(self, M, N, K, B, block_size, dtype, seed):
torch.manual_seed(seed)
factor_for_scale = 1e-2
fp8_info = torch.finfo(torch.float8_e4m3fn)
fp8_max, fp8_min = fp8_info.max, fp8_info.min
a_fp32 = torch.randn((B, M, K), dtype=torch.float32) / 10
a = a_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
w_fp32 = (torch.rand((B, N, K), dtype=torch.float32) - 0.5) * 2 * fp8_max
w = w_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
block_n, block_k = block_size[0], block_size[1]
n_tiles_w = (N + block_n - 1) // block_n
k_tiles_w = (K + block_k - 1) // block_k
w_s = (
torch.rand((B, n_tiles_w, k_tiles_w), dtype=torch.float32)
* factor_for_scale
)
a_s = torch.rand((B, M, k_tiles_w), dtype=torch.float32) * factor_for_scale
ae = a.new_empty(B, (M + 255) // 256 * 256, K)
ae_s = a_s.new_empty(B, (M + 255) // 256 * 256, k_tiles_w)
oe = torch.empty((B, (M + 255) // 256 * 256, N), dtype=dtype)
ae[:, :M, :] = a
ae_s[:, :M, :] = a_s
masked_m = torch.full((B,), M, dtype=torch.int)
expected_m = M
lhs = (
ae,
ae_s,
)
rhs = (
w,
w_s,
)
from deep_gemm import fp8_m_grouped_gemm_nt_masked
with torch.inference_mode():
ref_out = torch_w8a8_block_fp8_bmm(a, a_s, w, w_s, block_size, dtype)
fp8_m_grouped_gemm_nt_masked(lhs, rhs, oe, masked_m, expected_m)
out = oe[:, :M, :]
self.assertTrue(
torch.mean(torch.abs(out.to(torch.float32) - ref_out.to(torch.float32)))
/ torch.mean(torch.abs(ref_out.to(torch.float32)))
< 0.0001
)
def test_w8a8_block_fp8_batched_deep_gemm(self):
for params in itertools.product(
self.M,
self.N,
self.K,
self.BATCH,
self.BLOCK_SIZE,
self.DTYPES,
self.SEEDS,
):
with self.subTest(
M=params[0],
N=params[1],
K=params[2],
B=params[3],
block_size=params[4],
dtype=params[5],
seed=params[6],
):
self._w8a8_block_fp8_batched_deep_gemm(*params)
if __name__ == "__main__":
unittest.main(verbosity=2)
test/manual/quant/test_block_fp8_deep_gemm_blackwell.py
+251
View File
@@ -0,0 +1,251 @@
import itertools
import unittest
from typing import List, Tuple
import torch
from deep_gemm import fp8_gemm_nt
from sglang.test.test_utils import CustomTestCase
_is_cuda = torch.cuda.is_available() and torch.version.cuda
# Modify form DeepGEMM Blackwell
def ceil_div(x: int, y: int) -> int:
return (x + y - 1) // y
def align(x: int, y: int) -> int:
return ceil_div(x, y) * y
def per_token_group_quant_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
assert x.dim() == 2 and x.size(1) % 128 == 0
m, n = x.shape
x_view = x.view(m, -1, 128)
x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4)
sf = x_amax / 448.0
return (x_view * (1.0 / sf.unsqueeze(2))).to(torch.float8_e4m3fn).view(m, n), sf
def per_block_quant_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
assert x.dim() == 2
m, n = x.shape
x_padded = torch.zeros(
(align(m, 128), align(n, 128)), dtype=x.dtype, device=x.device
)
x_padded[:m, :n] = x
x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, 128)
x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
sf = x_amax / 448.0
x_scaled = (x_view * (1.0 / sf)).to(torch.float8_e4m3fn)
return x_scaled.view_as(x_padded)[:m, :n].contiguous(), sf.view(
x_view.size(0), x_view.size(2)
)
def ceil_to_ue8m0(x: torch.Tensor):
assert x.view(-1).amax().item() > 0
return torch.pow(2.0, torch.ceil(torch.log2(x.abs())))
def per_token_group_quant_mxfp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
assert x.dim() == 2 and x.size(1) % 128 == 0
m, n = x.shape
x_view = x.view(m, -1, 128)
x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4)
sf = ceil_to_ue8m0(x_amax / 448.0)
return (x_view * (1.0 / sf.unsqueeze(2))).to(torch.float8_e4m3fn).view(m, n), sf
def per_block_quant_mxfp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
assert x.dim() == 2
m, n = x.shape
x_padded = torch.zeros(
(align(m, 128), align(n, 128)), dtype=x.dtype, device=x.device
)
x_padded[:m, :n] = x
x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, 128)
x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
sf = ceil_to_ue8m0(x_amax / 448.0)
x_scaled = (x_view * (1.0 / sf)).to(torch.float8_e4m3fn)
return x_scaled.view_as(x_padded)[:m, :n].contiguous(), sf.view(
x_view.size(0), x_view.size(2)
)
# For test
def native_w8a8_block_fp8_matmul(A, B, As, Bs, block_size, output_dtype=torch.float16):
"""This function performs matrix multiplication with block-wise quantization using native torch.
It takes two input tensors `A` and `B` with scales `As` and `Bs`.
The output is returned in the specified `output_dtype`.
"""
A = A.to(torch.float32)
B = B.to(torch.float32)
assert A.shape[-1] == B.shape[-1]
assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2
assert len(block_size) == 2
block_n, block_k = block_size[0], block_size[1]
assert (A.shape[-1] + block_k - 1) // block_k == As.shape[-1]
assert A.shape[:-1] == As.shape[:-1]
M = A.numel() // A.shape[-1]
N, K = B.shape
origin_C_shape = A.shape[:-1] + (N,)
A = A.reshape(M, A.shape[-1])
As = As.reshape(M, As.shape[-1])
n_tiles = (N + block_n - 1) // block_n
k_tiles = (K + block_k - 1) // block_k
assert n_tiles == Bs.shape[0]
assert k_tiles == Bs.shape[1]
C_shape = (M, N)
C = torch.zeros(C_shape, dtype=torch.float32, device=A.device)
A_tiles = [A[:, i * block_k : min((i + 1) * block_k, K)] for i in range(k_tiles)]
B_tiles = [
[
B[
j * block_n : min((j + 1) * block_n, N),
i * block_k : min((i + 1) * block_k, K),
]
for i in range(k_tiles)
]
for j in range(n_tiles)
]
C_tiles = [C[:, j * block_n : min((j + 1) * block_n, N)] for j in range(n_tiles)]
As_tiles = [As[:, i : i + 1] for i in range(k_tiles)]
for i in range(k_tiles):
for j in range(n_tiles):
a = A_tiles[i]
b = B_tiles[j][i]
c = C_tiles[j]
s = As_tiles[i] * Bs[j][i]
c[:, :] += torch.matmul(a, b.t()) * s
C = C.reshape(origin_C_shape).to(output_dtype)
return C
def block_quant_dequant(
x_q_block: torch.Tensor,
x_s: torch.Tensor,
block_size: List[int],
dtype: torch.dtype,
) -> torch.Tensor:
"""This function converts block-wise quantization to unquantized.
The inputs are block-wise quantization tensor `x_q_block`, block-wise quantization scale
and the block size.
The output is an unquantized tensor with dtype.
"""
block_n, block_k = block_size[0], block_size[1]
n, k = x_q_block.shape
n_tiles = (n + block_n - 1) // block_n
k_tiles = (k + block_k - 1) // block_k
assert n_tiles == x_s.shape[0]
assert k_tiles == x_s.shape[1]
x_dq_block = torch.empty_like(x_q_block, dtype=dtype)
for j in range(n_tiles):
for i in range(k_tiles):
x_q_block_tile = x_q_block[
j * block_n : min((j + 1) * block_n, n),
i * block_k : min((i + 1) * block_k, k),
]
x_dq_block_tile = x_dq_block[
j * block_n : min((j + 1) * block_n, n),
i * block_k : min((i + 1) * block_k, k),
]
x_dq_block_tile[:, :] = x_q_block_tile.to(torch.float32) * x_s[j][i]
return x_dq_block
class TestDeepGemmBlackwell(CustomTestCase):
if not _is_cuda:
OUT_DTYPES = [torch.float32, torch.half, torch.bfloat16]
M = [1, 7, 83, 512, 2048]
NKs = [
(N, K)
for N in [128, 512, 1024, 4096, 7748, 13824]
for K in [256, 4096, 5120, 3884, 13824]
]
# BLOCK_SIZE = [[64, 64], [64, 128], [128, 64], [128, 128]]
BLOCK_SIZE = [[128, 128]]
SEEDS = [0]
else:
# use practical shape in DeepSeek V3 for test
OUT_DTYPES = [torch.bfloat16]
M = [64, 128, 512, 1024, 4096]
NKs = [
(2112, 7168),
(1536, 7168),
# (3072, 1536),
# (24576, 7168),
# (4096, 512),
# (7168, 2048),
# (4608, 7168),
# (512, 7168),
# (7168, 2304),
# (7168, 512),
]
BLOCK_SIZE = [[128, 128]]
SEEDS = [0]
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
torch.set_default_device("cuda")
def _test_deep_gemm_blackwell(self, M, NK, block_size, out_dtype, seed):
N, K = NK
torch.manual_seed(seed)
A = torch.empty((M, K), dtype=torch.bfloat16).normal_(0, 0.2)
B = torch.empty((N, K), dtype=torch.bfloat16).normal_(0, 0.2)
A_q, A_s = per_token_group_quant_fp8(A)
B_q, B_s = per_block_quant_fp8(B)
A_dq = block_quant_dequant(A_q, A_s, [1, block_size[1]], out_dtype)
B_dq = block_quant_dequant(B_q, B_s, block_size, out_dtype)
A_qu = per_token_group_quant_mxfp8(A_dq)
B_qu = per_block_quant_mxfp8(B_dq)
out = None
with torch.inference_mode():
ref_out = native_w8a8_block_fp8_matmul(
A_qu[0], B_qu[0], A_qu[1], B_qu[1], block_size, out_dtype
)
out = torch.empty_like(ref_out)
fp8_gemm_nt(A_qu, B_qu, out)
torch.testing.assert_close(out, ref_out, atol=1e-1, rtol=1e-2)
def test_deep_gemm_blackwell(self):
for params in itertools.product(
self.M,
self.NKs,
self.BLOCK_SIZE,
self.OUT_DTYPES,
self.SEEDS,
):
with self.subTest(
M=params[0],
NKs=params[1],
block_size=params[2],
out_dtype=params[3],
seed=params[4],
):
self._test_deep_gemm_blackwell(*params)
if __name__ == "__main__":
unittest.main(verbosity=2)
test/manual/quant/test_custom_ops.py
+152
View File
@@ -0,0 +1,152 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Adapted from https://github.com/vllm-project/vllm/blob/8ca7a71df787ad711ad3ac70a5bd2eb2bb398938/tests/quantization/test_fp8.py
import sys
import pytest
import torch
from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz, scaled_fp8_quant
from sglang.srt.utils import is_cuda, is_hip
_is_cuda = is_cuda()
_is_hip = is_hip()
_is_fp8_fnuz = is_fp8_fnuz()
fp8_dtype = torch.float8_e4m3fnuz if _is_fp8_fnuz else torch.float8_e4m3fn
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_scaled_fp8_quant_per_tensor(dtype) -> None:
def quantize_ref_per_tensor(tensor, inv_scale):
# The reference implementation that fully aligns to
# the kernel being tested.
finfo = torch.finfo(fp8_dtype)
scale = inv_scale.reciprocal()
qweight = (tensor.to(torch.float32) * scale).clamp(min=finfo.min, max=finfo.max)
qweight = qweight.to(fp8_dtype)
return qweight
def dequantize_per_tensor(tensor, inv_scale, dtype):
fake_qweight = tensor.to(dtype)
dq_weight = fake_qweight * inv_scale
return dq_weight
# Note that we use a shape % 8 != 0 to cover edge cases,
# because scaled_fp8_quant is vectorized by 8.
x = (torch.randn(size=(11, 11), device="cuda") * 13).to(dtype)
# Test Per Tensor Dynamic quantization
# scale = max(abs(x)) / FP8_E4M3_MAX
y, scale = scaled_fp8_quant(x, None)
ref_y = quantize_ref_per_tensor(x, scale)
torch.testing.assert_close(y, ref_y)
torch.testing.assert_close(
dequantize_per_tensor(y, scale, dtype),
dequantize_per_tensor(ref_y, scale, dtype),
)
# Test Per Tensor Static quantization
y, _ = scaled_fp8_quant(x, scale)
ref_y = quantize_ref_per_tensor(x, scale)
torch.testing.assert_close(y, ref_y)
torch.testing.assert_close(
dequantize_per_tensor(y, scale, dtype),
dequantize_per_tensor(ref_y, scale, dtype),
)
if _is_cuda or _is_hip:
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_scaled_fp8_quant_per_token_dynamic(dtype) -> None:
def quantize_ref_per_token(tensor, inv_scale):
# The reference implementation that fully aligns to
# the kernel being tested.
finfo = torch.finfo(fp8_dtype)
scale = inv_scale.reciprocal()
qweight = (tensor.to(torch.float32) * scale).clamp(
min=finfo.min, max=finfo.max
)
qweight = qweight.to(fp8_dtype)
return qweight
def dequantize_per_token(tensor, inv_scale, dtype):
fake_qweight = tensor.to(dtype)
dq_weight = fake_qweight * inv_scale
return dq_weight
# Note that we use a shape % 8 = 0,
# because per_token_quant_fp8 is vectorized by 8 elements.
x = (torch.randn(size=(11, 16), device="cuda") * 13).to(dtype)
# Test Per Tensor Dynamic quantization
# scale = max(abs(x)) / FP8_E4M3_MAX
y, scale = scaled_fp8_quant(x, None, use_per_token_if_dynamic=True)
ref_y = quantize_ref_per_token(x, scale)
torch.testing.assert_close(y, ref_y)
torch.testing.assert_close(
dequantize_per_token(y, scale, dtype),
dequantize_per_token(ref_y, scale, dtype),
)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_scaled_fp8_quant_with_padding(dtype) -> None:
original_rows = 5
x = (torch.randn(size=(original_rows, 16), device="cuda") * 13).to(dtype)
padding_size = 10
# Test with dynamic quantization
y_dynamic, scale_dynamic = scaled_fp8_quant(
x, None, num_token_padding=padding_size
)
# Verify output shape has the padded size
assert y_dynamic.shape[0] == padding_size
assert y_dynamic.shape[1] == x.shape[1]
# Verify that the actual data in the non-padded region is correctly quantized
y_without_padding, scale_without_padding = scaled_fp8_quant(x, None)
torch.testing.assert_close(y_dynamic[:original_rows], y_without_padding)
# Test with static quantization
# First get a scale
_, scale = scaled_fp8_quant(x, None)
# Then use it for static quantization with padding
y_static, _ = scaled_fp8_quant(x, scale, num_token_padding=padding_size)
# Verify output shape has the padded size
assert y_static.shape[0] == padding_size
assert y_static.shape[1] == x.shape[1]
# Verify that the actual data in the non-padded region is correctly quantized
y_static_without_padding, _ = scaled_fp8_quant(x, scale)
torch.testing.assert_close(y_static[:original_rows], y_static_without_padding)
# Test with per-token dynamic quantization
y_per_token, scale_per_token = scaled_fp8_quant(
x, None, num_token_padding=padding_size, use_per_token_if_dynamic=True
)
# Verify output shape has the padded size
assert y_per_token.shape[0] == padding_size
assert y_per_token.shape[1] == x.shape[1]
# Verify that the actual data in the non-padded region is correctly quantized
y_per_token_without_padding, scale_per_token_without_padding = scaled_fp8_quant(
x, None, use_per_token_if_dynamic=True
)
torch.testing.assert_close(
y_per_token[:original_rows], y_per_token_without_padding
)
torch.testing.assert_close(
scale_per_token[:original_rows], scale_per_token_without_padding
)
if __name__ == "__main__":
# Run the specific test function directly
sys.exit(pytest.main([__file__]))
test/manual/quant/test_cutlass_moe.py
Executable
+306
View File
@@ -0,0 +1,306 @@
import argparse
import torch
import triton # Added import
import triton.testing # Added import
from transformers import AutoConfig
from sglang.srt.layers.moe.cutlass_moe import cutlass_fused_experts_fp8
from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import fused_experts
from sglang.srt.layers.moe.topk import StandardTopKOutput
# Copy from: https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/utils.py
def calc_diff(x, y):
x, y = x.double(), y.double()
denominator = (x * x + y * y).sum()
sim = 2 * (x * y).sum() / denominator
return 1 - sim
def get_model_config(tp_size: int):
config = AutoConfig.from_pretrained(
"deepseek-ai/Deepseek-R1", trust_remote_code=True
)
E = config.n_routed_experts
topk = config.num_experts_per_tok
intermediate_size = config.moe_intermediate_size
shard_intermediate_size = 2 * intermediate_size // tp_size
return {
"num_experts": E,
"topk": topk,
"hidden_size": config.hidden_size,
"shard_intermediate_size": shard_intermediate_size,
"dtype": config.dtype,
"block_shape": config.quantization_config["weight_block_size"],
}
def to_fp8(tensor: torch.Tensor) -> torch.Tensor:
"""Converts tensor to FP8 E4M3, scaling values to fit the range."""
finfo = torch.finfo(torch.float8_e4m3fn)
# Calculate max absolute value safely
max_val = torch.max(torch.abs(tensor))
# Avoid division by zero if tensor is all zeros
if max_val == 0:
scale_factor = 1.0
else:
# Scale factor to bring the max value to finfo.max
scale_factor = finfo.max / max_val
# Apply scaling
scaled_tensor = tensor * scale_factor
# Clamp and convert
fp8_tensor = scaled_tensor.clamp(min=finfo.min, max=finfo.max).to(
dtype=torch.float8_e4m3fn
)
return fp8_tensor
def run_test(tp_size, batch_size, model_config, check=False):
print(f"\n--- Batch Size: {batch_size} ---")
torch.set_default_device("cuda")
torch.cuda.manual_seed_all(42) # For reproducible random numbers
E = model_config["num_experts"]
topk = model_config["topk"]
H = model_config["hidden_size"]
I = model_config["shard_intermediate_size"]
block_shape = model_config["block_shape"] # Tuple (BLOCK_N, BLOCK_K)
dtype = model_config["dtype"] # e.g., torch.bfloat16
print(
f"Config: E={E}, topk={topk}, H={H}, I_shard={I}, dtype={dtype}, block_shape={block_shape}"
)
# --- Input Data ---
# Use bf16/fp16 for input activation based on model config
x = torch.randn((batch_size, H), device="cuda", dtype=dtype)
# --- Weights (Generate in higher precision, then convert to FP8) ---
# Generate weights suitable for FP8 conversion (e.g., scaled appropriately)
w1_hp = torch.randn((E, I, H), device="cuda", dtype=torch.float32)
w2_hp = torch.randn((E, H, I // 2), device="cuda", dtype=torch.float32)
w1 = to_fp8(w1_hp)
w2 = to_fp8(w2_hp)
# --- Scales for FP8 Weights ---
block_n, block_k = block_shape
# Calculate number of blocks needed
w1_blocks_dim1 = (I + block_n - 1) // block_n
w1_blocks_dim2 = (H + block_k - 1) // block_k
w2_blocks_dim1 = (H + block_n - 1) // block_n
w2_blocks_dim2 = (I // 2 + block_k - 1) // block_k
# Scales are typically float32 or float16/bfloat16
scale_dtype = torch.float32 # Or dtype if scales match model dtype
w1_scale = torch.full(
(E, w1_blocks_dim1, w1_blocks_dim2), 1, device="cuda", dtype=scale_dtype
) # Avoid zero scales
w2_scale = torch.full(
(E, w2_blocks_dim1, w2_blocks_dim2), 1, device="cuda", dtype=scale_dtype
) # Avoid zero scales
# --- Routing Information ---
topk_weights = torch.softmax(
torch.rand(batch_size, topk, device="cuda", dtype=dtype), dim=-1
)
topk_ids = torch.randint(0, E, (batch_size, topk), dtype=torch.int32, device="cuda")
a1_strides = torch.full((E,), H, dtype=torch.int64, device="cuda")
c1_strides = torch.full((E,), I, dtype=torch.int64, device="cuda")
a2_strides = torch.full((E,), I // 2, dtype=torch.int64, device="cuda")
c2_strides = torch.full((E,), H, dtype=torch.int64, device="cuda")
workspace = torch.empty(
(7182 * 1024), device="cuda", dtype=torch.uint8
) # Allocate sufficient workspace
# Pointer arrays (often filled by the kernel or a prep step, but needed as args)
a_ptrs = torch.empty((E,), dtype=torch.int64, device="cuda")
b_ptrs = torch.empty((E,), dtype=torch.int64, device="cuda")
out_ptrs = torch.empty((E,), dtype=torch.int64, device="cuda")
a_scales_ptrs = torch.empty((E,), dtype=torch.int64, device="cuda")
b_scales_ptrs = torch.empty((E,), dtype=torch.int64, device="cuda")
expert_offsets = torch.empty((E + 1,), dtype=torch.int32, device="cuda")
problem_sizes1 = torch.empty((E, 3), dtype=torch.int32, device="cuda")
problem_sizes2 = torch.empty((E, 3), dtype=torch.int32, device="cuda")
enable_es = (False, False)
if torch.cuda.get_device_name(torch.cuda.current_device()) == "NVIDIA H200":
enable_es = (False, True)
elif torch.cuda.get_device_name(torch.cuda.current_device()) == "NVIDIA H20":
enable_es = (True, True)
# --- Lambdas for Benchmarking ---
cutlass_lambda = lambda: cutlass_fused_experts_fp8(
x,
w1.transpose(1, 2), # Transposed
w2.transpose(1, 2), # Transposed
w1_scale.transpose(1, 2),
w2_scale.transpose(1, 2),
topk_weights,
topk_ids,
a1_strides,
c1_strides,
a2_strides,
c2_strides,
workspace,
a_ptrs,
b_ptrs,
out_ptrs,
a_scales_ptrs,
b_scales_ptrs,
expert_offsets,
problem_sizes1,
problem_sizes2,
enable_es=enable_es,
)
topk_output = StandardTopKOutput(
topk_weights=topk_weights,
topk_ids=topk_ids,
router_logits=torch.randn(
(batch_size, topk), device=topk_weights.device, dtype=dtype
),
)
moe_runner_config = MoeRunnerConfig(
num_experts=E,
top_k=topk,
hidden_size=H,
intermediate_size_per_partition=I,
params_dtype=dtype,
activation="silu",
inplace=False,
)
# Note: Triton expects non-transposed weights
triton_lambda = lambda: fused_experts(
x,
w1,
w2,
topk_output,
moe_runner_config,
use_fp8_w8a8=True,
w1_scale=w1_scale,
w2_scale=w2_scale,
block_shape=block_shape,
)
# --- Warmup ---
print("Warming up...")
for _ in range(10):
_ = cutlass_lambda()
_ = triton_lambda()
torch.cuda.synchronize()
# --- Benchmarking ---
quantiles = [0.5, 0.2, 0.8]
print(f"Benchmarking Cutlass fused_experts...")
cutlass_ms, cutlass_min, cutlass_max = triton.testing.do_bench_cudagraph(
cutlass_lambda, rep=1000, quantiles=quantiles
)
print(f"Benchmarking Triton fused_experts...")
triton_ms, triton_min, triton_max = triton.testing.do_bench_cudagraph(
triton_lambda, rep=1000, quantiles=quantiles
)
print(
f"Cutlass fused_experts time: {cutlass_ms:.3f} ms (median) [{cutlass_min:.3f} - {cutlass_max:.3f}]"
)
print(
f"Triton fused_experts time: {triton_ms:.3f} ms (median) [{triton_min:.3f} - {triton_max:.3f}]"
)
# --- Correctness Check ---
if check:
print("Running correctness check...")
with torch.no_grad():
# Run CUTLASS version (requires transposed weights)
y_cutlass = cutlass_fused_experts_fp8(
x,
w1.transpose(1, 2), # Transposed
w2.transpose(1, 2), # Transposed
w1_scale.transpose(1, 2),
w2_scale.transpose(1, 2),
topk_weights,
topk_ids,
a1_strides,
c1_strides,
a2_strides,
c2_strides,
workspace,
a_ptrs,
b_ptrs,
out_ptrs,
a_scales_ptrs,
b_scales_ptrs,
expert_offsets,
problem_sizes1,
problem_sizes2,
enable_es=enable_es,
)
# Run Triton version (requires original shape weights, use inplace=False)
y_triton = fused_experts(
x,
w1, # Original shape
w2, # Original shape
topk_output,
moe_runner_config,
use_fp8_w8a8=True,
w1_scale=w1_scale,
w2_scale=w2_scale,
block_shape=block_shape,
)
diff = calc_diff(y_cutlass, y_triton)
print(f"Diff: {diff:.6f}")
# Tolerance might need adjustment based on FP8 specifics and kernel differences
# FP8 comparisons often require higher tolerance than FP16/BF16
assert diff < 1e-4, f"Diff too high! {diff}"
print("Correctness check passed.")
def main(tp_size=8, batch_sizes=[1, 4, 8, 16, 32, 64, 128, 256, 512], check=False):
model_config = get_model_config(tp_size)
print("Model Config:", model_config)
for batch_size in batch_sizes:
run_test(tp_size, batch_size, model_config, check)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--tp-size", type=int, default=8, help="Tensor Parallel size")
parser.add_argument(
"--batch-sizes",
type=int,
nargs="+",
default=[
1,
4,
8,
16,
32,
64,
128,
256,
512,
1024,
2048,
4096,
8192,
], # Adjusted default
help="List of batch sizes to test",
)
parser.add_argument("--check", action="store_true", help="Enable check mode")
args = parser.parse_args()
print(f"Running benchmarks with TP size: {args.tp_size}")
print(f"Testing batch sizes: {args.batch_sizes}")
main(tp_size=args.tp_size, batch_sizes=args.batch_sizes, check=args.check)
test/manual/quant/test_cutlass_w16a16_moe.py
+118
View File
@@ -0,0 +1,118 @@
# SPDX-License-Identifier: Apache-2.0
import pytest
import torch
from flashinfer.fused_moe import cutlass_fused_moe as flashinfer_cutlass_fused_moe
from sglang.srt.layers.activation import SiluAndMul
from sglang.srt.layers.moe.topk import TopKConfig, select_experts
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
MNK_FACTORS = [
(2, 1024, 1024),
(2, 1024, 1536),
(2, 3072, 1024),
(2, 3072, 1536),
(64, 1024, 1024),
(64, 1024, 1536),
(64, 3072, 1024),
(64, 2048, 1024),
(224, 1024, 1024),
(224, 1024, 1536),
]
# Reference implementation of torch_moe for unquantized weights
def torch_moe_reference(a, w13, w2, score, topk):
B, D = a.shape
set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
# Flip w13 layout
dim = -2
size = w13.size(dim)
assert size % 2 == 0, f"Expected even size in dim {dim}, got {size}"
half = size // 2
# Reorder weight
w1, w3 = w13.split(half, dim=dim)
w13 = torch.cat([w3, w1], dim=dim).contiguous()
a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
score = torch.softmax(score, dim=-1, dtype=torch.float32)
topk_weight, topk_ids = torch.topk(score, topk)
topk_weight = topk_weight.view(-1)
topk_ids = topk_ids.view(-1)
for i in range(w13.shape[0]):
mask = topk_ids == i
if mask.sum():
out[mask] = SiluAndMul()(a[mask] @ w13[i].transpose(0, 1)) @ w2[
i
].transpose(0, 1)
return (
out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype)
).sum(dim=1)
@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
@pytest.mark.parametrize("e", [40, 64, 256])
@pytest.mark.parametrize("topk", [1, 6, 8])
@torch.inference_mode()
def test_flashinfer_bf16_cutlass_moe(m: int, n: int, k: int, e: int, topk: int):
"""
Test the bf16 cutlass moe API.
Args:
m: number of tokens
n: intermediate size
k: hidden size
e: number of experts
topk: top-k experts per token
"""
torch.manual_seed(7)
dtype = torch.bfloat16
# Create unquantized weights
a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
# w13: fused gate_up projection [num_experts, 2*intermediate, hidden]
# FlashInfer CUTLASS expects [up, gate] layout
w13 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
# w2: down projection [num_experts, hidden, intermediate]
w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
# Generate router scores
score = torch.randn((m, e), device="cuda", dtype=dtype)
# Get topk routing
topk_output = select_experts(
hidden_states=a,
router_logits=score,
topk_config=TopKConfig(top_k=topk, renormalize=False),
)
topk_weights, topk_ids, _ = topk_output
# Test: Call FlashInfer CUTLASS fused_moe (unquantized version)
test_output = flashinfer_cutlass_fused_moe(
input=a,
token_selected_experts=topk_ids,
token_final_scales=topk_weights,
fc1_expert_weights=w13,
fc2_expert_weights=w2,
output_dtype=dtype,
quant_scales=None,
)[0]
# Reference: Torch implementation
torch_output = torch_moe_reference(a, w13, w2, score, topk)
# Compare outputs
torch.testing.assert_close(torch_output, test_output, rtol=1e-2, atol=1e-2)
if __name__ == "__main__":
# Run a simple test case
test_flashinfer_bf16_cutlass_moe(224, 1024, 1024, 8, 2)
test/manual/quant/test_cutlass_w4a8_moe.py
+285
View File
@@ -0,0 +1,285 @@
# SPDX-License-Identifier: Apache-2.0
from typing import Optional
import pytest
import torch
from sglang.srt.layers.moe.cutlass_w4a8_moe import cutlass_w4a8_moe
from sglang.srt.layers.moe.topk import TopKConfig, select_experts
def pack_int4_values_to_int8(int4_values_interleaved: torch.Tensor) -> torch.Tensor:
if int4_values_interleaved.shape[-1] % 2 != 0:
raise ValueError(
"the last dim size of int4_values_interleaved tensor must be even."
)
input_tensor_int8 = int4_values_interleaved.to(torch.int8)
low_nibbles = input_tensor_int8[..., 0::2]
high_nibbles = input_tensor_int8[..., 1::2]
packed_tensor = (high_nibbles << 4) | (low_nibbles & 0x0F)
return packed_tensor.to(torch.int8)
def pack_interleave(num_experts, ref_weight, ref_scale, alignment=4):
n, k = ref_weight.shape[1], ref_weight.shape[2]
weight = pack_int4_values_to_int8(ref_weight.cpu()).cuda()
w_q = weight.view((num_experts, n, k // 2)).view(torch.int8)
w_q = w_q.contiguous()
scale_interleaved = ref_scale.reshape(
ref_scale.shape[0],
ref_scale.shape[1],
(ref_scale.shape[2] // alignment),
alignment,
) # [E, N, K/4, 4]
scale_interleaved = scale_interleaved.permute(0, 2, 1, 3) # [E, K/4, N, 4]
scale_interleaved = scale_interleaved.reshape(
ref_scale.shape[0],
ref_scale.shape[2] // alignment,
ref_scale.shape[1] * alignment,
) # [E, K/4, N*4]
w_scale = scale_interleaved.contiguous()
return w_q, w_scale
@pytest.mark.parametrize("M", [1, 2, 4, 8, 16])
@pytest.mark.parametrize("N", [2048])
@pytest.mark.parametrize("K", [7168])
@pytest.mark.parametrize("E", [256])
@pytest.mark.parametrize("tp_size", [8])
@pytest.mark.parametrize("use_ep_moe", [True, False])
@pytest.mark.parametrize("topk", [8])
@pytest.mark.parametrize("group_size", [128])
@pytest.mark.parametrize("dtype", [torch.bfloat16])
def test_cutlass_w4a8_moe(M, N, K, E, tp_size, use_ep_moe, topk, group_size, dtype):
if use_ep_moe:
local_e = E // tp_size
else: # tp mode
local_e = E
N = N // tp_size
debug = False
if debug:
a = torch.ones((M, K), dtype=dtype, device="cuda") * 0.001
ref_weight_1 = torch.ones((local_e, N * 2, K), dtype=torch.int8, device="cuda")
ref_weight_2 = torch.ones((local_e, K, N), dtype=torch.int8, device="cuda")
a1_scale = torch.ones(1, dtype=torch.float32, device="cuda")
a2_scale = torch.ones(1, dtype=torch.float32, device="cuda")
scale_1 = torch.ones(
(local_e, N * 2, K // group_size), dtype=dtype, device="cuda"
)
scale_2 = torch.ones((local_e, K, N // group_size), dtype=dtype, device="cuda")
else:
a = torch.randn(M, K, dtype=dtype, device="cuda")
ref_weight_1 = torch.randint(
-8, 8, (local_e, N * 2, K), dtype=torch.int8, device="cuda"
)
ref_weight_2 = torch.randint(
-8, 8, (local_e, K, N), dtype=torch.int8, device="cuda"
)
affine_coeff = 0.005
a1_scale = torch.randn(1, dtype=torch.float32, device="cuda")
a2_scale = torch.randn(1, dtype=torch.float32, device="cuda")
scale_1 = (
torch.randn(local_e, N * 2, K // group_size, dtype=dtype, device="cuda")
* affine_coeff
)
scale_2 = (
torch.randn(local_e, K, N // group_size, dtype=dtype, device="cuda")
* affine_coeff
)
w1_q, w1_scale = pack_interleave(local_e, ref_weight_1, scale_1)
if use_ep_moe:
w2_q, w2_scale = pack_interleave(local_e, ref_weight_2, scale_2)
else:
w2_q, w2_scale = pack_interleave(local_e, ref_weight_2, scale_2, 1)
device = "cuda"
a_strides1 = torch.full((local_e, 3), K, device=device, dtype=torch.int64)
c_strides1 = torch.full((local_e, 3), 2 * N, device=device, dtype=torch.int64)
a_strides2 = torch.full((local_e, 3), N, device=device, dtype=torch.int64)
c_strides2 = torch.full((local_e, 3), K, device=device, dtype=torch.int64)
b_strides1 = a_strides1
s_strides13 = c_strides1
b_strides2 = a_strides2
s_strides2 = c_strides2
score = torch.randn((M, E), dtype=dtype, device=device)
topk_output = select_experts(
hidden_states=a,
router_logits=score,
topk_config=TopKConfig(top_k=topk, renormalize=False),
)
topk_weights, topk_ids, _ = topk_output
expert_map = torch.arange(E, dtype=torch.int32, device=device)
expert_map[local_e:] = -1
output = cutlass_moe(
a,
w1_q,
w2_q,
w1_scale,
w2_scale,
topk_weights,
topk_ids,
a_strides1,
b_strides1,
c_strides1,
a_strides2,
b_strides2,
c_strides2,
s_strides13,
s_strides2,
local_e,
a1_scale,
a2_scale,
expert_map,
)
ref_output = ref(
a,
local_e,
topk_weights,
topk_ids,
ref_weight_1,
ref_weight_2,
scale_1,
scale_2,
has_pre_quant=True,
has_alpha=True,
pre_quant_scale_1=a1_scale,
pre_quant_scale_2=a2_scale,
alpha_1=a1_scale,
alpha_2=a2_scale,
)
# compare
torch.cuda.synchronize()
# compare final output
torch.testing.assert_close(output, ref_output, rtol=1e-2, atol=0.1)
print("SUCCESS: Final output tensors are close.")
def cutlass_moe(
a: torch.Tensor,
w1_q: torch.Tensor,
w2_q: torch.Tensor,
w1_scale: torch.Tensor,
w2_scale: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
a_strides1: torch.Tensor,
b_strides1: torch.Tensor,
c_strides1: torch.Tensor,
a_strides2: torch.Tensor,
b_strides2: torch.Tensor,
c_strides2: torch.Tensor,
s_strides13: torch.Tensor,
s_strides2: torch.Tensor,
num_local_experts: int,
a1_scale: Optional[torch.Tensor] = None,
a2_scale: Optional[torch.Tensor] = None,
expert_map: Optional[torch.Tensor] = None,
apply_router_weight_on_input: bool = False,
):
topk_ids = expert_map[topk_ids]
device = a.device
expert_offsets = torch.empty(
(num_local_experts + 1), dtype=torch.int32, device=device
)
problem_sizes1 = torch.empty(
(num_local_experts, 3), dtype=torch.int32, device=device
)
problem_sizes2 = torch.empty(
(num_local_experts, 3), dtype=torch.int32, device=device
)
return cutlass_w4a8_moe(
a,
w1_q,
w2_q,
w1_scale,
w2_scale,
topk_weights,
topk_ids,
a_strides1,
b_strides1,
c_strides1,
a_strides2,
b_strides2,
c_strides2,
s_strides13,
s_strides2,
expert_offsets,
problem_sizes1,
problem_sizes2,
a1_scale,
a2_scale,
apply_router_weight_on_input,
)
def ref(
x: torch.Tensor,
num_experts: int,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
ref_weight_1: torch.Tensor,
ref_weight_2: torch.Tensor,
ref_weight_scale_1: torch.Tensor,
ref_weight_scale_2: torch.Tensor,
has_pre_quant: bool = False,
has_alpha: bool = False,
pre_quant_scale_1: Optional[torch.Tensor] = None,
pre_quant_scale_2: Optional[torch.Tensor] = None,
alpha_1: Optional[torch.Tensor] = None,
alpha_2: Optional[torch.Tensor] = None,
):
results = torch.zeros_like(x)
dtype = x.dtype
for e_idx in range(num_experts):
mask = topk_ids == e_idx
activated_tokens = mask.sum(1).bool()
act = x[activated_tokens, :]
if act.shape[0] == 0:
continue
final_scale = (topk_weights * mask).sum(1)[activated_tokens].unsqueeze(1)
act = (
torch.clamp((act / pre_quant_scale_1.float()), -448.0, 448.0)
.to(torch.float8_e4m3fn)
.to(dtype)
)
w3_w1 = ref_weight_1[e_idx]
ref_w_scale_repeat = (
ref_weight_scale_1[e_idx].repeat_interleave(128, dim=1).to(float)
)
w3_w1 = (w3_w1.to(float) * ref_w_scale_repeat).to(dtype)
fc1 = ((torch.matmul(act, w3_w1.T)) * alpha_1).to(torch.float16)
gate, fc1 = fc1.chunk(2, dim=-1)
fc1 = fc1 * torch.nn.functional.silu(gate)
act = torch.clamp((fc1 / pre_quant_scale_2.float()), -448.0, 448.0).to(
torch.float8_e4m3fn
)
act = act.to(dtype)
w2 = ref_weight_2[e_idx]
ref_w_scale_repeat = (
ref_weight_scale_2[e_idx].repeat_interleave(128, dim=1).to(float)
)
w2 = (w2.to(float) * ref_w_scale_repeat).to(dtype)
fc2 = (torch.matmul(act, w2.T) * alpha_2).to(torch.float16)
results[activated_tokens, :] += (fc2 * final_scale).to(results.dtype)
return results
test/manual/quant/test_kvfp4_quant_dequant.py
Executable
+116
View File
@@ -0,0 +1,116 @@
#!/usr/bin/env python3
import time
import numpy as np
import pytest
import torch
from sglang.srt.layers.quantization.kvfp4_tensor import KVFP4QuantizeUtil
def calculate_accuracy_metrics(
original: torch.Tensor, reconstructed: torch.Tensor
) -> dict[str, float]:
"""Calculate accuracy metrics between original and reconstructed tensors."""
mse = torch.mean((original - reconstructed) ** 2).item()
mae = torch.mean(torch.abs(original - reconstructed)).item()
# PSNR calculation
max_val = torch.max(torch.abs(original)).item()
psnr = 20 * np.log10(max_val / np.sqrt(mse)) if mse > 0 else float("inf")
# Relative error
rel_error = torch.mean(
torch.abs(original - reconstructed) / (torch.abs(original) + 1e-8)
).item()
return {"MSE": mse, "MAE": mae, "PSNR": psnr, "Relative Error": rel_error}
def run_benchmark(m, n, k, num_runs=100) -> dict[str, dict[str, float]]:
"""Run FP8 vs KVFP4 quantization benchmark and return metrics."""
tensor_bf16 = torch.randn(m, n, k, dtype=torch.bfloat16, device="cuda")
# --- FP8 ---
for _ in range(3): # warmup
_ = tensor_bf16 * 2
torch.cuda.synchronize()
start = time.time()
for _ in range(num_runs):
tensor_fp8 = tensor_bf16.to(torch.float8_e4m3fn)
torch.cuda.synchronize()
fp8_quant_time = (time.time() - start) / num_runs
start = time.time()
for _ in range(num_runs):
tensor_fp8_dequant = tensor_fp8.to(torch.bfloat16)
torch.cuda.synchronize()
fp8_dequant_time = (time.time() - start) / num_runs
fp8_metrics = calculate_accuracy_metrics(tensor_bf16, tensor_fp8_dequant)
# --- KVFP4 ---
tensor_fp4, scale_factors = KVFP4QuantizeUtil.batched_quantize(tensor_bf16)
_ = KVFP4QuantizeUtil.batched_dequantize(tensor_fp4, scale_factors)
start = time.time()
for _ in range(num_runs):
tensor_fp4, scale_factors = KVFP4QuantizeUtil.batched_quantize(tensor_bf16)
torch.cuda.synchronize()
fp4_quant_time = (time.time() - start) / num_runs
start = time.time()
for _ in range(num_runs):
tensor_fp4_dequant = KVFP4QuantizeUtil.batched_dequantize(
tensor_fp4, scale_factors
)
torch.cuda.synchronize()
fp4_dequant_time = (time.time() - start) / num_runs
fp4_metrics = calculate_accuracy_metrics(tensor_bf16, tensor_fp4_dequant)
return {
"fp8": {
"quant_time": fp8_quant_time,
"dequant_time": fp8_dequant_time,
**fp8_metrics,
},
"fp4": {
"quant_time": fp4_quant_time,
"dequant_time": fp4_dequant_time,
**fp4_metrics,
},
}
# default tensor shapes (m, n, k)
# [M, 1, 576]: DeepSeekR1-FP4 MLA
# [M, 8, 64]: gpt-oss-20b MHA
MNK_FACTORS = [
(64, 1, 576),
(512, 1, 576),
(1024, 1, 576),
(4096, 1, 576),
(2868672, 1, 576),
(64, 8, 64),
(512, 8, 64),
(1024, 8, 64),
(4096, 8, 64),
(2868672, 8, 64),
]
@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
def test_kvfp4_quant_dequant(m, n, k):
"""Benchmark FP8 vs KVFP4 for predefined tensor shapes."""
print(f"\n=== Running benchmark for tensor shape: [{m}, {n}, {k}] ===")
results = run_benchmark(m, n, k)
print("FP8:", results["fp8"])
print("FP4:", results["fp4"])
# Basic assertions to make sure metrics are reasonable
assert results["fp4"]["MSE"] < 1.0
assert results["fp8"]["MSE"] < 1.0