@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 46d2b23ac5

[ROCm][DSV4.1][Perf] Use AITER mHC for the delayed pre block (#56503)

Signed-off-by: Fangzhou Ai <fangzhou@semianalysis.ai> Co-authored-by: Fangzhou Ai <fangzhou@semianalysis.ai>

Fangzhou Ai committed Sep 11, 2026 at 20:35 UTC 46d2b23ac5047a813ebb082122166e4ae09b5f39
7 files changed +620 -19
tests/kernels/test_mhc_kernels.py
+95
@@ -23,6 +23,7 @@ from vllm.model_executor.layers.mhc import (
23 HAS_AITER_MHC_PRE_NORM,
24 HAS_TILELANG_MHC,
25 MHCFusedPostPreOp,
26 + MHCPreDelayedOp,
27 MHCPreOp,
28 )
29 from vllm.models.deepseek_v4.nvidia.model import (
@@ -697,6 +698,100 @@ def test_mhc_fused_rocm_fallback_applies_norm(monkeypatch):
698 torch.testing.assert_close(out[3], expected_layer_input)
699
700
701 +@pytest.mark.skipif(
702 + not (current_platform.is_rocm() and HAS_AITER_MHC),
703 + reason="AITER mHC required",
704 +)
705 +@pytest.mark.parametrize("num_tokens", [1, 2, 7, 128, 1024])
706 +@pytest.mark.parametrize("carried", [False, True])
707 +def test_mhc_pre_delayed_rocm_aiter(num_tokens, carried):
708 + """AITER must reproduce the delayed reference on both seam variants.
709 +
710 + ``num_tokens`` spans the split-k choices AITER makes for the projection,
711 + since the pre-mix is recovered from that unreduced output.
712 + """
713 + set_random_seed(0)
714 + hc_mult, hidden_size = 4, 5120
715 + residual, fn, hc_scale, hc_base, _ = _rocm_mhc_inputs(
716 + num_tokens=num_tokens, hidden_size=hidden_size, hc_mult=hc_mult
717 + )
718 + pre_mix = (
719 + torch.rand(num_tokens, hc_mult, dtype=torch.float32, device=DEVICE) + 0.5
720 + if carried
721 + else None
722 + )
723 + rms_eps = hc_pre_eps = hc_sinkhorn_eps = 1e-6
724 + args = (
725 + residual,
726 + fn,
727 + hc_scale,
728 + hc_base,
729 + rms_eps,
730 + hc_pre_eps,
731 + hc_sinkhorn_eps,
732 + 1.0,
733 + 20,
734 + )
735 +
736 + expected = mhc_pre_delayed_torch(*args, pre_mix=pre_mix)
737 + actual = object.__new__(MHCPreDelayedOp).forward_hip(*args, pre_mix=pre_mix)
738 +
739 + for i in (0, 1, 3):
740 + torch.testing.assert_close(actual[i], expected[i], atol=1e-4, rtol=1e-3)
741 + # The collapse is the same FP32 multiply-and-sum in both paths.
742 + torch.testing.assert_close(actual[2], expected[2], atol=0, rtol=0)
743 +
744 +
745 +@pytest.mark.skipif(
746 + not (current_platform.is_rocm() and HAS_AITER_MHC),
747 + reason="AITER mHC required",
748 +)
749 +def test_mhc_pre_delayed_rocm_aiter_declines_unsupported(monkeypatch):
750 + """The broadcast seam and a fused norm must not take the AITER path.
751 +
752 + Neither is expressible with AITER's pre kernels: the broadcast projects a
753 + narrower ``x``, and ``mhc_pre_gemm_sqrsum`` folds no RMSNorm. What is
754 + asserted here is the routing decision, which is what this gate owns; the
755 + numerics of whichever fallback it lands on are covered by
756 + ``test_deepseek_v41_mhc_pre_delayed``.
757 + """
758 + set_random_seed(0)
759 + hc_mult, hidden_size = 4, 5120
760 + residual, fn, hc_scale, hc_base, norm_weight = _rocm_mhc_inputs(
761 + num_tokens=4, hidden_size=hidden_size, hc_mult=hc_mult
762 + )
763 + args = (residual, fn, hc_scale, hc_base, 1e-6, 1e-6, 1e-6, 1.0, 20)
764 + op = object.__new__(MHCPreDelayedOp)
765 +
766 + aiter_op = torch.ops.vllm.mhc_pre_delayed_aiter
767 + took_aiter = False
768 +
769 + def spy(*spy_args, **spy_kwargs):
770 + nonlocal took_aiter
771 + took_aiter = True
772 + return aiter_op(*spy_args, **spy_kwargs)
773 +
774 + monkeypatch.setattr(torch.ops.vllm, "mhc_pre_delayed_aiter", spy)
775 +
776 + x = residual[:, 0].contiguous()
777 + broadcast_fn = fn.view(-1, hc_mult, hidden_size).sum(1)
778 + broadcast_residual = x.unsqueeze(1).expand(-1, hc_mult, -1).contiguous()
779 + expected = mhc_pre_delayed_torch(broadcast_residual, broadcast_fn, *args[2:], x=x)
780 + actual = op.forward_hip(broadcast_residual, broadcast_fn, *args[2:], x=x)
781 + assert not took_aiter, "the broadcast seam must not reach AITER"
782 + for i in range(4):
783 + torch.testing.assert_close(actual[i], expected[i], atol=1e-4, rtol=1e-3)
784 +
785 + op.forward_hip(*args, norm_weight=norm_weight, norm_eps=1e-6)
786 + assert not took_aiter, "a fused norm must not reach AITER"
787 +
788 + # Positive control: the same inputs without a norm do take the AITER path,
789 + # so the two declines above are the gate discriminating rather than the
790 + # path being unavailable in this environment.
791 + op.forward_hip(*args)
792 + assert took_aiter
793 +
794 +
795 @pytest.mark.skipif(
796 not (current_platform.is_rocm() and HAS_AITER_MHC and HAS_AITER_MHC_PRE_NORM),
797 reason="AITER mHC with fused RMSNorm required",
vllm/_aiter_ops.py
+131
@@ -3599,6 +3599,137 @@ class rocm_aiter_ops:
3599 layer_input.view(*outer_shape, hidden_size),
3600 )
3601
3602 + @staticmethod
3603 + def mhc_pre_delayed(
3604 + residual: torch.Tensor,
3605 + fn: torch.Tensor,
3606 + hc_scale: torch.Tensor,
3607 + hc_base: torch.Tensor,
3608 + rms_eps: float,
3609 + hc_pre_eps: float,
3610 + hc_sinkhorn_eps: float,
3611 + hc_post_mult_value: float,
3612 + sinkhorn_repeat: int,
3613 + pre_mix: torch.Tensor | None = None,
3614 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
3615 + """mHC pre using the pre-mix carried from the previous sublayer.
3616 +
3617 + Same gates as :meth:`mhc_pre`, but the stream collapse applies the
3618 + caller's ``pre_mix`` instead of the one computed here, and the one
3619 + computed here is returned for the next sublayer seam. AITER has no
3620 + single kernel for that shape, so drive its two stages directly:
3621 + ``mhc_pre_big_fuse`` still produces the post and comb gates (including
3622 + every sinkhorn iteration), while the pre gate is recovered from the
3623 + same split-k GEMM output and the collapse is done by the Triton
3624 + kernel. ``mhc_pre_big_fuse`` also writes a collapse we do not use;
3625 + that redundant store is the price of not having a native delayed
3626 + kernel, and is small next to the ~140 launches it replaces.
3627 +
3628 + Returns:
3629 + post_mix: shape (..., hc_mult, 1), dtype torch.float32
3630 + comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
3631 + layer_input: shape (..., hidden_size), dtype torch.bfloat16
3632 + next_pre_mix: shape (..., hc_mult), dtype torch.float32
3633 + """
3634 + from aiter.ops.mhc import (
3635 + get_mhc_pre_splitk,
3636 + mhc_pre_big_fuse,
3637 + mhc_pre_gemm_sqrsum,
3638 + )
3639 +
3640 + assert residual.dtype == torch.bfloat16
3641 + assert fn.dtype == torch.float32
3642 + assert hc_scale.dtype == torch.float32
3643 + assert hc_base.dtype == torch.float32
3644 +
3645 + hc_mult = residual.shape[-2]
3646 + hidden_size = residual.shape[-1]
3647 + hc_mult3 = hc_mult * 2 + hc_mult * hc_mult
3648 + hc_hidden_size = hc_mult * hidden_size
3649 +
3650 + assert fn.shape == (hc_mult3, hc_hidden_size)
3651 + assert hc_scale.shape == (3,)
3652 + assert hc_base.shape == (hc_mult3,)
3653 +
3654 + outer_shape = residual.shape[:-2]
3655 + residual_flat = residual.view(-1, hc_mult, hidden_size)
3656 + num_tokens = residual_flat.shape[0]
3657 + device = residual_flat.device
3658 +
3659 + if num_tokens == 0:
3660 + return (
3661 + torch.empty(0, hc_mult, 1, dtype=torch.float32, device=device),
3662 + torch.empty(0, hc_mult, hc_mult, dtype=torch.float32, device=device),
3663 + torch.empty(0, hidden_size, dtype=torch.bfloat16, device=device),
3664 + torch.empty(0, hc_mult, dtype=torch.float32, device=device),
3665 + )
3666 +
3667 + pre_mix_flat = None if pre_mix is None else pre_mix.view(-1, hc_mult)
3668 +
3669 + # AITER's Python wrappers allocate without explicit device arguments.
3670 + with torch.device(device):
3671 + splitk, tile_k = get_mhc_pre_splitk(num_tokens, hc_hidden_size)
3672 + # AITER pads the GEMM output to a multiple of 32 columns.
3673 + gemm_pad = torch.empty(
3674 + splitk,
3675 + num_tokens,
3676 + (hc_mult3 + 31) // 32 * 32,
3677 + dtype=torch.float32,
3678 + device=device,
3679 + )
3680 + gemm_out = gemm_pad[:, :, :hc_mult3]
3681 + 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 + )
3708 +
3709 + next_pre_mix = torch.ops.vllm.mhc_pre_mix_triton(
3710 + gemm_out,
3711 + sqrsum,
3712 + hc_scale,
3713 + hc_base,
3714 + hc_mult,
3715 + hc_hidden_size,
3716 + rms_eps,
3717 + hc_pre_eps,
3718 + )
3719 +
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 +
3726 + return (
3727 + post_mix.view(*outer_shape, hc_mult, 1),
3728 + comb_mix.view(*outer_shape, hc_mult, hc_mult),
3729 + layer_input.view(*outer_shape, hidden_size),
3730 + next_pre_mix.view(*outer_shape, hc_mult),
3731 + )
3732 +
3733 @staticmethod
3734 def hc_head(
3735 hs_flat: torch.Tensor,
vllm/model_executor/kernels/mhc/__init__.py
+2
@@ -12,6 +12,7 @@ __all__ = [
12 "mhc_fused_post_pre_cuda",
13 "hc_head_fused_kernel_cuda",
14 "mhc_pre_aiter",
15 + "mhc_pre_delayed_aiter",
16 "mhc_post_aiter",
17 "mhc_fused_post_pre_aiter",
18 "hc_head_fused_aiter",
@@ -26,6 +27,7 @@ __all__ = [
27 "mhc_post_torch",
28 "mhc_fused_post_pre_torch",
29 "hc_head_fused_torch",
30 + "mhc_pre_mix_triton",
31 "mhc_pre_triton",
32 "mhc_post_triton",
33 "mhc_fused_post_pre_triton",
vllm/model_executor/kernels/mhc/aiter.py
+90
@@ -104,6 +104,90 @@ def _mhc_pre_aiter_fake(
104 return post_mix, comb_mix, layer_input
105
106
107 +def mhc_pre_delayed_aiter(
108 + residual: torch.Tensor,
109 + fn: torch.Tensor,
110 + hc_scale: torch.Tensor,
111 + hc_base: torch.Tensor,
112 + rms_eps: float,
113 + hc_pre_eps: float,
114 + hc_sinkhorn_eps: float,
115 + hc_post_mult_value: float,
116 + sinkhorn_repeat: int,
117 + pre_mix: torch.Tensor | None = None,
118 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
119 + """mHC pre with the pre-mix carried in from the previous sublayer.
120 +
121 + Matches ``mhc_pre_delayed_torch``: the stream collapse uses *pre_mix*
122 + rather than the gate computed here, and that gate is returned as the
123 + pre-mix for the next sublayer seam.
124 +
125 + Args:
126 + residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16
127 + fn: shape (hc_mult3, hc_mult * hidden_size), dtype torch.float32
128 + hc_scale: shape (3,), dtype torch.float32
129 + hc_base: shape (hc_mult3,), dtype torch.float32
130 + rms_eps: RMS normalization epsilon
131 + hc_pre_eps: pre-mix epsilon
132 + hc_sinkhorn_eps: sinkhorn epsilon
133 + hc_post_mult_value: post-mix multiplier value
134 + sinkhorn_repeat: number of sinkhorn iterations
135 + pre_mix: shape (..., hc_mult) from the previous sublayer, or None at
136 + model entry to select residual stream zero.
137 +
138 + Returns:
139 + post_mix: shape (..., hc_mult, 1), dtype torch.float32
140 + comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
141 + layer_input: shape (..., hidden_size), dtype torch.bfloat16
142 + next_pre_mix: shape (..., hc_mult), dtype torch.float32
143 + """
144 + hidden_size = residual.shape[-1]
145 + assert hidden_size % 256 == 0
146 + from vllm._aiter_ops import rocm_aiter_ops
147 +
148 + return rocm_aiter_ops.mhc_pre_delayed(
149 + residual,
150 + fn,
151 + hc_scale,
152 + hc_base,
153 + rms_eps,
154 + hc_pre_eps,
155 + hc_sinkhorn_eps,
156 + hc_post_mult_value,
157 + sinkhorn_repeat,
158 + pre_mix,
159 + )
160 +
161 +
162 +def _mhc_pre_delayed_aiter_fake(
163 + residual: torch.Tensor,
164 + fn: torch.Tensor,
165 + hc_scale: torch.Tensor,
166 + hc_base: torch.Tensor,
167 + rms_eps: float,
168 + hc_pre_eps: float,
169 + hc_sinkhorn_eps: float,
170 + hc_post_mult_value: float,
171 + sinkhorn_repeat: int,
172 + pre_mix: torch.Tensor | None = None,
173 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
174 + hc_mult = residual.shape[-2]
175 + hidden_size = residual.shape[-1]
176 + outer_shape = residual.shape[:-2]
177 + return (
178 + torch.empty(
179 + *outer_shape, hc_mult, 1, dtype=torch.float32, device=residual.device
180 + ),
181 + torch.empty(
182 + *outer_shape, hc_mult, hc_mult, dtype=torch.float32, device=residual.device
183 + ),
184 + torch.empty(
185 + *outer_shape, hidden_size, dtype=torch.bfloat16, device=residual.device
186 + ),
187 + torch.empty(*outer_shape, hc_mult, dtype=torch.float32, device=residual.device),
188 + )
189 +
190 +
191 def mhc_post_aiter(
192 x: torch.Tensor,
193 residual: torch.Tensor,
@@ -226,6 +310,12 @@ direct_register_custom_op(
310 mutates_args=[],
311 fake_impl=_mhc_pre_aiter_fake,
312 )
313 +direct_register_custom_op(
314 + op_name="mhc_pre_delayed_aiter",
315 + op_func=mhc_pre_delayed_aiter,
316 + mutates_args=[],
317 + fake_impl=_mhc_pre_delayed_aiter_fake,
318 +)
319 direct_register_custom_op(
320 op_name="mhc_post_aiter",
321 op_func=mhc_post_aiter,
vllm/model_executor/kernels/mhc/triton.py
+131
@@ -141,6 +141,129 @@ def _hc_collapse_triton_fake(x: Tensor, pre_mix: Tensor) -> Tensor:
141 return torch.empty(x.shape[0], x.shape[2], dtype=x.dtype, device=x.device)
142
143
144 +@triton.jit
145 +def _mhc_pre_mix_kernel(
146 + gemm_ptr,
147 + sqrsum_ptr,
148 + hc_scale_ptr,
149 + hc_base_ptr,
150 + out_ptr,
151 + splitk,
152 + hc_mult: tl.constexpr,
153 + gemm_stride_k,
154 + gemm_stride_t,
155 + gemm_stride_j,
156 + sqrsum_stride_k,
157 + sqrsum_stride_t,
158 + out_stride_t,
159 + out_stride_j,
160 + inv_hc_hidden,
161 + rms_eps,
162 + hc_pre_eps,
163 + SPLITK_BLOCK: tl.constexpr,
164 + HC_BLOCK: tl.constexpr,
165 +):
166 + """Recover the pre-mix gate from a split-k mHC pre GEMM output."""
167 + token_idx = tl.program_id(0).to(tl.int64)
168 + ks = tl.arange(0, SPLITK_BLOCK)
169 + js = tl.arange(0, HC_BLOCK)
170 + kmask = ks < splitk
171 + jmask = js < hc_mult
172 +
173 + # Only the first hc_mult GEMM columns feed the pre gate.
174 + gemm = tl.load(
175 + gemm_ptr
176 + + ks[:, None] * gemm_stride_k
177 + + token_idx * gemm_stride_t
178 + + js[None, :] * gemm_stride_j,
179 + mask=kmask[:, None] & jmask[None, :],
180 + other=0.0,
181 + ).to(tl.float32)
182 + mixes = tl.sum(gemm, 0)
183 +
184 + sqrsum = tl.load(
185 + sqrsum_ptr + ks * sqrsum_stride_k + token_idx * sqrsum_stride_t,
186 + mask=kmask,
187 + other=0.0,
188 + ).to(tl.float32)
189 + rstd = tl.rsqrt(tl.sum(sqrsum, 0) * inv_hc_hidden + rms_eps)
190 +
191 + scale = tl.load(hc_scale_ptr).to(tl.float32)
192 + base = tl.load(hc_base_ptr + js, mask=jmask, other=0.0).to(tl.float32)
193 +
194 + pre = tl.sigmoid(mixes * rstd * scale + base) + hc_pre_eps
195 + tl.store(out_ptr + token_idx * out_stride_t + js * out_stride_j, pre, mask=jmask)
196 +
197 +
198 +def mhc_pre_mix_triton(
199 + gemm_out: Tensor,
200 + sqrsum: Tensor,
201 + hc_scale: Tensor,
202 + hc_base: Tensor,
203 + hc_mult: int,
204 + hc_hidden_size: int,
205 + rms_eps: float,
206 + hc_pre_eps: float,
207 +) -> Tensor:
208 + """Pre-mix gate for the delayed mHC pre, from AITER's split-k GEMM output.
209 +
210 + AITER's ``mhc_pre_big_fuse`` consumes the unreduced ``[splitk, tokens,
211 + hc_mult3]`` GEMM output and the matching row square-sums, but only returns
212 + the post and comb gates. The delayed formulation also needs the pre gate,
213 + to carry into the next sublayer seam. It is the same slice of the same
214 + numbers, so recover it here rather than repeating the projection.
215 + """
216 + assert gemm_out.ndim == 3 and gemm_out.dtype == torch.float32
217 + assert sqrsum.ndim == 2 and sqrsum.dtype == torch.float32
218 + splitk, num_tokens = gemm_out.shape[0], gemm_out.shape[1]
219 + assert sqrsum.shape == (splitk, num_tokens)
220 +
221 + out = torch.empty(num_tokens, hc_mult, dtype=torch.float32, device=gemm_out.device)
222 + if num_tokens == 0:
223 + return out
224 +
225 + _mhc_pre_mix_kernel[(num_tokens,)](
226 + gemm_out,
227 + sqrsum,
228 + hc_scale,
229 + hc_base,
230 + out,
231 + splitk,
232 + hc_mult,
233 + gemm_out.stride(0),
234 + gemm_out.stride(1),
235 + gemm_out.stride(2),
236 + sqrsum.stride(0),
237 + sqrsum.stride(1),
238 + out.stride(0),
239 + out.stride(1),
240 + 1.0 / hc_hidden_size,
241 + rms_eps,
242 + hc_pre_eps,
243 + SPLITK_BLOCK=triton.next_power_of_2(splitk),
244 + HC_BLOCK=triton.next_power_of_2(hc_mult),
245 + num_warps=1,
246 + # Match the separate FP32 multiply and add in the Torch reference.
247 + enable_fp_fusion=False,
248 + )
249 + return out
250 +
251 +
252 +def _mhc_pre_mix_triton_fake(
253 + gemm_out: Tensor,
254 + sqrsum: Tensor,
255 + hc_scale: Tensor,
256 + hc_base: Tensor,
257 + hc_mult: int,
258 + hc_hidden_size: int,
259 + rms_eps: float,
260 + hc_pre_eps: float,
261 +) -> Tensor:
262 + return torch.empty(
263 + gemm_out.shape[1], hc_mult, dtype=torch.float32, device=gemm_out.device
264 + )
265 +
266 +
267 def hc_head_reduce_triton_kernel(
268 x: torch.Tensor,
269 hc_fn: torch.Tensor,
@@ -216,3 +339,11 @@ direct_register_custom_op(
339 mutates_args=[],
340 fake_impl=_hc_collapse_triton_fake,
341 )
342 +
343 +
344 +direct_register_custom_op(
345 + op_name="mhc_pre_mix_triton",
346 + op_func=mhc_pre_mix_triton,
347 + mutates_args=[],
348 + fake_impl=_mhc_pre_mix_triton_fake,
349 +)
vllm/model_executor/layers/mhc.py
+155
@@ -292,6 +292,161 @@ class MHCPreOp(CustomOp):
292 )
293
294
295 +# --8<-- [start:mhc_pre_delayed]
296 +@CustomOp.register("mhc_pre_delayed")
297 +class MHCPreDelayedOp(CustomOp):
298 + """MHC pre block using the pre-mix carried from the previous sublayer.
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.
303 + """
304 +
305 + # --8<-- [end:mhc_pre_delayed]
306 + @classmethod
307 + def enabled(cls) -> bool:
308 + return True
309 +
310 + def forward_cuda(
311 + self,
312 + residual: torch.Tensor,
313 + fn: torch.Tensor,
314 + hc_scale: torch.Tensor,
315 + hc_base: torch.Tensor,
316 + rms_eps: float,
317 + hc_pre_eps: float,
318 + hc_sinkhorn_eps: float,
319 + hc_post_mult_value: float,
320 + sinkhorn_repeat: int,
321 + pre_mix: torch.Tensor | None = None,
322 + x: torch.Tensor | None = None,
323 + norm_weight: torch.Tensor | None = None,
324 + norm_eps: float = 1e-6,
325 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
326 + return torch.ops.vllm.mhc_pre_delayed_tilelang(
327 + residual,
328 + fn,
329 + hc_scale,
330 + hc_base,
331 + rms_eps,
332 + hc_pre_eps,
333 + hc_sinkhorn_eps,
334 + hc_post_mult_value,
335 + sinkhorn_repeat,
336 + pre_mix,
337 + x,
338 + norm_weight,
339 + norm_eps,
340 + )
341 +
342 + def forward_hip(
343 + self,
344 + residual: torch.Tensor,
345 + fn: torch.Tensor,
346 + hc_scale: torch.Tensor,
347 + hc_base: torch.Tensor,
348 + rms_eps: float,
349 + hc_pre_eps: float,
350 + hc_sinkhorn_eps: float,
351 + hc_post_mult_value: float,
352 + sinkhorn_repeat: int,
353 + pre_mix: torch.Tensor | None = None,
354 + x: torch.Tensor | None = None,
355 + norm_weight: torch.Tensor | None = None,
356 + norm_eps: float = 1e-6,
357 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
358 + # The AITER delayed path drives mhc_pre_gemm_sqrsum against `residual`
359 + # and folds no RMSNorm, so it cannot serve the model-entry broadcast
360 + # (which projects a narrower `x`) or a fused norm. Both are handled by
361 + # TileLang, or by the reference when TileLang is unavailable.
362 + if (
363 + x is None
364 + and norm_weight is None
365 + and _aiter_mhc_supported(residual, None, supports_norm=False)
366 + ):
367 + return torch.ops.vllm.mhc_pre_delayed_aiter(
368 + residual,
369 + fn,
370 + hc_scale,
371 + hc_base,
372 + rms_eps,
373 + hc_pre_eps,
374 + hc_sinkhorn_eps,
375 + hc_post_mult_value,
376 + sinkhorn_repeat,
377 + pre_mix,
378 + )
379 + if HAS_TILELANG_MHC:
380 + return torch.ops.vllm.mhc_pre_delayed_tilelang(
381 + residual,
382 + fn,
383 + hc_scale,
384 + hc_base,
385 + rms_eps,
386 + hc_pre_eps,
387 + hc_sinkhorn_eps,
388 + hc_post_mult_value,
389 + sinkhorn_repeat,
390 + pre_mix,
391 + x,
392 + norm_weight,
393 + norm_eps,
394 + )
395 + return self.forward_native(
396 + residual,
397 + fn,
398 + hc_scale,
399 + hc_base,
400 + rms_eps,
401 + hc_pre_eps,
402 + hc_sinkhorn_eps,
403 + hc_post_mult_value,
404 + sinkhorn_repeat,
405 + pre_mix,
406 + x,
407 + norm_weight,
408 + norm_eps,
409 + )
410 +
411 + def forward_native(
412 + self,
413 + residual: torch.Tensor,
414 + fn: torch.Tensor,
415 + hc_scale: torch.Tensor,
416 + hc_base: torch.Tensor,
417 + rms_eps: float,
418 + hc_pre_eps: float,
419 + hc_sinkhorn_eps: float,
420 + hc_post_mult_value: float,
421 + sinkhorn_repeat: int,
422 + pre_mix: torch.Tensor | None = None,
423 + x: torch.Tensor | None = None,
424 + norm_weight: torch.Tensor | None = None,
425 + norm_eps: float = 1e-6,
426 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
427 + post_mix, comb_mix, layer_input, next_pre_mix = (
428 + mhc_kernels.mhc_pre_delayed_torch(
429 + residual,
430 + fn,
431 + hc_scale,
432 + hc_base,
433 + rms_eps,
434 + hc_pre_eps,
435 + hc_sinkhorn_eps,
436 + hc_post_mult_value,
437 + sinkhorn_repeat,
438 + pre_mix=pre_mix,
439 + x=x,
440 + )
441 + )
442 + return (
443 + post_mix,
444 + comb_mix,
445 + _apply_mhc_norm(layer_input, norm_weight, norm_eps),
446 + next_pre_mix,
447 + )
448 +
449 +
450 # --8<-- [start:mhc_post]
451 @CustomOp.register("mhc_post")
452 class MHCPostOp(CustomOp):
vllm/models/deepseek_v4_1/amd/model.py
+16 -19
@@ -17,15 +17,12 @@ from vllm.distributed import (
17 )
18 from vllm.forward_context import get_forward_context, is_forward_context_available
19 from vllm.logger import init_logger
20 -from vllm.model_executor.kernels.mhc.torch import (
21 - mhc_post_torch,
22 - mhc_pre_delayed_torch,
23 -)
20 from vllm.model_executor.layers.fused_moe import (
21 fused_moe_make_expert_params_mapping,
22 )
23 from vllm.model_executor.layers.layernorm import RMSNorm
24 from vllm.model_executor.layers.logits_processor import LogitsProcessor
25 +from vllm.model_executor.layers.mhc import MHCPostOp, MHCPreDelayedOp
26 from vllm.model_executor.layers.quantization import QuantizationConfig
27 from vllm.model_executor.layers.vocab_parallel_embedding import (
28 ParallelLMHead,
@@ -248,11 +245,13 @@ class DeepseekV4DecoderLayer(nn.Module):
245 ),
246 requires_grad=False,
247 )
248 + self.mhc_pre_delayed = MHCPreDelayedOp()
249 + self.mhc_post = MHCPostOp()
250
251 @staticmethod
252 def _hc_collapse(x: torch.Tensor, pre_mix: torch.Tensor) -> torch.Tensor:
254 - """Reference hyper-connection collapse used by DSpark on ROCm."""
255 - return (pre_mix.unsqueeze(-1) * x.float()).sum(dim=1).to(x.dtype)
253 + """Hyper-connection collapse used by DSpark on ROCm."""
254 + return torch.ops.vllm.hc_collapse_triton(x, pre_mix)
255
256 def forward(
257 self,
@@ -275,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()
278 - post_mix, res_mix, x, attn_pre = mhc_pre_delayed_torch(
277 + post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed(
278 residual,
279 self.hc_attn_fn_broadcast,
280 self.hc_attn_scale,
@@ -289,7 +288,7 @@ class DeepseekV4DecoderLayer(nn.Module):
288 )
289 else:
290 residual = x
292 - post_mix, res_mix, x, attn_pre = mhc_pre_delayed_torch(
291 + post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed(
292 residual,
293 self.hc_attn_fn,
294 self.hc_attn_scale,
@@ -302,7 +301,7 @@ class DeepseekV4DecoderLayer(nn.Module):
301 pre_mix=pre_mix,
302 )
303 else:
305 - residual = mhc_post_torch(x, residual, post_mix, res_mix)
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
@@ -312,7 +311,7 @@ class DeepseekV4DecoderLayer(nn.Module):
311 engram_hashes[:, self.engram.layer_hash_index],
312 engram_mask,
313 )
315 - post_mix, res_mix, x, attn_pre = mhc_pre_delayed_torch(
314 + post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed(
315 residual,
316 self.hc_attn_fn,
317 self.hc_attn_scale,
@@ -333,8 +332,8 @@ class DeepseekV4DecoderLayer(nn.Module):
332 if self.use_sequence_parallel:
333 x = sp_reduce_scatter(x)
334
336 - residual = mhc_post_torch(x, residual, post_mix, res_mix)
337 - post_mix, res_mix, x, ffn_pre = mhc_pre_delayed_torch(
335 + residual = self.mhc_post(x, residual, post_mix, res_mix)
336 + post_mix, res_mix, x, ffn_pre = self.mhc_pre_delayed(
337 residual,
338 self.hc_ffn_fn,
339 self.hc_ffn_scale,
@@ -453,6 +452,8 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
452 else:
453 self.norm = PPMissingLayer()
454
455 + self.mhc_post = MHCPostOp()
456 +
457 spec_config = vllm_config.speculative_config
458 needs_mtp_hidden_states = spec_config is not None and (
459 spec_config.use_eagle() or spec_config.uses_draft_model()
@@ -603,7 +604,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
604 )
605 if idx + 1 in self.aux_hidden_state_layers:
606 # Reconstruct the aux hidden state for draft models
606 - aux_recon = mhc_post_torch(hidden_states, residual, post_mix, res_mix)
607 + aux_recon = self.mhc_post(hidden_states, residual, post_mix, res_mix)
608 aux_hidden_state = aux_recon.mean(dim=1)
609 if self.use_sequence_parallel:
610 aux_hidden_state = sp_all_gather(aux_hidden_state)[:full_num_tokens]
@@ -614,7 +615,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
615 if self.end_layer in self.aux_hidden_state_layers:
616 hidden_states = final_aux_recon
617 else:
617 - hidden_states = mhc_post_torch(
618 + hidden_states = self.mhc_post(
619 hidden_states, residual, post_mix, res_mix
620 )
621
@@ -635,11 +636,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
636 # mixes — the mix the reference applies via
637 # ``last_layer.hc_pre(h, pre_mix)`` (v4.1 has no learned hc_head).
638 assert pre_mix is not None
638 - hidden_states = (
639 - (pre_mix.unsqueeze(-1) * hidden_states.float())
640 - .sum(dim=1)
641 - .to(hidden_states.dtype)
642 - )
639 + hidden_states = torch.ops.vllm.hc_collapse_triton(hidden_states, pre_mix)
640 hidden_states = self.norm(hidden_states)
641 if len(aux_hidden_states) > 0:
642 return hidden_states, aux_hidden_states