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 8, 17:45 UTC
Fix RoutingMethodType logic (#33919)
Signed-off-by: Dimitrios Bariamis <12195802+dbari@users.noreply.github.com> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Dimitrios Bariamis <12195802+dbari@users.noreply.github.com> Co-authored-by: mgoin <mgoin64@gmail.com>
Dimitrios Bariamis committed
Feb 6, 2026 at 23:03 UTC
207c3a0c20c0fee6f375e37e067074b95053094f
9 files changed
+61
-20
docker/Dockerfile
+1
-1
@@ -582,7 +582,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
582
# This is ~1.1GB and only changes when FlashInfer version bumps
583
# https://docs.flashinfer.ai/installation.html
584
# From versions.json: .flashinfer.version
585
-ARG FLASHINFER_VERSION=0.6.2
585
+ARG FLASHINFER_VERSION=0.6.3
586
RUN --mount=type=cache,target=/root/.cache/uv \
587
uv pip install --system flashinfer-cubin==${FLASHINFER_VERSION} \
588
&& uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
docker/Dockerfile.nightly_torch
+2
-2
@@ -217,13 +217,13 @@ RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2.
217
218
219
# build flashinfer for torch nightly from source around 10 mins
220
-# release version: v0.6.2
220
+# release version: v0.6.3
221
# todo(elainewy): cache flashinfer build result for faster build
222
ENV CCACHE_DIR=/root/.cache/ccache
223
RUN --mount=type=cache,target=/root/.cache/ccache \
224
--mount=type=cache,target=/root/.cache/uv \
225
echo "git clone flashinfer..." \
226
- && git clone --depth 1 --branch v0.6.2 --recursive https://github.com/flashinfer-ai/flashinfer.git \
226
+ && git clone --depth 1 --branch v0.6.3 --recursive https://github.com/flashinfer-ai/flashinfer.git \
227
&& cd flashinfer \
228
&& git submodule update --init --recursive \
229
&& echo "finish git clone flashinfer..." \
docker/versions.json
+1
-1
@@ -68,7 +68,7 @@
68
"default": "true"
69
},
70
"FLASHINFER_VERSION": {
71
- "default": "0.6.2"
71
+ "default": "0.6.3"
72
},
73
"GDRCOPY_CUDA_VERSION": {
74
"default": "12.8"
requirements/cuda.txt
+1
-1
@@ -10,4 +10,4 @@ torchaudio==2.9.1
10
# These must be updated alongside torch
11
torchvision==0.24.1 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
12
# FlashInfer should be updated together with the Dockerfile
13
-flashinfer-python==0.6.2
13
+flashinfer-python==0.6.3
vllm/model_executor/layers/fused_moe/config.py
+17
@@ -124,6 +124,23 @@ class RoutingMethodType(IntEnum):
124
Unspecified = 8.0
125
126
127
+def get_routing_method_type(
128
+ scoring_func: str, top_k: int, renormalize: bool
129
+) -> RoutingMethodType:
130
+ if scoring_func == "sigmoid":
131
+ if top_k == 1:
132
+ return RoutingMethodType.Llama4
133
+ else:
134
+ return RoutingMethodType.DeepSeekV3
135
+ elif scoring_func == "softmax":
136
+ if renormalize:
137
+ return RoutingMethodType.Renormalize
138
+ else:
139
+ return RoutingMethodType.Default
140
+ else:
141
+ return RoutingMethodType.Unspecified
142
+
143
+
144
@dataclass
145
class FusedMoEQuantDesc:
146
"""
vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py
+9
-4
@@ -61,6 +61,8 @@ def _supports_routing_method(
61
routing_method: RoutingMethodType,
62
) -> bool:
63
"""Monolithic kernels need to express router support."""
64
+ # NOTE(dbari): TopK routing could also be enabled, but need to validate models
65
+ # NOTE(dbari): Default is not implemented and should not be enabled until it is
66
if (weight_key, activation_key) == (kFp8Static128BlockSym, kFp8Dynamic128Sym):
67
# NOTE(rob): potentially allow others here. This is a conservative list.
68
return routing_method in [
@@ -72,10 +74,8 @@ def _supports_routing_method(
74
# NOTE(dbari): as above, potentially allow others here.
75
return routing_method in [
76
RoutingMethodType.Llama4,
75
- # NOTE(mgoin): Disabled to investigate accuracy issues.
76
- # See https://github.com/vllm-project/vllm/issues/33532
77
- # RoutingMethodType.Renormalize,
78
- # RoutingMethodType.RenormalizeNaive,
77
+ RoutingMethodType.Renormalize,
78
+ RoutingMethodType.RenormalizeNaive,
79
]
80
else:
81
raise ValueError("Unsupported quantization scheme.")
@@ -302,6 +302,8 @@ def fi_trtllm_fp8_per_tensor_moe(
302
per_act_token_quant=False,
303
)
304
305
+ from flashinfer.fused_moe.core import ActivationType
306
+
307
from vllm.utils.flashinfer import flashinfer_trtllm_fp8_per_tensor_scale_moe
308
309
return flashinfer_trtllm_fp8_per_tensor_scale_moe(
@@ -323,6 +325,9 @@ def fi_trtllm_fp8_per_tensor_moe(
325
routed_scaling_factor=routed_scaling_factor,
326
use_routing_scales_on_input=use_routing_scales_on_input,
327
routing_method_type=routing_method_type,
328
+ # TODO: Required for flashinfer==0.6.3, remove with update
329
+ # https://github.com/flashinfer-ai/flashinfer/pull/2508
330
+ activation_type=ActivationType.Swiglu,
331
)
332
333
vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
+8
-5
@@ -10,7 +10,10 @@ from vllm.distributed.eplb.eplb_state import EplbLayerState
10
from vllm.model_executor.layers.batch_invariant import (
11
vllm_is_batch_invariant,
12
)
13
-from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
13
+from vllm.model_executor.layers.fused_moe.config import (
14
+ RoutingMethodType,
15
+ get_routing_method_type,
16
+)
17
from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
18
19
@@ -158,10 +161,10 @@ class FusedTopKBiasRouter(BaseRouter):
161
162
@property
163
def routing_method_type(self) -> RoutingMethodType:
161
- return (
162
- RoutingMethodType.Renormalize
163
- if not self.renormalize
164
- else RoutingMethodType.RenormalizeNaive
164
+ return get_routing_method_type(
165
+ scoring_func=self.scoring_func,
166
+ top_k=self.top_k,
167
+ renormalize=self.renormalize,
168
)
169
170
def _compute_routing(
vllm/model_executor/layers/fused_moe/router/fused_topk_router.py
+8
-5
@@ -7,7 +7,10 @@ import torch
7
import vllm._custom_ops as ops
8
from vllm._aiter_ops import rocm_aiter_ops
9
from vllm.distributed.eplb.eplb_state import EplbLayerState
10
-from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
10
+from vllm.model_executor.layers.fused_moe.config import (
11
+ RoutingMethodType,
12
+ get_routing_method_type,
13
+)
14
from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
15
16
@@ -135,10 +138,10 @@ class FusedTopKRouter(BaseRouter):
138
139
@property
140
def routing_method_type(self) -> RoutingMethodType:
138
- return (
139
- RoutingMethodType.Renormalize
140
- if not self.renormalize
141
- else RoutingMethodType.RenormalizeNaive
141
+ return get_routing_method_type(
142
+ scoring_func=self.scoring_func,
143
+ top_k=self.top_k,
144
+ renormalize=self.renormalize,
145
)
146
147
def _compute_routing(
vllm/model_executor/layers/fused_moe/router/router_factory.py
+14
-1
@@ -6,6 +6,7 @@ import torch
6
7
import vllm.envs as envs
8
from vllm.distributed.eplb.eplb_state import EplbLayerState
9
+from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
10
from vllm.model_executor.layers.fused_moe.router.custom_routing_router import (
11
CustomRoutingRouter,
12
)
@@ -106,7 +107,7 @@ def create_fused_moe_router(
107
"num_expert_group and topk_group must be provided when "
108
"use_grouped_topk is True"
109
)
109
- return GroupedTopKRouter(
110
+ grouped_topk_router = GroupedTopKRouter(
111
top_k=top_k,
112
global_num_experts=global_num_experts,
113
eplb_state=eplb_state,
@@ -120,6 +121,18 @@ def create_fused_moe_router(
121
enable_eplb=enable_eplb,
122
indices_type_getter=indices_type_getter,
123
)
124
+ if (
125
+ grouped_topk_router.routing_method_type != RoutingMethodType.Unspecified
126
+ or num_expert_group > 1
127
+ or topk_group > 1
128
+ ):
129
+ return grouped_topk_router
130
+
131
+ # If routing_method for GroupedTopKRouter is Unspecified and there is only
132
+ # one group, fallback to standard top-k routing
133
+ use_grouped_topk = False
134
+ num_expert_group = None
135
+ topk_group = None
136
137
if custom_routing_function is not None:
138
return CustomRoutingRouter(