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 14, 05:30 UTC
[UX] Separate marlin moe config logic from triton moe (#23006)
Michael Goin committed
Aug 16, 2025 at 22:16 UTC
94096a47c92c4a53ad44cfffdca918669c0f89e0
2 files changed
+7
-22
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
+6
-14
@@ -1,14 +1,12 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Fused MoE utilities for GPTQ."""
4
-import functools
4
from typing import Optional
5
6
import torch
7
8
import vllm._custom_ops as ops
10
-from vllm.model_executor.layers.fused_moe.fused_moe import (
11
- moe_align_block_size, try_get_optimal_moe_config)
9
+from vllm.model_executor.layers.fused_moe.fused_moe import moe_align_block_size
10
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
11
marlin_make_workspace_new, maybe_warn_marlin_atomic_add)
12
from vllm.scalar_type import ScalarType, scalar_types
@@ -98,17 +96,11 @@ def fused_marlin_moe(hidden_states: torch.Tensor,
96
N = w2.shape[1] * 16
97
topk = topk_ids.shape[1]
98
101
- get_config_func = functools.partial(
102
- try_get_optimal_moe_config,
103
- w1.shape,
104
- w2.shape,
105
- topk_ids.shape[1],
106
- None,
107
- is_marlin=True,
108
- )
109
- config = get_config_func(M)
110
-
111
- block_size_m = config["BLOCK_SIZE_M"]
99
+ # M block size selection logic
100
+ # TODO: tune this further for specific models
101
+ for block_size_m in [8, 16, 32, 48, 64]:
102
+ if M * topk / E / block_size_m < 0.9:
103
+ break
104
105
if global_num_experts == -1:
106
global_num_experts = E
vllm/model_executor/layers/fused_moe/fused_moe.py
+1
-8
@@ -801,7 +801,6 @@ def get_default_config(
801
K: int,
802
topk: int,
803
dtype: Optional[str],
804
- is_marlin: bool,
804
block_shape: Optional[list[int]] = None,
805
) -> dict[str, int]:
806
if dtype == "fp8_w8a8" and block_shape is not None:
@@ -832,11 +831,6 @@ def get_default_config(
831
config = {"BLOCK_SIZE_M": 32, "GROUP_SIZE_M": 1}
832
else:
833
config = {"BLOCK_SIZE_M": 64, "GROUP_SIZE_M": 1}
835
- elif is_marlin:
836
- for block_size_m in [8, 16, 32, 48, 64]:
837
- if M * topk / E / block_size_m < 0.9:
838
- break
839
- return {"BLOCK_SIZE_M": block_size_m}
834
elif M <= E:
835
config = {
836
"BLOCK_SIZE_M": 16,
@@ -860,7 +854,6 @@ def try_get_optimal_moe_config(
854
top_k: int,
855
dtype: Optional[str],
856
M: int,
863
- is_marlin: bool = False,
857
block_shape: Optional[list[int]] = None,
858
) -> dict[str, int]:
859
from vllm.model_executor.layers.fused_moe import get_config
@@ -883,7 +876,7 @@ def try_get_optimal_moe_config(
876
else:
877
# Else use the default config
878
config = get_default_config(M, E, N, w1_shape[2], top_k, dtype,
886
- is_marlin, block_shape)
879
+ block_shape)
880
return config
881
882