enable_cuda_graph fix

Author: ZJLi2013

vllm/model_executor/models/deepseek_v2.py

@@ -582,12 +582,18 @@ def forward(
     ) -> torch.Tensor:
         # Self Attention
         if is_rocm_aiter_fused_rms_quant_enabled():
-            q_hidden_states, residual, y_scale = aiter_ops.rocm_aiter_rmsnorm2d_fwd_with_add_quant(hidden_states,
-                                                            residual=residual, weight=self.input_layernorm.weight,
+            if residual is None: 
+                q_hidden_states, residual_out, y_scale = aiter_ops.rocm_aiter_rmsnorm2d_fwd_with_add_quant(hidden_states,
+                                                            residual=None, weight=self.input_layernorm.weight,
                                                             variance_epsilon=self.input_layernorm.variance_epsilon,
                                                             x_scale=None, y_scale_dtype=torch.float32)
-            if residual is None:
                 residual = hidden_states 
+            else: # tmp fix for enable cuda graph
+                q_hidden_states, residual_out, y_scale = aiter_ops.rocm_aiter_rmsnorm2d_fwd_with_add_quant(hidden_states,
+                                                            residual=residual, weight=self.input_layernorm.weight,
+                                                            variance_epsilon=self.input_layernorm.variance_epsilon,
+                                                            x_scale=None, y_scale_dtype=torch.float32)
+                residual = residual_out
         else: 
             if residual is None:
                 residual = hidden_states