A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 21:15 UTC
[MoE] Triton MoE Perf regression - restore low latency path (#39016)
milesial committed
Apr 20, 2026 at 23:37 UTC
257015d5e50029e8ebcbce4fe0c83f43b531552e
1 file changed
+72
-31
vllm/model_executor/layers/fused_moe/fused_moe.py
+72
-31
@@ -1551,6 +1551,55 @@ def dispatch_fused_experts_func(inplace: bool) -> Callable[..., torch.Tensor]:
1551
return torch_vllm_outplace_fused_experts
1552
1553
1554
+def _prepare_expert_assignment(
1555
+ topk_ids: torch.Tensor,
1556
+ config: dict[str, Any],
1557
+ num_tokens: int,
1558
+ top_k_num: int,
1559
+ global_num_experts: int,
1560
+ expert_map: torch.Tensor | None,
1561
+ *,
1562
+ use_int8_w8a16: bool = False,
1563
+ use_int4_w4a16: bool = False,
1564
+ block_shape: list[int] | None = None,
1565
+ ignore_invalid_experts: bool = False,
1566
+) -> tuple[torch.Tensor | None, torch.Tensor, torch.Tensor]:
1567
+ """Prepare expert assignments for the aligned and low-latency Triton paths."""
1568
+ # SPARSITY_FACTOR is a heuristic margin ensuring tokens_in_chunk * top_k
1569
+ # activates only a small fraction of total experts
1570
+ # Skips moe_align_block_size and activates the `sorted_token_ids is None`
1571
+ # path of the fused_moe_kernel kernel
1572
+ naive_block_assignment = (
1573
+ expert_map is None
1574
+ and num_tokens * top_k_num * 4 <= global_num_experts
1575
+ and not (
1576
+ (use_int8_w8a16 or use_int4_w4a16)
1577
+ and block_shape is not None
1578
+ and block_shape[1] > 0
1579
+ )
1580
+ )
1581
+
1582
+ if naive_block_assignment:
1583
+ return (
1584
+ None,
1585
+ topk_ids.view(-1),
1586
+ torch.full(
1587
+ (1,),
1588
+ topk_ids.numel() * config["BLOCK_SIZE_M"],
1589
+ dtype=torch.int32,
1590
+ device=topk_ids.device,
1591
+ ),
1592
+ )
1593
+
1594
+ return moe_align_block_size(
1595
+ topk_ids,
1596
+ config["BLOCK_SIZE_M"],
1597
+ global_num_experts,
1598
+ expert_map,
1599
+ ignore_invalid_experts=ignore_invalid_experts,
1600
+ )
1601
+
1602
+
1603
# TODO (bnell): replace this with modular op. Can get rid of inplace/outplace
1604
# torch ops.
1605
def fused_experts(
@@ -1791,37 +1840,19 @@ def fused_experts_impl(
1840
ocp_mx_scheme=ocp_mx_scheme,
1841
)
1842
1794
- # SPARSITY_FACTOR is a heuristic margin ensuring num_tokens * top_k
1795
- # activates only a small fraction of total experts
1796
- SPARSITY_FACTOR = 4
1797
- # block quantized code path is not implemented yet.
1798
- naive_block_assignment = (
1799
- expert_map is None
1800
- and num_tokens * top_k_num * SPARSITY_FACTOR <= global_num_experts
1801
- and not (
1802
- (use_int8_w8a16 or use_int4_w4a16)
1803
- and block_shape is not None
1804
- and block_shape[1] > 0
1805
- )
1843
+ sorted_token_ids, expert_ids, num_tokens_post_padded = _prepare_expert_assignment(
1844
+ topk_ids,
1845
+ config,
1846
+ num_tokens,
1847
+ top_k_num,
1848
+ global_num_experts,
1849
+ expert_map,
1850
+ use_int8_w8a16=use_int8_w8a16,
1851
+ use_int4_w4a16=use_int4_w4a16,
1852
+ block_shape=block_shape,
1853
+ ignore_invalid_experts=True,
1854
)
1855
1808
- if not naive_block_assignment:
1809
- sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
1810
- topk_ids,
1811
- config["BLOCK_SIZE_M"],
1812
- global_num_experts,
1813
- expert_map,
1814
- ignore_invalid_experts=True,
1815
- )
1816
- else:
1817
- max_num_tokens_padded = topk_ids.numel() * config["BLOCK_SIZE_M"]
1818
- expert_ids = topk_ids.view(-1)
1819
- num_tokens_post_padded = torch.empty(
1820
- (1), dtype=torch.int32, device=topk_ids.device
1821
- )
1822
- num_tokens_post_padded.fill_(max_num_tokens_padded)
1823
- sorted_token_ids = None
1824
-
1856
dispatch_fused_moe_kernel(
1857
qhidden_states,
1858
w1,
@@ -2073,8 +2104,18 @@ class TritonExperts(mk.FusedMoEExpertsModular):
2104
)
2105
intermediate_cache3 = _resize_cache(workspace2, (num_tokens, top_k_num, K))
2106
2076
- sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
2077
- topk_ids, config["BLOCK_SIZE_M"], global_num_experts, expert_map
2107
+ sorted_token_ids, expert_ids, num_tokens_post_padded = (
2108
+ _prepare_expert_assignment(
2109
+ topk_ids,
2110
+ config,
2111
+ num_tokens,
2112
+ top_k_num,
2113
+ global_num_experts,
2114
+ expert_map,
2115
+ use_int8_w8a16=self.quant_config.use_int8_w8a16,
2116
+ use_int4_w4a16=self.quant_config.use_int4_w4a16,
2117
+ block_shape=self.block_shape,
2118
+ )
2119
)
2120
2121
invoke_fused_moe_triton_kernel(