chore: import upstream snapshot with attribution

python/sglang/kernels/ops/memory/__init__.py

@@ -0,0 +1,30 @@
+"""Memory / KV-slot allocation kernels (Triton).
+
+The Triton kernels migrated here live in this package
+(``sglang.kernels.ops.memory.<module>``); import them from there. Their
+``KernelSpec`` metadata is registered below for inventory (backend = Triton).
+"""
+
+from sglang.kernels.registry import register_kernel
+from sglang.kernels.spec import KernelBackend, KernelSpec
+
+# (module, public_fn) migrated from mem_cache/triton_ops.
+_TRITON_KERNELS = [
+    ("allocator", "alloc_extend_kernel"),
+    ("allocator", "alloc_decode_kernel"),
+    ("common", "write_req_to_token_pool_triton"),
+    ("common", "get_last_loc_triton"),
+    ("common", "get_last_loc_triton_safe"),
+    ("virtual_slot", "alloc_bind_inplace"),
+]
+for _mod, _fn in _TRITON_KERNELS:
+    register_kernel(
+        KernelSpec(
+            op=f"memory.{_fn}",
+            backend=KernelBackend.TRITON,
+            target=f"sglang.kernels.ops.memory.{_mod}:{_fn}",
+        )
+    )
+del _mod, _fn
+
+__all__ = []

python/sglang/kernels/ops/memory/allocator.py

@@ -0,0 +1,135 @@
+import triton
+import triton.language as tl
+
+
+# free_page_ptr aliases self.free_pages, which the paged allocator re-slices
+# after every allocation (self.free_pages = self.free_pages[num_new_pages:]).
+# Slicing only advances data_ptr() by num_new_pages * 8 bytes, so the pointer
+# flips between 16-byte-aligned and unaligned across calls. Triton specializes
+# on pointer alignment by default and bakes it into the cache key, compiling two
+# kernel variants (one with tt.divisibility=16 on free_page_ptr, one without)
+# so the second prefill on a fresh DCP server hits the alternate alignment and
+# pays an extra ~100ms JIT for that kernel variant. do_not_specialize skips
+# that specialization so only one kernel is ever compiled; the perf cost is
+# negligible (this kernel runs in ~10us and only loads ~4KB through this ptr).
+@triton.jit(do_not_specialize=["free_page_ptr"])
+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,
+):
+    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 using a dynamic blocked loop.
+    # The loop bound is derived from num_part2 (runtime value), so Triton
+    # generates a real loop instead of unrolling -- no constexpr dependency
+    # on extend size and only one kernel compilation.
+    num_part2 = (
+        seq_len // page_size * page_size
+        - (pre_len + page_size - 1) // page_size * page_size
+    )
+    BLOCK_EXTEND: tl.constexpr = 4096
+    num_blocks = (num_part2 + 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
+
+    # 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,
+    )
+
+
+# Same free_page_ptr alignment rationale as alloc_extend_kernel above.
+@triton.jit(do_not_specialize=["free_page_ptr"])
+def alloc_decode_kernel(
+    seq_lens_ptr,
+    last_loc_ptr,
+    free_page_ptr,
+    out_indices,
+    bs_upper: tl.constexpr,
+    page_size: tl.constexpr,
+):
+    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.where(load_offset <= pid, seq_lens - 1, seq_lens)
+
+    seq_len = tl.load(seq_lens_ptr + pid)
+    pre_len = seq_len - 1
+
+    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
+
+    if num_page_start_loc_self == 0:
+        last_loc = tl.load(last_loc_ptr + pid)
+        tl.store(out_indices + pid, last_loc + 1)
+    else:
+        page = tl.load(free_page_ptr + new_page_start_loc)
+        tl.store(out_indices + pid, page * page_size)

python/sglang/kernels/ops/memory/common.py

