[ROCm][Perf] Fuse DeepSeek-V4 mHC post/pre and RMSNorm with AITER (#52737)

Signed-off-by: shen-shanshan <467638484@qq.com>
Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>

vllm/_aiter_ops.py

@@ -3299,6 +3299,8 @@ class rocm_aiter_ops:
         hc_sinkhorn_eps: float,
         hc_post_mult_value: float,
         sinkhorn_repeat: int,
+        norm_weight: torch.Tensor | None = None,
+        norm_eps: float = 0.0,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """
         Forward pass for mHC pre block.
@@ -3313,6 +3315,8 @@ class rocm_aiter_ops:
             hc_sinkhorn_eps: sinkhorn epsilon
             hc_post_mult_value: post-mix multiplier value
             sinkhorn_repeat: number of sinkhorn iterations
+            norm_weight: optional RMSNorm weight fused into the pre kernel
+            norm_eps: epsilon for the fused RMSNorm when norm_weight is set
 
         Returns:
             post_mix: shape (..., hc_mult), dtype torch.float32
@@ -3380,6 +3384,8 @@ class rocm_aiter_ops:
                 hc_sinkhorn_eps,
                 hc_post_mult_value,
                 sinkhorn_repeat,
+                norm_weight,
+                norm_eps,
             )
         return (
             post_mix.view(*outer_shape, hc_mult, 1),
@@ -3465,5 +3471,95 @@ class rocm_aiter_ops:
         )
         return out.view_as(residual)
 
+    @staticmethod
+    def mhc_fused_post_pre(
+        x: torch.Tensor,
+        residual: torch.Tensor,
+        post_layer_mix: torch.Tensor,
+        comb_res_mix: torch.Tensor,
+        fn: torch.Tensor,
+        hc_scale: torch.Tensor,
+        hc_base: torch.Tensor,
+        rms_eps: float,
+        hc_pre_eps: float,
+        hc_sinkhorn_eps: float,
+        hc_post_mult_value: float,
+        sinkhorn_repeat: int,
+        norm_weight: torch.Tensor | None = None,
+        norm_eps: float = 0.0,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Fused mHC post + next mHC pre via AITER.
+
+        Returns residual_cur, post_mix, comb_mix, layer_input in vLLM order
+        (AITER returns post_mix, comb_mix, layer_input, next_residual).
+        """
+        from aiter.ops.mhc import mhc_fused_post_pre
+
+        assert x.dtype == torch.bfloat16
+        assert residual.dtype == torch.bfloat16
+        assert fn.dtype == torch.float32
+        assert hc_scale.dtype == torch.float32
+        assert hc_base.dtype == torch.float32
+
+        hc_mult = residual.shape[-2]
+        hidden_size = residual.shape[-1]
+        outer_shape = residual.shape[:-2]
+
+        residual_flat = residual.view(-1, hc_mult, hidden_size)
+        num_tokens = residual_flat.shape[0]
+        x_flat = x.view(num_tokens, hidden_size)
+        post_flat = post_layer_mix.view(num_tokens, hc_mult, 1)
+        comb_flat = comb_res_mix.view(num_tokens, hc_mult, hc_mult)
+
+        if num_tokens == 0:
+            return (
+                torch.empty_like(residual_flat).view_as(residual),
+                torch.empty(
+                    *outer_shape,
+                    hc_mult,
+                    1,
+                    dtype=torch.float32,
+                    device=residual.device,
+                ),
+                torch.empty(
+                    *outer_shape,
+                    hc_mult,
+                    hc_mult,
+                    dtype=torch.float32,
+                    device=residual.device,
+                ),
+                torch.empty(
+                    *outer_shape,
+                    hidden_size,
+                    dtype=torch.bfloat16,
+                    device=residual.device,
+                ),
+            )
+
+        with torch.device(residual_flat.device):
+            post_mix, comb_mix, layer_input, next_residual = mhc_fused_post_pre(
+                x_flat,
+                residual_flat,
+                post_flat,
+                comb_flat,
+                fn,
+                hc_scale,
+                hc_base,
+                rms_eps,
+                hc_pre_eps,
+                hc_sinkhorn_eps,
+                hc_post_mult_value,
+                sinkhorn_repeat,
+                norm_weight,
+                norm_eps,
+            )
+
+        return (
+            next_residual.view_as(residual),
+            post_mix.view(*outer_shape, hc_mult, 1),
+            comb_mix.view(*outer_shape, hc_mult, hc_mult),
+            layer_input.view(*outer_shape, hidden_size),
+        )
+
 
 rocm_aiter_ops.register_ops_once()

vllm/model_executor/kernels/mhc/aiter.py

@@ -16,6 +16,8 @@ def mhc_pre_aiter(
     hc_post_mult_value: float,
     sinkhorn_repeat: int,
     n_splits: int = 1,
+    norm_weight: torch.Tensor | None = None,
+    norm_eps: float = 0.0,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """
     Forward pass for mHC pre block.
@@ -31,6 +33,8 @@ def mhc_pre_aiter(
         hc_post_mult_value: post-mix multiplier value
         sinkhorn_repeat: number of sinkhorn iterations
         n_splits: split-k factor;
+        norm_weight: optional RMSNorm weight fused into the pre kernel
+        norm_eps: epsilon for the fused RMSNorm when norm_weight is set
 
     Returns:
         post_mix: shape (..., hc_mult), dtype torch.float32
@@ -52,6 +56,8 @@ def mhc_pre_aiter(
         hc_sinkhorn_eps,
         hc_post_mult_value,
         sinkhorn_repeat,
+        norm_weight,
+        norm_eps,
     )
 
 
@@ -66,6 +72,8 @@ def _mhc_pre_aiter_fake(
     hc_post_mult_value: float,
     sinkhorn_repeat: int,
     n_splits: int = 1,
+    norm_weight: torch.Tensor | None = None,
+    norm_eps: float = 0.0,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     hc_mult = residual.shape[-2]
     hidden_size = residual.shape[-1]
@@ -124,6 +132,94 @@ def _mhc_post_aiter_fake(
     return torch.empty_like(residual)
 
 
+def mhc_fused_post_pre_aiter(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    post_layer_mix: torch.Tensor,
+    comb_res_mix: torch.Tensor,
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    rms_eps: float,
+    hc_pre_eps: float,
+    hc_sinkhorn_eps: float,
+    hc_post_mult_value: float,
+    sinkhorn_repeat: int,
+    n_splits: int = 1,
+    tile_n: int = 1,
+    norm_weight: torch.Tensor | None = None,
+    norm_eps: float = 0.0,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Fused mHC post + next mHC pre on ROCm via AITER.
+
+    Returns residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur.
+    """
+    hidden_size = residual.shape[-1]
+    assert hidden_size % 256 == 0
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    return rocm_aiter_ops.mhc_fused_post_pre(
+        x,
+        residual,
+        post_layer_mix,
+        comb_res_mix,
+        fn,
+        hc_scale,
+        hc_base,
+        rms_eps,
+        hc_pre_eps,
+        hc_sinkhorn_eps,
+        hc_post_mult_value,
+        sinkhorn_repeat,
+        norm_weight,
+        norm_eps,
+    )
+
+
+def _mhc_fused_post_pre_aiter_fake(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    post_layer_mix: torch.Tensor,
+    comb_res_mix: torch.Tensor,
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    rms_eps: float,
+    hc_pre_eps: float,
+    hc_sinkhorn_eps: float,
+    hc_post_mult_value: float,
+    sinkhorn_repeat: int,
+    n_splits: int = 1,
+    tile_n: int = 1,
+    norm_weight: torch.Tensor | None = None,
+    norm_eps: float = 0.0,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    hc_mult = residual.shape[-2]
+    hidden_size = residual.shape[-1]
+    outer_shape = residual.shape[:-2]
+    post_mix = torch.empty(
+        *outer_shape,
+        hc_mult,
+        1,
+        dtype=torch.float32,
+        device=residual.device,
+    )
+    comb_mix = torch.empty(
+        *outer_shape,
+        hc_mult,
+        hc_mult,
+        dtype=torch.float32,
+        device=residual.device,
+    )
+    layer_input = torch.empty(
+        *outer_shape,
+        hidden_size,
+        dtype=torch.bfloat16,
+        device=residual.device,
+    )
+    return torch.empty_like(residual), post_mix, comb_mix, layer_input
+
+
 direct_register_custom_op(
     op_name="mhc_pre_aiter",
     op_func=mhc_pre_aiter,
@@ -136,3 +232,9 @@ direct_register_custom_op(
     mutates_args=[],
     fake_impl=_mhc_post_aiter_fake,
 )
+direct_register_custom_op(
+    op_name="mhc_fused_post_pre_aiter",
+    op_func=mhc_fused_post_pre_aiter,
+    mutates_args=[],
+    fake_impl=_mhc_fused_post_pre_aiter_fake,
+)

vllm/model_executor/layers/mhc.py

@@ -90,7 +90,8 @@ class MHCPreOp(CustomOp):
         norm_eps: float = 0.0,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         hidden_size = residual.shape[-1]
-        if HAS_AITER_MHC and hidden_size % 256 == 0:
+        hc_mult = residual.shape[-2]
+        if HAS_AITER_MHC and hidden_size % 256 == 0 and hc_mult == 4:
             return torch.ops.vllm.mhc_pre_aiter(
                 residual,
                 fn,
@@ -101,6 +102,9 @@ class MHCPreOp(CustomOp):
                 hc_sinkhorn_eps,
                 hc_post_mult_value,
                 sinkhorn_repeat,
+                n_splits,
+                norm_weight,
+                norm_eps,
             )
         elif HAS_TILELANG_MHC:
             return torch.ops.vllm.mhc_pre_tilelang(
@@ -222,7 +226,8 @@ class MHCPostOp(CustomOp):
         comb_res_mix: torch.Tensor,
     ) -> torch.Tensor:
         hidden_size = residual.shape[-1]
-        if HAS_AITER_MHC and hidden_size % 256 == 0:
+        hc_mult = residual.shape[-2]
+        if HAS_AITER_MHC and hidden_size % 256 == 0 and hc_mult == 4:
             return torch.ops.vllm.mhc_post_aiter(
                 x,
                 residual,
@@ -444,6 +449,27 @@ class MHCFusedPostPreOp(CustomOp):
         norm_weight: torch.Tensor | None = None,
         norm_eps: float = 0.0,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        hidden_size = residual.shape[-1]
+        hc_mult = residual.shape[-2]
+        if HAS_AITER_MHC and hidden_size % 256 == 0 and hc_mult == 4:
+            return torch.ops.vllm.mhc_fused_post_pre_aiter(
+                x,
+                residual,
+                post_layer_mix,
+                comb_res_mix,
+                fn,
+                hc_scale,
+                hc_base,
+                rms_eps,
+                hc_pre_eps,
+                hc_sinkhorn_eps,
+                hc_post_mult_value,
+                sinkhorn_repeat,
+                n_splits,
+                tile_n,
+                norm_weight,
+                norm_eps,
+            )
         if HAS_TILELANG_MHC:
             return torch.ops.vllm.mhc_fused_post_pre_tilelang(
                 x,

vllm/models/deepseek_v4/amd/dspark.py

@@ -9,8 +9,8 @@ ROCm port of ``nvidia/dspark.py``. Follows the same nvidia->amd recipe used for
     attention + MHC CustomOp path) instead of the nvidia one;
   * route the MHC head through the ``HCHeadOp`` CustomOp dispatcher (aiter /
     tilelang / triton / torch) instead of calling the tilelang kernels directly,
-    and gate the trailing ``mhc_post`` on ``use_fused_mhc`` (False on the aiter
-    path, where the decoder layer already applies hc_post in-layer);
+    and gate the trailing ``mhc_post`` on ``use_fused_mhc`` (True when AITER
+    or TileLang fused MHC is available; False only on the torch fallback);
   * drop the mega-MoE weight path (``make_deepseek_v4_expert_params_mapping`` /
     ``use_mega_moe`` / ``finalize_mega_moe_weights`` do not exist in amd/model.py).
 

vllm/models/deepseek_v4/amd/model.py

@@ -305,6 +305,10 @@ class DeepseekV4MoE(nn.Module):
         return final_hidden_states.view(org_shape)
 
 
+# Hidden sizes supported by AITER mhc_pre_big_fuse_rmsnorm.
+_AITER_MHC_FUSED_RMSNORM_SIZES = frozenset({1280, 2560, 4096, 7168})
+
+
 class DeepseekV4DecoderLayer(nn.Module):
     def __init__(
         self,
@@ -384,9 +388,18 @@ class DeepseekV4DecoderLayer(nn.Module):
         self.mhc_pre = MHCPreOp()
         self.mhc_post = MHCPostOp()
         self.mhc_fused_post_pre = MHCFusedPostPreOp()
-        self.use_fused_mhc = HAS_TILELANG_MHC and not (
-            HAS_AITER_MHC and self.hidden_size % 256 == 0
+        # AITER mhc kernels (pre/post/fused) require hc_mult == 4.
+        use_aiter_mhc = (
+            HAS_AITER_MHC and self.hidden_size % 256 == 0 and self.hc_mult == 4
         )
+        # Prefer AITER fused post+pre when eligible; otherwise TileLang.
+        self.use_fused_mhc = use_aiter_mhc or HAS_TILELANG_MHC
+        # Fold attn/ffn RMSNorm into MHC only when the active backend's
+        # fused-rmsnorm path supports this hidden size.
+        if use_aiter_mhc:
+            self.fuse_mhc_rmsnorm = self.hidden_size in _AITER_MHC_FUSED_RMSNORM_SIZES
+        else:
+            self.fuse_mhc_rmsnorm = HAS_TILELANG_MHC and self.use_fused_mhc
 
     def hc_pre(
         self,
@@ -394,7 +407,13 @@ class DeepseekV4DecoderLayer(nn.Module):
         hc_fn: torch.Tensor,
         hc_scale: torch.Tensor,
         hc_base: torch.Tensor,
+        norm_weight: torch.Tensor | None = None,
+        norm_eps: float = 0.0,
     ):
+        """Reduce HC residual streams to the next sub-layer input.
+
+        When ``norm_weight`` is set, RMSNorm is fused into the pre kernel.
+        """
         post_mix, res_mix, layer_input = self.mhc_pre(
             residual=x,
             fn=hc_fn,
@@ -405,6 +424,8 @@ class DeepseekV4DecoderLayer(nn.Module):
             hc_sinkhorn_eps=self.hc_eps,
             hc_post_mult_value=self.hc_post_alpha,
             sinkhorn_repeat=self.hc_sinkhorn_iters,
+            norm_weight=norm_weight,
+            norm_eps=norm_eps,
         )
         return layer_input, post_mix, res_mix
 
@@ -426,11 +447,20 @@ class DeepseekV4DecoderLayer(nn.Module):
         res_mix: torch.Tensor | None = None,
         residual: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        attn_norm_weight = self.attn_norm.weight if self.fuse_mhc_rmsnorm else None
+        attn_norm_eps = (
+            self.attn_norm.variance_epsilon if self.fuse_mhc_rmsnorm else 0.0
+        )
         if residual is None:
             # Run standalone hc_pre on first layer
             residual = x
             x, post_mix, res_mix = self.hc_pre(
-                x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base
+                x,
+                self.hc_attn_fn,
+                self.hc_attn_scale,
+                self.hc_attn_base,
+                norm_weight=attn_norm_weight,
+                norm_eps=attn_norm_eps,
             )
         else:
             residual, post_mix, res_mix, x = self.mhc_fused_post_pre(
@@ -446,11 +476,16 @@ class DeepseekV4DecoderLayer(nn.Module):
                 self.hc_eps,
                 self.hc_post_alpha,
                 self.hc_sinkhorn_iters,
+                norm_weight=attn_norm_weight,
+                norm_eps=attn_norm_eps,
             )
 
-        x = self.attn_norm(x)
+        if not self.fuse_mhc_rmsnorm:
+            x = self.attn_norm(x)
         x = self.attn(positions, x, None)
 
+        ffn_norm_weight = self.ffn_norm.weight if self.fuse_mhc_rmsnorm else None
+        ffn_norm_eps = self.ffn_norm.variance_epsilon if self.fuse_mhc_rmsnorm else 0.0
         residual, post_mix, res_mix, x = self.mhc_fused_post_pre(
             x,
             residual,
@@ -464,8 +499,11 @@ class DeepseekV4DecoderLayer(nn.Module):
             self.hc_eps,
             self.hc_post_alpha,
             self.hc_sinkhorn_iters,
+            norm_weight=ffn_norm_weight,
+            norm_eps=ffn_norm_eps,
         )
-        x = self.ffn_norm(x)
+        if not self.fuse_mhc_rmsnorm:
+            x = self.ffn_norm(x)
         x = self.ffn(x, input_ids)
         return x, residual, post_mix, res_mix
 
@@ -678,7 +716,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
                 residual,
             )
             if (idx + 1) in self.aux_hidden_state_layers:
-                # On the unfused (aiter) path the layer already applied hc_post,
+                # On the unfused path the layer already applied hc_post,
                 # so hidden_states is the reconstructed stream; on the fused
                 # path reconstruct it via hc_post before averaging.
                 if layer.use_fused_mhc: