From d56706459c8e52ec3ab1c41dae778e4fe03e0da3 Mon Sep 17 00:00:00 2001
From: ming_wang <68357922+sigama-w@users.noreply.github.com>
Date: Fri, 28 Aug 2026 17:41:49 +0800
Subject: [PATCH] bugfix for index_fill_ on NPU (#36759)

---
 python/sglang/srt/mem_cache/allocator/swa.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py
index a95201664..568e77ab6 100644
--- a/python/sglang/srt/mem_cache/allocator/swa.py
+++ b/python/sglang/srt/mem_cache/allocator/swa.py
@@ -346,9 +346,14 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
     def clear_full_to_swa_mapping(self, full_indices: torch.Tensor) -> None:
         if full_indices.numel() == 0:
             return
-        # index_fill_ passes the 0 as a kernel argument; mapping[idx] = 0 copies a
-        # host-resident scalar and blocks until the stream drains.
-        self.full_to_swa_index_mapping.index_fill_(0, full_indices.to(torch.int64), 0)
+        full_indices = full_indices.to(torch.int64)
+        if _is_npu:
+            # NPU: aclnnIndexFill is unoptimized; direct assignment avoids the overhead.
+            self.full_to_swa_index_mapping[full_indices] = 0
+        else:
+            # CUDA: index_fill_ passes the 0 as a kernel argument; mapping[idx] = 0
+            # copies a host-resident scalar and blocks until the stream drains.
+            self.full_to_swa_index_mapping.index_fill_(0, full_indices, 0)
 
     def free_swa(self, free_index: torch.Tensor):
         if free_index.numel() == 0: