From 00972dfd72988942138a7a6089eaee08580210b8 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Mon, 14 Sep 2026 17:44:56 -0400 Subject: [PATCH] [ROCm][DSV4.1][Perf] Fold the mHC post step into the delayed pre projection (#56513) Signed-off-by: Fangzhou Ai Signed-off-by: Fangzhou Ai Co-authored-by: Cursor --- tests/kernels/test_mhc_kernels.py | 78 +++++++++- vllm/_aiter_ops.py | 175 ++++++++++++++++++----- vllm/model_executor/kernels/mhc/aiter.py | 21 ++- vllm/model_executor/layers/mhc.py | 82 ++++++++++- vllm/models/deepseek_v41/amd/model.py | 54 ++++--- 5 files changed, 346 insertions(+), 64 deletions(-) diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index 5049815380..f805fe39a5 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/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) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 454aba175e..4cfb3b3884 100644 --- a/vllm/_aiter_ops.py +++ b/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, diff --git a/vllm/model_executor/kernels/mhc/aiter.py b/vllm/model_executor/kernels/mhc/aiter.py index e04ce3322b..d34e29314d 100644 --- a/vllm/model_executor/kernels/mhc/aiter.py +++ b/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( diff --git a/vllm/model_executor/layers/mhc.py b/vllm/model_executor/layers/mhc.py index 14f1018597..73993fbc98 100644 --- a/vllm/model_executor/layers/mhc.py +++ b/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), diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index eaed15214f..b3eddc2024 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/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)