[ROCm][DSV4] Disable TileLang MHC dispatch on gfx942 (#45931)
Signed-off-by: Tuukka Sarvi <tuukka.sarvi@amd.com>
Tuukka Sarvi committed
Jun 22, 2026 at 12:26 UTC
89accad2cc9685bbd813ec0efab316b36cf123ca
4 files changed
+56
-22
tests/kernels/test_mhc_kernels.py
+11
-11
@@ -8,8 +8,8 @@ from vllm.model_executor.kernels.mhc.tilelang import (
8
_tilelang_hc_prenorm_gemm,
9
_torch_hc_prenorm_gemm,
10
)
11
+from vllm.model_executor.layers.mhc import HAS_TILELANG_MHC
12
from vllm.platforms import current_platform
12
-from vllm.utils.import_utils import has_tilelang
13
from vllm.utils.torch_utils import set_random_seed
14
15
DEVICE = current_platform.device_type
@@ -97,8 +97,8 @@ def hc_head_ref(
97
98
99
@pytest.mark.skipif(
100
- not (current_platform.is_cuda_alike() and has_tilelang()),
101
- reason="CUDA or ROCm and tilelang required",
100
+ not HAS_TILELANG_MHC,
101
+ reason="TileLang MHC support required",
102
)
103
@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
104
@pytest.mark.parametrize("hidden_size", [4096, 7168])
@@ -150,8 +150,8 @@ def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult):
150
151
152
@pytest.mark.skipif(
153
- not (current_platform.is_cuda_alike() and has_tilelang()),
154
- reason="CUDA or ROCm and tilelang required",
153
+ not HAS_TILELANG_MHC,
154
+ reason="TileLang MHC support required",
155
)
156
@pytest.mark.parametrize(
157
("num_tokens", "hidden_size"),
@@ -190,8 +190,8 @@ def test_hc_prenorm_gemm_tilelang(num_tokens, hidden_size):
190
191
192
@pytest.mark.skipif(
193
- not (current_platform.is_cuda_alike() and has_tilelang()),
194
- reason="CUDA or ROCm and tilelang required",
193
+ not HAS_TILELANG_MHC,
194
+ reason="TileLang MHC support required",
195
)
196
@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
197
@pytest.mark.parametrize("hidden_size", [4096, 7168])
@@ -217,8 +217,8 @@ def test_mhc_post_tilelang(num_tokens, hidden_size, hc_mult):
217
218
219
@pytest.mark.skipif(
220
- not (current_platform.is_cuda_alike() and has_tilelang()),
221
- reason="CUDA or ROCm and tilelang required",
220
+ not HAS_TILELANG_MHC,
221
+ reason="TileLang MHC support required",
222
)
223
@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
224
@pytest.mark.parametrize("hidden_size", [4096, 7168])
@@ -324,8 +324,8 @@ def test_hc_head_triton(num_tokens, hidden_size, hc_mult):
324
325
326
@pytest.mark.skipif(
327
- not (current_platform.is_cuda_alike() and has_tilelang()),
328
- reason="CUDA or ROCm and tilelang required",
327
+ not HAS_TILELANG_MHC,
328
+ reason="TileLang MHC support required",
329
)
330
@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
331
@pytest.mark.parametrize("hidden_size", [4096, 7168])
vllm/model_executor/layers/mhc.py
+40
-5
@@ -6,9 +6,25 @@ import torch
6
# import vllm.model_executor.kernels.mhc # noqa: F401
7
import vllm.model_executor.kernels.mhc as mhc_kernels
8
from vllm.model_executor.custom_op import CustomOp
9
+from vllm.platforms import current_platform
10
from vllm.utils.import_utils import has_tilelang
11
11
-HAS_TILELANG = has_tilelang()
12
+
13
+def _has_tilelang_mhc() -> bool:
14
+ if not has_tilelang():
15
+ return False
16
+ if current_platform.is_cuda():
17
+ return True
18
+ if current_platform.is_rocm():
19
+ from vllm.platforms.rocm import on_gfx942
20
+
21
+ # TileLang MHC currently produces incorrect results on gfx942. Keep
22
+ # gfx942 on the existing torch/triton fallbacks until that path is fixed.
23
+ return not on_gfx942()
24
+ return False
25
+
26
+
27
+HAS_TILELANG_MHC = _has_tilelang_mhc()
28
29
30
# --8<-- [start:mhc_pre]
@@ -89,7 +105,7 @@ class MHCPreOp(CustomOp):
105
# sinkhorn_repeat,
106
# )
107
# else:
92
- if HAS_TILELANG:
108
+ if HAS_TILELANG_MHC:
109
return torch.ops.vllm.mhc_pre_tilelang(
110
residual,
111
fn,
@@ -224,7 +240,7 @@ class MHCPostOp(CustomOp):
240
# comb_res_mix,
241
# )
242
# else:
227
- if HAS_TILELANG:
243
+ if HAS_TILELANG_MHC:
244
return torch.ops.vllm.mhc_post_tilelang(
245
x, residual, post_layer_mix, comb_res_mix
246
)
@@ -310,7 +326,7 @@ class HCHeadOp(CustomOp):
326
outer_shape = hidden_states.shape[:-2]
327
hs_flat = hidden_states.view(-1, hc_mult, hidden_size)
328
313
- if HAS_TILELANG:
329
+ if HAS_TILELANG_MHC:
330
out = torch.ops.vllm.hc_head_fused_kernel_tilelang(
331
hs_flat,
332
hc_fn,
@@ -447,7 +463,26 @@ class MHCFusedPostPreOp(CustomOp):
463
norm_weight: torch.Tensor | None = None,
464
norm_eps: float = 0.0,
465
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
450
- return torch.ops.vllm.mhc_fused_post_pre_tilelang(
466
+ if HAS_TILELANG_MHC:
467
+ return torch.ops.vllm.mhc_fused_post_pre_tilelang(
468
+ x,
469
+ residual,
470
+ post_layer_mix,
471
+ comb_res_mix,
472
+ fn,
473
+ hc_scale,
474
+ hc_base,
475
+ rms_eps,
476
+ hc_pre_eps,
477
+ hc_sinkhorn_eps,
478
+ hc_post_mult_value,
479
+ sinkhorn_repeat,
480
+ n_splits,
481
+ tile_n,
482
+ norm_weight,
483
+ norm_eps,
484
+ )
485
+ return self.forward_native(
486
x,
487
residual,
488
post_layer_mix,
vllm/models/deepseek_v4/amd/model.py
+3
-3
@@ -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_TILELANG_MHC,
31
HCHeadOp,
32
MHCFusedPostPreOp,
33
MHCPostOp,
@@ -51,7 +52,6 @@ from vllm.model_executor.models.utils import (
52
from vllm.models.deepseek_v4.amd.rocm import DeepseekV4ROCMAiterMLAAttention
53
from vllm.platforms import current_platform
54
from vllm.sequence import IntermediateTensors
54
-from vllm.utils.import_utils import has_tilelang
55
56
57
class DeepseekV4MLP(nn.Module):
@@ -303,7 +303,7 @@ class DeepseekV4DecoderLayer(nn.Module):
303
self.mhc_pre = MHCPreOp()
304
self.mhc_post = MHCPostOp()
305
self.mhc_fused_post_pre = MHCFusedPostPreOp()
306
- self.has_tilelang = has_tilelang()
306
+ self.has_tilelang = HAS_TILELANG_MHC
307
308
def hc_pre(
309
self,
@@ -513,7 +513,7 @@ class DeepseekV4Model(nn.Module):
513
requires_grad=False,
514
)
515
self.hc_head_op = HCHeadOp()
516
- self.has_tilelang = has_tilelang()
516
+ self.has_tilelang = HAS_TILELANG_MHC
517
# Pre-hc_head residual stream buffer for the MTP draft. Stable
518
# address (outside the cudagraph pool) so the copy_ in forward()
519
# refreshes it correctly across captured shapes.
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 HCHeadOp
31
+from vllm.model_executor.layers.mhc import HAS_TILELANG_MHC, HCHeadOp
32
from vllm.model_executor.layers.vocab_parallel_embedding import (
33
VocabParallelEmbedding,
34
)
@@ -42,7 +42,6 @@ from vllm.models.deepseek_v4.common.ops import (
42
)
43
from vllm.platforms import current_platform
44
from vllm.sequence import IntermediateTensors
45
-from vllm.utils.import_utils import has_tilelang
45
46
from .model import DeepseekV4DecoderLayer
47
@@ -124,7 +123,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
123
)
124
125
self.hc_head_op = HCHeadOp()
127
- self.has_tilelang = has_tilelang()
126
+ self.has_tilelang = HAS_TILELANG_MHC
127
128
def forward(
129
self,