[Kernel] Reclassify kernel tests by ops group + move helpers out of the package (RFC #29630) (#32128)
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Xiaoyu Zhang
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
a2935ce329
commit
2d1a7be8c4
@@ -0,0 +1,170 @@
|
||||
import itertools
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.kernels.ops.quantization.awq_dequantize import (
|
||||
awq_dequantize as jit_awq_dequantize,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=9, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
|
||||
try:
|
||||
from sgl_kernel import awq_dequantize as aot_awq_dequantize
|
||||
|
||||
AOT_AVAILABLE = True
|
||||
except ImportError:
|
||||
AOT_AVAILABLE = False
|
||||
|
||||
|
||||
def reverse_awq_order(t: torch.Tensor):
|
||||
bits = 4
|
||||
AWQ_REVERSE_ORDER = [0, 4, 1, 5, 2, 6, 3, 7]
|
||||
reverse_order_tensor = torch.arange(
|
||||
t.shape[-1],
|
||||
dtype=torch.int32,
|
||||
device=t.device,
|
||||
)
|
||||
reverse_order_tensor = reverse_order_tensor.view(-1, 32 // bits)
|
||||
reverse_order_tensor = reverse_order_tensor[:, AWQ_REVERSE_ORDER]
|
||||
reverse_order_tensor = reverse_order_tensor.view(-1)
|
||||
|
||||
t = t[:, reverse_order_tensor] & 0xF
|
||||
return t
|
||||
|
||||
|
||||
# qweights - [R , C // 8], int32
|
||||
# scales - [R // G, C ], float16
|
||||
# zeros - [R // G, C // 8], int32
|
||||
def awq_dequantize_torch(
|
||||
qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor, group_size: int
|
||||
) -> torch.Tensor:
|
||||
if group_size == -1:
|
||||
group_size = qweight.shape[0]
|
||||
|
||||
bits = 4
|
||||
shifts = torch.arange(0, 32, bits, device=qzeros.device)
|
||||
|
||||
iweights = torch.bitwise_right_shift(qweight[:, :, None], shifts[None, None, :]).to(
|
||||
torch.int8
|
||||
)
|
||||
|
||||
iweights = iweights.view(iweights.shape[0], -1)
|
||||
|
||||
zeros = torch.bitwise_right_shift(qzeros[:, :, None], shifts[None, None, :]).to(
|
||||
torch.int8
|
||||
)
|
||||
zeros = zeros.view(qzeros.shape[0], -1)
|
||||
zeros = reverse_awq_order(zeros)
|
||||
|
||||
iweights = reverse_awq_order(iweights)
|
||||
|
||||
iweights = torch.bitwise_and(iweights, (2**bits) - 1)
|
||||
zeros = torch.bitwise_and(zeros, (2**bits) - 1)
|
||||
|
||||
scales = scales.repeat_interleave(group_size, dim=0)
|
||||
zeros = zeros.repeat_interleave(group_size, dim=0)
|
||||
return (iweights - zeros) * scales
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"qweight_row,qweight_col,is_bf16_act",
|
||||
list(
|
||||
itertools.product(
|
||||
[128, 256, 512, 1024, 3584],
|
||||
[16, 32, 64, 128, 448],
|
||||
[True, False],
|
||||
)
|
||||
),
|
||||
)
|
||||
def test_awq_dequantize_jit_vs_torch(
|
||||
qweight_row: int, qweight_col: int, is_bf16_act: bool
|
||||
):
|
||||
device = torch.device("cuda")
|
||||
qweight = torch.randint(
|
||||
0,
|
||||
torch.iinfo(torch.int32).max,
|
||||
(qweight_row, qweight_col),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
group_size = qweight_row
|
||||
scales_row = qweight_row // group_size
|
||||
scales_col = qweight_col * 8
|
||||
|
||||
if is_bf16_act:
|
||||
scales = torch.rand(scales_row, scales_col, dtype=torch.bfloat16, device=device)
|
||||
else:
|
||||
scales = torch.rand(scales_row, scales_col, dtype=torch.float16, device=device)
|
||||
|
||||
qzeros = torch.randint(
|
||||
0,
|
||||
torch.iinfo(torch.int32).max,
|
||||
(scales_row, qweight_col),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
|
||||
# Run both implementations
|
||||
torch_out = awq_dequantize_torch(qweight, scales, qzeros, group_size)
|
||||
jit_out = jit_awq_dequantize(qweight, scales, qzeros)
|
||||
|
||||
# Compare results (approximate due to different computation paths)
|
||||
torch.testing.assert_close(
|
||||
torch_out.to(torch.float32), jit_out.to(torch.float32), rtol=1e-3, atol=1e-5
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"qweight_row,qweight_col,is_bf16_act",
|
||||
list(
|
||||
itertools.product(
|
||||
[128, 256, 512, 1024, 3584],
|
||||
[16, 32, 64, 128, 448],
|
||||
[True, False],
|
||||
)
|
||||
),
|
||||
)
|
||||
def test_awq_dequantize_jit_vs_aot(
|
||||
qweight_row: int, qweight_col: int, is_bf16_act: bool
|
||||
):
|
||||
if not AOT_AVAILABLE:
|
||||
pytest.skip("sgl_kernel AOT not available")
|
||||
|
||||
device = torch.device("cuda")
|
||||
qweight = torch.randint(
|
||||
0,
|
||||
torch.iinfo(torch.int32).max,
|
||||
(qweight_row, qweight_col),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
group_size = qweight_row
|
||||
scales_row = qweight_row // group_size
|
||||
scales_col = qweight_col * 8
|
||||
|
||||
if is_bf16_act:
|
||||
scales = torch.rand(scales_row, scales_col, dtype=torch.bfloat16, device=device)
|
||||
else:
|
||||
scales = torch.rand(scales_row, scales_col, dtype=torch.float16, device=device)
|
||||
|
||||
qzeros = torch.randint(
|
||||
0,
|
||||
torch.iinfo(torch.int32).max,
|
||||
(scales_row, qweight_col),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
|
||||
# Run both implementations
|
||||
aot_out = aot_awq_dequantize(qweight, scales, qzeros)
|
||||
jit_out = jit_awq_dequantize(qweight, scales, qzeros)
|
||||
|
||||
# Bitwise equality
|
||||
torch.testing.assert_close(jit_out, aot_out, rtol=0, atol=0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,124 @@
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from sgl_kernel.scalar_type import scalar_types
|
||||
|
||||
from sglang.kernels.ops.quantization.awq_marlin_repack import (
|
||||
awq_marlin_moe_repack as jit_awq_marlin_moe_repack,
|
||||
)
|
||||
from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
|
||||
|
||||
def _has_aot_awq_marlin_moe_repack() -> bool:
|
||||
return hasattr(torch.ops.sgl_kernel, "awq_marlin_moe_repack") and hasattr(
|
||||
torch.ops.sgl_kernel.awq_marlin_moe_repack, "default"
|
||||
)
|
||||
|
||||
|
||||
AOT_AVAILABLE = _has_aot_awq_marlin_moe_repack()
|
||||
|
||||
|
||||
def awq_pack(
|
||||
q_w: torch.Tensor,
|
||||
num_bits: int,
|
||||
size_k: int,
|
||||
size_n: int,
|
||||
):
|
||||
assert q_w.shape == (size_k, size_n)
|
||||
|
||||
if num_bits == 4:
|
||||
interleave = np.array([0, 2, 4, 6, 1, 3, 5, 7])
|
||||
elif num_bits == 8:
|
||||
interleave = np.array([0, 2, 1, 3])
|
||||
else:
|
||||
raise Exception("num_bits must be 4 or 8, got {}".format(num_bits))
|
||||
|
||||
q_w = q_w.reshape((-1, len(interleave)))[:, interleave].ravel()
|
||||
q_w = q_w.reshape((-1, size_n)).contiguous()
|
||||
|
||||
return pack_cols(q_w, num_bits, size_k, size_n)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_bits", [4])
|
||||
@pytest.mark.parametrize("num_experts", [2, 4, 8])
|
||||
@pytest.mark.parametrize("k_tiles,n_tiles", [(1, 1), (2, 2), (4, 4)])
|
||||
@pytest.mark.parametrize("group_size", [16, 32])
|
||||
def test_awq_marlin_moe_repack_jit_vs_aot(
|
||||
num_bits, num_experts, k_tiles, n_tiles, group_size
|
||||
):
|
||||
if not AOT_AVAILABLE:
|
||||
pytest.skip("sgl_kernel AOT not available")
|
||||
|
||||
tile_k, tile_n = 16, 64
|
||||
size_k = k_tiles * tile_k
|
||||
size_n = n_tiles * tile_n
|
||||
pack_factor = 32 // num_bits
|
||||
|
||||
# Create per-expert AWQ-packed weights
|
||||
b_q_weight = torch.empty(
|
||||
(num_experts, size_k, size_n // pack_factor),
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
)
|
||||
for e in range(num_experts):
|
||||
b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda")
|
||||
w_ref, q_w, s, zp = quantize_weights(
|
||||
b_weight, scalar_types.uint4, group_size, zero_points=True
|
||||
)
|
||||
b_q_weight[e] = awq_pack(q_w, num_bits, size_k, size_n)
|
||||
|
||||
perm = torch.empty((num_experts, 0), dtype=torch.int32, device="cuda")
|
||||
|
||||
out_jit = jit_awq_marlin_moe_repack(b_q_weight, perm, size_k, size_n, num_bits)
|
||||
out_aot = torch.ops.sgl_kernel.awq_marlin_moe_repack.default(
|
||||
b_q_weight, perm, size_k, size_n, num_bits
|
||||
)
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# Bitwise equality
|
||||
torch.testing.assert_close(out_jit, out_aot, rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_bits", [4])
|
||||
@pytest.mark.parametrize("num_experts", [2, 4])
|
||||
@pytest.mark.parametrize("k_tiles,n_tiles", [(1, 1), (2, 2)])
|
||||
@pytest.mark.parametrize("group_size", [16, 32])
|
||||
def test_awq_marlin_moe_repack_shape(
|
||||
num_bits, num_experts, k_tiles, n_tiles, group_size
|
||||
):
|
||||
tile_k, tile_n = 16, 64
|
||||
size_k = k_tiles * tile_k
|
||||
size_n = n_tiles * tile_n
|
||||
pack_factor = 32 // num_bits
|
||||
|
||||
# Create per-expert AWQ-packed weights
|
||||
b_q_weight = torch.empty(
|
||||
(num_experts, size_k, size_n // pack_factor),
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
)
|
||||
for e in range(num_experts):
|
||||
b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda")
|
||||
w_ref, q_w, s, zp = quantize_weights(
|
||||
b_weight, scalar_types.uint4, group_size, zero_points=True
|
||||
)
|
||||
b_q_weight[e] = awq_pack(q_w, num_bits, size_k, size_n)
|
||||
|
||||
perm = torch.empty((num_experts, 0), dtype=torch.int32, device="cuda")
|
||||
|
||||
out = jit_awq_marlin_moe_repack(b_q_weight, perm, size_k, size_n, num_bits)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
assert out.is_cuda and out.dtype == torch.int32
|
||||
expected_shape = (num_experts, size_k // 16, size_n * (num_bits // 2))
|
||||
assert list(out.shape) == list(expected_shape)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,110 @@
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from sgl_kernel.scalar_type import scalar_types
|
||||
|
||||
from sglang.kernels.ops.quantization.awq_marlin_repack import (
|
||||
awq_marlin_repack as jit_awq_marlin_repack,
|
||||
)
|
||||
from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_marlin_utils import get_weight_perm, marlin_weights
|
||||
|
||||
register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
|
||||
|
||||
def _has_aot_awq_marlin_repack() -> bool:
|
||||
return hasattr(torch.ops.sgl_kernel, "awq_marlin_repack") and hasattr(
|
||||
torch.ops.sgl_kernel.awq_marlin_repack, "default"
|
||||
)
|
||||
|
||||
|
||||
AOT_AVAILABLE = _has_aot_awq_marlin_repack()
|
||||
|
||||
|
||||
def awq_pack(
|
||||
q_w: torch.Tensor,
|
||||
num_bits: int,
|
||||
size_k: int,
|
||||
size_n: int,
|
||||
):
|
||||
assert q_w.shape == (size_k, size_n)
|
||||
|
||||
if num_bits == 4:
|
||||
interleave = np.array([0, 2, 4, 6, 1, 3, 5, 7])
|
||||
elif num_bits == 8:
|
||||
interleave = np.array([0, 2, 1, 3])
|
||||
else:
|
||||
raise Exception("num_bits must be 4 or 8, got {}".format(num_bits))
|
||||
|
||||
q_w = q_w.reshape((-1, len(interleave)))[:, interleave].ravel()
|
||||
q_w = q_w.reshape((-1, size_n)).contiguous()
|
||||
|
||||
return pack_cols(q_w, num_bits, size_k, size_n)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_bits", [4, 8])
|
||||
@pytest.mark.parametrize("k_tiles,n_tiles", [(1, 1), (2, 2), (4, 4)])
|
||||
@pytest.mark.parametrize("group_size", [16, 32])
|
||||
def test_awq_marlin_repack_jit_vs_aot(num_bits, k_tiles, n_tiles, group_size):
|
||||
if not AOT_AVAILABLE:
|
||||
pytest.skip("sgl_kernel AOT not available")
|
||||
|
||||
tile_k, tile_n = 16, 64
|
||||
size_k = k_tiles * tile_k
|
||||
size_n = n_tiles * tile_n
|
||||
|
||||
b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda")
|
||||
|
||||
w_ref, q_w, s, zp = quantize_weights(
|
||||
b_weight, scalar_types.uint4, group_size, zero_points=True
|
||||
)
|
||||
|
||||
q_w_awq = awq_pack(q_w, num_bits, size_k, size_n)
|
||||
|
||||
out_jit = jit_awq_marlin_repack(q_w_awq, size_k, size_n, num_bits)
|
||||
out_aot = torch.ops.sgl_kernel.awq_marlin_repack.default(
|
||||
q_w_awq, size_k, size_n, num_bits
|
||||
)
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# Bitwise equality
|
||||
torch.testing.assert_close(out_jit, out_aot, rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_bits", [4, 8])
|
||||
@pytest.mark.parametrize("k_tiles,n_tiles", [(1, 1), (2, 2)])
|
||||
@pytest.mark.parametrize("group_size", [16, 32])
|
||||
def test_awq_marlin_repack_correct(num_bits, k_tiles, n_tiles, group_size):
|
||||
tile_k, tile_n = 16, 64
|
||||
size_k = k_tiles * tile_k
|
||||
size_n = n_tiles * tile_n
|
||||
pack_factor = 32 // num_bits
|
||||
|
||||
b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda")
|
||||
|
||||
w_ref, q_w, s, zp = quantize_weights(
|
||||
b_weight, scalar_types.uint4, group_size, zero_points=True
|
||||
)
|
||||
|
||||
q_w_awq = awq_pack(q_w, num_bits, size_k, size_n)
|
||||
|
||||
weight_perm = get_weight_perm(num_bits)
|
||||
q_w_marlin = marlin_weights(q_w, size_k, size_n, num_bits, weight_perm)
|
||||
|
||||
out_gpu = jit_awq_marlin_repack(q_w_awq, size_k, size_n, num_bits)
|
||||
assert out_gpu.is_cuda and out_gpu.dtype == torch.int32
|
||||
|
||||
expected_cols = size_n * tile_k // pack_factor
|
||||
assert list(out_gpu.shape) == [size_k // tile_k, expected_cols]
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
torch.testing.assert_close(out_gpu, q_w_marlin)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,191 @@
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from sgl_kernel.scalar_type import scalar_types
|
||||
|
||||
from sglang.kernels.ops.quantization.gptq_marlin import gptq_marlin_gemm
|
||||
from sglang.srt.layers.quantization.marlin_utils import (
|
||||
check_marlin_supported,
|
||||
marlin_make_workspace,
|
||||
)
|
||||
from sglang.srt.layers.quantization.marlin_utils_fp4 import (
|
||||
apply_fp4_marlin_linear,
|
||||
nvfp4_marlin_process_global_scale,
|
||||
prepare_nvfp4_layer_for_marlin,
|
||||
)
|
||||
from sglang.srt.utils.common import is_sm80_supported, is_sm90_supported
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_marlin_utils import (
|
||||
awq_marlin_quantize,
|
||||
make_nvfp4_weight_and_ref,
|
||||
marlin_quantize,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=13, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
|
||||
MNK_FACTORS = [
|
||||
(1, 1, 1),
|
||||
(1, 4, 8),
|
||||
(13, 17, 67),
|
||||
(257, 13, 11),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("k_chunk", [128])
|
||||
@pytest.mark.parametrize("n_chunk", [64, 256])
|
||||
@pytest.mark.parametrize("quant_type", [scalar_types.uint4, scalar_types.uint4b8])
|
||||
@pytest.mark.parametrize("group_size", [-1, 128])
|
||||
@pytest.mark.parametrize("mnk_factors", MNK_FACTORS)
|
||||
@pytest.mark.parametrize("act_order", [False, True])
|
||||
def test_gptq_marlin_gemm(
|
||||
k_chunk,
|
||||
n_chunk,
|
||||
quant_type,
|
||||
group_size,
|
||||
mnk_factors,
|
||||
act_order,
|
||||
):
|
||||
m_factor, n_factor, k_factor = mnk_factors
|
||||
has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8]
|
||||
|
||||
size_m = m_factor
|
||||
size_k = k_chunk * k_factor
|
||||
size_n = n_chunk * n_factor
|
||||
|
||||
if act_order:
|
||||
if group_size == -1:
|
||||
return
|
||||
if group_size == size_k:
|
||||
return
|
||||
if has_zp:
|
||||
return
|
||||
|
||||
if size_k % group_size != 0:
|
||||
return
|
||||
|
||||
a_input = torch.randn((size_m, size_k), dtype=torch.float16, device="cuda")
|
||||
b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda")
|
||||
|
||||
if has_zp:
|
||||
w_ref, marlin_q_w, marlin_s, marlin_zp = awq_marlin_quantize(
|
||||
b_weight, quant_type, group_size
|
||||
)
|
||||
g_idx = None
|
||||
sort_indices = None
|
||||
marlin_s2 = None
|
||||
else:
|
||||
w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, _ = marlin_quantize(
|
||||
b_weight, quant_type, group_size, act_order
|
||||
)
|
||||
marlin_zp = None
|
||||
marlin_s2 = None
|
||||
|
||||
workspace = marlin_make_workspace(w_ref.device)
|
||||
|
||||
output = gptq_marlin_gemm(
|
||||
a_input,
|
||||
None,
|
||||
marlin_q_w,
|
||||
marlin_s,
|
||||
marlin_s2,
|
||||
marlin_zp,
|
||||
g_idx,
|
||||
sort_indices,
|
||||
workspace,
|
||||
quant_type,
|
||||
a_input.shape[0],
|
||||
b_weight.shape[1],
|
||||
a_input.shape[1],
|
||||
is_k_full=True,
|
||||
use_atomic_add=False,
|
||||
use_fp32_reduce=False,
|
||||
is_zp_float=False,
|
||||
)
|
||||
|
||||
output_ref = torch.matmul(a_input, w_ref)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# JIT kernel should produce approximately correct results vs torch.matmul
|
||||
max_diff = torch.mean(torch.abs(output - output_ref)) / torch.mean(
|
||||
torch.abs(output_ref)
|
||||
)
|
||||
assert max_diff < 0.04
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not (is_sm80_supported() or is_sm90_supported()),
|
||||
reason="NVFP4 Marlin fallback tests require CUDA SM8X/SM9X",
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
def test_nvfp4_marlin_support_and_scale_transforms_sm80_sm90(dtype):
|
||||
major, minor = torch.cuda.get_device_capability()
|
||||
capability = major * 10 + minor
|
||||
assert check_marlin_supported(
|
||||
scalar_types.float4_e2m1f,
|
||||
group_size=16,
|
||||
has_zp=False,
|
||||
device_capability=capability,
|
||||
)
|
||||
|
||||
global_scale = torch.tensor(1.0, dtype=dtype, device="cuda")
|
||||
actual_global_scale = nvfp4_marlin_process_global_scale(global_scale)
|
||||
assert actual_global_scale.is_cuda
|
||||
assert actual_global_scale.ndim == 1
|
||||
assert actual_global_scale.numel() == 1
|
||||
if dtype == torch.float16:
|
||||
assert actual_global_scale.item() == 128.0
|
||||
else:
|
||||
assert actual_global_scale.item() == 2.0**119
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not (is_sm80_supported() or is_sm90_supported()),
|
||||
reason="NVFP4 Marlin dense numeric test requires CUDA SM80, SM86, or SM90",
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
def test_nvfp4_marlin_dense_matches_dequant_reference(dtype):
|
||||
torch.manual_seed(0)
|
||||
|
||||
size_m = 17
|
||||
size_k = 256
|
||||
size_n = 192
|
||||
group_size = 16
|
||||
|
||||
a_input = torch.randn((size_m, size_k), dtype=dtype, device="cuda") / 10
|
||||
fp4_weight, scales, global_scale, weight_ref = make_nvfp4_weight_and_ref(
|
||||
size_n, size_k, dtype, group_size=group_size
|
||||
)
|
||||
|
||||
layer = torch.nn.Module()
|
||||
layer.quant_config = SimpleNamespace(group_size=group_size)
|
||||
layer.output_size_per_partition = size_n
|
||||
layer.input_size_per_partition = size_k
|
||||
layer.params_dtype = dtype
|
||||
layer.weight = torch.nn.Parameter(fp4_weight, requires_grad=False)
|
||||
layer.weight_scale = torch.nn.Parameter(scales, requires_grad=False)
|
||||
layer.weight_global_scale = torch.nn.Parameter(
|
||||
global_scale.reshape(1), requires_grad=False
|
||||
)
|
||||
prepare_nvfp4_layer_for_marlin(layer)
|
||||
|
||||
output = apply_fp4_marlin_linear(
|
||||
a_input,
|
||||
layer.weight,
|
||||
layer.weight_scale,
|
||||
layer.weight_global_scale,
|
||||
layer.workspace,
|
||||
size_n,
|
||||
size_k,
|
||||
use_fp32_reduce=True,
|
||||
)
|
||||
|
||||
output_ref = torch.matmul(a_input, weight_ref.T)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
torch.testing.assert_close(output, output_ref, rtol=0.04, atol=0.04)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,95 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from sgl_kernel.scalar_type import scalar_types
|
||||
|
||||
from sglang.kernels.ops.quantization.gptq_marlin_repack import gptq_marlin_repack
|
||||
from sglang.srt.layers.quantization.utils import (
|
||||
gptq_quantize_weights,
|
||||
pack_rows,
|
||||
sort_weights,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_marlin_utils import get_weight_perm, marlin_weights
|
||||
|
||||
register_cuda_ci(est_time=16, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
|
||||
MARLIN_K_CHUNKS = [128]
|
||||
MARLIN_N_CHUNKS = [64, 256]
|
||||
|
||||
MNK_FACTORS = [
|
||||
(1, 1, 1),
|
||||
(1, 4, 8),
|
||||
(1, 7, 5),
|
||||
(13, 17, 67),
|
||||
(26, 37, 13),
|
||||
(67, 13, 11),
|
||||
(257, 13, 11),
|
||||
(658, 13, 11),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("k_chunk", MARLIN_K_CHUNKS)
|
||||
@pytest.mark.parametrize("n_chunk", MARLIN_N_CHUNKS)
|
||||
@pytest.mark.parametrize("quant_type", [scalar_types.uint4b8])
|
||||
@pytest.mark.parametrize("group_size", [-1, 32, 64, 128])
|
||||
@pytest.mark.parametrize("act_order", [False, True])
|
||||
@pytest.mark.parametrize("mnk_factors", MNK_FACTORS)
|
||||
def test_gptq_marlin_repack(
|
||||
k_chunk, n_chunk, quant_type, group_size, act_order, mnk_factors
|
||||
):
|
||||
m_factor, n_factor, k_factor = mnk_factors
|
||||
|
||||
size_k = k_chunk * k_factor
|
||||
size_n = n_chunk * n_factor
|
||||
|
||||
# Filter act_order
|
||||
if act_order:
|
||||
if group_size == -1:
|
||||
return
|
||||
if group_size == size_k:
|
||||
return
|
||||
|
||||
# Normalize group_size
|
||||
if group_size == -1:
|
||||
group_size = size_k
|
||||
assert group_size <= size_k
|
||||
|
||||
if size_k % group_size != 0:
|
||||
pytest.skip("size_k must be divisible by group_size")
|
||||
|
||||
# Create input
|
||||
b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda")
|
||||
|
||||
# Quantize (and apply act_order if provided)
|
||||
w_ref, q_w, s, g_idx, rand_perm = gptq_quantize_weights(
|
||||
b_weight, quant_type, group_size, act_order
|
||||
)
|
||||
|
||||
q_w_gptq = pack_rows(q_w, quant_type.size_bits, size_k, size_n)
|
||||
|
||||
# For act_order, sort the "weights" and "g_idx" so that group ids are
|
||||
# increasing
|
||||
sort_indices = torch.empty(0, dtype=torch.int, device=b_weight.device)
|
||||
if act_order:
|
||||
q_w, g_idx, sort_indices = sort_weights(q_w, g_idx)
|
||||
|
||||
marlin_layout_perm = get_weight_perm(quant_type.size_bits)
|
||||
q_w_marlin_ref = marlin_weights(
|
||||
q_w, size_k, size_n, quant_type.size_bits, marlin_layout_perm
|
||||
)
|
||||
|
||||
# Run JIT repack kernel
|
||||
jit_output = gptq_marlin_repack(
|
||||
q_w_gptq, sort_indices, size_k, size_n, quant_type.size_bits
|
||||
)
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# JIT should match the reference (computed from CPU marlin_weights)
|
||||
torch.testing.assert_close(jit_output, q_w_marlin_ref)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,152 @@
|
||||
import random
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.kernels.ops.quantization.mxfp8 import (
|
||||
es_sm100_mxfp8_blockscaled_grouped_quant,
|
||||
es_sm100_mxfp8_blockscaled_moe_grouped_gemm,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=5, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
|
||||
|
||||
def align(val: int, alignment: int = 128) -> int:
|
||||
return int((val + alignment - 1) // alignment * alignment)
|
||||
|
||||
|
||||
# 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 is_sm100_supported(device=None) -> bool:
|
||||
return (torch.cuda.get_device_capability(device)[0] == 10) and (
|
||||
torch.version.cuda >= "12.8"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not is_sm100_supported(),
|
||||
reason="test_mxfp8_moe at jit kernen is only supported on sm100",
|
||||
)
|
||||
@pytest.mark.parametrize("num_experts", [8, 16, 32, 64])
|
||||
@pytest.mark.parametrize("out_dtype", [torch.half, torch.bfloat16])
|
||||
def test_es_sm100_mxfp8_blockscaled_grouped_mm(num_experts, out_dtype):
|
||||
device = "cuda"
|
||||
alignment = 128
|
||||
n_g = random.randint(1, 64) * alignment
|
||||
k_g = random.randint(1, 64) * alignment
|
||||
|
||||
expert_offset = 0
|
||||
expert_offsets = []
|
||||
aux_expert_offset = 0
|
||||
aux_expert_offsets = []
|
||||
a_blockscale_offset = 0
|
||||
a_blockscale_offsets = []
|
||||
b_blockscale_offset = 0
|
||||
b_blockscale_offsets = []
|
||||
a_list = []
|
||||
b_list = []
|
||||
ref_d_list = []
|
||||
tokens_per_expert = []
|
||||
|
||||
for g in range(num_experts):
|
||||
m_g = random.randint(1, 512)
|
||||
tokens_per_expert.append(m_g)
|
||||
expert_offsets.append(expert_offset)
|
||||
expert_offset += m_g
|
||||
aux_expert_offsets.append(aux_expert_offset)
|
||||
aux_expert_offset += n_g
|
||||
a_blockscale_offsets.append(a_blockscale_offset)
|
||||
a_blockscale_offset += align(m_g, 128)
|
||||
b_blockscale_offsets.append(b_blockscale_offset)
|
||||
b_blockscale_offset += n_g # n_g already align to 128
|
||||
|
||||
a = torch.normal(
|
||||
0.0, std=1.0, size=(m_g, k_g), device=device, dtype=out_dtype
|
||||
) # (M, K):(K, 1)
|
||||
b = torch.normal(
|
||||
0.0, std=1.0, size=(n_g, k_g), device=device, dtype=out_dtype
|
||||
) # (N, K):(K, 1)
|
||||
|
||||
a_list.append(a)
|
||||
b_list.append(b)
|
||||
ref_d = a @ b.T
|
||||
ref_d_list.append(ref_d)
|
||||
a = torch.concat(a_list, dim=0)
|
||||
b = torch.concat(b_list, dim=0)
|
||||
|
||||
_expert_offsets = torch.tensor(expert_offsets).to(device=device, dtype=torch.int32)
|
||||
_aux_expert_offsets = torch.tensor(aux_expert_offsets).to(
|
||||
device=device, dtype=torch.int32
|
||||
)
|
||||
_a_blockscale_offsets = torch.tensor(a_blockscale_offsets).to(
|
||||
device=device, dtype=torch.int32
|
||||
)
|
||||
_b_blockscale_offsets = torch.tensor(b_blockscale_offsets).to(
|
||||
device=device, dtype=torch.int32
|
||||
)
|
||||
|
||||
a_quant = torch.zeros_like(a, dtype=torch.float8_e4m3fn, device=device)
|
||||
a_scale_factor = torch.zeros(
|
||||
(a_blockscale_offset, k_g // 32), dtype=torch.uint8, device=device
|
||||
)
|
||||
|
||||
b_quant = torch.zeros_like(b, dtype=torch.float8_e4m3fn, device=device)
|
||||
b_scale_factor = torch.zeros(
|
||||
(num_experts * n_g, k_g // 32), dtype=torch.uint8, device=device
|
||||
)
|
||||
tokens_per_expert = torch.tensor(tokens_per_expert).to(
|
||||
device=device, dtype=torch.int32
|
||||
)
|
||||
workspace = torch.empty((1024, 1024, 1024), dtype=torch.uint8, device=device)
|
||||
|
||||
es_sm100_mxfp8_blockscaled_grouped_quant(
|
||||
a,
|
||||
tokens_per_expert,
|
||||
_expert_offsets,
|
||||
_a_blockscale_offsets,
|
||||
a_quant,
|
||||
a_scale_factor,
|
||||
)
|
||||
es_sm100_mxfp8_blockscaled_grouped_quant(
|
||||
b,
|
||||
torch.ones_like(tokens_per_expert) * n_g,
|
||||
_aux_expert_offsets,
|
||||
_b_blockscale_offsets,
|
||||
b_quant,
|
||||
b_scale_factor,
|
||||
)
|
||||
|
||||
b_quant = b_quant.view(num_experts, n_g, k_g)
|
||||
b_scale_factor = b_scale_factor.view(num_experts, n_g, k_g // 32)
|
||||
d = es_sm100_mxfp8_blockscaled_moe_grouped_gemm(
|
||||
b_quant,
|
||||
a_quant,
|
||||
b_scale_factor,
|
||||
a_scale_factor,
|
||||
_expert_offsets,
|
||||
_a_blockscale_offsets,
|
||||
tokens_per_expert,
|
||||
workspace,
|
||||
a.dtype,
|
||||
)
|
||||
|
||||
for g in range(num_experts):
|
||||
baseline = ref_d_list[g]
|
||||
actual = d[expert_offsets[g] : (expert_offsets[g] + tokens_per_expert[g])]
|
||||
diff = calc_diff(actual, baseline)
|
||||
assert diff < 0.001
|
||||
print(
|
||||
f"m_g={baseline.shape[0]} n_g={n_g} k_g={k_g} num_experts={num_experts}, out_dtype={out_dtype}, diff={diff:.5f}: OK"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -0,0 +1,103 @@
|
||||
import itertools
|
||||
import sys
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.kernels.jit.utils import get_ci_test_range
|
||||
from sglang.kernels.ops.quantization._jit_per_tensor_quant_fp8 import (
|
||||
per_tensor_quant_fp8,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=16, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
# Nightly is not redundant here: it sets SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1 to expand get_ci_test_range sweeps.
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
try:
|
||||
from sglang.srt.utils import is_hip
|
||||
|
||||
_is_hip = is_hip()
|
||||
except ImportError:
|
||||
_is_hip = False
|
||||
|
||||
fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
|
||||
|
||||
|
||||
def sglang_scaled_fp8_quant(
|
||||
input: torch.Tensor,
|
||||
scale: Optional[torch.Tensor] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
fp8_type_: torch.dtype = torch.float8_e4m3fn
|
||||
output = torch.empty_like(input, device=input.device, dtype=fp8_type_)
|
||||
is_static = True
|
||||
if scale is None:
|
||||
scale = torch.zeros(1, device=input.device, dtype=torch.float32)
|
||||
is_static = False
|
||||
per_tensor_quant_fp8(input, output, scale, is_static)
|
||||
|
||||
return output, scale
|
||||
|
||||
|
||||
def torch_scaled_fp8_quant(tensor, inv_scale):
|
||||
finfo = torch.finfo(torch.float8_e4m3fn)
|
||||
scale = inv_scale.reciprocal()
|
||||
qweight = (tensor.to(torch.float32) * scale).clamp(min=finfo.min, max=finfo.max)
|
||||
qweight = qweight.to(torch.float8_e4m3fn)
|
||||
return qweight
|
||||
|
||||
|
||||
PER_TENSOR_QUANT_CASES = get_ci_test_range(
|
||||
list(
|
||||
itertools.product(
|
||||
[38, 39, 128, 256, 512, 1392],
|
||||
[512, 1536, 2048, 4096, 7168],
|
||||
)
|
||||
),
|
||||
[(38, 7168), (39, 1536), (128, 512), (512, 4096), (1392, 1536)],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens,hidden_dim", PER_TENSOR_QUANT_CASES)
|
||||
def test_jit_per_tensor_quant_compare_implementations(
|
||||
num_tokens: int,
|
||||
hidden_dim: int,
|
||||
):
|
||||
device = torch.device("cuda")
|
||||
x = torch.rand((num_tokens, hidden_dim), dtype=torch.float16, device=device)
|
||||
|
||||
sglang_out, sglang_scale = sglang_scaled_fp8_quant(x)
|
||||
torch_out = torch_scaled_fp8_quant(x, sglang_scale)
|
||||
|
||||
torch.testing.assert_close(
|
||||
sglang_out.float(), torch_out.float(), rtol=1e-3, atol=1e-3
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", [(4, 8, 64), (2, 16, 128), (19260817, 1, 1)])
|
||||
def test_jit_per_tensor_quant_supports_3d(shape):
|
||||
device = torch.device("cuda")
|
||||
x = torch.rand(shape, dtype=torch.bfloat16, device=device)
|
||||
out = torch.empty_like(x, device=x.device, dtype=fp8_type_)
|
||||
scale = torch.zeros(1, device=x.device, dtype=torch.float32)
|
||||
|
||||
per_tensor_quant_fp8(x, out, scale, is_static=False)
|
||||
|
||||
x_2d = x.flatten(0, -2)
|
||||
out_ref_2d = torch_scaled_fp8_quant(x_2d, scale)
|
||||
out_ref = out_ref_2d.reshape(shape)
|
||||
|
||||
torch.testing.assert_close(out.float(), out_ref.float(), rtol=1e-3, atol=1e-3)
|
||||
|
||||
scale = torch.rand(1, dtype=torch.float32, device=device)
|
||||
sglang_out, _ = sglang_scaled_fp8_quant(x, scale)
|
||||
torch_out = torch_scaled_fp8_quant(x, scale)
|
||||
|
||||
torch.testing.assert_close(
|
||||
sglang_out.float(), torch_out.float(), rtol=1e-3, atol=1e-3
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,475 @@
|
||||
"""Correctness tests for the trait-driven per_token_group_quant JIT kernel.
|
||||
|
||||
The reference is computed in pure PyTorch (the quantization math itself), NOT by
|
||||
calling the v2 / minimax kernels -- those are being deprecated, so the tests
|
||||
must outlive them.
|
||||
|
||||
Two guard strengths, chosen by what the kernel's numerics can actually pin:
|
||||
- UE8M0 paths: the quant multiplier is an exact power of two (a bit shift, no
|
||||
division), so codes and packed exponent bytes are compared BIT-EXACT
|
||||
against the torch reference. These are the production paths (DeepGEMM dense,
|
||||
EP-MoE), so this is where bit-exactness matters.
|
||||
- fp32 / int8 scale paths: the kernel divides under ``--use_fast_math`` (fast
|
||||
reciprocal), so codes are not bit-reproducible from an exact torch divide.
|
||||
Those tests pin the exactly-reproducible parts -- the stored scale (a single
|
||||
multiply) -- and the dequant round-trip error, which is what downstream
|
||||
actually consumes.
|
||||
"""
|
||||
|
||||
import itertools
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.kernels.jit.utils import get_ci_test_range
|
||||
from sglang.kernels.ops.quantization._jit_per_token_group_quant import (
|
||||
per_token_group_quant,
|
||||
)
|
||||
from sglang.kernels.ops.quantization.fp8_kernel import (
|
||||
create_per_token_group_quant_fp8_output_scale,
|
||||
fp8_dtype,
|
||||
fp8_max,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=90, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
register_cuda_ci(est_time=90, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
|
||||
|
||||
G = 128
|
||||
FMAX = float(fp8_max) # 448 for e4m3
|
||||
I8_MAX, I8_MIN = 127.0, -128.0
|
||||
EPS = 1e-10
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Pure-torch references (match the kernel's expression order).
|
||||
# --------------------------------------------------------------------------- #
|
||||
def _group_amax(x: torch.Tensor, gs: int) -> torch.Tensor:
|
||||
"""Per-group absmax over the last dim, floored at EPS. Returns [..., ng]."""
|
||||
xf = x.float().unflatten(-1, (-1, gs))
|
||||
return xf.abs().amax(-1).clamp_min(EPS)
|
||||
|
||||
|
||||
def _quantize(x: torch.Tensor, gs: int, quant_scale: torch.Tensor, out_dtype, lo, hi):
|
||||
xf = x.float().unflatten(-1, (-1, gs))
|
||||
q = (xf * quant_scale.unsqueeze(-1)).clamp(lo, hi).to(out_dtype)
|
||||
return q.flatten(-2)
|
||||
|
||||
|
||||
def ref_fp8_fp32_scale(x, gs):
|
||||
"""fp8 codes + fp32 stored scale (scale = amax / FMAX, a single multiply)."""
|
||||
amax = _group_amax(x, gs)
|
||||
scale_inv = amax * (1.0 / FMAX)
|
||||
q = _quantize(x, gs, FMAX / amax, fp8_dtype, -FMAX, FMAX)
|
||||
return q, scale_inv
|
||||
|
||||
|
||||
def ref_int8(x, gs):
|
||||
amax = _group_amax(x, gs)
|
||||
scale_inv = amax * (1.0 / I8_MAX)
|
||||
q = _quantize(x, gs, I8_MAX / amax, torch.int8, I8_MIN, I8_MAX)
|
||||
return q, scale_inv
|
||||
|
||||
|
||||
def ref_fp8_ue8m0(x, gs):
|
||||
"""fp8 codes + UE8M0 exponent bytes [..., ng]. The multiplier 2^-e is exact
|
||||
in fp32, so codes are bit-reproducible (unlike the fp32-scale path)."""
|
||||
amax = _group_amax(x, gs)
|
||||
raw = (amax / FMAX).contiguous()
|
||||
bits = raw.view(torch.int32)
|
||||
exp = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(
|
||||
torch.int32
|
||||
) # ceil to ue8m0
|
||||
quant_scale = ((127 + 127 - exp) << 23).view(torch.float32) # 2^(127 - (exp-127))
|
||||
q = _quantize(x, gs, quant_scale, fp8_dtype, -FMAX, FMAX)
|
||||
return q, exp.to(torch.uint8)
|
||||
|
||||
|
||||
def _decode_packed_exp(s_int32: torch.Tensor, ng: int) -> torch.Tensor:
|
||||
"""Decode an int32 packed-UE8M0 scale (logical [..., ceil(ng/4)]) to the
|
||||
[..., ng] exponent grid, independent of the physical (row/col-major) layout:
|
||||
exponent[..., g] = byte (g % 4) of int32[..., g // 4]."""
|
||||
g = torch.arange(ng, device=s_int32.device)
|
||||
col = s_int32.index_select(-1, g // 4)
|
||||
return ((col >> (8 * (g % 4))) & 0xFF).to(torch.uint8)
|
||||
|
||||
|
||||
def _dequant_rel_err(q, scale_inv, x, gs) -> float:
|
||||
deq = (q.float().unflatten(-1, (-1, gs)) * scale_inv.unsqueeze(-1)).flatten(-2)
|
||||
return ((x.float() - deq).abs() / (x.float().abs() + 1e-6)).mean().item()
|
||||
|
||||
|
||||
def _packed_exp_to_dequant_scale(x_s, ng) -> torch.Tensor:
|
||||
"""Decode a packed-UE8M0 scale buffer to the fp32 dequant scale 2^(e-127)."""
|
||||
exp = _decode_packed_exp(x_s, ng).to(torch.int32)
|
||||
return torch.exp2(exp.float() - 127.0)
|
||||
|
||||
|
||||
def _alloc_scale(x_shape, *, column_major, scale_ue8m0):
|
||||
s = create_per_token_group_quant_fp8_output_scale(
|
||||
x_shape=x_shape,
|
||||
device="cuda",
|
||||
group_size=G,
|
||||
column_major_scales=column_major,
|
||||
scale_tma_aligned=column_major,
|
||||
scale_ue8m0=scale_ue8m0,
|
||||
)
|
||||
s.zero_()
|
||||
return s
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# UE8M0 paths: bit-exact vs the torch reference.
|
||||
# --------------------------------------------------------------------------- #
|
||||
# hidden 768 (Qwen3-30B-A3B moe_intermediate: 6 groups) exercises the non-4
|
||||
# aligned col-packed tail; 128 is a single group.
|
||||
UE8M0_CASES = get_ci_test_range(
|
||||
list(
|
||||
itertools.product(
|
||||
[torch.bfloat16, torch.float16],
|
||||
[1, 7, 38, 333],
|
||||
[128, 768, 2048, 7168],
|
||||
)
|
||||
),
|
||||
[
|
||||
(torch.bfloat16, 1, 128),
|
||||
(torch.bfloat16, 38, 768),
|
||||
(torch.bfloat16, 333, 7168),
|
||||
(torch.float16, 7, 2048),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype,num_tokens,hidden", UE8M0_CASES)
|
||||
def test_ue8m0_bitexact(dtype, num_tokens, hidden):
|
||||
"""Col-major packed UE8M0: fp8 codes and decoded exponent bytes are
|
||||
bit-exact with the torch reference (exact pow-2 multiplier). Covers the
|
||||
aligned and non-4-aligned (hidden=768) pack-tail layouts."""
|
||||
torch.manual_seed(hidden * 10 + num_tokens)
|
||||
x = torch.randn(num_tokens, hidden, device="cuda", dtype=dtype)
|
||||
q_ref, exp_ref = ref_fp8_ue8m0(x, G)
|
||||
|
||||
x_q = torch.zeros_like(x, dtype=fp8_dtype)
|
||||
x_s = _alloc_scale((num_tokens, hidden), column_major=True, scale_ue8m0=True)
|
||||
per_token_group_quant(x, x_q, x_s, G, scale_ue8m0=True)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
assert torch.equal(x_q.view(torch.int8), q_ref.view(torch.int8)), "codes differ"
|
||||
exp = _decode_packed_exp(x_s, hidden // G)
|
||||
assert torch.equal(exp, exp_ref), "exponent bytes differ"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("group_size", get_ci_test_range([16, 32, 64, 128], [16, 64]))
|
||||
def test_ue8m0_group_sizes(group_size):
|
||||
"""Group size is a template axis (v2 dispatched a runtime switch). Each size
|
||||
maps a group onto a different subwarp lane count; codes/exponents must stay
|
||||
bit-exact -- a wrong lane span would fold the wrong elements into absmax."""
|
||||
torch.manual_seed(group_size)
|
||||
num_tokens, hidden = 9, 4096
|
||||
x = torch.randn(num_tokens, hidden, device="cuda", dtype=torch.bfloat16)
|
||||
q_ref, exp_ref = ref_fp8_ue8m0(x, group_size)
|
||||
|
||||
x_q = torch.zeros_like(x, dtype=fp8_dtype)
|
||||
x_s = create_per_token_group_quant_fp8_output_scale(
|
||||
x_shape=(num_tokens, hidden),
|
||||
device="cuda",
|
||||
group_size=group_size,
|
||||
column_major_scales=True,
|
||||
scale_tma_aligned=True,
|
||||
scale_ue8m0=True,
|
||||
)
|
||||
x_s.zero_()
|
||||
per_token_group_quant(x, x_q, x_s, group_size, scale_ue8m0=True)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
assert torch.equal(x_q.view(torch.int8), q_ref.view(torch.int8)), "codes differ"
|
||||
exp = _decode_packed_exp(x_s, hidden // group_size)
|
||||
assert torch.equal(exp, exp_ref), "exponent bytes differ"
|
||||
|
||||
|
||||
# hidden 4096 -> 32 groups (aligned); 768 -> 6 groups (6 % 4 = 2, unaligned:
|
||||
# the last int32 holds 2 real exponent bytes + 2 zero-padded tail bytes).
|
||||
@pytest.mark.parametrize("hidden", [4096, 768])
|
||||
def test_ue8m0_row_packed_bitexact(hidden):
|
||||
"""Row-major packed UE8M0 (int32 [T, ceil(G/4)] contiguous, the minimax
|
||||
layout): bit-exact vs the torch reference. The unaligned hidden exercises
|
||||
the row-major pack-tail zeroing (fill_unaligned)."""
|
||||
torch.manual_seed(hidden)
|
||||
num_tokens = 17
|
||||
x = torch.randn(num_tokens, hidden, device="cuda", dtype=torch.bfloat16)
|
||||
q_ref, exp_ref = ref_fp8_ue8m0(x, G)
|
||||
|
||||
x_q = torch.zeros_like(x, dtype=fp8_dtype)
|
||||
x_s = torch.zeros(
|
||||
num_tokens, (hidden // G + 3) // 4, device="cuda", dtype=torch.int32
|
||||
)
|
||||
per_token_group_quant(x, x_q, x_s, G, scale_ue8m0=True)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
assert torch.equal(x_q.view(torch.int8), q_ref.view(torch.int8)), "codes differ"
|
||||
exp = _decode_packed_exp(x_s, hidden // G)
|
||||
assert torch.equal(exp, exp_ref), "exponent bytes differ"
|
||||
# unaligned tail bytes of the last int32 must be zero-padded, not garbage.
|
||||
ng = hidden // G
|
||||
if ng % 4:
|
||||
last_bytes = x_s[:, -1].contiguous().view(torch.uint8).view(num_tokens, 4)
|
||||
assert torch.all(last_bytes[:, ng % 4 :] == 0), "pack-tail bytes not zeroed"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# fp32 / int8 scale paths: exact stored scale + dequant round-trip (the codes
|
||||
# are not bit-reproducible under fast-math division).
|
||||
# --------------------------------------------------------------------------- #
|
||||
@pytest.mark.parametrize("hidden", [4096, 768])
|
||||
@pytest.mark.parametrize("column_major", [False, True])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
def test_fp32_scale(dtype, column_major, hidden):
|
||||
"""fp32 scale (row-major contiguous / col-major TMA view): the stored scale
|
||||
is amax/FMAX (a single multiply, bit-exact) and dequant round-trips within
|
||||
fp8 error.
|
||||
|
||||
hidden=768 (6 groups, ng % 4 != 0) is a bug regression: the host check
|
||||
used to apply the ue8m0 pack-tail alignment requirement to fp32 scales,
|
||||
which have no packing, and rejected this shape outright."""
|
||||
torch.manual_seed(int(column_major) + 2 * (dtype == torch.float16))
|
||||
num_tokens = 128
|
||||
x = torch.randn(num_tokens, hidden, device="cuda", dtype=dtype)
|
||||
_, scale_ref = ref_fp8_fp32_scale(x, G)
|
||||
|
||||
x_q = torch.zeros_like(x, dtype=fp8_dtype)
|
||||
x_s = _alloc_scale(
|
||||
(num_tokens, hidden), column_major=column_major, scale_ue8m0=False
|
||||
)
|
||||
per_token_group_quant(x, x_q, x_s, G)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
torch.testing.assert_close(x_s, scale_ref, rtol=0, atol=0)
|
||||
assert _dequant_rel_err(x_q, x_s, x, G) < 0.05
|
||||
|
||||
|
||||
@pytest.mark.parametrize("column_major", [False, True])
|
||||
def test_int8_scale(column_major):
|
||||
"""int8 output (row-major / col-major fp32 scale): exact stored scale
|
||||
(amax/127) + dequant round-trip. Pins the multiply-by-inverse family (v1
|
||||
divided by the scale and differed by a ULP) without depending on v2."""
|
||||
torch.manual_seed(2 + int(column_major))
|
||||
num_tokens, hidden = 33, 4096
|
||||
x = torch.randn(num_tokens, hidden, device="cuda", dtype=torch.bfloat16)
|
||||
_, scale_ref = ref_int8(x, G)
|
||||
|
||||
x_q = torch.zeros(num_tokens, hidden, device="cuda", dtype=torch.int8)
|
||||
x_s = _alloc_scale(
|
||||
(num_tokens, hidden), column_major=column_major, scale_ue8m0=False
|
||||
)
|
||||
per_token_group_quant(x, x_q, x_s, G)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
torch.testing.assert_close(x_s, scale_ref, rtol=0, atol=0)
|
||||
# int8 group-quant is coarser than fp8 (127 vs 448 levels), so its mean
|
||||
# relative round-trip error on randn data sits a little above the fp8 0.05.
|
||||
assert _dequant_rel_err(x_q, x_s, x, G) < 0.08
|
||||
|
||||
|
||||
def test_group_size_256_roundtrip():
|
||||
"""Group 256 (32 lanes on H100, 16 on Blackwell) is above v2's old cap of
|
||||
128. Pin the derived property: the fp32 scale equals absmax/FMAX and dequant
|
||||
round-trips. A mis-mapped wide subwarp would fold the wrong elements into
|
||||
the group absmax and move the scale."""
|
||||
torch.manual_seed(256)
|
||||
num_tokens, hidden, gs = 9, 4096, 256
|
||||
x = torch.randn(num_tokens, hidden, device="cuda", dtype=torch.bfloat16)
|
||||
_, scale_ref = ref_fp8_fp32_scale(x, gs)
|
||||
|
||||
x_q = torch.zeros_like(x, dtype=fp8_dtype)
|
||||
x_s = torch.zeros(num_tokens, hidden // gs, device="cuda", dtype=torch.float32)
|
||||
per_token_group_quant(x, x_q, x_s, gs)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
torch.testing.assert_close(x_s, scale_ref, rtol=0, atol=0)
|
||||
assert _dequant_rel_err(x_q, x_s, x, gs) < 0.05
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Fused silu+mul.
|
||||
# --------------------------------------------------------------------------- #
|
||||
def _ref_silu_mul(x, hidden):
|
||||
"""silu in fp32, round to the input dtype, multiply in the input dtype --
|
||||
matching the kernel's fused path exactly."""
|
||||
gate, up = x[..., :hidden], x[..., hidden:]
|
||||
return torch.nn.functional.silu(gate.float()).to(x.dtype) * up
|
||||
|
||||
|
||||
@pytest.mark.parametrize("column_major", [True, False])
|
||||
@pytest.mark.parametrize("scale_ue8m0", [True, False])
|
||||
def test_fused_silu(scale_ue8m0, column_major):
|
||||
"""fuse_silu_and_mul quantizes ``silu(x[..., :h]) * x[..., h:]`` (SGLang's
|
||||
SiluAndMul: first half is the gated half). Covered across all four scale
|
||||
layouts so the fused [gate | up] input layout is pinned everywhere.
|
||||
|
||||
The kernel's silu uses the fast ``__tanhf`` intrinsic on Blackwell, which
|
||||
is not bit-reproducible from torch's sigmoid-based silu, so this is a
|
||||
property test: dequant the kernel output through its own scale and check it
|
||||
round-trips to the torch activation within fp8 error. (The quant math is
|
||||
pinned bit-exact by the non-fused ue8m0 tests; a wrong gate/up split or
|
||||
offset would move the round-trip well past tolerance.)"""
|
||||
torch.manual_seed(int(scale_ue8m0) * 2 + int(column_major))
|
||||
num_tokens, hidden = 37, 4096
|
||||
x = torch.randn(num_tokens, hidden * 2, device="cuda", dtype=torch.bfloat16)
|
||||
act = _ref_silu_mul(x, hidden)
|
||||
|
||||
x_q = torch.zeros(num_tokens, hidden, device="cuda", dtype=fp8_dtype)
|
||||
if scale_ue8m0 and not column_major:
|
||||
x_s = torch.zeros(
|
||||
num_tokens, hidden // G // 4, device="cuda", dtype=torch.int32
|
||||
)
|
||||
else:
|
||||
x_s = _alloc_scale(
|
||||
(num_tokens, hidden), column_major=column_major, scale_ue8m0=scale_ue8m0
|
||||
)
|
||||
per_token_group_quant(
|
||||
x, x_q, x_s, G, scale_ue8m0=scale_ue8m0, fuse_silu_and_mul=True
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
deq_scale = _packed_exp_to_dequant_scale(x_s, hidden // G) if scale_ue8m0 else x_s
|
||||
assert _dequant_rel_err(x_q, deq_scale, act, G) < 0.05
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Masked EP-MoE schedule.
|
||||
# --------------------------------------------------------------------------- #
|
||||
MASKED_CASES = get_ci_test_range(
|
||||
list(itertools.product([2, 5], [2048, 4096], [128, 384])),
|
||||
[(2, 2048, 128), (5, 4096, 384)],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("masked_m_dtype", [torch.int32, torch.int64])
|
||||
@pytest.mark.parametrize("expected_m", [None, 4])
|
||||
@pytest.mark.parametrize("num_experts,hidden,tokens_pad", MASKED_CASES)
|
||||
def test_masked(num_experts, hidden, tokens_pad, expected_m, masked_m_dtype):
|
||||
"""Masked EP-MoE schedule (col-packed ue8m0, plain quant -- no silu, so the
|
||||
quant is bit-reproducible): rows < masked_m[e] are bit-exact vs the torch
|
||||
reference; rows >= masked_m[e] stay zero (untouched). Fusion numerics are
|
||||
covered by test_fused_silu; here the schedule is what's under test.
|
||||
|
||||
masked_m is accepted as int32 or int64 (the latter read as its low word),
|
||||
so both dtypes are exercised.
|
||||
|
||||
expected_m=4 shrinks the grid's token axis far below masked_m, so the
|
||||
grid-stride token loop must still cover every valid token -- guards the
|
||||
host-hint-only contract (a wrong hint can never drop tokens)."""
|
||||
torch.manual_seed(num_experts * 1000 + hidden + tokens_pad)
|
||||
x = torch.randn(
|
||||
num_experts, tokens_pad, hidden, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
masked_m = torch.randint(
|
||||
0, tokens_pad + 1, (num_experts,), device="cuda", dtype=masked_m_dtype
|
||||
)
|
||||
out_shape = (num_experts, tokens_pad, hidden)
|
||||
|
||||
x_q = torch.zeros(out_shape, device="cuda", dtype=fp8_dtype)
|
||||
x_s = _alloc_scale(out_shape, column_major=True, scale_ue8m0=True)
|
||||
per_token_group_quant(
|
||||
x, x_q, x_s, G, scale_ue8m0=True, masked_m=masked_m, expected_m=expected_m
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
q_ref, exp_ref = ref_fp8_ue8m0(x, G)
|
||||
exp = _decode_packed_exp(x_s, hidden // G)
|
||||
for e in range(num_experts):
|
||||
m = int(masked_m[e])
|
||||
assert torch.equal(
|
||||
x_q[e, :m].view(torch.int8), q_ref[e, :m].view(torch.int8)
|
||||
), "written codes differ"
|
||||
assert torch.equal(exp[e, :m], exp_ref[e, :m]), "written exponents differ"
|
||||
assert torch.all(x_q[e, m:].view(torch.int8) == 0), "padding codes touched"
|
||||
|
||||
|
||||
def _as_int32(t: torch.Tensor) -> torch.Tensor:
|
||||
return t.view(torch.int32) if t.dtype == torch.int32 else t
|
||||
|
||||
|
||||
# (out_dtype, column_major_scales, scale_ue8m0); ue8m0 implies fp8 output.
|
||||
AUTO_ALLOC_CASES = [
|
||||
(torch.float8_e4m3fn, True, True),
|
||||
(torch.float8_e4m3fn, False, True),
|
||||
(torch.float8_e4m3fn, True, False),
|
||||
(torch.float8_e4m3fn, False, False),
|
||||
(torch.int8, True, False),
|
||||
(torch.int8, False, False),
|
||||
]
|
||||
|
||||
|
||||
def test_masked_fused():
|
||||
"""The production EP-MoE path: masked schedule + fuse_silu_and_mul +
|
||||
col-packed ue8m0. silu is not bit-reproducible, so check the written rows
|
||||
round-trip to the torch activation and padding rows stay zero."""
|
||||
torch.manual_seed(7)
|
||||
num_experts, tokens_pad, hidden = 3, 256, 2048
|
||||
x = torch.randn(
|
||||
num_experts, tokens_pad, hidden * 2, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
masked_m = torch.randint(
|
||||
0, tokens_pad + 1, (num_experts,), device="cuda", dtype=torch.int32
|
||||
)
|
||||
out_shape = (num_experts, tokens_pad, hidden)
|
||||
|
||||
x_q = torch.zeros(out_shape, device="cuda", dtype=fp8_dtype)
|
||||
x_s = _alloc_scale(out_shape, column_major=True, scale_ue8m0=True)
|
||||
per_token_group_quant(
|
||||
x, x_q, x_s, G, scale_ue8m0=True, fuse_silu_and_mul=True, masked_m=masked_m
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
act = _ref_silu_mul(x, hidden)
|
||||
deq_scale = _packed_exp_to_dequant_scale(x_s, hidden // G)
|
||||
for e in range(num_experts):
|
||||
m = int(masked_m[e])
|
||||
if m > 0:
|
||||
assert _dequant_rel_err(x_q[e, :m], deq_scale[e, :m], act[e, :m], G) < 0.05
|
||||
assert torch.all(x_q[e, m:].view(torch.int8) == 0), "padding touched"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("out_dtype,column_major_scales,scale_ue8m0", AUTO_ALLOC_CASES)
|
||||
def test_auto_allocation(out_dtype, column_major_scales, scale_ue8m0):
|
||||
"""Omitting output_q/output_s allocates them per out_dtype / major mode /
|
||||
scale format and returns (q, s). The auto-allocated run must be bit-
|
||||
identical to quantizing into caller-supplied buffers of the same layout --
|
||||
guards that _allocate_outputs picks the layout the kernel decodes."""
|
||||
torch.manual_seed(int(column_major_scales) * 2 + int(scale_ue8m0))
|
||||
num_tokens, hidden = 38, 2048 # 16 groups, %4 == 0 for row-packed ue8m0
|
||||
x = torch.randn(num_tokens, hidden, device="cuda", dtype=torch.bfloat16)
|
||||
|
||||
q_buf = torch.zeros(num_tokens, hidden, device="cuda", dtype=out_dtype)
|
||||
if scale_ue8m0 and not column_major_scales:
|
||||
s_buf = torch.zeros(
|
||||
num_tokens, hidden // G // 4, device="cuda", dtype=torch.int32
|
||||
)
|
||||
else:
|
||||
s_buf = _alloc_scale(
|
||||
(num_tokens, hidden),
|
||||
column_major=column_major_scales,
|
||||
scale_ue8m0=scale_ue8m0,
|
||||
)
|
||||
per_token_group_quant(x, q_buf, s_buf, G, scale_ue8m0=scale_ue8m0)
|
||||
|
||||
q_auto, s_auto = per_token_group_quant(
|
||||
x,
|
||||
group_size=G,
|
||||
scale_ue8m0=scale_ue8m0,
|
||||
column_major_scales=column_major_scales,
|
||||
out_dtype=out_dtype,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
assert q_auto.dtype == out_dtype and q_auto.shape == x.shape
|
||||
assert s_auto.dtype == s_buf.dtype and s_auto.shape == s_buf.shape
|
||||
assert torch.equal(q_auto.view(torch.int8), q_buf.view(torch.int8)), "codes differ"
|
||||
assert torch.equal(_as_int32(s_auto), _as_int32(s_buf)), "scales differ"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,191 @@
|
||||
import itertools
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.kernels.jit.utils import get_ci_test_range
|
||||
from sglang.kernels.ops.quantization._jit_per_token_group_quant_8bit_v2 import (
|
||||
per_token_group_quant_8bit_v2,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=90, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
|
||||
try:
|
||||
from sgl_kernel import sgl_per_token_group_quant_8bit # AOT v2 reference op
|
||||
except ImportError:
|
||||
sgl_per_token_group_quant_8bit = None
|
||||
|
||||
if sgl_per_token_group_quant_8bit is None and not torch.cuda.is_available():
|
||||
pytest.skip("sgl_kernel AOT reference op is unavailable", allow_module_level=True)
|
||||
if sgl_per_token_group_quant_8bit is None:
|
||||
raise ImportError("sgl_kernel AOT reference op is unavailable")
|
||||
|
||||
from sglang.kernels.ops.quantization.fp8_kernel import ( # noqa: E402
|
||||
create_per_token_group_quant_fp8_output_scale,
|
||||
fp8_dtype,
|
||||
fp8_max,
|
||||
fp8_min,
|
||||
)
|
||||
|
||||
G = 128
|
||||
|
||||
|
||||
def _alloc(x_shape, scale_ue8m0):
|
||||
"""Pre-allocated (zeroed) output_q + output_s for a given input/output shape.
|
||||
Zeroing makes the unwritten (padding / aligned) regions compare equal."""
|
||||
x_q = torch.zeros(x_shape, device="cuda", dtype=fp8_dtype)
|
||||
x_s = create_per_token_group_quant_fp8_output_scale(
|
||||
x_shape=x_shape,
|
||||
device="cuda",
|
||||
group_size=G,
|
||||
column_major_scales=True,
|
||||
scale_tma_aligned=True,
|
||||
scale_ue8m0=scale_ue8m0,
|
||||
)
|
||||
x_s.zero_()
|
||||
return x_q, x_s
|
||||
|
||||
|
||||
V2_QUANT_CASES = get_ci_test_range(
|
||||
list(
|
||||
itertools.product(
|
||||
[torch.bfloat16, torch.float16],
|
||||
[1, 7, 38, 64, 333],
|
||||
[128, 2048, 4096, 7168],
|
||||
[False, True],
|
||||
[False, True],
|
||||
)
|
||||
),
|
||||
[
|
||||
(torch.bfloat16, 1, 128, False, False),
|
||||
(torch.bfloat16, 17, 1536, False, False),
|
||||
(torch.bfloat16, 38, 7168, False, True),
|
||||
(torch.bfloat16, 38, 4096, False, True),
|
||||
(torch.bfloat16, 64, 4096, True, True),
|
||||
(torch.float16, 333, 2048, True, False),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"dtype,num_tokens,hidden,fuse_silu_and_mul,scale_ue8m0", V2_QUANT_CASES
|
||||
)
|
||||
def test_v2_jit_matches_aot(dtype, num_tokens, hidden, fuse_silu_and_mul, scale_ue8m0):
|
||||
"""JIT v2 must be bit-exact with the AOT v2 across vanilla/silu and float/ue8m0
|
||||
scales (NaiveScheduler)."""
|
||||
torch.manual_seed(
|
||||
hidden + num_tokens + int(fuse_silu_and_mul) + 7 * int(scale_ue8m0)
|
||||
)
|
||||
in_hidden = hidden * (2 if fuse_silu_and_mul else 1)
|
||||
x = torch.randn(num_tokens, in_hidden, device="cuda", dtype=dtype)
|
||||
out_shape = (num_tokens, hidden)
|
||||
|
||||
q_ref, s_ref = _alloc(out_shape, scale_ue8m0)
|
||||
sgl_per_token_group_quant_8bit(
|
||||
x,
|
||||
q_ref,
|
||||
s_ref,
|
||||
G,
|
||||
1e-10,
|
||||
float(fp8_min),
|
||||
float(fp8_max),
|
||||
scale_ue8m0,
|
||||
fuse_silu_and_mul,
|
||||
None,
|
||||
enable_v2=True,
|
||||
)
|
||||
|
||||
x_q, x_s = _alloc(out_shape, scale_ue8m0)
|
||||
per_token_group_quant_8bit_v2(
|
||||
x,
|
||||
x_q,
|
||||
x_s,
|
||||
G,
|
||||
1e-10,
|
||||
float(fp8_min),
|
||||
float(fp8_max),
|
||||
scale_ue8m0=scale_ue8m0,
|
||||
fuse_silu_and_mul=fuse_silu_and_mul,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
assert torch.equal(x_q.view(torch.int8), q_ref.view(torch.int8)), "fp8 codes differ"
|
||||
assert torch.equal(x_s, s_ref), "scales differ"
|
||||
|
||||
|
||||
# NOTE: "row-major + scale_ue8m0=True" names two different formats:
|
||||
# 1. packed int32 [T, ceil(G/4)] (4 exponent bytes per int32) -- supported by
|
||||
# the JIT per_token_group_quant kernel and pinned bit-exact in test_per_token_group_quant
|
||||
# (test_v3_ue8m0_row_packed_bitexact);
|
||||
# 2. fp32 [T, G] storing power-of-two VALUES (the deep_gemm.fp8_einsum
|
||||
# format) -- v2-only. No srt caller requests it (production ties
|
||||
# scale_ue8m0 and column_major_scales to the same DEEPGEMM_SCALE_UE8M0
|
||||
# flag), so the srt entry `sglang_per_token_group_quant_fp8`, which now
|
||||
# routes to the JIT kernel, rejects it loudly instead of allocating an fp32
|
||||
# buffer the kernel cannot fill. The v2 JIT kernel itself still implements it and is
|
||||
# covered by test_v2_jit_matches_aot above.
|
||||
|
||||
|
||||
# Masked (EP-MoE) path: the v2 op only has a masked scheduler for the
|
||||
# column-major + ue8m0 + fused-silu+mul + masked combination. Input is 3D
|
||||
# [num_experts, tokens_padded, hidden*2]; only tokens < masked_m[e] are processed
|
||||
# (padding left untouched → zeros in both). Compare JIT vs AOT bit-exact.
|
||||
MASKED_V2_CASES = get_ci_test_range(
|
||||
# MaskedLayoutScheduler requires hidden / group_size to be divisible by 16.
|
||||
list(itertools.product([2, 5], [2048, 4096, 8192], [128, 384])),
|
||||
[(2, 2048, 128), (5, 4096, 384), (2, 8192, 128)],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_experts,hidden,tokens_pad", MASKED_V2_CASES)
|
||||
def test_v2_jit_masked_matches_aot(num_experts, hidden, tokens_pad):
|
||||
torch.manual_seed(num_experts * 1000 + hidden + tokens_pad)
|
||||
x = torch.randn(
|
||||
num_experts, tokens_pad, hidden * 2, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
masked_m = torch.randint(
|
||||
0, tokens_pad + 1, (num_experts,), device="cuda", dtype=torch.int32
|
||||
)
|
||||
out_shape = (num_experts, tokens_pad, hidden)
|
||||
|
||||
q_ref, s_ref = _alloc(out_shape, scale_ue8m0=True)
|
||||
sgl_per_token_group_quant_8bit(
|
||||
x,
|
||||
q_ref,
|
||||
s_ref,
|
||||
G,
|
||||
1e-10,
|
||||
float(fp8_min),
|
||||
float(fp8_max),
|
||||
True,
|
||||
True,
|
||||
masked_m,
|
||||
enable_v2=True,
|
||||
)
|
||||
|
||||
x_q, x_s = _alloc(out_shape, scale_ue8m0=True)
|
||||
per_token_group_quant_8bit_v2(
|
||||
x,
|
||||
x_q,
|
||||
x_s,
|
||||
G,
|
||||
1e-10,
|
||||
float(fp8_min),
|
||||
float(fp8_max),
|
||||
scale_ue8m0=True,
|
||||
fuse_silu_and_mul=True,
|
||||
masked_m=masked_m,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
assert torch.equal(
|
||||
x_q.view(torch.int8), q_ref.view(torch.int8)
|
||||
), "masked fp8 differ"
|
||||
assert torch.equal(x_s, s_ref), "masked scales differ"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
Reference in New Issue
Block a user