[ROCm][DSV4.1][Perf] Use AITER mHC for the delayed pre block

tests/kernels/test_mhc_kernels.py

@@ -23,6 +23,7 @@
     HAS_AITER_MHC_PRE_NORM,
     HAS_TILELANG_MHC,
     MHCFusedPostPreOp,
+    MHCPreDelayedOp,
     MHCPreOp,
 )
 from vllm.models.deepseek_v4.nvidia.model import (
@@ -697,6 +698,100 @@ def test_mhc_fused_rocm_fallback_applies_norm(monkeypatch):
     torch.testing.assert_close(out[3], expected_layer_input)
 
 
+@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, 1024])
+@pytest.mark.parametrize("carried", [False, True])
+def test_mhc_pre_delayed_rocm_aiter(num_tokens, carried):
+    """AITER must reproduce the delayed reference on both seam variants.
+
+    ``num_tokens`` spans the split-k choices AITER makes for the projection,
+    since the pre-mix is recovered from that unreduced output.
+    """
+    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
+    )
+    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
+    args = (
+        residual,
+        fn,
+        hc_scale,
+        hc_base,
+        rms_eps,
+        hc_pre_eps,
+        hc_sinkhorn_eps,
+        1.0,
+        20,
+    )
+
+    expected = mhc_pre_delayed_torch(*args, pre_mix=pre_mix)
+    actual = object.__new__(MHCPreDelayedOp).forward_hip(*args, pre_mix=pre_mix)
+
+    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.
+    torch.testing.assert_close(actual[2], expected[2], atol=0, rtol=0)
+
+
+@pytest.mark.skipif(
+    not (current_platform.is_rocm() and HAS_AITER_MHC),
+    reason="AITER mHC required",
+)
+def test_mhc_pre_delayed_rocm_aiter_declines_unsupported(monkeypatch):
+    """The broadcast seam and a fused norm must not take the AITER path.
+
+    Neither is expressible with AITER's pre kernels: the broadcast projects a
+    narrower ``x``, and ``mhc_pre_gemm_sqrsum`` folds no RMSNorm. What is
+    asserted here is the routing decision, which is what this gate owns; the
+    numerics of whichever fallback it lands on are covered by
+    ``test_deepseek_v41_mhc_pre_delayed``.
+    """
+    set_random_seed(0)
+    hc_mult, hidden_size = 4, 5120
+    residual, fn, hc_scale, hc_base, norm_weight = _rocm_mhc_inputs(
+        num_tokens=4, hidden_size=hidden_size, hc_mult=hc_mult
+    )
+    args = (residual, fn, hc_scale, hc_base, 1e-6, 1e-6, 1e-6, 1.0, 20)
+    op = object.__new__(MHCPreDelayedOp)
+
+    aiter_op = torch.ops.vllm.mhc_pre_delayed_aiter
+    took_aiter = False
+
+    def spy(*spy_args, **spy_kwargs):
+        nonlocal took_aiter
+        took_aiter = True
+        return aiter_op(*spy_args, **spy_kwargs)
+
+    monkeypatch.setattr(torch.ops.vllm, "mhc_pre_delayed_aiter", spy)
+
+    x = residual[:, 0].contiguous()
+    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)
+    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)
+
+    op.forward_hip(*args, norm_weight=norm_weight, norm_eps=1e-6)
+    assert not took_aiter, "a fused norm must not reach AITER"
+
+    # Positive control: the same inputs without a norm do take the AITER path,
+    # so the two declines above are the gate discriminating rather than the
+    # path being unavailable in this environment.
+    op.forward_hip(*args)
+    assert took_aiter
+
+
 @pytest.mark.skipif(
     not (current_platform.is_rocm() and HAS_AITER_MHC and HAS_AITER_MHC_PRE_NORM),
     reason="AITER mHC with fused RMSNorm required",

vllm/_aiter_ops.py

@@ -3599,6 +3599,137 @@ def mhc_pre(
             layer_input.view(*outer_shape, hidden_size),
         )
 
+    @staticmethod
+    def mhc_pre_delayed(
+        residual: 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,
+        pre_mix: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """mHC pre using the pre-mix carried from the previous sublayer.
+
+        Same gates as :meth:`mhc_pre`, but the stream collapse applies the
+        caller's ``pre_mix`` instead of the one computed here, and the one
+        computed here is returned for the next sublayer seam. AITER has no
+        single kernel for that shape, so drive its two stages directly:
+        ``mhc_pre_big_fuse`` still produces the post and comb gates (including
+        every sinkhorn iteration), while the pre gate is recovered from the
+        same split-k GEMM output and the collapse is done by the Triton
+        kernel. ``mhc_pre_big_fuse`` also writes a collapse we do not use;
+        that redundant store is the price of not having a native delayed
+        kernel, and is small next to the ~140 launches it replaces.
+
+        Returns:
+            post_mix: shape (..., hc_mult, 1), dtype torch.float32
+            comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
+            layer_input: shape (..., hidden_size), dtype torch.bfloat16
+            next_pre_mix: shape (..., hc_mult), dtype torch.float32
+        """
+        from aiter.ops.mhc import (
+            get_mhc_pre_splitk,
+            mhc_pre_big_fuse,
+            mhc_pre_gemm_sqrsum,
+        )
+
+        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]
+        hc_mult3 = hc_mult * 2 + hc_mult * hc_mult
+        hc_hidden_size = hc_mult * hidden_size
+
+        assert fn.shape == (hc_mult3, hc_hidden_size)
+        assert hc_scale.shape == (3,)
+        assert hc_base.shape == (hc_mult3,)
+
+        outer_shape = residual.shape[:-2]
+        residual_flat = residual.view(-1, hc_mult, hidden_size)
+        num_tokens = residual_flat.shape[0]
+        device = residual_flat.device
+
+        if num_tokens == 0:
+            return (
+                torch.empty(0, hc_mult, 1, dtype=torch.float32, device=device),
+                torch.empty(0, hc_mult, hc_mult, dtype=torch.float32, device=device),
+                torch.empty(0, hidden_size, dtype=torch.bfloat16, device=device),
+                torch.empty(0, hc_mult, dtype=torch.float32, device=device),
+            )
+
+        pre_mix_flat = None if pre_mix is None else pre_mix.view(-1, hc_mult)
+
+        # 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)
+            # AITER pads the GEMM output to a multiple of 32 columns.
+            gemm_pad = torch.empty(
+                splitk,
+                num_tokens,
+                (hc_mult3 + 31) // 32 * 32,
+                dtype=torch.float32,
+                device=device,
+            )
+            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)
+
+            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(
+            gemm_out,
+            sqrsum,
+            hc_scale,
+            hc_base,
+            hc_mult,
+            hc_hidden_size,
+            rms_eps,
+            hc_pre_eps,
+        )
+
+        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),
+            layer_input.view(*outer_shape, hidden_size),
+            next_pre_mix.view(*outer_shape, hc_mult),
+        )
+
     @staticmethod
     def hc_head(
         hs_flat: torch.Tensor,

vllm/model_executor/kernels/mhc/__init__.py

@@ -12,6 +12,7 @@
     "mhc_fused_post_pre_cuda",
     "hc_head_fused_kernel_cuda",
     "mhc_pre_aiter",
+    "mhc_pre_delayed_aiter",
     "mhc_post_aiter",
     "mhc_fused_post_pre_aiter",
     "hc_head_fused_aiter",
@@ -26,6 +27,7 @@
     "mhc_post_torch",
     "mhc_fused_post_pre_torch",
     "hc_head_fused_torch",
+    "mhc_pre_mix_triton",
     "mhc_pre_triton",
     "mhc_post_triton",
     "mhc_fused_post_pre_triton",

vllm/model_executor/kernels/mhc/aiter.py

@@ -104,6 +104,90 @@ def _mhc_pre_aiter_fake(
     return post_mix, comb_mix, layer_input
 
 
+def mhc_pre_delayed_aiter(
+    residual: 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,
+    pre_mix: 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.
+
+    Matches ``mhc_pre_delayed_torch``: the stream collapse uses *pre_mix*
+    rather than the gate computed here, and that gate is returned as the
+    pre-mix for the next sublayer seam.
+
+    Args:
+        residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16
+        fn: shape (hc_mult3, hc_mult * hidden_size), dtype torch.float32
+        hc_scale: shape (3,), dtype torch.float32
+        hc_base: shape (hc_mult3,), dtype torch.float32
+        rms_eps: RMS normalization epsilon
+        hc_pre_eps: pre-mix epsilon
+        hc_sinkhorn_eps: sinkhorn epsilon
+        hc_post_mult_value: post-mix multiplier value
+        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.
+
+    Returns:
+        post_mix: shape (..., hc_mult, 1), dtype torch.float32
+        comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
+        layer_input: shape (..., hidden_size), dtype torch.bfloat16
+        next_pre_mix: shape (..., hc_mult), dtype torch.float32
+    """
+    hidden_size = residual.shape[-1]
+    assert hidden_size % 256 == 0
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    return rocm_aiter_ops.mhc_pre_delayed(
+        residual,
+        fn,
+        hc_scale,
+        hc_base,
+        rms_eps,
+        hc_pre_eps,
+        hc_sinkhorn_eps,
+        hc_post_mult_value,
+        sinkhorn_repeat,
+        pre_mix,
+    )
+
+
+def _mhc_pre_delayed_aiter_fake(
+    residual: 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,
+    pre_mix: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    hc_mult = residual.shape[-2]
+    hidden_size = residual.shape[-1]
+    outer_shape = residual.shape[:-2]
+    return (
+        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
+        ),
+        torch.empty(*outer_shape, hc_mult, dtype=torch.float32, device=residual.device),
+    )
+
+
 def mhc_post_aiter(
     x: torch.Tensor,
     residual: torch.Tensor,
@@ -226,6 +310,12 @@ def _mhc_fused_post_pre_aiter_fake(
     mutates_args=[],
     fake_impl=_mhc_pre_aiter_fake,
 )
+direct_register_custom_op(
+    op_name="mhc_pre_delayed_aiter",
+    op_func=mhc_pre_delayed_aiter,
+    mutates_args=[],
+    fake_impl=_mhc_pre_delayed_aiter_fake,
+)
 direct_register_custom_op(
     op_name="mhc_post_aiter",
     op_func=mhc_post_aiter,