chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Mohammad Miadh Angkad
2026-03-05 15:22:28 +08:00
committed by GitHub
parent 1bbfed0539
commit 2bdd89a6cd
27 changed files with 2458 additions and 989 deletions
Download Patch File Download Diff File Expand all files Collapse all files

250
python/sglang/jit_kernel/benchmark/bench_nvfp4_blockwise_moe.py Normal file
View File

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

181
python/sglang/jit_kernel/benchmark/bench_nvfp4_quant.py Normal file
View File

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

175
python/sglang/jit_kernel/benchmark/bench_nvfp4_scaled_mm.py Normal file
View File

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

260
sgl-kernel/csrc/gemm/nvfp4_expert_quant.cu → python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_expert_quant.cuh
View File

@@ -1,15 +1,19 @@
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <cuda_runtime_api.h>
#include <torch/all.h>
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/runtime.cuh>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include "nvfp4_quant.cuh"
#include "utils.h"
#include <cuda_runtime.h>
#include <cuda_runtime_api.h>
using namespace host;
// Quantizes the provided PackedVec into the uint32_t output
template <class Type, bool UE8M0_SF = false>
__device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal, uint8_t* SFout) {
SGL_DEVICE uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal, uint8_t* SFout) {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
// Get absolute maximum values among the local 8 values.
auto localMax = __habs2(vec.elts[0]);
@@ -62,11 +66,7 @@ __device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal,
#pragma unroll
for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
if constexpr (std::is_same_v<Type, half>) {
fp2Vals[i] = __half22float2(vec.elts[i]);
} else {
fp2Vals[i] = __bfloat1622float2(vec.elts[i]);
}
fp2Vals[i] = device::cast<float2>(vec.elts[i]);
fp2Vals[i].x *= outputScale;
fp2Vals[i].y *= outputScale;
}
@@ -81,30 +81,22 @@ __device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal,
#endif
}
__device__ __forceinline__ float silu(const float& val) {
SGL_DEVICE float silu(const float& val) {
return val / (1.0f + __expf(-val));
}
template <class Type>
inline __device__ void silu_and_mul(PackedVec<Type>& x_vec, const PackedVec<Type>& y_vec) {
SGL_DEVICE void silu_and_mul(PackedVec<Type>& x_vec, const PackedVec<Type>& y_vec) {
float2 x[CVT_FP4_ELTS_PER_THREAD / 2];
float2 y[CVT_FP4_ELTS_PER_THREAD / 2];
#pragma unroll
for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
if constexpr (std::is_same_v<Type, half>) {
x[i] = __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<float2>(x_vec.elts[i]);
y[i] = device::cast<float2>(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<packed_t<Type>>(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<fp16_t, bf16_t>()
.template with_device<kDLCUDA>(device)
.verify(input);
TensorMatcher({MTopK, OutputCols}) //
.with_dtype<uint8_t>()
.with_device(device)
.verify(output);
TensorMatcher({OutputScaleRows, OutputScaleCols}) //
.with_dtype<int32_t>()
.with_device(device)
.verify(output_scale);
TensorMatcher({NExperts}) //
.with_dtype<float>()
.with_device(device)
.verify(input_global_scale);
TensorMatcher({OffsetSize}) //
.with_dtype<int32_t>()
.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<int>(MTopK.unwrap());
const auto k = static_cast<int>(K.unwrap());
RuntimeCheck(k % BLOCK_SIZE == 0, "k must be a multiple of 16");
const auto n_experts = static_cast<int>(NExperts.unwrap());
const auto offset_size = static_cast<int>(OffsetSize.unwrap());
RuntimeCheck(offset_size == n_experts + 1, "input/output offset size mismatch");
RuntimeCheck(static_cast<int>(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<int>(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<fp16_t>(input.dtype())) {
quant_impl<half>(
output.data_ptr(),
output_scale.data_ptr(),
@@ -625,7 +609,7 @@ void scaled_fp4_experts_quant_sm100a(
k,
n_experts,
stream);
} else if (in_dtype == at::ScalarType::BFloat16) {
} else {
quant_impl<__nv_bfloat16>(
output.data_ptr(),
output_scale.data_ptr(),
@@ -639,62 +623,64 @@ void scaled_fp4_experts_quant_sm100a(
k,
n_experts,
stream);
} else {
TORCH_CHECK(false, "Expected input data type to be half or bfloat16");
}
}
void silu_and_mul_scaled_fp4_experts_quant_sm100a(
torch::Tensor& output,
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<fp16_t, bf16_t>()
.template with_device<kDLCUDA>(device)
.verify(input);
TensorMatcher({MTopK, OutputCols}) //
.with_dtype<uint8_t>()
.with_device(device)
.verify(output);
TensorMatcher({OutputScaleRows, OutputScaleCols}) //
.with_dtype<int32_t>()
.with_device(device)
.verify(output_scale);
TensorMatcher({NExperts}) //
.with_dtype<float>()
.with_device(device)
.verify(input_global_scale);
TensorMatcher({NExperts}) //
.with_dtype<int32_t>()
.with_device(device)
.verify(mask);
TORCH_CHECK(output.dim() == 2);
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<int>(MTopK.unwrap());
const auto k_by_2 = static_cast<int>(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<int>(NExperts.unwrap());
RuntimeCheck(static_cast<int>(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<int>(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<fp16_t>(input.dtype())) {
quant_impl<half>(
output.data_ptr(),
output_scale.data_ptr(),
@@ -708,7 +694,7 @@ void silu_and_mul_scaled_fp4_experts_quant_sm100a(
k,
n_experts,
stream);
} else if (in_dtype == at::ScalarType::BFloat16) {
} else {
quant_impl<__nv_bfloat16>(
output.data_ptr(),
output_scale.data_ptr(),
@@ -722,7 +708,5 @@ void silu_and_mul_scaled_fp4_experts_quant_sm100a(
k,
n_experts,
stream);
} else {
TORCH_CHECK(false, "Expected input data type to be half or bfloat16");
}
}

42
sgl-kernel/csrc/gemm/nvfp4_quant.cuh → python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant.cuh
View File

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

68
python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant_entry.cuh Normal file
View File

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

103
sgl-kernel/csrc/gemm/nvfp4_quant_kernels.cu → python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant_kernels.cuh
View File

@@ -13,18 +13,21 @@ See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <cuda_runtime_api.h>
#include <torch/all.h>
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/runtime.cuh>
#include <sgl_kernel/utils.cuh>
#include "nvfp4_quant.cuh"
#include "utils.h"
#include <cuda_runtime.h>
#include <cuda_runtime_api.h>
using namespace host;
// Quantizes the provided PackedVec into the uint32_t output
template <class Type, bool UE8M0_SF = false>
__device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal, uint8_t* SFout) {
SGL_DEVICE uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal, uint8_t* SFout) {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
// Get absolute maximum values among the local 8 values.
auto localMax = __habs2(vec.elts[0]);
@@ -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<uint8_t>(output.dtype()), "output must be uint8");
RuntimeCheck(host::is_type<int32_t>(output_sf.dtype()), "output_sf must be int32");
RuntimeCheck(host::is_type<float>(input_sf.dtype()), "input_sf must be float32");
RuntimeCheck(
host::is_type<fp16_t>(input.dtype()) || host::is_type<bf16_t>(input.dtype()), "input dtype must be fp16 or bf16");
int32_t m = input.size(0);
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<int32_t>(input.size(0));
const int32_t n = static_cast<int32_t>(input.size(1));
int multiProcessorCount = getMultiProcessorCount();
RuntimeCheck(output.size(0) == m, "output row size mismatch");
RuntimeCheck(output.size(1) == n / 2, "output column size mismatch");
RuntimeCheck(n % 16 == 0, "The N dimension must be multiple of 16.");
const int multiProcessorCount = static_cast<int>(runtime::get_sm_count(device_id));
auto input_sf_ptr = static_cast<float const*>(input_sf.data_ptr());
auto sf_out = static_cast<int32_t*>(output_sf.data_ptr());
auto output_ptr = static_cast<int64_t*>(output.data_ptr());
at::cuda::CUDAGuard device_guard{(char)input.get_device()};
const cudaStream_t stream = 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<half const*>(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<fp16_t>(input.dtype())) {
auto input_ptr = reinterpret_cast<half const*>(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);
}
}

34
python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_entry.cuh Normal file
View File

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

351
sgl-kernel/csrc/gemm/nvfp4_scaled_mm_kernels.cu → python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_kernels.cuh
View File

@@ -13,11 +13,18 @@ See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/all.h>
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include "utils.h"
#include <sgl_kernel/runtime.cuh>
#include <sgl_kernel/utils.cuh>
#include <cstddef>
#include <cstdint>
#include <cuda_runtime.h>
#include <unordered_map>
using namespace host;
// clang-format off
#include "cutlass/cutlass.h"
@@ -31,10 +38,10 @@ limitations under the License.
/**
* Helper function for checking CUTLASS errors
*/
#define CUTLASS_CHECK(status) \
{ \
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<int>{}(key.device_id);
size_t h2 = std::hash<uintptr_t>{}(key.stream);
return h1 ^ (h2 + 0x9e3779b97f4a7c15ULL + (h1 << 6) + (h1 >> 2));
}
};
struct WorkspaceState {
void* ptr = nullptr;
size_t bytes = 0;
};
inline auto get_cached_workspace(size_t required_bytes, int device_id, cudaStream_t stream) -> void* {
if (required_bytes == 0) {
return nullptr;
}
thread_local std::unordered_map<WorkspaceKey, WorkspaceState, WorkspaceKeyHash> cache;
WorkspaceKey key{device_id, reinterpret_cast<uintptr_t>(stream)};
auto& ws = cache[key];
if (ws.ptr != nullptr && ws.bytes >= required_bytes) {
return ws.ptr;
}
RuntimeDeviceCheck(cudaSetDevice(device_id));
if (ws.ptr != nullptr) {
RuntimeDeviceCheck(cudaFreeAsync(ws.ptr, stream));
ws.ptr = nullptr;
ws.bytes = 0;
}
RuntimeDeviceCheck(cudaMallocAsync(&ws.ptr, required_bytes, stream));
ws.bytes = required_bytes;
return ws.ptr;
}
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || \
defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED)
// Config(half_t/bfloat16_t) for M <= 128
@@ -277,12 +327,12 @@ struct Fp4GemmSm120 {
template <typename T>
typename T::Gemm::Arguments args_from_options(
at::Tensor& D,
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 <typename T>
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<T>(D, A, B, A_sf, B_sf, alpha, m, n, k);
size_t workspace_size = T::Gemm::get_workspace_size(arguments);
auto const workspace_options = torch::TensorOptions().dtype(torch::kUInt8).device(A.device());
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>
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 <typename Gemm>
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<Gemm>(D, A, B, A_sf, B_sf, alpha, M, N, K);
size_t workspace_size = Gemm::get_workspace_size(arguments);
auto const workspace_options = torch::TensorOptions().dtype(torch::kUInt8).device(A.device());
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 <typename OutType>
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<float>(
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<float>(
// 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 <typename T>
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<uint8_t>(A.dtype()), "a must be uint8");
RuntimeCheck(host::is_type<uint8_t>(B.dtype()), "b must be uint8");
RuntimeCheck(host::is_type<fp8_e4m3_t>(A_sf.dtype()), "scale_a must be float8_e4m3fn");
RuntimeCheck(host::is_type<fp8_e4m3_t>(B_sf.dtype()), "scale_b must be float8_e4m3fn");
RuntimeCheck(host::is_type<float>(alpha.dtype()), "alpha must be float32");
RuntimeCheck(A.dim() == 2, "a must be a matrix");
RuntimeCheck(B.dim() == 2, "b must be a matrix");
RuntimeCheck(A_sf.dim() == 2, "scale_a must be a matrix");
RuntimeCheck(B_sf.dim() == 2, "scale_b must be a matrix");
RuntimeCheck(alpha.numel() == 1, "alpha must have exactly one element");
RuntimeCheck(
A.size(1) == B.size(1),
"a and b shapes cannot be multiplied (",
A.size(0),
@@ -587,42 +652,24 @@ void cutlass_scaled_fp4_mm_sm100a_sm120a(
B.size(1),
")");
auto const m = A.size(0);
auto const n = B.size(0);
auto const k = A.size(1) * 2;
const auto m = static_cast<int64_t>(A.size(0));
const auto n = static_cast<int64_t>(B.size(0));
const auto k = static_cast<int64_t>(A.size(1) * 2);
RuntimeCheck(D.dim() == 2, "out must be 2D");
RuntimeCheck(D.size(0) == m, "out first dim must equal m");
RuntimeCheck(D.size(1) == n, "out second dim must equal n");
constexpr int alignment = 32;
TORCH_CHECK(
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<fp16_t>(D.dtype())) {
cutlass_fp4_f16_gemm_dispatch_sm120(
D, A, B, A_sf, B_sf, alpha, static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), stream);
} else if (host::is_type<bf16_t>(D.dtype())) {
cutlass_fp4_bf16_gemm_dispatch_sm120(
D, A, B, A_sf, B_sf, alpha, static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), stream);
} else {
TORCH_CHECK(false, "Unsupported output data type of nvfp4 mm sm120 (", out_dtype, ")");
Panic("Unsupported output data type of nvfp4 mm sm120");
}
} else {
// Use SM100 dispatch for other architectures
if (out_dtype == at::ScalarType::Half) {
if (host::is_type<fp16_t>(D.dtype())) {
cutlassFp4GemmDispatch<cutlass::half_t>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
} else if (out_dtype == at::ScalarType::BFloat16) {
} else if (host::is_type<bf16_t>(D.dtype())) {
cutlassFp4GemmDispatch<cutlass::bfloat16_t>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
} else if (out_dtype == at::ScalarType::Float) {
} else if (host::is_type<float>(D.dtype())) {
cutlassFp4GemmDispatch<float>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
} else {
TORCH_CHECK(false, "Unsupported output data type of nvfp4 mm");
Panic("Unsupported output data type of nvfp4 mm");
}
}
}

492
sgl-kernel/csrc/moe/nvfp4_blockwise_moe.cu → python/sglang/jit_kernel/csrc/moe/nvfp4_blockwise_moe.cuh
View File

@@ -1,10 +1,11 @@
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <cutlass/arch/arch.h>
#include <torch/all.h>
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <cassert>
#include <sgl_kernel/runtime.cuh>
#include <sgl_kernel/utils.cuh>
#include <cutlass/arch/arch.h>
#include <cutlass/cutlass.h>
#include "cute/tensor.hpp"
#include "cutlass/epilogue/collective/collective_builder.hpp"
@@ -27,10 +28,66 @@
#include "cutlass/util/reference/host/tensor_fill.h"
#include "cutlass/util/reference/host/tensor_norm.h"
#include "cutlass/util/tensor_view_io.h"
#include "utils.h"
#include <algorithm>
#include <cassert>
#include <cstdint>
#include <limits>
#include <unordered_map>
using namespace host;
using namespace cute;
struct WorkspaceKey {
int device_id;
uintptr_t stream;
auto operator==(const WorkspaceKey&) const -> bool = default;
};
struct WorkspaceKeyHash {
auto operator()(const WorkspaceKey& key) const -> size_t {
size_t h1 = std::hash<int>{}(key.device_id);
size_t h2 = std::hash<uintptr_t>{}(key.stream);
return h1 ^ (h2 + 0x9e3779b97f4a7c15ULL + (h1 << 6) + (h1 >> 2));
}
};
struct WorkspaceState {
void* ptr = nullptr;
size_t bytes = 0;
};
inline auto get_cached_workspace(size_t required_bytes, int device_id, cudaStream_t stream) -> void* {
if (required_bytes == 0) {
return nullptr;
}
thread_local std::unordered_map<WorkspaceKey, WorkspaceState, WorkspaceKeyHash> cache;
WorkspaceKey key{device_id, reinterpret_cast<uintptr_t>(stream)};
auto& ws = cache[key];
if (ws.ptr != nullptr && ws.bytes >= required_bytes) {
return ws.ptr;
}
RuntimeDeviceCheck(cudaSetDevice(device_id));
if (ws.ptr != nullptr) {
RuntimeDeviceCheck(cudaFreeAsync(ws.ptr, stream));
ws.ptr = nullptr;
ws.bytes = 0;
}
RuntimeDeviceCheck(cudaMallocAsync(&ws.ptr, required_bytes, stream));
ws.bytes = required_bytes;
return ws.ptr;
}
inline int getSMVersion(int device_id) {
int sm_major = 0;
int sm_minor = 0;
RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device_id));
RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_minor, cudaDevAttrComputeCapabilityMinor, device_id));
return sm_major * 10 + sm_minor;
}
template <
typename ElementAB,
typename ElementC,
@@ -104,8 +161,8 @@ __global__ void __get_group_gemm_starts(
}
#define __CALL_GET_STARTS_KERNEL_BLOCKSCALE( \
ELEMENT_AB_TYPE, SF_TYPE, TENSOR_C_TYPE, C_TYPE, LayoutSFA, LayoutSFB, ScaleConfig) \
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<ELEMENT_AB_TYPE, C_TYPE, SF_TYPE, float, LayoutSFA, LayoutSFB, ScaleConfig> \
<<<1, num_experts, 0, stream>>>( \
static_cast<ELEMENT_AB_TYPE**>(a_starts.data_ptr()), \
@@ -131,32 +188,32 @@ __global__ void __get_group_gemm_starts(
template <typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
void run_get_group_gemm_starts(
const torch::Tensor& a_starts,
const 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<int>(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<bf16_t>(out_tensors.dtype()),
cutlass::bfloat16_t,
LayoutSFA,
LayoutSFB,
ScaleConfig)
__CALL_GET_STARTS_KERNEL_BLOCKSCALE(
cutlass::float_e2m1_t, cutlass::float_ue4m3_t, torch::kFloat16, half, LayoutSFA, LayoutSFB, ScaleConfig)
cutlass::float_e2m1_t,
cutlass::float_ue4m3_t,
host::is_type<fp16_t>(out_tensors.dtype()),
cutlass::half_t,
LayoutSFA,
LayoutSFB,
ScaleConfig)
else {
TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
Panic("Invalid output type (must be float16 or bfloat16)");
}
}
void run_fp4_blockwise_scaled_group_mm_sm120(
torch::Tensor& output,
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<int>(expert_offsets.size(0));
auto options_int = torch::TensorOptions().dtype(torch::kInt64).device(a.device());
torch::Tensor a_ptrs = torch::empty(num_experts, options_int);
torch::Tensor b_ptrs = torch::empty(num_experts, options_int);
torch::Tensor out_ptrs = torch::empty(num_experts, options_int);
torch::Tensor a_scales_ptrs = torch::empty(num_experts, options_int);
torch::Tensor b_scales_ptrs = torch::empty(num_experts, options_int);
torch::Tensor alpha_ptrs = torch::empty(num_experts, options_int);
torch::Tensor layout_sfa = torch::empty({num_experts, 5}, options_int);
torch::Tensor layout_sfb = torch::empty({num_experts, 5}, options_int);
run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
a_ptrs,
@@ -320,13 +381,13 @@ void run_fp4_blockwise_scaled_group_mm_sm120(
using RasterOrderOptions = cutlass::gemm::kernel::detail::RasterOrderOptions;
typename Gemm::GemmKernel::TileSchedulerArguments scheduler;
scheduler.raster_order = RasterOrderOptions::AlongM;
hw_info.device_id = a.get_device();
hw_info.device_id = a.device().device_id;
static std::unordered_map<int, int> cached_sm_counts;
if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) {
cached_sm_counts[hw_info.device_id] =
cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
}
hw_info.sm_count = min(cached_sm_counts[hw_info.device_id], INT_MAX);
hw_info.sm_count = std::min(cached_sm_counts[hw_info.device_id], std::numeric_limits<int>::max());
// Mainloop Arguments
typename GemmKernel::MainloopArguments mainloop_args{
@@ -361,34 +422,44 @@ void run_fp4_blockwise_scaled_group_mm_sm120(
scheduler};
size_t workspace_size = Gemm::get_workspace_size(args);
auto const workspace_options = torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
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 <typename OutType>
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<int>(expert_offsets.size(0));
auto options_int = torch::TensorOptions().dtype(torch::kInt64).device(a.device());
torch::Tensor a_ptrs = torch::empty(num_experts, options_int);
torch::Tensor b_ptrs = torch::empty(num_experts, options_int);
torch::Tensor out_ptrs = torch::empty(num_experts, options_int);
torch::Tensor a_scales_ptrs = torch::empty(num_experts, options_int);
torch::Tensor b_scales_ptrs = torch::empty(num_experts, options_int);
torch::Tensor alpha_ptrs = torch::empty(num_experts, options_int);
torch::Tensor layout_sfa = torch::empty({num_experts, 5}, options_int);
torch::Tensor layout_sfb = torch::empty({num_experts, 5}, options_int);
run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
a_ptrs,
@@ -519,13 +580,13 @@ void run_fp4_blockwise_scaled_group_mm_sm100(
typename ProblemShape::UnderlyingProblemShape>::RasterOrderOptions;
typename Gemm::GemmKernel::TileSchedulerArguments scheduler;
scheduler.raster_order = RasterOrderOptions::AlongM;
hw_info.device_id = a.get_device();
hw_info.device_id = a.device().device_id;
static std::unordered_map<int, int> cached_sm_counts;
if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) {
cached_sm_counts[hw_info.device_id] =
cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
}
hw_info.sm_count = min(cached_sm_counts[hw_info.device_id], INT_MAX);
hw_info.sm_count = std::min(cached_sm_counts[hw_info.device_id], std::numeric_limits<int>::max());
// Mainloop Arguments
typename GemmKernel::MainloopArguments mainloop_args{
@@ -559,80 +620,144 @@ void run_fp4_blockwise_scaled_group_mm_sm100(
scheduler};
size_t workspace_size = Gemm::get_workspace_size(args);
auto const workspace_options = torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
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<uint8_t>(a.dtype()), "a must be uint8");
RuntimeCheck(host::is_type<uint8_t>(b.dtype()), "b must be uint8");
RuntimeCheck(host::is_type<fp8_e4m3_t>(a_blockscale.dtype()), "a_blockscale must be float8_e4m3fn");
RuntimeCheck(host::is_type<fp8_e4m3_t>(b_blockscales.dtype()), "b_blockscales must be float8_e4m3fn");
RuntimeCheck(host::is_type<float>(alphas.dtype()), "alphas must be float32");
RuntimeCheck(host::is_type<int64_t>(ab_strides.dtype()), "ab_strides must be int64");
RuntimeCheck(host::is_type<int64_t>(c_strides.dtype()), "c_strides must be int64");
RuntimeCheck(host::is_type<int32_t>(problem_sizes.dtype()), "problem_sizes must be int32");
RuntimeCheck(host::is_type<int32_t>(expert_offsets.dtype()), "expert_offsets must be int32");
RuntimeCheck(host::is_type<int32_t>(sf_offsets.dtype()), "sf_offsets must be int32");
RuntimeCheck(host::is_type<int64_t>(a_ptrs.dtype()), "a_ptrs must be int64");
RuntimeCheck(host::is_type<int64_t>(b_ptrs.dtype()), "b_ptrs must be int64");
RuntimeCheck(host::is_type<int64_t>(out_ptrs.dtype()), "out_ptrs must be int64");
RuntimeCheck(host::is_type<int64_t>(a_scales_ptrs.dtype()), "a_scales_ptrs must be int64");
RuntimeCheck(host::is_type<int64_t>(b_scales_ptrs.dtype()), "b_scales_ptrs must be int64");
RuntimeCheck(host::is_type<int64_t>(alpha_ptrs.dtype()), "alpha_ptrs must be int64");
RuntimeCheck(host::is_type<int64_t>(layout_sfa.dtype()), "layout_sfa must be int64");
RuntimeCheck(host::is_type<int64_t>(layout_sfb.dtype()), "layout_sfb must be int64");
RuntimeCheck(
host::is_type<bf16_t>(output.dtype()) || host::is_type<fp16_t>(output.dtype()),
"output must be bfloat16 or float16");
TORCH_CHECK(
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<int>(expert_offsets.size(0));
RuntimeCheck(problem_sizes.size(0) == num_experts, "problem_sizes size mismatch with expert_offsets");
RuntimeCheck(sf_offsets.size(0) == num_experts, "sf_offsets size mismatch with expert_offsets");
RuntimeCheck(alphas.size(0) == num_experts, "alphas size mismatch with expert_offsets");
RuntimeCheck(ab_strides.size(0) == num_experts, "ab_strides size mismatch with expert_offsets");
RuntimeCheck(c_strides.size(0) == num_experts, "c_strides size mismatch with expert_offsets");
RuntimeCheck(a_ptrs.size(0) == num_experts, "a_ptrs size mismatch with expert_offsets");
RuntimeCheck(b_ptrs.size(0) == num_experts, "b_ptrs size mismatch with expert_offsets");
RuntimeCheck(out_ptrs.size(0) == num_experts, "out_ptrs size mismatch with expert_offsets");
RuntimeCheck(a_scales_ptrs.size(0) == num_experts, "a_scales_ptrs size mismatch with expert_offsets");
RuntimeCheck(b_scales_ptrs.size(0) == num_experts, "b_scales_ptrs size mismatch with expert_offsets");
RuntimeCheck(alpha_ptrs.size(0) == num_experts, "alpha_ptrs size mismatch with expert_offsets");
RuntimeCheck(layout_sfa.size(0) == num_experts && layout_sfa.size(1) == 5, "layout_sfa must be [num_experts, 5]");
RuntimeCheck(layout_sfb.size(0) == num_experts && layout_sfb.size(1) == 5, "layout_sfb must be [num_experts, 5]");
int M = static_cast<int>(a.size(0));
int N = static_cast<int>(b.size(1));
int E = static_cast<int>(b.size(0));
int K = static_cast<int>(2 * b.size(2));
RuntimeCheck(output.dim() == 2, "output must be 2D");
RuntimeCheck(output.size(0) == M && output.size(1) == N, "output shape mismatch");
auto sm_version = getSMVersion();
auto sm_version = getSMVersion(a.device().device_id);
if (sm_version == 100 || sm_version == 103) {
if (output.scalar_type() == torch::kBFloat16) {
if (host::is_type<bf16_t>(output.dtype())) {
run_fp4_blockwise_scaled_group_mm_sm100<cutlass::bfloat16_t>(
output,
a,
@@ -645,6 +770,14 @@ void cutlass_fp4_group_mm(
problem_sizes,
expert_offsets,
sf_offsets,
a_ptrs,
b_ptrs,
out_ptrs,
a_scales_ptrs,
b_scales_ptrs,
alpha_ptrs,
layout_sfa,
layout_sfb,
M,
N,
K);
@@ -661,12 +794,20 @@ void cutlass_fp4_group_mm(
problem_sizes,
expert_offsets,
sf_offsets,
a_ptrs,
b_ptrs,
out_ptrs,
a_scales_ptrs,
b_scales_ptrs,
alpha_ptrs,
layout_sfa,
layout_sfb,
M,
N,
K);
}
} else if (sm_version >= 120) {
if (output.scalar_type() == torch::kBFloat16) {
if (host::is_type<bf16_t>(output.dtype())) {
run_fp4_blockwise_scaled_group_mm_sm120(
output,
a,
@@ -679,20 +820,63 @@ void cutlass_fp4_group_mm(
problem_sizes,
expert_offsets,
sf_offsets,
a_ptrs,
b_ptrs,
out_ptrs,
a_scales_ptrs,
b_scales_ptrs,
alpha_ptrs,
layout_sfa,
layout_sfb,
M,
N,
K);
} else {
std::cout << "run_fp4_blockwise_scaled_group_mm_sm120 half no implementation" << std::endl;
Panic("SM120 path currently supports only bfloat16 output");
}
} else {
TORCH_CHECK_NOT_IMPLEMENTED(false, "Unsupported SM version: " + std::to_string(sm_version));
RuntimeCheck(false, "Unsupported SM version: ", sm_version);
}
#else
TORCH_CHECK_NOT_IMPLEMENTED(
false,
"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);
}

490
python/sglang/jit_kernel/nvfp4.py Normal file
View File

@@ -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

127
python/sglang/jit_kernel/tests/test_nvfp4_blockwise_moe.py Normal file
View File

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

142
python/sglang/jit_kernel/tests/test_nvfp4_gemm.py Normal file
View File

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

214
python/sglang/jit_kernel/tests/test_nvfp4_quant.py Normal file
View File

@@ -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]])

7
python/sglang/srt/layers/moe/cutlass_moe.py
View File

@@ -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,

38
python/sglang/srt/layers/moe/cutlass_moe_params.py
View File

@@ -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,
}

2
python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
View File

@@ -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

2
python/sglang/srt/layers/moe/token_dispatcher/standard.py
View File

@@ -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:

6
python/sglang/srt/layers/quantization/modelopt_quant.py
View File

@@ -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

6
sgl-kernel/CMakeLists.txt
View File

@@ -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"

30
sgl-kernel/csrc/common_extension.cc
View File

@@ -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);

77
sgl-kernel/csrc/gemm/nvfp4_quant_entry.cu
View File

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

64
sgl-kernel/csrc/gemm/nvfp4_scaled_mm_entry.cu
View File

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

38
sgl-kernel/include/sgl_kernel_ops.h
View File

@@ -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
*/

2
sgl-kernel/python/sgl_kernel/__init__.py
View File

@@ -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,

197
sgl-kernel/python/sgl_kernel/gemm.py
View File

@@ -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

49
sgl-kernel/python/sgl_kernel/moe.py
View File

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