From b35352718c8a701816525a3dfd5d272f5b39b8e7 Mon Sep 17 00:00:00 2001
From: Yifan Qiao <yifanqiao@inferact.ai>
Date: Sat, 25 Apr 2026 18:15:34 +0000
Subject: [PATCH] chore: fix pre-commit

Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
---
 tests/models/test_deepseek_v4_mega_moe.py |  8 +++++++-
 vllm/model_executor/layers/mhc.py         | 18 ++++++++++++++++--
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py
index 0687cb8d6d0..b1fe5abed23 100644
--- a/tests/models/test_deepseek_v4_mega_moe.py
+++ b/tests/models/test_deepseek_v4_mega_moe.py
@@ -5,13 +5,19 @@ from types import SimpleNamespace
 
 import pytest
 import torch
+from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8
 
 from vllm.model_executor.models.deepseek_v4 import (
     DeepseekV4MegaMoEExperts,
     _stage_deepseek_v4_mega_moe_inputs,
     make_deepseek_v4_expert_params_mapping,
 )
-from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cuda(),
+    reason="DeepSeek V4 MegaMoE requires CUDA",
+)
 
 
 def test_deepseek_v4_mega_moe_expert_mapping():
diff --git a/vllm/model_executor/layers/mhc.py b/vllm/model_executor/layers/mhc.py
index b6e4374032b..1521a6b601b 100644
--- a/vllm/model_executor/layers/mhc.py
+++ b/vllm/model_executor/layers/mhc.py
@@ -2,14 +2,28 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import math
 from functools import cache
+from typing import TYPE_CHECKING
 
-import tilelang
-import tilelang.language as T
 import torch
 
+from vllm.platforms import current_platform
+from vllm.utils.import_utils import has_tilelang
 from vllm.utils.math_utils import cdiv
 from vllm.utils.torch_utils import direct_register_custom_op
 
+# tilelang is only available on CUDA platforms
+if TYPE_CHECKING or current_platform.is_cuda_alike():
+    if not has_tilelang():
+        raise ImportError(
+            "tilelang is required for mhc but is not installed. Install it with "
+            "`pip install tilelang`."
+        )
+    import tilelang
+    import tilelang.language as T
+else:
+    tilelang = None  # type: ignore[assignment]
+    T = None  # type: ignore[assignment]
+
 
 @cache
 def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int: