[ROCm][DSV4] Disable TileLang MHC dispatch on gfx942 (#45931)

Signed-off-by: Tuukka Sarvi <[email protected]>

tests/kernels/test_mhc_kernels.py

@@ -8,8 +8,8 @@ from vllm.model_executor.kernels.mhc.tilelang import (
     _tilelang_hc_prenorm_gemm,
     _torch_hc_prenorm_gemm,
 )
+from vllm.model_executor.layers.mhc import HAS_TILELANG_MHC
 from vllm.platforms import current_platform
-from vllm.utils.import_utils import has_tilelang
 from vllm.utils.torch_utils import set_random_seed
 
 DEVICE = current_platform.device_type
@@ -97,8 +97,8 @@ def hc_head_ref(
 
 
 @pytest.mark.skipif(
-    not (current_platform.is_cuda_alike() and has_tilelang()),
-    reason="CUDA or ROCm and tilelang required",
+    not HAS_TILELANG_MHC,
+    reason="TileLang MHC support required",
 )
 @pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
 @pytest.mark.parametrize("hidden_size", [4096, 7168])
@@ -150,8 +150,8 @@ def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult):
 
 
 @pytest.mark.skipif(
-    not (current_platform.is_cuda_alike() and has_tilelang()),
-    reason="CUDA or ROCm and tilelang required",
+    not HAS_TILELANG_MHC,
+    reason="TileLang MHC support required",
 )
 @pytest.mark.parametrize(
     ("num_tokens", "hidden_size"),
@@ -190,8 +190,8 @@ def test_hc_prenorm_gemm_tilelang(num_tokens, hidden_size):
 
 
 @pytest.mark.skipif(
-    not (current_platform.is_cuda_alike() and has_tilelang()),
-    reason="CUDA or ROCm and tilelang required",
+    not HAS_TILELANG_MHC,
+    reason="TileLang MHC support required",
 )
 @pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
 @pytest.mark.parametrize("hidden_size", [4096, 7168])
@@ -217,8 +217,8 @@ def test_mhc_post_tilelang(num_tokens, hidden_size, hc_mult):
 
 
 @pytest.mark.skipif(
-    not (current_platform.is_cuda_alike() and has_tilelang()),
-    reason="CUDA or ROCm and tilelang required",
+    not HAS_TILELANG_MHC,
+    reason="TileLang MHC support required",
 )
 @pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
 @pytest.mark.parametrize("hidden_size", [4096, 7168])
@@ -324,8 +324,8 @@ def test_hc_head_triton(num_tokens, hidden_size, hc_mult):
 
 
 @pytest.mark.skipif(
-    not (current_platform.is_cuda_alike() and has_tilelang()),
-    reason="CUDA or ROCm and tilelang required",
+    not HAS_TILELANG_MHC,
+    reason="TileLang MHC support required",
 )
 @pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
 @pytest.mark.parametrize("hidden_size", [4096, 7168])

vllm/model_executor/layers/mhc.py

@@ -6,9 +6,25 @@ import torch
 # import vllm.model_executor.kernels.mhc  # noqa: F401
 import vllm.model_executor.kernels.mhc as mhc_kernels
 from vllm.model_executor.custom_op import CustomOp
+from vllm.platforms import current_platform
 from vllm.utils.import_utils import has_tilelang
 
-HAS_TILELANG = has_tilelang()
+
+def _has_tilelang_mhc() -> bool:
+    if not has_tilelang():
+        return False
+    if current_platform.is_cuda():
+        return True
+    if current_platform.is_rocm():
+        from vllm.platforms.rocm import on_gfx942
+
+        # TileLang MHC currently produces incorrect results on gfx942. Keep
+        # gfx942 on the existing torch/triton fallbacks until that path is fixed.
+        return not on_gfx942()
+    return False
+
+
+HAS_TILELANG_MHC = _has_tilelang_mhc()
 
 
 # --8<-- [start:mhc_pre]
@@ -89,7 +105,7 @@ class MHCPreOp(CustomOp):
         #         sinkhorn_repeat,
         #     )
         # else:
-        if HAS_TILELANG:
+        if HAS_TILELANG_MHC:
             return torch.ops.vllm.mhc_pre_tilelang(
                 residual,
                 fn,
@@ -224,7 +240,7 @@ class MHCPostOp(CustomOp):
         #         comb_res_mix,
         #     )
         # else:
-        if HAS_TILELANG:
+        if HAS_TILELANG_MHC:
             return torch.ops.vllm.mhc_post_tilelang(
                 x, residual, post_layer_mix, comb_res_mix
             )
@@ -310,7 +326,7 @@ class HCHeadOp(CustomOp):
         outer_shape = hidden_states.shape[:-2]
         hs_flat = hidden_states.view(-1, hc_mult, hidden_size)
 
-        if HAS_TILELANG:
+        if HAS_TILELANG_MHC:
             out = torch.ops.vllm.hc_head_fused_kernel_tilelang(
                 hs_flat,
                 hc_fn,
@@ -447,7 +463,26 @@ class MHCFusedPostPreOp(CustomOp):
         norm_weight: torch.Tensor | None = None,
         norm_eps: float = 0.0,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        return torch.ops.vllm.mhc_fused_post_pre_tilelang(
+        if HAS_TILELANG_MHC:
+            return torch.ops.vllm.mhc_fused_post_pre_tilelang(
+                x,
+                residual,
+                post_layer_mix,
+                comb_res_mix,
+                fn,
+                hc_scale,
+                hc_base,
+                rms_eps,
+                hc_pre_eps,
+                hc_sinkhorn_eps,
+                hc_post_mult_value,
+                sinkhorn_repeat,
+                n_splits,
+                tile_n,
+                norm_weight,
+                norm_eps,
+            )
+        return self.forward_native(
             x,
             residual,
             post_layer_mix,

vllm/models/deepseek_v4/amd/model.py

@@ -27,6 +27,7 @@ from vllm.model_executor.layers.linear import (
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.mhc import (
+    HAS_TILELANG_MHC,
     HCHeadOp,
     MHCFusedPostPreOp,
     MHCPostOp,
@@ -51,7 +52,6 @@ from vllm.model_executor.models.utils import (
 from vllm.models.deepseek_v4.amd.rocm import DeepseekV4ROCMAiterMLAAttention
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
-from vllm.utils.import_utils import has_tilelang
 
 
 class DeepseekV4MLP(nn.Module):
@@ -303,7 +303,7 @@ class DeepseekV4DecoderLayer(nn.Module):
         self.mhc_pre = MHCPreOp()
         self.mhc_post = MHCPostOp()
         self.mhc_fused_post_pre = MHCFusedPostPreOp()
-        self.has_tilelang = has_tilelang()
+        self.has_tilelang = HAS_TILELANG_MHC
 
     def hc_pre(
         self,
@@ -513,7 +513,7 @@ class DeepseekV4Model(nn.Module):
             requires_grad=False,
         )
         self.hc_head_op = HCHeadOp()
-        self.has_tilelang = has_tilelang()
+        self.has_tilelang = HAS_TILELANG_MHC
         # Pre-hc_head residual stream buffer for the MTP draft. Stable
         # address (outside the cudagraph pool) so the copy_ in forward()
         # refreshes it correctly across captured shapes.

vllm/models/deepseek_v4/amd/mtp.py

@@ -28,7 +28,7 @@ from vllm.model_executor.layers.fused_moe import fused_moe_make_expert_params_ma
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import ReplicatedLinear
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
-from vllm.model_executor.layers.mhc import HCHeadOp
+from vllm.model_executor.layers.mhc import HAS_TILELANG_MHC, HCHeadOp
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     VocabParallelEmbedding,
 )
@@ -42,7 +42,6 @@ from vllm.models.deepseek_v4.common.ops import (
 )
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
-from vllm.utils.import_utils import has_tilelang
 
 from .model import DeepseekV4DecoderLayer
 
@@ -124,7 +123,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
         )
 
         self.hc_head_op = HCHeadOp()
-        self.has_tilelang = has_tilelang()
+        self.has_tilelang = HAS_TILELANG_MHC
 
     def forward(
         self,