@@ -0,0 +1,163 @@
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def write_req_to_token_pool_triton(
+    req_to_token_ptr,  # [max_batch, max_context_len]
+    req_pool_indices,
+    prefix_tensors,
+    pre_lens,
+    seq_lens,
+    extend_lens,
+    out_cache_loc,
+    req_to_token_ptr_stride: tl.constexpr,
+):
+    BLOCK_SIZE: tl.constexpr = 512
+    pid = tl.program_id(0)
+
+    req_pool_index = tl.load(req_pool_indices + pid)
+    pre_len = tl.load(pre_lens + pid)
+    seq_len = tl.load(seq_lens + pid)
+    prefix_tensor = tl.load(prefix_tensors + pid).to(tl.pointer_type(tl.int64))
+
+    # write prefix
+    num_loop = tl.cdiv(pre_len, BLOCK_SIZE)
+    for i in range(num_loop):
+        offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
+        mask = offset < pre_len
+        value = tl.load(prefix_tensor + offset, mask=mask)
+        tl.store(
+            req_to_token_ptr + req_pool_index * req_to_token_ptr_stride + offset,
+            value,
+            mask=mask,
+        )
+
+    # NOTE: This can be slow for large bs
+    cumsum_start = tl.cast(0, tl.int64)
+    for i in range(pid):
+        cumsum_start += tl.load(extend_lens + i)
+
+    num_loop = tl.cdiv(seq_len - pre_len, BLOCK_SIZE)
+    for i in range(num_loop):
+        offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
+        mask = offset < (seq_len - pre_len)
+        value = tl.load(out_cache_loc + cumsum_start + offset, mask=mask)
+        tl.store(
+            req_to_token_ptr
+            + req_pool_index * req_to_token_ptr_stride
+            + offset
+            + pre_len,
+            value,
+            mask=mask,
+        )
+
+
+@triton.jit
+def _get_last_loc_safe_kernel(
+    req_to_token,
+    req_pool_indices_tensor,
+    prefix_lens_tensor,
+    result_i32,
+    num_tokens,
+    req_to_token_stride,
+    BLOCK_SIZE: tl.constexpr,
+    PREFIX_DTYPE_IS_I64: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offset = tl.arange(0, BLOCK_SIZE) + pid * BLOCK_SIZE
+    mask = offset < num_tokens
+
+    if PREFIX_DTYPE_IS_I64:
+        prefix_lens = tl.load(prefix_lens_tensor + offset, mask=mask, other=0)
+        req_pool_indices = tl.load(req_pool_indices_tensor + offset, mask=mask, other=0)
+        token_index = req_pool_indices * req_to_token_stride + (prefix_lens - 1)
+    else:
+        prefix_lens = tl.load(prefix_lens_tensor + offset, mask=mask, other=0)
+        req_pool_indices = tl.load(req_pool_indices_tensor + offset, mask=mask, other=0)
+        token_index = req_pool_indices.to(tl.int64) * req_to_token_stride + (
+            prefix_lens.to(tl.int64) - 1
+        )
+
+    token_mask = mask & (prefix_lens > 0)
+    tokens = tl.load(req_to_token + token_index, mask=token_mask, other=-1)
+    # Result stays int32 (req_to_token dtype); caller promotes after return.
+    tl.store(result_i32 + offset, tokens, mask=mask)
+
+
+def get_last_loc_triton_safe(
+    req_to_token: torch.Tensor,
+    req_pool_indices_tensor: torch.Tensor,
+    prefix_lens_tensor: torch.Tensor,
+) -> torch.Tensor:
+    """Fused `last_loc` Triton kernel whose in-kernel result buffer is int32
+    (the dtype of req_to_token). The consumer-dtype promotion happens in
+    torch after the kernel returns, so Triton never issues a mixed-width
+    store -- avoiding the HIP int32->int64 store bug hit by the legacy kernel.
+    """
+    num_tokens = prefix_lens_tensor.shape[0]
+    BLOCK_SIZE = 256
+    result_i32 = torch.empty(
+        num_tokens, dtype=torch.int32, device=prefix_lens_tensor.device
+    )
+    grid = (triton.cdiv(num_tokens, BLOCK_SIZE),)
+    _get_last_loc_safe_kernel[grid](
+        req_to_token,
+        req_pool_indices_tensor,
+        prefix_lens_tensor,
+        result_i32,
+        num_tokens,
+        req_to_token.stride(0),
+        BLOCK_SIZE=BLOCK_SIZE,
+        PREFIX_DTYPE_IS_I64=(prefix_lens_tensor.dtype == torch.int64),
+    )
+    return result_i32.to(prefix_lens_tensor.dtype)
+
+
+@triton.jit
+def get_last_loc_kernel(
+    req_to_token,
+    req_pool_indices_tensor,
+    prefix_lens_tensor,
+    result,
+    num_tokens,
+    req_to_token_stride,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offset = tl.arange(0, BLOCK_SIZE) + pid * BLOCK_SIZE
+    mask = offset < num_tokens
+
+    prefix_lens = tl.load(prefix_lens_tensor + offset, mask=mask, other=0)
+    req_pool_indices = tl.load(req_pool_indices_tensor + offset, mask=mask, other=0)
+
+    token_mask = prefix_lens > 0
+    token_index = req_pool_indices * req_to_token_stride + (prefix_lens - 1)
+    tokens = tl.load(req_to_token + token_index, mask=token_mask, other=-1)
+
+    tl.store(result + offset, tokens, mask=mask)
+
+
+def get_last_loc_triton(
+    req_to_token: torch.Tensor,
+    req_pool_indices_tensor: torch.Tensor,
+    prefix_lens_tensor: torch.Tensor,
+) -> torch.Tensor:
+    BLOCK_SIZE = 256
+    num_tokens = prefix_lens_tensor.shape[0]
+    result = torch.empty_like(prefix_lens_tensor)
+    grid = (triton.cdiv(num_tokens, BLOCK_SIZE),)
+
+    get_last_loc_kernel[grid](
+        req_to_token,
+        req_pool_indices_tensor,
+        prefix_lens_tensor,
+        result,
+        num_tokens,
+        req_to_token.stride(0),
+        BLOCK_SIZE,
+    )
+    return result

python/sglang/kernels/ops/memory/virtual_slot.py

@@ -0,0 +1,96 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Virtual<->physical slot Triton kernels for the unified memory pool."""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+# Fused take-physical-pages + bind for the alloc fast path. Invoked ONLY when
+# `_hole_count == 0`; otherwise the slow path drains holes first (Invariant B,
+# greedy hole reuse). Caller advances `watermark_physical` and checks overflow
+# BEFORE launch, passing the PRE-extension watermark. Cuda-graph safe (no
+# `.item()`, no tensor branching); runs on the scheduler thread.
+
+
+@triton.jit
+def alloc_bind_inplace_kernel(
+    v_pages_ptr,  # in: [N] int64 — virtual page ids
+    v2p_ptr,  # in/out: int64 — virtual_to_physical table
+    p2v_ptr,  # in/out: int64 — physical_to_virtual table
+    out_phys_ptr,  # out: [N] int64 — physical page ids
+    N,  # runtime: number of pages to allocate
+    start_phys,  # runtime: lowest physical page id in the new range
+    BLOCK: tl.constexpr,
+):
+    """Fused: ascending arange + out_phys/v2p/p2v scatter.
+
+    Caller pre-adjusts `start_phys` per direction so the range is always
+    ascending (grow-up: start_wm; grow-down: start_wm - N + 1), making the
+    v->p mapping byte-identical to the `torch.arange` slow path.
+    """
+    pid = tl.program_id(0)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < N
+
+    phys = (start_phys + offs).to(tl.int64)
+    v = tl.load(v_pages_ptr + offs, mask=mask, other=0).to(tl.int64)
+
+    # Masked stores skip out-of-range lanes, and `other=0` keeps us off the
+    # v2p[0]/p2v[0] padding-sink slot.
+    tl.store(out_phys_ptr + offs, phys, mask=mask)
+    tl.store(v2p_ptr + v, phys, mask=mask)
+    tl.store(p2v_ptr + phys, v, mask=mask)
+
+
+ALLOC_BIND_BLOCK = 128
+
+
+def alloc_bind_inplace(
+    v_pages: torch.Tensor,
+    v2p: torch.Tensor,
+    p2v: torch.Tensor,
+    start_phys: int,
+) -> torch.Tensor:
+    """Allocate N ascending physical pages from `start_phys` and bind to `v_pages`.
+
+    Caller must advance `watermark_physical` by N and verify overflow BEFORE
+    calling; this launcher does neither.
+    """
+    N = int(v_pages.numel())
+    if N == 0:
+        return torch.empty(0, dtype=torch.int64, device=v_pages.device)
+    if not v_pages.is_cuda:
+        # Pure-torch CPU reference for the CUDA-only kernel.
+        phys_pages = torch.arange(
+            start_phys, start_phys + N, dtype=torch.int64, device=v_pages.device
+        )
+        v = v_pages.to(torch.int64)
+        v2p[v] = phys_pages
+        p2v[phys_pages] = v
+        return phys_pages
+    phys_pages = torch.empty(N, dtype=torch.int64, device=v_pages.device)
+    grid = (triton.cdiv(N, ALLOC_BIND_BLOCK),)
+    alloc_bind_inplace_kernel[grid](
+        v_pages,
+        v2p,
+        p2v,
+        phys_pages,
+        N,
+        start_phys,
+        BLOCK=ALLOC_BIND_BLOCK,
+    )
+    return phys_pages