@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 00972dfd72

[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>

Fangzhou Ai committed Sep 14, 2026 at 17:44 UTC 00972dfd72988942138a7a6089eaee08580210b8

5 files changed +346 -64

tests/kernels/test_mhc_kernels.py

+76 -2

@@ -983,8 +983,12 @@ def test_mhc_pre_delayed_rocm_aiter(num_tokens, carried):

983 )

984

985 expected = mhc_pre_delayed_torch(*args, pre_mix=pre_mix)

986 - actual = object.__new__(MHCPreDelayedOp).forward_hip(*args, pre_mix=pre_mix)

986 + residual_out, *actual = object.__new__(MHCPreDelayedOp).forward_hip(

987 + *args, pre_mix=pre_mix

988 + )

989

990 + # No post was requested, so the residual comes straight back.

991 + assert residual_out is residual

992 for i in (0, 1, 3):

993 torch.testing.assert_close(actual[i], expected[i], atol=1e-4, rtol=1e-3)

994 # 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):

1030 broadcast_fn = fn.view(-1, hc_mult, hidden_size).sum(1)

1031 broadcast_residual = x.unsqueeze(1).expand(-1, hc_mult, -1).contiguous()

1032 expected = mhc_pre_delayed_torch(broadcast_residual, broadcast_fn, *args[2:], x=x)

1029 - actual = op.forward_hip(broadcast_residual, broadcast_fn, *args[2:], x=x)

1033 + # forward_hip leads with the residual; no post was requested here, so it is

1034 + # the caller's own tensor and the reference's four outputs follow it.

1035 + residual_out, *actual = op.forward_hip(

1036 + broadcast_residual, broadcast_fn, *args[2:], x=x

1037 + )

1038 + assert residual_out is broadcast_residual

1039 assert not took_aiter, "the broadcast seam must not reach AITER"

1040 for i in range(4):

1041 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):

1277 assert layer.hc_attn_fn_broadcast is buffer

1278 expected = layer.hc_attn_fn.detach().view(-1, 2, 8).sum(dim=1)

1279 assert torch.equal(layer.hc_attn_fn_broadcast, expected)

1280 +

1281 +

1282 +@pytest.mark.skipif(

1283 + not (current_platform.is_rocm() and HAS_AITER_MHC),

1284 + reason="AITER mHC required",

1285 +)

1286 +@pytest.mark.parametrize("num_tokens", [1, 2, 7, 128])

1287 +@pytest.mark.parametrize("carried", [False, True])

1288 +def test_mhc_fused_post_pre_delayed_rocm_aiter(

1289 + num_tokens, carried, default_vllm_config

1290 +):

1291 + """Folding the post into the pre projection must match doing them apart.

1292 +

1293 + The residual is bfloat16, so the two accumulation orders are allowed to

1294 + land a rounding step apart; the tolerance here is one bfloat16 ULP.

1295 + """

1296 + set_random_seed(0)

1297 + hc_mult, hidden_size = 4, 5120

1298 + residual, fn, hc_scale, hc_base, _ = _rocm_mhc_inputs(

1299 + num_tokens=num_tokens, hidden_size=hidden_size, hc_mult=hc_mult

1300 + )

1301 + x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16, device=DEVICE)

1302 + post_layer_mix = (

1303 + torch.rand(num_tokens, hc_mult, 1, dtype=torch.float32, device=DEVICE) + 0.5

1304 + )

1305 + comb_res_mix = (

1306 + torch.rand(num_tokens, hc_mult, hc_mult, dtype=torch.float32, device=DEVICE)

1307 + + 0.2

1308 + )

1309 + pre_mix = (

1310 + torch.rand(num_tokens, hc_mult, dtype=torch.float32, device=DEVICE) + 0.5

1311 + if carried

1312 + else None

1313 + )

1314 + rms_eps = hc_pre_eps = hc_sinkhorn_eps = 1e-6

1315 + pre_args = (fn, hc_scale, hc_base, rms_eps, hc_pre_eps, hc_sinkhorn_eps, 1.0, 20)

1316 +

1317 + next_residual = mhc_post_torch(x, residual, post_layer_mix, comb_res_mix)

