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

[Perf][Kernel] Fuse SiLU+Mul into NVFP4 Expert Quantization for CUTLASS MoE (#18612)

Browse Source
This commit is contained in:
Ho-Ren (Jack) Chuang
2026-06-29 16:51:01 -07:00
committed by GitHub
parent f0bf96390b
commit a6bc432fd8
7 changed files with 568 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_expert_quant.cuh
+107 -13
View File
@@ -118,7 +118,8 @@ cvt_fp16_to_fp4(
uint32_t* output_scale_offset_by_experts,
int32_t* mask,
int n_experts,
bool low_latency) {
bool low_latency,
bool use_silu_and_mul) {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
using PackedVec = PackedVec<Type>;
static constexpr int CVT_FP4_NUM_THREADS_PER_SF = (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
@@ -127,11 +128,9 @@ cvt_fp16_to_fp4(
// Input tensor row/col loops.
int tid = blockIdx.x * blockDim.x + threadIdx.x;
int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
// TODO(kaixih@nvidia): For now, we assume mask is used together with
// silu_and_mal. Maybe we want a more general behavior of mask later. In the
// silu case, the input last dim doubles.
bool use_mask = mask != nullptr;
int actualColsPerRow = use_mask ? colsPerRow * 2 : colsPerRow;
// When use_silu_and_mul is true, input last dim is 2*k (gate+up concatenated).
int actualColsPerRow = (use_mask || use_silu_and_mul) ? colsPerRow * 2 : colsPerRow;
// Each global thread processes one element
for (int globalIdx = tid; globalIdx < numRows * colsPerRow; globalIdx += gridDim.x * blockDim.x) {
@@ -188,7 +187,7 @@ cvt_fp16_to_fp4(
int64_t inOffset = rowIdx * actualColsPerRow + colIdx;
PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
if (use_mask) {
if (use_mask || use_silu_and_mul) {
PackedVec in_vec_mul = reinterpret_cast<PackedVec const*>(in)[inOffset + colsPerRow];
silu_and_mul(in_vec, in_vec_mul);
}
@@ -335,7 +334,8 @@ cvt_fp16_to_fp4(
uint32_t* input_offset_by_experts,
uint32_t* output_scale_offset_by_experts,
int32_t* mask,
int n_experts) {
int n_experts,
bool use_silu_and_mul) {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
using PackedVec = PackedVec<Type>;
static constexpr int CVT_FP4_NUM_THREADS_PER_SF = (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
@@ -363,7 +363,8 @@ cvt_fp16_to_fp4(
int tid = blockIdx.x * blockDim.x + threadIdx.x;
int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
bool use_mask = mask != nullptr;
int actualColsPerRow = use_mask ? colsPerRow * 2 : colsPerRow;
// When use_silu_and_mul is true, input last dim is 2*k (gate+up concatenated).
int actualColsPerRow = (use_mask || use_silu_and_mul) ? colsPerRow * 2 : colsPerRow;
// Each global thread processes one element
for (int globalIdx = tid; globalIdx < numRows * colsPerRow; globalIdx += gridDim.x * blockDim.x) {
@@ -402,7 +403,7 @@ cvt_fp16_to_fp4(
int64_t inOffset = rowIdx * actualColsPerRow + colIdx;
PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
if (use_mask) {
if (use_mask || use_silu_and_mul) {
PackedVec in_vec_mul = reinterpret_cast<PackedVec const*>(in)[inOffset + colsPerRow];
silu_and_mul(in_vec, in_vec_mul);
}
@@ -488,7 +489,8 @@ void quant_impl(
reinterpret_cast<uint32_t*>(input_offset_by_experts),
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
reinterpret_cast<int32_t*>(mask),
n_experts);
n_experts,
use_silu_and_mul);
} else {
cvt_fp16_to_fp4<T, false, true><<<grid, block, shared_mem_size, stream>>>(
m_topk,
@@ -500,7 +502,8 @@ void quant_impl(
reinterpret_cast<uint32_t*>(input_offset_by_experts),
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
reinterpret_cast<int32_t*>(mask),
n_experts);
n_experts,
use_silu_and_mul);
}
} else {
if (n_experts >= 16) {
@@ -515,7 +518,8 @@ void quant_impl(
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
reinterpret_cast<int32_t*>(mask),
n_experts,
/* bool low_latency */ true);
/* bool low_latency */ true,
use_silu_and_mul);
} else {
cvt_fp16_to_fp4<T, false, true><<<grid, block, 0, stream>>>(
m_topk,
@@ -528,7 +532,8 @@ void quant_impl(
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
reinterpret_cast<int32_t*>(mask),
n_experts,
/* bool low_latency */ true);
/* bool low_latency */ true,
use_silu_and_mul);
}
}
}
@@ -710,3 +715,92 @@ void silu_and_mul_scaled_fp4_experts_quant_sm100a(
stream);
}
}
void silu_and_mul_scaled_fp4_experts_quant_packed_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) {
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 OffsetSize = SymbolicSize{"offset_size"};
auto device = SymbolicDevice{};
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({OffsetSize}) //
.with_dtype<int32_t>()
.with_device(device)
.verify(input_offset_by_experts)
.verify(output_scale_offset_by_experts);
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;
const auto m_topk = static_cast<int>(MTopK.unwrap());
const auto k_by_2 = static_cast<int>(KBy2.unwrap());
// Input last dim is 2*k (gate+up concatenated). The kernel does SiLU(gate)*up
// then FP4-quantizes the k-dim result.
RuntimeCheck(k_by_2 % 2 == 0, "input last dim must be even (2*k)");
const int k = k_by_2 / 2;
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");
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(),
input.data_ptr(),
input_global_scale.data_ptr(),
input_offset_by_experts.data_ptr(),
output_scale_offset_by_experts.data_ptr(),
nullptr, // mask
true, // use_silu_and_mul
m_topk,
k,
n_experts,
stream);
} else {
quant_impl<__nv_bfloat16>(
output.data_ptr(),
output_scale.data_ptr(),
input.data_ptr(),
input_global_scale.data_ptr(),
input_offset_by_experts.data_ptr(),
output_scale_offset_by_experts.data_ptr(),
nullptr, // mask
true, // use_silu_and_mul
m_topk,
k,
n_experts,
stream);
}
}
python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant_entry.cuh
+19
View File
@@ -38,6 +38,14 @@ void silu_and_mul_scaled_fp4_experts_quant_sm100a(
tvm::ffi::TensorView mask,
bool use_silu_and_mul);
void silu_and_mul_scaled_fp4_experts_quant_packed_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 scaled_fp4_quant(
tvm::ffi::TensorView output,
tvm::ffi::TensorView input,
@@ -66,3 +74,14 @@ void silu_and_mul_scaled_fp4_experts_quant(
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);
}
void silu_and_mul_scaled_fp4_experts_quant_packed(
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) {
silu_and_mul_scaled_fp4_experts_quant_packed_sm100a(
output, output_scale, input, input_global_scale, input_offset_by_experts, output_scale_offset_by_experts);
}
python/sglang/jit_kernel/nvfp4.py
+93
View File
@@ -90,6 +90,10 @@ def _jit_nvfp4_expert_quant_module() -> Module:
"silu_and_mul_scaled_fp4_experts_quant",
"silu_and_mul_scaled_fp4_experts_quant_sm100a",
),
(
"silu_and_mul_scaled_fp4_experts_quant_packed",
"silu_and_mul_scaled_fp4_experts_quant_packed_sm100a",
),
],
extra_dependencies=["cutlass"],
extra_cuda_cflags=_nvfp4_cuda_flags(),
@@ -355,6 +359,95 @@ def scaled_fp4_experts_quant(
return output, output_scales
@register_custom_op(
op_name="silu_and_mul_scaled_fp4_experts_quant_packed",
mutates_args=["output", "output_scales"],
)
def _silu_and_mul_scaled_fp4_experts_quant_packed_custom_op(
output: torch.Tensor,
output_scales: torch.Tensor,
input_tensor: torch.Tensor,
input_global_scale: torch.Tensor,
expert_offsets: torch.Tensor,
blockscale_offsets: torch.Tensor,
) -> None:
module = _jit_nvfp4_expert_quant_module()
module.silu_and_mul_scaled_fp4_experts_quant_packed(
output,
output_scales,
input_tensor,
input_global_scale,
expert_offsets,
blockscale_offsets,
)
@debug_kernel_api
def silu_and_mul_scaled_fp4_experts_quant_packed(
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]:
"""Fused SiLU+mul then FP4 quant for packed MoE inputs (expert_offsets aware).
Input shape is (m, 2*k) — gate+up concatenated. The kernel does SiLU(gate)*up
then FP4-quantizes the k-dim result.
"""
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_doubled = input_tensor.shape
k = k_input_doubled // 2
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
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,
)
_silu_and_mul_scaled_fp4_experts_quant_packed_custom_op(
output,
output_scales,
input_tensor,
input_global_scale,
expert_offsets,
blockscale_offsets,
)
output_scales = output_scales.view(torch.float8_e4m3fn)
return output, output_scales
@register_custom_op(
op_name="scaled_fp4_grouped_quant",
mutates_args=["output", "output_scales"],
python/sglang/srt/layers/moe/cutlass_moe.py
+5 -8
View File
@@ -23,6 +23,7 @@ if _is_cuda:
from sglang.jit_kernel.nvfp4 import (
cutlass_fp4_group_mm,
scaled_fp4_experts_quant,
silu_and_mul_scaled_fp4_experts_quant_packed,
)
@@ -468,19 +469,15 @@ def cutlass_moe_fp4(
)
del rep_a_fp4, rep_a_blockscale
# hidden size dimension is split to one half sized tensor.
intermediate = torch.empty(
(m_a * num_topk, w1_fp4.shape[1] // 2), device=device, dtype=out_dtype
)
silu_and_mul(c1, intermediate)
int_fp4, int_blockscale = scaled_fp4_experts_quant(
intermediate,
# fused: SiLU + mul then FP4 quant (expert-packed)
int_fp4, int_blockscale = silu_and_mul_scaled_fp4_experts_quant_packed(
c1,
a2_gscale,
params.expert_offsets,
params.blockscale_offsets,
num_topk,
)
c2 = cutlass_fp4_group_mm(
int_fp4,
w2_fp4,
python/sglang/srt/layers/moe/moe_runner/runner.py
+2
View File
@@ -65,6 +65,8 @@ class MoeRunner:
self.runner_core = None # FlashInfer CUTLASS only supports fused path
elif runner_backend.is_flashinfer_mxfp4():
self.runner_core = None # FlashInfer MXFP4 only supports fused path
elif runner_backend.is_cutlass():
self.runner_core = None # CUTLASS uses the direct cutlass_moe_fp4 path
else:
raise NotImplementedError(f"Unsupported runner backend: {runner_backend}")
python/sglang/srt/layers/quantization/modelopt_quant.py
+5 -1
View File
@@ -2268,7 +2268,11 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
if moe_runner_backend.is_flashinfer_cutlass():
import sglang.srt.layers.moe.moe_runner.flashinfer_cutlass # noqa: F401
self.runner = MoeRunner(moe_runner_backend, moe_runner_config)
# The plain CUTLASS backend uses the direct cutlass_moe_fp4 fused path
# (see apply()), not a registered MoeRunner fused func, so skip creating
# a MoeRunner for it -- constructing one would fail the fused-func check.
if not moe_runner_backend.is_cutlass():
self.runner = MoeRunner(moe_runner_backend, moe_runner_config)
def apply(
self,