From e90e359e77e9eeaa31a44b23ee255600c6f0b4ee Mon Sep 17 00:00:00 2001 From: Mohit Khatwani Date: Wed, 11 Mar 2026 22:44:00 +0000 Subject: [PATCH] fix expert_parallelism_name assignment --- src/maxtext/layers/moe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 70e67490f5..c8253d5203 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -368,7 +368,7 @@ def __init__( else: self._tensor_parallelism_name = "tensor" - if self.config.attention == "vllm_rpa": + if self.config.attention == "vllm_rpa" and self.config.enable_dp_attention: self._expert_parallelism_name = "attn_dp_expert" else: self._expert_parallelism_name = "expert"