[ROCm][DSV4.1][Perf] Fold the mHC post step into the delayed pre projection (#56513)

Signed-off-by: Fangzhou Ai <fangzhou@semianalysis.ai>
Signed-off-by: Fangzhou Ai <fangzhou.ai@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

tests/kernels/test_mhc_kernels.py

@@ -983,8 +983,12 @@ def test_mhc_pre_delayed_rocm_aiter(num_tokens, carried):
     )
 
     expected = mhc_pre_delayed_torch(*args, pre_mix=pre_mix)
-    actual = object.__new__(MHCPreDelayedOp).forward_hip(*args, pre_mix=pre_mix)
+    residual_out, *actual = object.__new__(MHCPreDelayedOp).forward_hip(
+        *args, pre_mix=pre_mix
+    )
 
+    # No post was requested, so the residual comes straight back.
+    assert residual_out is residual
     for i in (0, 1, 3):
         torch.testing.assert_close(actual[i], expected[i], atol=1e-4, rtol=1e-3)
     # The collapse is the same FP32 multiply-and-sum in both paths.
@@ -1026,7 +1030,12 @@ def test_mhc_pre_delayed_rocm_aiter_declines_unsupported(monkeypatch):
     broadcast_fn = fn.view(-1, hc_mult, hidden_size).sum(1)
     broadcast_residual = x.unsqueeze(1).expand(-1, hc_mult, -1).contiguous()
     expected = mhc_pre_delayed_torch(broadcast_residual, broadcast_fn, *args[2:], x=x)
-    actual = op.forward_hip(broadcast_residual, broadcast_fn, *args[2:], x=x)
+    # forward_hip leads with the residual; no post was requested here, so it is
+    # the caller's own tensor and the reference's four outputs follow it.
+    residual_out, *actual = op.forward_hip(
+        broadcast_residual, broadcast_fn, *args[2:], x=x
+    )
+    assert residual_out is broadcast_residual
     assert not took_aiter, "the broadcast seam must not reach AITER"
     for i in range(4):
         torch.testing.assert_close(actual[i], expected[i], atol=1e-4, rtol=1e-3)
@@ -1268,3 +1277,68 @@ def test_deepseek_v4_mhc_broadcast_refit_refreshes_in_place(monkeypatch):
     assert layer.hc_attn_fn_broadcast is buffer
     expected = layer.hc_attn_fn.detach().view(-1, 2, 8).sum(dim=1)
     assert torch.equal(layer.hc_attn_fn_broadcast, expected)
+
+
+@pytest.mark.skipif(
+    not (current_platform.is_rocm() and HAS_AITER_MHC),
+    reason="AITER mHC required",
+)
+@pytest.mark.parametrize("num_tokens", [1, 2, 7, 128])
+@pytest.mark.parametrize("carried", [False, True])
+def test_mhc_fused_post_pre_delayed_rocm_aiter(
+    num_tokens, carried, default_vllm_config
+):
+    """Folding the post into the pre projection must match doing them apart.
+
+    The residual is bfloat16, so the two accumulation orders are allowed to
+    land a rounding step apart; the tolerance here is one bfloat16 ULP.
+    """
+    set_random_seed(0)
+    hc_mult, hidden_size = 4, 5120
+    residual, fn, hc_scale, hc_base, _ = _rocm_mhc_inputs(
+        num_tokens=num_tokens, hidden_size=hidden_size, hc_mult=hc_mult
+    )
+    x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16, device=DEVICE)
+    post_layer_mix = (
+        torch.rand(num_tokens, hc_mult, 1, dtype=torch.float32, device=DEVICE) + 0.5
+    )
+    comb_res_mix = (
+        torch.rand(num_tokens, hc_mult, hc_mult, dtype=torch.float32, device=DEVICE)
+        + 0.2
+    )
+    pre_mix = (
+        torch.rand(num_tokens, hc_mult, dtype=torch.float32, device=DEVICE) + 0.5
+        if carried
+        else None
+    )
+    rms_eps = hc_pre_eps = hc_sinkhorn_eps = 1e-6
+    pre_args = (fn, hc_scale, hc_base, rms_eps, hc_pre_eps, hc_sinkhorn_eps, 1.0, 20)
+
+    next_residual = mhc_post_torch(x, residual, post_layer_mix, comb_res_mix)
+    expected = (
+        next_residual,
+        *mhc_pre_delayed_torch(next_residual, *pre_args, pre_mix=pre_mix),
+    )
+    actual = MHCPreDelayedOp().forward_hip(
+        residual,
+        *pre_args,
+        pre_mix=pre_mix,
+        sublayer_out=x,
+        post_layer_mix=post_layer_mix,
+        comb_res_mix=comb_res_mix,
+    )
+
+    for got, want in zip(actual, expected, strict=True):
+        torch.testing.assert_close(got, want.reshape(got.shape), atol=2e-2, rtol=8e-3)
+
+
+@pytest.mark.skipif(
+    not (current_platform.is_rocm() and HAS_AITER_MHC),
+    reason="AITER mHC required",
+)
+def test_mhc_fused_post_pre_delayed_falls_back_for_large_batches():
+    """AITER's heuristic hands large batches back to the unfused path."""
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    assert not rocm_aiter_ops.mhc_fused_post_pre_delayed_prefers_unfused(1)
+    assert rocm_aiter_ops.mhc_fused_post_pre_delayed_prefers_unfused(1 << 20)

vllm/_aiter_ops.py

@@ -1676,6 +1676,75 @@ def _rocm_aiter_fp8_attn_fake(
     )
 
 
+def _mhc_delayed_pre_tail(
+    gemm_out: torch.Tensor,
+    sqrsum: torch.Tensor,
+    residual: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    pre_mix: torch.Tensor | None,
+    rms_eps: float,
+    hc_pre_eps: float,
+    hc_sinkhorn_eps: float,
+    hc_post_mult_value: float,
+    sinkhorn_repeat: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Finish a delayed mHC pre once AITER has produced the projection.
+
+    Shared by the plain and the post-fused entry points, which differ only in
+    which kernel computes ``gemm_out`` / ``sqrsum`` and the residual they
+    project. ``mhc_pre_big_fuse`` also writes a collapse against its own
+    pre-mix, which the delayed formulation cannot use and so discards.
+    Returns ``(post_mix, comb_mix, layer_input, next_pre_mix)``.
+    """
+    from aiter.ops.mhc import mhc_pre_big_fuse
+
+    num_tokens, hc_mult, hidden_size = residual.shape
+    device = residual.device
+    with torch.device(device):
+        post_mix = torch.empty(
+            num_tokens, hc_mult, 1, dtype=torch.float32, device=device
+        )
+        comb_mix = torch.empty(
+            num_tokens, hc_mult, hc_mult, dtype=torch.float32, device=device
+        )
+        unused_collapse = torch.empty(
+            num_tokens, hidden_size, dtype=torch.bfloat16, device=device
+        )
+        mhc_pre_big_fuse(
+            post_mix,
+            comb_mix,
+            unused_collapse,
+            gemm_out,
+            sqrsum,
+            hc_scale,
+            hc_base,
+            residual,
+            rms_eps,
+            hc_pre_eps,
+            hc_sinkhorn_eps,
+            hc_post_mult_value,
+            sinkhorn_repeat,
+        )
+
+    next_pre_mix = torch.ops.vllm.mhc_pre_mix_triton(
+        gemm_out,
+        sqrsum,
+        hc_scale,
+        hc_base,
+        hc_mult,
+        hc_mult * hidden_size,
+        rms_eps,
+        hc_pre_eps,
+    )
+    if pre_mix is None:
+        # Model entry selects residual stream zero.
+        layer_input = residual[:, 0]
+    else:
+        layer_input = torch.ops.vllm.hc_collapse_triton(residual, pre_mix)
+    return post_mix, comb_mix, layer_input, next_pre_mix
+
+
 # Global flag to ensure ops are registered only once
 _OPS_REGISTERED = False
 
@@ -3611,6 +3680,10 @@ class rocm_aiter_ops:
         hc_post_mult_value: float,
         sinkhorn_repeat: int,
         pre_mix: torch.Tensor | None = None,
+        sublayer_out: torch.Tensor | None = None,
+        post_layer_mix: torch.Tensor | None = None,
+        comb_res_mix: torch.Tensor | None = None,
+        residual_out: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         """mHC pre using the pre-mix carried from the previous sublayer.
 
@@ -3625,6 +3698,15 @@ class rocm_aiter_ops:
         that redundant store is the price of not having a native delayed
         kernel, and is small next to the ~140 launches it replaces.
 
+        Passing ``sublayer_out`` / ``post_layer_mix`` / ``comb_res_mix``
+        also applies the preceding post block. AITER then projects the new
+        residual in the same kernel that computes it
+        (``mhc_fused_post_pre_gemm_sqrsum``), which saves one launch; every
+        later stage is identical either way. That new residual is written to
+        ``residual_out``, which the caller must supply in that case; it is an
+        output buffer rather than a return value so the op never has to hand
+        back one of its own inputs on the unfused path.
+
         Returns:
             post_mix: shape (..., hc_mult, 1), dtype torch.float32
             comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
@@ -3632,8 +3714,9 @@ class rocm_aiter_ops:
             next_pre_mix: shape (..., hc_mult), dtype torch.float32
         """
         from aiter.ops.mhc import (
+            get_mhc_fused_post_pre_config,
             get_mhc_pre_splitk,
-            mhc_pre_big_fuse,
+            mhc_fused_post_pre_gemm_sqrsum,
             mhc_pre_gemm_sqrsum,
         )
 
@@ -3666,9 +3749,16 @@ class rocm_aiter_ops:
 
         pre_mix_flat = None if pre_mix is None else pre_mix.view(-1, hc_mult)
 
+        fuse_post = sublayer_out is not None
         # AITER's Python wrappers allocate without explicit device arguments.
         with torch.device(device):
-            splitk, tile_k = get_mhc_pre_splitk(num_tokens, hc_hidden_size)
+            if fuse_post:
+                # Keyed on the per-stream K, i.e. hidden_size, not hc_mult * it.
+                splitk, tile_m, tile_n, tile_k = get_mhc_fused_post_pre_config(
+                    num_tokens, hidden_size
+                )
+            else:
+                splitk, tile_k = get_mhc_pre_splitk(num_tokens, hc_hidden_size)
             # AITER pads the GEMM output to a multiple of 32 columns.
             gemm_pad = torch.empty(
                 splitk,
@@ -3679,50 +3769,46 @@ class rocm_aiter_ops:
             )
             gemm_out = gemm_pad[:, :, :hc_mult3]
             sqrsum = torch.empty(splitk, num_tokens, dtype=torch.float32, device=device)
-            mhc_pre_gemm_sqrsum(gemm_out, sqrsum, residual_flat, fn, tile_k, 0)
+            if fuse_post:
+                # Restated for the type checker, which cannot narrow the
+                # optional arguments through the fuse_post flag.
+                assert sublayer_out is not None
+                assert post_layer_mix is not None and comb_res_mix is not None
+                assert residual_out is not None
+                projected = residual_out.view(-1, hc_mult, hidden_size)
+                mhc_fused_post_pre_gemm_sqrsum(
+                    gemm_out,
+                    sqrsum,
+                    projected,
+                    sublayer_out.view(-1, hidden_size),
+                    residual_flat,
+                    # The kernel wants the post mix squeezed to (tokens, hc_mult).
+                    post_layer_mix.view(-1, hc_mult, 1).squeeze(-1),
+                    comb_res_mix.view(-1, hc_mult, hc_mult),
+                    fn,
+                    tile_m,
+                    tile_n,
+                    tile_k,
+                    0,
+                )
+            else:
+                projected = residual_flat
+                mhc_pre_gemm_sqrsum(gemm_out, sqrsum, projected, fn, tile_k, 0)
 
-            post_mix = torch.empty(
-                num_tokens, hc_mult, 1, dtype=torch.float32, device=device
-            )
-            comb_mix = torch.empty(
-                num_tokens, hc_mult, hc_mult, dtype=torch.float32, device=device
-            )
-            unused_collapse = torch.empty(
-                num_tokens, hidden_size, dtype=torch.bfloat16, device=device
-            )
-            mhc_pre_big_fuse(
-                post_mix,
-                comb_mix,
-                unused_collapse,
-                gemm_out,
-                sqrsum,
-                hc_scale,
-                hc_base,
-                residual_flat,
-                rms_eps,
-                hc_pre_eps,
-                hc_sinkhorn_eps,
-                hc_post_mult_value,
-                sinkhorn_repeat,
-            )
-
-        next_pre_mix = torch.ops.vllm.mhc_pre_mix_triton(
+        post_mix, comb_mix, layer_input, next_pre_mix = _mhc_delayed_pre_tail(
             gemm_out,
             sqrsum,
+            projected,
             hc_scale,
             hc_base,
-            hc_mult,
-            hc_hidden_size,
+            pre_mix_flat,
             rms_eps,
             hc_pre_eps,
+            hc_sinkhorn_eps,
+            hc_post_mult_value,
+            sinkhorn_repeat,
         )
 
-        if pre_mix_flat is None:
-            # Model entry selects residual stream zero.
-            layer_input = residual_flat[:, 0]
-        else:
-            layer_input = torch.ops.vllm.hc_collapse_triton(residual_flat, pre_mix_flat)
-
         return (
             post_mix.view(*outer_shape, hc_mult, 1),
             comb_mix.view(*outer_shape, hc_mult, hc_mult),
@@ -3730,6 +3816,21 @@ class rocm_aiter_ops:
             next_pre_mix.view(*outer_shape, hc_mult),
         )
 
+    @staticmethod
+    def mhc_fused_post_pre_delayed_prefers_unfused(num_tokens: int) -> bool:
+        """True when AITER's own heuristic favours separate post and pre.
+
+        Mirrors the ``fused_m_upper_bound`` table in ``aiter.ops.mhc``: the
+        fused GEMM wins for decode-sized batches and loses for large ones.
+        """
+        from aiter.jit.utils.chip_info import get_gfx_runtime
+
+        return num_tokens >= {
+            "gfx950": 1024,
+            "gfx942": 128,
+            "gfx1250": 1024,
+        }.get(get_gfx_runtime(), 1024)
+
     @staticmethod
     def hc_head(
         hs_flat: torch.Tensor,

vllm/model_executor/kernels/mhc/aiter.py

@@ -115,6 +115,10 @@ def mhc_pre_delayed_aiter(
     hc_post_mult_value: float,
     sinkhorn_repeat: int,
     pre_mix: torch.Tensor | None = None,
+    sublayer_out: torch.Tensor | None = None,
+    post_layer_mix: torch.Tensor | None = None,
+    comb_res_mix: torch.Tensor | None = None,
+    residual_out: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
     """mHC pre with the pre-mix carried in from the previous sublayer.
 
@@ -134,6 +138,13 @@ def mhc_pre_delayed_aiter(
         sinkhorn_repeat: number of sinkhorn iterations
         pre_mix: shape (..., hc_mult) from the previous sublayer, or None at
             model entry to select residual stream zero.
+        sublayer_out: attention or FFN output, shape (..., hidden_size). When
+            given with the two mixes below, the preceding post block is
+            applied here so AITER can fold it into the pre projection.
+        post_layer_mix: shape (..., hc_mult, 1), post gate for that block.
+        comb_res_mix: shape (..., hc_mult, hc_mult), residual comb for it.
+        residual_out: shape (..., hc_mult, hidden_size), written with the post
+            block's new residual. Required exactly when the post is folded in.
 
     Returns:
         post_mix: shape (..., hc_mult, 1), dtype torch.float32
@@ -156,6 +167,10 @@ def mhc_pre_delayed_aiter(
         hc_post_mult_value,
         sinkhorn_repeat,
         pre_mix,
+        sublayer_out,
+        post_layer_mix,
+        comb_res_mix,
+        residual_out,
     )
 
 
@@ -170,6 +185,10 @@ def _mhc_pre_delayed_aiter_fake(
     hc_post_mult_value: float,
     sinkhorn_repeat: int,
     pre_mix: torch.Tensor | None = None,
+    sublayer_out: torch.Tensor | None = None,
+    post_layer_mix: torch.Tensor | None = None,
+    comb_res_mix: torch.Tensor | None = None,
+    residual_out: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
     hc_mult = residual.shape[-2]
     hidden_size = residual.shape[-1]
@@ -313,7 +332,7 @@ direct_register_custom_op(
 direct_register_custom_op(
     op_name="mhc_pre_delayed_aiter",
     op_func=mhc_pre_delayed_aiter,
-    mutates_args=[],
+    mutates_args=["residual_out"],
     fake_impl=_mhc_pre_delayed_aiter_fake,
 )
 direct_register_custom_op(

vllm/model_executor/layers/mhc.py

@@ -299,7 +299,13 @@ class MHCPreDelayedOp(CustomOp):
 
     Same gates as :class:`MHCPreOp`, but the stream collapse applies the
     caller's ``pre_mix`` and this sublayer's pre-mix is returned for the next
-    sublayer seam. Returns post_mix, comb_mix, layer_input, next_pre_mix.
+    sublayer seam.
+
+    Passing ``sublayer_out`` / ``post_layer_mix`` / ``comb_res_mix`` also
+    applies the preceding post block, which lets AITER fold it into the pre
+    projection and saves a launch. Returns residual, post_mix, comb_mix,
+    layer_input, next_pre_mix; ``residual`` is the caller's own tensor when
+    no post is requested.
     """
 
     # --8<-- [end:mhc_pre_delayed]
@@ -307,6 +313,24 @@ class MHCPreDelayedOp(CustomOp):
     def enabled(cls) -> bool:
         return True
 
+    def __init__(self) -> None:
+        super().__init__()
+        # Built here, not lazily: CustomOp resolves dispatch against the
+        # current vLLM config, which is only set during model construction.
+        self._post = MHCPostOp()
+
+    def _maybe_post(
+        self,
+        residual: torch.Tensor,
+        sublayer_out: torch.Tensor | None,
+        post_layer_mix: torch.Tensor | None,
+        comb_res_mix: torch.Tensor | None,
+    ) -> torch.Tensor:
+        if sublayer_out is None:
+            return residual
+        assert post_layer_mix is not None and comb_res_mix is not None
+        return self._post(sublayer_out, residual, post_layer_mix, comb_res_mix)
+
     def forward_cuda(
         self,
         residual: torch.Tensor,
@@ -322,8 +346,14 @@ class MHCPreDelayedOp(CustomOp):
         x: torch.Tensor | None = None,
         norm_weight: torch.Tensor | None = None,
         norm_eps: float = 1e-6,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        return torch.ops.vllm.mhc_pre_delayed_tilelang(
+        sublayer_out: torch.Tensor | None = None,
+        post_layer_mix: torch.Tensor | None = None,
+        comb_res_mix: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        residual = self._maybe_post(
+            residual, sublayer_out, post_layer_mix, comb_res_mix
+        )
+        return residual, *torch.ops.vllm.mhc_pre_delayed_tilelang(
             residual,
             fn,
             hc_scale,
@@ -354,7 +384,10 @@ class MHCPreDelayedOp(CustomOp):
         x: torch.Tensor | None = None,
         norm_weight: torch.Tensor | None = None,
         norm_eps: float = 1e-6,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        sublayer_out: torch.Tensor | None = None,
+        post_layer_mix: torch.Tensor | None = None,
+        comb_res_mix: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         # The AITER delayed path drives mhc_pre_gemm_sqrsum against `residual`
         # and folds no RMSNorm, so it cannot serve the model-entry broadcast
         # (which projects a narrower `x`) or a fused norm. Both are handled by
@@ -364,7 +397,26 @@ class MHCPreDelayedOp(CustomOp):
             and norm_weight is None
             and _aiter_mhc_supported(residual, None, supports_norm=False)
         ):
-            return torch.ops.vllm.mhc_pre_delayed_aiter(
+            from vllm._aiter_ops import rocm_aiter_ops
+
+            num_tokens = residual.numel() // (residual.shape[-1] * residual.shape[-2])
+            # Folding the post in only pays while the residual still fits in
+            # cache; past that AITER's heuristic wants the separate kernels,
+            # which can ask for a non-temporal store and a large-m split-k.
+            if sublayer_out is not None and (
+                rocm_aiter_ops.mhc_fused_post_pre_delayed_prefers_unfused(num_tokens)
+            ):
+                residual = self._maybe_post(
+                    residual, sublayer_out, post_layer_mix, comb_res_mix
+                )
+                sublayer_out = None
+            # The op writes the folded post's residual into this buffer rather
+            # than returning it, since on the unfused path it has none of its
+            # own and an op must not hand back one of its inputs.
+            residual_out = (
+                torch.empty_like(residual) if sublayer_out is not None else None
+            )
+            rest = torch.ops.vllm.mhc_pre_delayed_aiter(
                 residual,
                 fn,
                 hc_scale,
@@ -375,9 +427,17 @@ class MHCPreDelayedOp(CustomOp):
                 hc_post_mult_value,
                 sinkhorn_repeat,
                 pre_mix,
+                sublayer_out,
+                post_layer_mix,
+                comb_res_mix,
+                residual_out,
             )
+            return (residual if residual_out is None else residual_out), *rest
+        residual = self._maybe_post(
+            residual, sublayer_out, post_layer_mix, comb_res_mix
+        )
         if HAS_TILELANG_MHC:
-            return torch.ops.vllm.mhc_pre_delayed_tilelang(
+            return residual, *torch.ops.vllm.mhc_pre_delayed_tilelang(
                 residual,
                 fn,
                 hc_scale,
@@ -392,6 +452,7 @@ class MHCPreDelayedOp(CustomOp):
                 norm_weight,
                 norm_eps,
             )
+        # The post is already applied, so it must not be requested again.
         return self.forward_native(
             residual,
             fn,
@@ -423,7 +484,13 @@ class MHCPreDelayedOp(CustomOp):
         x: torch.Tensor | None = None,
         norm_weight: torch.Tensor | None = None,
         norm_eps: float = 1e-6,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        sublayer_out: torch.Tensor | None = None,
+        post_layer_mix: torch.Tensor | None = None,
+        comb_res_mix: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        residual = self._maybe_post(
+            residual, sublayer_out, post_layer_mix, comb_res_mix
+        )
         post_mix, comb_mix, layer_input, next_pre_mix = (
             mhc_kernels.mhc_pre_delayed_torch(
                 residual,
@@ -440,6 +507,7 @@ class MHCPreDelayedOp(CustomOp):
             )
         )
         return (
+            residual,
             post_mix,
             comb_mix,
             _apply_mhc_norm(layer_input, norm_weight, norm_eps),

vllm/models/deepseek_v41/amd/model.py

@@ -274,7 +274,7 @@ class DeepseekV4DecoderLayer(nn.Module):
                 # copies and the identity pre-mix selects copy 0.
                 assert self.hc_attn_fn_broadcast is not None
                 residual = x.unsqueeze(1).expand(-1, self.hc_mult, -1).contiguous()
-                post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed(
+                residual, post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed(
                     residual,
                     self.hc_attn_fn_broadcast,
                     self.hc_attn_scale,
@@ -288,7 +288,7 @@ class DeepseekV4DecoderLayer(nn.Module):
                 )
             else:
                 residual = x
-                post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed(
+                residual, post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed(
                     residual,
                     self.hc_attn_fn,
                     self.hc_attn_scale,
@@ -301,18 +301,7 @@ class DeepseekV4DecoderLayer(nn.Module):
                     pre_mix=pre_mix,
                 )
         else:
-            residual = self.mhc_post(x, residual, post_mix, res_mix)
-            if self.engram is not None and engram_hashes is not None:
-                # Engram injection happens between the previous sublayer's
-                # post and this block's pre, on the full hc stream, so the
-                # mix coefficients see the injected stream.
-                residual = self.engram(
-                    residual,
-                    engram_hashes[:, self.engram.layer_hash_index],
-                    engram_mask,
-                )
-            post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed(
-                residual,
+            pre_args = (
                 self.hc_attn_fn,
                 self.hc_attn_scale,
                 self.hc_attn_base,
@@ -321,8 +310,37 @@ class DeepseekV4DecoderLayer(nn.Module):
                 self.hc_eps,
                 self.hc_post_alpha,
                 self.hc_sinkhorn_iters,
-                pre_mix=pre_mix,
             )
+            if self.engram is not None and engram_hashes is not None:
+                # Engram injection happens between the previous sublayer's
+                # post and this block's pre, on the full hc stream, so the
+                # mix coefficients see the injected stream. That read of the
+                # residual is what stops the post from fusing into the pre
+                # here, unlike the attention seam below.
+                residual = self.mhc_post(x, residual, post_mix, res_mix)
+                residual = self.engram(
+                    residual,
+                    engram_hashes[:, self.engram.layer_hash_index],
+                    engram_mask,
+                )
+                residual, post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed(
+                    residual, *pre_args, pre_mix=pre_mix
+                )
+            else:
+                (
+                    residual,
+                    post_mix,
+                    res_mix,
+                    x,
+                    attn_pre,
+                ) = self.mhc_pre_delayed(
+                    residual,
+                    *pre_args,
+                    pre_mix=pre_mix,
+                    sublayer_out=x,
+                    post_layer_mix=post_mix,
+                    comb_res_mix=res_mix,
+                )
         x = self.attn_norm(x)
 
         if self.use_sequence_parallel:
@@ -332,8 +350,7 @@ class DeepseekV4DecoderLayer(nn.Module):
         if self.use_sequence_parallel:
             x = sp_reduce_scatter(x)
 
-        residual = self.mhc_post(x, residual, post_mix, res_mix)
-        post_mix, res_mix, x, ffn_pre = self.mhc_pre_delayed(
+        residual, post_mix, res_mix, x, ffn_pre = self.mhc_pre_delayed(
             residual,
             self.hc_ffn_fn,
             self.hc_ffn_scale,
@@ -344,6 +361,9 @@ class DeepseekV4DecoderLayer(nn.Module):
             self.hc_post_alpha,
             self.hc_sinkhorn_iters,
             pre_mix=attn_pre,
+            sublayer_out=x,
+            post_layer_mix=post_mix,
+            comb_res_mix=res_mix,
         )
         x = self.ffn_norm(x)
         x = self.ffn(x, input_ids)