diff --git a/python/sglang/jit_kernel/benchmark/bench_nvfp4_blockwise_moe.py b/python/sglang/jit_kernel/benchmark/bench_nvfp4_blockwise_moe.py new file mode 100644 index 000000000..54d37ddeb --- /dev/null +++ b/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) diff --git a/python/sglang/jit_kernel/benchmark/bench_nvfp4_quant.py b/python/sglang/jit_kernel/benchmark/bench_nvfp4_quant.py new file mode 100644 index 000000000..b1a7a24e6 --- /dev/null +++ b/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) diff --git a/python/sglang/jit_kernel/benchmark/bench_nvfp4_scaled_mm.py b/python/sglang/jit_kernel/benchmark/bench_nvfp4_scaled_mm.py new file mode 100644 index 000000000..8c0bdc322 --- /dev/null +++ b/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) diff --git a/sgl-kernel/csrc/gemm/nvfp4_expert_quant.cu b/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_expert_quant.cuh similarity index 77% rename from sgl-kernel/csrc/gemm/nvfp4_expert_quant.cu rename to python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_expert_quant.cuh index 5c9eeae80..f76936782 100644 --- a/sgl-kernel/csrc/gemm/nvfp4_expert_quant.cu +++ b/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_expert_quant.cuh @@ -1,15 +1,19 @@ -#include -#include -#include -#include -#include +#include +#include + +#include +#include +#include #include "nvfp4_quant.cuh" -#include "utils.h" +#include +#include + +using namespace host; // Quantizes the provided PackedVec into the uint32_t output template -__device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec& vec, float SFScaleVal, uint8_t* SFout) { +SGL_DEVICE uint32_t cvt_warp_fp16_to_fp4(PackedVec& 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& vec, float SFScaleVal, #pragma unroll for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { - if constexpr (std::is_same_v) { - fp2Vals[i] = __half22float2(vec.elts[i]); - } else { - fp2Vals[i] = __bfloat1622float2(vec.elts[i]); - } + fp2Vals[i] = device::cast(vec.elts[i]); fp2Vals[i].x *= outputScale; fp2Vals[i].y *= outputScale; } @@ -81,30 +81,22 @@ __device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec& 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 -inline __device__ void silu_and_mul(PackedVec& x_vec, const PackedVec& y_vec) { +SGL_DEVICE void silu_and_mul(PackedVec& x_vec, const PackedVec& 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) { - x[i] = __half22float2(x_vec.elts[i]); - y[i] = __half22float2(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] = __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]); - } + x[i] = device::cast(x_vec.elts[i]); + y[i] = device::cast(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] = device::cast>(x[i]); } } @@ -541,77 +533,69 @@ void quant_impl( } } -// Avoid redefinition warnings -#undef CHECK_CONTIGUOUS -#undef CHECK_TH_CUDA -#undef CHECK_INPUT - -/*Quantization entry for fp4 experts quantization*/ -#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; +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; +} 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) { - auto sm_version = getSMVersion(); - TORCH_CHECK(sm_version >= 100, "fp4_quant is only supported on sm100+"); + 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) { + auto MTopK = SymbolicSize{"m_topk"}; + 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"); - CHECK_INPUT(output_scale, "output_scale must be a CUDA tensor"); - CHECK_INPUT(input, "input must be a CUDA tensor"); - CHECK_INPUT(input_global_scale, "input_global_scale must be a CUDA tensor"); - CHECK_INPUT(input_offset_by_experts, "input_offset_by_experts must be a CUDA tensor"); - CHECK_INPUT(output_scale_offset_by_experts, "output_scale_offset_by_experts must be a CUDA tensor"); + TensorMatcher({MTopK, K}) // + .with_dtype() + .template with_device(device) + .verify(input); + TensorMatcher({MTopK, OutputCols}) // + .with_dtype() + .with_device(device) + .verify(output); + TensorMatcher({OutputScaleRows, OutputScaleCols}) // + .with_dtype() + .with_device(device) + .verify(output_scale); + TensorMatcher({NExperts}) // + .with_dtype() + .with_device(device) + .verify(input_global_scale); + TensorMatcher({OffsetSize}) // + .with_dtype() + .with_device(device) + .verify(input_offset_by_experts) + .verify(output_scale_offset_by_experts); - TORCH_CHECK(output.dim() == 2); - TORCH_CHECK(output_scale.dim() == 2); - 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 device_id = input.device().device_id; + RuntimeCheck(getSMVersion(device_id) >= 100, "fp4_quant is only supported on sm100+"); const int BLOCK_SIZE = 16; - auto m_topk = input.size(0); - auto k = input.size(1); - TORCH_CHECK(k % BLOCK_SIZE == 0, "k must be a multiple of 16"); - auto n_experts = input_global_scale.size(0); - TORCH_CHECK(input_offset_by_experts.size(0) == n_experts + 1); - TORCH_CHECK(output_scale_offset_by_experts.size(0) == n_experts + 1); - TORCH_CHECK(output.size(0) == m_topk); - TORCH_CHECK(output.size(1) == k / 2); - int scales_k = k / BLOCK_SIZE; - // 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); + const auto m_topk = static_cast(MTopK.unwrap()); + const auto k = static_cast(K.unwrap()); + RuntimeCheck(k % BLOCK_SIZE == 0, "k must be a multiple of 16"); + const auto n_experts = static_cast(NExperts.unwrap()); + const auto offset_size = static_cast(OffsetSize.unwrap()); + RuntimeCheck(offset_size == n_experts + 1, "input/output offset size mismatch"); + RuntimeCheck(static_cast(OutputCols.unwrap()) == k / 2, "output second dim mismatch"); + const int scales_k = k / BLOCK_SIZE; + const int padded_k = (scales_k + 3) / 4 * 4; + RuntimeCheck(static_cast(OutputScaleCols.unwrap()) * 4 == padded_k, "output_scale second dim mismatch"); - auto in_dtype = input.dtype(); - at::cuda::CUDAGuard device_guard{(char)input.get_device()}; - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(input.get_device()); - if (in_dtype == at::ScalarType::Half) { + const cudaStream_t stream = LaunchKernel::resolve_device(input.device()); + if (host::is_type(input.dtype())) { quant_impl( 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, - torch::Tensor& output_scale, - torch::Tensor const& input, - torch::Tensor const& input_global_scale, - torch::Tensor const& mask, + 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) { - auto sm_version = getSMVersion(); - TORCH_CHECK(sm_version >= 100, "fp4_quant is only supported on sm100+"); + auto MTopK = SymbolicSize{"m_topk"}; + 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"); - CHECK_INPUT(output_scale, "output_scale must be a CUDA tensor"); - CHECK_INPUT(input, "input must be a CUDA tensor"); - CHECK_INPUT(input_global_scale, "input_global_scale must be a CUDA tensor"); - CHECK_INPUT(mask, "mask must be a CUDA tensor"); + TensorMatcher({MTopK, KBy2}) // + .with_dtype() + .template with_device(device) + .verify(input); + TensorMatcher({MTopK, OutputCols}) // + .with_dtype() + .with_device(device) + .verify(output); + TensorMatcher({OutputScaleRows, OutputScaleCols}) // + .with_dtype() + .with_device(device) + .verify(output_scale); + TensorMatcher({NExperts}) // + .with_dtype() + .with_device(device) + .verify(input_global_scale); + TensorMatcher({NExperts}) // + .with_dtype() + .with_device(device) + .verify(mask); - TORCH_CHECK(output.dim() == 2); - TORCH_CHECK(output_scale.dim() == 2); - 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 device_id = input.device().device_id; + RuntimeCheck(getSMVersion(device_id) >= 100, "fp4_quant is only supported on sm100+"); const int BLOCK_SIZE = 16; - auto m_topk = input.size(0); - auto k_by_2 = input.size(1); - auto k = k_by_2; + const auto m_topk = static_cast(MTopK.unwrap()); + const auto k_by_2 = static_cast(KBy2.unwrap()); + 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); - TORCH_CHECK(mask.size(0) == n_experts); - TORCH_CHECK(output.size(0) == m_topk); - TORCH_CHECK(output.size(1) == k / 2); - int scales_k = k / BLOCK_SIZE; - // 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); + const auto n_experts = static_cast(NExperts.unwrap()); + RuntimeCheck(static_cast(OutputCols.unwrap()) == k / 2, "output second dim mismatch"); + const int scales_k = k / BLOCK_SIZE; + const int padded_k = (scales_k + 3) / 4 * 4; + RuntimeCheck(static_cast(OutputScaleCols.unwrap()) * 4 == padded_k, "output_scale second dim mismatch"); - auto in_dtype = input.dtype(); - at::cuda::CUDAGuard device_guard{(char)input.get_device()}; - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(input.get_device()); - if (in_dtype == at::ScalarType::Half) { + const cudaStream_t stream = LaunchKernel::resolve_device(input.device()); + if (host::is_type(input.dtype())) { quant_impl( 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"); } } diff --git a/sgl-kernel/csrc/gemm/nvfp4_quant.cuh b/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant.cuh similarity index 86% rename from sgl-kernel/csrc/gemm/nvfp4_quant.cuh rename to python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant.cuh index 82f22f479..e2d696673 100644 --- a/sgl-kernel/csrc/gemm/nvfp4_quant.cuh +++ b/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 -#include +#include +#include + #include -// Get type2 from type or vice versa (applied to half and bfloat16) -template -struct TypeConverter { - using Type = half2; -}; // keep for generality - -template <> -struct TypeConverter { - using Type = half; -}; - -template <> -struct TypeConverter { - using Type = half2; -}; - -template <> -struct TypeConverter<__nv_bfloat162> { - using Type = __nv_bfloat16; -}; - -template <> -struct TypeConverter<__nv_bfloat16> { - using Type = __nv_bfloat162; -}; +#include +#include #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 -__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 struct PackedVec { - typename TypeConverter::Type elts[4]; + packed_t elts[4]; }; template <> diff --git a/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant_entry.cuh b/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant_entry.cuh new file mode 100644 index 000000000..29b06dfc0 --- /dev/null +++ b/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 +#include + +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); +} diff --git a/sgl-kernel/csrc/gemm/nvfp4_quant_kernels.cu b/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant_kernels.cuh similarity index 69% rename from sgl-kernel/csrc/gemm/nvfp4_quant_kernels.cu rename to python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant_kernels.cuh index bb99d76cc..bac38d83e 100644 --- a/sgl-kernel/csrc/gemm/nvfp4_quant_kernels.cu +++ b/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 -#include -#include -#include -#include +#include +#include + +#include +#include #include "nvfp4_quant.cuh" -#include "utils.h" +#include +#include + +using namespace host; // Quantizes the provided PackedVec into the uint32_t output template -__device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec& vec, float SFScaleVal, uint8_t* SFout) { +SGL_DEVICE uint32_t cvt_warp_fp16_to_fp4(PackedVec& 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() { - static int multi_processor_count = []() { - int device_id = 0; - int count = 0; - - // Get the current CUDA device ID - 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 +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; } void scaled_fp4_quant_sm100a_sm120a( - torch::Tensor const& output, - torch::Tensor const& input, - torch::Tensor const& output_sf, - torch::Tensor const& input_sf) { - auto sm_version = getSMVersion(); - TORCH_CHECK(sm_version >= 100, "fp4_quant is only supported on sm100+"); + tvm::ffi::TensorView output, + tvm::ffi::TensorView input, + tvm::ffi::TensorView output_sf, + tvm::ffi::TensorView input_sf) { + RuntimeCheck(input.device().device_type == kDLCUDA, "input must be a CUDA tensor"); + 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(output.dtype()), "output must be uint8"); + RuntimeCheck(host::is_type(output_sf.dtype()), "output_sf must be int32"); + RuntimeCheck(host::is_type(input_sf.dtype()), "input_sf must be float32"); + RuntimeCheck( + host::is_type(input.dtype()) || host::is_type(input.dtype()), "input dtype must be fp16 or bf16"); - int32_t m = input.size(0); - int32_t n = input.size(1); + const int device_id = input.device().device_id; + 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(input.size(0)); + const int32_t n = static_cast(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(runtime::get_sm_count(device_id)); auto input_sf_ptr = static_cast(input_sf.data_ptr()); auto sf_out = static_cast(output_sf.data_ptr()); auto output_ptr = static_cast(output.data_ptr()); - at::cuda::CUDAGuard device_guard{(char)input.get_device()}; - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(input.get_device()); + const cudaStream_t stream = LaunchKernel::resolve_device(input.device()); - // We don't support e8m0 scales at this moment. - bool useUE8M0 = false; - - switch (input.scalar_type()) { - case torch::kHalf: { - auto input_ptr = reinterpret_cast(input.data_ptr()); - 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."); - } + constexpr bool useUE8M0 = false; + if (host::is_type(input.dtype())) { + auto input_ptr = reinterpret_cast(input.data_ptr()); + invokeFP4Quantization(m, n, input_ptr, input_sf_ptr, output_ptr, sf_out, useUE8M0, multiProcessorCount, stream); + } else { + 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); } } diff --git a/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_entry.cuh b/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_entry.cuh new file mode 100644 index 000000000..72d68f7d5 --- /dev/null +++ b/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 + +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); +} diff --git a/sgl-kernel/csrc/gemm/nvfp4_scaled_mm_kernels.cu b/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_kernels.cuh similarity index 69% rename from sgl-kernel/csrc/gemm/nvfp4_scaled_mm_kernels.cu rename to python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_kernels.cuh index 28e86f7c8..9cc309f14 100644 --- a/sgl-kernel/csrc/gemm/nvfp4_scaled_mm_kernels.cu +++ b/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 -#include -#include +#include +#include -#include "utils.h" +#include +#include + +#include +#include +#include +#include + +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) \ - { \ - cutlass::Status error = status; \ - TORCH_CHECK(error == cutlass::Status::kSuccess, cutlassGetStatusString(error)); \ +#define CUTLASS_CHECK(status) \ + { \ + cutlass::Status error = status; \ + 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{}(key.device_id); + size_t h2 = std::hash{}(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 cache; + WorkspaceKey key{device_id, reinterpret_cast(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::Gemm::Arguments args_from_options( - at::Tensor& D, - at::Tensor const& A, - at::Tensor const& B, - at::Tensor const& A_sf, - at::Tensor const& B_sf, - at::Tensor const& alpha, + 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, int64_t M, int64_t N, int64_t K) { @@ -335,12 +385,12 @@ typename T::Gemm::Arguments args_from_options( template void runGemm( - at::Tensor& D, - at::Tensor const& A, - at::Tensor const& B, - at::Tensor const& A_sf, - at::Tensor const& B_sf, - at::Tensor const& alpha, + 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, int64_t m, int64_t n, int64_t k, @@ -349,25 +399,25 @@ void runGemm( auto arguments = args_from_options(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()); - auto workspace = torch::empty(workspace_size, workspace_options); + int device_id = A.device().device_id; + 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::Arguments args_from_options_sm120( - at::Tensor& D, - at::Tensor const& A, - at::Tensor const& B, - at::Tensor const& A_sf, - at::Tensor const& B_sf, - torch::Tensor const& alpha, + 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, int M, int N, int K) { @@ -413,12 +463,12 @@ typename Gemm::Arguments args_from_options_sm120( // SM120 specific runGemm function template void runGemmSm120( - at::Tensor& D, - at::Tensor const& A, - at::Tensor const& B, - at::Tensor const& A_sf, - at::Tensor const& B_sf, - torch::Tensor const& alpha, + 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, int M, int N, int K, @@ -428,25 +478,25 @@ void runGemmSm120( auto arguments = args_from_options_sm120(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()); - auto workspace = torch::empty(workspace_size, workspace_options); + int device_id = A.device().device_id; + 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 void cutlassFp4GemmDispatch( - 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, + 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, 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( - 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, + 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, int64_t m, int64_t n, int64_t k, @@ -481,12 +531,12 @@ void cutlassFp4GemmDispatch( // 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, + 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, 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, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& A_sf, - torch::Tensor const& B_sf, - torch::Tensor const& alpha, + 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, int m, int n, int k, @@ -525,17 +575,17 @@ void cutlass_fp4_f16_gemm_dispatch_sm120( #else template void cutlassFp4GemmDispatch( - at::Tensor& D, - at::Tensor const& A, - at::Tensor const& B, - at::Tensor const& A_sf, - at::Tensor const& B_sf, - at::Tensor const& alpha, + 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, 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 -#undef CHECK_CONTIGUOUS -#undef CHECK_INPUT - -#define CHECK_TYPE(x, st, m) TORCH_CHECK(x.scalar_type() == st, "Inconsistency of Tensor type:", m) -#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, 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; +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; +} 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) { - CHECK_INPUT(A, FLOAT4_E2M1X2, "a"); - CHECK_INPUT(B, FLOAT4_E2M1X2, "b"); + 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) { + RuntimeCheck(A.device().device_type == kDLCUDA, "a must be a CUDA tensor"); + 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"); - CHECK_INPUT(B_sf, SF_DTYPE, "scale_b"); + RuntimeCheck(A.device() == B.device(), "a and b must be on same device"); + 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"); - TORCH_CHECK(B.dim() == 2, "b must be a matrix"); - TORCH_CHECK( + RuntimeCheck(host::is_type(A.dtype()), "a must be uint8"); + RuntimeCheck(host::is_type(B.dtype()), "b must be uint8"); + RuntimeCheck(host::is_type(A_sf.dtype()), "scale_a must be float8_e4m3fn"); + RuntimeCheck(host::is_type(B_sf.dtype()), "scale_b must be float8_e4m3fn"); + RuntimeCheck(host::is_type(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); - auto const n = B.size(0); - auto const k = A.size(1) * 2; + const auto m = static_cast(A.size(0)); + const auto n = static_cast(B.size(0)); + const auto k = static_cast(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( - k % alignment == 0, - "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), - ")."); + RuntimeCheck(k % alignment == 0, "Expected k to be divisible by ", alignment, ", but got k: ", k); + RuntimeCheck(n % alignment == 0, "Expected n to be divisible by ", alignment, ", but got n: ", n); - auto round_up = [](int x, int y) { return (x + y - 1) / y * y; }; - int rounded_m = round_up(m, 128); - int rounded_n = round_up(n, 128); - // Since k is divisible by 32 (alignment), k / 16 is guaranteed to be an - // integer. - int rounded_k = round_up(k / 16, 4); + auto round_up = [](int64_t x, int64_t y) { return (x + y - 1) / y * y; }; + const int64_t rounded_m = round_up(m, 128); + const int64_t rounded_n = round_up(n, 128); + const int64_t rounded_k = round_up(k / 16, 4); - TORCH_CHECK(A_sf.dim() == 2, "scale_a must be a matrix"); - TORCH_CHECK(B_sf.dim() == 2, "scale_b must be a matrix"); - TORCH_CHECK( + RuntimeCheck( 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(); - at::cuda::CUDAGuard device_guard{(char)A.get_device()}; - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(A.get_device()); - - // Check SM version and dispatch accordingly - auto sm_version = getSMVersion(); + const cudaStream_t stream = LaunchKernel::resolve_device(A.device()); + const int sm_version = getSMVersion(A.device().device_id); if (sm_version >= 120) { - // Use SM120 specific dispatch - if (out_dtype == at::ScalarType::Half) { - cutlass_fp4_f16_gemm_dispatch_sm120(D, A, B, A_sf, B_sf, alpha, m, n, k, stream); - } else if (out_dtype == at::ScalarType::BFloat16) { - cutlass_fp4_bf16_gemm_dispatch_sm120(D, A, B, A_sf, B_sf, alpha, m, n, k, stream); + if (host::is_type(D.dtype())) { + cutlass_fp4_f16_gemm_dispatch_sm120( + D, A, B, A_sf, B_sf, alpha, static_cast(m), static_cast(n), static_cast(k), stream); + } else if (host::is_type(D.dtype())) { + cutlass_fp4_bf16_gemm_dispatch_sm120( + D, A, B, A_sf, B_sf, alpha, static_cast(m), static_cast(n), static_cast(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 (out_dtype == at::ScalarType::Half) { + if (host::is_type(D.dtype())) { cutlassFp4GemmDispatch(D, A, B, A_sf, B_sf, alpha, m, n, k, stream); - } else if (out_dtype == at::ScalarType::BFloat16) { + } else if (host::is_type(D.dtype())) { cutlassFp4GemmDispatch(D, A, B, A_sf, B_sf, alpha, m, n, k, stream); - } else if (out_dtype == at::ScalarType::Float) { + } else if (host::is_type(D.dtype())) { cutlassFp4GemmDispatch(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"); } } } diff --git a/sgl-kernel/csrc/moe/nvfp4_blockwise_moe.cu b/python/sglang/jit_kernel/csrc/moe/nvfp4_blockwise_moe.cuh similarity index 56% rename from sgl-kernel/csrc/moe/nvfp4_blockwise_moe.cu rename to python/sglang/jit_kernel/csrc/moe/nvfp4_blockwise_moe.cuh index ba9da1d77..c3293fbfd 100644 --- a/sgl-kernel/csrc/moe/nvfp4_blockwise_moe.cu +++ b/python/sglang/jit_kernel/csrc/moe/nvfp4_blockwise_moe.cuh @@ -1,10 +1,11 @@ -#include -#include -#include -#include -#include +#include +#include -#include +#include +#include + +#include +#include #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 +#include +#include +#include +#include +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{}(key.device_id); + size_t h2 = std::hash{}(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 cache; + WorkspaceKey key{device_id, reinterpret_cast(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) \ - else if (out_tensors.dtype() == TENSOR_C_TYPE) { \ + ELEMENT_AB_TYPE, SF_TYPE, TYPE_CHECK, C_TYPE, LayoutSFA, LayoutSFB, ScaleConfig) \ + else if (TYPE_CHECK) { \ __get_group_gemm_starts \ <<<1, num_experts, 0, stream>>>( \ static_cast(a_starts.data_ptr()), \ @@ -131,32 +188,32 @@ __global__ void __get_group_gemm_starts( template void run_get_group_gemm_starts( - const torch::Tensor& a_starts, - const torch::Tensor& b_starts, - const torch::Tensor& out_starts, - const torch::Tensor& a_scales_starts, - const torch::Tensor& b_scales_starts, - const torch::Tensor& alpha_starts, - const torch::Tensor& layout_sfa, - const torch::Tensor& layout_sfb, + const tvm::ffi::TensorView a_starts, + const tvm::ffi::TensorView b_starts, + const tvm::ffi::TensorView out_starts, + const tvm::ffi::TensorView a_scales_starts, + const tvm::ffi::TensorView b_scales_starts, + const tvm::ffi::TensorView alpha_starts, + const tvm::ffi::TensorView layout_sfa, + const tvm::ffi::TensorView layout_sfb, /*these are used for their base addresses*/ - torch::Tensor const& a_tensors, - torch::Tensor const& b_tensors, - torch::Tensor const& out_tensors, - torch::Tensor const& a_scales, - torch::Tensor const& b_scales, - torch::Tensor const& alphas, - torch::Tensor const& expert_offsets, - torch::Tensor const& sf_offsets, - torch::Tensor const& problem_sizes, + tvm::ffi::TensorView const& a_tensors, + tvm::ffi::TensorView const& b_tensors, + tvm::ffi::TensorView const& out_tensors, + tvm::ffi::TensorView const& a_scales, + tvm::ffi::TensorView const& b_scales, + tvm::ffi::TensorView const& alphas, + tvm::ffi::TensorView const& expert_offsets, + tvm::ffi::TensorView const& sf_offsets, + tvm::ffi::TensorView const& problem_sizes, int M, int N, int K) { - int num_experts = (int)expert_offsets.size(0); - auto stream = at::cuda::getCurrentCUDAStream(a_tensors.device().index()); + int num_experts = static_cast(expert_offsets.size(0)); + auto stream = LaunchKernel::resolve_device(a_tensors.device()); - TORCH_CHECK(out_tensors.size(1) == N, "Output tensor shape doesn't match expected shape"); - TORCH_CHECK( + RuntimeCheck(out_tensors.size(1) == N, "Output tensor shape doesn't match expected shape"); + 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(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(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, - 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, + tvm::ffi::TensorView output, + 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, 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(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( 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 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::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()); - auto workspace = torch::empty(workspace_size, workspace_options); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(a.get_device()); + const cudaStream_t stream = LaunchKernel::resolve_device(a.device()); + void* workspace = get_cached_workspace(workspace_size, hw_info.device_id, stream); 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()); - TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to initialize GEMM"); + auto status = gemm_op.initialize(args, workspace); + RuntimeCheck(status == cutlass::Status::kSuccess, "Failed to initialize GEMM: ", cutlassGetStatusString(status)); - status = gemm_op.run(args, workspace.data_ptr(), stream); - TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM"); + status = gemm_op.run(args, workspace, stream); + RuntimeCheck(status == cutlass::Status::kSuccess, "Failed to run GEMM: ", cutlassGetStatusString(status)); } template void run_fp4_blockwise_scaled_group_mm_sm100( - 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, + tvm::ffi::TensorView output, + 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, 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(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( 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 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::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()); - auto workspace = torch::empty(workspace_size, workspace_options); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(a.get_device()); + const cudaStream_t stream = LaunchKernel::resolve_device(a.device()); + void* workspace = get_cached_workspace(workspace_size, hw_info.device_id, stream); 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()); - TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to initialize GEMM"); + auto status = gemm_op.initialize(args, workspace); + RuntimeCheck(status == cutlass::Status::kSuccess, "Failed to initialize GEMM: ", cutlassGetStatusString(status)); - status = gemm_op.run(args, workspace.data_ptr(), stream); - TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM"); + status = gemm_op.run(args, workspace, stream); + RuntimeCheck(status == cutlass::Status::kSuccess, "Failed to run GEMM: ", cutlassGetStatusString(status)); } -// Undefine macros from utils.h to redefine with custom signatures -#undef CHECK_CONTIGUOUS -#undef CHECK_INPUT +void cutlass_fp4_group_mm_sm100a_sm120a( + tvm::ffi::TensorView output, + 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) -#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, st, m) \ - CHECK_TH_CUDA(x, m); \ - CHECK_CONTIGUOUS(x, m); \ - CHECK_TYPE(x, st, m) + check_cuda_contig(output, "output"); + check_cuda_contig(a, "a"); + check_cuda_contig(b, "b"); + check_cuda_contig(a_blockscale, "a_blockscale"); + check_cuda_contig(b_blockscales, "b_blockscales"); + check_cuda_contig(alphas, "alphas"); + 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( - 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) { -#if defined ENABLE_NVFP4 && ENABLE_NVFP4 + RuntimeCheck( + output.device() == a.device() && a.device() == b.device() && a.device() == a_blockscale.device() && + a.device() == b_blockscales.device() && a.device() == alphas.device() && a.device() == ab_strides.device() && + a.device() == c_strides.device() && a.device() == problem_sizes.device() && + a.device() == expert_offsets.device() && a.device() == sf_offsets.device() && a.device() == a_ptrs.device() && + a.device() == b_ptrs.device() && a.device() == out_ptrs.device() && a.device() == a_scales_ptrs.device() && + a.device() == b_scales_ptrs.device() && a.device() == alpha_ptrs.device() && + a.device() == layout_sfa.device() && a.device() == layout_sfb.device(), + "all tensors must be on the same device"); - constexpr auto FLOAT4_E2M1X2 = at::ScalarType::Byte; - constexpr auto SF_DTYPE = at::ScalarType::Float8_e4m3fn; - // Input validation - CHECK_INPUT(a, FLOAT4_E2M1X2, "a"); - CHECK_INPUT(b, FLOAT4_E2M1X2, "b"); - CHECK_INPUT(a_blockscale, SF_DTYPE, "a_blockscale"); - CHECK_INPUT(b_blockscales, SF_DTYPE, "b_blockscales"); - CHECK_INPUT(alphas, at::ScalarType::Float, "alphas"); + RuntimeCheck(host::is_type(a.dtype()), "a must be uint8"); + RuntimeCheck(host::is_type(b.dtype()), "b must be uint8"); + RuntimeCheck(host::is_type(a_blockscale.dtype()), "a_blockscale must be float8_e4m3fn"); + RuntimeCheck(host::is_type(b_blockscales.dtype()), "b_blockscales must be float8_e4m3fn"); + RuntimeCheck(host::is_type(alphas.dtype()), "alphas must be float32"); + RuntimeCheck(host::is_type(ab_strides.dtype()), "ab_strides must be int64"); + RuntimeCheck(host::is_type(c_strides.dtype()), "c_strides must be int64"); + RuntimeCheck(host::is_type(problem_sizes.dtype()), "problem_sizes must be int32"); + RuntimeCheck(host::is_type(expert_offsets.dtype()), "expert_offsets must be int32"); + RuntimeCheck(host::is_type(sf_offsets.dtype()), "sf_offsets must be int32"); + RuntimeCheck(host::is_type(a_ptrs.dtype()), "a_ptrs must be int64"); + RuntimeCheck(host::is_type(b_ptrs.dtype()), "b_ptrs must be int64"); + RuntimeCheck(host::is_type(out_ptrs.dtype()), "out_ptrs must be int64"); + RuntimeCheck(host::is_type(a_scales_ptrs.dtype()), "a_scales_ptrs must be int64"); + RuntimeCheck(host::is_type(b_scales_ptrs.dtype()), "b_scales_ptrs must be int64"); + RuntimeCheck(host::is_type(alpha_ptrs.dtype()), "alpha_ptrs must be int64"); + RuntimeCheck(host::is_type(layout_sfa.dtype()), "layout_sfa must be int64"); + RuntimeCheck(host::is_type(layout_sfb.dtype()), "layout_sfb must be int64"); + RuntimeCheck( + host::is_type(output.dtype()) || host::is_type(output.dtype()), + "output must be bfloat16 or float16"); - TORCH_CHECK( - a_blockscale.dim() == 2, - "expected a_blockscale to be of shape [num_experts, rounded_m," - " k // group_size], observed rank: ", - a_blockscale.dim()) - TORCH_CHECK( - b_blockscales.dim() == 3, - "expected b_blockscale to be of shape: " - " [num_experts, n, k // group_size], observed rank: ", - b_blockscales.dim()) - TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be a 2D tensor"); - TORCH_CHECK(problem_sizes.size(1) == 3, "problem_sizes must have the shape (num_experts, 3)"); - TORCH_CHECK( - problem_sizes.size(0) == expert_offsets.size(0), "Number of experts in problem_sizes must match expert_offsets"); - TORCH_CHECK(problem_sizes.dtype() == torch::kInt32, "problem_sizes must be int32."); + RuntimeCheck(a.dim() == 2, "a must be 2D"); + RuntimeCheck(b.dim() == 3, "b must be 3D"); + RuntimeCheck(a_blockscale.dim() == 2, "a_blockscale must be 2D"); + RuntimeCheck(b_blockscales.dim() == 3, "b_blockscales must be 3D"); + RuntimeCheck(alphas.dim() == 1, "alphas must be 1D"); + RuntimeCheck(ab_strides.dim() == 1, "ab_strides must be 1D"); + RuntimeCheck(c_strides.dim() == 1, "c_strides must be 1D"); + RuntimeCheck(problem_sizes.dim() == 2, "problem_sizes must be 2D"); + RuntimeCheck(expert_offsets.dim() == 1, "expert_offsets must be 1D"); + RuntimeCheck(sf_offsets.dim() == 1, "sf_offsets must be 1D"); + RuntimeCheck(a_ptrs.dim() == 1, "a_ptrs must be 1D"); + RuntimeCheck(b_ptrs.dim() == 1, "b_ptrs must be 1D"); + RuntimeCheck(out_ptrs.dim() == 1, "out_ptrs must be 1D"); + RuntimeCheck(a_scales_ptrs.dim() == 1, "a_scales_ptrs must be 1D"); + 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(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(a.size(0)); int N = static_cast(b.size(1)); - int E = static_cast(b.size(0)); int K = static_cast(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(output.dtype())) { run_fp4_blockwise_scaled_group_mm_sm100( 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(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, - "No compiled cutlass_fp4_group_mm kernel, sgl-kernel must " - "be compiled with ENABLE_NVFP4 for SM100+ and CUDA " - "12.8 or above."); -#endif +} + +void cutlass_fp4_group_mm( + tvm::ffi::TensorView output, + 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) { + 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); } diff --git a/python/sglang/jit_kernel/nvfp4.py b/python/sglang/jit_kernel/nvfp4.py new file mode 100644 index 000000000..dde345adf --- /dev/null +++ b/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 diff --git a/python/sglang/jit_kernel/tests/test_nvfp4_blockwise_moe.py b/python/sglang/jit_kernel/tests/test_nvfp4_blockwise_moe.py new file mode 100644 index 000000000..6eb809d33 --- /dev/null +++ b/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) diff --git a/python/sglang/jit_kernel/tests/test_nvfp4_gemm.py b/python/sglang/jit_kernel/tests/test_nvfp4_gemm.py new file mode 100644 index 000000000..9a76cd800 --- /dev/null +++ b/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) diff --git a/python/sglang/jit_kernel/tests/test_nvfp4_quant.py b/python/sglang/jit_kernel/tests/test_nvfp4_quant.py new file mode 100644 index 000000000..eb7816408 --- /dev/null +++ b/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]]) diff --git a/python/sglang/srt/layers/moe/cutlass_moe.py b/python/sglang/srt/layers/moe/cutlass_moe.py index 9fbc2764d..57569c73a 100755 --- a/python/sglang/srt/layers/moe/cutlass_moe.py +++ b/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, diff --git a/python/sglang/srt/layers/moe/cutlass_moe_params.py b/python/sglang/srt/layers/moe/cutlass_moe_params.py index f3de60e04..6a71d606e 100644 --- a/python/sglang/srt/layers/moe/cutlass_moe_params.py +++ b/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, - # "b_ptrs": self.b_ptrs, - # "out_ptrs": self.out_ptrs, - # "a_scales_ptrs": self.a_scales_ptrs, - # "b_scales_ptrs": self.b_scales_ptrs, + "a_ptrs": self.a_ptrs, + "b_ptrs": self.b_ptrs, + "out_ptrs": self.out_ptrs, + "a_scales_ptrs": self.a_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, - # "b_ptrs": self.b_ptrs, - # "out_ptrs": self.out_ptrs, - # "a_scales_ptrs": self.a_scales_ptrs, - # "b_scales_ptrs": self.b_scales_ptrs, + "a_ptrs": self.a_ptrs, + "b_ptrs": self.b_ptrs, + "out_ptrs": self.out_ptrs, + "a_scales_ptrs": self.a_scales_ptrs, + "b_scales_ptrs": self.b_scales_ptrs, + "alpha_ptrs": self.alpha_ptrs, + "layout_sfa": self.layout_sfa, + "layout_sfb": self.layout_sfb, } diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index 589f0c0ed..f87c39606 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/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 diff --git a/python/sglang/srt/layers/moe/token_dispatcher/standard.py b/python/sglang/srt/layers/moe/token_dispatcher/standard.py index 0a1270098..8ec839991 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/standard.py +++ b/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: diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index d23bd075c..e47305e45 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/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 diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt index fbd913ff4..ec9f10ef0 100644 --- a/sgl-kernel/CMakeLists.txt +++ b/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" diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc index fb3e00fa3..5a572e092 100644 --- a/sgl-kernel/csrc/common_extension.cc +++ b/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); diff --git a/sgl-kernel/csrc/gemm/nvfp4_quant_entry.cu b/sgl-kernel/csrc/gemm/nvfp4_quant_entry.cu deleted file mode 100644 index 67044d015..000000000 --- a/sgl-kernel/csrc/gemm/nvfp4_quant_entry.cu +++ /dev/null @@ -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 - -#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"); -} diff --git a/sgl-kernel/csrc/gemm/nvfp4_scaled_mm_entry.cu b/sgl-kernel/csrc/gemm/nvfp4_scaled_mm_entry.cu deleted file mode 100644 index 3a050bbc2..000000000 --- a/sgl-kernel/csrc/gemm/nvfp4_scaled_mm_entry.cu +++ /dev/null @@ -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 - -#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."); -} diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h index 43e9dece4..92493599f 100644 --- a/sgl-kernel/include/sgl_kernel_ops.h +++ b/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 */ diff --git a/sgl-kernel/python/sgl_kernel/__init__.py b/sgl-kernel/python/sgl_kernel/__init__.py index bfeb03f15..e84ff12f5 100644 --- a/sgl-kernel/python/sgl_kernel/__init__.py +++ b/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, diff --git a/sgl-kernel/python/sgl_kernel/gemm.py b/sgl-kernel/python/sgl_kernel/gemm.py index e9e11c834..0d2389c96 100644 --- a/sgl-kernel/python/sgl_kernel/gemm.py +++ b/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 - m, n = a.shape[0], b.shape[0] - 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 + from sglang.jit_kernel.nvfp4 import ( + cutlass_scaled_fp4_mm as jit_cutlass_scaled_fp4_mm, + ) + + 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}." - 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 + from sglang.jit_kernel.nvfp4 import scaled_fp4_quant as jit_scaled_fp4_quant - 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}." - - # 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 + return jit_scaled_fp4_quant(input, input_global_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 - 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 + from sglang.jit_kernel.nvfp4 import ( + scaled_fp4_grouped_quant as jit_scaled_fp4_grouped_quant, ) - torch.ops.sgl_kernel.silu_and_mul_scaled_fp4_experts_quant.default( - 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 + return jit_scaled_fp4_grouped_quant(input_tensor, input_global_scale, mask) 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 - 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 + from sglang.jit_kernel.nvfp4 import ( + silu_and_mul_scaled_fp4_grouped_quant as jit_silu_and_mul_scaled_fp4_grouped_quant, ) - torch.ops.sgl_kernel.silu_and_mul_scaled_fp4_experts_quant.default( - 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, + return jit_silu_and_mul_scaled_fp4_grouped_quant( input_tensor, input_global_scale, - expert_offsets, - blockscale_offsets, + mask, ) - output_scales = output_scales.view(torch.float8_e4m3fn) - return output, output_scales # GPTQ kernels diff --git a/sgl-kernel/python/sgl_kernel/moe.py b/sgl-kernel/python/sgl_kernel/moe.py index d85e4b602..bb8271749 100755 --- a/sgl-kernel/python/sgl_kernel/moe.py +++ b/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)