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 9, 07:30 UTC
[Quantization][ROCm] Fix MoE weight loading to be robust (Qwen3_MoE/Qwen3_next as example models) (#33173)
Signed-off-by: xuebwang-amd <xuebwang@amd.com>
xuebwang-amd committed
Jan 31, 2026 at 01:50 UTC
f451b4558b2bb42dafcdd19f7b5c0fc58194af5a
1 file changed
+6
-2
vllm/model_executor/layers/fused_moe/layer.py
+6
-2
@@ -996,7 +996,9 @@ class FusedMoE(CustomOp):
996
shard_size = expert_data.shape[shard_dim] // 2
997
else:
998
shard_size = expert_data.shape[shard_dim]
999
- if not load_full:
999
+ # Only narrow if the loaded_weight is not a scalar (0-dim tensor)
1000
+ # and we're not loading the full weight
1001
+ if not load_full and loaded_weight.ndim > 0:
1002
loaded_weight = loaded_weight.narrow(
1003
shard_dim, shard_size * tp_rank, shard_size
1004
)
@@ -1022,7 +1024,9 @@ class FusedMoE(CustomOp):
1024
# down_proj: "RowParallel" so tp sharding on input_dim
1025
# Narrow parameter and load.
1026
shard_size = expert_data.shape[shard_dim]
1025
- if not load_full:
1027
+ # Only narrow if the loaded_weight is not a scalar (0-dim tensor)
1028
+ # and we're not loading the full weight
1029
+ if not load_full and loaded_weight.ndim > 0:
1030
loaded_weight = loaded_weight.narrow(
1031
shard_dim, shard_size * tp_rank, shard_size
1032
)