From fc55b45e5fb352c345b24b3cb409b4fa54d0015a Mon Sep 17 00:00:00 2001
From: Vitaly Tuzov <13571155+terfendail@users.noreply.github.com>
Date: Fri, 14 Nov 2025 21:40:17 +0300
Subject: [PATCH] Optimized prefill cache allocation for NPU (#13288)

---
 .../sglang/srt/mem_cache/allocator_ascend.py  | 34 +++++++++++++++----
 python/sglang/srt/utils/common.py             |  2 +-
 2 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/python/sglang/srt/mem_cache/allocator_ascend.py b/python/sglang/srt/mem_cache/allocator_ascend.py
index 4adbf592a..1b1b84fe0 100644
--- a/python/sglang/srt/mem_cache/allocator_ascend.py
+++ b/python/sglang/srt/mem_cache/allocator_ascend.py
@@ -1,7 +1,12 @@
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import torch
 
+if TYPE_CHECKING:
+    from sglang.srt.mem_cache.memory_pool import KVCache
+
 from sglang.srt.mem_cache.allocator import PagedTokenToKVPoolAllocator
 from sglang.srt.utils import get_num_new_pages
 
@@ -61,6 +66,17 @@ def alloc_extend_kernel_ascend(
 
 
 class AscendPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
+    def __init__(
+        self,
+        size: int,
+        page_size: int,
+        dtype: torch.dtype,
+        device: str,
+        kvcache: KVCache,
+        need_sort: bool,
+    ):
+        super().__init__(size, page_size, dtype, device, kvcache, need_sort)
+        self.roundup = page_size - 1
 
     def alloc_extend(
         self,
@@ -77,8 +93,8 @@ class AscendPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
             )
 
         num_new_pages = (
-            (seq_lens + self.page_size - 1) // self.page_size
-            - (prefix_lens + self.page_size - 1) // self.page_size
+            (seq_lens + self.roundup) // self.page_size
+            - (prefix_lens + self.roundup) // self.page_size
         ).sum()
         num_new_pages_item = num_new_pages.item()
         if self.need_sort and num_new_pages_item > len(self.free_pages):
@@ -87,13 +103,14 @@ class AscendPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
         if num_new_pages_item > len(self.free_pages):
             return None
 
-        out_indices = torch.empty(
-            (extend_num_tokens,), dtype=torch.int64, device=self.device
-        )
-
         if num_new_pages_item < 200:
             import sgl_kernel_npu  # noqa: F401
 
+            out_indices = torch.empty(
+                (extend_num_tokens,),
+                dtype=torch.int64,
+                device=self.device,
+            )
             torch.ops.npu.alloc_extend(
                 prefix_lens,
                 seq_lens,
@@ -105,6 +122,11 @@ class AscendPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
             )
 
         else:
+            out_indices = torch.empty(
+                (extend_num_tokens,),
+                dtype=torch.int32,
+                device=self.device,
+            )
             alloc_extend_kernel_ascend(
                 prefix_lens,
                 seq_lens,
diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py
index f1bee1b14..babd57888 100644
--- a/python/sglang/srt/utils/common.py
+++ b/python/sglang/srt/utils/common.py
@@ -267,7 +267,7 @@ def get_int_env_var(name: str, default: int = 0) -> int:
 
 
 def support_triton(backend: str) -> bool:
-    return backend not in ["torch_native", "intel_amx", "ascend"]
+    return backend not in ["torch_native", "intel_amx"]
 
 
 try: