Revert "[kernel] add fused silu mul quant fp8" (#38381)
This commit is contained in:
Liangsheng Yin
@@ -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()
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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"]))
|
||||
Reference in New Issue
Block a user