[NPU] Adapt cann 8.5: use sfa and lightning indexer op from cann and CI update (#17615)

Co-authored-by: Kelon <kelonlu@163.com>

python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py

@@ -546,7 +546,7 @@ class AscendAttnBackend(AttentionBackend):
         actual_seq_qlen_prev, actual_seq_qlen_next = actual_seq_qlen
         actual_seq_lengths_kv_prev, actual_seq_lengths_kv_next = actual_seq_lengths_kv
 
-        attn_out_prev = torch.ops.custom.npu_sparse_flash_attention(
+        attn_out_prev, _, _ = torch_npu.npu_sparse_flash_attention(
             query=q_nope_prev,
             key=k_nope,
             value=k_nope,
@@ -565,8 +565,10 @@ class AscendAttnBackend(AttentionBackend):
             layout_query="TND",
             layout_kv="PA_BSND",
             sparse_mode=3,
+            attention_mode=2,
+            return_softmax_lse=False,
         )
-        attn_out_next = torch.ops.custom.npu_sparse_flash_attention(
+        attn_out_next, _, _ = torch_npu.npu_sparse_flash_attention(
             query=q_nope_next,
             key=k_nope,
             value=k_nope,
@@ -585,6 +587,8 @@ class AscendAttnBackend(AttentionBackend):
             layout_query="TND",
             layout_kv="PA_BSND",
             sparse_mode=3,
+            attention_mode=2,
+            return_softmax_lse=False,
         )
         return torch.cat([attn_out_prev, attn_out_next], dim=0)
 
@@ -675,7 +679,7 @@ class AscendAttnBackend(AttentionBackend):
                 actual_seq_lengths_kv,
             )
         else:
-            attn_out = torch.ops.custom.npu_sparse_flash_attention(
+            attn_out, _, _ = torch_npu.npu_sparse_flash_attention(
                 query=q_nope,
                 key=k_nope,
                 value=k_nope,
@@ -694,6 +698,8 @@ class AscendAttnBackend(AttentionBackend):
                 layout_query="TND",
                 layout_kv="PA_BSND",
                 sparse_mode=3,
+                attention_mode=2,
+                return_softmax_lse=False,
             )
 
         return attn_out

python/sglang/srt/hardware_backend/npu/utils.py

@@ -92,14 +92,6 @@ def init_npu_backend():
     assert _is_npu, "NPU backend initialization called on non-NPU device."
 
     import sgl_kernel_npu  # noqa: F401
-
-    try:
-        import custom_ops  # noqa: F401
-    except ImportError:
-        logger.warning(
-            f"custom_ops not found, dsv3.2 requires this package, which includes the npu_lightning_indexer and npu_sparse_flash_attention operators."
-        )
-
     import torch_npu
     from torch_npu.contrib import transfer_to_npu  # noqa: F401
 

python/sglang/srt/layers/attention/nsa/nsa_indexer.py

@@ -23,8 +23,7 @@ if _is_cuda:
     except ImportError as e:
         deep_gemm = e
 
-if _is_npu:
-    import custom_ops  # noqa: F401
+if is_npu():
     import torch_npu
     from sglang.srt.hardware_backend.npu.utils import get_indexer_weight_stream
 
@@ -1277,6 +1276,7 @@ class Indexer(MultiPlatformOp):
                 actual_seq_lengths_kv,
                 block_table,
             )
+            return topk_indices
         else:
             block_table = (
                 block_table[: actual_seq_lengths_q.size()[0]]
@@ -1284,7 +1284,7 @@ class Indexer(MultiPlatformOp):
                 else block_table
             )
 
-            topk_indices = torch.ops.custom.npu_lightning_indexer(
+            topk_indices = torch_npu.npu_lightning_indexer(
                 query=q.view(-1, self.n_heads, self.head_dim),
                 key=past_key_states,
                 weights=weights,
@@ -1298,8 +1298,7 @@ class Indexer(MultiPlatformOp):
                 sparse_count=self.index_topk,
                 sparse_mode=3,
             )
-
-        return topk_indices
+            return topk_indices[0]
 
     def do_npu_cp_balance_indexer(
         self,
@@ -1322,7 +1321,7 @@ class Indexer(MultiPlatformOp):
         actual_seq_lengths_q_prev, actual_seq_lengths_q_next = actual_seq_lengths_q
         actual_seq_lengths_kv_prev, actual_seq_lengths_kv_next = actual_seq_lengths_kv
 
-        topk_indices_prev = torch.ops.custom.npu_lightning_indexer(
+        topk_indices_prev = torch_npu.npu_lightning_indexer(
             query=q_prev,
             key=past_key_states,
             weights=weights_prev,
@@ -1338,7 +1337,7 @@ class Indexer(MultiPlatformOp):
             sparse_count=self.index_topk,
             sparse_mode=3,
         )
-        topk_indices_next = torch.ops.custom.npu_lightning_indexer(
+        topk_indices_next = torch_npu.npu_lightning_indexer(
             query=q_next,
             key=past_key_states,
             weights=weights_next,
@@ -1354,4 +1353,4 @@ class Indexer(MultiPlatformOp):
             sparse_count=self.index_topk,
             sparse_mode=3,
         )
-        return topk_indices_prev, topk_indices_next
+        return topk_indices_prev[0], topk_indices_next[0]