From 494629a20d91dee018ed8e3a3893cb4a04a2c909 Mon Sep 17 00:00:00 2001 From: At1a8 Date: Tue, 7 Jul 2026 08:40:26 +0000 Subject: [PATCH 1/3] fix --- python/sglang/jit_kernel/dsa/__init__.py | 22 ++++++++++++++++++- .../srt/layers/attention/dsa/dsa_indexer.py | 4 +++- 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/python/sglang/jit_kernel/dsa/__init__.py b/python/sglang/jit_kernel/dsa/__init__.py index 4dc1b47ea36a..5896a5bb1274 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 @@ 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", @@ -14,3 +19,18 @@ "deepgemm_paged_mqa_logits_native", "deepgemm_paged_mqa_logits_split", ] + + +if is_hip(): + + def __getattr__(name): + if name in ("CuteDSLPagedMQALogitsRunner", "pick_dsl_expand"): + from .cutedsl_paged_mqa_logits import ( + CuteDSLPagedMQALogitsRunner, + pick_dsl_expand, + ) + + globals()["CuteDSLPagedMQALogitsRunner"] = CuteDSLPagedMQALogitsRunner + globals()["pick_dsl_expand"] = pick_dsl_expand + return globals()[name] + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index f256639696e7..d11a4d03cb14 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 @@ 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,9 @@ _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 _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip _is_fp8_fnuz = is_fp8_fnuz() _is_gfx95_supported = is_gfx95_supported() From 619b7d3c76ac45c189980bfa12e130c1f232c3d2 Mon Sep 17 00:00:00 2001 From: At1a8 Date: Tue, 7 Jul 2026 08:44:01 +0000 Subject: [PATCH 2/3] rm --- python/sglang/jit_kernel/dsa/__init__.py | 15 --------------- 1 file changed, 15 deletions(-) diff --git a/python/sglang/jit_kernel/dsa/__init__.py b/python/sglang/jit_kernel/dsa/__init__.py index 5896a5bb1274..9edc55a2420a 100644 --- a/python/sglang/jit_kernel/dsa/__init__.py +++ b/python/sglang/jit_kernel/dsa/__init__.py @@ -19,18 +19,3 @@ "deepgemm_paged_mqa_logits_native", "deepgemm_paged_mqa_logits_split", ] - - -if is_hip(): - - def __getattr__(name): - if name in ("CuteDSLPagedMQALogitsRunner", "pick_dsl_expand"): - from .cutedsl_paged_mqa_logits import ( - CuteDSLPagedMQALogitsRunner, - pick_dsl_expand, - ) - - globals()["CuteDSLPagedMQALogitsRunner"] = CuteDSLPagedMQALogitsRunner - globals()["pick_dsl_expand"] = pick_dsl_expand - return globals()[name] - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") From 769827ed3cc649c13ba8dc6f4c78d9996b521ab0 Mon Sep 17 00:00:00 2001 From: At1a8 Date: Tue, 7 Jul 2026 14:28:10 +0000 Subject: [PATCH 3/3] fix --- python/sglang/srt/layers/attention/dsa/dsa_indexer.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index d11a4d03cb14..5e5c93ef0d0f 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py @@ -67,6 +67,8 @@ 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() @@ -851,6 +853,7 @@ def _get_topk_paged( 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,