From 89accad2cc9685bbd813ec0efab316b36cf123ca Mon Sep 17 00:00:00 2001 From: Tuukka Sarvi Date: Mon, 22 Jun 2026 12:26:54 +0300 Subject: [PATCH] [ROCm][DSV4] Disable TileLang MHC dispatch on gfx942 (#45931) Signed-off-by: Tuukka Sarvi --- tests/kernels/test_mhc_kernels.py | 22 +++++++------- vllm/model_executor/layers/mhc.py | 45 ++++++++++++++++++++++++---- vllm/models/deepseek_v4/amd/model.py | 6 ++-- vllm/models/deepseek_v4/amd/mtp.py | 5 ++-- 4 files changed, 56 insertions(+), 22 deletions(-) diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index 0e0e3769f49..2bdce9f9c14 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/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]) diff --git a/vllm/model_executor/layers/mhc.py b/vllm/model_executor/layers/mhc.py index de1b2a0c617..fd9d287e9d5 100644 --- a/vllm/model_executor/layers/mhc.py +++ b/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, diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 24c88bb8eb9..edb92351150 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/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. diff --git a/vllm/models/deepseek_v4/amd/mtp.py b/vllm/models/deepseek_v4/amd/mtp.py index 37ce8074af4..5757035cb63 100644 --- a/vllm/models/deepseek_v4/amd/mtp.py +++ b/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,