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 14, 05:30 UTC
[BugFix] GPT-OSS Attention DP + MoE TP weight loading issue (#24032)
Signed-off-by: Po-Han Huang <pohanh@nvidia.com>
Po-Han Huang (NVIDIA) committed
Oct 21, 2025 at 12:03 UTC
aef368aa08572505b820db01da82e2fbb3d43a72
2 files changed
+30
-13
vllm/model_executor/layers/fused_moe/config.py
+14
-9
@@ -662,6 +662,17 @@ class FusedMoEParallelConfig:
662
def use_deepep_ll_kernels(self):
663
return self.use_all2all_kernels and self.all2all_backend == "deepep_low_latency"
664
665
+ @staticmethod
666
+ def flatten_tp_across_dp(
667
+ tp_size: int, dp_size: int, dp_rank: int
668
+ ) -> tuple[int, int]:
669
+ tp_rank = 0 if tp_size == 1 else get_tensor_model_parallel_rank()
670
+ # There are actually dp_size * tp_size devices. Update tp_size
671
+ # and tp_rank so we shard across all devices.
672
+ flatten_tp_size = dp_size * tp_size
673
+ flatten_tp_rank = dp_rank * tp_size + tp_rank
674
+ return flatten_tp_size, flatten_tp_rank
675
+
676
@staticmethod
677
def make(
678
tp_size_: int, dp_size_: int, vllm_parallel_config: ParallelConfig
@@ -737,19 +748,13 @@ class FusedMoEParallelConfig:
748
between the 4 devices.
749
"""
750
740
- def flatten_tp_across_dp(dp_rank: int):
741
- tp_rank = 0 if tp_size_ == 1 else get_tensor_model_parallel_rank()
742
- # There are actually dp_size_ * tp_size_ devices. Update tp_size
743
- # and tp_rank so we shard across all devices.
744
- tp_size = dp_size_ * tp_size_
745
- tp_rank = dp_rank * tp_size_ + tp_rank
746
- return tp_size, tp_rank
747
-
751
use_ep = dp_size_ * tp_size_ > 1 and vllm_parallel_config.enable_expert_parallel
752
753
dp_size = dp_size_
754
dp_rank = get_dp_group().rank_in_group if dp_size > 1 else 0
752
- tp_size, tp_rank = flatten_tp_across_dp(dp_rank)
755
+ tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp(
756
+ tp_size_, dp_size_, dp_rank
757
+ )
758
759
if not use_ep:
760
return FusedMoEParallelConfig(
vllm/model_executor/models/gpt_oss.py
+16
-4
@@ -11,6 +11,7 @@ from vllm.attention import Attention, AttentionType
11
from vllm.compilation.decorators import support_torch_compile
12
from vllm.config import CacheConfig, VllmConfig
13
from vllm.distributed import (
14
+ get_dp_group,
15
get_ep_group,
16
get_pp_group,
17
get_tensor_model_parallel_rank,
@@ -18,6 +19,7 @@ from vllm.distributed import (
19
tensor_model_parallel_all_gather,
20
)
21
from vllm.model_executor.layers.fused_moe import FusedMoE
22
+from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
23
from vllm.model_executor.layers.layernorm import RMSNorm
24
from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
25
from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -305,8 +307,13 @@ class GptOssModel(nn.Module):
307
use_ep = self.parallel_config.enable_expert_parallel
308
num_experts = self.config.num_local_experts
309
308
- tp_rank = get_tensor_model_parallel_rank()
309
- tp_size = get_tensor_model_parallel_world_size()
310
+ # In MoE, we need to flatten the tensor parallel size across the data
311
+ # parallel size when EP is disabled.
312
+ tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp(
313
+ tp_size=get_tensor_model_parallel_world_size(),
314
+ dp_size=get_dp_group().world_size,
315
+ dp_rank=get_dp_group().rank_in_group,
316
+ )
317
318
intermediate_size = self.config.intermediate_size
319
intermediate_size_block = intermediate_size // mxfp4_block
@@ -488,8 +495,13 @@ class GptOssModel(nn.Module):
495
496
use_ep = self.parallel_config.enable_expert_parallel
497
491
- tp_rank = get_tensor_model_parallel_rank()
492
- tp_size = get_tensor_model_parallel_world_size()
498
+ # In MoE, we need to flatten the tensor parallel size across the data
499
+ # parallel size when EP is disabled.
500
+ tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp(
501
+ tp_size=get_tensor_model_parallel_world_size(),
502
+ dp_size=get_dp_group().world_size,
503
+ dp_rank=get_dp_group().rank_in_group,
504
+ )
505
506
intermediate_size = self.config.intermediate_size
507
per_rank_intermediate_size = cdiv(intermediate_size, tp_size)