diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index 427058691..4ece982c2 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py @@ -8,12 +8,6 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union import torch from einops import rearrange -from sglang.jit_kernel.dsa import ( - aiter_paged_mqa_logits, - cutedsl_paged_mqa_logits, - deepgemm_paged_mqa_logits_native, - deepgemm_paged_mqa_logits_split, -) from sglang.jit_kernel.fused_store_index_cache import ( can_use_dsa_fused_store, fused_store_index_k_cache, @@ -64,7 +58,20 @@ global _use_multi_stream _is_cuda = is_cuda() _is_hip = is_hip() _is_npu = is_npu() -if not _is_hip: +if not _is_npu: + from sglang.jit_kernel.dsa import ( + aiter_paged_mqa_logits, + cutedsl_paged_mqa_logits, + deepgemm_paged_mqa_logits_native, + deepgemm_paged_mqa_logits_split, + ) +else: + aiter_paged_mqa_logits = None + cutedsl_paged_mqa_logits = None + 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. from sglang.jit_kernel.dsa import pick_dsl_expand else: