diff --git a/python/sglang/jit_kernel/dsa/__init__.py b/python/sglang/jit_kernel/dsa/__init__.py index 9edc55a24..9c2f92d76 100644 --- a/python/sglang/jit_kernel/dsa/__init__.py +++ b/python/sglang/jit_kernel/dsa/__init__.py @@ -1,4 +1,12 @@ -from sglang.srt.utils import is_hip +from sglang.srt.utils import is_cuda + +_is_cuda = is_cuda() + +if _is_cuda: + from .cutedsl_paged_mqa_logits import CuteDSLPagedMQALogitsRunner, pick_dsl_expand +else: + CuteDSLPagedMQALogitsRunner = None + pick_dsl_expand = None from .paged_mqa_logits import ( aiter_paged_mqa_logits, @@ -7,10 +15,6 @@ 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 4ece982c2..c3d6bcc51 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py @@ -58,6 +58,7 @@ global _use_multi_stream _is_cuda = is_cuda() _is_hip = is_hip() _is_npu = is_npu() + if not _is_npu: from sglang.jit_kernel.dsa import ( aiter_paged_mqa_logits, @@ -71,11 +72,11 @@ else: deepgemm_paged_mqa_logits_native = None deepgemm_paged_mqa_logits_split = None -if not _is_hip and not _is_npu: - # Preserve the original eager import behavior on non-ROCm platforms. +if _is_cuda: 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() @@ -928,7 +929,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. " + assert pick_dsl_expand is not None, "CuTe DSL paged MQA is CUDA-only." dsl_expand_factor, dsl_atom = pick_dsl_expand( next_n, batch_size=B,