From 4dcecc78916b00837b0ff951caf7f9bcd611e350 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Mon, 7 Sep 2026 17:36:10 -0700 Subject: [PATCH] Revert "[kernel] add fused silu mul quant fp8" (#38381) --- .../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, 7 insertions(+), 695 deletions(-) delete mode 100644 python/sglang/kernels/aot/benchmark/bench_silu_quant_fp8.py delete 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 deleted file mode 100644 index 81b27207e..000000000 --- a/python/sglang/kernels/aot/benchmark/bench_silu_quant_fp8.py +++ /dev/null @@ -1,112 +0,0 @@ -"""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 c98f25dfc..e90e0fb08 100644 --- a/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py +++ b/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py @@ -1177,102 +1177,6 @@ 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 564c70316..ce665e936 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,7 +17,6 @@ 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, @@ -126,41 +125,6 @@ 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, @@ -618,26 +582,6 @@ 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) @@ -716,24 +660,14 @@ def _fused_moe_kernel_sequence( ) if not fuse_swiglu_interleaved: - if not use_fused_silu_mul_quant_fp8: - intermediate_cache2 = torch.empty( - (total_tokens, N // 2), - device=hidden_states.device, - dtype=hidden_states.dtype, - ) + intermediate_cache2 = torch.empty( + (total_tokens, N // 2), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) # Activation function with multiplication - 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: + if fuse_swiglu_interleaved: # silu(gate) * up was already applied by the up-GEMM epilogue. pass elif activation == "silu" and is_gated: @@ -900,7 +834,7 @@ def _fused_moe_kernel_sequence( else out_hidden_states.unsqueeze(0) ) ), - fused_a2_scale if use_fused_silu_mul_quant_fp8 else a2_scale, + 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 deleted file mode 100644 index e0b6302f9..000000000 --- a/test/registered/kernels/ops/quantization/test_fused_silu_mul_quant_fp8.py +++ /dev/null @@ -1,414 +0,0 @@ -"""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"]))