From a8b2f36dee1a85105b95504c6abde690115932d8 Mon Sep 17 00:00:00 2001 From: xieminghe1 <141820649+xieminghe1@users.noreply.github.com> Date: Mon, 7 Sep 2026 11:39:11 +0800 Subject: [PATCH] [kernel] add fused silu mul quant fp8 (#37376) Co-authored-by: undefined Co-authored-by: xq25478 Co-authored-by: xieminghe.simon Co-authored-by: Xiaoyu Zhang <1182563586@qq.com> --- .../aot/benchmark/bench_silu_quant_fp8.py | 112 +++++ .../ops/moe/fused_moe_triton_kernels.py | 96 ++++ .../moe/moe_runner/triton_utils/fused_moe.py | 80 +++- .../test_fused_silu_mul_quant_fp8.py | 414 ++++++++++++++++++ 4 files changed, 695 insertions(+), 7 deletions(-) create mode 100644 python/sglang/kernels/aot/benchmark/bench_silu_quant_fp8.py create mode 100644 test/registered/kernels/ops/quantization/test_fused_silu_mul_quant_fp8.py diff --git a/python/sglang/kernels/aot/benchmark/bench_silu_quant_fp8.py b/python/sglang/kernels/aot/benchmark/bench_silu_quant_fp8.py new file mode 100644 index 000000000..81b27207e --- /dev/null +++ b/python/sglang/kernels/aot/benchmark/bench_silu_quant_fp8.py @@ -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() diff --git a/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py b/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py index e90e0fb08..c98f25dfc 100644 --- a/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py +++ b/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py @@ -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( diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py index ce665e936..564c70316 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py @@ -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, diff --git a/test/registered/kernels/ops/quantization/test_fused_silu_mul_quant_fp8.py b/test/registered/kernels/ops/quantization/test_fused_silu_mul_quant_fp8.py new file mode 100644 index 000000000..e0b6302f9 --- /dev/null +++ b/test/registered/kernels/ops/quantization/test_fused_silu_mul_quant_fp8.py @@ -0,0 +1,414 @@ +"""Correctness tests for fused_silu_mul_quant_fp8 kernel. + +The reference is computed in pure PyTorch: silu(gate) * up followed by +per_token_group_quant_fp8. The fused kernel performs both steps in a single +Triton launch, so we verify that the output fp8 codes and scales match the +two-step baseline within FP8 quantization tolerance. + +Test strategy: + - scale_diff: max abs difference of per-group scales (should be ~0) + - fp8_match: percentage of FP8 outputs that are bit-exact + - cosine_sim: cosine similarity of dequantized outputs (should be > 0.9999) + - rel_err: mean relative error of dequantized outputs (should be < 0.01) +""" + +import sys + +import pytest +import torch + +import sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe as fused_moe_module +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, +) +from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import ( + _can_use_fused_silu_mul_quant_fp8, + fused_experts_impl, +) +from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.layer_ut_utils import init_single_process_dist + +register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +G = 128 +EPS = 1e-10 + + +# --------------------------------------------------------------------------- # +# Pure-torch reference (two-step baseline). +# --------------------------------------------------------------------------- # +def ref_silu_mul_quant(x: torch.Tensor, gs: int): + """Two-step baseline: silu_and_mul + per_token_group_quant_fp8.""" + d = x.shape[-1] // 2 + result = torch.nn.functional.silu(x[..., :d]) * x[..., d:] + fp8_out, scale = per_token_group_quant_fp8(result, gs) + return fp8_out, scale + + +def ref_silu_mul_quant_swiglu_clamp(x: torch.Tensor, gs: int, limit: float): + """Two-step baseline with DeepSeek-V4 SwiGLU clamp. + + gate is clamped to [-inf, limit]; up is clamped to [-limit, limit]. + Then silu(gate) * up is computed and quantized. + """ + d = x.shape[-1] // 2 + gate = x[..., :d].float() + up = x[..., d:].float() + gate = gate.clamp(min=None, max=limit) + up = up.clamp(min=-limit, max=limit) + result = (torch.nn.functional.silu(gate) * up).to(x.dtype) + fp8_out, scale = per_token_group_quant_fp8(result, gs) + return fp8_out, scale + + +def dequantize(fp8_codes: torch.Tensor, scales: torch.Tensor, gs: int): + """Dequantize fp8 codes back to float32 using per-group scales.""" + return fp8_codes.float() * scales.repeat_interleave(gs, dim=1).float() + + +# --------------------------------------------------------------------------- # +# Parametrized test configurations. +# --------------------------------------------------------------------------- @@ +SIZES = [ + # (num_tokens, hidden_dim, group_size) + (1, 2048, 128), + (16, 2048, 128), + (256, 2048, 128), + (4096, 2048, 128), + (8192, 2048, 128), + (256, 4096, 128), + (256, 8192, 128), + (256, 2048, 64), + (256, 2048, 256), +] + +DTYPES = [torch.bfloat16, torch.float16] + + +def test_fused_silu_mul_quant_fp8_auto_dispatch(): + """The fast path is automatic only when all caller contracts are satisfied.""" + compatible = dict( + hidden_size=4096, + hidden_dtype=torch.bfloat16, + use_fp8_w8a8=True, + block_shape=[128, 128], + filter_expert=False, + activation="silu", + is_gated=True, + gemm1_alpha=None, + gemm1_limit=None, + hooks=None, + fuse_swiglu_interleaved=False, + ) + assert _can_use_fused_silu_mul_quant_fp8(**compatible) + + incompatible_overrides = ( + {"hidden_size": 4100}, + {"hidden_dtype": torch.float32}, + {"use_fp8_w8a8": False}, + {"block_shape": None}, + {"block_shape": [128, 0]}, + {"filter_expert": True}, + {"activation": "gelu"}, + {"is_gated": False}, + {"gemm1_alpha": 1.0}, + {"gemm1_limit": 7.0}, + {"hooks": object()}, + {"fuse_swiglu_interleaved": True}, + ) + for override in incompatible_overrides: + inputs = compatible | override + assert not _can_use_fused_silu_mul_quant_fp8(**inputs), override + + +def test_fused_silu_mul_quant_fp8_auto_dispatch_moe_path(monkeypatch): + """The compatible MoE path invokes the fusion and matches its fallback.""" + set_global_server_args_for_scheduler(ServerArgs(model_path="dummy")) + init_single_process_dist(master_port=29676) + + torch.manual_seed(42) + num_tokens, hidden_size, intermediate_size = 8, 256, 256 + num_experts, topk, group_size = 4, 2, 128 + block_shape = [group_size, group_size] + + hidden_states = torch.randn( + num_tokens, hidden_size, device="cuda", dtype=torch.bfloat16 + ) + w1 = ( + torch.randn( + num_experts, + 2 * intermediate_size, + hidden_size, + device="cuda", + dtype=torch.float32, + ) + .mul_(0.02) + .to(torch.float8_e4m3fn) + ) + w2 = ( + torch.randn( + num_experts, + hidden_size, + intermediate_size, + device="cuda", + dtype=torch.float32, + ) + .mul_(0.02) + .to(torch.float8_e4m3fn) + ) + w1_scale = torch.ones( + num_experts, + 2 * intermediate_size // group_size, + hidden_size // group_size, + device="cuda", + dtype=torch.float32, + ) + w2_scale = torch.ones( + num_experts, + hidden_size // group_size, + intermediate_size // group_size, + device="cuda", + dtype=torch.float32, + ) + topk_ids = torch.tensor( + [[0, 1], [1, 2], [2, 3], [3, 0]] * 2, device="cuda", dtype=torch.int32 + ) + topk_weights = torch.full( + (num_tokens, topk), 1.0 / topk, device="cuda", dtype=torch.float32 + ) + + def run_moe(): + return fused_experts_impl( + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + use_fp8_w8a8=True, + w1_scale=w1_scale, + w2_scale=w2_scale, + block_shape=block_shape, + filter_expert=False, + ) + + original_fused_kernel = fused_moe_module.fused_silu_mul_quant_fp8 + fused_kernel_calls = 0 + + def record_fused_kernel(*args, **kwargs): + nonlocal fused_kernel_calls + fused_kernel_calls += 1 + return original_fused_kernel(*args, **kwargs) + + monkeypatch.setattr( + fused_moe_module, "fused_silu_mul_quant_fp8", record_fused_kernel + ) + fused_output = run_moe() + assert fused_kernel_calls == 1 + + monkeypatch.setattr( + fused_moe_module, "_can_use_fused_silu_mul_quant_fp8", lambda **_: False + ) + fallback_output = run_moe() + assert fused_kernel_calls == 1 + torch.testing.assert_close(fused_output, fallback_output, rtol=0.02, atol=0.02) + + +@pytest.mark.parametrize("num_tokens, hidden_dim, group_size", SIZES) +@pytest.mark.parametrize("dtype", DTYPES) +def test_fused_silu_mul_quant_fp8_correctness( + num_tokens, hidden_dim, group_size, dtype +): + """Verify fused kernel matches the two-step baseline.""" + torch.manual_seed(42) + x = torch.randn(num_tokens, 2 * hidden_dim, device="cuda", dtype=dtype) * 3.0 + + # Baseline + ref_fp8, ref_scale = ref_silu_mul_quant(x, group_size) + + # Fused + fused_fp8, fused_scale = fused_silu_mul_quant_fp8(x, group_size) + + # Scale difference + scale_diff = (ref_scale.float() - fused_scale.float()).abs().max().item() + assert scale_diff < 0.01, ( + f"scale_diff={scale_diff} exceeds tolerance for " + f"shape=({num_tokens}, {hidden_dim}), gs={group_size}, dtype={dtype}" + ) + + # FP8 bit-exact match rate + fp8_match = (ref_fp8 == fused_fp8).float().mean().item() + assert fp8_match > 0.95, ( + f"fp8_match={fp8_match:.1%} below 95% for " + f"shape=({num_tokens}, {hidden_dim}), gs={group_size}, dtype={dtype}" + ) + + # Dequantized cosine similarity + ref_dq = dequantize(ref_fp8, ref_scale, group_size) + fused_dq = dequantize(fused_fp8, fused_scale, group_size) + cosine_sim = torch.nn.functional.cosine_similarity( + ref_dq.flatten().unsqueeze(0), fused_dq.flatten().unsqueeze(0) + ).item() + assert cosine_sim > 0.9999, ( + f"cosine_sim={cosine_sim} below 0.9999 for " + f"shape=({num_tokens}, {hidden_dim}), gs={group_size}, dtype={dtype}" + ) + + # Relative error + denom = ref_dq.flatten().abs().clamp(min=1e-6) + rel_err = ((ref_dq.flatten() - fused_dq.flatten()).abs() / denom).mean().item() + assert rel_err < 0.01, ( + f"rel_err={rel_err} exceeds 0.01 for " + f"shape=({num_tokens}, {hidden_dim}), gs={group_size}, dtype={dtype}" + ) + + +def test_fused_silu_mul_quant_fp8_zero_input(): + """All-zero input should produce all-zero output and near-zero scales.""" + x = torch.zeros(16, 2 * 2048, device="cuda", dtype=torch.bfloat16) + fp8_out, scale = fused_silu_mul_quant_fp8(x, G) + max_scale = scale.max().item() + assert max_scale < 1e-5, f"Expected near-zero scale, got {max_scale}" + assert fp8_out.float().abs().max().item() == 0.0, "Expected all-zero fp8 output" + + +def test_fused_silu_mul_quant_fp8_negative_input(): + """Negative input values should be handled correctly.""" + x = -torch.randn(16, 2 * 2048, device="cuda", dtype=torch.bfloat16) * 3.0 + ref_fp8, ref_scale = ref_silu_mul_quant(x, G) + fused_fp8, fused_scale = fused_silu_mul_quant_fp8(x, G) + scale_diff = (ref_scale.float() - fused_scale.float()).abs().max().item() + assert scale_diff < 0.01, f"Negative input scale_diff={scale_diff}" + + +@pytest.mark.parametrize("group_size", [32, 64, 128, 256]) +def test_fused_silu_mul_quant_fp8_group_sizes(group_size): + """Test various group sizes.""" + hidden_dim = group_size * 4 # ensure divisibility + x = torch.randn(8, 2 * hidden_dim, device="cuda", dtype=torch.bfloat16) + ref_fp8, ref_scale = ref_silu_mul_quant(x, group_size) + fused_fp8, fused_scale = fused_silu_mul_quant_fp8(x, group_size) + assert ref_fp8.shape == fused_fp8.shape + assert ref_scale.shape == fused_scale.shape + scale_diff = (ref_scale.float() - fused_scale.float()).abs().max().item() + assert scale_diff < 0.01, f"gs={group_size} scale_diff={scale_diff}" + + +# --------------------------------------------------------------------------- # +# P2: swiglu_limit clamp correctness. +# +# DeepSeek-V4 uses swiglu_limit=10.0 (see config.json "swiglu_limit": 10.0). +# The non-fused path asserts swiglu_limit == 10. We test primarily with the +# DSV4 value, plus a few edge cases. +# --------------------------------------------------------------------------- # +DSV4_SWIGLU_LIMIT = 10.0 # DeepSeek-V4 actual value from config.json + + +@pytest.mark.parametrize("swiglu_limit", [1.0, 5.0, DSV4_SWIGLU_LIMIT]) +def test_fused_silu_mul_quant_fp8_swiglu_limit(swiglu_limit): + """swiglu_limit must clamp gate to [-inf, L] and up to [-L, L] before silu*up. + + This exercises the DeepSeek-V4 activation contract. We use inputs with + large magnitudes so the clamp is guaranteed to trigger. + """ + torch.manual_seed(42) + x = torch.randn(64, 2 * 2048, device="cuda", dtype=torch.bfloat16) * 20.0 + + # Reference with clamp + ref_fp8, ref_scale = ref_silu_mul_quant_swiglu_clamp(x, G, swiglu_limit) + # Fused with clamp + fused_fp8, fused_scale = fused_silu_mul_quant_fp8(x, G, swiglu_limit=swiglu_limit) + + # Scales should match closely + scale_diff = (ref_scale.float() - fused_scale.float()).abs().max().item() + assert scale_diff < 0.01, f"swiglu_limit={swiglu_limit} scale_diff={scale_diff}" + + # FP8 codes should be mostly bit-exact + fp8_match = (ref_fp8 == fused_fp8).float().mean().item() + assert fp8_match > 0.95, f"swiglu_limit={swiglu_limit} fp8_match={fp8_match:.1%}" + + # Dequantized cosine similarity + ref_dq = dequantize(ref_fp8, ref_scale, G) + fused_dq = dequantize(fused_fp8, fused_scale, G) + cosine_sim = torch.nn.functional.cosine_similarity( + ref_dq.flatten().unsqueeze(0), fused_dq.flatten().unsqueeze(0) + ).item() + assert cosine_sim > 0.9999, f"swiglu_limit={swiglu_limit} cosine_sim={cosine_sim}" + + +def test_fused_silu_mul_quant_fp8_swiglu_limit_dsv4(): + """DeepSeek-V4 exact configuration: swiglu_limit=10.0, hidden_dim=2048, + group_size=128, bf16. + + DSV4 config.json: swiglu_limit=10.0, moe_intermediate_size=2048, + weight_block_size=[128, 128]. This test uses the exact production values. + """ + torch.manual_seed(42) + # DSV4 expert weights are [2048, 2048] (w1) and [4096, 1024] (w2). + # hidden_dim=4096, intermediate=2048 -> gate_up dim = 2*2048 = 4096. + # But the fused kernel operates on the gate_up output (intermediate_cache1) + # which has shape [num_tokens, 2 * intermediate_size]. + num_tokens = 256 + intermediate_size = 2048 # DSV4 moe_intermediate_size + x = ( + torch.randn( + num_tokens, 2 * intermediate_size, device="cuda", dtype=torch.bfloat16 + ) + * 20.0 + ) + + ref_fp8, ref_scale = ref_silu_mul_quant_swiglu_clamp(x, G, DSV4_SWIGLU_LIMIT) + fused_fp8, fused_scale = fused_silu_mul_quant_fp8( + x, G, swiglu_limit=DSV4_SWIGLU_LIMIT + ) + + scale_diff = (ref_scale.float() - fused_scale.float()).abs().max().item() + assert scale_diff < 0.01, f"DSV4 config scale_diff={scale_diff}" + + fp8_match = (ref_fp8 == fused_fp8).float().mean().item() + assert fp8_match > 0.95, f"DSV4 config fp8_match={fp8_match:.1%}" + + ref_dq = dequantize(ref_fp8, ref_scale, G) + fused_dq = dequantize(fused_fp8, fused_scale, G) + cosine_sim = torch.nn.functional.cosine_similarity( + ref_dq.flatten().unsqueeze(0), fused_dq.flatten().unsqueeze(0) + ).item() + assert cosine_sim > 0.9999, f"DSV4 config cosine_sim={cosine_sim}" + + +def test_fused_silu_mul_quant_fp8_swiglu_limit_changes_output(): + """swiglu_limit=10.0 (DSV4 value) must produce different output than no clamp + when inputs exceed the limit.""" + torch.manual_seed(42) + x = torch.randn(64, 2 * 2048, device="cuda", dtype=torch.bfloat16) * 20.0 + + no_clamp_fp8, no_clamp_scale = fused_silu_mul_quant_fp8(x, G, swiglu_limit=0.0) + clamp_fp8, clamp_scale = fused_silu_mul_quant_fp8( + x, G, swiglu_limit=DSV4_SWIGLU_LIMIT + ) + + # With large inputs and clamp=10.0, outputs must differ + diff_rate = (no_clamp_fp8 != clamp_fp8).float().mean().item() + assert diff_rate > 0.01, f"Expected clamp to change outputs, diff_rate={diff_rate}" + + # Clamped scales should be smaller (values are bounded by limit) + assert clamp_scale.max().item() <= no_clamp_scale.max().item() + 1e-6 + + +def test_fused_silu_mul_quant_fp8_swiglu_limit_zero_is_noop(): + """swiglu_limit=0.0 should behave identically to no clamp (default).""" + torch.manual_seed(42) + x = torch.randn(32, 2 * 2048, device="cuda", dtype=torch.bfloat16) * 5.0 + + default_fp8, default_scale = fused_silu_mul_quant_fp8(x, G) + zero_fp8, zero_scale = fused_silu_mul_quant_fp8(x, G, swiglu_limit=0.0) + + assert torch.equal(default_fp8, zero_fp8), "swiglu_limit=0 should be a no-op" + assert torch.equal(default_scale, zero_scale), "scales should match" + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v", "-s"]))