[NPU] Remove paged attention & Change fia to default attention (#17394)

Co-authored-by: Liwansi <62291011+Liwansi@users.noreply.github.com>
Co-authored-by: chenxu214 <justin_cc2025@163.com>
Co-authored-by: chenyang08056032 <chenyang08056032@163.com>

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

@@ -18,7 +18,6 @@ from sglang.srt.hardware_backend.npu.attention.mla_preprocess import (
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
 from sglang.srt.layers.attention.nsa.utils import is_nsa_enable_prefill_cp
 from sglang.srt.layers.attention.torch_native_backend import TorchNativeAttnBackend
-from sglang.srt.layers.dp_attention import get_attention_tp_size
 from sglang.srt.layers.radix_attention import AttentionType
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.speculative.spec_info import SpecInput
@@ -1098,9 +1097,6 @@ class AscendAttnBackend(AttentionBackend):
             return attn_out
 
         if not self.use_mla:
-            num_tokens = q.shape[0]
-            """PA will support bs<tp in the later version of CANN"""
-            if num_tokens < get_attention_tp_size():
             k_cache = forward_batch.token_to_kv_pool.get_key_buffer(
                 layer.layer_id
             ).view(-1, self.page_size, layer.tp_k_head_num * layer.qk_head_dim)
@@ -1115,8 +1111,7 @@ class AscendAttnBackend(AttentionBackend):
                     self.forward_metadata.seq_lens_cpu_int.cpu().int().tolist()
                 )
             num_tokens = query.shape[0]
-                workspace = (
+            workspace = torch_npu._npu_fused_infer_attention_score_get_max_workspace(
-                    torch_npu._npu_fused_infer_attention_score_get_max_workspace(
                 query,
                 k_cache,
                 v_cache,
@@ -1128,7 +1123,6 @@ class AscendAttnBackend(AttentionBackend):
                 scale=layer.scaling,
                 actual_seq_lengths_kv=actual_seq_len_kv,
             )
-                )
             output = torch.empty(
                 (num_tokens, 1, layer.tp_q_head_num * layer.v_head_dim),
                 dtype=q.dtype,
@@ -1150,41 +1144,6 @@ class AscendAttnBackend(AttentionBackend):
                 out=[output, softmax_lse],
             )
             return output.view(num_tokens, layer.tp_q_head_num * layer.v_head_dim)
-            else:
-                k_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id)
-                v_cache = forward_batch.token_to_kv_pool.get_value_buffer(
-                    layer.layer_id
-                )
-                query = q.reshape(-1, layer.tp_q_head_num, layer.qk_head_dim)
-                num_tokens = query.shape[0]
-                attn_output = torch.empty(
-                    (num_tokens, layer.tp_q_head_num, layer.v_head_dim),
-                    dtype=query.dtype,
-                    device=query.device,
-                )
-                if self.forward_metadata.seq_lens_cpu_int is None:
-                    actual_seq_len_kv = torch.from_numpy(
-                        np.array(self.forward_metadata.seq_lens_cpu_list).astype(
-                            np.int32
-                        )
-                    )
-                else:
-                    actual_seq_len_kv = self.forward_metadata.seq_lens_cpu_int
-
-                torch_npu._npu_paged_attention(
-                    query=query,
-                    key_cache=k_cache,
-                    value_cache=v_cache,
-                    num_heads=layer.tp_q_head_num,
-                    num_kv_heads=layer.tp_k_head_num,
-                    scale_value=layer.scaling,
-                    block_table=self.forward_metadata.block_tables,
-                    context_lens=actual_seq_len_kv,
-                    out=attn_output,
-                )
-                return attn_output.view(
-                    num_tokens, layer.tp_q_head_num * layer.v_head_dim
-                )
         else:
             c_kv, k_rope = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id)
             if is_fia_nz():

python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py

@@ -23,7 +23,6 @@ import numpy as np
 import torch
 
 from sglang.srt.configs.model_config import AttentionArch, is_deepseek_nsa
-from sglang.srt.layers.dp_attention import get_attention_tp_size
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
     EAGLEDraftCudaGraphRunner,
@@ -78,15 +77,11 @@ class EAGLEDraftNpuGraphRunner(EAGLEDraftCudaGraphRunner):
             out = run_once_fn()
         return out
 
-    def _get_update_attr_name(self, model_runner):
+    def _get_update_attr_name(self):
-        if self.bs < get_attention_tp_size():
         return self.attr_name[AttentionArch.MLA]
-        return self.attr_name[model_runner.model_config.attention_arch]
 
-    def _get_update_attr_type(self, model_runner):
+    def _get_update_attr_type(self):
-        if self.bs < get_attention_tp_size():
         return self.attr_type[AttentionArch.MLA]
-        return self.attr_type[model_runner.model_config.attention_arch]
 
     def _replay_update(self, seq_lens):
         if isinstance(self.update_attr_type, torch.Tensor):
@@ -97,8 +92,8 @@ class EAGLEDraftNpuGraphRunner(EAGLEDraftCudaGraphRunner):
         )
 
     def _replay(self, forward_batch: ForwardBatch):
-        self.update_attr_name = self._get_update_attr_name(self.model_runner)
+        self.update_attr_name = self._get_update_attr_name()
-        self.update_attr_type = self._get_update_attr_type(self.model_runner)
+        self.update_attr_type = self._get_update_attr_type()
         if not is_deepseek_nsa(self.model_runner.model_config.hf_config):
             seq_lens = forward_batch.seq_lens_cpu.tolist() + [0] * (
                 self.bs - self.raw_bs

python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py

@@ -28,7 +28,6 @@ import torch
 import sglang
 from sglang.srt.configs.model_config import AttentionArch, is_deepseek_nsa
 from sglang.srt.distributed.parallel_state import GroupCoordinator
-from sglang.srt.layers.dp_attention import get_attention_tp_size
 from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner
 from sglang.srt.utils import (
     empty_context,
@@ -111,23 +110,11 @@ class NPUGraphRunner(CudaGraphRunner):
             out = run_once_fn()
         return out
 
-    def _get_update_attr_name(self, model_runner, forward_batch):
+    def _get_update_attr_name(self):
-        if (
-            self.bs < get_attention_tp_size()
-            or forward_batch.forward_mode.is_target_verify()
-            or self.use_fia
-        ):
         return self.attr_name[AttentionArch.MLA]
-        return self.attr_name[model_runner.model_config.attention_arch]
 
-    def _get_update_attr_type(self, model_runner, forward_batch):
+    def _get_update_attr_type(self):
-        if (
-            self.bs < get_attention_tp_size()
-            or forward_batch.forward_mode.is_target_verify()
-            or self.use_fia
-        ):
         return self.attr_type[AttentionArch.MLA]
-        return self.attr_type[model_runner.model_config.attention_arch]
 
     def _update_inputs(self, seq_lens):
         if isinstance(self.update_attr_type, torch.Tensor):
@@ -181,12 +168,8 @@ class NPUGraphRunner(CudaGraphRunner):
             self.buffers.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids)
             self.buffers.positions[: self.raw_num_token].copy_(forward_batch.positions)
 
-        self.update_attr_name = self._get_update_attr_name(
+        self.update_attr_name = self._get_update_attr_name()
-            self.model_runner, forward_batch
+        self.update_attr_type = self._get_update_attr_type()
-        )
-        self.update_attr_type = self._get_update_attr_type(
-            self.model_runner, forward_batch
-        )
         # Replay
         if not is_deepseek_nsa(self.model_runner.model_config.hf_config):
             if forward_batch.forward_mode.is_target_verify():