1318 + expected = (

1319 + next_residual,

1320 + *mhc_pre_delayed_torch(next_residual, *pre_args, pre_mix=pre_mix),

1321 + )

1322 + actual = MHCPreDelayedOp().forward_hip(

1323 + residual,

1324 + *pre_args,

1325 + pre_mix=pre_mix,

1326 + sublayer_out=x,

1327 + post_layer_mix=post_layer_mix,

1328 + comb_res_mix=comb_res_mix,

1329 + )

1330 +

1331 + for got, want in zip(actual, expected, strict=True):

1332 + torch.testing.assert_close(got, want.reshape(got.shape), atol=2e-2, rtol=8e-3)

1333 +

1334 +

1335 +@pytest.mark.skipif(

1336 + not (current_platform.is_rocm() and HAS_AITER_MHC),

1337 + reason="AITER mHC required",

1338 +)

1339 +def test_mhc_fused_post_pre_delayed_falls_back_for_large_batches():

1340 + """AITER's heuristic hands large batches back to the unfused path."""

1341 + from vllm._aiter_ops import rocm_aiter_ops

1342 +

1343 + assert not rocm_aiter_ops.mhc_fused_post_pre_delayed_prefers_unfused(1)

1344 + assert rocm_aiter_ops.mhc_fused_post_pre_delayed_prefers_unfused(1 << 20)

vllm/_aiter_ops.py

+138 -37

@@ -1676,6 +1676,75 @@ def _rocm_aiter_fp8_attn_fake(

1676 )

1677

1678

