[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)

python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_expert_quant.cuh

@@ -1,15 +1,19 @@
-#include <ATen/cuda/CUDAContext.h>
+#include <sgl_kernel/tensor.h>
-#include <c10/cuda/CUDAGuard.h>
+#include <sgl_kernel/utils.h>
-#include <cuda_runtime.h>
+
-#include <cuda_runtime_api.h>
+#include <sgl_kernel/runtime.cuh>
-#include <torch/all.h>
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
 
 #include "nvfp4_quant.cuh"
-#include "utils.h"
+#include <cuda_runtime.h>
+#include <cuda_runtime_api.h>
+
+using namespace host;
 
 // Quantizes the provided PackedVec into the uint32_t output
 template <class Type, bool UE8M0_SF = false>
-__device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal, uint8_t* SFout) {
+SGL_DEVICE uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal, uint8_t* SFout) {
 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
   // Get absolute maximum values among the local 8 values.
   auto localMax = __habs2(vec.elts[0]);
@@ -62,11 +66,7 @@ __device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal,
 
 #pragma unroll
   for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
-    if constexpr (std::is_same_v<Type, half>) {
+    fp2Vals[i] = device::cast<float2>(vec.elts[i]);
-      fp2Vals[i] = __half22float2(vec.elts[i]);
-    } else {
-      fp2Vals[i] = __bfloat1622float2(vec.elts[i]);
-    }
     fp2Vals[i].x *= outputScale;
     fp2Vals[i].y *= outputScale;
   }
@@ -81,30 +81,22 @@ __device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal,
 #endif
 }
 
-__device__ __forceinline__ float silu(const float& val) {
+SGL_DEVICE float silu(const float& val) {
   return val / (1.0f + __expf(-val));
 }
 
 template <class Type>
-inline __device__ void silu_and_mul(PackedVec<Type>& x_vec, const PackedVec<Type>& y_vec) {
+SGL_DEVICE void silu_and_mul(PackedVec<Type>& x_vec, const PackedVec<Type>& y_vec) {
   float2 x[CVT_FP4_ELTS_PER_THREAD / 2];
   float2 y[CVT_FP4_ELTS_PER_THREAD / 2];
 
 #pragma unroll
   for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
-    if constexpr (std::is_same_v<Type, half>) {
+    x[i] = device::cast<float2>(x_vec.elts[i]);
-      x[i] = __half22float2(x_vec.elts[i]);
+    y[i] = device::cast<float2>(y_vec.elts[i]);
-      y[i] = __half22float2(y_vec.elts[i]);
+    x[i].x = silu(x[i].x) * y[i].x;
-      x[i].x = silu(x[i].x) * y[i].x;
+    x[i].y = silu(x[i].y) * y[i].y;
-      x[i].y = silu(x[i].y) * y[i].y;
+    x_vec.elts[i] = device::cast<packed_t<Type>>(x[i]);
-      x_vec.elts[i] = __float22half2_rn(x[i]);
-    } else {
-      x[i] = __bfloat1622float2(x_vec.elts[i]);
-      y[i] = __bfloat1622float2(y_vec.elts[i]);
-      x[i].x = silu(x[i].x) * y[i].x;
-      x[i].y = silu(x[i].y) * y[i].y;
-      x_vec.elts[i] = __float22bfloat162_rn(x[i]);
-    }
   }
 }
 
@@ -541,77 +533,69 @@ void quant_impl(
   }
 }
 
