[Kernel Slimming] Migrate NVFP4 kernels to JIT (#19437)

python/sglang/jit_kernel/benchmark/bench_nvfp4_blockwise_moe.py

@@ -0,0 +1,250 @@
+from __future__ import annotations
+
+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,
+)
+
+FLOAT4_E2M1_MAX = 6.0
+FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
+
+
+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."
+    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 _AOT_GROUP_MM_AVAILABLE:
+        print(
+            f"[info] legacy AOT grouped_mm baseline unavailable: {_AOT_GROUP_MM_REASON}"
+        )
+    benchmark.run(print_data=True)

python/sglang/jit_kernel/benchmark/bench_nvfp4_quant.py

@@ -0,0 +1,181 @@
+from __future__ import annotations
+
+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
+
+FLOAT4_E2M1_MAX = 6.0
+FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
+BLOCK_SIZE = 16
+
+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."
+    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."
+    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 _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)

python/sglang/jit_kernel/benchmark/bench_nvfp4_scaled_mm.py

@@ -0,0 +1,175 @@
+from __future__ import annotations
+
+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
+
+FLOAT4_E2M1_MAX = 6.0
+FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
+BLOCK_SIZE = 16
+
+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."
+    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 _AOT_SCALED_MM_AVAILABLE:
+        print(
+            f"[info] legacy AOT scaled_mm baseline unavailable: {_AOT_SCALED_MM_REASON}"
+        )
+    benchmark.run(print_data=True)