From ac1f07487a6acdcf499d5edb6c8659aa809abba2 Mon Sep 17 00:00:00 2001
From: ybyang <10629930+whybeyoung@users.noreply.github.com>
Date: Wed, 4 Mar 2026 13:01:16 +0800
Subject: [PATCH] Fix triton alloc extend kernel (#19780)

---
 python/sglang/srt/mem_cache/allocator.py | 30 ++++++++++++++----------
 1 file changed, 18 insertions(+), 12 deletions(-)

diff --git a/python/sglang/srt/mem_cache/allocator.py b/python/sglang/srt/mem_cache/allocator.py
index fa08bb66a..d0f646a18 100644
--- a/python/sglang/srt/mem_cache/allocator.py
+++ b/python/sglang/srt/mem_cache/allocator.py
@@ -280,22 +280,28 @@ def alloc_extend_kernel(
     if pre_len + num_part1 == seq_len:
         return
 
-    # Part 2: fill the new full pages
+    # Part 2: fill the new full pages. Use blocked loop (tl.arange(0, BLOCK_EXTEND))
+    # instead of tl.arange(0, max_num_extend_tokens): Triton does not handle very
+    # large arange well; block size 4096 avoids the bottleneck.
     num_part2 = (
         seq_len // page_size * page_size
         - (pre_len + page_size - 1) // page_size * page_size
     )
-
-    offset_many_page = tl.arange(0, max_num_extend_tokens)
-    page_start = tl.load(
-        free_page_ptr + new_page_start_loc + offset_many_page // page_size,
-        mask=offset_many_page < num_part2,
-    )
-    tl.store(
-        out_indices + output_start_loc + num_part1 + offset_many_page,
-        page_start * page_size + offset_many_page % page_size,
-        mask=offset_many_page < num_part2,
-    )
+    BLOCK_EXTEND: tl.constexpr = 4096
+    num_blocks = (max_num_extend_tokens + BLOCK_EXTEND - 1) // BLOCK_EXTEND
+    for block_id in range(num_blocks):
+        offset_in_block = tl.arange(0, BLOCK_EXTEND)
+        offset = block_id * BLOCK_EXTEND + offset_in_block
+        mask = offset < num_part2
+        page_start = tl.load(
+            free_page_ptr + new_page_start_loc + offset // page_size,
+            mask=mask,
+        )
+        tl.store(
+            out_indices + output_start_loc + num_part1 + offset,
+            page_start * page_size + offset % page_size,
+            mask=mask,
+        )
     if pre_len + num_part1 + num_part2 == seq_len:
         return