-// Avoid redefinition warnings
+inline int getSMVersion(int device_id) {
-#undef CHECK_CONTIGUOUS
+  int sm_major = 0;
-#undef CHECK_TH_CUDA
+  int sm_minor = 0;
-#undef CHECK_INPUT
+  RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device_id));
-
+  RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_minor, cudaDevAttrComputeCapabilityMinor, device_id));
-/*Quantization entry for fp4 experts quantization*/
+  return sm_major * 10 + sm_minor;
-#define CHECK_TH_CUDA(x, m) TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor")
+}
-#define CHECK_CONTIGUOUS(x, m) TORCH_CHECK(x.is_contiguous(), m, "must be contiguous")
-#define CHECK_INPUT(x, m) \
-  CHECK_TH_CUDA(x, m);    \
-  CHECK_CONTIGUOUS(x, m);
-
-// constexpr auto FP8 = at::ScalarType::Float8_e4m3fn;
-constexpr auto HALF = at::ScalarType::Half;
-constexpr auto BF16 = at::ScalarType::BFloat16;
-constexpr auto FLOAT = at::ScalarType::Float;
-constexpr auto INT = at::ScalarType::Int;
-constexpr auto UINT8 = at::ScalarType::Byte;
 
 void scaled_fp4_experts_quant_sm100a(
-    torch::Tensor& output,
+    tvm::ffi::TensorView output,
-    torch::Tensor& output_scale,
+    tvm::ffi::TensorView output_scale,
-    torch::Tensor const& input,
+    tvm::ffi::TensorView input,
-    torch::Tensor const& input_global_scale,
+    tvm::ffi::TensorView input_global_scale,
-    torch::Tensor const& input_offset_by_experts,
+    tvm::ffi::TensorView input_offset_by_experts,
-    torch::Tensor const& output_scale_offset_by_experts) {
+    tvm::ffi::TensorView output_scale_offset_by_experts) {
-  auto sm_version = getSMVersion();
+  auto MTopK = SymbolicSize{"m_topk"};
-  TORCH_CHECK(sm_version >= 100, "fp4_quant is only supported on sm100+");
+  auto K = SymbolicSize{"k"};
+  auto OutputCols = SymbolicSize{"output_cols"};
+  auto OutputScaleRows = SymbolicSize{"output_scale_rows"};
+  auto OutputScaleCols = SymbolicSize{"output_scale_cols"};
+  auto NExperts = SymbolicSize{"n_experts"};
+  auto OffsetSize = SymbolicSize{"offset_size"};
+  auto device = SymbolicDevice{};
 
-  CHECK_INPUT(output, "output must be a CUDA tensor");
+  TensorMatcher({MTopK, K})  //
-  CHECK_INPUT(output_scale, "output_scale must be a CUDA tensor");
+      .with_dtype<fp16_t, bf16_t>()
-  CHECK_INPUT(input, "input must be a CUDA tensor");
+      .template with_device<kDLCUDA>(device)
-  CHECK_INPUT(input_global_scale, "input_global_scale must be a CUDA tensor");
+      .verify(input);
-  CHECK_INPUT(input_offset_by_experts, "input_offset_by_experts must be a CUDA tensor");
+  TensorMatcher({MTopK, OutputCols})  //
-  CHECK_INPUT(output_scale_offset_by_experts, "output_scale_offset_by_experts must be a CUDA tensor");
+      .with_dtype<uint8_t>()
+      .with_device(device)
+      .verify(output);
+  TensorMatcher({OutputScaleRows, OutputScaleCols})  //
+      .with_dtype<int32_t>()
+      .with_device(device)
+      .verify(output_scale);
+  TensorMatcher({NExperts})  //
+      .with_dtype<float>()
+      .with_device(device)
+      .verify(input_global_scale);
+  TensorMatcher({OffsetSize})  //
+      .with_dtype<int32_t>()
+      .with_device(device)
+      .verify(input_offset_by_experts)
+      .verify(output_scale_offset_by_experts);
 
-  TORCH_CHECK(output.dim() == 2);
+  const int device_id = input.device().device_id;
-  TORCH_CHECK(output_scale.dim() == 2);
+  RuntimeCheck(getSMVersion(device_id) >= 100, "fp4_quant is only supported on sm100+");
-  TORCH_CHECK(input.dim() == 2);
-  TORCH_CHECK(input_global_scale.dim() == 1);
-  TORCH_CHECK(input_offset_by_experts.dim() == 1);
-  TORCH_CHECK(output_scale_offset_by_experts.dim() == 1);
-
-  TORCH_CHECK(input.scalar_type() == HALF || input.scalar_type() == BF16);
-  TORCH_CHECK(input_global_scale.scalar_type() == FLOAT);
-  TORCH_CHECK(input_offset_by_experts.scalar_type() == INT);
-  TORCH_CHECK(output_scale_offset_by_experts.scalar_type() == INT);
-  // output is uint8 (two nvfp4 values are packed into one uint8)
-  // output_scale is int32 (four fp8 values are packed into one int32)
-  TORCH_CHECK(output.scalar_type() == UINT8);
-  TORCH_CHECK(output_scale.scalar_type() == INT);
 
   const int BLOCK_SIZE = 16;
-  auto m_topk = input.size(0);
+  const auto m_topk = static_cast<int>(MTopK.unwrap());
-  auto k = input.size(1);
+  const auto k = static_cast<int>(K.unwrap());
-  TORCH_CHECK(k % BLOCK_SIZE == 0, "k must be a multiple of 16");
+  RuntimeCheck(k % BLOCK_SIZE == 0, "k must be a multiple of 16");
-  auto n_experts = input_global_scale.size(0);
+  const auto n_experts = static_cast<int>(NExperts.unwrap());
-  TORCH_CHECK(input_offset_by_experts.size(0) == n_experts + 1);
+  const auto offset_size = static_cast<int>(OffsetSize.unwrap());
-  TORCH_CHECK(output_scale_offset_by_experts.size(0) == n_experts + 1);
+  RuntimeCheck(offset_size == n_experts + 1, "input/output offset size mismatch");
-  TORCH_CHECK(output.size(0) == m_topk);
+  RuntimeCheck(static_cast<int>(OutputCols.unwrap()) == k / 2, "output second dim mismatch");
-  TORCH_CHECK(output.size(1) == k / 2);
+  const int scales_k = k / BLOCK_SIZE;
-  int scales_k = k / BLOCK_SIZE;
+  const int padded_k = (scales_k + 3) / 4 * 4;
-  // 4 means the swizzle requirement by nvidia nvfp4.
+  RuntimeCheck(static_cast<int>(OutputScaleCols.unwrap()) * 4 == padded_k, "output_scale second dim mismatch");
-  int padded_k = (scales_k + (4 - 1)) / 4 * 4;
-  // 4 means 4 fp8 values are packed into one int32
-  TORCH_CHECK(output_scale.size(1) * 4 == padded_k);
 
-  auto in_dtype = input.dtype();
+  const cudaStream_t stream = LaunchKernel::resolve_device(input.device());
-  at::cuda::CUDAGuard device_guard{(char)input.get_device()};
+  if (host::is_type<fp16_t>(input.dtype())) {
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream(input.get_device());
-  if (in_dtype == at::ScalarType::Half) {
     quant_impl<half>(
         output.data_ptr(),
         output_scale.data_ptr(),
@@ -625,7 +609,7 @@ void scaled_fp4_experts_quant_sm100a(
         k,
         n_experts,
         stream);
-  } else if (in_dtype == at::ScalarType::BFloat16) {
+  } else {
     quant_impl<__nv_bfloat16>(
         output.data_ptr(),
         output_scale.data_ptr(),
@@ -639,62 +623,64 @@ void scaled_fp4_experts_quant_sm100a(
         k,
         n_experts,
         stream);
-  } else {
-    TORCH_CHECK(false, "Expected input data type to be half or bfloat16");
   }
 }
 
 void silu_and_mul_scaled_fp4_experts_quant_sm100a(
-    torch::Tensor& output,
+    tvm::ffi::TensorView output,
-    torch::Tensor& output_scale,
+    tvm::ffi::TensorView output_scale,
-    torch::Tensor const& input,
+    tvm::ffi::TensorView input,
-    torch::Tensor const& input_global_scale,
+    tvm::ffi::TensorView input_global_scale,
-    torch::Tensor const& mask,
+    tvm::ffi::TensorView mask,
     bool use_silu_and_mul) {
-  auto sm_version = getSMVersion();
+  auto MTopK = SymbolicSize{"m_topk"};
-  TORCH_CHECK(sm_version >= 100, "fp4_quant is only supported on sm100+");
+  auto KBy2 = SymbolicSize{"k_by_2"};
+  auto OutputCols = SymbolicSize{"output_cols"};
+  auto OutputScaleRows = SymbolicSize{"output_scale_rows"};
+  auto OutputScaleCols = SymbolicSize{"output_scale_cols"};
+  auto NExperts = SymbolicSize{"n_experts"};
+  auto device = SymbolicDevice{};
 
-  CHECK_INPUT(output, "output must be a CUDA tensor");
+  TensorMatcher({MTopK, KBy2})  //
-  CHECK_INPUT(output_scale, "output_scale must be a CUDA tensor");
+      .with_dtype<fp16_t, bf16_t>()
-  CHECK_INPUT(input, "input must be a CUDA tensor");
+      .template with_device<kDLCUDA>(device)
-  CHECK_INPUT(input_global_scale, "input_global_scale must be a CUDA tensor");
+      .verify(input);
-  CHECK_INPUT(mask, "mask must be a CUDA tensor");
+  TensorMatcher({MTopK, OutputCols})  //
+      .with_dtype<uint8_t>()
+      .with_device(device)
+      .verify(output);
+  TensorMatcher({OutputScaleRows, OutputScaleCols})  //
+      .with_dtype<int32_t>()
+      .with_device(device)
+      .verify(output_scale);
+  TensorMatcher({NExperts})  //
+      .with_dtype<float>()
+      .with_device(device)
+      .verify(input_global_scale);
+  TensorMatcher({NExperts})  //
+      .with_dtype<int32_t>()
+      .with_device(device)
+      .verify(mask);
 
-  TORCH_CHECK(output.dim() == 2);
+  const int device_id = input.device().device_id;
-  TORCH_CHECK(output_scale.dim() == 2);
+  RuntimeCheck(getSMVersion(device_id) >= 100, "fp4_quant is only supported on sm100+");
-  TORCH_CHECK(input.dim() == 2);
-  TORCH_CHECK(input_global_scale.dim() == 1);
-
-  TORCH_CHECK(input.scalar_type() == HALF || input.scalar_type() == BF16);
-  TORCH_CHECK(input_global_scale.scalar_type() == FLOAT);
-  TORCH_CHECK(mask.scalar_type() == INT);
-  // output is uint8 (two nvfp4 values are packed into one uint8)
-  // output_scale is int32 (four fp8 values are packed into one int32)
-  TORCH_CHECK(output.scalar_type() == UINT8);
-  TORCH_CHECK(output_scale.scalar_type() == INT);
 
   const int BLOCK_SIZE = 16;
-  auto m_topk = input.size(0);
+  const auto m_topk = static_cast<int>(MTopK.unwrap());
-  auto k_by_2 = input.size(1);
+  const auto k_by_2 = static_cast<int>(KBy2.unwrap());
-  auto k = k_by_2;
+  int k = k_by_2;
   if (use_silu_and_mul) {
-    TORCH_CHECK(k_by_2 % 2 == 0, "k must be a multiple of 2");
+    RuntimeCheck(k_by_2 % 2 == 0, "k must be a multiple of 2");
     k = k_by_2 / 2;
   }
-  auto n_experts = input_global_scale.size(0);
+  const auto n_experts = static_cast<int>(NExperts.unwrap());
-  TORCH_CHECK(mask.size(0) == n_experts);
+  RuntimeCheck(static_cast<int>(OutputCols.unwrap()) == k / 2, "output second dim mismatch");
-  TORCH_CHECK(output.size(0) == m_topk);
+  const int scales_k = k / BLOCK_SIZE;
-  TORCH_CHECK(output.size(1) == k / 2);
+  const int padded_k = (scales_k + 3) / 4 * 4;
-  int scales_k = k / BLOCK_SIZE;
+  RuntimeCheck(static_cast<int>(OutputScaleCols.unwrap()) * 4 == padded_k, "output_scale second dim mismatch");
-  // 4 means the swizzle requirement by nvidia nvfp4.
-  int padded_k = (scales_k + (4 - 1)) / 4 * 4;
-  // 4 means 4 fp8 values are packed into one int32
-  TORCH_CHECK(output_scale.size(1) * 4 == padded_k);
 
-  auto in_dtype = input.dtype();
+  const cudaStream_t stream = LaunchKernel::resolve_device(input.device());
-  at::cuda::CUDAGuard device_guard{(char)input.get_device()};
+  if (host::is_type<fp16_t>(input.dtype())) {
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream(input.get_device());
-  if (in_dtype == at::ScalarType::Half) {
     quant_impl<half>(
         output.data_ptr(),
         output_scale.data_ptr(),
@@ -708,7 +694,7 @@ void silu_and_mul_scaled_fp4_experts_quant_sm100a(
         k,
         n_experts,
         stream);
-  } else if (in_dtype == at::ScalarType::BFloat16) {
+  } else {
     quant_impl<__nv_bfloat16>(
         output.data_ptr(),
         output_scale.data_ptr(),
@@ -722,7 +708,5 @@ void silu_and_mul_scaled_fp4_experts_quant_sm100a(
         k,
         n_experts,
         stream);
-  } else {
-    TORCH_CHECK(false, "Expected input data type to be half or bfloat16");
   }
 }

python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant.cuh

@@ -13,35 +13,13 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include <cuda.h>
+#include <sgl_kernel/type.cuh>
-#include <cuda_fp8.h>
+#include <sgl_kernel/utils.cuh>
+
 #include <cutlass/arch/config.h>
 
-// Get type2 from type or vice versa (applied to half and bfloat16)
+#include <cuda.h>
-template <typename T>
+#include <cuda_fp8.h>
-struct TypeConverter {
-  using Type = half2;
-};  // keep for generality
-
-template <>
-struct TypeConverter<half2> {
-  using Type = half;
-};
-
-template <>
-struct TypeConverter<half> {
-  using Type = half2;
-};
-
-template <>
-struct TypeConverter<__nv_bfloat162> {
-  using Type = __nv_bfloat16;
-};
-
-template <>
-struct TypeConverter<__nv_bfloat16> {
-  using Type = __nv_bfloat162;
-};
 
 #define ELTS_PER_THREAD 8
 
@@ -49,7 +27,7 @@ constexpr int CVT_FP4_ELTS_PER_THREAD = 8;
 constexpr int CVT_FP4_SF_VEC_SIZE = 16;
 
 // Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t).
-inline __device__ uint32_t fp32_vec_to_e2m1(float (&array)[8]) {
+SGL_DEVICE uint32_t fp32_vec_to_e2m1(float (&array)[8]) {
   // PTX instructions used here requires >= sm100f.
 #if CUTLASS_ARCH_MMA_SM100A_ENABLED || CUTLASS_ARCH_MMA_SM103A_ENABLED || CUTLASS_ARCH_MMA_SM120A_ENABLED || \
     (defined(__CUDA_ARCH_FAMILY_SPECIFIC__) && (__CUDA_ARCH_FAMILY_SPECIFIC__ >= 1000))
@@ -84,7 +62,7 @@ inline __device__ uint32_t fp32_vec_to_e2m1(float (&array)[8]) {
 }
 
 // Convert 4 float2 values into 8 e2m1 values (represented as one uint32_t).
-inline __device__ uint32_t fp32_vec_to_e2m1(float2 (&array)[4]) {
+SGL_DEVICE uint32_t fp32_vec_to_e2m1(float2 (&array)[4]) {
   // PTX instructions used here requires >= sm100f.
 #if CUTLASS_ARCH_MMA_SM100A_ENABLED || CUTLASS_ARCH_MMA_SM103A_ENABLED || CUTLASS_ARCH_MMA_SM120A_ENABLED || \
     (defined(__CUDA_ARCH_FAMILY_SPECIFIC__) && (__CUDA_ARCH_FAMILY_SPECIFIC__ >= 1000))
@@ -119,14 +97,14 @@ inline __device__ uint32_t fp32_vec_to_e2m1(float2 (&array)[4]) {
 }
 
 // Fast reciprocal.
-inline __device__ float reciprocal_approximate_ftz(float a) {
+SGL_DEVICE float reciprocal_approximate_ftz(float a) {
   float b;
   asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a));
   return b;
 }
 
 template <class SFType, int CVT_FP4_NUM_THREADS_PER_SF>
-__device__ uint8_t* cvt_quant_to_fp4_get_sf_out_offset(int rowIdx, int colIdx, int numCols, SFType* SFout) {
+SGL_DEVICE uint8_t* cvt_quant_to_fp4_get_sf_out_offset(int rowIdx, int colIdx, int numCols, SFType* SFout) {
 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
   static_assert(CVT_FP4_NUM_THREADS_PER_SF == 1 || CVT_FP4_NUM_THREADS_PER_SF == 2);
 
@@ -173,7 +151,7 @@ __device__ uint8_t* cvt_quant_to_fp4_get_sf_out_offset(int rowIdx, int colIdx, i
 // Define a 16 bytes packed data type.
 template <class Type>
 struct PackedVec {
-  typename TypeConverter<Type>::Type elts[4];
+  packed_t<Type> elts[4];
 };
 
 template <>

python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant_entry.cuh

@@ -0,0 +1,68 @@
+/* Copyright 2025 SGLang Team. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+void scaled_fp4_quant_sm100a_sm120a(
+    tvm::ffi::TensorView output,
+    tvm::ffi::TensorView input,
+    tvm::ffi::TensorView output_sf,
+    tvm::ffi::TensorView input_sf);
+
+void scaled_fp4_experts_quant_sm100a(
+    tvm::ffi::TensorView output,
+    tvm::ffi::TensorView output_scale,
+    tvm::ffi::TensorView input,
+    tvm::ffi::TensorView input_global_scale,
+    tvm::ffi::TensorView input_offset_by_experts,
+    tvm::ffi::TensorView output_scale_offset_by_experts);
+
+void silu_and_mul_scaled_fp4_experts_quant_sm100a(
+    tvm::ffi::TensorView output,
+    tvm::ffi::TensorView output_scale,
+    tvm::ffi::TensorView input,
+    tvm::ffi::TensorView input_global_scale,
+    tvm::ffi::TensorView mask,
+    bool use_silu_and_mul);
+
+void scaled_fp4_quant(
+    tvm::ffi::TensorView output,
+    tvm::ffi::TensorView input,
+    tvm::ffi::TensorView output_sf,
+    tvm::ffi::TensorView input_sf) {
+  scaled_fp4_quant_sm100a_sm120a(output, input, output_sf, input_sf);
+}
+
+void scaled_fp4_experts_quant(
+    tvm::ffi::TensorView output,
+    tvm::ffi::TensorView output_scale,
+    tvm::ffi::TensorView input,
+    tvm::ffi::TensorView input_global_scale,
+    tvm::ffi::TensorView input_offset_by_experts,
+    tvm::ffi::TensorView output_scale_offset_by_experts) {
+  scaled_fp4_experts_quant_sm100a(
+      output, output_scale, input, input_global_scale, input_offset_by_experts, output_scale_offset_by_experts);
+}
+
+void silu_and_mul_scaled_fp4_experts_quant(
+    tvm::ffi::TensorView output,
+    tvm::ffi::TensorView output_scale,
+    tvm::ffi::TensorView input,
+    tvm::ffi::TensorView input_global_scale,
+    tvm::ffi::TensorView mask,
+    bool use_silu_and_mul) {
+  silu_and_mul_scaled_fp4_experts_quant_sm100a(output, output_scale, input, input_global_scale, mask, use_silu_and_mul);
+}

python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant_kernels.cuh

@@ -13,18 +13,21 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include <ATen/cuda/CUDAContext.h>
+#include <sgl_kernel/tensor.h>
-#include <c10/cuda/CUDAGuard.h>
+#include <sgl_kernel/utils.h>
-#include <cuda_runtime.h>
+
-#include <cuda_runtime_api.h>
+#include <sgl_kernel/runtime.cuh>
-#include <torch/all.h>
+#include <sgl_kernel/utils.cuh>
 
 #include "nvfp4_quant.cuh"
-#include "utils.h"
+#include <cuda_runtime.h>
+#include <cuda_runtime_api.h>
+
+using namespace host;
 
 // Quantizes the provided PackedVec into the uint32_t output
 template <class Type, bool UE8M0_SF = false>
-__device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal, uint8_t* SFout) {
+SGL_DEVICE uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal, uint8_t* SFout) {
 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
   // Get absolute maximum values among the local 8 values.
   auto localMax = __habs2(vec.elts[0]);
@@ -182,61 +185,57 @@ template void invokeFP4Quantization(
     int multiProcessorCount,
     cudaStream_t stream);
 
-inline int getMultiProcessorCount() {
+inline int getSMVersion(int device_id) {
-  static int multi_processor_count = []() {
+  int sm_major = 0;
-    int device_id = 0;
+  int sm_minor = 0;
-    int count = 0;
+  RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device_id));
-
+  RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_minor, cudaDevAttrComputeCapabilityMinor, device_id));
-    // Get the current CUDA device ID
+  return sm_major * 10 + sm_minor;
-    CHECK_CUDA_SUCCESS(cudaGetDevice(&device_id));
-
-    // Get the number of multiprocessors for the current device
-    CHECK_CUDA_SUCCESS(cudaDeviceGetAttribute(&count, cudaDevAttrMultiProcessorCount, device_id));
-
-    return count;  // Initialize the static variable
-  }();
-
-  return multi_processor_count;  // Return the cached value on subsequent calls
 }
 
 void scaled_fp4_quant_sm100a_sm120a(
-    torch::Tensor const& output,
+    tvm::ffi::TensorView output,
-    torch::Tensor const& input,
+    tvm::ffi::TensorView input,
-    torch::Tensor const& output_sf,
+    tvm::ffi::TensorView output_sf,
-    torch::Tensor const& input_sf) {
+    tvm::ffi::TensorView input_sf) {
-  auto sm_version = getSMVersion();
+  RuntimeCheck(input.device().device_type == kDLCUDA, "input must be a CUDA tensor");
-  TORCH_CHECK(sm_version >= 100, "fp4_quant is only supported on sm100+");
+  RuntimeCheck(output.device() == input.device(), "output and input must be on same device");
+  RuntimeCheck(output_sf.device() == input.device(), "output_sf and input must be on same device");
+  RuntimeCheck(input_sf.device() == input.device(), "input_sf and input must be on same device");
+  RuntimeCheck(input.dim() == 2, "input must be a 2D tensor");
+  RuntimeCheck(output.dim() == 2, "output must be a 2D tensor");
+  RuntimeCheck(output_sf.dim() == 2, "output_sf must be a 2D tensor");
+  RuntimeCheck(input_sf.numel() == 1, "input_sf must have exactly one element");
+  RuntimeCheck(host::is_type<uint8_t>(output.dtype()), "output must be uint8");
+  RuntimeCheck(host::is_type<int32_t>(output_sf.dtype()), "output_sf must be int32");
+  RuntimeCheck(host::is_type<float>(input_sf.dtype()), "input_sf must be float32");
+  RuntimeCheck(
+      host::is_type<fp16_t>(input.dtype()) || host::is_type<bf16_t>(input.dtype()), "input dtype must be fp16 or bf16");
 
-  int32_t m = input.size(0);
+  const int device_id = input.device().device_id;
-  int32_t n = input.size(1);
+  const auto sm_version = getSMVersion(device_id);
+  RuntimeCheck(sm_version >= 100, "fp4_quant is only supported on sm100+");
 
-  TORCH_CHECK(n % 16 == 0, "The N dimension must be multiple of 16.");
+  const int32_t m = static_cast<int32_t>(input.size(0));
+  const int32_t n = static_cast<int32_t>(input.size(1));
 
-  int multiProcessorCount = getMultiProcessorCount();
+  RuntimeCheck(output.size(0) == m, "output row size mismatch");
+  RuntimeCheck(output.size(1) == n / 2, "output column size mismatch");
+  RuntimeCheck(n % 16 == 0, "The N dimension must be multiple of 16.");
+
+  const int multiProcessorCount = static_cast<int>(runtime::get_sm_count(device_id));
 
   auto input_sf_ptr = static_cast<float const*>(input_sf.data_ptr());
   auto sf_out = static_cast<int32_t*>(output_sf.data_ptr());
   auto output_ptr = static_cast<int64_t*>(output.data_ptr());
-  at::cuda::CUDAGuard device_guard{(char)input.get_device()};
+  const cudaStream_t stream = LaunchKernel::resolve_device(input.device());
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream(input.get_device());
 
-  // We don't support e8m0 scales at this moment.
+  constexpr bool useUE8M0 = false;
-  bool useUE8M0 = false;
+  if (host::is_type<fp16_t>(input.dtype())) {
-
+    auto input_ptr = reinterpret_cast<half const*>(input.data_ptr());
-  switch (input.scalar_type()) {
+    invokeFP4Quantization(m, n, input_ptr, input_sf_ptr, output_ptr, sf_out, useUE8M0, multiProcessorCount, stream);
-    case torch::kHalf: {
+  } else {
-      auto input_ptr = reinterpret_cast<half const*>(input.data_ptr());
+    auto input_ptr = reinterpret_cast<__nv_bfloat16 const*>(input.data_ptr());
-      invokeFP4Quantization(m, n, input_ptr, input_sf_ptr, output_ptr, sf_out, useUE8M0, multiProcessorCount, stream);
+    invokeFP4Quantization(m, n, input_ptr, input_sf_ptr, output_ptr, sf_out, useUE8M0, multiProcessorCount, stream);
-      break;
-    }
-    case torch::kBFloat16: {
-      auto input_ptr = reinterpret_cast<__nv_bfloat16 const*>(input.data_ptr());
-      invokeFP4Quantization(m, n, input_ptr, input_sf_ptr, output_ptr, sf_out, useUE8M0, multiProcessorCount, stream);
-      break;
-    }
-    default: {
-      std::cerr << "Observing: " << input.scalar_type() << " for the input datatype which is invalid";
-      throw std::runtime_error("Unsupported input data type for quantize_to_fp4.");
-    }
   }
 }

python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_entry.cuh

@@ -0,0 +1,34 @@
+/* Copyright 2025 SGLang Team. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <sgl_kernel/tensor.h>
+
+void cutlass_scaled_fp4_mm_sm100a_sm120a(
+    tvm::ffi::TensorView D,
+    tvm::ffi::TensorView A,
+    tvm::ffi::TensorView B,
+    tvm::ffi::TensorView A_sf,
+    tvm::ffi::TensorView B_sf,
+    tvm::ffi::TensorView alpha);
+
+void cutlass_scaled_fp4_mm(
+    tvm::ffi::TensorView D,
+    tvm::ffi::TensorView A,
+    tvm::ffi::TensorView B,
+    tvm::ffi::TensorView A_sf,
+    tvm::ffi::TensorView B_sf,
+    tvm::ffi::TensorView alpha) {
+  cutlass_scaled_fp4_mm_sm100a_sm120a(D, A, B, A_sf, B_sf, alpha);
+}

python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_kernels.cuh

@@ -13,11 +13,18 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include <ATen/cuda/CUDAContext.h>
+#include <sgl_kernel/tensor.h>
-#include <c10/cuda/CUDAGuard.h>
+#include <sgl_kernel/utils.h>
-#include <torch/all.h>
 
-#include "utils.h"
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/utils.cuh>
+
+#include <cstddef>
+#include <cstdint>
+#include <cuda_runtime.h>
+#include <unordered_map>
+
+using namespace host;
 
 // clang-format off
 #include "cutlass/cutlass.h"
@@ -31,10 +38,10 @@ limitations under the License.
 /**
  * Helper function for checking CUTLASS errors
  */
-#define CUTLASS_CHECK(status)                                                       \
+#define CUTLASS_CHECK(status)                                                        \
-  {                                                                                 \
+  {                                                                                  \
-    cutlass::Status error = status;                                                 \
+    cutlass::Status error = status;                                                  \
-    TORCH_CHECK(error == cutlass::Status::kSuccess, cutlassGetStatusString(error)); \
+    RuntimeCheck(error == cutlass::Status::kSuccess, cutlassGetStatusString(error)); \
   }
 
 using namespace cute;
@@ -51,6 +58,49 @@ inline uint32_t next_pow_2(uint32_t x) {
   return x + 1;
 }
 
+struct WorkspaceKey {
+  int device_id;
+  uintptr_t stream;
+  auto operator==(const WorkspaceKey&) const -> bool = default;
+};
+
+struct WorkspaceKeyHash {
+  auto operator()(const WorkspaceKey& key) const -> size_t {
+    size_t h1 = std::hash<int>{}(key.device_id);
+    size_t h2 = std::hash<uintptr_t>{}(key.stream);
+    return h1 ^ (h2 + 0x9e3779b97f4a7c15ULL + (h1 << 6) + (h1 >> 2));
+  }
+};
+
+struct WorkspaceState {
+  void* ptr = nullptr;
+  size_t bytes = 0;
+};
+
+inline auto get_cached_workspace(size_t required_bytes, int device_id, cudaStream_t stream) -> void* {
+  if (required_bytes == 0) {
+    return nullptr;
+  }
+
+  thread_local std::unordered_map<WorkspaceKey, WorkspaceState, WorkspaceKeyHash> cache;
+  WorkspaceKey key{device_id, reinterpret_cast<uintptr_t>(stream)};
+  auto& ws = cache[key];
+
+  if (ws.ptr != nullptr && ws.bytes >= required_bytes) {
+    return ws.ptr;
+  }
+
+  RuntimeDeviceCheck(cudaSetDevice(device_id));
+  if (ws.ptr != nullptr) {
+    RuntimeDeviceCheck(cudaFreeAsync(ws.ptr, stream));
+    ws.ptr = nullptr;
+    ws.bytes = 0;
+  }
+  RuntimeDeviceCheck(cudaMallocAsync(&ws.ptr, required_bytes, stream));
+  ws.bytes = required_bytes;
+  return ws.ptr;
+}
+
 #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || \
     defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED)
 // Config(half_t/bfloat16_t) for M <= 128
@@ -277,12 +327,12 @@ struct Fp4GemmSm120 {
 
 template <typename T>
 typename T::Gemm::Arguments args_from_options(
-    at::Tensor& D,
+    tvm::ffi::TensorView D,
-    at::Tensor const& A,
+    tvm::ffi::TensorView A,
-    at::Tensor const& B,
+    tvm::ffi::TensorView B,
-    at::Tensor const& A_sf,
+    tvm::ffi::TensorView A_sf,
-    at::Tensor const& B_sf,
+    tvm::ffi::TensorView B_sf,
-    at::Tensor const& alpha,
+    tvm::ffi::TensorView alpha,
     int64_t M,
     int64_t N,
     int64_t K) {
@@ -335,12 +385,12 @@ typename T::Gemm::Arguments args_from_options(
 
 template <typename T>
 void runGemm(
-    at::Tensor& D,
+    tvm::ffi::TensorView D,
-    at::Tensor const& A,
+    tvm::ffi::TensorView A,
-    at::Tensor const& B,
+    tvm::ffi::TensorView B,
-    at::Tensor const& A_sf,
+    tvm::ffi::TensorView A_sf,
-    at::Tensor const& B_sf,
+    tvm::ffi::TensorView B_sf,
-    at::Tensor const& alpha,
+    tvm::ffi::TensorView alpha,
     int64_t m,
     int64_t n,
     int64_t k,
@@ -349,25 +399,25 @@ void runGemm(
   auto arguments = args_from_options<T>(D, A, B, A_sf, B_sf, alpha, m, n, k);
 
   size_t workspace_size = T::Gemm::get_workspace_size(arguments);
-  auto const workspace_options = torch::TensorOptions().dtype(torch::kUInt8).device(A.device());
+  int device_id = A.device().device_id;
-  auto workspace = torch::empty(workspace_size, workspace_options);
+  void* workspace = get_cached_workspace(workspace_size, device_id, stream);
 
   CUTLASS_CHECK(gemm.can_implement(arguments));
 
-  CUTLASS_CHECK(gemm.initialize(arguments, workspace.data_ptr(), stream));
+  CUTLASS_CHECK(gemm.initialize(arguments, workspace, stream));
 
-  CUTLASS_CHECK(gemm.run(arguments, workspace.data_ptr(), stream));
+  CUTLASS_CHECK(gemm.run(arguments, workspace, stream));
 }
 
 // SM120 specific args_from_options function
 template <typename Gemm>
 typename Gemm::Arguments args_from_options_sm120(
-    at::Tensor& D,
+    tvm::ffi::TensorView D,
-    at::Tensor const& A,
+    tvm::ffi::TensorView A,
-    at::Tensor const& B,
+    tvm::ffi::TensorView B,
-    at::Tensor const& A_sf,
+    tvm::ffi::TensorView A_sf,
-    at::Tensor const& B_sf,
+    tvm::ffi::TensorView B_sf,
-    torch::Tensor const& alpha,
+    tvm::ffi::TensorView alpha,
     int M,
     int N,
     int K) {
@@ -413,12 +463,12 @@ typename Gemm::Arguments args_from_options_sm120(
 // SM120 specific runGemm function
 template <typename Gemm>
 void runGemmSm120(
-    at::Tensor& D,
+    tvm::ffi::TensorView D,
-    at::Tensor const& A,
+    tvm::ffi::TensorView A,
-    at::Tensor const& B,
+    tvm::ffi::TensorView B,
-    at::Tensor const& A_sf,
+    tvm::ffi::TensorView A_sf,
-    at::Tensor const& B_sf,
+    tvm::ffi::TensorView B_sf,
-    torch::Tensor const& alpha,
+    tvm::ffi::TensorView alpha,
     int M,
     int N,
     int K,
@@ -428,25 +478,25 @@ void runGemmSm120(
   auto arguments = args_from_options_sm120<Gemm>(D, A, B, A_sf, B_sf, alpha, M, N, K);
 
   size_t workspace_size = Gemm::get_workspace_size(arguments);
-  auto const workspace_options = torch::TensorOptions().dtype(torch::kUInt8).device(A.device());
+  int device_id = A.device().device_id;
-  auto workspace = torch::empty(workspace_size, workspace_options);
+  void* workspace = get_cached_workspace(workspace_size, device_id, stream);
 
   CUTLASS_CHECK(gemm.can_implement(arguments));
 
-  CUTLASS_CHECK(gemm.initialize(arguments, workspace.data_ptr(), stream));
+  CUTLASS_CHECK(gemm.initialize(arguments, workspace, stream));
 
-  CUTLASS_CHECK(gemm.run(arguments, workspace.data_ptr(), stream));
+  CUTLASS_CHECK(gemm.run(arguments, workspace, stream));
 }
 
 // Dispatch function to select appropriate config based on M
 template <typename OutType>
 void cutlassFp4GemmDispatch(
-    torch::Tensor& D,
+    tvm::ffi::TensorView D,
-    torch::Tensor const& A,
+    tvm::ffi::TensorView A,
-    torch::Tensor const& B,
+    tvm::ffi::TensorView B,
-    torch::Tensor const& A_sf,
+    tvm::ffi::TensorView A_sf,
-    torch::Tensor const& B_sf,
+    tvm::ffi::TensorView B_sf,
-    torch::Tensor const& alpha,
+    tvm::ffi::TensorView alpha,
     int64_t m,
     int64_t n,
     int64_t k,
@@ -466,12 +516,12 @@ void cutlassFp4GemmDispatch(
 // Dispatch function to select appropriate config based on M
 template <>
 void cutlassFp4GemmDispatch<float>(
-    torch::Tensor& D,
+    tvm::ffi::TensorView D,
-    torch::Tensor const& A,
+    tvm::ffi::TensorView A,
-    torch::Tensor const& B,
+    tvm::ffi::TensorView B,
-    torch::Tensor const& A_sf,
+    tvm::ffi::TensorView A_sf,
-    torch::Tensor const& B_sf,
+    tvm::ffi::TensorView B_sf,
-    torch::Tensor const& alpha,
+    tvm::ffi::TensorView alpha,
     int64_t m,
     int64_t n,
     int64_t k,
@@ -481,12 +531,12 @@ void cutlassFp4GemmDispatch<float>(
 
 // SM120 specific dispatch functions
 void cutlass_fp4_bf16_gemm_dispatch_sm120(
-    torch::Tensor& D,
+    tvm::ffi::TensorView D,
-    torch::Tensor const& A,
+    tvm::ffi::TensorView A,
-    torch::Tensor const& B,
+    tvm::ffi::TensorView B,
-    torch::Tensor const& A_sf,
+    tvm::ffi::TensorView A_sf,
-    torch::Tensor const& B_sf,
+    tvm::ffi::TensorView B_sf,
-    torch::Tensor const& alpha,
+    tvm::ffi::TensorView alpha,
     int m,
     int n,
     int k,
@@ -502,12 +552,12 @@ void cutlass_fp4_bf16_gemm_dispatch_sm120(
 }
 
 void cutlass_fp4_f16_gemm_dispatch_sm120(
-    torch::Tensor& D,
+    tvm::ffi::TensorView D,
-    torch::Tensor const& A,
+    tvm::ffi::TensorView A,
-    torch::Tensor const& B,
+    tvm::ffi::TensorView B,
-    torch::Tensor const& A_sf,
+    tvm::ffi::TensorView A_sf,
-    torch::Tensor const& B_sf,
+    tvm::ffi::TensorView B_sf,
-    torch::Tensor const& alpha,
+    tvm::ffi::TensorView alpha,
     int m,
     int n,
     int k,
@@ -525,17 +575,17 @@ void cutlass_fp4_f16_gemm_dispatch_sm120(
 #else
 template <typename T>
 void cutlassFp4GemmDispatch(
-    at::Tensor& D,
+    tvm::ffi::TensorView D,
-    at::Tensor const& A,
+    tvm::ffi::TensorView A,
-    at::Tensor const& B,
+    tvm::ffi::TensorView B,
-    at::Tensor const& A_sf,
+    tvm::ffi::TensorView A_sf,
-    at::Tensor const& B_sf,
+    tvm::ffi::TensorView B_sf,
-    at::Tensor const& alpha,
+    tvm::ffi::TensorView alpha,
     int64_t m,
     int64_t n,
     int64_t k,
     cudaStream_t stream) {
-  TORCH_CHECK(
+  RuntimeCheck(
       false,
       "Unsupported CUTLASS version. Set VLLM_CUTLASS_SRC_DIR to "
       "a CUTLASS 3.8 source directory to enable support.");
@@ -543,39 +593,54 @@ void cutlassFp4GemmDispatch(
 #endif  // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) ||
         // defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED)
 
-// Undefine macros from utils.h to redefine with custom signatures
+inline int getSMVersion(int device_id) {
-#undef CHECK_CONTIGUOUS
+  int sm_major = 0;
-#undef CHECK_INPUT
+  int sm_minor = 0;
-
+  RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device_id));
-#define CHECK_TYPE(x, st, m) TORCH_CHECK(x.scalar_type() == st, "Inconsistency of Tensor type:", m)
+  RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_minor, cudaDevAttrComputeCapabilityMinor, device_id));
-#define CHECK_TH_CUDA(x, m) TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor")
+  return sm_major * 10 + sm_minor;
-#define CHECK_CONTIGUOUS(x, m) TORCH_CHECK(x.is_contiguous(), m, "must be contiguous")
+}
-#define CHECK_INPUT(x, st, m) \
-  CHECK_TH_CUDA(x, m);        \
-  CHECK_CONTIGUOUS(x, m);     \
-  CHECK_TYPE(x, st, m)
-
-constexpr auto FLOAT4_E2M1X2 = at::ScalarType::Byte;
-constexpr auto SF_DTYPE = at::ScalarType::Float8_e4m3fn;
 
 void cutlass_scaled_fp4_mm_sm100a_sm120a(
-    torch::Tensor& D,
+    tvm::ffi::TensorView D,
-    torch::Tensor const& A,
+    tvm::ffi::TensorView A,
-    torch::Tensor const& B,
+    tvm::ffi::TensorView B,
-    torch::Tensor const& A_sf,
+    tvm::ffi::TensorView A_sf,
-    torch::Tensor const& B_sf,
+    tvm::ffi::TensorView B_sf,
-    torch::Tensor const& alpha) {
+    tvm::ffi::TensorView alpha) {
-  CHECK_INPUT(A, FLOAT4_E2M1X2, "a");
+  RuntimeCheck(A.device().device_type == kDLCUDA, "a must be a CUDA tensor");
-  CHECK_INPUT(B, FLOAT4_E2M1X2, "b");
+  RuntimeCheck(B.device().device_type == kDLCUDA, "b must be a CUDA tensor");
+  RuntimeCheck(A_sf.device().device_type == kDLCUDA, "scale_a must be a CUDA tensor");
+  RuntimeCheck(B_sf.device().device_type == kDLCUDA, "scale_b must be a CUDA tensor");
+  RuntimeCheck(alpha.device().device_type == kDLCUDA, "alpha must be a CUDA tensor");
+  RuntimeCheck(D.device().device_type == kDLCUDA, "out must be a CUDA tensor");
 
-  CHECK_INPUT(A_sf, SF_DTYPE, "scale_a");
+  RuntimeCheck(A.device() == B.device(), "a and b must be on same device");
-  CHECK_INPUT(B_sf, SF_DTYPE, "scale_b");
+  RuntimeCheck(A.device() == A_sf.device(), "a and scale_a must be on same device");
+  RuntimeCheck(A.device() == B_sf.device(), "a and scale_b must be on same device");
+  RuntimeCheck(A.device() == alpha.device(), "a and alpha must be on same device");
+  RuntimeCheck(A.device() == D.device(), "a and out must be on same device");
 
-  CHECK_INPUT(alpha, at::ScalarType::Float, "alpha");
+  RuntimeCheck(A.is_contiguous(), "a must be contiguous");
+  RuntimeCheck(B.is_contiguous(), "b must be contiguous");
+  RuntimeCheck(A_sf.is_contiguous(), "scale_a must be contiguous");
+  RuntimeCheck(B_sf.is_contiguous(), "scale_b must be contiguous");
+  RuntimeCheck(alpha.is_contiguous(), "alpha must be contiguous");
+  RuntimeCheck(D.is_contiguous(), "out must be contiguous");
 
-  TORCH_CHECK(A.dim() == 2, "a must be a matrix");
+  RuntimeCheck(host::is_type<uint8_t>(A.dtype()), "a must be uint8");
-  TORCH_CHECK(B.dim() == 2, "b must be a matrix");
+  RuntimeCheck(host::is_type<uint8_t>(B.dtype()), "b must be uint8");
-  TORCH_CHECK(
+  RuntimeCheck(host::is_type<fp8_e4m3_t>(A_sf.dtype()), "scale_a must be float8_e4m3fn");
+  RuntimeCheck(host::is_type<fp8_e4m3_t>(B_sf.dtype()), "scale_b must be float8_e4m3fn");
+  RuntimeCheck(host::is_type<float>(alpha.dtype()), "alpha must be float32");
+
+  RuntimeCheck(A.dim() == 2, "a must be a matrix");
+  RuntimeCheck(B.dim() == 2, "b must be a matrix");
+  RuntimeCheck(A_sf.dim() == 2, "scale_a must be a matrix");
+  RuntimeCheck(B_sf.dim() == 2, "scale_b must be a matrix");
+  RuntimeCheck(alpha.numel() == 1, "alpha must have exactly one element");
+
+  RuntimeCheck(
       A.size(1) == B.size(1),
       "a and b shapes cannot be multiplied (",
       A.size(0),
@@ -587,42 +652,24 @@ void cutlass_scaled_fp4_mm_sm100a_sm120a(
       B.size(1),
       ")");
 
-  auto const m = A.size(0);
+  const auto m = static_cast<int64_t>(A.size(0));
-  auto const n = B.size(0);
+  const auto n = static_cast<int64_t>(B.size(0));
-  auto const k = A.size(1) * 2;
+  const auto k = static_cast<int64_t>(A.size(1) * 2);
+
+  RuntimeCheck(D.dim() == 2, "out must be 2D");
+  RuntimeCheck(D.size(0) == m, "out first dim must equal m");
+  RuntimeCheck(D.size(1) == n, "out second dim must equal n");
 
   constexpr int alignment = 32;
-  TORCH_CHECK(
+  RuntimeCheck(k % alignment == 0, "Expected k to be divisible by ", alignment, ", but got k: ", k);
-      k % alignment == 0,
+  RuntimeCheck(n % alignment == 0, "Expected n to be divisible by ", alignment, ", but got n: ", n);
-      "Expected k to be divisible by ",
-      alignment,
-      ", but got a shape: (",
-      A.size(0),
-      "x",
-      A.size(1),
-      "), k: ",
-      k,
-      ".");
-  TORCH_CHECK(
-      n % alignment == 0,
-      "Expected n to be divisible by ",
-      alignment,
-      ", but got b shape: (",
-      B.size(0),
-      "x",
-      B.size(1),
-      ").");
 
-  auto round_up = [](int x, int y) { return (x + y - 1) / y * y; };
+  auto round_up = [](int64_t x, int64_t y) { return (x + y - 1) / y * y; };
-  int rounded_m = round_up(m, 128);
+  const int64_t rounded_m = round_up(m, 128);
-  int rounded_n = round_up(n, 128);
+  const int64_t rounded_n = round_up(n, 128);
-  // Since k is divisible by 32 (alignment), k / 16 is guaranteed to be an
+  const int64_t rounded_k = round_up(k / 16, 4);
-  // integer.
-  int rounded_k = round_up(k / 16, 4);
 
-  TORCH_CHECK(A_sf.dim() == 2, "scale_a must be a matrix");
+  RuntimeCheck(
-  TORCH_CHECK(B_sf.dim() == 2, "scale_b must be a matrix");
-  TORCH_CHECK(
       A_sf.size(1) == B_sf.size(1),
       "scale_a and scale_b shapes cannot be multiplied (",
       A_sf.size(0),
@@ -633,55 +680,51 @@ void cutlass_scaled_fp4_mm_sm100a_sm120a(
       "x",
       B_sf.size(1),
       ")");
-  TORCH_CHECK(
+  RuntimeCheck(
       A_sf.size(0) == rounded_m && A_sf.size(1) == rounded_k,
-      "scale_a must be padded and swizzled to a shape (",
+      "scale_a must be padded/swizzled to shape (",
       rounded_m,
       "x",
       rounded_k,
-      "), but got a shape (",
+      "), got (",
       A_sf.size(0),
       "x",
       A_sf.size(1),
       ")");
-  TORCH_CHECK(
+  RuntimeCheck(
       B_sf.size(0) == rounded_n && B_sf.size(1) == rounded_k,
-      "scale_b must be padded and swizzled to a shape (",
+      "scale_b must be padded/swizzled to shape (",
       rounded_n,
       "x",
       rounded_k,
-      "), but got a shape (",
+      "), got (",
       B_sf.size(0),
       "x",
       B_sf.size(1),
       ")");
 
-  auto out_dtype = D.dtype();
+  const cudaStream_t stream = LaunchKernel::resolve_device(A.device());
-  at::cuda::CUDAGuard device_guard{(char)A.get_device()};
+  const int sm_version = getSMVersion(A.device().device_id);
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream(A.get_device());
-
-  // Check SM version and dispatch accordingly
-  auto sm_version = getSMVersion();
 
   if (sm_version >= 120) {
-    // Use SM120 specific dispatch
+    if (host::is_type<fp16_t>(D.dtype())) {
-    if (out_dtype == at::ScalarType::Half) {
+      cutlass_fp4_f16_gemm_dispatch_sm120(
-      cutlass_fp4_f16_gemm_dispatch_sm120(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
+          D, A, B, A_sf, B_sf, alpha, static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), stream);
-    } else if (out_dtype == at::ScalarType::BFloat16) {
+    } else if (host::is_type<bf16_t>(D.dtype())) {
-      cutlass_fp4_bf16_gemm_dispatch_sm120(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
+      cutlass_fp4_bf16_gemm_dispatch_sm120(
+          D, A, B, A_sf, B_sf, alpha, static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), stream);
     } else {
-      TORCH_CHECK(false, "Unsupported output data type of nvfp4 mm sm120 (", out_dtype, ")");
+      Panic("Unsupported output data type of nvfp4 mm sm120");
     }
   } else {
-    // Use SM100 dispatch for other architectures
+    if (host::is_type<fp16_t>(D.dtype())) {
-    if (out_dtype == at::ScalarType::Half) {
       cutlassFp4GemmDispatch<cutlass::half_t>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
-    } else if (out_dtype == at::ScalarType::BFloat16) {
+    } else if (host::is_type<bf16_t>(D.dtype())) {
       cutlassFp4GemmDispatch<cutlass::bfloat16_t>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
-    } else if (out_dtype == at::ScalarType::Float) {
+    } else if (host::is_type<float>(D.dtype())) {
       cutlassFp4GemmDispatch<float>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
     } else {
-      TORCH_CHECK(false, "Unsupported output data type of nvfp4 mm");
+      Panic("Unsupported output data type of nvfp4 mm");
     }
   }
 }

python/sglang/jit_kernel/csrc/moe/nvfp4_blockwise_moe.cuh

@@ -1,10 +1,11 @@
-#include <ATen/cuda/CUDAContext.h>
+#include <sgl_kernel/tensor.h>
-#include <c10/cuda/CUDAGuard.h>
+#include <sgl_kernel/utils.h>
-#include <c10/cuda/CUDAStream.h>
-#include <cutlass/arch/arch.h>
-#include <torch/all.h>
 
-#include <cassert>
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/utils.cuh>
+
+#include <cutlass/arch/arch.h>
+#include <cutlass/cutlass.h>
 
 #include "cute/tensor.hpp"
 #include "cutlass/epilogue/collective/collective_builder.hpp"
@@ -27,10 +28,66 @@
 #include "cutlass/util/reference/host/tensor_fill.h"
 #include "cutlass/util/reference/host/tensor_norm.h"
 #include "cutlass/util/tensor_view_io.h"
-#include "utils.h"
+#include <algorithm>
+#include <cassert>
+#include <cstdint>
+#include <limits>
+#include <unordered_map>
 
+using namespace host;
 using namespace cute;
 
+struct WorkspaceKey {
+  int device_id;
+  uintptr_t stream;
+  auto operator==(const WorkspaceKey&) const -> bool = default;
+};
+
+struct WorkspaceKeyHash {
+  auto operator()(const WorkspaceKey& key) const -> size_t {
+    size_t h1 = std::hash<int>{}(key.device_id);
+    size_t h2 = std::hash<uintptr_t>{}(key.stream);
+    return h1 ^ (h2 + 0x9e3779b97f4a7c15ULL + (h1 << 6) + (h1 >> 2));
+  }
+};
+
+struct WorkspaceState {
+  void* ptr = nullptr;
+  size_t bytes = 0;
+};
+
+inline auto get_cached_workspace(size_t required_bytes, int device_id, cudaStream_t stream) -> void* {
+  if (required_bytes == 0) {
+    return nullptr;
+  }
+
+  thread_local std::unordered_map<WorkspaceKey, WorkspaceState, WorkspaceKeyHash> cache;
+  WorkspaceKey key{device_id, reinterpret_cast<uintptr_t>(stream)};
+  auto& ws = cache[key];
+
+  if (ws.ptr != nullptr && ws.bytes >= required_bytes) {
+    return ws.ptr;
+  }
+
+  RuntimeDeviceCheck(cudaSetDevice(device_id));
+  if (ws.ptr != nullptr) {
+    RuntimeDeviceCheck(cudaFreeAsync(ws.ptr, stream));
+    ws.ptr = nullptr;
+    ws.bytes = 0;
+  }
+  RuntimeDeviceCheck(cudaMallocAsync(&ws.ptr, required_bytes, stream));
+  ws.bytes = required_bytes;
+  return ws.ptr;
+}
+
+inline int getSMVersion(int device_id) {
+  int sm_major = 0;
+  int sm_minor = 0;
+  RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device_id));
+  RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_minor, cudaDevAttrComputeCapabilityMinor, device_id));
+  return sm_major * 10 + sm_minor;
+}
+
 template <
     typename ElementAB,
     typename ElementC,
@@ -104,8 +161,8 @@ __global__ void __get_group_gemm_starts(
 }
 
 #define __CALL_GET_STARTS_KERNEL_BLOCKSCALE(                                                            \
-    ELEMENT_AB_TYPE, SF_TYPE, TENSOR_C_TYPE, C_TYPE, LayoutSFA, LayoutSFB, ScaleConfig)                 \
+    ELEMENT_AB_TYPE, SF_TYPE, TYPE_CHECK, C_TYPE, LayoutSFA, LayoutSFB, ScaleConfig)                    \
-  else if (out_tensors.dtype() == TENSOR_C_TYPE) {                                                      \
+  else if (TYPE_CHECK) {                                                                                \
     __get_group_gemm_starts<ELEMENT_AB_TYPE, C_TYPE, SF_TYPE, float, LayoutSFA, LayoutSFB, ScaleConfig> \
         <<<1, num_experts, 0, stream>>>(                                                                \
             static_cast<ELEMENT_AB_TYPE**>(a_starts.data_ptr()),                                        \
@@ -131,32 +188,32 @@ __global__ void __get_group_gemm_starts(
 
 template <typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
 void run_get_group_gemm_starts(
-    const torch::Tensor& a_starts,
+    const tvm::ffi::TensorView a_starts,
-    const torch::Tensor& b_starts,
+    const tvm::ffi::TensorView b_starts,
-    const torch::Tensor& out_starts,
+    const tvm::ffi::TensorView out_starts,
-    const torch::Tensor& a_scales_starts,
+    const tvm::ffi::TensorView a_scales_starts,
-    const torch::Tensor& b_scales_starts,
+    const tvm::ffi::TensorView b_scales_starts,
-    const torch::Tensor& alpha_starts,
+    const tvm::ffi::TensorView alpha_starts,
-    const torch::Tensor& layout_sfa,
+    const tvm::ffi::TensorView layout_sfa,
-    const torch::Tensor& layout_sfb,
+    const tvm::ffi::TensorView layout_sfb,
     /*these are used for their base addresses*/
-    torch::Tensor const& a_tensors,
+    tvm::ffi::TensorView const& a_tensors,
-    torch::Tensor const& b_tensors,
+    tvm::ffi::TensorView const& b_tensors,
-    torch::Tensor const& out_tensors,
+    tvm::ffi::TensorView const& out_tensors,
-    torch::Tensor const& a_scales,
+    tvm::ffi::TensorView const& a_scales,
-    torch::Tensor const& b_scales,
+    tvm::ffi::TensorView const& b_scales,
-    torch::Tensor const& alphas,
+    tvm::ffi::TensorView const& alphas,
-    torch::Tensor const& expert_offsets,
+    tvm::ffi::TensorView const& expert_offsets,
-    torch::Tensor const& sf_offsets,
+    tvm::ffi::TensorView const& sf_offsets,
-    torch::Tensor const& problem_sizes,
+    tvm::ffi::TensorView const& problem_sizes,
     int M,
     int N,
     int K) {
-  int num_experts = (int)expert_offsets.size(0);
+  int num_experts = static_cast<int>(expert_offsets.size(0));
-  auto stream = at::cuda::getCurrentCUDAStream(a_tensors.device().index());
+  auto stream = LaunchKernel::resolve_device(a_tensors.device());
 
-  TORCH_CHECK(out_tensors.size(1) == N, "Output tensor shape doesn't match expected shape");
+  RuntimeCheck(out_tensors.size(1) == N, "Output tensor shape doesn't match expected shape");
-  TORCH_CHECK(
+  RuntimeCheck(
       K / 2 == b_tensors.size(2),
       "b_tensors(dim = 2) and a_tensors(dim = 1) trailing"
       " dimension must match");
@@ -167,30 +224,44 @@ void run_get_group_gemm_starts(
   __CALL_GET_STARTS_KERNEL_BLOCKSCALE(
       cutlass::float_e2m1_t,
       cutlass::float_ue4m3_t,
-      torch::kBFloat16,
+      host::is_type<bf16_t>(out_tensors.dtype()),
       cutlass::bfloat16_t,
       LayoutSFA,
       LayoutSFB,
       ScaleConfig)
   __CALL_GET_STARTS_KERNEL_BLOCKSCALE(
-      cutlass::float_e2m1_t, cutlass::float_ue4m3_t, torch::kFloat16, half, LayoutSFA, LayoutSFB, ScaleConfig)
+      cutlass::float_e2m1_t,
+      cutlass::float_ue4m3_t,
+      host::is_type<fp16_t>(out_tensors.dtype()),
+      cutlass::half_t,
+      LayoutSFA,
+      LayoutSFB,
+      ScaleConfig)
   else {
-    TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
+    Panic("Invalid output type (must be float16 or bfloat16)");
   }
 }
 
 void run_fp4_blockwise_scaled_group_mm_sm120(
-    torch::Tensor& output,
+    tvm::ffi::TensorView output,
-    const torch::Tensor& a,
+    const tvm::ffi::TensorView a,
-    const torch::Tensor& b,
+    const tvm::ffi::TensorView b,
-    const torch::Tensor& a_blockscale,
+    const tvm::ffi::TensorView a_blockscale,
-    const torch::Tensor& b_blockscales,
+    const tvm::ffi::TensorView b_blockscales,
-    const torch::Tensor& alphas,
+    const tvm::ffi::TensorView alphas,
-    const torch::Tensor& ab_strides,
+    const tvm::ffi::TensorView ab_strides,
-    const torch::Tensor& c_strides,
+    const tvm::ffi::TensorView c_strides,
-    const torch::Tensor& problem_sizes,
+    const tvm::ffi::TensorView problem_sizes,
-    const torch::Tensor& expert_offsets,
+    const tvm::ffi::TensorView expert_offsets,
-    const torch::Tensor& sf_offsets,
+    const tvm::ffi::TensorView sf_offsets,
+    const tvm::ffi::TensorView a_ptrs,
+    const tvm::ffi::TensorView b_ptrs,
+    const tvm::ffi::TensorView out_ptrs,
+    const tvm::ffi::TensorView a_scales_ptrs,
+    const tvm::ffi::TensorView b_scales_ptrs,
+    const tvm::ffi::TensorView alpha_ptrs,
+    const tvm::ffi::TensorView layout_sfa,
+    const tvm::ffi::TensorView layout_sfb,
     int M,
     int N,
     int K) {
@@ -275,16 +346,6 @@ void run_fp4_blockwise_scaled_group_mm_sm120(
 
   using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape;
   int num_experts = static_cast<int>(expert_offsets.size(0));
-  auto options_int = torch::TensorOptions().dtype(torch::kInt64).device(a.device());
-
-  torch::Tensor a_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor b_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor out_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor a_scales_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor b_scales_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor alpha_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor layout_sfa = torch::empty({num_experts, 5}, options_int);
-  torch::Tensor layout_sfb = torch::empty({num_experts, 5}, options_int);
 
   run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
       a_ptrs,
@@ -320,13 +381,13 @@ void run_fp4_blockwise_scaled_group_mm_sm120(
   using RasterOrderOptions = cutlass::gemm::kernel::detail::RasterOrderOptions;
   typename Gemm::GemmKernel::TileSchedulerArguments scheduler;
   scheduler.raster_order = RasterOrderOptions::AlongM;
-  hw_info.device_id = a.get_device();
+  hw_info.device_id = a.device().device_id;
   static std::unordered_map<int, int> cached_sm_counts;
   if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) {
     cached_sm_counts[hw_info.device_id] =
         cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
   }
-  hw_info.sm_count = min(cached_sm_counts[hw_info.device_id], INT_MAX);
+  hw_info.sm_count = std::min(cached_sm_counts[hw_info.device_id], std::numeric_limits<int>::max());
 
   // Mainloop Arguments
   typename GemmKernel::MainloopArguments mainloop_args{
@@ -361,34 +422,44 @@ void run_fp4_blockwise_scaled_group_mm_sm120(
       scheduler};
 
   size_t workspace_size = Gemm::get_workspace_size(args);
-  auto const workspace_options = torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
+  const cudaStream_t stream = LaunchKernel::resolve_device(a.device());
-  auto workspace = torch::empty(workspace_size, workspace_options);
+  void* workspace = get_cached_workspace(workspace_size, hw_info.device_id, stream);
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream(a.get_device());
 
   auto can_implement_status = gemm_op.can_implement(args);
-  TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess, "Failed to implement GEMM");
+  RuntimeCheck(
+      can_implement_status == cutlass::Status::kSuccess,
+      "Failed to implement GEMM: ",
+      cutlassGetStatusString(can_implement_status));
 
   // Run the GEMM
-  auto status = gemm_op.initialize(args, workspace.data_ptr());
+  auto status = gemm_op.initialize(args, workspace);
-  TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to initialize GEMM");
+  RuntimeCheck(status == cutlass::Status::kSuccess, "Failed to initialize GEMM: ", cutlassGetStatusString(status));
 
-  status = gemm_op.run(args, workspace.data_ptr(), stream);
+  status = gemm_op.run(args, workspace, stream);
-  TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
+  RuntimeCheck(status == cutlass::Status::kSuccess, "Failed to run GEMM: ", cutlassGetStatusString(status));
 }
 
 template <typename OutType>
 void run_fp4_blockwise_scaled_group_mm_sm100(
-    torch::Tensor& output,
+    tvm::ffi::TensorView output,
-    const torch::Tensor& a,
+    const tvm::ffi::TensorView a,
-    const torch::Tensor& b,
+    const tvm::ffi::TensorView b,
-    const torch::Tensor& a_blockscale,
+    const tvm::ffi::TensorView a_blockscale,
-    const torch::Tensor& b_blockscales,
+    const tvm::ffi::TensorView b_blockscales,
-    const torch::Tensor& alphas,
+    const tvm::ffi::TensorView alphas,
-    const torch::Tensor& ab_strides,
+    const tvm::ffi::TensorView ab_strides,
-    const torch::Tensor& c_strides,
+    const tvm::ffi::TensorView c_strides,
-    const torch::Tensor& problem_sizes,
+    const tvm::ffi::TensorView problem_sizes,
-    const torch::Tensor& expert_offsets,
+    const tvm::ffi::TensorView expert_offsets,
-    const torch::Tensor& sf_offsets,
+    const tvm::ffi::TensorView sf_offsets,
+    const tvm::ffi::TensorView a_ptrs,
+    const tvm::ffi::TensorView b_ptrs,
+    const tvm::ffi::TensorView out_ptrs,
+    const tvm::ffi::TensorView a_scales_ptrs,
+    const tvm::ffi::TensorView b_scales_ptrs,
+    const tvm::ffi::TensorView alpha_ptrs,
+    const tvm::ffi::TensorView layout_sfa,
+    const tvm::ffi::TensorView layout_sfb,
     int M,
     int N,
     int K) {
@@ -474,16 +545,6 @@ void run_fp4_blockwise_scaled_group_mm_sm100(
 
   using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape;
   int num_experts = static_cast<int>(expert_offsets.size(0));
-  auto options_int = torch::TensorOptions().dtype(torch::kInt64).device(a.device());
-
-  torch::Tensor a_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor b_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor out_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor a_scales_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor b_scales_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor alpha_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor layout_sfa = torch::empty({num_experts, 5}, options_int);
-  torch::Tensor layout_sfb = torch::empty({num_experts, 5}, options_int);
 
   run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
       a_ptrs,
@@ -519,13 +580,13 @@ void run_fp4_blockwise_scaled_group_mm_sm100(
       typename ProblemShape::UnderlyingProblemShape>::RasterOrderOptions;
   typename Gemm::GemmKernel::TileSchedulerArguments scheduler;
   scheduler.raster_order = RasterOrderOptions::AlongM;
-  hw_info.device_id = a.get_device();
+  hw_info.device_id = a.device().device_id;
   static std::unordered_map<int, int> cached_sm_counts;
   if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) {
     cached_sm_counts[hw_info.device_id] =
         cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
   }
-  hw_info.sm_count = min(cached_sm_counts[hw_info.device_id], INT_MAX);
+  hw_info.sm_count = std::min(cached_sm_counts[hw_info.device_id], std::numeric_limits<int>::max());
 
   // Mainloop Arguments
   typename GemmKernel::MainloopArguments mainloop_args{
@@ -559,80 +620,144 @@ void run_fp4_blockwise_scaled_group_mm_sm100(
       scheduler};
 
   size_t workspace_size = Gemm::get_workspace_size(args);
-  auto const workspace_options = torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
+  const cudaStream_t stream = LaunchKernel::resolve_device(a.device());
-  auto workspace = torch::empty(workspace_size, workspace_options);
+  void* workspace = get_cached_workspace(workspace_size, hw_info.device_id, stream);
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream(a.get_device());
 
   auto can_implement_status = gemm_op.can_implement(args);
-  TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess, "Failed to implement GEMM");
+  RuntimeCheck(
+      can_implement_status == cutlass::Status::kSuccess,
+      "Failed to implement GEMM: ",
+      cutlassGetStatusString(can_implement_status));
 
   // Run the GEMM
-  auto status = gemm_op.initialize(args, workspace.data_ptr());
+  auto status = gemm_op.initialize(args, workspace);
-  TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to initialize GEMM");
+  RuntimeCheck(status == cutlass::Status::kSuccess, "Failed to initialize GEMM: ", cutlassGetStatusString(status));
 
-  status = gemm_op.run(args, workspace.data_ptr(), stream);
+  status = gemm_op.run(args, workspace, stream);
-  TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
+  RuntimeCheck(status == cutlass::Status::kSuccess, "Failed to run GEMM: ", cutlassGetStatusString(status));
 }
 
-// Undefine macros from utils.h to redefine with custom signatures
+void cutlass_fp4_group_mm_sm100a_sm120a(
-#undef CHECK_CONTIGUOUS
+    tvm::ffi::TensorView output,
-#undef CHECK_INPUT
+    const tvm::ffi::TensorView a,
+    const tvm::ffi::TensorView b,
+    const tvm::ffi::TensorView a_blockscale,
+    const tvm::ffi::TensorView b_blockscales,
+    const tvm::ffi::TensorView alphas,
+    const tvm::ffi::TensorView ab_strides,
+    const tvm::ffi::TensorView c_strides,
+    const tvm::ffi::TensorView problem_sizes,
+    const tvm::ffi::TensorView expert_offsets,
+    const tvm::ffi::TensorView sf_offsets,
+    const tvm::ffi::TensorView a_ptrs,
+    const tvm::ffi::TensorView b_ptrs,
+    const tvm::ffi::TensorView out_ptrs,
+    const tvm::ffi::TensorView a_scales_ptrs,
+    const tvm::ffi::TensorView b_scales_ptrs,
+    const tvm::ffi::TensorView alpha_ptrs,
+    const tvm::ffi::TensorView layout_sfa,
+    const tvm::ffi::TensorView layout_sfb) {
+  auto check_cuda_contig = [](const tvm::ffi::TensorView t, const char* name) {
+    RuntimeCheck(t.device().device_type == kDLCUDA, name, " must be a CUDA tensor");
+    RuntimeCheck(t.is_contiguous(), name, " must be contiguous");
+  };
 
-#define CHECK_TYPE(x, st, m) TORCH_CHECK(x.scalar_type() == st, ": Inconsistency of Tensor type:", m)
+  check_cuda_contig(output, "output");
-#define CHECK_TH_CUDA(x, m) TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor.")
+  check_cuda_contig(a, "a");
-#define CHECK_CONTIGUOUS(x, m) TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous.")
+  check_cuda_contig(b, "b");
-#define CHECK_INPUT(x, st, m) \
+  check_cuda_contig(a_blockscale, "a_blockscale");
-  CHECK_TH_CUDA(x, m);        \
+  check_cuda_contig(b_blockscales, "b_blockscales");
-  CHECK_CONTIGUOUS(x, m);     \
+  check_cuda_contig(alphas, "alphas");
-  CHECK_TYPE(x, st, m)
+  check_cuda_contig(ab_strides, "ab_strides");
+  check_cuda_contig(c_strides, "c_strides");
+  check_cuda_contig(problem_sizes, "problem_sizes");
+  check_cuda_contig(expert_offsets, "expert_offsets");
+  check_cuda_contig(sf_offsets, "sf_offsets");
+  check_cuda_contig(a_ptrs, "a_ptrs");
+  check_cuda_contig(b_ptrs, "b_ptrs");
+  check_cuda_contig(out_ptrs, "out_ptrs");
+  check_cuda_contig(a_scales_ptrs, "a_scales_ptrs");
+  check_cuda_contig(b_scales_ptrs, "b_scales_ptrs");
+  check_cuda_contig(alpha_ptrs, "alpha_ptrs");
+  check_cuda_contig(layout_sfa, "layout_sfa");
+  check_cuda_contig(layout_sfb, "layout_sfb");
 
-void cutlass_fp4_group_mm(
+  RuntimeCheck(
-    torch::Tensor& output,
+      output.device() == a.device() && a.device() == b.device() && a.device() == a_blockscale.device() &&
-    const torch::Tensor& a,
+          a.device() == b_blockscales.device() && a.device() == alphas.device() && a.device() == ab_strides.device() &&
-    const torch::Tensor& b,
+          a.device() == c_strides.device() && a.device() == problem_sizes.device() &&
-    const torch::Tensor& a_blockscale,
+          a.device() == expert_offsets.device() && a.device() == sf_offsets.device() && a.device() == a_ptrs.device() &&
-    const torch::Tensor& b_blockscales,
+          a.device() == b_ptrs.device() && a.device() == out_ptrs.device() && a.device() == a_scales_ptrs.device() &&
-    const torch::Tensor& alphas,
+          a.device() == b_scales_ptrs.device() && a.device() == alpha_ptrs.device() &&
-    const torch::Tensor& ab_strides,
+          a.device() == layout_sfa.device() && a.device() == layout_sfb.device(),
-    const torch::Tensor& c_strides,
+      "all tensors must be on the same device");
-    const torch::Tensor& problem_sizes,
-    const torch::Tensor& expert_offsets,
-    const torch::Tensor& sf_offsets) {
-#if defined ENABLE_NVFP4 && ENABLE_NVFP4
 
-  constexpr auto FLOAT4_E2M1X2 = at::ScalarType::Byte;
+  RuntimeCheck(host::is_type<uint8_t>(a.dtype()), "a must be uint8");
-  constexpr auto SF_DTYPE = at::ScalarType::Float8_e4m3fn;
+  RuntimeCheck(host::is_type<uint8_t>(b.dtype()), "b must be uint8");
-  // Input validation
+  RuntimeCheck(host::is_type<fp8_e4m3_t>(a_blockscale.dtype()), "a_blockscale must be float8_e4m3fn");
-  CHECK_INPUT(a, FLOAT4_E2M1X2, "a");
+  RuntimeCheck(host::is_type<fp8_e4m3_t>(b_blockscales.dtype()), "b_blockscales must be float8_e4m3fn");
-  CHECK_INPUT(b, FLOAT4_E2M1X2, "b");
+  RuntimeCheck(host::is_type<float>(alphas.dtype()), "alphas must be float32");
-  CHECK_INPUT(a_blockscale, SF_DTYPE, "a_blockscale");
+  RuntimeCheck(host::is_type<int64_t>(ab_strides.dtype()), "ab_strides must be int64");
-  CHECK_INPUT(b_blockscales, SF_DTYPE, "b_blockscales");
+  RuntimeCheck(host::is_type<int64_t>(c_strides.dtype()), "c_strides must be int64");
-  CHECK_INPUT(alphas, at::ScalarType::Float, "alphas");
+  RuntimeCheck(host::is_type<int32_t>(problem_sizes.dtype()), "problem_sizes must be int32");
+  RuntimeCheck(host::is_type<int32_t>(expert_offsets.dtype()), "expert_offsets must be int32");
+  RuntimeCheck(host::is_type<int32_t>(sf_offsets.dtype()), "sf_offsets must be int32");
+  RuntimeCheck(host::is_type<int64_t>(a_ptrs.dtype()), "a_ptrs must be int64");
+  RuntimeCheck(host::is_type<int64_t>(b_ptrs.dtype()), "b_ptrs must be int64");
+  RuntimeCheck(host::is_type<int64_t>(out_ptrs.dtype()), "out_ptrs must be int64");
+  RuntimeCheck(host::is_type<int64_t>(a_scales_ptrs.dtype()), "a_scales_ptrs must be int64");
+  RuntimeCheck(host::is_type<int64_t>(b_scales_ptrs.dtype()), "b_scales_ptrs must be int64");
+  RuntimeCheck(host::is_type<int64_t>(alpha_ptrs.dtype()), "alpha_ptrs must be int64");
+  RuntimeCheck(host::is_type<int64_t>(layout_sfa.dtype()), "layout_sfa must be int64");
+  RuntimeCheck(host::is_type<int64_t>(layout_sfb.dtype()), "layout_sfb must be int64");
+  RuntimeCheck(
+      host::is_type<bf16_t>(output.dtype()) || host::is_type<fp16_t>(output.dtype()),
+      "output must be bfloat16 or float16");
 
-  TORCH_CHECK(
+  RuntimeCheck(a.dim() == 2, "a must be 2D");
-      a_blockscale.dim() == 2,
+  RuntimeCheck(b.dim() == 3, "b must be 3D");
-      "expected a_blockscale to be of shape [num_experts, rounded_m,"
+  RuntimeCheck(a_blockscale.dim() == 2, "a_blockscale must be 2D");
-      " k // group_size], observed rank: ",
+  RuntimeCheck(b_blockscales.dim() == 3, "b_blockscales must be 3D");
-      a_blockscale.dim())
+  RuntimeCheck(alphas.dim() == 1, "alphas must be 1D");
-  TORCH_CHECK(
+  RuntimeCheck(ab_strides.dim() == 1, "ab_strides must be 1D");
-      b_blockscales.dim() == 3,
+  RuntimeCheck(c_strides.dim() == 1, "c_strides must be 1D");
-      "expected b_blockscale to be of shape: "
+  RuntimeCheck(problem_sizes.dim() == 2, "problem_sizes must be 2D");
-      " [num_experts, n, k // group_size], observed rank: ",
+  RuntimeCheck(expert_offsets.dim() == 1, "expert_offsets must be 1D");
-      b_blockscales.dim())
+  RuntimeCheck(sf_offsets.dim() == 1, "sf_offsets must be 1D");
-  TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be  a 2D tensor");
+  RuntimeCheck(a_ptrs.dim() == 1, "a_ptrs must be 1D");
-  TORCH_CHECK(problem_sizes.size(1) == 3, "problem_sizes must have the shape (num_experts, 3)");
+  RuntimeCheck(b_ptrs.dim() == 1, "b_ptrs must be 1D");
-  TORCH_CHECK(
+  RuntimeCheck(out_ptrs.dim() == 1, "out_ptrs must be 1D");
-      problem_sizes.size(0) == expert_offsets.size(0), "Number of experts in problem_sizes must match expert_offsets");
+  RuntimeCheck(a_scales_ptrs.dim() == 1, "a_scales_ptrs must be 1D");
-  TORCH_CHECK(problem_sizes.dtype() == torch::kInt32, "problem_sizes must be int32.");
+  RuntimeCheck(b_scales_ptrs.dim() == 1, "b_scales_ptrs must be 1D");
+  RuntimeCheck(alpha_ptrs.dim() == 1, "alpha_ptrs must be 1D");
+  RuntimeCheck(layout_sfa.dim() == 2, "layout_sfa must be 2D");
+  RuntimeCheck(layout_sfb.dim() == 2, "layout_sfb must be 2D");
+  RuntimeCheck(problem_sizes.size(1) == 3, "problem_sizes must have shape (num_experts, 3)");
+
+  const int num_experts = static_cast<int>(expert_offsets.size(0));
+  RuntimeCheck(problem_sizes.size(0) == num_experts, "problem_sizes size mismatch with expert_offsets");
+  RuntimeCheck(sf_offsets.size(0) == num_experts, "sf_offsets size mismatch with expert_offsets");
+  RuntimeCheck(alphas.size(0) == num_experts, "alphas size mismatch with expert_offsets");
+  RuntimeCheck(ab_strides.size(0) == num_experts, "ab_strides size mismatch with expert_offsets");
+  RuntimeCheck(c_strides.size(0) == num_experts, "c_strides size mismatch with expert_offsets");
+  RuntimeCheck(a_ptrs.size(0) == num_experts, "a_ptrs size mismatch with expert_offsets");
+  RuntimeCheck(b_ptrs.size(0) == num_experts, "b_ptrs size mismatch with expert_offsets");
+  RuntimeCheck(out_ptrs.size(0) == num_experts, "out_ptrs size mismatch with expert_offsets");
+  RuntimeCheck(a_scales_ptrs.size(0) == num_experts, "a_scales_ptrs size mismatch with expert_offsets");
+  RuntimeCheck(b_scales_ptrs.size(0) == num_experts, "b_scales_ptrs size mismatch with expert_offsets");
+  RuntimeCheck(alpha_ptrs.size(0) == num_experts, "alpha_ptrs size mismatch with expert_offsets");
+  RuntimeCheck(layout_sfa.size(0) == num_experts && layout_sfa.size(1) == 5, "layout_sfa must be [num_experts, 5]");
+  RuntimeCheck(layout_sfb.size(0) == num_experts && layout_sfb.size(1) == 5, "layout_sfb must be [num_experts, 5]");
 
   int M = static_cast<int>(a.size(0));
   int N = static_cast<int>(b.size(1));
-  int E = static_cast<int>(b.size(0));
   int K = static_cast<int>(2 * b.size(2));
+  RuntimeCheck(output.dim() == 2, "output must be 2D");
+  RuntimeCheck(output.size(0) == M && output.size(1) == N, "output shape mismatch");
 
-  auto sm_version = getSMVersion();
+  auto sm_version = getSMVersion(a.device().device_id);
   if (sm_version == 100 || sm_version == 103) {
-    if (output.scalar_type() == torch::kBFloat16) {
+    if (host::is_type<bf16_t>(output.dtype())) {
       run_fp4_blockwise_scaled_group_mm_sm100<cutlass::bfloat16_t>(
           output,
           a,
@@ -645,6 +770,14 @@ void cutlass_fp4_group_mm(
           problem_sizes,
           expert_offsets,
           sf_offsets,
+          a_ptrs,
+          b_ptrs,
+          out_ptrs,
+          a_scales_ptrs,
+          b_scales_ptrs,
+          alpha_ptrs,
+          layout_sfa,
+          layout_sfb,
           M,
           N,
           K);
@@ -661,12 +794,20 @@ void cutlass_fp4_group_mm(
           problem_sizes,
           expert_offsets,
           sf_offsets,
+          a_ptrs,
+          b_ptrs,
+          out_ptrs,
+          a_scales_ptrs,
+          b_scales_ptrs,
+          alpha_ptrs,
+          layout_sfa,
+          layout_sfb,
           M,
           N,
           K);
     }
   } else if (sm_version >= 120) {
-    if (output.scalar_type() == torch::kBFloat16) {
+    if (host::is_type<bf16_t>(output.dtype())) {
       run_fp4_blockwise_scaled_group_mm_sm120(
           output,
           a,
@@ -679,20 +820,63 @@ void cutlass_fp4_group_mm(
           problem_sizes,
           expert_offsets,
           sf_offsets,
+          a_ptrs,
+          b_ptrs,
+          out_ptrs,
+          a_scales_ptrs,
+          b_scales_ptrs,
+          alpha_ptrs,
+          layout_sfa,
+          layout_sfb,
           M,
           N,
           K);
     } else {
-      std::cout << "run_fp4_blockwise_scaled_group_mm_sm120 half no implementation" << std::endl;
+      Panic("SM120 path currently supports only bfloat16 output");
     }
   } else {
-    TORCH_CHECK_NOT_IMPLEMENTED(false, "Unsupported SM version: " + std::to_string(sm_version));
+    RuntimeCheck(false, "Unsupported SM version: ", sm_version);
   }
-#else
+}
-  TORCH_CHECK_NOT_IMPLEMENTED(
+
-      false,
+void cutlass_fp4_group_mm(
-      "No compiled cutlass_fp4_group_mm kernel, sgl-kernel must "
+    tvm::ffi::TensorView output,
-      "be compiled with ENABLE_NVFP4 for SM100+ and CUDA "
+    const tvm::ffi::TensorView a,
-      "12.8 or above.");
+    const tvm::ffi::TensorView b,
-#endif
+    const tvm::ffi::TensorView a_blockscale,
+    const tvm::ffi::TensorView b_blockscales,
+    const tvm::ffi::TensorView alphas,
+    const tvm::ffi::TensorView ab_strides,
+    const tvm::ffi::TensorView c_strides,
+    const tvm::ffi::TensorView problem_sizes,
+    const tvm::ffi::TensorView expert_offsets,
+    const tvm::ffi::TensorView sf_offsets,
+    const tvm::ffi::TensorView a_ptrs,
+    const tvm::ffi::TensorView b_ptrs,
+    const tvm::ffi::TensorView out_ptrs,
+    const tvm::ffi::TensorView a_scales_ptrs,
+    const tvm::ffi::TensorView b_scales_ptrs,
+    const tvm::ffi::TensorView alpha_ptrs,
+    const tvm::ffi::TensorView layout_sfa,
+    const tvm::ffi::TensorView layout_sfb) {
+  cutlass_fp4_group_mm_sm100a_sm120a(
+      output,
+      a,
+      b,
+      a_blockscale,
+      b_blockscales,
+      alphas,
+      ab_strides,
+      c_strides,
+      problem_sizes,
+      expert_offsets,
+      sf_offsets,
+      a_ptrs,
+      b_ptrs,
+      out_ptrs,
+      a_scales_ptrs,
+      b_scales_ptrs,
+      alpha_ptrs,
+      layout_sfa,
+      layout_sfb);
 }

python/sglang/jit_kernel/nvfp4.py

@@ -0,0 +1,490 @@
+from __future__ import annotations
+
+import importlib.util
+import os
+import pathlib
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Optional, Tuple
+
+import torch
+
+from sglang.jit_kernel.utils import cache_once, load_jit
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+_FLOAT4_E2M1_MAX = 6.0
+_FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
+
+
+def _find_package_root(package: str) -> Optional[pathlib.Path]:
+    spec = importlib.util.find_spec(package)
+    if spec is None or spec.origin is None:
+        return None
+    return pathlib.Path(spec.origin).resolve().parent
+
+
+def _resolve_cutlass_include_paths() -> list[str]:
+    include_paths: list[str] = []
+
+    flashinfer_root = _find_package_root("flashinfer")
+    if flashinfer_root is not None:
+        candidates = [
+            flashinfer_root / "data" / "cutlass" / "include",
+            flashinfer_root / "data" / "cutlass" / "tools" / "util" / "include",
+        ]
+        for path in candidates:
+            if path.exists():
+                include_paths.append(str(path))
+
+    deep_gemm_root = _find_package_root("deep_gemm")
+    if deep_gemm_root is not None:
+        candidate = deep_gemm_root / "include"
+        if candidate.exists():
+            include_paths.append(str(candidate))
+
+    # De-duplicate while preserving order.
+    unique_paths = []
+    seen = set()
+    for path in include_paths:
+        if path in seen:
+            continue
+        seen.add(path)
+        unique_paths.append(path)
+    return unique_paths
+
+
+def _nvfp4_cuda_flags() -> list[str]:
+    return [
+        "-DNDEBUG",
+        "-DFLASHINFER_ENABLE_F16",
+        "-DCUTE_USE_PACKED_TUPLE=1",
+        "-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1",
+        "-DCUTLASS_VERSIONS_GENERATED",
+        "-DCUTLASS_TEST_LEVEL=0",
+        "-DCUTLASS_TEST_ENABLE_CACHED_RESULTS=1",
+        "-DCUTLASS_DEBUG_TRACE_LEVEL=0",
+        "--expt-extended-lambda",
+    ]
+
+
+def _parse_cuda_version() -> tuple[int, int]:
+    v = torch.version.cuda
+    if not v:
+        return (0, 0)
+    parts = v.split(".")
+    if len(parts) < 2:
+        return (0, 0)
+    try:
+        return int(parts[0]), int(parts[1])
+    except ValueError:
+        return (0, 0)
+
+
+def _get_nvfp4_cuda_arch_list() -> str:
+    if not torch.cuda.is_available():
+        raise RuntimeError("NVFP4 JIT kernels require CUDA.")
+    major, minor = torch.cuda.get_device_capability()
+    if major < 10:
+        raise RuntimeError(
+            f"NVFP4 JIT kernels require compute capability >= 10.0, got {major}.{minor}."
+        )
+    # NVFP4 kernels use architecture-family-specific instructions and must be
+    # compiled for `sm_*a` targets (e.g. sm_100a), not plain sm_100.
+    archs = [f"{major}.{minor}a"]
+    cuda_major, _cuda_minor = _parse_cuda_version()
+    if cuda_major >= 13 and "10.3a" not in archs:
+        # Match sgl-kernel AOT fatbin behavior on CUDA 13+ for Blackwell.
+        archs.append("10.3a")
+    # Preserve order while de-duplicating.
+    seen = set()
+    ordered_archs: list[str] = []
+    for arch in archs:
+        if arch in seen:
+            continue
+        seen.add(arch)
+        ordered_archs.append(arch)
+    return " ".join(ordered_archs)
+
+
+@contextmanager
+def _nvfp4_arch_env():
+    key = "TVM_FFI_CUDA_ARCH_LIST"
+    old_val = os.environ.get(key)
+    os.environ[key] = _get_nvfp4_cuda_arch_list()
+    try:
+        yield
+    finally:
+        if old_val is None:
+            os.environ.pop(key, None)
+        else:
+            os.environ[key] = old_val
+
+
+@cache_once
+def _jit_nvfp4_quant_module() -> Module:
+    extra_include_paths = _resolve_cutlass_include_paths()
+    if not extra_include_paths:
+        raise RuntimeError(
+            "Cannot find CUTLASS headers required for NVFP4 JIT quantization. "
+            "Please install flashinfer or deep_gemm with CUTLASS headers."
+        )
+
+    with _nvfp4_arch_env():
+        return load_jit(
+            "nvfp4_quant",
+            cuda_files=[
+                "gemm/nvfp4/nvfp4_quant_kernels.cuh",
+            ],
+            cuda_wrappers=[
+                ("scaled_fp4_quant", "scaled_fp4_quant_sm100a_sm120a"),
+            ],
+            extra_include_paths=extra_include_paths,
+            extra_cuda_cflags=_nvfp4_cuda_flags(),
+        )
+
+
+@cache_once
+def _jit_nvfp4_expert_quant_module() -> Module:
+    extra_include_paths = _resolve_cutlass_include_paths()
+    if not extra_include_paths:
+        raise RuntimeError(
+            "Cannot find CUTLASS headers required for NVFP4 JIT expert quantization. "
+            "Please install flashinfer or deep_gemm with CUTLASS headers."
+        )
+
+    with _nvfp4_arch_env():
+        return load_jit(
+            "nvfp4_expert_quant",
+            cuda_files=[
+                "gemm/nvfp4/nvfp4_expert_quant.cuh",
+            ],
+            cuda_wrappers=[
+                ("scaled_fp4_experts_quant", "scaled_fp4_experts_quant_sm100a"),
+                (
+                    "silu_and_mul_scaled_fp4_experts_quant",
+                    "silu_and_mul_scaled_fp4_experts_quant_sm100a",
+                ),
+            ],
+            extra_include_paths=extra_include_paths,
+            extra_cuda_cflags=_nvfp4_cuda_flags(),
+        )
+
+
+@cache_once
+def _jit_nvfp4_scaled_mm_module() -> Module:
+    extra_include_paths = _resolve_cutlass_include_paths()
+    if not extra_include_paths:
+        raise RuntimeError(
+            "Cannot find CUTLASS headers required for NVFP4 JIT GEMM. "
+            "Please install flashinfer or deep_gemm with CUTLASS headers."
+        )
+
+    with _nvfp4_arch_env():
+        return load_jit(
+            "nvfp4_scaled_mm",
+            cuda_files=[
+                "gemm/nvfp4/nvfp4_scaled_mm_kernels.cuh",
+                "gemm/nvfp4/nvfp4_scaled_mm_entry.cuh",
+            ],
+            cuda_wrappers=[("cutlass_scaled_fp4_mm", "cutlass_scaled_fp4_mm")],
+            extra_include_paths=extra_include_paths,
+            extra_cuda_cflags=_nvfp4_cuda_flags(),
+        )
+
+
+@cache_once
+def _jit_nvfp4_blockwise_moe_module() -> Module:
+    extra_include_paths = _resolve_cutlass_include_paths()
+    if not extra_include_paths:
+        raise RuntimeError(
+            "Cannot find CUTLASS headers required for NVFP4 JIT MoE grouped GEMM. "
+            "Please install flashinfer or deep_gemm with CUTLASS headers."
+        )
+
+    with _nvfp4_arch_env():
+        return load_jit(
+            "nvfp4_blockwise_moe",
+            cuda_files=[
+                "moe/nvfp4_blockwise_moe.cuh",
+            ],
+            cuda_wrappers=[
+                ("cutlass_fp4_group_mm", "cutlass_fp4_group_mm_sm100a_sm120a")
+            ],
+            extra_include_paths=extra_include_paths,
+            extra_cuda_cflags=_nvfp4_cuda_flags(),
+        )
+
+
+def 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:
+    assert a.ndim == 2 and b.ndim == 2
+    m, n = a.shape[0], b.shape[0]
+    out = torch.empty((m, n), dtype=out_dtype, device=a.device)
+    module = _jit_nvfp4_scaled_mm_module()
+    module.cutlass_scaled_fp4_mm(out, a, b, block_scale_a, block_scale_b, alpha)
+    return out
+
+
+def cutlass_fp4_group_mm(
+    a_fp4: torch.Tensor,
+    b_fp4: torch.Tensor,
+    a_blockscale: torch.Tensor,
+    b_blockscale: torch.Tensor,
+    alphas: torch.Tensor,
+    out_dtype: torch.dtype,
+    params: dict[str, torch.Tensor],
+) -> torch.Tensor:
+    m_topk = a_fp4.shape[0]
+    n = b_fp4.shape[1]
+    output = torch.empty((m_topk, n), device=a_fp4.device, dtype=out_dtype)
+    num_experts = int(params["expert_offsets"].numel())
+    device = a_fp4.device
+
+    # Backward compatibility: older callers may not pass scratch tensors.
+    a_ptrs = params.get(
+        "a_ptrs", torch.empty((num_experts,), dtype=torch.int64, device=device)
+    )
+    b_ptrs = params.get(
+        "b_ptrs", torch.empty((num_experts,), dtype=torch.int64, device=device)
+    )
+    out_ptrs = params.get(
+        "out_ptrs", torch.empty((num_experts,), dtype=torch.int64, device=device)
+    )
+    a_scales_ptrs = params.get(
+        "a_scales_ptrs", torch.empty((num_experts,), dtype=torch.int64, device=device)
+    )
+    b_scales_ptrs = params.get(
+        "b_scales_ptrs", torch.empty((num_experts,), dtype=torch.int64, device=device)
+    )
+    alpha_ptrs = params.get(
+        "alpha_ptrs", torch.empty((num_experts,), dtype=torch.int64, device=device)
+    )
+    layout_sfa = params.get(
+        "layout_sfa", torch.empty((num_experts, 5), dtype=torch.int64, device=device)
+    )
+    layout_sfb = params.get(
+        "layout_sfb", torch.empty((num_experts, 5), dtype=torch.int64, device=device)
+    )
+
+    module = _jit_nvfp4_blockwise_moe_module()
+    module.cutlass_fp4_group_mm(
+        output,
+        a_fp4,
+        b_fp4,
+        a_blockscale,
+        b_blockscale,
+        alphas,
+        params["ab_strides"],
+        params["c_strides"],
+        params["problem_sizes"],
+        params["expert_offsets"],
+        params["blockscale_offsets"],
+        a_ptrs,
+        b_ptrs,
+        out_ptrs,
+        a_scales_ptrs,
+        b_scales_ptrs,
+        alpha_ptrs,
+        layout_sfa,
+        layout_sfb,
+    )
+    return output
+
+
+def scaled_fp4_quant(
+    input: torch.Tensor, input_global_scale: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Quantize input tensor to FP4 and return packed FP4 tensor + swizzled scales."""
+    assert input.ndim >= 1, f"input.ndim needs to be >= 1, but got {input.ndim}."
+    other_dims = 1 if input.ndim == 1 else -1
+    input = input.reshape(other_dims, input.shape[-1])
+    m, n = input.shape
+    block_size = 16
+    device = input.device
+
+    assert n % block_size == 0, f"last dim has to be multiple of 16, but got {n}."
+    assert input.dtype in (
+        torch.float16,
+        torch.bfloat16,
+    ), f"input.dtype needs to be fp16 or bf16 but got {input.dtype}."
+
+    output = torch.empty((m, n // 2), device=device, dtype=torch.uint8)
+
+    rounded_m = ((m + 128 - 1) // 128) * 128
+    scale_n = n // block_size
+    rounded_n = ((scale_n + 4 - 1) // 4) * 4
+    if rounded_n > scale_n:
+        output_scale = torch.zeros(
+            (rounded_m, rounded_n // 4), device=device, dtype=torch.int32
+        )
+    else:
+        output_scale = torch.empty(
+            (rounded_m, rounded_n // 4), device=device, dtype=torch.int32
+        )
+
+    module = _jit_nvfp4_quant_module()
+    module.scaled_fp4_quant(output, input, output_scale, input_global_scale)
+    output_scale = output_scale.view(torch.float8_e4m3fn)
+    return output, output_scale
+
+
+def _shuffle_rows_torch(
+    input_tensor: torch.Tensor,
+    dst2src_map: torch.Tensor,
+    output_tensor_shape: tuple[int, int],
+) -> torch.Tensor:
+    # Keep compatibility when sgl-kernel is slimmed and shuffle_rows may not be present.
+    output = input_tensor.index_select(0, dst2src_map.to(dtype=torch.int64))
+    return output.view(output_tensor_shape)
+
+
+def scaled_fp4_experts_quant(
+    input_tensor: torch.Tensor,
+    input_global_scale: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    blockscale_offsets: torch.Tensor,
+    topk: int,
+    expert_map: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize packed MoE activations to NVFP4."""
+    assert (
+        input_tensor.ndim == 2
+    ), f"input.ndim needs to be == 2, but got {input_tensor.ndim}."
+    if expert_map is not None:
+        m, k = input_tensor.shape
+        output_tensor_shape = (m * topk, k)
+        input_tensor = _shuffle_rows_torch(
+            input_tensor, expert_map, output_tensor_shape
+        )
+
+    m_numtopk, k = input_tensor.shape
+    max_tokens_per_expert = int(os.environ.get("MODELOPT_MAX_TOKENS_PER_EXPERT", 65536))
+    assert m_numtopk <= max_tokens_per_expert * topk, (
+        f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT({max_tokens_per_expert})"
+        f" for cutlass_moe_fp4, observed m_numtopk = {m_numtopk}. Use"
+        " MODELOPT_MAX_TOKENS_PER_EXPERT to set this value."
+    )
+    scales_k = k // 16
+    # output_scales is int32-packed FP8 scales, so second dim is in int32 units.
+    padded_k_in_int32 = (scales_k + 3) // 4
+
+    output = torch.empty(
+        m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
+    )
+    if padded_k_in_int32 * 4 > scales_k:
+        output_scales = torch.zeros(
+            max_tokens_per_expert * topk,
+            padded_k_in_int32,
+            dtype=torch.int32,
+            device=input_tensor.device,
+        )
+    else:
+        output_scales = torch.empty(
+            max_tokens_per_expert * topk,
+            padded_k_in_int32,
+            dtype=torch.int32,
+            device=input_tensor.device,
+        )
+
+    module = _jit_nvfp4_expert_quant_module()
+    module.scaled_fp4_experts_quant(
+        output,
+        output_scales,
+        input_tensor,
+        input_global_scale,
+        expert_offsets,
+        blockscale_offsets,
+    )
+    output_scales = output_scales.view(torch.float8_e4m3fn)
+    return output, output_scales
+
+
+def scaled_fp4_grouped_quant(
+    input_tensor: torch.Tensor,
+    input_global_scale: torch.Tensor,
+    mask: torch.Tensor,
+):
+    """Quantize grouped GEMM inputs to FP4 and return logical (m, k//2, l)."""
+    device = input_tensor.device
+    l, m, k = input_tensor.shape
+    sf_vec_size = 16
+    assert k % sf_vec_size == 0, f"k must be multiple of 16, but got {k}."
+
+    scale_k = k // sf_vec_size
+    padded_k = (scale_k + (4 - 1)) // 4 * 4
+    padded_k_int32 = padded_k // 4
+    padded_m = (m + (128 - 1)) // 128 * 128
+    output = torch.empty(l, m, k // 2, device=device, dtype=torch.uint8)
+    output_scales = torch.empty(
+        l, padded_m, padded_k_int32, device=device, dtype=torch.int32
+    )
+
+    module = _jit_nvfp4_expert_quant_module()
+    module.silu_and_mul_scaled_fp4_experts_quant(
+        output.view(l * m, k // 2),
+        output_scales.view(l * padded_m, padded_k_int32),
+        input_tensor.view(l * m, k),
+        input_global_scale,
+        mask,
+        False,
+    )
+
+    output = output.permute(1, 2, 0)
+    output_scales = output_scales.view(torch.float8_e4m3fn).view(
+        l, padded_m // 128, padded_k // 4, 32, 4, 4
+    )
+    output_scales = output_scales.permute(3, 4, 1, 5, 2, 0)
+    return output, output_scales
+
+
+def silu_and_mul_scaled_fp4_grouped_quant(
+    input_tensor: torch.Tensor,
+    input_global_scale: torch.Tensor,
+    mask: torch.Tensor,
+):
+    """Apply SiLU-and-mul then quantize grouped GEMM inputs to FP4."""
+    device = input_tensor.device
+    l, m, k_by_2 = input_tensor.shape
+    k = k_by_2 // 2
+    sf_vec_size = 16
+    assert k % sf_vec_size == 0, f"k must be multiple of 16, but got {k}."
+
+    scale_k = k // sf_vec_size
+    padded_k = (scale_k + (4 - 1)) // 4 * 4
+    padded_k_int32 = padded_k // 4
+    padded_m = (m + (128 - 1)) // 128 * 128
+    output = torch.empty(l, m, k // 2, device=device, dtype=torch.uint8)
+    output_scales = torch.empty(
+        l, padded_m, padded_k_int32, device=device, dtype=torch.int32
+    )
+
+    module = _jit_nvfp4_expert_quant_module()
+    module.silu_and_mul_scaled_fp4_experts_quant(
+        output.view(l * m, k // 2),
+        output_scales.view(l * padded_m, padded_k_int32),
+        input_tensor.view(l * m, k_by_2),
+        input_global_scale,
+        mask,
+        True,
+    )
+
+    output = output.permute(1, 2, 0)
+    output_scales = output_scales.view(torch.float8_e4m3fn).view(
+        l, padded_m // 128, padded_k // 4, 32, 4, 4
+    )
+    output_scales = output_scales.permute(3, 4, 1, 5, 2, 0)
+    return output, output_scales
+
+
+def suggest_nvfp4_global_scale(x: torch.Tensor) -> torch.Tensor:
+    """Utility for tests/benchmarks: return global scale used by NVFP4 quantization."""
+    tensor_amax = torch.abs(x).max().to(torch.float32)
+    return _FLOAT8_E4M3_MAX * _FLOAT4_E2M1_MAX / tensor_amax

python/sglang/jit_kernel/tests/test_nvfp4_blockwise_moe.py

@@ -0,0 +1,127 @@
+import pytest
+import torch
+
+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 _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)

python/sglang/jit_kernel/tests/test_nvfp4_gemm.py

@@ -0,0 +1,142 @@
+import pytest
+import torch
+
+from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm, scaled_fp4_quant
+
+
+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)

python/sglang/jit_kernel/tests/test_nvfp4_quant.py

@@ -0,0 +1,214 @@
+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
+
+
+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]])

python/sglang/srt/layers/moe/cutlass_moe.py

@@ -11,17 +11,20 @@ _is_cuda = is_cuda()
 if _is_cuda:
     from sgl_kernel import (
         apply_shuffle_mul_sum,
-        cutlass_fp4_group_mm,
         es_fp8_blockwise_scaled_grouped_mm,
         es_sm100_mxfp8_blockscaled_grouped_mm,
         es_sm100_mxfp8_blockscaled_grouped_quant,
         fp8_blockwise_scaled_grouped_mm,
         prepare_moe_input,
-        scaled_fp4_experts_quant,
         shuffle_rows,
         silu_and_mul,
     )
 
+    from sglang.jit_kernel.nvfp4 import (
+        cutlass_fp4_group_mm,
+        scaled_fp4_experts_quant,
+    )
+
 
 def cutlass_fused_experts_fp8(
     a: torch.Tensor,

python/sglang/srt/layers/moe/cutlass_moe_params.py

@@ -72,6 +72,11 @@ class CutlassMoEParams:
     # b_scales_ptrs: [e] dtype: int64
     a_scales_ptrs: torch.Tensor
     b_scales_ptrs: torch.Tensor
+    # Pointers for per-expert alpha values
+    alpha_ptrs: torch.Tensor
+    # CUTLASS blockscale layouts for A and B operands
+    layout_sfa: torch.Tensor
+    layout_sfb: torch.Tensor
 
     # Offsets that mark at which token index each expert begins its computation
     # The number of tokens computed with expert E is expert_offsets[E + 1] - expert_offsets[E]
@@ -139,6 +144,13 @@ class CutlassMoEParams:
         self.b_scales_ptrs = torch.empty(
             (self.e,), dtype=torch.int64, device=self.device
         )
+        self.alpha_ptrs = torch.empty((self.e,), dtype=torch.int64, device=self.device)
+        self.layout_sfa = torch.empty(
+            (self.e, 5), dtype=torch.int64, device=self.device
+        )
+        self.layout_sfb = torch.empty(
+            (self.e, 5), dtype=torch.int64, device=self.device
+        )
 
     def to_gemm1_args(self) -> dict:
         return {
@@ -147,11 +159,14 @@ class CutlassMoEParams:
             "problem_sizes": self.problem_sizes1,
             "expert_offsets": self.expert_offsets[:-1],
             "blockscale_offsets": self.blockscale_offsets[:-1],
-            #    "a_ptrs": self.a_ptrs,
+            "a_ptrs": self.a_ptrs,
-            #    "b_ptrs": self.b_ptrs,
+            "b_ptrs": self.b_ptrs,
-            #    "out_ptrs": self.out_ptrs,
+            "out_ptrs": self.out_ptrs,
-            #    "a_scales_ptrs": self.a_scales_ptrs,
+            "a_scales_ptrs": self.a_scales_ptrs,
-            #    "b_scales_ptrs": self.b_scales_ptrs,
+            "b_scales_ptrs": self.b_scales_ptrs,
+            "alpha_ptrs": self.alpha_ptrs,
+            "layout_sfa": self.layout_sfa,
+            "layout_sfb": self.layout_sfb,
         }
 
     def to_gemm2_args(self) -> dict:
@@ -161,9 +176,12 @@ class CutlassMoEParams:
             "problem_sizes": self.problem_sizes2,
             "expert_offsets": self.expert_offsets[:-1],
             "blockscale_offsets": self.blockscale_offsets[:-1],
-            #    "a_ptrs": self.a_ptrs,
+            "a_ptrs": self.a_ptrs,
-            #    "b_ptrs": self.b_ptrs,
+            "b_ptrs": self.b_ptrs,
-            #    "out_ptrs": self.out_ptrs,
+            "out_ptrs": self.out_ptrs,
-            #    "a_scales_ptrs": self.a_scales_ptrs,
+            "a_scales_ptrs": self.a_scales_ptrs,
-            #    "b_scales_ptrs": self.b_scales_ptrs,
+            "b_scales_ptrs": self.b_scales_ptrs,
+            "alpha_ptrs": self.alpha_ptrs,
+            "layout_sfa": self.layout_sfa,
+            "layout_sfb": self.layout_sfb,
         }

python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py

@@ -38,7 +38,7 @@ if TYPE_CHECKING:
 if is_flashinfer_available() and is_sm120_supported():
     from flashinfer import fp4_quantize
 elif is_cuda_alike():
-    from sgl_kernel import scaled_fp4_quant as fp4_quantize
+    from sglang.jit_kernel.nvfp4 import scaled_fp4_quant as fp4_quantize
 else:
     fp4_quantize = None
 

python/sglang/srt/layers/moe/token_dispatcher/standard.py

@@ -43,7 +43,7 @@ try:
     if is_sm120_supported():
         from flashinfer import fp4_quantize
     else:
-        from sgl_kernel import scaled_fp4_quant as fp4_quantize
+        from sglang.jit_kernel.nvfp4 import scaled_fp4_quant as fp4_quantize
 
     from flashinfer import fp4_quantize as fp4_quantize_flashinfer
 except ImportError:

python/sglang/srt/layers/quantization/modelopt_quant.py

@@ -74,9 +74,9 @@ try:
         try:
             from flashinfer import fp4_quantize
         except ImportError:
-            from sgl_kernel import scaled_fp4_quant as fp4_quantize
+            from sglang.jit_kernel.nvfp4 import scaled_fp4_quant as fp4_quantize
     else:
-        from sgl_kernel import scaled_fp4_quant as fp4_quantize
+        from sglang.jit_kernel.nvfp4 import scaled_fp4_quant as fp4_quantize
 except ImportError:
     fp4_quantize = None
 
@@ -87,7 +87,7 @@ try:
     enable_flashinfer_fp4_gemm = True
 except ImportError:
     if is_cuda():
-        from sgl_kernel import cutlass_scaled_fp4_mm as cutlass_fp4_gemm
+        from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm as cutlass_fp4_gemm
     enable_flashinfer_fp4_gemm = False
     reorder_rows_for_gated_act_gemm = None
     shuffle_matrix_a = None

sgl-kernel/CMakeLists.txt

@@ -286,11 +286,6 @@ set(SOURCES
     "csrc/gemm/fp8_blockwise_gemm_kernel.cu"
     "csrc/gemm/fp8_gemm_kernel.cu"
     "csrc/gemm/int8_gemm_kernel.cu"
-    "csrc/gemm/nvfp4_expert_quant.cu"
-    "csrc/gemm/nvfp4_quant_entry.cu"
-    "csrc/gemm/nvfp4_quant_kernels.cu"
-    "csrc/gemm/nvfp4_scaled_mm_entry.cu"
-    "csrc/gemm/nvfp4_scaled_mm_kernels.cu"
     "csrc/gemm/per_tensor_quant_fp8.cu"
     "csrc/gemm/per_token_group_quant_8bit.cu"
     "csrc/gemm/per_token_group_quant_8bit_v2.cu"
@@ -316,7 +311,6 @@ set(SOURCES
     "csrc/moe/moe_sum_reduce.cu"
     "csrc/moe/moe_topk_softmax_kernels.cu"
     "csrc/moe/moe_topk_sigmoid_kernels.cu"
-    "csrc/moe/nvfp4_blockwise_moe.cu"
     "csrc/moe/fp8_blockwise_moe_kernel.cu"
     "csrc/moe/prepare_moe_input.cu"
 

sgl-kernel/csrc/common_extension.cc

@@ -157,39 +157,9 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
   m.def("sgl_per_token_quant_fp8(Tensor input, Tensor! output_q, Tensor! output_s) -> ()");
   m.impl("sgl_per_token_quant_fp8", torch::kCUDA, &sgl_per_token_quant_fp8);
 
-  m.def(
-      "cutlass_scaled_fp4_mm(Tensor! out, Tensor a, Tensor b,"
-      "                      Tensor block_scale_a, Tensor block_scale_b,"
-      "                      Tensor alpha) -> ()");
-  m.impl("cutlass_scaled_fp4_mm", torch::kCUDA, &cutlass_scaled_fp4_mm);
-
-  m.def(
-      "scaled_fp4_quant(Tensor! output, Tensor! input,"
-      "                 Tensor! output_scale, Tensor! input_scale) -> ()");
-  m.impl("scaled_fp4_quant", torch::kCUDA, &scaled_fp4_quant);
-
   m.def("dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
   m.impl("dsv3_fused_a_gemm", torch::kCUDA, &dsv3_fused_a_gemm);
 
-  // Compute NVFP4 experts quantization.
-  m.def(
-      "scaled_fp4_experts_quant(Tensor! output, Tensor! output_scale,"
-      "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts,"
-      "Tensor output_scale_offset_by_experts) -> ()");
-  m.impl("scaled_fp4_experts_quant", torch::kCUDA, &scaled_fp4_experts_quant);
-
-  m.def(
-      "silu_and_mul_scaled_fp4_experts_quant(Tensor! output, Tensor! output_scale,"
-      "Tensor input, Tensor input_global_scale, Tensor mask, bool use_silu_and_mul) -> ()");
-  m.impl("silu_and_mul_scaled_fp4_experts_quant", torch::kCUDA, &silu_and_mul_scaled_fp4_experts_quant);
-
-  m.def(
-      "cutlass_fp4_group_mm(Tensor! output, Tensor a, Tensor b,"
-      "Tensor a_blockscale, Tensor b_blockscale, Tensor alphas,"
-      "Tensor ab_strides, Tensor c_strides, Tensor problem_sizes,"
-      " Tensor expert_offsets, Tensor sf_offsets) -> ()");
-  m.impl("cutlass_fp4_group_mm", torch::kCUDA, &cutlass_fp4_group_mm);
-
   m.def("dsv3_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
   m.impl("dsv3_router_gemm", torch::kCUDA, &dsv3_router_gemm);
 

sgl-kernel/csrc/gemm/nvfp4_quant_entry.cu

@@ -1,77 +0,0 @@
-/* Copyright 2025 SGLang Team. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-==============================================================================*/
-
-#include <torch/all.h>
-
-#if defined ENABLE_NVFP4 && ENABLE_NVFP4
-void scaled_fp4_quant_sm100a_sm120a(
-    torch::Tensor const& output,
-    torch::Tensor const& input,
-    torch::Tensor const& output_sf,
-    torch::Tensor const& input_sf);
-
-void scaled_fp4_experts_quant_sm100a(
-    torch::Tensor& output,
-    torch::Tensor& output_scale,
-    torch::Tensor const& input,
-    torch::Tensor const& input_global_scale,
-    torch::Tensor const& input_offset_by_experts,
-    torch::Tensor const& output_scale_offset_by_experts);
-
-void silu_and_mul_scaled_fp4_experts_quant_sm100a(
-    torch::Tensor& output,
-    torch::Tensor& output_scale,
-    torch::Tensor const& input,
-    torch::Tensor const& input_global_scale,
-    torch::Tensor const& mask,
-    bool use_silu_and_mul);
-
-#endif
-
-void scaled_fp4_quant(
-    torch::Tensor& output, torch::Tensor const& input, torch::Tensor& output_sf, torch::Tensor const& input_sf) {
-#if defined ENABLE_NVFP4 && ENABLE_NVFP4
-  return scaled_fp4_quant_sm100a_sm120a(output, input, output_sf, input_sf);
-#endif
-  TORCH_CHECK_NOT_IMPLEMENTED(false, "No compiled nvfp4 quantization");
-}
-
-void scaled_fp4_experts_quant(
-    torch::Tensor& output,
-    torch::Tensor& output_scale,
-    torch::Tensor const& input,
-    torch::Tensor const& input_global_scale,
-    torch::Tensor const& input_offset_by_experts,
-    torch::Tensor const& output_scale_offset_by_experts) {
-#if defined ENABLE_NVFP4 && ENABLE_NVFP4
-  return scaled_fp4_experts_quant_sm100a(
-      output, output_scale, input, input_global_scale, input_offset_by_experts, output_scale_offset_by_experts);
-#endif
-  TORCH_CHECK_NOT_IMPLEMENTED(false, "No compiled nvfp4 experts quantization kernel");
-}
-
-void silu_and_mul_scaled_fp4_experts_quant(
-    torch::Tensor& output,
-    torch::Tensor& output_scale,
-    torch::Tensor const& input,
-    torch::Tensor const& input_global_scale,
-    torch::Tensor const& mask,
-    bool use_silu_and_mul) {
-#if defined ENABLE_NVFP4 && ENABLE_NVFP4
-  return silu_and_mul_scaled_fp4_experts_quant_sm100a(
-      output, output_scale, input, input_global_scale, mask, use_silu_and_mul);
-#endif
-  TORCH_CHECK_NOT_IMPLEMENTED(false, "No compiled nvfp4 experts quantization kernel");
-}

sgl-kernel/csrc/gemm/nvfp4_scaled_mm_entry.cu

@@ -1,64 +0,0 @@
-/* Copyright 2025 SGLang Team. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-==============================================================================*/
-
-#include <torch/all.h>
-
-#if defined ENABLE_NVFP4 && ENABLE_NVFP4
-void cutlass_scaled_fp4_mm_sm100a_sm120a(
-    torch::Tensor& D,
-    torch::Tensor const& A,
-    torch::Tensor const& B,
-    torch::Tensor const& A_sf,
-    torch::Tensor const& B_sf,
-    torch::Tensor const& alpha);
-
-// SM120 specific dispatch functions
-void cutlass_fp4_bf16_gemm_dispatch_sm120(
-    torch::Tensor& D,
-    torch::Tensor const& A,
-    torch::Tensor const& B,
-    torch::Tensor const& A_sf,
-    torch::Tensor const& B_sf,
-    torch::Tensor const& alpha,
-    int m,
-    int n,
-    int k,
-    cudaStream_t stream);
-
-void cutlass_fp4_f16_gemm_dispatch_sm120(
-    torch::Tensor& D,
-    torch::Tensor const& A,
-    torch::Tensor const& B,
-    torch::Tensor const& A_sf,
-    torch::Tensor const& B_sf,
-    torch::Tensor const& alpha,
-    int m,
-    int n,
-    int k,
-    cudaStream_t stream);
-#endif
-
-void cutlass_scaled_fp4_mm(
-    torch::Tensor& D,
-    torch::Tensor const& A,
-    torch::Tensor const& B,
-    torch::Tensor const& A_sf,
-    torch::Tensor const& B_sf,
-    torch::Tensor const& alpha) {
-#if defined ENABLE_NVFP4 && ENABLE_NVFP4
-  return cutlass_scaled_fp4_mm_sm100a_sm120a(D, A, B, A_sf, B_sf, alpha);
-#endif
-  TORCH_CHECK_NOT_IMPLEMENTED(false, "No compiled nvfp4 mm kernel.");
-}

sgl-kernel/include/sgl_kernel_ops.h

@@ -199,13 +199,6 @@ void gelu_quick(at::Tensor& out, const at::Tensor& input);
  * From csrc/gemm
  */
 torch::Tensor awq_dequantize(torch::Tensor qweight, torch::Tensor scales, torch::Tensor qzeros);
-void cutlass_scaled_fp4_mm(
-    torch::Tensor& D,
-    torch::Tensor const& A,
-    torch::Tensor const& B,
-    torch::Tensor const& A_sf,
-    torch::Tensor const& B_sf,
-    torch::Tensor const& alpha);
 torch::Tensor int8_scaled_mm(
     const torch::Tensor& mat_a,
     const torch::Tensor& mat_b,
@@ -226,8 +219,6 @@ torch::Tensor fp8_blockwise_scaled_mm(
     const torch::Tensor& scales_a,
     const torch::Tensor& scales_b,
     const torch::Dtype& out_dtype);
-void scaled_fp4_quant(
-    torch::Tensor& output, torch::Tensor const& input, torch::Tensor& output_scale, torch::Tensor const& input_scale);
 void sgl_per_token_group_quant_8bit(
     at::Tensor input,
     at::Tensor output_q,
@@ -380,35 +371,6 @@ void fused_qk_norm_rope(
     double attention_factor,
     int64_t rotary_dim);
 
-void cutlass_fp4_group_mm(
-    torch::Tensor& output,
-    const torch::Tensor& a,
-    const torch::Tensor& b,
-    const torch::Tensor& a_blockscale,
-    const torch::Tensor& b_blockscales,
-    const torch::Tensor& alphas,
-    const torch::Tensor& ab_strides,
-    const torch::Tensor& c_strides,
-    const torch::Tensor& problem_sizes,
-    const torch::Tensor& expert_offsets,
-    const torch::Tensor& sf_offsets);
-
-void scaled_fp4_experts_quant(
-    torch::Tensor& output,
-    torch::Tensor& output_scale,
-    torch::Tensor const& input,
-    torch::Tensor const& input_global_scale,
-    torch::Tensor const& input_offset_by_experts,
-    torch::Tensor const& output_scale_offset_by_experts);
-
-void silu_and_mul_scaled_fp4_experts_quant(
-    torch::Tensor& output,
-    torch::Tensor& output_scale,
-    torch::Tensor const& input,
-    torch::Tensor const& input_global_scale,
-    torch::Tensor const& mask,
-    bool use_silu_and_mul);
-
 /*
  * From csrc/moe/cutlass_moe/w4a8
  */

sgl-kernel/python/sgl_kernel/__init__.py

@@ -51,7 +51,6 @@ from sgl_kernel.gemm import (
     int8_scaled_mm,
     qserve_w4a8_per_chn_gemm,
     qserve_w4a8_per_group_gemm,
-    scaled_fp4_experts_quant,
     scaled_fp4_grouped_quant,
     scaled_fp4_quant,
     sgl_per_tensor_quant_fp8,
@@ -79,7 +78,6 @@ from sgl_kernel.mamba import (
 from sgl_kernel.memory import set_kv_buffer_kernel, weak_ref_tensor
 from sgl_kernel.moe import (
     apply_shuffle_mul_sum,
-    cutlass_fp4_group_mm,
     fp8_blockwise_scaled_grouped_mm,
     fused_qk_norm_rope,
     kimi_k2_moe_fused_gate,

sgl-kernel/python/sgl_kernel/gemm.py

@@ -167,13 +167,18 @@ def cutlass_scaled_fp4_mm(
     alpha: torch.Tensor,
     out_dtype: torch.dtype,
 ) -> torch.Tensor:
-    assert a.ndim == 2 and b.ndim == 2
+    from sglang.jit_kernel.nvfp4 import (
-    m, n = a.shape[0], b.shape[0]
+        cutlass_scaled_fp4_mm as jit_cutlass_scaled_fp4_mm,
-    out = torch.empty((m, n), 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 jit_cutlass_scaled_fp4_mm(
+        a,
+        b,
+        block_scale_a,
+        block_scale_b,
+        alpha,
+        out_dtype,
     )
-    return out
 
 
 def scaled_fp4_quant(
@@ -197,42 +202,9 @@ def scaled_fp4_quant(
             two values are packed into a uint8 and float8_e4m3 scaling factors
             in a sizzled layout.
     """
-    assert input.ndim >= 1, f"input.ndim needs to be >= 1, but got {input.ndim}."
+    from sglang.jit_kernel.nvfp4 import scaled_fp4_quant as jit_scaled_fp4_quant
-    other_dims = 1 if input.ndim == 1 else -1
-    input = input.reshape(other_dims, input.shape[-1])
-    m, n = input.shape
-    block_size = 16
-    device = input.device
 
-    assert n % block_size == 0, f"last dim has to be multiple of 16, but got {n}."
+    return jit_scaled_fp4_quant(input, input_global_scale)
-    assert input.dtype in (
-        torch.float16,
-        torch.bfloat16,
-    ), f"input.dtype needs to be fp16 or bf16 but got {input.dtype}."
-
-    # Two fp4 values will be packed into an uint8.
-    output = torch.empty((m, n // 2), device=device, dtype=torch.uint8)
-
-    # We use the rounded values to store the swizzled values. Then, the scaling
-    # factors in float8_e4m3fn are packed into an int32 for every 4 values.
-    rounded_m = ((m + 128 - 1) // 128) * 128
-    scale_n = n // block_size
-    rounded_n = ((scale_n + 4 - 1) // 4) * 4
-    # padded part should be zeroed out
-    if rounded_n > scale_n:
-        output_scale = torch.zeros(
-            (rounded_m, rounded_n // 4), device=device, dtype=torch.int32
-        )
-    else:
-        output_scale = torch.empty(
-            (rounded_m, rounded_n // 4), device=device, dtype=torch.int32
-        )
-
-    torch.ops.sgl_kernel.scaled_fp4_quant.default(
-        output, input, output_scale, input_global_scale
-    )
-    output_scale = output_scale.view(torch.float8_e4m3fn)
-    return output, output_scale
 
 
 def qserve_w4a8_per_chn_gemm(
@@ -332,39 +304,11 @@ def scaled_fp4_grouped_quant(
         `4 * rk` is a padded `k // 16` to nearest multiple of 4. These layout constants are
         required by the NVIDIA Blackwell MMA operations.
     """
-    device = input_tensor.device
+    from sglang.jit_kernel.nvfp4 import (
-    l, m, k = input_tensor.shape
+        scaled_fp4_grouped_quant as jit_scaled_fp4_grouped_quant,
-    sf_vec_size = 16
-    assert k % sf_vec_size == 0, f"k must be multiple of 16, but got {k}."
-
-    scale_k = k // sf_vec_size
-    padded_k = (scale_k + (4 - 1)) // 4 * 4
-    padded_k_int32 = padded_k // 4
-    padded_m = (m + (128 - 1)) // 128 * 128
-    output = torch.empty(l, m, k // 2, device=device, dtype=torch.uint8)
-    output_scales = torch.empty(
-        l, padded_m, padded_k_int32, device=device, dtype=torch.int32
     )
 
-    torch.ops.sgl_kernel.silu_and_mul_scaled_fp4_experts_quant.default(
+    return jit_scaled_fp4_grouped_quant(input_tensor, input_global_scale, mask)
-        output.view(l * m, k // 2),
-        output_scales.view(l * padded_m, padded_k_int32),
-        input_tensor.view(l * m, k),
-        input_global_scale,
-        mask,
-        use_silu_and_mul=False,
-    )
-    # The physical layout of the output is (l, m, k // 2), but we want to return a
-    # logical layout (m, k // 2, l) required by the flashinfer masked group gemm.
-    output = output.permute(1, 2, 0)
-    # The physical layout of the output scales is already swizzled as (l, rm, rk, 32, 4, 4), a
-    # requirement for the flashinfer masked group gemm, where rm=m/128 and rk=k/4. The logic
-    # layout is (32, 4, rm, 4, rk, l).
-    output_scales = output_scales.view(torch.float8_e4m3fn).view(
-        l, padded_m // 128, padded_k // 4, 32, 4, 4
-    )
-    output_scales = output_scales.permute(3, 4, 1, 5, 2, 0)
-    return output, output_scales
 
 
 def silu_and_mul_scaled_fp4_grouped_quant(
@@ -392,116 +336,15 @@ def silu_and_mul_scaled_fp4_grouped_quant(
         `4 * rk` is a padded `k // 16` to nearest multiple of 4. These layout constants are
         required by the NVIDIA Blackwell MMA operations.
     """
-    device = input_tensor.device
+    from sglang.jit_kernel.nvfp4 import (
-    l, m, k_by_2 = input_tensor.shape
+        silu_and_mul_scaled_fp4_grouped_quant as jit_silu_and_mul_scaled_fp4_grouped_quant,
-    k = k_by_2 // 2
-    sf_vec_size = 16
-    assert k % sf_vec_size == 0, f"k must be multiple of 16, but got {k}."
-
-    scale_k = k // sf_vec_size
-    padded_k = (scale_k + (4 - 1)) // 4 * 4
-    padded_k_int32 = padded_k // 4
-    padded_m = (m + (128 - 1)) // 128 * 128
-    output = torch.empty(l, m, k // 2, device=device, dtype=torch.uint8)
-    output_scales = torch.empty(
-        l, padded_m, padded_k_int32, device=device, dtype=torch.int32
     )
 
-    torch.ops.sgl_kernel.silu_and_mul_scaled_fp4_experts_quant.default(
+    return jit_silu_and_mul_scaled_fp4_grouped_quant(
-        output.view(l * m, k // 2),
-        output_scales.view(l * padded_m, padded_k_int32),
-        input_tensor.view(l * m, k_by_2),
-        input_global_scale,
-        mask,
-        use_silu_and_mul=True,
-    )
-    # The physical layout of the output is (l, m, k // 2), but we want to return a
-    # logical layout (m, k // 2, l) required by the flashinfer masked group gemm.
-    output = output.permute(1, 2, 0)
-    # The physical layout of the output scales is already swizzled as (l, rm, rk, 32, 4, 4), a
-    # requirement for the flashinfer masked group gemm, where rm=m/128 and rk=k/4. The logic
-    # layout is (32, 4, rm, 4, rk, l).
-    output_scales = output_scales.view(torch.float8_e4m3fn).view(
-        l, padded_m // 128, padded_k // 4, 32, 4, 4
-    )
-    output_scales = output_scales.permute(3, 4, 1, 5, 2, 0)
-    return output, output_scales
-
-
-def scaled_fp4_experts_quant(
-    input_tensor: torch.Tensor,
-    input_global_scale: torch.Tensor,
-    expert_offsets: torch.Tensor,
-    blockscale_offsets: torch.Tensor,
-    topk: int,
-    expert_map: Optional[torch.Tensor] = None,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Quantize input tensor to FP4 and return quantized tensor and scale, for
-    packed MoE Inputs.
-    Args:
-        input: The input tensor to be quantized to FP4
-        expert_map: The expert map tensor
-        input_global_scale: A scalar scaling factor for the entire tensor.
-        expert_offsets: The expert offsets tensor
-        blockscale_offsets: The blockscale offsets tensor
-    Outputs:
-        output: The quantized tensor in FP4
-        output_scales: The blockscale tensor in FP8-E4M3
-    """
-    assert (
-        input_tensor.ndim == 2
-    ), f"input.ndim needs to be == 2, but got {input_tensor.ndim}."
-    if expert_map is not None:
-        m, k = input_tensor.shape
-        output_tensor_shape = (m * topk, k)
-        input_tensor = shuffle_rows(input_tensor, expert_map, output_tensor_shape)
-    m_numtopk, k = input_tensor.shape
-    # Control the maximum number of tokens per expert supported by the
-    # NVFP4 MoE Expert Quantization. This is used to prevent the kernel
-    # from running out of memory. This value can also be increased to support
-    # larger models.
-    import os
-
-    MAX_TOKENS_PER_EXPERT = int(os.environ.get("MODELOPT_MAX_TOKENS_PER_EXPERT", 65536))
-    assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk, (
-        f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT("
-        f"{MAX_TOKENS_PER_EXPERT})"
-        f" for cutlass_moe_fp4, observed m_numtopk = {m_numtopk}. Use"
-        f" MODELOPT_MAX_TOKENS_PER_EXPERT to set this value."
-    )
-    scales_k = k // 16
-    padded_k = (scales_k + (4 - 1)) // 4
-
-    # output is uint8 and packed fp4 values
-    output = torch.empty(
-        m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
-    )
-    # padded part should be zeroed out
-    if padded_k > scales_k:
-        output_scales = torch.zeros(
-            MAX_TOKENS_PER_EXPERT * topk,
-            padded_k,
-            dtype=torch.int32,
-            device=input_tensor.device,
-        )
-    else:
-        output_scales = torch.empty(
-            MAX_TOKENS_PER_EXPERT * topk,
-            padded_k,
-            dtype=torch.int32,
-            device=input_tensor.device,
-        )
-    torch.ops.sgl_kernel.scaled_fp4_experts_quant.default(
-        output,
-        output_scales,
         input_tensor,
         input_global_scale,
-        expert_offsets,
+        mask,
-        blockscale_offsets,
     )
-    output_scales = output_scales.view(torch.float8_e4m3fn)
-    return output, output_scales
 
 
 # GPTQ kernels

sgl-kernel/python/sgl_kernel/moe.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Optional
+from typing import Optional
 
 import torch
 
@@ -287,50 +287,3 @@ def fused_qk_norm_rope(
         attention_factor,
         rotary_dim if rotary_dim is not None else head_dim,
     )
-
-
-def cutlass_fp4_group_mm(
-    a_fp4,
-    b_fp4,
-    a_blockscale,
-    b_blockscale,
-    alphas,
-    out_dtype,
-    device,
-    params: Dict[str, Any],
-):
-    """
-    An FP4 Blockscaled Group Gemm that takes in  a_tensors, b_tensors and runs
-    the gemms for each combination based on the specified problem sizes.
-
-    This is used as the MoE gemm during NVFP4 Quantized FusedMoE forward.
-    - a/b_tensors: the NVFP4 a_ptrs and b_ptrs tensors which are quantized
-                     input and expert weights.
-    - a_/b_scales: The blockscales in FP8-E4M3 precision
-    - ab_strides/c_strides: Strides for the a/b tensors between rows.
-    - expert_offsets/sf_offsets: Indices that mark at which token index
-                    each expert begins its computation. The number of tokens
-                    computed with expert E is expert_offsets[E + 1] -
-                    expert_offsets[E] And the sf_size per expert is
-                    sf_offset[E+1] - sf_offset[E]
-    - problem_sizes: MxNxK sizes of each expert's multiplication in two grouped
-                     MMs used in the fused MoE operation.
-    """
-    m_topk = a_fp4.shape[0]
-    n = b_fp4.shape[1]
-    c_shape = (m_topk, n)
-    c = torch.empty(c_shape, device=device, dtype=out_dtype)
-    torch.ops.sgl_kernel.cutlass_fp4_group_mm.default(
-        c,
-        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 c.to(dtype=out_dtype)