[Bugfix][ROCm][GLM-5.2] Stop forcing skip_final_all_reduce on AMD DSA model path (#58589)

Signed-off-by: avininjamay8 <Avinash.Paul@amd.com>
Signed-off-by: avininjamay8 <avpaul@amd.com>
Co-authored-by: avininjamay8 <avpaul@amd.com>
Co-authored-by: Kevin H. Luu <khluu000@gmail.com>
Co-authored-by: vllm-agent <claw@inferact.ai>

vllm/models/deepseek_v32/amd/model.py

@@ -8,7 +8,7 @@ from itertools import islice
 import torch
 
 from vllm.config import VllmConfig
-from vllm.distributed import get_pp_group
+from vllm.distributed import get_pp_group, get_tensor_model_parallel_rank
 from vllm.model_executor.layers.fused_moe import (
     fused_moe_make_expert_params_mapping,
 )
@@ -71,6 +71,7 @@ class DeepseekV32DecoderLayer(torch.nn.Module):
             topk_indices_buffer=topk_indices_buffer,
         )
 
+        self.zero_mlp_output = False
         if (
             config.n_routed_experts is not None
             and layer_idx >= config.first_k_dense_replace
@@ -82,7 +83,15 @@ class DeepseekV32DecoderLayer(torch.nn.Module):
                 quant_config=quant_config,
                 prefix=f"{prefix}.mlp",
             )
-            self.mlp.experts.moe_config.skip_final_all_reduce = True
+            moe_cfg = self.mlp.experts.moe_config
+            pc = moe_cfg.moe_parallel_config
+            # Validated on MoRI only. DeepEP, pplx and the other combines
+            # keep the previous always-skip behavior.
+            if pc.use_mori_kernels:
+                # Combine is already all-reduced; the next fused norm sums TP.
+                self.zero_mlp_output = get_tensor_model_parallel_rank() > 0
+            else:
+                moe_cfg.skip_final_all_reduce = True
         else:
             self.mlp = DeepseekV2MLP(
                 hidden_size=config.hidden_size,
@@ -116,6 +125,8 @@ class DeepseekV32DecoderLayer(torch.nn.Module):
             hidden_states, residual, self.post_attention_layernorm
         )
         hidden_states = self.mlp(hidden_states)
+        if self.zero_mlp_output:
+            hidden_states = torch.zeros_like(hidden_states)
         return hidden_states, residual