[Perf][MoE] Skip top-k slots routed to non-local experts in TritonExp… (#58051)
Signed-off-by: Shuolei Wang <shuoleiwang123@gmail.com> Signed-off-by: Shuolei Wang <948904026@qq.com>
Shuolei Wang committed
Sep 23, 2026 at 21:03 UTC
9ed81581a896a5bf0448b26d43f52486a435f799
2 files changed
+95
-4
tests/kernels/moe/test_block_fp8.py
+61
-1
@@ -107,7 +107,16 @@ SEEDS = [0]
107
108
109
def torch_w8a8_block_fp8_moe(
110
- a, w1, w2, w1_s, w2_s, topk_weight, topk_ids, block_shape, silu_fp32=False
110
+ a,
111
+ w1,
112
+ w2,
113
+ w1_s,
114
+ w2_s,
115
+ topk_weight,
116
+ topk_ids,
117
+ block_shape,
118
+ silu_fp32=False,
119
+ expert_map=None,
120
):
121
"""Fused MoE with block-wise fp8 quantization using native torch.
122
@@ -123,6 +132,9 @@ def torch_w8a8_block_fp8_moe(
132
133
topk_weight = topk_weight.view(-1)
134
topk_ids = topk_ids.view(-1)
135
+ if expert_map is not None:
136
+ # w1/w2 hold this rank's experts; slots mapped to -1 contribute nothing.
137
+ topk_ids = expert_map[topk_ids]
138
139
_, block_k = block_shape[0], block_shape[1]
140
# Quantize with the production per-token-group fp8 kernel (same HIP/CUDA op the
@@ -301,6 +313,54 @@ def test_w8a8_block_fp8_clamp_skips_fused_silu_quant(
313
)
314
315
316
+@pytest.mark.parametrize(("M", "N", "K"), [(83, 512, 512), (2048, 4608, 512)])
317
+@pytest.mark.parametrize("topk", [2, 6])
318
+@pytest.mark.parametrize("ep_size", [2, 4])
319
+@torch.inference_mode()
320
+def test_w8a8_block_fp8_fused_moe_expert_map(M, N, K, topk, ep_size, workspace_init):
321
+ """One EP rank: local expert shard plus an expert_map with -1 entries."""
322
+ E, block_size, dtype = 16, [128, 128], torch.bfloat16
323
+ torch.manual_seed(0)
324
+
325
+ a = torch.randn((M, K), dtype=dtype) / 10
326
+ score = torch.randn((M, E), dtype=dtype)
327
+ w1, w2, quant_config = make_test_quant_config(
328
+ E, N, K, dtype, quant_dtype=torch.float8_e4m3fn, block_shape=block_size
329
+ )
330
+
331
+ local_e = E // ep_size
332
+ e_ids = torch.randperm(E, dtype=torch.int32)[:local_e]
333
+ e_map = torch.full((E,), -1, dtype=torch.int32)
334
+ e_map[e_ids] = torch.arange(local_e, dtype=torch.int32)
335
+ w1, w2 = w1[e_ids].contiguous(), w2[e_ids].contiguous()
336
+ w1_s = quant_config.w1_scale[e_ids].contiguous()
337
+ w2_s = quant_config.w2_scale[e_ids].contiguous()
338
+ quant_config = fp8_w8a8_moe_quant_config(
339
+ w1_scale=w1_s, w2_scale=w2_s, block_shape=block_size
340
+ )
341
+
342
+ m_fused_moe = modular_triton_fused_moe(make_dummy_moe_config(), quant_config)
343
+ topk_weights, topk_ids, _ = fused_topk(a, score.float(), topk, False)
344
+
345
+ with set_current_vllm_config(vllm_config):
346
+ ref_out = torch_w8a8_block_fp8_moe(
347
+ a, w1, w2, w1_s, w2_s, topk_weights, topk_ids, block_size, expert_map=e_map
348
+ )
349
+ m_out = m_fused_moe.apply(
350
+ a,
351
+ w1,
352
+ w2,
353
+ topk_weights,
354
+ topk_ids,
355
+ activation=MoEActivation.SILU,
356
+ apply_router_weight_on_input=False,
357
+ expert_map=e_map,
358
+ global_num_experts=E,
359
+ )
360
+
361
+ torch.testing.assert_close(m_out, ref_out, atol=0.035, rtol=0.035)
362
+
363
+
364
@pytest.mark.parametrize(("M", "N", "K"), MNK_FACTORS_DG)
365
@pytest.mark.parametrize("E", E)
366
@pytest.mark.parametrize("topk", TOP_KS)
vllm/model_executor/layers/fused_moe/experts/triton_moe.py
+34
-3
@@ -337,6 +337,25 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
337
)
338
intermediate_cache3 = _resize_cache(workspace2, (num_tokens, top_k_num, K))
339
340
+ # Under EP, drop the top-k slots routed to remote experts at alignment
341
+ # time instead of launching `off_experts == -1` GEMM blocks that only
342
+ # write zeros; the pad-aware `moe_sum` below skips the same slots.
343
+ # Their rows in the workspace caches are then never written, so every
344
+ # consumer of those caches must be row-local: a dynamic per-tensor
345
+ # scale for the second GEMM's activation would take its amax over the
346
+ # untouched rows, and LoRA aligns the full `topk_ids` itself.
347
+ a2_scale_is_global = (
348
+ self.quant_dtype is not None
349
+ and a2_scale is None
350
+ and self.block_shape is None
351
+ and (not self.per_act_token_quant or self.quantization_emulation)
352
+ )
353
+ skip_invalid = (
354
+ expert_map is not None
355
+ and self._lora_context is None
356
+ and not a2_scale_is_global
357
+ )
358
+
359
# Include fused shared-expert rows while preserving EP remapping.
360
num_align_experts = w1.shape[0] if expert_map is None else global_num_experts
361
sorted_token_ids, expert_ids, num_tokens_post_padded = (
@@ -350,6 +369,7 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
369
use_int8_w8a16=self.quant_config.use_int8_w8a16,
370
use_int4_w4a16=self.quant_config.use_int4_w4a16,
371
block_shape=self.block_shape,
372
+ ignore_invalid_experts=skip_invalid,
373
)
374
)
375
@@ -585,10 +605,21 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
605
)
606
607
# separate function is required for MoE + LoRA
588
- self.moe_sum(intermediate_cache3, output)
608
+ self.moe_sum(intermediate_cache3, output, topk_ids, expert_map)
609
590
- def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None:
591
- ops.moe_sum(input, output)
610
+ def moe_sum(
611
+ self,
612
+ input: torch.Tensor,
613
+ output: torch.Tensor,
614
+ topk_ids: torch.Tensor | None = None,
615
+ expert_map: torch.Tensor | None = None,
616
+ ) -> None:
617
+ if expert_map is not None:
618
+ # Skip the slots whose expert is not on this rank: the rows the
619
+ # alignment dropped, or the zeros the `-1` blocks wrote.
620
+ ops.moe_sum(input, output, topk_ids, expert_map)
621
+ else:
622
+ ops.moe_sum(input, output)
623
624
625
class TritonWNA16Experts(TritonExperts):