From 40a68521c9c325cf2757c05e7a476ad4e54f8038 Mon Sep 17 00:00:00 2001
From: "Wang, FangYuan" <39615225+At1a8@users.noreply.github.com>
Date: Wed, 8 Jul 2026 05:43:05 +0800
Subject: [PATCH] [AMD] Fix DeepSeekV4 server cutlass error (#30374)

---
 python/sglang/jit_kernel/dsa/__init__.py              | 7 ++++++-
 python/sglang/srt/layers/attention/dsa/dsa_indexer.py | 7 ++++++-
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/python/sglang/jit_kernel/dsa/__init__.py b/python/sglang/jit_kernel/dsa/__init__.py
index 4dc1b47ea..9edc55a24 100644
--- a/python/sglang/jit_kernel/dsa/__init__.py
+++ b/python/sglang/jit_kernel/dsa/__init__.py
@@ -1,4 +1,5 @@
-from .cutedsl_paged_mqa_logits import CuteDSLPagedMQALogitsRunner, pick_dsl_expand
+from sglang.srt.utils import is_hip
+
 from .paged_mqa_logits import (
     aiter_paged_mqa_logits,
     cutedsl_paged_mqa_logits,
@@ -6,6 +7,10 @@ from .paged_mqa_logits import (
     deepgemm_paged_mqa_logits_split,
 )
 
+if not is_hip():
+    # Preserve the original eager import behavior on non-ROCm platforms.
+    from .cutedsl_paged_mqa_logits import CuteDSLPagedMQALogitsRunner, pick_dsl_expand
+
 __all__ = [
     "CuteDSLPagedMQALogitsRunner",
     "pick_dsl_expand",
diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py
index 153973f77..c367779cb 100644
--- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py
+++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py
@@ -13,7 +13,6 @@ from sglang.jit_kernel.dsa import (
     cutedsl_paged_mqa_logits,
     deepgemm_paged_mqa_logits_native,
     deepgemm_paged_mqa_logits_split,
-    pick_dsl_expand,
 )
 from sglang.jit_kernel.fused_store_index_cache import (
     can_use_dsa_fused_store,
@@ -65,6 +64,11 @@ global _use_multi_stream
 _is_cuda = is_cuda()
 _is_hip = is_hip()
 _is_npu = is_npu()
+if not _is_hip:
+    # Preserve the original eager import behavior on non-ROCm platforms.
+    from sglang.jit_kernel.dsa import pick_dsl_expand
+else:
+    pick_dsl_expand = None
 _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
 _is_fp8_fnuz = is_fp8_fnuz()
 _is_gfx95_supported = is_gfx95_supported()
@@ -906,6 +910,7 @@ class Indexer(MultiPlatformOp):
             and forward_batch.forward_mode.is_target_verify()
             and next_n >= 2
         ):
+            assert pick_dsl_expand is not None, "Not supported on AMD/ROCm. "
             dsl_expand_factor, dsl_atom = pick_dsl_expand(
                 next_n,
                 batch_size=B,