[Bugfix][Kernel] Skip the fused silu-mul block-quant fast path when a swiglu clamp is set (#57984)
Signed-off-by: Garrett Goon <garrett@primeintellect.ai> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Garrett Goon committed
Sep 21, 2026 at 18:38 UTC
8bf051cd51adff79e1bc94e910adaea543aaa2e1
2 files changed
+58
tests/kernels/moe/test_block_fp8.py
+55
@@ -25,6 +25,7 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
25
from vllm.model_executor.layers.fused_moe.config import (
26
fp8_w8a8_moe_quant_config,
27
)
28
+from vllm.model_executor.layers.fused_moe.experts import triton_moe
29
from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import (
30
_valid_deep_gemm_shape,
31
)
@@ -246,6 +247,60 @@ def test_w8a8_block_fp8_fused_moe(
247
torch.testing.assert_close(m_out, ref_out_m, atol=tol, rtol=tol)
248
249
250
+@pytest.mark.parametrize("clamp_limit", [None, 10.0])
251
+@torch.inference_mode()
252
+def test_w8a8_block_fp8_clamp_skips_fused_silu_quant(
253
+ clamp_limit, monkeypatch, workspace_init, disable_deepgemm_ue8m0
254
+):
255
+ """A SwiGLU clamp must bypass the fused SiLU+quant kernel, which cannot clamp."""
256
+ torch.manual_seed(0)
257
+ M, N, K, E, topk = 16, 256, 512, 8, 2
258
+ block_size = [128, 128]
259
+
260
+ a = torch.randn((M, K), dtype=torch.bfloat16) / 10
261
+ score = torch.randn((M, E), dtype=torch.bfloat16)
262
+ (_, w1, w1_s, _), (_, w2, w2_s, _) = make_test_weights(
263
+ E, N, K, torch.bfloat16, torch.float8_e4m3fn, block_shape=block_size
264
+ )
265
+ quant_config = fp8_w8a8_moe_quant_config(
266
+ w1_scale=w1_s,
267
+ w2_scale=w2_s,
268
+ block_shape=block_size,
269
+ gemm1_clamp_limit=clamp_limit,
270
+ )
271
+ m_fused_moe = modular_triton_fused_moe(make_dummy_moe_config(), quant_config)
272
+ topk_weights, topk_ids, _ = fused_topk(a, score.float(), topk, False)
273
+
274
+ fused_calls = 0
275
+ fused_op = triton_moe.ops.silu_and_mul_per_block_quant
276
+
277
+ def counting_fused_op(*args, **kwargs):
278
+ nonlocal fused_calls
279
+ fused_calls += 1
280
+ return fused_op(*args, **kwargs)
281
+
282
+ monkeypatch.setattr(
283
+ triton_moe.ops, "silu_and_mul_per_block_quant", counting_fused_op
284
+ )
285
+
286
+ with set_current_vllm_config(vllm_config):
287
+ m_fused_moe.apply(
288
+ a,
289
+ w1,
290
+ w2,
291
+ topk_weights,
292
+ topk_ids,
293
+ activation=MoEActivation.SILU,
294
+ apply_router_weight_on_input=False,
295
+ expert_map=None,
296
+ global_num_experts=E,
297
+ )
298
+
299
+ assert (fused_calls > 0) == (clamp_limit is None), (
300
+ f"fused fast path ran {fused_calls} times with clamp_limit={clamp_limit}"
301
+ )
302
+
303
+
304
@pytest.mark.parametrize(("M", "N", "K"), MNK_FACTORS_DG)
305
@pytest.mark.parametrize("E", E)
306
@pytest.mark.parametrize("topk", TOP_KS)
vllm/model_executor/layers/fused_moe/experts/triton_moe.py
+3
@@ -479,8 +479,11 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
479
# Fuse SiLU+Mul + FP8 block quantize into a single kernel
480
# when conditions permit (gated SiLU, fp8 block quant with
481
# group_size=128, no LoRA requiring the BF16 intermediate).
482
+ # The fused kernel has no clamp parameter, so a configured
483
+ # SwiGLU clamp limit falls through to the unfused path.
484
if (
485
activation == MoEActivation.SILU
486
+ and self.activation_config.clamp_limit is None
487
and self.quant_config.use_fp8_w8a8
488
and self.block_shape == [128, 128]
489
and lora_context is None