[NPU] fix npu import cutlass error (#30454)
This commit is contained in:
@@ -8,12 +8,6 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
|
|||||||
import torch
|
import torch
|
||||||
from einops import rearrange
|
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 (
|
from sglang.jit_kernel.fused_store_index_cache import (
|
||||||
can_use_dsa_fused_store,
|
can_use_dsa_fused_store,
|
||||||
fused_store_index_k_cache,
|
fused_store_index_k_cache,
|
||||||
@@ -64,7 +58,20 @@ global _use_multi_stream
|
|||||||
_is_cuda = is_cuda()
|
_is_cuda = is_cuda()
|
||||||
_is_hip = is_hip()
|
_is_hip = is_hip()
|
||||||
_is_npu = is_npu()
|
_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.
|
# Preserve the original eager import behavior on non-ROCm platforms.
|
||||||
from sglang.jit_kernel.dsa import pick_dsl_expand
|
from sglang.jit_kernel.dsa import pick_dsl_expand
|
||||||
else:
|
else:
|
||||||
|
|||||||
Reference in New Issue
Block a user