Add alloc_extend_kernel (#196)

Author: hw-csong

python/sgl_kernel_npu/sgl_kernel_npu/mem_cache/allocator.py

@@ -0,0 +1,89 @@
+import torch
+import triton
+import triton.language as tl
+from sgl_kernel_npu.utils.triton_utils import get_device_properties
+
+
+@triton.jit
+def alloc_extend_kernel(
+    pre_lens_ptr,
+    seq_lens_ptr,
+    last_loc_ptr,
+    free_page_ptr,
+    out_indices,
+    bs_upper: tl.constexpr,
+    page_size: tl.constexpr,
+    max_num_extend_tokens: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr = 2048,
+):
+    pid = tl.program_id(0)
+
+    load_offset = tl.arange(0, bs_upper)
+    seq_lens = tl.load(seq_lens_ptr + load_offset, mask=load_offset <= pid)
+    pre_lens = tl.load(pre_lens_ptr + load_offset, mask=load_offset <= pid)
+    extend_lens = seq_lens - pre_lens
+
+    seq_len = tl.load(seq_lens_ptr + pid)
+    pre_len = tl.load(pre_lens_ptr + pid)
+    extend_len = seq_len - pre_len
+
+    sum_extend_lens = tl.sum(extend_lens)
+    output_start_loc = sum_extend_lens - extend_len
+
+    num_pages_after = (seq_lens + page_size - 1) // page_size
+    num_pages_before = (pre_lens + page_size - 1) // page_size
+    num_new_pages = num_pages_after - num_pages_before
+
+    num_page_start_loc_self = (seq_len + page_size - 1) // page_size - (
+        pre_len + page_size - 1
+    ) // page_size
+    sum_num_new_pages = tl.sum(num_new_pages)
+    new_page_start_loc = sum_num_new_pages - num_page_start_loc_self
+
+    # Part 1: fill the old partial page
+    last_loc = tl.load(last_loc_ptr + pid)
+    num_part1 = (
+        min(seq_len, (pre_len + page_size - 1) // page_size * page_size) - pre_len
+    )
+    offset_one_page = tl.arange(0, page_size)
+    tl.store(
+        out_indices + output_start_loc + offset_one_page,
+        last_loc + 1 + offset_one_page,
+        mask=offset_one_page < num_part1,
+    )
+    if pre_len + num_part1 == seq_len:
+        return
+
+    # Part 2: fill the new full pages
+    num_part2 = (
+        seq_len // page_size * page_size
+        - (pre_len + page_size - 1) // page_size * page_size
+    )
+
+    num_loop = tl.cdiv(max_num_extend_tokens, BLOCK_SIZE)
+    blk_offset = tl.arange(0, BLOCK_SIZE)
+    for i in range(num_loop):
+        offset_many_page = blk_offset + i * BLOCK_SIZE
+        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,
+        )
+
+    if pre_len + num_part1 + num_part2 == seq_len:
+        return
+
+    # Part 3: fill the new partial page
+    num_part3 = seq_len - seq_len // page_size * page_size
+    start_loc = tl.load(
+        free_page_ptr + new_page_start_loc + num_page_start_loc_self - 1
+    )
+    tl.store(
+        out_indices + output_start_loc + num_part1 + num_part2 + offset_one_page,
+        start_loc * page_size + offset_one_page,
+        mask=offset_one_page < num_part3,
+    )