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 15, 11:00 UTC
[Bugfix] Fix marlin moe fallback logic for llama4 (#18042)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
May 13, 2025 at 03:53 UTC
ea6ae8cb45f7c1f8bbe76e1166893adae43881ae
3 files changed
+12
-5
tests/weight_loading/models-large.txt
+2
-1
@@ -4,4 +4,5 @@ compressed-tensors, nm-testing/Mixtral-8x7B-Instruct-v0.1-W8A16-quantized, main
4
compressed-tensors, nm-testing/test-w4a16-mixtral-actorder-group, main
5
gptq_marlin, TheBloke/Mixtral-8x7B-v0.1-GPTQ, main
6
gptq_marlin, TheBloke/Mixtral-8x7B-v0.1-GPTQ, gptq-8bit-128g-actorder_True
7
-awq_marlin, casperhansen/deepseek-coder-v2-instruct-awq, main
\ No newline at end of file
7
+awq_marlin, casperhansen/deepseek-coder-v2-instruct-awq, main
8
+compressed-tensors, RedHatAI/Llama-4-Scout-17B-16E-Instruct-quantized.w4a16, main
\ No newline at end of file
vllm/model_executor/layers/fused_moe/layer.py
+1
-1
@@ -480,6 +480,7 @@ class FusedMoE(torch.nn.Module):
480
self.custom_routing_function = custom_routing_function
481
self.scoring_func = scoring_func
482
self.e_score_correction_bias = e_score_correction_bias
483
+ self.apply_router_weight_on_input = apply_router_weight_on_input
484
self.activation = activation
485
486
if self.scoring_func != "softmax" and not self.use_grouped_topk:
@@ -498,7 +499,6 @@ class FusedMoE(torch.nn.Module):
499
self.quant_method = quant_config.get_quant_method(self, prefix)
500
assert self.quant_method is not None
501
501
- self.apply_router_weight_on_input = apply_router_weight_on_input
502
moe_quant_params = {
503
"num_experts": self.local_num_experts,
504
"hidden_size": hidden_size,
vllm/model_executor/layers/quantization/utils/marlin_utils.py
+9
-3
@@ -171,13 +171,19 @@ def check_moe_marlin_supports_layer(layer: LinearBase, group_size: int) \
171
-> bool:
172
hidden_size = layer.hidden_size
173
intermediate_size_per_partition = layer.intermediate_size_per_partition
174
+ # apply_router_weight_on_input is not supported for moe marlin
175
+ supports_router_weight = not layer.apply_router_weight_on_input
176
+ # moe marlin requires the activation to be silu
177
+ supports_activation = layer.activation == "silu"
178
179
# gate-up: (n, k) = (intermediate_size_per_partition * 2, hidden_size)
180
# down: (n, k) = (hidden_size, intermediate_size_per_partition)
181
# moe marlin requires n % 128 == 0 and k % 64 == 0
178
- return hidden_size % 128 == 0 and \
179
- intermediate_size_per_partition % max(64, group_size) == 0 and \
180
- group_size in [-1, 32, 64, 128]
182
+ supports_shape = hidden_size % 128 == 0 and \
183
+ intermediate_size_per_partition % max(64, group_size) == 0
184
+ supports_group_size = group_size in [-1, 32, 64, 128]
185
+ return supports_shape and supports_group_size and \
186
+ supports_router_weight and supports_activation
187
188
189
def marlin_make_workspace(output_size_per_partition: int,