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

[kernel] add fused silu mul quant fp8 (#37376)

Browse Source 
Co-authored-by: undefined <zhouchen.arrebol@jd.com>
Co-authored-by: xq25478 <xq25478@qq.com>
Co-authored-by: xieminghe.simon <xieminghe.simon@jd.com>
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
This commit is contained in:
xieminghe1
2026-09-07 11:39:11 +08:00
committed by GitHub
co-authored by undefined xq25478 xieminghe.simon Xiaoyu Zhang
parent 503511c963
commit a8b2f36dee
4 changed files with 695 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/aot/benchmark/bench_silu_quant_fp8.py
+112
View File
@@ -0,0 +1,112 @@
"""Benchmark: fused_silu_mul_quant_fp8 vs separate silu_and_mul + per_token_group_quant_fp8 (H200).
Compares two paths for the activation + quantization step after MoE gate-up GEMM:
- baseline: silu_and_mul -> per_token_group_quant_fp8 (two kernel launches)
- fused: fused_silu_mul_quant_fp8 (one kernel launch)
"""
import sys
import time
import torch
import triton
import triton.language as tl
sys.path.insert(0, "/tmp")
def bench():
from sglang.kernels.ops.moe.fused_moe_triton_kernels import fused_silu_mul_quant_fp8
from sglang.kernels.ops.quantization.fp8_kernel import per_token_group_quant_fp8
SWIGLU_LIMIT = 0.0 # DSV4 uses swiglu_limit=10; test without clamp first
configs = [
# (num_tokens, hidden_dim, group_size)
(128, 2048, 128),
(256, 2048, 128),
(512, 2048, 128),
(1024, 2048, 128),
(2048, 2048, 128),
(4096, 2048, 128),
(8192, 2048, 128),
(128, 4096, 128),
(1024, 4096, 128),
(8192, 4096, 128),
(128, 8192, 128),
(1024, 8192, 128),
(8192, 8192, 128),
]
# Correctness check
torch.manual_seed(42)
x = torch.randn(256, 2 * 2048, device="cuda", dtype=torch.bfloat16)
d = 2048
result_ref = torch.nn.functional.silu(x[:, :d]) * x[:, d:]
ref_fp8, ref_scale = per_token_group_quant_fp8(result_ref, 128)
fused_fp8, fused_scale = fused_silu_mul_quant_fp8(x, 128)
p_diff = (ref_fp8.float() - fused_fp8.float()).abs().max().item()
s_diff = (ref_scale - fused_scale).abs().max().item()
print(f"Correctness: fp8_diff={p_diff}, scale_diff={s_diff:.6f}")
print()
print("=" * 85)
print(f"fused_silu_mul_quant_fp8 vs separate silu_and_mul + quant (H200)")
print(f"swiglu_limit={SWIGLU_LIMIT}, group_size=128")
print("=" * 85)
print(
f"{'tokens':>7} {'hidden':>8} |{'baseline(ms)':>13}{'fused(ms)':>11}{'speedup':>8} | {'scale_diff':>11}"
)
print("-" * 65)
for num_tokens, hidden_dim, group_size in configs:
torch.manual_seed(42)
x = torch.randn(num_tokens, 2 * hidden_dim, device="cuda", dtype=torch.bfloat16)
n_iter = 200 if num_tokens <= 1024 else 50
# baseline: silu_and_mul + per_token_group_quant_fp8
def run_baseline():
d = hidden_dim
result = torch.nn.functional.silu(x[:, :d]) * x[:, d:]
result_fp8, result_scale = per_token_group_quant_fp8(result, group_size)
return result_fp8, result_scale
for _ in range(10):
run_baseline()
torch.cuda.synchronize()
t0 = time.time()
for _ in range(n_iter):
run_baseline()
torch.cuda.synchronize()
lat_base = (time.time() - t0) / n_iter * 1000
# fused
def run_fused():
return fused_silu_mul_quant_fp8(x, group_size)
for _ in range(10):
run_fused()
torch.cuda.synchronize()
t0 = time.time()
for _ in range(n_iter):
run_fused()
torch.cuda.synchronize()
lat_fused = (time.time() - t0) / n_iter * 1000
# Per-config correctness
b_fp8, b_scale = run_baseline()
f_fp8, f_scale = run_fused()
s_diff = (b_scale.float() - f_scale.float()).abs().max().item()
speedup = lat_base / lat_fused if lat_fused > 0 else 0
print(
f"{num_tokens:>7} {hidden_dim:>8} |{lat_base:>12.4f}ms{lat_fused:>10.4f}ms{speedup:>7.2f}x | {s_diff:>10.6f}"
)
print("-" * 65)
print("baseline = F.silu(gate)*up + per_token_group_quant_fp8 (2 launches)")
print("fused = fused_silu_mul_quant_fp8 (1 launch)")
if __name__ == "__main__":
bench()
python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py
+96
View File
@@ -1177,6 +1177,102 @@ def act_and_mul_triton(
)
# ============================================================
# Fused silu_and_mul + per_token_group_quant_fp8 kernel
# ============================================================
_fp8_type = torch.float8_e4m3fnuz if is_hip() else torch.float8_e4m3fn
_FP8_MAX = torch.finfo(_fp8_type).max
@triton.jit
def _fused_silu_mul_quant_fp8_kernel(
input_ptr,
output_ptr,
scale_ptr,
num_tokens,
hidden_dim,
FP8_MAX: tl.constexpr,
EPS: tl.constexpr,
GROUP_SIZE: tl.constexpr,
BLOCK_M: tl.constexpr,
SWIGLU_LIMIT: tl.constexpr = 0.0,
HAS_SWIGLU_LIMIT: tl.constexpr = False,
):
"""Fused kernel: silu(gate) * up -> fp8 quantize with block-wise scales."""
pid_m = tl.program_id(0)
pid_g = tl.program_id(1)
offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
offs_k = pid_g * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask_m = offs_m < num_tokens
mask_k = offs_k < hidden_dim
mask = mask_m[:, None] & mask_k[None, :]
two_d = hidden_dim * 2
base_ptrs = input_ptr + offs_m[:, None] * two_d + offs_k[None, :]
gate = tl.load(base_ptrs, mask=mask, other=0.0).to(tl.float32)
up = tl.load(base_ptrs + hidden_dim, mask=mask, other=0.0).to(tl.float32)
# DeepSeek-V4 SwiGLU clamp: gate clamped to [-inf, L], up clamped to [-L, L]
if HAS_SWIGLU_LIMIT:
gate = tl.minimum(gate, SWIGLU_LIMIT)
up = tl.maximum(tl.minimum(up, SWIGLU_LIMIT), -SWIGLU_LIMIT)
result = (gate * tl.sigmoid(gate)) * up
group_max = tl.max(tl.abs(result), axis=1)
scale = group_max / FP8_MAX
scale = tl.where(scale > EPS, scale, EPS)
result_scaled = result / scale[:, None]
result_fp8 = tl.clamp(result_scaled, -FP8_MAX, FP8_MAX).to(
output_ptr.dtype.element_ty
)
out_ptrs = output_ptr + offs_m[:, None] * hidden_dim + offs_k[None, :]
tl.store(out_ptrs, result_fp8, mask=mask)
num_groups = hidden_dim // GROUP_SIZE
scale_mask = mask_m & (pid_g < num_groups)
scale_ptrs = scale_ptr + offs_m * num_groups + pid_g
tl.store(scale_ptrs, scale, mask=scale_mask)
def fused_silu_mul_quant_fp8(x, group_size, swiglu_limit=0.0):
"""Fused Triton kernel: silu_and_mul + per_token_group_quant_fp8 in one launch.
Args:
x: [num_tokens, 2 * hidden_dim], bf16/fp16, contiguous row-major
group_size: quantization group size (e.g. 128 for DeepSeek-V4 block-wise FP8)
swiglu_limit: SwiGLU clamp limit (0 = no clamp, 10.0 for DeepSeek-V4).
When > 0, gate is clamped to [-inf, L] and up to [-L, L] before
silu(gate) * up, matching the DeepSeek-V4 activation contract.
Returns:
(x_fp8, x_scale):
x_fp8: [num_tokens, hidden_dim], fp8
x_scale: [num_tokens, hidden_dim // group_size], float32, row-major
"""
assert x.is_contiguous(), "Input must be contiguous"
num_tokens = x.shape[0]
hidden_dim = x.shape[1] // 2
num_groups = hidden_dim // group_size
assert hidden_dim % group_size == 0
x_fp8 = torch.empty(num_tokens, hidden_dim, device=x.device, dtype=_fp8_type)
x_scale = torch.empty(num_tokens, num_groups, device=x.device, dtype=torch.float32)
BLOCK_M = 128
grid = (triton.cdiv(num_tokens, BLOCK_M), num_groups)
has_swiglu_limit = swiglu_limit is not None and swiglu_limit > 0
_fused_silu_mul_quant_fp8_kernel[grid](
x,
x_fp8,
x_scale,
num_tokens,
hidden_dim,
FP8_MAX=_FP8_MAX,
EPS=1e-10,
GROUP_SIZE=group_size,
BLOCK_M=BLOCK_M,
SWIGLU_LIMIT=float(swiglu_limit) if has_swiglu_limit else 0.0,
HAS_SWIGLU_LIMIT=has_swiglu_limit,
num_warps=4,
)
return x_fp8, x_scale
# _moe_sum_reduce_kernel kernel modified from https://github.com/ModelTC/lightllm/blob/main/lightllm/common/fused_moe/moe_sum_reduce.py
@triton.jit
def _moe_sum_reduce_kernel(
python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py
+73 -7
View File
@@ -17,6 +17,7 @@ import triton.language as tl
from sglang.kernels.jit.utils import is_arch_support_pdl
from sglang.kernels.ops.moe.fused_moe_triton_kernels import (
act_and_mul_triton,
fused_silu_mul_quant_fp8,
invoke_fused_moe_kernel,
moe_sum_reduce_triton,
support_tensor_descriptor,
@@ -125,6 +126,41 @@ def _use_moe_sum_reduce_torch_compile(num_tokens: int) -> bool:
return num_tokens <= 32 and not is_batch_invariant_mode_enabled()
def _can_use_fused_silu_mul_quant_fp8(
*,
hidden_size: int,
hidden_dtype: torch.dtype,
use_fp8_w8a8: bool,
block_shape: Optional[List[int]],
filter_expert: bool,
activation: str,
is_gated: bool,
gemm1_alpha: Optional[float],
gemm1_limit: Optional[float],
hooks: Optional[Any],
fuse_swiglu_interleaved: bool,
) -> bool:
"""Whether the fused activation/quantization kernel is a safe replacement."""
if block_shape is None or len(block_shape) != 2 or block_shape[1] <= 0:
return False
activation_size = hidden_size // 2
return (
_is_cuda
and hidden_dtype in (torch.bfloat16, torch.float16)
and use_fp8_w8a8
and hidden_size % 2 == 0
and activation_size % block_shape[1] == 0
and not filter_expert
and activation == "silu"
and is_gated
and gemm1_alpha is None
and gemm1_limit is None
and hooks is None
and not fuse_swiglu_interleaved
)
@register_custom_op(mutates_args=["hidden_states"])
def inplace_fused_experts(
hidden_states: torch.Tensor,
@@ -582,6 +618,26 @@ def _fused_moe_kernel_sequence(
):
out_hidden_states = torch.empty_like(hidden_states)
# Automatically fuse the activation and down-input quantization when the
# CUDA block-wise FP8 path satisfies every kernel and caller contract.
# Unsupported shapes, activation modifiers, expert filtering, and LoRA
# hooks keep using the existing two-kernel path below. DeepSeek-V4's
# swiglu_limit is supported and applied inside the fused kernel.
use_fused_silu_mul_quant_fp8 = _can_use_fused_silu_mul_quant_fp8(
hidden_size=N,
hidden_dtype=hidden_states.dtype,
use_fp8_w8a8=use_fp8_w8a8,
block_shape=block_shape,
filter_expert=filter_expert,
activation=activation,
is_gated=is_gated,
gemm1_alpha=gemm1_alpha,
gemm1_limit=gemm1_limit,
hooks=hooks,
fuse_swiglu_interleaved=fuse_swiglu_interleaved,
)
fused_a2_scale = None
use_fused_moe_sum_all_reduce = (
get_exec().moe.enable_fused_moe_sum_all_reduce
and (not no_combine)
@@ -660,14 +716,24 @@ def _fused_moe_kernel_sequence(
)
if not fuse_swiglu_interleaved:
intermediate_cache2 = torch.empty(
(total_tokens, N // 2),
device=hidden_states.device,
dtype=hidden_states.dtype,
)
if not use_fused_silu_mul_quant_fp8:
intermediate_cache2 = torch.empty(
(total_tokens, N // 2),
device=hidden_states.device,
dtype=hidden_states.dtype,
)
# Activation function with multiplication
if fuse_swiglu_interleaved:
if use_fused_silu_mul_quant_fp8:
# Fused path: silu_and_mul + fp8_quant in one kernel launch.
# Pass swiglu_limit so the DeepSeek-V4 clamp is applied inside the
# Triton kernel before silu(gate)*up, matching the non-fused path.
intermediate_cache2, fused_a2_scale = fused_silu_mul_quant_fp8(
intermediate_cache1.view(-1, N),
block_shape[1],
swiglu_limit=swiglu_limit if swiglu_limit is not None else 0.0,
)
elif fuse_swiglu_interleaved:
# silu(gate) * up was already applied by the up-GEMM epilogue.
pass
elif activation == "silu" and is_gated:
@@ -834,7 +900,7 @@ def _fused_moe_kernel_sequence(
else out_hidden_states.unsqueeze(0)
)
),
a2_scale,
fused_a2_scale if use_fused_silu_mul_quant_fp8 else a2_scale,
w2_scale,
w2_zp,
topk_weights,