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
[Hardware/NVIDIA/Modelopt] Fix modelopt forward method for v1 torch.compile (#18101)
Signed-off-by: Pavani Majety <pmajety@nvidia.com>
Pavani Majety committed
May 13, 2025 at 19:33 UTC
65f0f74b6686103cbe6300fc576ed60d3ab0411e
2 files changed
+14
-10
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+5
-2
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
""" CUTLASS based Fused MoE kernels."""
3
+import os
4
from typing import Optional
5
6
import torch
@@ -183,7 +184,8 @@ def cutlass_moe_fp8(
184
185
FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max()
186
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
186
-MAX_TOKENS_PER_EXPERT = 65536
187
+MAX_TOKENS_PER_EXPERT = int(
188
+ os.environ.get('VLLM_MODELOPT_MAX_TOKENS_PER_EXPERT', '65536'))
189
190
191
def cutlass_moe_fp4(a: torch.Tensor, a1_gscale: torch.Tensor,
@@ -243,7 +245,8 @@ def cutlass_moe_fp4(a: torch.Tensor, a1_gscale: torch.Tensor,
245
== m), ("topk must be provided for each row of a")
246
assert (m <= MAX_TOKENS_PER_EXPERT), (
247
f"m must be less than MAX_TOKENS_PER_EXPERT({MAX_TOKENS_PER_EXPERT})"
246
- f" for cutlass_moe_fp4, observed m = {m}")
248
+ f" for cutlass_moe_fp4, observed m = {m}. Use"
249
+ f" VLLM_MODELOPT_MAX_TOKENS_PER_EXPERT to set this value.")
250
out_dtype = a.dtype
251
num_topk = topk_ids.shape[1]
252
vllm/model_executor/layers/quantization/modelopt.py
+9
-8
@@ -401,6 +401,7 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
401
402
layer.weight_scale_swizzled = Parameter(swizzled_weight_scale,
403
requires_grad=False)
404
+ layer.weight = Parameter(layer.weight.data, requires_grad=False)
405
406
if self.use_marlin:
407
prepare_fp4_layer_for_marlin(layer)
@@ -426,11 +427,7 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
427
bias=bias)
428
429
output_dtype = x.dtype
429
-
430
- # for input only the contracting dimension has a constraint.
431
- x_m, _ = x.shape
432
- w_n, _ = layer.weight.shape
433
- output_shape = [x_m, w_n]
430
+ output_shape = [x.shape[0], layer.weight.shape[0]]
431
432
# quantize BF16 or FP16 to (FP4 and interleaved block scale)
433
s_quant = 1 / layer.input_scale
@@ -586,11 +583,11 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
583
if scale_ndim == 2 else swizzled_scale.reshape(B, M, K))
584
585
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
589
- # GEMM 1
586
587
+ # GEMM 1
588
assert torch.allclose(
589
layer.w13_weight_scale_2[:, 0], layer.w13_weight_scale_2[:, 1]), (
593
- "Expected w1_weight_scale_2 to equal w3_weight_scale_2")
590
+ "w1_weight_scale_2 must match w3_weight_scale_2")
591
592
w13_weight_scale_2 = layer.w13_weight_scale_2[:, 0]
593
layer.w13_weight_scale_2 = Parameter(w13_weight_scale_2,
@@ -616,6 +613,9 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
613
layer.w13_input_scale_quant = Parameter(
614
(1 / w13_input_scale).to(torch.float32), requires_grad=False)
615
616
+ layer.w13_weight = Parameter(layer.w13_weight.data,
617
+ requires_grad=False)
618
+
619
# GEMM 2
620
layer.g2_alphas = Parameter(
621
(layer.w2_input_scale * layer.w2_weight_scale_2).to(torch.float32),
@@ -633,6 +633,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
633
634
layer.w2_blockscale_swizzled = Parameter(w2_blockscale_swizzled,
635
requires_grad=False)
636
+ layer.w2_weight = Parameter(layer.w2_weight.data, requires_grad=False)
637
638
if self.use_marlin:
639
prepare_moe_fp4_layer_for_marlin(layer)
@@ -694,7 +695,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
695
assert not apply_router_weight_on_input, (
696
"Router weight on input is not "
697
"supported for ModelOptNvFp4FusedMoE.")
697
- assert expert_map is None, ("Expert Parallelism /expert_map "
698
+ assert expert_map is None, ("Expert Parallelism / expert_map "
699
"is currently not supported for "
700
"ModelOptNvFp4FusedMoE.")
701