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, 18:15 UTC
[Bugfix] Fix FP8 MoE EP Weight Loading for ModelOpt Llama4 (#32886)
Signed-off-by: baonudesifeizhai <baonudesifeizhai@gmail.com>
baonudesifeizhai committed
Jan 23, 2026 at 10:31 UTC
1fb648bf107e9c24e7f99671c77406d126506d7a
1 file changed
+21
-1
vllm/model_executor/models/llama4.py
+21
-1
@@ -51,6 +51,8 @@ from vllm.model_executor.model_loader.weight_utils import (
51
)
52
from vllm.model_executor.models.interfaces import MixtureOfExperts
53
from vllm.model_executor.models.utils import sequence_parallel_chunk
54
+from vllm.platforms import current_platform
55
+from vllm.utils.torch_utils import is_torch_equal_or_newer
56
57
from .llama import LlamaForCausalLM, LlamaMLP, LlamaModel
58
from .utils import (
@@ -504,7 +506,25 @@ class Llama4Model(LlamaModel):
506
.flatten()
507
.to(new_loaded_weight.device)
508
)
507
- new_loaded_weight = new_loaded_weight[local_expert_indices]
509
+ # Workaround for FP8 CPU indexing on older PyTorch:
510
+ # https://github.com/vllm-project/vllm/issues/32862
511
+ is_fp8_dtype = new_loaded_weight.dtype == (
512
+ current_platform.fp8_dtype()
513
+ ) or (
514
+ new_loaded_weight.dtype.is_floating_point
515
+ and new_loaded_weight.element_size() == 1
516
+ )
517
+ if (
518
+ new_loaded_weight.device.type == "cpu"
519
+ and is_fp8_dtype
520
+ and not is_torch_equal_or_newer("2.11.0")
521
+ ):
522
+ # PyTorch < 2.11 doesn't support CPU float8 indexing.
523
+ new_loaded_weight = new_loaded_weight.to(torch.float16)[
524
+ local_expert_indices
525
+ ].to(new_loaded_weight.dtype)
526
+ else:
527
+ new_loaded_weight = new_loaded_weight[local_expert_indices]
528
expert_id = local_expert_indices[0].item()
529
else:
530
# TODO: add EP support for non fused weights