1679 +def _mhc_delayed_pre_tail(

1680 + gemm_out: torch.Tensor,

1681 + sqrsum: torch.Tensor,

1682 + residual: torch.Tensor,

1683 + hc_scale: torch.Tensor,

1684 + hc_base: torch.Tensor,

1685 + pre_mix: torch.Tensor | None,

1686 + rms_eps: float,

1687 + hc_pre_eps: float,

1688 + hc_sinkhorn_eps: float,

1689 + hc_post_mult_value: float,

1690 + sinkhorn_repeat: int,

1691 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

1692 + """Finish a delayed mHC pre once AITER has produced the projection.

1693 +

1694 + Shared by the plain and the post-fused entry points, which differ only in

1695 + which kernel computes ``gemm_out`` / ``sqrsum`` and the residual they

1696 + project. ``mhc_pre_big_fuse`` also writes a collapse against its own

1697 + pre-mix, which the delayed formulation cannot use and so discards.

1698 + Returns ``(post_mix, comb_mix, layer_input, next_pre_mix)``.

1699 + """

1700 + from aiter.ops.mhc import mhc_pre_big_fuse

1701 +

1702 + num_tokens, hc_mult, hidden_size = residual.shape

1703 + device = residual.device

1704 + with torch.device(device):

1705 + post_mix = torch.empty(

1706 + num_tokens, hc_mult, 1, dtype=torch.float32, device=device

1707 + )

1708 + comb_mix = torch.empty(

1709 + num_tokens, hc_mult, hc_mult, dtype=torch.float32, device=device

1710 + )

1711 + unused_collapse = torch.empty(

1712 + num_tokens, hidden_size, dtype=torch.bfloat16, device=device

1713 + )

1714 + mhc_pre_big_fuse(

1715 + post_mix,

1716 + comb_mix,

1717 + unused_collapse,

1718 + gemm_out,

1719 + sqrsum,

1720 + hc_scale,

1721 + hc_base,

1722 + residual,

1723 + rms_eps,

1724 + hc_pre_eps,

1725 + hc_sinkhorn_eps,

1726 + hc_post_mult_value,

1727 + sinkhorn_repeat,

1728 + )

1729 +

1730 + next_pre_mix = torch.ops.vllm.mhc_pre_mix_triton(

1731 + gemm_out,

1732 + sqrsum,

1733 + hc_scale,

1734 + hc_base,

1735 + hc_mult,

1736 + hc_mult * hidden_size,

1737 + rms_eps,

1738 + hc_pre_eps,

1739 + )

1740 + if pre_mix is None:

1741 + # Model entry selects residual stream zero.

1742 + layer_input = residual[:, 0]

1743 + else:

1744 + layer_input = torch.ops.vllm.hc_collapse_triton(residual, pre_mix)

1745 + return post_mix, comb_mix, layer_input, next_pre_mix

1746 +

1747 +

1748 # Global flag to ensure ops are registered only once

1749 _OPS_REGISTERED = False

1750

@@ -3611,6 +3680,10 @@ class rocm_aiter_ops:

3680 hc_post_mult_value: float,

3681 sinkhorn_repeat: int,

3682 pre_mix: torch.Tensor | None = None,

3683 + sublayer_out: torch.Tensor | None = None,

3684 + post_layer_mix: torch.Tensor | None = None,

3685 + comb_res_mix: torch.Tensor | None = None,

3686 + residual_out: torch.Tensor | None = None,

3687 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

3688 """mHC pre using the pre-mix carried from the previous sublayer.

3689

@@ -3625,6 +3698,15 @@ class rocm_aiter_ops:

3698 that redundant store is the price of not having a native delayed

3699 kernel, and is small next to the ~140 launches it replaces.

3700

3701 + Passing ``sublayer_out`` / ``post_layer_mix`` / ``comb_res_mix``

3702 + also applies the preceding post block. AITER then projects the new

3703 + residual in the same kernel that computes it

3704 + (``mhc_fused_post_pre_gemm_sqrsum``), which saves one launch; every

3705 + later stage is identical either way. That new residual is written to

3706 + ``residual_out``, which the caller must supply in that case; it is an

3707 + output buffer rather than a return value so the op never has to hand

3708 + back one of its own inputs on the unfused path.

3709 +

3710 Returns:

3711 post_mix: shape (..., hc_mult, 1), dtype torch.float32

3712 comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32

@@ -3632,8 +3714,9 @@ class rocm_aiter_ops:

3714 next_pre_mix: shape (..., hc_mult), dtype torch.float32

3715 """

3716 from aiter.ops.mhc import (

3717 + get_mhc_fused_post_pre_config,

3718 get_mhc_pre_splitk,

3636 - mhc_pre_big_fuse,

3719 + mhc_fused_post_pre_gemm_sqrsum,

3720 mhc_pre_gemm_sqrsum,

3721 )

3722

@@ -3666,9 +3749,16 @@ class rocm_aiter_ops:

3749

3750 pre_mix_flat = None if pre_mix is None else pre_mix.view(-1, hc_mult)

3751

3752 + fuse_post = sublayer_out is not None

3753 # AITER's Python wrappers allocate without explicit device arguments.

3754 with torch.device(device):

3671 - splitk, tile_k = get_mhc_pre_splitk(num_tokens, hc_hidden_size)

3755 + if fuse_post:

3756 + # Keyed on the per-stream K, i.e. hidden_size, not hc_mult * it.

3757 + splitk, tile_m, tile_n, tile_k = get_mhc_fused_post_pre_config(

3758 + num_tokens, hidden_size

3759 + )

3760 + else:

3761 + splitk, tile_k = get_mhc_pre_splitk(num_tokens, hc_hidden_size)

3762 # AITER pads the GEMM output to a multiple of 32 columns.

3763 gemm_pad = torch.empty(

3764 splitk,

@@ -3679,50 +3769,46 @@ class rocm_aiter_ops:

3769 )

3770 gemm_out = gemm_pad[:, :, :hc_mult3]

3771 sqrsum = torch.empty(splitk, num_tokens, dtype=torch.float32, device=device)

3682 - mhc_pre_gemm_sqrsum(gemm_out, sqrsum, residual_flat, fn, tile_k, 0)

3683 -

3684 - post_mix = torch.empty(

3685 - num_tokens, hc_mult, 1, dtype=torch.float32, device=device

3686 - )

3687 - comb_mix = torch.empty(

3688 - num_tokens, hc_mult, hc_mult, dtype=torch.float32, device=device

3689 - )

3690 - unused_collapse = torch.empty(

3691 - num_tokens, hidden_size, dtype=torch.bfloat16, device=device

3692 - )

3693 - mhc_pre_big_fuse(

3694 - post_mix,

3695 - comb_mix,

3696 - unused_collapse,

3697 - gemm_out,

3698 - sqrsum,

3699 - hc_scale,

3700 - hc_base,

3701 - residual_flat,

3702 - rms_eps,

3703 - hc_pre_eps,

3704 - hc_sinkhorn_eps,

3705 - hc_post_mult_value,

3706 - sinkhorn_repeat,

3707 - )

3772 + if fuse_post:

3773 + # Restated for the type checker, which cannot narrow the

3774 + # optional arguments through the fuse_post flag.

3775 + assert sublayer_out is not None

3776 + assert post_layer_mix is not None and comb_res_mix is not None

3777 + assert residual_out is not None

3778 + projected = residual_out.view(-1, hc_mult, hidden_size)

3779 + mhc_fused_post_pre_gemm_sqrsum(

3780 + gemm_out,

3781 + sqrsum,

3782 + projected,

3783 + sublayer_out.view(-1, hidden_size),

3784 + residual_flat,

3785 + # The kernel wants the post mix squeezed to (tokens, hc_mult).

3786 + post_layer_mix.view(-1, hc_mult, 1).squeeze(-1),

3787 + comb_res_mix.view(-1, hc_mult, hc_mult),

3788 + fn,

3789 + tile_m,

3790 + tile_n,

3791 + tile_k,

3792 + 0,

3793 + )

3794 + else:

3795 + projected = residual_flat

3796 + mhc_pre_gemm_sqrsum(gemm_out, sqrsum, projected, fn, tile_k, 0)

3797

3709 - next_pre_mix = torch.ops.vllm.mhc_pre_mix_triton(

3798 + post_mix, comb_mix, layer_input, next_pre_mix = _mhc_delayed_pre_tail(

3799 gemm_out,

3800 sqrsum,

3801 + projected,

3802 hc_scale,

3803 hc_base,

3714 - hc_mult,

3715 - hc_hidden_size,

3804 + pre_mix_flat,

3805 rms_eps,

3806 hc_pre_eps,

3807 + hc_sinkhorn_eps,

3808 + hc_post_mult_value,

3809 + sinkhorn_repeat,

3810 )

3811

3720 - if pre_mix_flat is None:

3721 - # Model entry selects residual stream zero.

3722 - layer_input = residual_flat[:, 0]

3723 - else:

3724 - layer_input = torch.ops.vllm.hc_collapse_triton(residual_flat, pre_mix_flat)

3725 -

3812 return (

3813 post_mix.view(*outer_shape, hc_mult, 1),

3814 comb_mix.view(*outer_shape, hc_mult, hc_mult),

@@ -3730,6 +3816,21 @@ class rocm_aiter_ops:

3816 next_pre_mix.view(*outer_shape, hc_mult),

3817 )

3818

3819 + @staticmethod

3820 + def mhc_fused_post_pre_delayed_prefers_unfused(num_tokens: int) -> bool:

3821 + """True when AITER's own heuristic favours separate post and pre.

3822 +

3823 + Mirrors the ``fused_m_upper_bound`` table in ``aiter.ops.mhc``: the

3824 + fused GEMM wins for decode-sized batches and loses for large ones.

3825 + """

3826 + from aiter.jit.utils.chip_info import get_gfx_runtime

3827 +

3828 + return num_tokens >= {

3829 + "gfx950": 1024,

3830 + "gfx942": 128,

3831 + "gfx1250": 1024,

3832 + }.get(get_gfx_runtime(), 1024)

3833 +

3834 @staticmethod

3835 def hc_head(

3836 hs_flat: torch.Tensor,

vllm/model_executor/kernels/mhc/aiter.py

+20 -1

@@ -115,6 +115,10 @@ def mhc_pre_delayed_aiter(

115 hc_post_mult_value: float,

116 sinkhorn_repeat: int,

117 pre_mix: torch.Tensor | None = None,

118 + sublayer_out: torch.Tensor | None = None,

119 + post_layer_mix: torch.Tensor | None = None,

120 + comb_res_mix: torch.Tensor | None = None,

121 + residual_out: torch.Tensor | None = None,

122 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

123 """mHC pre with the pre-mix carried in from the previous sublayer.

124

@@ -134,6 +138,13 @@ def mhc_pre_delayed_aiter(

138 sinkhorn_repeat: number of sinkhorn iterations

139 pre_mix: shape (..., hc_mult) from the previous sublayer, or None at

140 model entry to select residual stream zero.

141 + sublayer_out: attention or FFN output, shape (..., hidden_size). When

142 + given with the two mixes below, the preceding post block is

143 + applied here so AITER can fold it into the pre projection.

144 + post_layer_mix: shape (..., hc_mult, 1), post gate for that block.

145 + comb_res_mix: shape (..., hc_mult, hc_mult), residual comb for it.

146 + residual_out: shape (..., hc_mult, hidden_size), written with the post

147 + block's new residual. Required exactly when the post is folded in.

148

149 Returns:

150 post_mix: shape (..., hc_mult, 1), dtype torch.float32

@@ -156,6 +167,10 @@ def mhc_pre_delayed_aiter(

167 hc_post_mult_value,

168 sinkhorn_repeat,

169 pre_mix,

170 + sublayer_out,

171 + post_layer_mix,

172 + comb_res_mix,

173 + residual_out,

174 )

175

176

@@ -170,6 +185,10 @@ def _mhc_pre_delayed_aiter_fake(

185 hc_post_mult_value: float,

186 sinkhorn_repeat: int,

187 pre_mix: torch.Tensor | None = None,

188 + sublayer_out: torch.Tensor | None = None,

189 + post_layer_mix: torch.Tensor | None = None,

190 + comb_res_mix: torch.Tensor | None = None,

191 + residual_out: torch.Tensor | None = None,

192 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

193 hc_mult = residual.shape[-2]

194 hidden_size = residual.shape[-1]

@@ -313,7 +332,7 @@ direct_register_custom_op(

332 direct_register_custom_op(

333 op_name="mhc_pre_delayed_aiter",

334 op_func=mhc_pre_delayed_aiter,

316 - mutates_args=[],

335 + mutates_args=["residual_out"],

336 fake_impl=_mhc_pre_delayed_aiter_fake,

337 )

338 direct_register_custom_op(

vllm/model_executor/layers/mhc.py

+75 -7

@@ -299,7 +299,13 @@ class MHCPreDelayedOp(CustomOp):

299

300 Same gates as :class:`MHCPreOp`, but the stream collapse applies the

301 caller's ``pre_mix`` and this sublayer's pre-mix is returned for the next

302 - sublayer seam. Returns post_mix, comb_mix, layer_input, next_pre_mix.

302 + sublayer seam.

303 +

304 + Passing ``sublayer_out`` / ``post_layer_mix`` / ``comb_res_mix`` also

305 + applies the preceding post block, which lets AITER fold it into the pre

306 + projection and saves a launch. Returns residual, post_mix, comb_mix,

307 + layer_input, next_pre_mix; ``residual`` is the caller's own tensor when

308 + no post is requested.

309 """

310

311 # --8<-- [end:mhc_pre_delayed]

@@ -307,6 +313,24 @@ class MHCPreDelayedOp(CustomOp):

313 def enabled(cls) -> bool:

314 return True

315

316 + def __init__(self) -> None:

317 + super().__init__()

318 + # Built here, not lazily: CustomOp resolves dispatch against the

319 + # current vLLM config, which is only set during model construction.

320 + self._post = MHCPostOp()

321 +

322 + def _maybe_post(

323 + self,

324 + residual: torch.Tensor,

325 + sublayer_out: torch.Tensor | None,

326 + post_layer_mix: torch.Tensor | None,

327 + comb_res_mix: torch.Tensor | None,

328 + ) -> torch.Tensor:

329 + if sublayer_out is None:

330 + return residual

331 + assert post_layer_mix is not None and comb_res_mix is not None

332 + return self._post(sublayer_out, residual, post_layer_mix, comb_res_mix)

333 +

334 def forward_cuda(

335 self,

336 residual: torch.Tensor,

@@ -322,8 +346,14 @@ class MHCPreDelayedOp(CustomOp):

346 x: torch.Tensor | None = None,

347 norm_weight: torch.Tensor | None = None,

348 norm_eps: float = 1e-6,

325 - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

326 - return torch.ops.vllm.mhc_pre_delayed_tilelang(

349 + sublayer_out: torch.Tensor | None = None,

350 + post_layer_mix: torch.Tensor | None = None,

351 + comb_res_mix: torch.Tensor | None = None,

352 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

353 + residual = self._maybe_post(

354 + residual, sublayer_out, post_layer_mix, comb_res_mix

355 + )

356 + return residual, *torch.ops.vllm.mhc_pre_delayed_tilelang(

357 residual,

358 fn,

359 hc_scale,

@@ -354,7 +384,10 @@ class MHCPreDelayedOp(CustomOp):

384 x: torch.Tensor | None = None,

385 norm_weight: torch.Tensor | None = None,

386 norm_eps: float = 1e-6,

357 - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

387 + sublayer_out: torch.Tensor | None = None,

388 + post_layer_mix: torch.Tensor | None = None,

389 + comb_res_mix: torch.Tensor | None = None,

390 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

391 # The AITER delayed path drives mhc_pre_gemm_sqrsum against `residual`

392 # and folds no RMSNorm, so it cannot serve the model-entry broadcast

393 # (which projects a narrower `x`) or a fused norm. Both are handled by

@@ -364,7 +397,26 @@ class MHCPreDelayedOp(CustomOp):

397 and norm_weight is None

398 and _aiter_mhc_supported(residual, None, supports_norm=False)

399 ):

367 - return torch.ops.vllm.mhc_pre_delayed_aiter(

400 + from vllm._aiter_ops import rocm_aiter_ops

401 +

402 + num_tokens = residual.numel() // (residual.shape[-1] * residual.shape[-2])

403 + # Folding the post in only pays while the residual still fits in

404 + # cache; past that AITER's heuristic wants the separate kernels,

405 + # which can ask for a non-temporal store and a large-m split-k.

406 + if sublayer_out is not None and (

407 + rocm_aiter_ops.mhc_fused_post_pre_delayed_prefers_unfused(num_tokens)

408 + ):

409 + residual = self._maybe_post(

410 + residual, sublayer_out, post_layer_mix, comb_res_mix

411 + )

412 + sublayer_out = None

413 + # The op writes the folded post's residual into this buffer rather

414 + # than returning it, since on the unfused path it has none of its

415 + # own and an op must not hand back one of its inputs.

416 + residual_out = (

417 + torch.empty_like(residual) if sublayer_out is not None else None

418 + )

419 + rest = torch.ops.vllm.mhc_pre_delayed_aiter(

420 residual,

421 fn,

422 hc_scale,

@@ -375,9 +427,17 @@ class MHCPreDelayedOp(CustomOp):

427 hc_post_mult_value,

428 sinkhorn_repeat,

429 pre_mix,

430 + sublayer_out,

431 + post_layer_mix,

432 + comb_res_mix,

433 + residual_out,

434 )

435 + return (residual if residual_out is None else residual_out), *rest

436 + residual = self._maybe_post(

437 + residual, sublayer_out, post_layer_mix, comb_res_mix

438 + )

439 if HAS_TILELANG_MHC:

380 - return torch.ops.vllm.mhc_pre_delayed_tilelang(

440 + return residual, *torch.ops.vllm.mhc_pre_delayed_tilelang(

441 residual,

442 fn,

443 hc_scale,

@@ -392,6 +452,7 @@ class MHCPreDelayedOp(CustomOp):

452 norm_weight,

453 norm_eps,

454 )

455 + # The post is already applied, so it must not be requested again.

456 return self.forward_native(

457 residual,

458 fn,

@@ -423,7 +484,13 @@ class MHCPreDelayedOp(CustomOp):

484 x: torch.Tensor | None = None,

485 norm_weight: torch.Tensor | None = None,

486 norm_eps: float = 1e-6,

426 - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

487 + sublayer_out: torch.Tensor | None = None,

488 + post_layer_mix: torch.Tensor | None = None,

489 + comb_res_mix: torch.Tensor | None = None,

490 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

491 + residual = self._maybe_post(

492 + residual, sublayer_out, post_layer_mix, comb_res_mix

493 + )

494 post_mix, comb_mix, layer_input, next_pre_mix = (

495 mhc_kernels.mhc_pre_delayed_torch(

496 residual,

@@ -440,6 +507,7 @@ class MHCPreDelayedOp(CustomOp):

507 )

508 )

509 return (

510 + residual,

511 post_mix,

512 comb_mix,

513 _apply_mhc_norm(layer_input, norm_weight, norm_eps),

vllm/models/deepseek_v41/amd/model.py

+37 -17

@@ -274,7 +274,7 @@ class DeepseekV4DecoderLayer(nn.Module):

274 # copies and the identity pre-mix selects copy 0.

275 assert self.hc_attn_fn_broadcast is not None

276 residual = x.unsqueeze(1).expand(-1, self.hc_mult, -1).contiguous()

277 - post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed(

277 + residual, post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed(

278 residual,

279 self.hc_attn_fn_broadcast,

280 self.hc_attn_scale,

@@ -288,7 +288,7 @@ class DeepseekV4DecoderLayer(nn.Module):

288 )

289 else:

290 residual = x

291 - post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed(

291 + residual, post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed(

292 residual,

293 self.hc_attn_fn,

294 self.hc_attn_scale,

@@ -301,18 +301,7 @@ class DeepseekV4DecoderLayer(nn.Module):

301 pre_mix=pre_mix,

302 )

303 else:

304 - residual = self.mhc_post(x, residual, post_mix, res_mix)

305 - if self.engram is not None and engram_hashes is not None:

306 - # Engram injection happens between the previous sublayer's

307 - # post and this block's pre, on the full hc stream, so the

308 - # mix coefficients see the injected stream.

309 - residual = self.engram(

310 - residual,

311 - engram_hashes[:, self.engram.layer_hash_index],

312 - engram_mask,

313 - )

314 - post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed(

315 - residual,

304 + pre_args = (

305 self.hc_attn_fn,

306 self.hc_attn_scale,

307 self.hc_attn_base,

@@ -321,8 +310,37 @@ class DeepseekV4DecoderLayer(nn.Module):

310 self.hc_eps,

311 self.hc_post_alpha,

312 self.hc_sinkhorn_iters,

324 - pre_mix=pre_mix,

313 )

314 + if self.engram is not None and engram_hashes is not None:

315 + # Engram injection happens between the previous sublayer's

316 + # post and this block's pre, on the full hc stream, so the

317 + # mix coefficients see the injected stream. That read of the

318 + # residual is what stops the post from fusing into the pre

319 + # here, unlike the attention seam below.

320 + residual = self.mhc_post(x, residual, post_mix, res_mix)

321 + residual = self.engram(

322 + residual,

323 + engram_hashes[:, self.engram.layer_hash_index],

324 + engram_mask,

325 + )

326 + residual, post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed(

327 + residual, *pre_args, pre_mix=pre_mix

328 + )

329 + else:

330 + (

331 + residual,

332 + post_mix,

333 + res_mix,

334 + x,

335 + attn_pre,

336 + ) = self.mhc_pre_delayed(

337 + residual,

338 + *pre_args,

339 + pre_mix=pre_mix,

340 + sublayer_out=x,

341 + post_layer_mix=post_mix,

342 + comb_res_mix=res_mix,

343 + )

344 x = self.attn_norm(x)

345

346 if self.use_sequence_parallel:

@@ -332,8 +350,7 @@ class DeepseekV4DecoderLayer(nn.Module):

350 if self.use_sequence_parallel:

351 x = sp_reduce_scatter(x)

352

335 - residual = self.mhc_post(x, residual, post_mix, res_mix)

336 - post_mix, res_mix, x, ffn_pre = self.mhc_pre_delayed(

353 + residual, post_mix, res_mix, x, ffn_pre = self.mhc_pre_delayed(

354 residual,

355 self.hc_ffn_fn,

356 self.hc_ffn_scale,

@@ -344,6 +361,9 @@ class DeepseekV4DecoderLayer(nn.Module):

361 self.hc_post_alpha,

362 self.hc_sinkhorn_iters,

363 pre_mix=attn_pre,

364 + sublayer_out=x,

365 + post_layer_mix=post_mix,

366 + comb_res_mix=res_mix,

367 )

368 x = self.ffn_norm(x)

369 x = self.ffn(x, input_ids)