[ROCm][DSV4] Use aiter mHC pre/post as the default ROCm path (#43950)
Signed-off-by: Fangzhou Ai <fangzhou.ai@amd.com> Signed-off-by: Fangzhou-Ai <fangzhouai@gmail.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Claude <noreply@anthropic.com>
Fangzhou Ai committed
Jul 1, 2026 at 04:27 UTC
ed41aa270a9e01320e55ff9a069a826acb764101
3 files changed
+32
-39
vllm/model_executor/layers/mhc.py
+24
-32
@@ -5,6 +5,7 @@ import torch
5
# this import will also register the custom ops
6
# import vllm.model_executor.kernels.mhc # noqa: F401
7
import vllm.model_executor.kernels.mhc as mhc_kernels
8
+from vllm._aiter_ops import is_aiter_found_and_supported
9
from vllm.model_executor.custom_op import CustomOp
10
from vllm.platforms import current_platform
11
from vllm.utils.import_utils import has_tilelang
@@ -25,6 +26,7 @@ def _has_tilelang_mhc() -> bool:
26
27
28
HAS_TILELANG_MHC = _has_tilelang_mhc()
29
+HAS_AITER_MHC = is_aiter_found_and_supported()
30
31
32
# --8<-- [start:mhc_pre]
@@ -87,25 +89,20 @@ class MHCPreOp(CustomOp):
89
norm_weight: torch.Tensor | None = None,
90
norm_eps: float = 0.0,
91
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
90
- # TODO: Reenable aiter after we are at the aiter
91
- # version that has this bugfix
92
- # https://github.com/ROCm/aiter/commit/b639cb63bcac4672dce33a731fad042a65cb3649
93
- # It has accuracy problem at large number of tokens.
94
- # hidden_size = residual.shape[-1]
95
- # if hidden_size % 256 == 0:
96
- # return torch.ops.vllm.mhc_pre_aiter(
97
- # residual,
98
- # fn,
99
- # hc_scale,
100
- # hc_base,
101
- # rms_eps,
102
- # hc_pre_eps,
103
- # hc_sinkhorn_eps,
104
- # hc_post_mult_value,
105
- # sinkhorn_repeat,
106
- # )
107
- # else:
108
- if HAS_TILELANG_MHC:
92
+ hidden_size = residual.shape[-1]
93
+ if HAS_AITER_MHC and hidden_size % 256 == 0:
94
+ return torch.ops.vllm.mhc_pre_aiter(
95
+ residual,
96
+ fn,
97
+ hc_scale,
98
+ hc_base,
99
+ rms_eps,
100
+ hc_pre_eps,
101
+ hc_sinkhorn_eps,
102
+ hc_post_mult_value,
103
+ sinkhorn_repeat,
104
+ )
105
+ elif HAS_TILELANG_MHC:
106
return torch.ops.vllm.mhc_pre_tilelang(
107
residual,
108
fn,
@@ -227,19 +224,14 @@ class MHCPostOp(CustomOp):
224
post_layer_mix: torch.Tensor,
225
comb_res_mix: torch.Tensor,
226
) -> torch.Tensor:
230
- # TODO: Reenable aiter after we are at the aiter
231
- # version that has this bugfix
232
- # https://github.com/ROCm/aiter/commit/b639cb63bcac4672dce33a731fad042a65cb3649
233
- # It has accuracy problem at large number of tokens.
234
- # hidden_size = residual.shape[-1]
235
- # if hidden_size % 256 == 0:
236
- # return torch.ops.vllm.mhc_post_aiter(
237
- # x,
238
- # residual,
239
- # post_layer_mix,
240
- # comb_res_mix,
241
- # )
242
- # else:
227
+ hidden_size = residual.shape[-1]
228
+ if HAS_AITER_MHC and hidden_size % 256 == 0:
229
+ return torch.ops.vllm.mhc_post_aiter(
230
+ x,
231
+ residual,
232
+ post_layer_mix,
233
+ comb_res_mix,
234
+ )
235
if HAS_TILELANG_MHC:
236
return torch.ops.vllm.mhc_post_tilelang(
237
x, residual, post_layer_mix, comb_res_mix
vllm/models/deepseek_v4/amd/model.py
+6
-4
@@ -27,6 +27,7 @@ from vllm.model_executor.layers.linear import (
27
)
28
from vllm.model_executor.layers.logits_processor import LogitsProcessor
29
from vllm.model_executor.layers.mhc import (
30
+ HAS_AITER_MHC,
31
HAS_TILELANG_MHC,
32
HCHeadOp,
33
MHCFusedPostPreOp,
@@ -303,7 +304,9 @@ class DeepseekV4DecoderLayer(nn.Module):
304
self.mhc_pre = MHCPreOp()
305
self.mhc_post = MHCPostOp()
306
self.mhc_fused_post_pre = MHCFusedPostPreOp()
306
- self.has_tilelang = HAS_TILELANG_MHC
307
+ self.use_fused_mhc = HAS_TILELANG_MHC and not (
308
+ HAS_AITER_MHC and self.hidden_size % 256 == 0
309
+ )
310
311
def hc_pre(
312
self,
@@ -425,7 +428,7 @@ class DeepseekV4DecoderLayer(nn.Module):
428
) -> tuple[
429
torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None
430
]:
428
- if not self.has_tilelang:
431
+ if not self.use_fused_mhc:
432
return self._forward_unfused_post_pre(
433
x, positions, input_ids, post_mix, res_mix, residual
434
)
@@ -513,7 +516,6 @@ class DeepseekV4Model(nn.Module):
516
requires_grad=False,
517
)
518
self.hc_head_op = HCHeadOp()
516
- self.has_tilelang = HAS_TILELANG_MHC
519
# Pre-hc_head residual stream buffer for the MTP draft. Stable
520
# address (outside the cudagraph pool) so the copy_ in forward()
521
# refreshes it correctly across captured shapes.
@@ -580,7 +582,7 @@ class DeepseekV4Model(nn.Module):
582
res_mix,
583
residual,
584
)
583
- if layer is not None and self.has_tilelang:
585
+ if layer is not None and layer.use_fused_mhc:
586
hidden_states = layer.hc_post(hidden_states, residual, post_mix, res_mix)
587
588
if not get_pp_group().is_last_rank:
vllm/models/deepseek_v4/amd/mtp.py
+2
-3
@@ -28,7 +28,7 @@ from vllm.model_executor.layers.fused_moe import fused_moe_make_expert_params_ma
28
from vllm.model_executor.layers.layernorm import RMSNorm
29
from vllm.model_executor.layers.linear import ReplicatedLinear
30
from vllm.model_executor.layers.logits_processor import LogitsProcessor
31
-from vllm.model_executor.layers.mhc import HAS_TILELANG_MHC, HCHeadOp
31
+from vllm.model_executor.layers.mhc import HCHeadOp
32
from vllm.model_executor.layers.vocab_parallel_embedding import (
33
VocabParallelEmbedding,
34
)
@@ -123,7 +123,6 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
123
)
124
125
self.hc_head_op = HCHeadOp()
126
- self.has_tilelang = HAS_TILELANG_MHC
126
127
def forward(
128
self,
@@ -156,7 +155,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
155
hidden_states, residual, post_mix, res_mix = self.mtp_block(
156
positions=positions, x=hidden_states, input_ids=None
157
)
159
- if self.has_tilelang:
158
+ if self.mtp_block.use_fused_mhc:
159
hidden_states = self.mtp_block.hc_post(
160
hidden_states, residual, post_mix, res_mix
161
)