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

Refactor FP4 quantization and remove deprecated JIT kernels (#30448)

Browse Source 
Co-authored-by: root <root@sgl-b300-inference.datacrunch.io>
This commit is contained in:
Brayden Zhong
2026-07-14 09:22:07 +08:00
committed by GitHub
co-authored by root
parent 0c01971eeb
commit 9756f768a6
38 changed files with 120 additions and 5705 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/jit/benchmark/bench_nvfp4_blockwise_moe.py
-263
View File
@@ -1,263 +0,0 @@
from __future__ import annotations
import sys
from typing import Any
import torch
import triton
from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
from sglang.jit_kernel.nvfp4 import (
cutlass_fp4_group_mm,
scaled_fp4_experts_quant,
scaled_fp4_quant,
)
from sglang.srt.utils import is_sm100_supported
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(
est_time=5, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
)
FLOAT4_E2M1_MAX = 6.0
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
_NVFP4_SUPPORTED = is_sm100_supported()
def _round_up(x: int, y: int) -> int:
return ((x + y - 1) // y) * y
def _expert_offsets(m_per_expert: list[int], device: torch.device) -> torch.Tensor:
offsets = [0]
for m in m_per_expert:
offsets.append(offsets[-1] + m)
return torch.tensor(offsets, dtype=torch.int32, device=device)
def _blockscale_offsets(m_per_expert: list[int], device: torch.device) -> torch.Tensor:
offsets = [0]
for m in m_per_expert:
offsets.append(offsets[-1] + _round_up(m, 128))
return torch.tensor(offsets, dtype=torch.int32, device=device)
def _prepare_case(
total_tokens: int, n: int, k: int, num_experts: int, dtype: torch.dtype
) -> dict[str, Any]:
device = torch.device("cuda")
base = total_tokens // num_experts
rem = total_tokens % num_experts
m_per_expert = [base + (1 if i < rem else 0) for i in range(num_experts)]
expert_offsets_full = _expert_offsets(m_per_expert, device)
blockscale_offsets_full = _blockscale_offsets(m_per_expert, device)
a = torch.randn((total_tokens, k), device=device, dtype=dtype) * 0.1
b = torch.randn((num_experts, n, k), device=device, dtype=dtype) * 0.1
a_global_scale = torch.empty((num_experts,), device=device, dtype=torch.float32)
for i in range(num_experts):
start = int(expert_offsets_full[i].item())
end = int(expert_offsets_full[i + 1].item())
a_global_scale[i] = (
FLOAT8_E4M3_MAX
* FLOAT4_E2M1_MAX
/ a[start:end].abs().max().to(torch.float32)
)
b_global_scale = torch.empty((num_experts,), device=device, dtype=torch.float32)
for i in range(num_experts):
b_global_scale[i] = (
FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / b[i].abs().max().to(torch.float32)
)
a_fp4, a_blockscale = scaled_fp4_experts_quant(
a,
a_global_scale,
expert_offsets_full,
blockscale_offsets_full,
topk=1,
)
b_fp4 = torch.empty((num_experts, n, k // 2), device=device, dtype=torch.uint8)
b_blockscale = torch.empty(
(num_experts, _round_up(n, 128), _round_up(k // 16, 4)),
device=device,
dtype=torch.float8_e4m3fn,
)
for i in range(num_experts):
b_fp4_i, b_scale_i = scaled_fp4_quant(b[i], b_global_scale[i])
b_fp4[i].copy_(b_fp4_i)
b_blockscale[i].copy_(b_scale_i)
alphas = (1.0 / (a_global_scale * b_global_scale)).to(torch.float32)
params = {
"ab_strides": torch.full((num_experts,), k, dtype=torch.int64, device=device),
"c_strides": torch.full((num_experts,), n, dtype=torch.int64, device=device),
"problem_sizes": torch.tensor(
[[m, n, k] for m in m_per_expert], dtype=torch.int32, device=device
),
"expert_offsets": expert_offsets_full[:-1].contiguous(),
"blockscale_offsets": blockscale_offsets_full[:-1].contiguous(),
"a_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
"b_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
"out_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
"a_scales_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
"b_scales_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
"alpha_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
"layout_sfa": torch.empty((num_experts, 5), dtype=torch.int64, device=device),
"layout_sfb": torch.empty((num_experts, 5), dtype=torch.int64, device=device),
}
expert_ranges: list[tuple[int, int]] = []
start = 0
for m in m_per_expert:
end = start + m
expert_ranges.append((start, end))
start = end
return {
"a": a,
"b": b,
"a_fp4": a_fp4,
"b_fp4": b_fp4,
"a_blockscale": a_blockscale,
"b_blockscale": b_blockscale,
"alphas": alphas,
"params": params,
"expert_offsets_full": expert_offsets_full,
"expert_ranges": expert_ranges,
"dtype": dtype,
}
def _torch_ref_group_mm(case: dict[str, Any]) -> torch.Tensor:
a = case["a"]
b = case["b"]
dtype = case["dtype"]
expert_ranges = case["expert_ranges"]
total_tokens = a.shape[0]
n = b.shape[1]
out = torch.empty((total_tokens, n), device=a.device, dtype=dtype)
for i, (start, end) in enumerate(expert_ranges):
out[start:end] = torch.matmul(a[start:end], b[i].t())
return out
def _aot_cutlass_fp4_group_mm(case: dict[str, Any]) -> torch.Tensor:
a_fp4 = case["a_fp4"]
b_fp4 = case["b_fp4"]
a_blockscale = case["a_blockscale"]
b_blockscale = case["b_blockscale"]
alphas = case["alphas"]
params = case["params"]
out_dtype = case["dtype"]
out = torch.empty(
(a_fp4.shape[0], b_fp4.shape[1]), device=a_fp4.device, dtype=out_dtype
)
torch.ops.sgl_kernel.cutlass_fp4_group_mm.default(
out,
a_fp4,
b_fp4,
a_blockscale,
b_blockscale,
alphas,
params["ab_strides"],
params["c_strides"],
params["problem_sizes"],
params["expert_offsets"],
params["blockscale_offsets"],
)
return out
def _probe_legacy_aot_group_mm() -> tuple[bool, str]:
if not torch.cuda.is_available():
return False, "CUDA is not available."
if not _NVFP4_SUPPORTED:
return False, "NVFP4 benchmarks require sm100+ with CUDA 12.8+."
try:
import sgl_kernel # noqa: F401
except Exception as e:
return False, f"import sgl_kernel failed: {e}"
if not hasattr(torch.ops, "sgl_kernel"):
return False, "torch.ops.sgl_kernel is not registered."
op = getattr(torch.ops.sgl_kernel, "cutlass_fp4_group_mm", None)
if op is None or not hasattr(op, "default"):
return False, "torch.ops.sgl_kernel.cutlass_fp4_group_mm.default is missing."
try:
case = _prepare_case(64, 256, 128, 4, torch.bfloat16)
_aot_cutlass_fp4_group_mm(case)
torch.cuda.synchronize()
except Exception as e:
return False, f"calling AOT grouped_mm op failed: {e}"
return True, ""
_AOT_GROUP_MM_AVAILABLE, _AOT_GROUP_MM_REASON = _probe_legacy_aot_group_mm()
shape_range = get_benchmark_range(
full_range=[(128, 256, 128, 4), (256, 512, 128, 8), (512, 512, 256, 8)],
ci_range=[(128, 256, 128, 4)],
)
line_vals = ["jit"]
line_names = ["JIT NVFP4 MoE GroupMM"]
styles = [("green", "-")]
if _AOT_GROUP_MM_AVAILABLE:
line_vals.append("aot_sgl_kernel")
line_names.append("AOT NVFP4 MoE GroupMM")
styles.append(("orange", "-"))
line_vals.append("torch_ref")
line_names.append("Torch Ref")
styles.append(("blue", "-"))
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["total_tokens", "n", "k", "num_experts"],
x_vals=shape_range,
x_log=False,
line_arg="provider",
line_vals=line_vals,
line_names=line_names,
styles=styles,
ylabel="us",
plot_name="nvfp4-blockwise-moe-groupmm-performance",
args={},
)
)
def benchmark(total_tokens, n, k, num_experts, provider):
case = _prepare_case(total_tokens, n, k, num_experts, torch.bfloat16)
if provider == "jit":
fn = lambda: cutlass_fp4_group_mm(
case["a_fp4"],
case["b_fp4"],
case["a_blockscale"],
case["b_blockscale"],
case["alphas"],
case["dtype"],
case["params"],
)
elif provider == "aot_sgl_kernel":
fn = lambda: _aot_cutlass_fp4_group_mm(case)
elif provider == "torch_ref":
fn = lambda: _torch_ref_group_mm(case)
else:
raise ValueError(f"Unknown provider: {provider}")
return run_benchmark(fn)
if __name__ == "__main__":
if not _NVFP4_SUPPORTED:
print("[skip] NVFP4 blockwise MoE benchmark requires sm100+ with CUDA 12.8+.")
sys.exit(0)
if not _AOT_GROUP_MM_AVAILABLE:
print(
f"[info] legacy AOT grouped_mm baseline unavailable: {_AOT_GROUP_MM_REASON}"
)
benchmark.run(print_data=True)
test/registered/jit/benchmark/bench_nvfp4_quant.py
-197
View File
@@ -1,197 +0,0 @@
from __future__ import annotations
import sys
import torch
import triton
from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
from sglang.jit_kernel.nvfp4 import scaled_fp4_quant
from sglang.srt.utils import is_sm100_supported
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(
est_time=5, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
)
FLOAT4_E2M1_MAX = 6.0
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
BLOCK_SIZE = 16
_NVFP4_SUPPORTED = is_sm100_supported()
try:
from flashinfer import fp4_quantize as flashinfer_fp4_quantize
except Exception:
flashinfer_fp4_quantize = None
def _torch_ref_quant(input: torch.Tensor, input_global_scale: torch.Tensor):
m, n = input.shape
x = input.view(m, n // BLOCK_SIZE, BLOCK_SIZE)
vec_max = torch.max(torch.abs(x), dim=-1, keepdim=True)[0].to(torch.float32)
scale = input_global_scale * (vec_max / FLOAT4_E2M1_MAX)
scale = scale.to(torch.float8_e4m3fn).to(torch.float32)
output_scale = torch.where(scale == 0, torch.zeros_like(scale), 1.0 / scale)
scaled_x = x.to(torch.float32) * output_scale
clipped = torch.clamp(scaled_x, -6.0, 6.0).reshape(m, n)
rounded = clipped.clone()
rounded[(rounded >= 0.0) & (rounded <= 0.25)] = 0.0
rounded[(rounded > 0.25) & (rounded < 0.75)] = 0.5
rounded[(rounded >= 0.75) & (rounded <= 1.25)] = 1.0
rounded[(rounded > 1.25) & (rounded < 1.75)] = 1.5
rounded[(rounded >= 1.75) & (rounded <= 2.5)] = 2.0
rounded[(rounded > 2.5) & (rounded < 3.5)] = 3.0
rounded[(rounded >= 3.5) & (rounded <= 5.0)] = 4.0
rounded[rounded > 5.0] = 6.0
# This baseline intentionally keeps work on GPU but does not pack to uint8.
return rounded, scale
def _aot_scaled_fp4_quant(input: torch.Tensor, input_global_scale: torch.Tensor):
m, n = input.shape
output = torch.empty((m, n // 2), device=input.device, dtype=torch.uint8)
rounded_m = ((m + 128 - 1) // 128) * 128
scale_n = n // BLOCK_SIZE
rounded_n = ((scale_n + 4 - 1) // 4) * 4
output_scale = torch.empty(
(rounded_m, rounded_n // 4), device=input.device, dtype=torch.int32
)
torch.ops.sgl_kernel.scaled_fp4_quant.default(
output, input, output_scale, input_global_scale
)
return output, output_scale.view(torch.float8_e4m3fn)
def _probe_legacy_aot_quant() -> tuple[bool, str]:
if not torch.cuda.is_available():
return False, "CUDA is not available."
if not _NVFP4_SUPPORTED:
return False, "NVFP4 benchmarks require sm100+ with CUDA 12.8+."
try:
import sgl_kernel # noqa: F401
except Exception as e:
return False, f"import sgl_kernel failed: {e}"
if not hasattr(torch.ops, "sgl_kernel"):
return False, "torch.ops.sgl_kernel is not registered."
op = getattr(torch.ops.sgl_kernel, "scaled_fp4_quant", None)
if op is None or not hasattr(op, "default"):
return False, "torch.ops.sgl_kernel.scaled_fp4_quant.default is missing."
try:
x = torch.randn((16, 64), dtype=torch.bfloat16, device="cuda")
global_scale = (
FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.abs(x).max().to(torch.float32)
)
_aot_scaled_fp4_quant(x, global_scale)
torch.cuda.synchronize()
except Exception as e:
return False, f"calling AOT quant op failed: {e}"
return True, ""
_AOT_QUANT_AVAILABLE, _AOT_QUANT_REASON = _probe_legacy_aot_quant()
def _probe_flashinfer_quant() -> tuple[bool, str]:
if flashinfer_fp4_quantize is None:
return False, "import flashinfer.fp4_quantize failed."
if not torch.cuda.is_available():
return False, "CUDA is not available."
if not _NVFP4_SUPPORTED:
return False, "NVFP4 benchmarks require sm100+ with CUDA 12.8+."
try:
x = torch.randn((16, 64), dtype=torch.bfloat16, device="cuda")
global_scale = (
FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.abs(x).max().to(torch.float32)
)
flashinfer_fp4_quantize(
x,
global_scale,
BLOCK_SIZE, # sf_vec_size
False, # use_ue8m0
True, # is_sf_swizzled_layout
)
torch.cuda.synchronize()
except Exception as e:
return False, f"calling flashinfer.fp4_quantize failed: {e}"
return True, ""
_FLASHINFER_QUANT_AVAILABLE, _FLASHINFER_QUANT_REASON = _probe_flashinfer_quant()
shape_range = get_benchmark_range(
full_range=[(128, 2048), (512, 4096), (1024, 4096), (2048, 8192)],
ci_range=[(128, 2048)],
)
line_vals = []
line_names = []
styles = []
if _FLASHINFER_QUANT_AVAILABLE:
line_vals.append("flashinfer")
line_names.append("FlashInfer FP4 Quant")
styles.append(("purple", "-"))
line_vals.append("jit")
line_names.append("JIT NVFP4 Quant")
styles.append(("green", "-"))
if _AOT_QUANT_AVAILABLE:
line_vals.append("aot_sgl_kernel")
line_names.append("AOT NVFP4 Quant")
styles.append(("orange", "-"))
line_vals.append("torch_ref")
line_names.append("Torch Ref")
styles.append(("blue", "-"))
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["m", "n"],
x_vals=shape_range,
x_log=False,
line_arg="provider",
line_vals=line_vals,
line_names=line_names,
styles=styles,
ylabel="us",
plot_name="nvfp4-quant-performance",
args={},
)
)
def benchmark(m, n, provider):
x = torch.randn((m, n), dtype=torch.bfloat16, device="cuda")
tensor_amax = torch.abs(x).max().to(torch.float32)
global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax
if provider == "jit":
fn = lambda: scaled_fp4_quant(x, global_scale)
elif provider == "flashinfer":
fn = lambda: flashinfer_fp4_quantize(
x,
global_scale,
BLOCK_SIZE, # sf_vec_size
False, # use_ue8m0
True, # is_sf_swizzled_layout
)
elif provider == "aot_sgl_kernel":
fn = lambda: _aot_scaled_fp4_quant(x, global_scale)
elif provider == "torch_ref":
fn = lambda: _torch_ref_quant(x, global_scale)
else:
raise ValueError(f"Unknown provider: {provider}")
return run_benchmark(fn)
if __name__ == "__main__":
if not _NVFP4_SUPPORTED:
print("[skip] NVFP4 quant benchmark requires sm100+ with CUDA 12.8+.")
sys.exit(0)
if not _FLASHINFER_QUANT_AVAILABLE:
print(
f"[info] flashinfer quant baseline unavailable: {_FLASHINFER_QUANT_REASON}"
)
if not _AOT_QUANT_AVAILABLE:
print(f"[info] legacy AOT quant baseline unavailable: {_AOT_QUANT_REASON}")
benchmark.run(print_data=True)
test/registered/jit/benchmark/bench_nvfp4_scaled_mm.py
-189
View File
@@ -1,189 +0,0 @@
from __future__ import annotations
import sys
import torch
import triton
from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm, scaled_fp4_quant
from sglang.srt.utils import is_sm100_supported, is_sm120_supported
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(
est_time=5, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
)
FLOAT4_E2M1_MAX = 6.0
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
BLOCK_SIZE = 16
_NVFP4_SUPPORTED = is_sm100_supported() or is_sm120_supported()
K_E2M1_TO_FLOAT = [
0.0,
0.5,
1.0,
1.5,
2.0,
3.0,
4.0,
6.0,
0.0,
-0.5,
-1.0,
-1.5,
-2.0,
-3.0,
-4.0,
-6.0,
]
def _dequantize_to_fp16(
tensor_fp4: torch.Tensor, tensor_sf: torch.Tensor, global_scale: torch.Tensor
):
m, packed_k = tensor_fp4.shape
k = packed_k * 2
flat = tensor_fp4.flatten()
high = (flat & 0xF0) >> 4
low = flat & 0x0F
f_h = torch.tensor([K_E2M1_TO_FLOAT[x] for x in high], device=tensor_fp4.device)
f_l = torch.tensor([K_E2M1_TO_FLOAT[x] for x in low], device=tensor_fp4.device)
val = torch.stack((f_l, f_h), dim=-1).reshape(m, k)
rounded_m = ((m + 128 - 1) // 128) * 128
scale_n = k // BLOCK_SIZE
rounded_n = ((scale_n + 4 - 1) // 4) * 4
sf = tensor_sf.view(torch.float8_e4m3fn)
tmp = torch.reshape(sf, (1, rounded_m // 128, rounded_n // 4, 32, 4, 4))
tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5))
scale = torch.reshape(tmp, (rounded_m, rounded_n))[:m, :scale_n].to(torch.float32)
scale = scale / global_scale
return (val.view(m, scale_n, BLOCK_SIZE) * scale.unsqueeze(-1)).reshape(m, k)
def _aot_cutlass_scaled_fp4_mm(
a: torch.Tensor,
b: torch.Tensor,
block_scale_a: torch.Tensor,
block_scale_b: torch.Tensor,
alpha: torch.Tensor,
out_dtype: torch.dtype,
) -> torch.Tensor:
out = torch.empty((a.shape[0], b.shape[0]), dtype=out_dtype, device=a.device)
torch.ops.sgl_kernel.cutlass_scaled_fp4_mm.default(
out, a, b, block_scale_a, block_scale_b, alpha
)
return out
def _probe_legacy_aot_scaled_mm() -> tuple[bool, str]:
if not torch.cuda.is_available():
return False, "CUDA is not available."
if not _NVFP4_SUPPORTED:
return False, "NVFP4 benchmarks require sm100+ with CUDA 12.8+."
try:
import sgl_kernel # noqa: F401
except Exception as e:
return False, f"import sgl_kernel failed: {e}"
if not hasattr(torch.ops, "sgl_kernel"):
return False, "torch.ops.sgl_kernel is not registered."
op = getattr(torch.ops.sgl_kernel, "cutlass_scaled_fp4_mm", None)
if op is None or not hasattr(op, "default"):
return False, "torch.ops.sgl_kernel.cutlass_scaled_fp4_mm.default is missing."
try:
m, n, k = 16, 32, 64
a = torch.randn((m, k), dtype=torch.bfloat16, device="cuda")
b = torch.randn((n, k), dtype=torch.bfloat16, device="cuda")
a_global_scale = (
FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.amax(a.flatten(), dim=-1)
).to(torch.float32)
b_global_scale = (
FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.amax(b.flatten(), dim=-1)
).to(torch.float32)
alpha = 1.0 / (a_global_scale * b_global_scale)
a_fp4, a_sf = scaled_fp4_quant(a, a_global_scale)
b_fp4, b_sf = scaled_fp4_quant(b, b_global_scale)
_aot_cutlass_scaled_fp4_mm(a_fp4, b_fp4, a_sf, b_sf, alpha, torch.bfloat16)
torch.cuda.synchronize()
except Exception as e:
return False, f"calling AOT scaled_mm op failed: {e}"
return True, ""
_AOT_SCALED_MM_AVAILABLE, _AOT_SCALED_MM_REASON = _probe_legacy_aot_scaled_mm()
shape_range = get_benchmark_range(
full_range=[(128, 4096, 4096), (512, 4096, 4096), (1024, 8192, 4096)],
ci_range=[(128, 4096, 4096)],
)
line_vals = ["jit"]
line_names = ["JIT NVFP4 GEMM"]
styles = [("green", "-")]
if _AOT_SCALED_MM_AVAILABLE:
line_vals.append("aot_sgl_kernel")
line_names.append("AOT NVFP4 GEMM")
styles.append(("orange", "-"))
line_vals.append("torch_ref")
line_names.append("Torch Ref")
styles.append(("blue", "-"))
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["m", "n", "k"],
x_vals=shape_range,
x_log=False,
line_arg="provider",
line_vals=line_vals,
line_names=line_names,
styles=styles,
ylabel="us",
plot_name="nvfp4-scaled-mm-performance",
args={},
)
)
def benchmark(m, n, k, provider):
a = torch.randn((m, k), dtype=torch.bfloat16, device="cuda")
b = torch.randn((n, k), dtype=torch.bfloat16, device="cuda")
a_global_scale = (
FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.amax(a.flatten(), dim=-1)
).to(torch.float32)
b_global_scale = (
FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.amax(b.flatten(), dim=-1)
).to(torch.float32)
alpha = 1.0 / (a_global_scale * b_global_scale)
a_fp4, a_sf = scaled_fp4_quant(a, a_global_scale)
b_fp4, b_sf = scaled_fp4_quant(b, b_global_scale)
if provider == "jit":
fn = lambda: cutlass_scaled_fp4_mm(
a_fp4, b_fp4, a_sf, b_sf, alpha, torch.bfloat16
)
elif provider == "aot_sgl_kernel":
fn = lambda: _aot_cutlass_scaled_fp4_mm(
a_fp4, b_fp4, a_sf, b_sf, alpha, torch.bfloat16
)
elif provider == "torch_ref":
a_ref = _dequantize_to_fp16(a_fp4, a_sf, a_global_scale)
b_ref = _dequantize_to_fp16(b_fp4, b_sf, b_global_scale)
fn = lambda: torch.matmul(a_ref, b_ref.t())
else:
raise ValueError(f"Unknown provider: {provider}")
return run_benchmark(fn)
if __name__ == "__main__":
if not _NVFP4_SUPPORTED:
print("[skip] NVFP4 scaled_mm benchmark requires sm100/sm120 with CUDA 12.8+.")
sys.exit(0)
if not _AOT_SCALED_MM_AVAILABLE:
print(
f"[info] legacy AOT scaled_mm baseline unavailable: {_AOT_SCALED_MM_REASON}"
)
benchmark.run(print_data=True)
test/registered/jit/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py
+1 -10
View File
@@ -8,7 +8,6 @@ from pathlib import Path
from typing import Any, Callable
import flashinfer
import sgl_kernel
import torch
from sglang.jit_kernel.benchmark.utils import DEFAULT_DTYPE
@@ -38,7 +37,7 @@ WARMUP = 8
ITERS = 20
FLOAT4_E2M1_MAX = 6.0
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
METHODS = ("cutlass", "flashinfer_auto", "flashinfer_cudnn")
METHODS = ("flashinfer_auto", "flashinfer_cudnn")
def benchmark_provider(
@@ -256,14 +255,6 @@ def run_shape_suite(shape_cases: list[dict[str, Any]]) -> list[dict[str, Any]]:
}
providers: dict[str, Callable[[], torch.Tensor]] = {
"cutlass": lambda: sgl_kernel.cutlass_scaled_fp4_mm(
quantized["x_fp4"],
quantized["w_fp4"],
quantized["x_sf"],
quantized["w_sf"],
quantized["alpha"],
DTYPE,
),
"flashinfer_auto": lambda: flashinfer.mm_fp4(
quantized["x_fp4"],
quantized["w_fp4"].T,
test/registered/jit/diffusion/test_diffusion_nvfp4_scaled_mm.py
+10 -23
View File
@@ -4,7 +4,6 @@ import flashinfer
import pytest
import torch
from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm, scaled_fp4_quant
from sglang.multimodal_gen.runtime.layers.quantization import (
modelopt_quant as diffusion_modelopt_quant,
)
@@ -241,8 +240,6 @@ def _build_layer(
def _resolve_mode(mode: str):
if mode == "jit_cutlass":
return scaled_fp4_quant, cutlass_scaled_fp4_mm, None
if mode == "flashinfer2":
return flashinfer.fp4_quantize, flashinfer.mm_fp4, "cudnn"
if mode == "flashinfer_trtllm":
@@ -281,7 +278,7 @@ def test_checkpoint_processing(
not _nvfp4_supported(),
reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs",
)
@pytest.mark.parametrize("mode", ["jit_cutlass", "flashinfer2"])
@pytest.mark.parametrize("mode", ["flashinfer2"])
def test_flux2_shape_correctness(mode: str) -> None:
m, n, k = FLUX2_PROJECTION_SHAPE
quantize_op, gemm_op, gemm_backend = _resolve_mode(mode)
@@ -306,25 +303,15 @@ def test_flux2_shape_correctness(mode: str) -> None:
_dequantize_nvfp4(weight_fp4, weight_scale_swizzled, weight_global_scale).t(),
)
if gemm_backend is None:
actual = gemm_op(
x_fp4,
weight_fp4,
x_scale_swizzled,
weight_scale_swizzled,
alpha,
DTYPE,
)
else:
actual = gemm_op(
x_fp4,
weight_fp4.t(),
x_scale_swizzled,
weight_scale_swizzled.t(),
alpha,
DTYPE,
backend=gemm_backend,
)
actual = gemm_op(
x_fp4,
weight_fp4.t(),
x_scale_swizzled,
weight_scale_swizzled.t(),
alpha,
DTYPE,
backend=gemm_backend,
)
diff = _calc_diff(actual, expected.to(dtype=DTYPE))
assert diff < DEEPGEMM_FP4_MAX_DIFF, f"{mode=}, {m=}, {n=}, {k=}, {diff=:.6f}"
test/registered/jit/test_nvfp4_blockwise_moe.py
-137
View File
@@ -1,137 +0,0 @@
import sys
import pytest
import torch
from sglang.jit_kernel.nvfp4 import (
cutlass_fp4_group_mm,
scaled_fp4_experts_quant,
scaled_fp4_quant,
)
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")
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
FLOAT4_E2M1_MAX = 6.0
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
def _nvfp4_supported() -> bool:
return torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0)
def _round_up(x: int, y: int) -> int:
return ((x + y - 1) // y) * y
def _build_expert_offsets(
m_per_expert: list[int], device: torch.device
) -> torch.Tensor:
offsets = [0]
for m in m_per_expert:
offsets.append(offsets[-1] + m)
return torch.tensor(offsets, dtype=torch.int32, device=device)
def _build_blockscale_offsets(
m_per_expert: list[int], device: torch.device
) -> torch.Tensor:
offsets = [0]
for m in m_per_expert:
offsets.append(offsets[-1] + _round_up(m, 128))
return torch.tensor(offsets, dtype=torch.int32, device=device)
@pytest.mark.skipif(
not _nvfp4_supported(), reason="NVFP4 requires compute capability >= 10.0"
)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_nvfp4_blockwise_moe_grouped_mm(dtype: torch.dtype) -> None:
torch.manual_seed(0)
device = torch.device("cuda")
num_experts = 4
m_per_expert = [33, 17, 48, 29]
n = 256
k = 128
expert_offsets_full = _build_expert_offsets(m_per_expert, device)
blockscale_offsets_full = _build_blockscale_offsets(m_per_expert, device)
total_m = int(expert_offsets_full[-1].item())
a = torch.randn((total_m, k), device=device, dtype=dtype) * 0.1
b = torch.randn((num_experts, n, k), device=device, dtype=dtype) * 0.1
a_global_scale = torch.empty((num_experts,), device=device, dtype=torch.float32)
for i in range(num_experts):
start = int(expert_offsets_full[i].item())
end = int(expert_offsets_full[i + 1].item())
amax = a[start:end].abs().max().to(torch.float32)
a_global_scale[i] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / amax
b_global_scale = torch.empty((num_experts,), device=device, dtype=torch.float32)
for i in range(num_experts):
bmax = b[i].abs().max().to(torch.float32)
b_global_scale[i] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / bmax
a_fp4, a_blockscale = scaled_fp4_experts_quant(
a,
a_global_scale,
expert_offsets_full,
blockscale_offsets_full,
topk=1,
)
b_fp4 = torch.empty((num_experts, n, k // 2), device=device, dtype=torch.uint8)
b_blockscale = torch.empty(
(num_experts, _round_up(n, 128), _round_up(k // 16, 4)),
device=device,
dtype=torch.float8_e4m3fn,
)
for i in range(num_experts):
b_fp4_i, b_scale_i = scaled_fp4_quant(b[i], b_global_scale[i])
b_fp4[i].copy_(b_fp4_i)
b_blockscale[i].copy_(b_scale_i)
alphas = (1.0 / (a_global_scale * b_global_scale)).to(torch.float32)
params = {
"ab_strides": torch.full((num_experts,), k, dtype=torch.int64, device=device),
"c_strides": torch.full((num_experts,), n, dtype=torch.int64, device=device),
"problem_sizes": torch.tensor(
[[m, n, k] for m in m_per_expert], dtype=torch.int32, device=device
),
"expert_offsets": expert_offsets_full[:-1].contiguous(),
"blockscale_offsets": blockscale_offsets_full[:-1].contiguous(),
"a_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
"b_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
"out_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
"a_scales_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
"b_scales_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
"alpha_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device),
"layout_sfa": torch.empty((num_experts, 5), dtype=torch.int64, device=device),
"layout_sfb": torch.empty((num_experts, 5), dtype=torch.int64, device=device),
}
out = cutlass_fp4_group_mm(
a_fp4,
b_fp4,
a_blockscale,
b_blockscale,
alphas,
dtype,
params,
)
ref = torch.empty((total_m, n), device=device, dtype=dtype)
for i in range(num_experts):
start = int(expert_offsets_full[i].item())
end = int(expert_offsets_full[i + 1].item())
ref[start:end] = torch.matmul(a[start:end], b[i].t())
torch.testing.assert_close(out, ref, atol=1e-1, rtol=1e-1)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/jit/test_nvfp4_gemm.py
-152
View File
@@ -1,152 +0,0 @@
import sys
import pytest
import torch
from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm, scaled_fp4_quant
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")
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
def _nvfp4_supported() -> bool:
return torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0)
DTYPES = [torch.float16, torch.bfloat16]
SHAPES = [
(128, 128, 64),
(128, 128, 128),
(256, 128, 64),
(128, 256, 128),
(150, 128, 64),
]
FLOAT4_E2M1_MAX = 6.0
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
K_E2M1_TO_FLOAT = [
0.0,
0.5,
1.0,
1.5,
2.0,
3.0,
4.0,
6.0,
]
def e2m1_to_fp32(int4_value: int) -> float:
sign_bit = int4_value & 0x8
int4_abs_value = int4_value & 0x7
float_result = K_E2M1_TO_FLOAT[int4_abs_value]
return -float_result if sign_bit else float_result
def break_fp4_bytes(a: torch.Tensor) -> torch.Tensor:
assert a.dtype == torch.uint8
m, n = a.shape
a = a.flatten()
high_half_byte = (a & 0xF0) >> 4
low_half_byte = a & 0x0F
f_h = torch.tensor([e2m1_to_fp32(x) for x in high_half_byte], device=a.device)
f_l = torch.tensor([e2m1_to_fp32(x) for x in low_half_byte], device=a.device)
return torch.stack((f_l, f_h), dim=-1).reshape(m, n * 2)
def convert_swizzled_to_linear(
a_sf_swizzled: torch.Tensor, m: int, k: int, block_size: int
) -> torch.Tensor:
sf_m, sf_k = a_sf_swizzled.shape
del sf_m, sf_k
m_tiles = (m + 128 - 1) // 128
f = block_size * 4
k_tiles = (k + f - 1) // f
tmp = torch.reshape(a_sf_swizzled, (1, m_tiles, k_tiles, 32, 4, 4))
tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5))
out = tmp.reshape(m_tiles * 128, k_tiles * f // block_size)
return out[0:m, 0 : k // block_size]
def dequantize_to_dtype(
tensor_fp4: torch.Tensor,
tensor_sf: torch.Tensor,
global_scale: torch.Tensor,
block_size: int = 16,
) -> torch.Tensor:
assert tensor_fp4.dtype == torch.uint8
m, packed_k = tensor_fp4.shape
k = packed_k * 2
tensor_f32 = break_fp4_bytes(tensor_fp4)
tensor_f32 = tensor_f32.reshape(m, k // block_size, block_size)
tensor_sf = tensor_sf.view(torch.float8_e4m3fn)
tensor_sf = convert_swizzled_to_linear(tensor_sf, m, k, block_size)
tensor_sf_dtype = tensor_sf.to(torch.float32) / global_scale
return (tensor_f32 * tensor_sf_dtype.unsqueeze(-1)).reshape(m, k)
def get_ref_results(
a_fp4: torch.Tensor,
b_fp4: torch.Tensor,
a_sf: torch.Tensor,
b_sf: torch.Tensor,
a_global_scale: torch.Tensor,
b_global_scale: torch.Tensor,
block_size: int,
) -> torch.Tensor:
a_in_dtype = dequantize_to_dtype(a_fp4, a_sf, a_global_scale, block_size=block_size)
b_in_dtype = dequantize_to_dtype(b_fp4, b_sf, b_global_scale, block_size=block_size)
return torch.matmul(a_in_dtype, b_in_dtype.t())
@pytest.mark.skipif(
not _nvfp4_supported(), reason="NVFP4 requires compute capability >= 10.0"
)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("shape", SHAPES)
def test_nvfp4_gemm(dtype: torch.dtype, shape: tuple[int, int, int]) -> None:
m, n, packed_k = shape
k = packed_k * 2
block_size = 16
a_dtype = torch.randn((m, k), dtype=dtype, device="cuda")
b_dtype = torch.randn((n, k), dtype=dtype, device="cuda")
a_global_scale = (
(FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a_dtype.flatten(), dim=-1)
).to(torch.float32)
b_global_scale = (
(FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(b_dtype.flatten(), dim=-1)
).to(torch.float32)
alpha = 1.0 / (a_global_scale * b_global_scale)
a_fp4, a_scale_interleaved = scaled_fp4_quant(a_dtype, a_global_scale)
b_fp4, b_scale_interleaved = scaled_fp4_quant(b_dtype, b_global_scale)
expected_out = get_ref_results(
a_fp4,
b_fp4,
a_scale_interleaved,
b_scale_interleaved,
a_global_scale,
b_global_scale,
block_size,
)
out = cutlass_scaled_fp4_mm(
a_fp4,
b_fp4,
a_scale_interleaved,
b_scale_interleaved,
alpha,
dtype,
)
torch.testing.assert_close(out, expected_out.to(dtype=dtype), atol=1e-1, rtol=1e-1)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/jit/test_nvfp4_quant.py
-225
View File
@@ -1,225 +0,0 @@
import sys
import pytest
import torch
from sglang.jit_kernel.nvfp4 import (
scaled_fp4_grouped_quant,
scaled_fp4_quant,
silu_and_mul_scaled_fp4_grouped_quant,
)
try:
from sgl_kernel import silu_and_mul as _sgl_silu_and_mul
except Exception:
_sgl_silu_and_mul = None
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")
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
def _nvfp4_supported() -> bool:
return torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0)
def _silu_and_mul_reference(x: torch.Tensor) -> torch.Tensor:
if _sgl_silu_and_mul is not None:
return _sgl_silu_and_mul(x)
k = x.shape[-1] // 2
return torch.nn.functional.silu(x[:, :, :k]) * x[:, :, k:]
DTYPES = [torch.float16, torch.bfloat16]
SHAPES = [(128, 64), (128, 128), (256, 64), (256, 128)]
PAD_SHAPES = [
(90, 64),
(150, 64),
(128, 48),
(128, 80),
]
FLOAT4_E2M1_MAX = 6.0
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
BLOCK_SIZE = 16
E2M1_TO_FLOAT32 = [
0.0,
0.5,
1.0,
1.5,
2.0,
3.0,
4.0,
6.0,
0.0,
-0.5,
-1.0,
-1.5,
-2.0,
-3.0,
-4.0,
-6.0,
]
def cast_from_fp4(x: torch.Tensor, m: int, n: int) -> torch.Tensor:
v_2nd = (x & 0xF).to(torch.long)
v_1st = ((x >> 4) & 0xF).to(torch.long)
c = torch.stack((v_2nd, v_1st), dim=-1).flatten()
lut = torch.tensor(E2M1_TO_FLOAT32, device=x.device, dtype=torch.float32)
return lut[c].reshape(m, n)
def cast_to_fp4(x: torch.Tensor) -> torch.Tensor:
sign = torch.sign(x)
x = torch.abs(x)
x[(x >= 0.0) & (x <= 0.25)] = 0.0
x[(x > 0.25) & (x < 0.75)] = 0.5
x[(x >= 0.75) & (x <= 1.25)] = 1.0
x[(x > 1.25) & (x < 1.75)] = 1.5
x[(x >= 1.75) & (x <= 2.5)] = 2.0
x[(x > 2.5) & (x < 3.5)] = 3.0
x[(x >= 3.5) & (x <= 5.0)] = 4.0
x[x > 5.0] = 6.0
return x * sign
def get_reciprocal(x):
if isinstance(x, torch.Tensor):
return torch.where(x == 0, torch.tensor(0.0, dtype=x.dtype), 1.0 / x)
return 0.0 if x == 0 else 1.0 / x
def ref_nvfp4_quant(x: torch.Tensor, global_scale: torch.Tensor):
assert global_scale.dtype == torch.float32
assert x.ndim == 2
m, n = x.shape
x = torch.reshape(x, (m, n // BLOCK_SIZE, BLOCK_SIZE))
vec_max = torch.max(torch.abs(x), dim=-1, keepdim=True)[0].to(torch.float32)
scale = global_scale * (vec_max * get_reciprocal(FLOAT4_E2M1_MAX))
scale = scale.to(torch.float8_e4m3fn).to(torch.float32)
output_scale = get_reciprocal(scale * get_reciprocal(global_scale))
scaled_x = x.to(torch.float32) * output_scale
clipped_x = torch.clamp(scaled_x, -6.0, 6.0).reshape(m, n)
return cast_to_fp4(clipped_x), scale.squeeze(-1)
def recover_swizzled_scales(scale: torch.Tensor, m: int, n: int) -> torch.Tensor:
rounded_m = ((m + 128 - 1) // 128) * 128
scale_n = n // BLOCK_SIZE
rounded_n = ((scale_n + 4 - 1) // 4) * 4
tmp = torch.reshape(scale, (1, rounded_m // 128, rounded_n // 4, 32, 4, 4))
tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5))
result = torch.reshape(tmp, (rounded_m, rounded_n)).to(torch.float32)
return result[:m, :scale_n]
@pytest.mark.skipif(
not _nvfp4_supported(), reason="NVFP4 requires compute capability >= 10.0"
)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("shape", SHAPES)
def test_quantize_to_fp4(dtype: torch.dtype, shape: tuple[int, int]) -> None:
torch.manual_seed(42)
m, n = shape
x = torch.randn((m, n), dtype=dtype, device="cuda")
tensor_amax = torch.abs(x).max().to(torch.float32)
global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax
out_ref, scale_ref = ref_nvfp4_quant(x, global_scale)
out, out_scale = scaled_fp4_quant(x, global_scale)
scale_ans = recover_swizzled_scales(out_scale, m, n)
out_ans = cast_from_fp4(out, m, n)
torch.testing.assert_close(out_ans, out_ref)
torch.testing.assert_close(scale_ans, scale_ref)
@pytest.mark.skipif(
not _nvfp4_supported(), reason="NVFP4 requires compute capability >= 10.0"
)
@pytest.mark.parametrize("shape", PAD_SHAPES)
def test_quantize_to_fp4_padded(shape: tuple[int, int]) -> None:
torch.manual_seed(42)
m, n = shape
x = torch.randn((m, n), dtype=torch.float16, device="cuda")
tensor_amax = torch.abs(x).max().to(torch.float32)
global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax
out_ref, scale_ref = ref_nvfp4_quant(x, global_scale)
out, out_scale = scaled_fp4_quant(x, global_scale)
scale_ans = recover_swizzled_scales(out_scale, m, n)
out_ans = cast_from_fp4(out, m, n)
torch.testing.assert_close(out_ans, out_ref)
torch.testing.assert_close(scale_ans, scale_ref)
@pytest.mark.skipif(
not _nvfp4_supported(), reason="NVFP4 requires compute capability >= 10.0"
)
@pytest.mark.parametrize("shape", [(2, 128, 512), (2, 100, 128)])
def test_quantize_to_fp4_grouped(shape: tuple[int, int, int]) -> None:
torch.manual_seed(42)
l, m, k = shape
x = torch.randn((l, m, k), dtype=torch.bfloat16, device="cuda")
mask = torch.randint(1, max(2, m // 2), (l,), dtype=torch.int32, device="cuda")
tensor_amax = x.abs().amax(dim=(1, 2)).to(torch.float32)
x_sf_global = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax
output, output_scales = scaled_fp4_grouped_quant(x, x_sf_global, mask)
output = output.permute(2, 0, 1)
padded_m = ((m + 128 - 1) // 128) * 128
output_scales = output_scales.permute(5, 2, 4, 0, 1, 3).view(l, padded_m, -1)
for i in range(l):
a_fp4, a_scale_interleaved = scaled_fp4_quant(x[i], x_sf_global[i])
torch.testing.assert_close(a_fp4[: mask[i]], output[i][: mask[i]])
scale_ref = recover_swizzled_scales(a_scale_interleaved, m, k)
scale_ans = recover_swizzled_scales(output_scales[i], m, k)
torch.testing.assert_close(scale_ref[: mask[i]], scale_ans[: mask[i]])
@pytest.mark.skipif(
not _nvfp4_supported(), reason="NVFP4 requires compute capability >= 10.0"
)
@pytest.mark.parametrize("shape", [(4, 96, 256), (8, 128, 512)])
def test_silu_and_mul_quantize_to_fp4_grouped(shape: tuple[int, int, int]) -> None:
torch.manual_seed(42)
l, m, k = shape
x = torch.randn((l, m, k * 2), dtype=torch.bfloat16, device="cuda")
mask = torch.randint(1, max(2, m // 2), (l,), dtype=torch.int32, device="cuda")
ref_y = _silu_and_mul_reference(x)
tensor_amax = ref_y.abs().amax(dim=(1, 2)).to(torch.float32)
y_sf_global = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax
ref_output, ref_output_scales = scaled_fp4_grouped_quant(ref_y, y_sf_global, mask)
output, output_scales = silu_and_mul_scaled_fp4_grouped_quant(x, y_sf_global, mask)
output = output.permute(2, 0, 1)
ref_output = ref_output.permute(2, 0, 1)
padded_m = ((m + 128 - 1) // 128) * 128
output_scales = output_scales.permute(5, 2, 4, 0, 1, 3).view(l, padded_m, -1)
ref_output_scales = ref_output_scales.permute(5, 2, 4, 0, 1, 3).view(
l, padded_m, -1
)
for i in range(l):
torch.testing.assert_close(ref_output[i, : mask[i]], output[i, : mask[i]])
scale_ref = recover_swizzled_scales(ref_output_scales[i], m, k)
scale_ans = recover_swizzled_scales(output_scales[i], m, k)
torch.testing.assert_close(scale_ref[: mask[i]], scale_ans[: mask[i]])
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/jit/test_silu_and_mul_scaled_fp4_experts_quant_packed.py
-337
View File
@@ -1,337 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""Unit test for the fused JIT op ``silu_and_mul_scaled_fp4_experts_quant_packed``
(introduced in PR #18612).
On the CUTLASS NVFP4 MoE intermediate, the op fuses the previous two-step path
intermediate = silu_and_mul(c1) # SiLU(gate) * up
fp4, sf = scaled_fp4_experts_quant(intermediate) # NVFP4 expert quant
into a single kernel
fp4, sf = silu_and_mul_scaled_fp4_experts_quant_packed(c1, ...)
This test compares the fused op against that exact unfused
``silu_and_mul`` + ``scaled_fp4_experts_quant`` path **with uneven expert offsets**:
experts deliberately receive very different token counts, including tiny experts and
experts whose row count is not a multiple of the 128-row block-scale padding. That is
precisely the regime that stresses the per-expert ``expert_offsets`` /
``blockscale_offsets`` indexing the fusion has to get right.
It follows the two existing siblings:
* ``test_silu_and_mul_quantize_to_fp4_grouped`` (the grouped/masked variant) -- the
unfused path is the reference, and the fused output must match it bit-exactly
(packed FP4 nibbles + recovered block scales), and
* ``test_nvfp4_blockwise_moe`` (the expert-offset variant) -- offsets are built from
an explicit, non-uniform per-expert token list.
A high-precision ``F.silu(gate) * up`` check additionally grounds the unfused path so a
bug shared by both kernels cannot produce a false (vacuous) pass.
pytest python/sglang/jit_kernel/tests/test_silu_and_mul_scaled_fp4_experts_quant_packed.py -v
"""
import sys
import pytest
import torch
import triton
from torch.nn import functional as F
from sglang.jit_kernel.activation import silu_and_mul
from sglang.jit_kernel.nvfp4 import (
scaled_fp4_experts_quant,
silu_and_mul_scaled_fp4_experts_quant_packed,
)
from sglang.test.ci.ci_register import register_cuda_ci
# The NVFP4 expert-quant kernels are Blackwell-only (sm100a), so this runs on
# the B200 unit suite.
register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
FLOAT8_E4M3_MAX = 448.0
FLOAT4_E2M1_MAX = 6.0
BLOCK_SIZE = 16
kE2M1ToFloat = torch.tensor(
[0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32
)
def _nvfp4_supported() -> bool:
return torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0)
def _round_up(x: int, y: int) -> int:
return ((x + y - 1) // y) * y
# --------------------------------------------------------------------------- #
# Offset builders (mirror test_nvfp4_blockwise_moe.py).
# expert_offsets: cumulative *actual* per-expert rows ([E+1] int32)
# blockscale_offsets: cumulative rows padded up to 128 per expert ([E+1] int32)
# A non-uniform ``m_per_expert`` makes both offset tensors uneven.
# --------------------------------------------------------------------------- #
def _build_expert_offsets(m_per_expert, device) -> torch.Tensor:
offsets = [0]
for m in m_per_expert:
offsets.append(offsets[-1] + m)
return torch.tensor(offsets, dtype=torch.int32, device=device)
def _build_blockscale_offsets(m_per_expert, device) -> torch.Tensor:
offsets = [0]
for m in m_per_expert:
offsets.append(offsets[-1] + _round_up(m, 128))
return torch.tensor(offsets, dtype=torch.int32, device=device)
# --------------------------------------------------------------------------- #
# FP4 dequant / scale-recovery helpers (mirror test/registered/kernels/test_fp4_moe.py)
# --------------------------------------------------------------------------- #
def break_fp4_bytes(a: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
assert a.dtype == torch.uint8
m, n = a.shape
a_flat = a.flatten()
high = (a_flat & 0xF0) >> 4
low = a_flat & 0x0F
combined = torch.stack((low, high), dim=1).flatten()
signs = (combined & 0x08).to(torch.bool)
abs_vals = (combined & 0x07).to(torch.long)
kE2M1 = kE2M1ToFloat.to(device=a.device)
values = kE2M1[abs_vals] * torch.where(signs, -1.0, 1.0)
return values.reshape(m, n * 2).to(dtype=dtype)
def convert_swizzled_to_linear(
a_sf_swizzled: torch.Tensor, m: int, k: int, block_size: int
) -> torch.Tensor:
"""De-swizzle one expert's block-scale region and drop the 128-row padding tail."""
m_tiles = (m + 128 - 1) // 128
f = block_size * 4
k_tiles = (k + f - 1) // f
tmp = torch.reshape(a_sf_swizzled, (1, m_tiles, k_tiles, 32, 4, 4))
tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5))
out = tmp.reshape(m_tiles * 128, k_tiles * f // block_size)
return out[0:m, 0:k]
def dequantize_nvfp4_to_dtype(
tensor_fp4: torch.Tensor,
tensor_sf: torch.Tensor,
global_scale: torch.Tensor,
dtype: torch.dtype,
device: torch.device,
block_size: int = 16,
) -> torch.Tensor:
"""Dequantize one expert's packed FP4 (m, k//2) + swizzled block scales."""
assert tensor_fp4.dtype == torch.uint8
m, packed_k = tensor_fp4.shape
k = packed_k * 2
tensor_f32 = break_fp4_bytes(tensor_fp4, dtype)
tensor_f32 = tensor_f32.reshape(m, k // block_size, block_size)
tensor_sf = tensor_sf.view(torch.float8_e4m3fn)
tensor_sf = convert_swizzled_to_linear(tensor_sf, m, k, block_size)
tensor_sf_dtype = tensor_sf.to(torch.float32) / global_scale
out = (tensor_f32 * tensor_sf_dtype.unsqueeze(-1)).reshape(m, k)
return out.to(dtype=dtype)
def _recover_block_scales(
sf: torch.Tensor, s0: int, s1: int, m_e: int, n: int
) -> torch.Tensor:
"""De-swizzled, un-padded block scales (float32) for one expert's region."""
block = sf[s0:s1].contiguous().view(torch.float8_e4m3fn)
return convert_swizzled_to_linear(block, m_e, n, BLOCK_SIZE).to(torch.float32)
def _rel_l2(a: torch.Tensor, b: torch.Tensor) -> float:
return (a.float() - b.float()).norm().item() / b.float().norm().clamp_min(
1e-9
).item()
# --------------------------------------------------------------------------- #
# Uneven per-expert token counts. Each list is deliberately NON-uniform so that
# expert_offsets / blockscale_offsets are uneven, exercising:
# * tiny experts (1, 5, 7 tokens),
# * an exactly-128 expert (no padding),
# * experts straddling the 128-row block-scale padding (130, 200, 384).
# --------------------------------------------------------------------------- #
UNEVEN_M_PER_EXPERT = [
[33, 17, 48, 29], # all < 128 (matches test_nvfp4_blockwise_moe)
[1, 128, 200, 5, 64], # tiny + exactly-128 + cross-128
[130, 1, 384, 17, 96, 7], # heavy skew, large dynamic range
]
NS = [256, 768] # 768 == Qwen3-30B-A3B moe_intermediate_size
DTYPES = [torch.bfloat16, torch.float16]
@pytest.mark.skipif(
not _nvfp4_supported(),
reason="NVFP4 fused expert-quant kernel requires compute capability >= 10.0 (B200/SM100).",
)
@pytest.mark.parametrize("m_per_expert", UNEVEN_M_PER_EXPERT)
@pytest.mark.parametrize("n", NS)
@pytest.mark.parametrize("dtype", DTYPES)
@torch.inference_mode()
def test_fused_matches_unfused_uneven_offsets(m_per_expert, n, dtype):
torch.manual_seed(0)
device = torch.device("cuda")
num_experts = len(m_per_expert)
# --- uneven expert offsets ---
expert_offsets = _build_expert_offsets(m_per_expert, device)
blockscale_offsets = _build_blockscale_offsets(m_per_expert, device)
total_m = int(expert_offsets[-1].item())
counts = torch.tensor(m_per_expert)
assert (
counts.max() >= 2 * counts.min()
), "expert offsets must be uneven for this test"
# gate+up concatenated input (m, 2n); /5 keeps values in a sane FP4 range.
c1 = torch.randn((total_m, 2 * n), dtype=dtype, device=device) / 5.0
gate, up = c1[:, :n].float(), c1[:, n:].float()
ref = F.silu(gate) * up # high-precision SiLU(gate) * up, (total_m, n) fp32
# Per-expert global scale, exactly like cutlass_moe builds a2_gscale.
gscale = torch.empty(num_experts, dtype=torch.float32, device=device)
for e in range(num_experts):
r0, r1 = int(expert_offsets[e]), int(expert_offsets[e + 1])
amax = ref[r0:r1].abs().max().clamp_min(1e-6)
gscale[e] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / amax
# topk only gates the buffer-size assertion in the wrapper; the per-expert
# layout is driven entirely by the offsets. Both paths use the same value.
topk = 1
# ---- fused (new op) ----
fused_fp4, fused_sf = silu_and_mul_scaled_fp4_experts_quant_packed(
c1, gscale, expert_offsets, blockscale_offsets, topk
)
# ---- unfused (the exact path the op replaced) ----
intermediate = torch.empty((total_m, n), dtype=dtype, device=device)
silu_and_mul(c1, intermediate)
unf_fp4, unf_sf = scaled_fp4_experts_quant(
intermediate, gscale, expert_offsets, blockscale_offsets, topk
)
assert fused_fp4.shape == unf_fp4.shape == (total_m, n // 2)
# Per-expert, bit-exact comparison honoring the uneven offsets.
for e in range(num_experts):
r0, r1 = int(expert_offsets[e]), int(expert_offsets[e + 1])
s0, s1 = int(blockscale_offsets[e]), int(blockscale_offsets[e + 1])
m_e = r1 - r0
# (1) Packed FP4 nibbles are identical: same NVFP4 quantizer, same fused
# SiLU(gate)*up rounded to the storage dtype before quantization.
torch.testing.assert_close(
fused_fp4[r0:r1], unf_fp4[r0:r1], msg=f"FP4 bytes differ for expert {e}"
)
# (2) Recovered (de-swizzled, un-padded) block scales are identical.
torch.testing.assert_close(
_recover_block_scales(fused_sf, s0, s1, m_e, n),
_recover_block_scales(unf_sf, s0, s1, m_e, n),
msg=f"block scales differ for expert {e}",
)
# (3) Grounding: the unfused path really reproduces SiLU(gate)*up within FP4
# error, so (1)/(2) cannot pass vacuously on a bug shared by both kernels.
deq = dequantize_nvfp4_to_dtype(
unf_fp4[r0:r1].contiguous(),
unf_sf[s0:s1].contiguous(),
gscale[e],
dtype,
device,
BLOCK_SIZE,
)
assert (
_rel_l2(deq, ref[r0:r1]) < 0.2
), f"expert {e}: unfused dequant does not match SiLU(gate)*up reference"
# --------------------------------------------------------------------------- #
# Performance. The fusion removes, on the MoE down-projection input, one
# intermediate buffer allocation, one extra kernel launch, and a full HBM
# round-trip of the SiLU(gate)*up result. The speedup is measured under CUDA
# graphs -- the steady-state GPU memory-traffic saving, matching how SGLang
# executes graphed decode (the credible, low-noise number; eager wall-clock is
# dominated by launch/dispatch overhead and is too noisy to assert on). The
# assert is only a conservative regression floor; the printed speedup is the real
# result. Mirrors test_cutedsl_gdn_performance, in the same kernel unit suite.
#
# Tokens are spread evenly across experts here (the representative throughput
# case) at realistic Qwen3-30B-A3B MoE dims (n=768, 128 experts), swept from a
# decode batch up to a prefill chunk; the uneven-offset corner cases are covered
# by the correctness test above.
# --------------------------------------------------------------------------- #
PERF_SHAPES = [
(1024, 768, 128),
(4096, 768, 128),
(16384, 768, 128),
]
def _even_offsets(total_tokens, num_experts, device):
base, rem = divmod(total_tokens, num_experts)
m_per_expert = [base + (1 if i < rem else 0) for i in range(num_experts)]
return (
_build_expert_offsets(m_per_expert, device),
_build_blockscale_offsets(m_per_expert, device),
)
@pytest.mark.skipif(
not _nvfp4_supported(),
reason="NVFP4 fused expert-quant kernel requires compute capability >= 10.0 (B200/SM100).",
)
@pytest.mark.parametrize("total_tokens,n,num_experts", PERF_SHAPES)
@torch.inference_mode()
def test_fused_perf_not_regressed(total_tokens, n, num_experts):
device = torch.device("cuda")
dtype = torch.bfloat16
expert_offsets, blockscale_offsets = _even_offsets(
total_tokens, num_experts, device
)
c1 = torch.randn((total_tokens, 2 * n), dtype=dtype, device=device) / 5.0
gscale = torch.empty(num_experts, dtype=torch.float32, device=device)
for e in range(num_experts):
r0, r1 = int(expert_offsets[e]), int(expert_offsets[e + 1])
amax = c1[r0:r1].abs().max().to(torch.float32).clamp_min(1e-6)
gscale[e] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / amax
topk = 1
def fused():
silu_and_mul_scaled_fp4_experts_quant_packed(
c1, gscale, expert_offsets, blockscale_offsets, topk
)
def unfused():
# The exact path the op replaced: alloc the intermediate, SiLU*mul into
# it, then quantize it -- one extra buffer + kernel + HBM round-trip.
intermediate = torch.empty((total_tokens, n), dtype=dtype, device=device)
silu_and_mul(c1, intermediate)
scaled_fp4_experts_quant(
intermediate, gscale, expert_offsets, blockscale_offsets, topk
)
g_f = triton.testing.do_bench_cudagraph(fused) # ms, median
g_u = triton.testing.do_bench_cudagraph(unfused)
cuda_graph_speedup = g_u / g_f
print(
f"\n [PERF] tokens={total_tokens:>6} n={n} E={num_experts}: "
f"unfused {g_u * 1e3:6.1f}us fused {g_f * 1e3:6.1f}us "
f"cuda-graph speedup = {cuda_graph_speedup:.2f}x"
)
# Regression guard only: the fusion must not make this op slower. The actual
# win carries a wide margin over this floor, so shared-runner noise cannot
# flake it.
assert (
cuda_graph_speedup >= 1.05
), f"fused regressed under cuda-graph: {cuda_graph_speedup:.2f}x"
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/kernels/test_fp4_moe.py
+3 -58
View File
@@ -8,9 +8,6 @@ from flashinfer.fused_moe import cutlass_fused_moe as flashinfer_cutlass_fused_m
from sgl_kernel import silu_and_mul
from torch.nn import functional as F
from sglang.jit_kernel.nvfp4 import scaled_fp4_quant
from sglang.srt.layers.moe.cutlass_moe import cutlass_moe_fp4
from sglang.srt.layers.moe.cutlass_moe_params import CutlassMoEParams, CutlassMoEType
from sglang.srt.layers.moe.topk import TopKConfig, select_experts
from sglang.test.ci.ci_register import register_cuda_ci
@@ -282,13 +279,9 @@ def check_moe(
w1_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w1_amax
w2_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w2_amax
w1_q[expert], w1_blockscale[expert] = scaled_fp4_quant(
w1[expert], w1_gs[expert]
)
w1_q[expert], w1_blockscale[expert] = fp4_quantize(w1[expert], w1_gs[expert])
w2_q[expert], w2_blockscale[expert] = scaled_fp4_quant(
w2[expert], w2_gs[expert]
)
w2_q[expert], w2_blockscale[expert] = fp4_quantize(w2[expert], w2_gs[expert])
score = torch.randn((m, e), device="cuda", dtype=dtype)
@@ -319,7 +312,7 @@ def check_moe(
a_global_scale = (
(FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a.flatten(), dim=-1)
).to(torch.float32)
a_fp4, a_scale_interleaved = scaled_fp4_quant(a, a_global_scale)
a_fp4, a_scale_interleaved = fp4_quantize(a, a_global_scale)
_, m_k = a_fp4.shape
a_in_dtype = dequantize_nvfp4_to_dtype(
a_fp4,
@@ -365,53 +358,6 @@ def check_moe(
torch.testing.assert_close(torch_output, test_output, atol=1e-1, rtol=1e-1)
@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
@pytest.mark.parametrize("e", [40, 64, 256])
@pytest.mark.parametrize("topk", [1, 6, 8])
@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16])
@torch.inference_mode()
def test_cutlass_fp4_moe_no_graph(
m: int, n: int, k: int, e: int, topk: int, dtype: torch.dtype
):
def cutlass_moe_impl(
a,
topk_weights,
topk_ids,
w1_q,
w2_q,
a1_gs,
w1_blockscale,
w1_alphas,
a2_gs,
w2_blockscale,
w2_alphas,
):
params = CutlassMoEParams(
CutlassMoEType.BlockscaledFP4,
device=a.device,
num_experts=e,
intermediate_size_per_partition=n, # n
hidden_size=k,
) # k
return cutlass_moe_fp4(
a=a,
a1_gscale=a1_gs,
w1_fp4=w1_q,
w1_blockscale=w1_blockscale,
w1_alphas=w1_alphas,
a2_gscale=a2_gs,
w2_fp4=w2_q,
w2_blockscale=w2_blockscale,
w2_alphas=w2_alphas,
topk_weights=topk_weights,
topk_ids=topk_ids,
params=params,
apply_router_weight_on_input=False,
)
check_moe(m, n, k, e, topk, dtype, cutlass_moe_impl, flip_w13=False)
@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
@pytest.mark.parametrize("e", [40, 64, 256])
@pytest.mark.parametrize("topk", [1, 6, 8])
@@ -454,5 +400,4 @@ def test_flashinfer_fp4_moe_no_graph(
if __name__ == "__main__":
test_cutlass_fp4_moe_no_graph(224, 1024, 1024, 256, 8, torch.half)
test_flashinfer_fp4_moe_no_graph(224, 1024, 1024, 256, 8, torch.half)
test/registered/unit/test_server_args_migration.py
+2 -2
View File
@@ -45,10 +45,10 @@ class TestServerArgsAnnotatedCli(CustomTestCase):
def test_cli_name_differs_from_field_name(self):
"""cli_name maps a different CLI flag to the dataclass field via dest."""
sa = self._parse(
["--fp8-gemm-backend", "triton", "--fp4-gemm-backend", "cutlass"]
["--fp8-gemm-backend", "triton", "--fp4-gemm-backend", "marlin"]
)
self.assertEqual(sa.fp8_gemm_runner_backend, "triton")
self.assertEqual(sa.fp4_gemm_runner_backend, "cutlass")
self.assertEqual(sa.fp4_gemm_runner_backend, "marlin")
def test_nargs_question_with_const(self):
"""nargs='?' + const='' for --model-checksum."""