From a95c9f5b815cc4fed373e1f3483868652a67a778 Mon Sep 17 00:00:00 2001
From: JiaruiChang5268 <71279366+JiaruiChang5268@users.noreply.github.com>
Date: Wed, 21 Jan 2026 23:58:25 +0800
Subject: [PATCH] [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>
---
 .../npu/attention/ascend_backend.py           | 131 ++++++------------
 .../eagle_draft_npu_graph_runner.py           |  17 +--
 .../npu/graph_runner/npu_graph_runner.py      |  29 +---
 3 files changed, 57 insertions(+), 120 deletions(-)

diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py
index fa82245ed..45aa241e5 100644
--- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py
+++ b/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,93 +1097,53 @@ 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)
-                v_cache = forward_batch.token_to_kv_pool.get_value_buffer(
-                    layer.layer_id
-                ).view(-1, self.page_size, layer.tp_v_head_num * layer.v_head_dim)
-                query = q.reshape(-1, 1, layer.tp_q_head_num * layer.qk_head_dim)
-                if self.forward_metadata.seq_lens_cpu_int is None:
-                    actual_seq_len_kv = self.forward_metadata.seq_lens_cpu_list
-                else:
-                    actual_seq_len_kv = (
-                        self.forward_metadata.seq_lens_cpu_int.cpu().int().tolist()
-                    )
-                num_tokens = query.shape[0]
-                workspace = (
-                    torch_npu._npu_fused_infer_attention_score_get_max_workspace(
-                        query,
-                        k_cache,
-                        v_cache,
-                        block_table=self.forward_metadata.block_tables,
-                        block_size=self.page_size,
-                        num_heads=layer.tp_q_head_num,
-                        num_key_value_heads=layer.tp_k_head_num,
-                        input_layout="BSH",
-                        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,
-                    device=q.device,
-                )
-                softmax_lse = torch.empty(1, dtype=q.dtype, device=q.device)
-                torch_npu.npu_fused_infer_attention_score.out(
-                    query,
-                    k_cache,
-                    v_cache,
-                    block_table=self.forward_metadata.block_tables,
-                    block_size=self.page_size,
-                    num_heads=layer.tp_q_head_num,
-                    num_key_value_heads=layer.tp_k_head_num,
-                    input_layout="BSH",
-                    scale=layer.scaling,
-                    actual_seq_lengths_kv=actual_seq_len_kv,
-                    workspace=workspace,
-                    out=[output, softmax_lse],
-                )
-                return output.view(num_tokens, layer.tp_q_head_num * layer.v_head_dim)
+            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)
+            v_cache = forward_batch.token_to_kv_pool.get_value_buffer(
+                layer.layer_id
+            ).view(-1, self.page_size, layer.tp_v_head_num * layer.v_head_dim)
+            query = q.reshape(-1, 1, layer.tp_q_head_num * layer.qk_head_dim)
+            if self.forward_metadata.seq_lens_cpu_int is None:
+                actual_seq_len_kv = self.forward_metadata.seq_lens_cpu_list
             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
+                actual_seq_len_kv = (
+                    self.forward_metadata.seq_lens_cpu_int.cpu().int().tolist()
                 )
+            num_tokens = query.shape[0]
+            workspace = torch_npu._npu_fused_infer_attention_score_get_max_workspace(
+                query,
+                k_cache,
+                v_cache,
+                block_table=self.forward_metadata.block_tables,
+                block_size=self.page_size,
+                num_heads=layer.tp_q_head_num,
+                num_key_value_heads=layer.tp_k_head_num,
+                input_layout="BSH",
+                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,
+                device=q.device,
+            )
+            softmax_lse = torch.empty(1, dtype=q.dtype, device=q.device)
+            torch_npu.npu_fused_infer_attention_score.out(
+                query,
+                k_cache,
+                v_cache,
+                block_table=self.forward_metadata.block_tables,
+                block_size=self.page_size,
+                num_heads=layer.tp_q_head_num,
+                num_key_value_heads=layer.tp_k_head_num,
+                input_layout="BSH",
+                scale=layer.scaling,
+                actual_seq_lengths_kv=actual_seq_len_kv,
+                workspace=workspace,
+                out=[output, softmax_lse],
+            )
+            return 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():
diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py
index adca2c6c2..bad71cf56 100644
--- a/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py
+++ b/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):
-        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_name(self):
+        return self.attr_name[AttentionArch.MLA]
 
-    def _get_update_attr_type(self, model_runner):
-        if self.bs < get_attention_tp_size():
-            return self.attr_type[AttentionArch.MLA]
-        return self.attr_type[model_runner.model_config.attention_arch]
+    def _get_update_attr_type(self):
+        return self.attr_type[AttentionArch.MLA]
 
     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_type = self._get_update_attr_type(self.model_runner)
+        self.update_attr_name = self._get_update_attr_name()
+        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
diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py
index d089fcc63..896667d33 100644
--- a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py
+++ b/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):
-        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_name(self):
+        return self.attr_name[AttentionArch.MLA]
 
-    def _get_update_attr_type(self, model_runner, forward_batch):
-        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 _get_update_attr_type(self):
+        return self.attr_type[AttentionArch.MLA]
 
     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.model_runner, forward_batch
-        )
-        self.update_attr_type = self._get_update_attr_type(
-            self.model_runner, forward_batch
-        )
+        self.update_attr_name = self._get_update_attr_name()
+        self.update_attr_type = self._get_update_attr_type()
         # Replay
         if not is_deepseek_nsa(self.model_runner.model_config.hf_config):
             if forward_batch.forward_mode.is_target_verify():