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(