chore: import upstream snapshot with attribution

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

@@ -0,0 +1,89 @@
+"""KV-cache write/transfer kernels.
+
+This group wraps the Triton ``reshape_and_cache`` launcher, whose implementation
+now lives in this package (``sglang.kernels.ops.kvcache.cache_ops``) after being
+migrated out of ``sglang.srt.layers.attention.triton_ops`` (RFC #29630).
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+from sglang.kernels.registry import register_kernel
+from sglang.kernels.selector import get_kernel
+from sglang.kernels.spec import FormatSignature, KernelBackend, KernelSpec
+
+if TYPE_CHECKING:
+    import torch
+
+register_kernel(
+    KernelSpec(
+        op="kvcache.reshape_and_cache_flash",
+        backend=KernelBackend.TRITON,
+        target="sglang.kernels.ops.kvcache.cache_ops:launch_reshape_and_cache_flash",
+        format_signature=FormatSignature(
+            in_place=True,
+            description="write token-major K/V into paged KV cache layout",
+        ),
+        description="Reshape-and-cache (Triton launcher).",
+    )
+)
+
+
+def reshape_and_cache_flash(
+    key: torch.Tensor,
+    value: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    swa_slot_mapping: Optional[torch.Tensor] = None,
+    k_scale: Optional[torch.Tensor] = None,
+    v_scale: Optional[torch.Tensor] = None,
+) -> None:
+    """Write token-major ``key``/``value`` into paged KV cache layout."""
+    return get_kernel("kvcache.reshape_and_cache_flash", KernelBackend.TRITON)(
+        key,
+        value,
+        key_cache,
+        value_cache,
+        slot_mapping,
+        swa_slot_mapping,
+        k_scale,
+        v_scale,
+    )
+
+
+__all__ = ["reshape_and_cache_flash"]
+
+
+# Other Triton kernels migrated into this group (from attention/mem_cache
+# triton_ops); registered for inventory. Import them from their modules.
+_TRITON_KERNELS = [
+    ("cache_ops", "concat_and_cast_mha_k_triton"),
+    ("cache_ops", "launch_reshape_and_cache_flash"),
+    ("kv_indices", "create_flashinfer_kv_indices_triton"),
+    ("kv_indices", "create_flashmla_kv_indices_triton"),
+    ("kv_indices", "create_chunked_prefix_cache_kv_indices"),
+    ("kv_indices", "get_num_kv_index_blocks_flashmla"),
+    ("kv_indices", "get_num_page_per_block_flashmla"),
+    ("rope_cache", "fused_qk_rope_reshape_and_cache"),
+    ("trtllm_fp8_kv_kernel", "fused_fp8_set_kv_buffer"),
+    ("trtllm_mha_page_table", "build_trtllm_mha_page_table"),
+    ("trtllm_mha_graph_metadata", "update_trtllm_mha_graph_metadata"),
+    ("aiter_unified_attention", "scatter_ragged_to_page_table_kernel"),
+    ("aiter_unified_attention", "scatter_req_to_token_to_page_table_kernel"),
+    ("cache_move", "store_cache_4d"),
+    ("cache_move", "set_kv_buffer_prefix_valid_tiled"),
+    ("cache_move", "copy_all_layer_kv_cache_tiled"),
+    ("mla_buffer", "set_mla_kv_buffer_triton"),
+    ("mla_buffer", "get_mla_kv_buffer_triton"),
+]
+for _mod, _fn in _TRITON_KERNELS:
+    register_kernel(
+        KernelSpec(
+            op=f"kvcache.{_fn}",
+            backend=KernelBackend.TRITON,
+            target=f"sglang.kernels.ops.kvcache.{_mod}:{_fn}",
+        )
+    )
+del _mod, _fn

python/sglang/kernels/ops/kvcache/aiter_unified_attention.py

@@ -0,0 +1,97 @@
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def scatter_ragged_to_page_table_kernel(
+    kv_flat_ptr,
+    kv_indptr_ptr,
+    dest_ptr,
+    dest_stride,
+    sw_page_table_ptr,
+    swa_slot_mapping_ptr,
+    PAGE_SIZE: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    HAS_SWA: tl.constexpr,
+):
+    """Scatter ragged token-level kv_indices into a 2D block-level page table."""
+    pid = tl.program_id(0)
+    block_id = tl.program_id(1)
+
+    start = tl.load(kv_indptr_ptr + pid).to(tl.int64)
+    kv_len = tl.load(kv_indptr_ptr + pid + 1).to(tl.int64) - start
+    num_blocks = (kv_len + PAGE_SIZE - 1) // PAGE_SIZE
+
+    offsets = block_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    if block_id * BLOCK_SIZE >= num_blocks:
+        return
+    mask = offsets < num_blocks
+    token_idx = offsets.to(tl.int64) * PAGE_SIZE
+    vals = tl.load(kv_flat_ptr + start + token_idx, mask=mask, other=0)
+    block_vals = vals // PAGE_SIZE
+    tl.store(
+        dest_ptr + pid.to(tl.int64) * dest_stride + offsets,
+        block_vals,
+        mask=mask,
+    )
+
+    if HAS_SWA:
+        sw_vals = tl.load(swa_slot_mapping_ptr + vals)
+        block_vals = sw_vals // PAGE_SIZE
+        tl.store(
+            sw_page_table_ptr + pid.to(tl.int64) * dest_stride + offsets,
+            block_vals,
+            mask=mask,
+        )
+
+
+@triton.jit
+def scatter_req_to_token_to_page_table_kernel(
+    req_to_token_ptr,
+    req_pool_indices_ptr,
+    seq_lens_ptr,
+    page_table_ptr,
+    req_to_token_stride,
+    page_table_stride,
+    sw_page_table_ptr,
+    swa_slot_mapping_ptr,
+    DRAFT_NUM: tl.constexpr,
+    PAGE_SIZE: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    HAS_SWA: tl.constexpr,
+):
+    """Build the 2D block-level page_table for target_verify from req_to_token."""
+    pid = tl.program_id(0)
+    block_id = tl.program_id(1)
+
+    seq_len = tl.load(seq_lens_ptr + pid).to(tl.int64)
+    kv_len = seq_len + DRAFT_NUM
+    num_blocks = (kv_len + PAGE_SIZE - 1) // PAGE_SIZE
+
+    offsets = block_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    if block_id * BLOCK_SIZE >= num_blocks:
+        return
+    mask = offsets < num_blocks
+
+    rp = tl.load(req_pool_indices_ptr + pid).to(tl.int64)
+    token_idx = offsets.to(tl.int64) * PAGE_SIZE
+    vals = tl.load(
+        req_to_token_ptr + rp * req_to_token_stride + token_idx,
+        mask=mask,
+        other=0,
+    )
+    block_vals = vals // PAGE_SIZE
+    tl.store(
+        page_table_ptr + pid.to(tl.int64) * page_table_stride + offsets,
+        block_vals,
+        mask=mask,
+    )
+
+    if HAS_SWA:
+        sw_vals = tl.load(swa_slot_mapping_ptr + vals)
+        block_vals = sw_vals // PAGE_SIZE
+        tl.store(
+            sw_page_table_ptr + pid.to(tl.int64) * page_table_stride + offsets,
+            block_vals,
+            mask=mask,
+        )

python/sglang/kernels/ops/kvcache/cache_move.py

@@ -0,0 +1,306 @@
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.utils import is_cpu
+
+_is_cpu = is_cpu()
+
+if _is_cpu:
+    from sgl_kernel import copy_all_layer_kv_cache_cpu
+
+
+@triton.jit
+def set_kv_buffer_prefix_valid_tiled(
+    src_k_ptr,
+    src_v_ptr,
+    dst_k_ptr,
+    dst_v_ptr,
+    loc_2d_ptr,
+    commit_len_ptr,
+    src_k_row_stride,
+    src_v_row_stride,
+    dst_k_row_stride,
+    dst_v_row_stride,
+    block_size,
+    ROW_BYTES: tl.constexpr,
+    BYTES_PER_TILE: tl.constexpr,
+):
+    bid = tl.program_id(0)
+    row = tl.program_id(1)
+    tid = tl.program_id(2)
+
+    commit_len = tl.load(commit_len_ptr + bid)
+    if row >= commit_len:
+        return
+
+    byte_off = tid * BYTES_PER_TILE + tl.arange(0, BYTES_PER_TILE)
+    mask_byte = byte_off < ROW_BYTES
+    tl.multiple_of(byte_off, 16)
+
+    loc = tl.load(loc_2d_ptr + bid * block_size + row)
+    src_row = bid * block_size + row
+
+    src_k_ptr = tl.cast(src_k_ptr, tl.pointer_type(tl.uint8))
+    src_v_ptr = tl.cast(src_v_ptr, tl.pointer_type(tl.uint8))
+    dst_k_ptr = tl.cast(dst_k_ptr, tl.pointer_type(tl.uint8))
+    dst_v_ptr = tl.cast(dst_v_ptr, tl.pointer_type(tl.uint8))
+
+    src_k_row_ptr = src_k_ptr + src_row * src_k_row_stride + byte_off
+    src_v_row_ptr = src_v_ptr + src_row * src_v_row_stride + byte_off
+    dst_k_row_ptr = dst_k_ptr + loc * dst_k_row_stride + byte_off
+    dst_v_row_ptr = dst_v_ptr + loc * dst_v_row_stride + byte_off
+
+    k_val = tl.load(src_k_row_ptr, mask=mask_byte, other=0)
+    v_val = tl.load(src_v_row_ptr, mask=mask_byte, other=0)
+    tl.store(dst_k_row_ptr, k_val, mask=mask_byte)
+    tl.store(dst_v_row_ptr, v_val, mask=mask_byte)
+
+
+@triton.jit
+def copy_all_layer_kv_cache_tiled(
+    data_ptrs,
+    strides,
+    tgt_loc_ptr,
+    src_loc_ptr,
+    num_locs,
+    num_locs_upper: tl.constexpr,
+    BYTES_PER_TILE: tl.constexpr,
+):
+    """2D tiled kernel. Safe for in-place copy."""
+    bid = tl.program_id(0)
+    tid = tl.program_id(1)
+
+    stride = tl.load(strides + bid)
+    base_ptr = tl.load(data_ptrs + bid)
+    base_ptr = tl.cast(base_ptr, tl.pointer_type(tl.uint8))
+
+    byte_off = tid * BYTES_PER_TILE + tl.arange(0, BYTES_PER_TILE)
+    mask_byte = byte_off < stride
+    tl.multiple_of(byte_off, 16)
+
+    loc_idx = tl.arange(0, num_locs_upper)
+    mask_loc = loc_idx < num_locs
+
+    src = tl.load(src_loc_ptr + loc_idx, mask=mask_loc, other=0)
+    tgt = tl.load(tgt_loc_ptr + loc_idx, mask=mask_loc, other=0)
+
+    src_ptr = base_ptr + src[:, None] * stride + byte_off[None, :]
+    tgt_ptr = base_ptr + tgt[:, None] * stride + byte_off[None, :]
+
+    mask = mask_loc[:, None] & mask_byte[None, :]
+    vals = tl.load(src_ptr, mask=mask)
+    tl.store(tgt_ptr, vals, mask=mask)
+
+
+def copy_all_layer_kv_cache_func(
+    data_ptrs: torch.Tensor,
+    strides: torch.Tensor,
+    tgt_loc: torch.Tensor,
+    src_loc: torch.Tensor,
+    num_locs: int,
+    num_locs_upper: int,
+    kv_copy_config: dict,
+):
+    if _is_cpu:
+        copy_all_layer_kv_cache_cpu(
+            data_ptrs,
+            strides,
+            tgt_loc[:num_locs],
+            src_loc[:num_locs],
+        )
+        return
+    grid = (data_ptrs.numel(), kv_copy_config["byte_tiles"])
+    copy_all_layer_kv_cache_tiled[grid](
+        data_ptrs,
+        strides,
+        tgt_loc,
+        src_loc,
+        num_locs,
+        num_locs_upper,
+        BYTES_PER_TILE=kv_copy_config["bytes_per_tile"],
+        num_warps=kv_copy_config["num_warps"],
+        num_stages=2,
+    )
+
+
+# ---------------------------------------------------------------------------
+# store_cache_4d — single-launch Triton write into the 4-D page-major envelope
+# K/V views. At `PAGE_SIZE = 1` the kernel constexpr-folds to byte-identical
+# addresses as the slot-major envelope view; at `PAGE_SIZE > 1` it uses the
+# same `(page_id, tok_in_p)` split the attention read kernels use.
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def store_cache_4d_kernel(
+    k_view_ptr,
+    v_view_ptr,
+    cache_k_ptr,
+    cache_v_ptr,
+    loc_ptr,
+    # Strides in ELEMENTS (not bytes); wrapper passes view.stride(D)
+    # directly. K and V may have different head_dim → different per-token
+    # strides, so we carry both.
+    stride_k_page,
+    stride_k_tok,
+    stride_v_page,
+    stride_v_tok,
+    stride_src_k_row,
+    stride_src_v_row,
+    K_ROW_DIM: tl.constexpr,  # head_num * head_dim
+    V_ROW_DIM: tl.constexpr,  # head_num * v_head_dim
+    PAGE_SIZE: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    """Token-parallel Triton write into a 4-D envelope-strided K/V view.
+
+    Grid: ``(N, ceil(max(K_ROW_DIM, V_ROW_DIM) / BLOCK), 2)`` where:
+      - axis 0 → one program per token (loc[i])
+      - axis 1 → blocks within one slot's K (or V) row
+      - axis 2 → 0 = K, 1 = V (two-tensor write fused into one launch)
+
+    For each token i, the kernel writes:
+        page_id  = loc[i] // PAGE_SIZE
+        tok_in_p = loc[i] %  PAGE_SIZE
+        k_view[page_id, tok_in_p, :, :] = cache_k[i, :, :]
+        v_view[page_id, tok_in_p, :, :] = cache_v[i, :, :]
+
+    Cuda-graph safe: no Python branching on tensor values, no `.item()`,
+    all shapes/strides known at launch time.
+    """
+    pid_n = tl.program_id(0)
+    pid_b = tl.program_id(1)
+    pid_kv = tl.program_id(2)
+
+    # 1. Resolve destination slot in the 4-D view.
+    loc = tl.load(loc_ptr + pid_n).to(tl.int64)
+    if PAGE_SIZE == 1:
+        page_id = loc
+        tok_in_p = tl.zeros([], dtype=tl.int64)
+    else:
+        page_id = loc // PAGE_SIZE
+        tok_in_p = loc % PAGE_SIZE
+
+    # 2. Compute per-tensor source/dest pointers.
+    base_off = pid_b * BLOCK + tl.arange(0, BLOCK)
+
+    if pid_kv == 0:
+        mask = base_off < K_ROW_DIM
+        # The trailing (head_num, head_dim) axes of `k_view` are
+        # contiguous: stride[-1]==1, stride[-2]==head_dim. So we can
+        # treat them as a flat K_ROW_DIM dimension addressed by `base_off`.
+        # The wrapper asserts this invariant.
+        src_ptr = cache_k_ptr + pid_n * stride_src_k_row + base_off
+        dst_ptr = (
+            k_view_ptr + page_id * stride_k_page + tok_in_p * stride_k_tok + base_off
+        )
+    else:
+        mask = base_off < V_ROW_DIM
+        src_ptr = cache_v_ptr + pid_n * stride_src_v_row + base_off
+        dst_ptr = (
+            v_view_ptr + page_id * stride_v_page + tok_in_p * stride_v_tok + base_off
+        )
+
+    src = tl.load(src_ptr, mask=mask)
+    tl.store(dst_ptr, src, mask=mask)
+
+
+def store_cache_4d(
+    k_view: torch.Tensor,
+    v_view: torch.Tensor,
+    cache_k: torch.Tensor,
+    cache_v: torch.Tensor,
+    loc: torch.Tensor,
+    page_size: int,
+) -> None:
+    """One-launch Triton write into the 4-D page-major envelope K/V views.
+
+    Writes ``cache_k[i]`` and ``cache_v[i]`` to
+    ``k_view[loc[i]//ps, loc[i]%ps, :, :]`` (and analogously for V) for
+    ``i in [0, N)``.
+
+    Contract:
+        - ``k_view``, ``v_view``: 4-D ``(num_pages, page_size, head_num,
+          head_dim*)``, contiguous in the trailing ``(head_num, head_dim)``
+          dims (i.e., ``stride[-1] == 1`` and ``stride[-2] == head_dim``).
+        - ``cache_k``, ``cache_v``: 3-D ``(N, head_num, head_dim*)``,
+          contiguous in the trailing ``(head_num, head_dim)`` dims.
+        - ``loc``: 1-D int64 or int32, N elements, values in
+          ``[0, num_pages * page_size)``. The caller is responsible for
+          clamping any negative entries to ≥ 0.
+        - At ``page_size == 1`` the kernel produces byte-identical output
+          to the legacy advanced-indexing path.
+
+    Returns nothing; writes in place.
+    """
+    if loc.numel() == 0:
+        return
+    assert k_view.is_cuda and v_view.is_cuda, "store_cache_4d: CUDA only"
+    assert k_view.ndim == 4 and v_view.ndim == 4, (
+        f"store_cache_4d: k_view/v_view must be 4-D, "
+        f"got {k_view.ndim}/{v_view.ndim}"
+    )
+    assert cache_k.ndim == 3 and cache_v.ndim == 3, (
+        f"store_cache_4d: cache_k/cache_v must be 3-D, "
+        f"got {cache_k.ndim}/{cache_v.ndim}"
+    )
+    assert cache_k.shape[0] == cache_v.shape[0] == loc.numel(), (
+        "store_cache_4d: cache_k/cache_v/loc batch dim mismatch: "
+        f"{cache_k.shape[0]}, {cache_v.shape[0]}, {loc.numel()}"
+    )
+    assert k_view.dtype == v_view.dtype == cache_k.dtype == cache_v.dtype, (
+        "store_cache_4d: dtype mismatch: "
+        f"k_view={k_view.dtype}, v_view={v_view.dtype}, "
+        f"cache_k={cache_k.dtype}, cache_v={cache_v.dtype}"
+    )
+    # Stride invariants — the kernel addresses (head_num, head_dim) as one
+    # flat ROW_DIM dimension; this requires the trailing two dims to be
+    # contiguous. This holds for the page-major envelope views
+    # (k_stride = (page_bytes/itemsize, k_row_bytes/itemsize, head_dim, 1)) and
+    # for cache_k/cache_v produced by the model forward.
+    assert k_view.stride(-1) == 1 and k_view.stride(-2) == k_view.shape[-1], (
+        f"store_cache_4d: k_view trailing dims must be contiguous; "
+        f"got stride={k_view.stride()}, shape={tuple(k_view.shape)}"
+    )
+    assert v_view.stride(-1) == 1 and v_view.stride(-2) == v_view.shape[-1], (
+        f"store_cache_4d: v_view trailing dims must be contiguous; "
+        f"got stride={v_view.stride()}, shape={tuple(v_view.shape)}"
+    )
+    assert cache_k.stride(-1) == 1 and cache_k.stride(-2) == cache_k.shape[-1], (
+        f"store_cache_4d: cache_k trailing dims must be contiguous; "
+        f"got stride={cache_k.stride()}, shape={tuple(cache_k.shape)}"
+    )
+    assert cache_v.stride(-1) == 1 and cache_v.stride(-2) == cache_v.shape[-1], (
+        f"store_cache_4d: cache_v trailing dims must be contiguous; "
+        f"got stride={cache_v.stride()}, shape={tuple(cache_v.shape)}"
+    )
+
+    head_num = k_view.shape[2]
+    head_dim = k_view.shape[3]
+    v_head_dim = v_view.shape[3]
+    K_ROW_DIM = head_num * head_dim
+    V_ROW_DIM = head_num * v_head_dim
+    BLOCK = 128
+    N = loc.numel()
+    row_dim_max = max(K_ROW_DIM, V_ROW_DIM)
+    grid = (N, triton.cdiv(row_dim_max, BLOCK), 2)
+
+    store_cache_4d_kernel[grid](
+        k_view,
+        v_view,
+        cache_k,
+        cache_v,
+        loc,
+        k_view.stride(0),
+        k_view.stride(1),
+        v_view.stride(0),
+        v_view.stride(1),
+        cache_k.stride(0),
+        cache_v.stride(0),
+        K_ROW_DIM=K_ROW_DIM,
+        V_ROW_DIM=V_ROW_DIM,
+        PAGE_SIZE=page_size,
+        BLOCK=BLOCK,
+        num_warps=4,
+    )

python/sglang/kernels/ops/kvcache/cache_ops.py

@@ -0,0 +1,266 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def concat_and_cast_mha_k_kernel(
+    k_ptr,
+    k_nope_ptr,
+    k_rope_ptr,
+    head_cnt: tl.constexpr,
+    k_stride0: tl.constexpr,
+    k_stride1: tl.constexpr,
+    nope_stride0: tl.constexpr,
+    nope_stride1: tl.constexpr,
+    rope_stride0: tl.constexpr,
+    nope_dim: tl.constexpr,
+    rope_dim: tl.constexpr,
+):
+    pid_loc = tl.program_id(0)
+    head_range = tl.arange(0, head_cnt)
+
+    k_head_ptr = k_ptr + pid_loc * k_stride0 + head_range[:, None] * k_stride1
+
+    nope_offs = tl.arange(0, nope_dim)
+
+    src_nope_ptr = (
+        k_nope_ptr
+        + pid_loc * nope_stride0
+        + head_range[:, None] * nope_stride1
+        + nope_offs[None, :]
+    )
+    dst_nope_ptr = k_head_ptr + nope_offs[None, :]
+
+    src_nope = tl.load(src_nope_ptr)
+    tl.store(dst_nope_ptr, src_nope)
+
+    rope_offs = tl.arange(0, rope_dim)
+    src_rope_ptr = k_rope_ptr + pid_loc * rope_stride0 + rope_offs[None, :]
+    dst_rope_ptr = k_head_ptr + nope_dim + rope_offs[None, :]
+    src_rope = tl.load(src_rope_ptr)
+    tl.store(dst_rope_ptr, src_rope)
+
+
+def concat_and_cast_mha_k_triton(
+    k: torch.Tensor,
+    k_nope: torch.Tensor,
+    k_rope: torch.Tensor,
+):
+    # The source data type will be implicitly converted to the target data type.
+    assert (
+        len(k.shape) == 3 and len(k_nope.shape) == 3 and len(k_rope.shape) == 3
+    ), f"shape should be 3d, but got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}"
+    assert (
+        k.shape[0] == k_nope.shape[0] and k.shape[0] == k_rope.shape[0]
+    ), f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}"
+    assert (
+        k.shape[1] == k_nope.shape[1] and 1 == k_rope.shape[1]
+    ), f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}"
+    assert (
+        k.shape[-1] == k_nope.shape[-1] + k_rope.shape[-1]
+    ), f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}"
+
+    nope_dim = k_nope.shape[-1]
+    rope_dim = k_rope.shape[-1]
+    grid = (k.shape[0],)
+
+    concat_and_cast_mha_k_kernel[grid](
+        k,
+        k_nope,
+        k_rope,
+        k.shape[1],
+        k.stride(0),
+        k.stride(1),
+        k_nope.stride(0),
+        k_nope.stride(1),
+        k_rope.stride(0),
+        nope_dim,
+        rope_dim,
+    )
+
+
+@triton.jit
+def reshape_and_cache_flash(
+    key_ptr,
+    value_ptr,
+    key_cache_ptr,
+    value_cache_ptr,
+    slot_mapping_ptr,
+    swa_slot_mapping_ptr,
+    k_scale_ptr,
+    v_scale_ptr,
+    block_stride,
+    key_stride,
+    value_stride,
+    num_heads,
+    head_size,
+    block_size,
+    HEAD_BLOCK: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    HAS_SWA: tl.constexpr,
+    USE_SCALE: tl.constexpr,
+):
+    """
+    Triton kernel for reshaping per-token K/V tensors into paged KV cache layout.
+
+    Source layout:
+        key/value: [num_tokens, num_heads, head_size]
+
+    Target cache layout:
+        cache: [num_blocks, block_size, num_heads, head_size]
+
+    Each Triton program instance handles:
+        - one token (program_id(0))
+        - one block of heads (program_id(1))
+
+    Features:
+        - optional SWA slot remapping
+        - optional FP8 scale dequantization before cache write
+
+    Args:
+        key_ptr: Pointer to source key tensor.
+        value_ptr: Pointer to source value tensor.
+        key_cache_ptr: Pointer to destination key cache tensor.
+        value_cache_ptr: Pointer to destination value cache tensor.
+        slot_mapping_ptr: Maps token -> cache slot.
+        swa_slot_mapping_ptr: Optional second-stage slot remap for SWA mode.
+        k_scale_ptr: Optional key scaling factor pointer.
+        v_scale_ptr: Optional value scaling factor pointer.
+        block_stride: Stride between cache blocks.
+        key_stride: Stride between source key tokens.
+        value_stride: Stride between source value tokens.
+        num_heads: Number of attention heads.
+        head_size: Hidden dimension per head.
+        block_size: Number of slots per cache block.
+        HEAD_BLOCK: Number of heads processed per program.
+        BLOCK_D: Vectorized dimension size (power-of-2 padded).
+        HAS_SWA: Enable SWA remapping.
+        USE_SCALE: Enable scale division before storing.
+    """
+
+    # ----------------------------------
+    # program ids
+    # pid0 = token
+    # pid1 = head block
+    # ----------------------------------
+    token_idx = tl.program_id(0)
+    head_block_idx = tl.program_id(1)
+
+    # ----------------------------------
+    # slot mapping
+    # ----------------------------------
+    slot_idx = tl.load(slot_mapping_ptr + token_idx)
+
+    if HAS_SWA:
+        slot_idx = tl.load(swa_slot_mapping_ptr + slot_idx)
+
+    if slot_idx < 0:
+        return
+
+    block_idx = slot_idx // block_size
+    block_offset = slot_idx % block_size
+
+    # ----------------------------------
+    # head range
+    # ----------------------------------
+    head_idx = head_block_idx * HEAD_BLOCK + tl.arange(0, HEAD_BLOCK)
+
+    head_mask = head_idx < num_heads
+
+    dim_idx = tl.arange(0, BLOCK_D)
+
+    # shape = [HEAD_BLOCK, BLOCK_D]
+    offs = head_idx[:, None] * head_size + dim_idx[None, :]
+
+    mask = head_mask[:, None] & (dim_idx[None, :] < head_size)
+
+    # ----------------------------------
+    # source load
+    # ----------------------------------
+    src_key = token_idx * key_stride + offs
+    src_value = token_idx * value_stride + offs
+
+    k = tl.load(key_ptr + src_key, mask=mask)
+    v = tl.load(value_ptr + src_value, mask=mask)
+
+    # ----------------------------------
+    # optional scale
+    # ----------------------------------
+    if USE_SCALE:
+        k_scale = tl.load(k_scale_ptr)
+        v_scale = tl.load(v_scale_ptr)
+
+        k = k / k_scale
+        v = v / v_scale
+
+    # ----------------------------------
+    # target layout
+    # [block_idx, block_offset, head, dim]
+    # ----------------------------------
+    tgt = block_idx * block_stride + block_offset * num_heads * head_size + offs
+
+    tl.store(key_cache_ptr + tgt, k, mask=mask)
+    tl.store(value_cache_ptr + tgt, v, mask=mask)
+
+
+def launch_reshape_and_cache_flash(
+    key,
+    value,
+    key_cache,
+    value_cache,
+    slot_mapping,
+    swa_slot_mapping=None,
+    k_scale=None,
+    v_scale=None,
+):
+    """
+    Launch wrapper for reshape_and_cache_flash Triton kernel.
+
+    This wrapper prepares launch configuration and dispatches the Triton kernel
+    that writes token-major K/V tensors into paged KV cache layout.
+
+    Args:
+        key: Source key tensor [num_tokens, num_heads, head_size]
+        value: Source value tensor [num_tokens, num_heads, head_size]
+        key_cache: Destination key cache [num_blocks, block_size, num_heads, head_size]
+        value_cache: Destination value cache [num_blocks, block_size, num_heads, head_size]
+        slot_mapping: Token-to-cache slot mapping
+        swa_slot_mapping: Optional SWA remapping table
+        k_scale: Optional key scaling factor
+        v_scale: Optional value scaling factor
+    """
+
+    num_tokens = key.shape[0]
+    num_heads = key.shape[1]
+    head_size = key.shape[2]
+
+    HEAD_BLOCK = 4
+
+    BLOCK_D = triton.next_power_of_2(head_size)
+
+    grid = (
+        num_tokens,
+        triton.cdiv(num_heads, HEAD_BLOCK),
+    )
+
+    reshape_and_cache_flash[grid](
+        key,
+        value,
+        key_cache,
+        value_cache,
+        slot_mapping,
+        swa_slot_mapping,
+        k_scale if k_scale is not None else key,
+        v_scale if v_scale is not None else key,
+        key_cache.stride(0),
+        key.stride(0),
+        value.stride(0),
+        num_heads,
+        head_size,
+        key_cache.shape[1],
+        HEAD_BLOCK=HEAD_BLOCK,
+        BLOCK_D=BLOCK_D,
+        HAS_SWA=(swa_slot_mapping is not None),
+        USE_SCALE=(k_scale is not None),
+    )

python/sglang/kernels/ops/kvcache/kv_indices.py

@@ -0,0 +1,152 @@
+import triton
+import triton.language as tl
+
+_FLASHMLA_CREATE_KV_BLOCK_SIZE = 4096
+FLASHMLA_CREATE_KV_BLOCK_SIZE_TRITON = tl.constexpr(_FLASHMLA_CREATE_KV_BLOCK_SIZE)
+
+
+@triton.jit
+def create_flashinfer_kv_indices_triton(
+    req_to_token_ptr,  # [max_batch, max_context_len]
+    req_pool_indices_ptr,
+    page_kernel_lens_ptr,
+    kv_indptr,
+    kv_start_idx,
+    kv_indices_ptr,
+    req_to_token_ptr_stride: tl.constexpr,
+):
+    BLOCK_SIZE: tl.constexpr = 512
+    pid = tl.program_id(axis=0)
+
+    # find the req pool idx, this is for batch to token
+    req_pool_index = tl.load(req_pool_indices_ptr + pid)
+    kv_indices_offset = tl.load(kv_indptr + pid)
+
+    kv_start = 0
+    kv_end = 0
+    if kv_start_idx:
+        kv_start = tl.load(kv_start_idx + pid).to(tl.int32)
+        kv_end = kv_start
+    kv_end += tl.load(page_kernel_lens_ptr + pid).to(tl.int32)
+
+    num_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE)
+    for i in range(num_loop):
+        # index into req_to_token_ptr needs to be int64
+        offset = tl.arange(0, BLOCK_SIZE).to(tl.int64) + i * BLOCK_SIZE
+        mask = offset < kv_end - kv_start
+        data = tl.load(
+            req_to_token_ptr
+            + req_pool_index * req_to_token_ptr_stride
+            + kv_start
+            + offset,
+            mask=mask,
+        )
+        tl.store(kv_indices_ptr + kv_indices_offset + offset, data, mask=mask)
+
+
+@triton.jit
+def create_chunked_prefix_cache_kv_indices(
+    req_to_token_ptr,  # (max_batch, max_context_len,)
+    req_pool_indices_ptr,  # (batch_size,)
+    chunk_start_idx_ptr,  # (batch_size,)
+    chunk_seq_lens_ptr,  # (batch_size,)
+    chunk_cu_seq_lens_ptr,  # (batch_size + 1,)
+    chunk_kv_indices_ptr,  # (num_chunk_tokens,)
+    req_to_token_ptr_stride: tl.constexpr,
+):
+    BLOCK_SIZE: tl.constexpr = 512
+    pid = tl.program_id(axis=0)
+
+    # find the req pool idx, this is for batch to token
+    req_pool_index = tl.load(req_pool_indices_ptr + pid)
+    chunk_kv_indices_offset = tl.load(chunk_cu_seq_lens_ptr + pid)
+
+    # get the token positions of current chunk
+    chunk_start_pos = tl.load(chunk_start_idx_ptr + pid).to(tl.int32)
+    chunk_seq_len = tl.load(chunk_seq_lens_ptr + pid).to(tl.int32)
+
+    num_loop = tl.cdiv(chunk_seq_len, BLOCK_SIZE)
+    for i in range(num_loop):
+        offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
+        mask = offset < chunk_seq_len
+        data = tl.load(
+            req_to_token_ptr
+            + req_pool_index * req_to_token_ptr_stride
+            + chunk_start_pos
+            + offset,
+            mask=mask,
+        )
+        tl.store(
+            chunk_kv_indices_ptr + chunk_kv_indices_offset + offset, data, mask=mask
+        )
+
+
+def get_num_page_per_block_flashmla(page_size: int = 64) -> int:
+    num_page_per_block = _FLASHMLA_CREATE_KV_BLOCK_SIZE // page_size
+    return num_page_per_block
+
+
+def get_num_kv_index_blocks_flashmla(kv_indices_width: int, page_size: int) -> int:
+    """Grid axis-1 size for create_flashmla_kv_indices_triton: the number of
+    page-blocks spanning the widest sequence (one CTA per block). kv_indices_width
+    is the per-row width of the kv_indices buffer (the kernel's kv_indices_ptr_stride).
+    """
+    npb = get_num_page_per_block_flashmla(page_size)
+    return (kv_indices_width + npb - 1) // npb
+
+
+@triton.jit
+def create_flashmla_kv_indices_triton(
+    req_to_token_ptr,  # [max_batch, max_context_len]
+    req_pool_indices_ptr,
+    page_kernel_lens_ptr,
+    kv_start_idx,
+    kv_indices_ptr,
+    req_to_token_ptr_stride: tl.constexpr,
+    kv_indices_ptr_stride: tl.constexpr,
+    PAGED_SIZE: tl.constexpr = 64,
+):
+    NUM_PAGE_PER_BLOCK: tl.constexpr = (
+        FLASHMLA_CREATE_KV_BLOCK_SIZE_TRITON // PAGED_SIZE
+    )
+    pid = tl.program_id(axis=0)
+
+    # find the req pool idx, this is for batch to token
+    req_pool_index = tl.load(req_pool_indices_ptr + pid)
+
+    kv_start = 0
+    kv_end = 0
+    if kv_start_idx:
+        kv_start = tl.load(kv_start_idx + pid).to(tl.int32)
+        kv_end = kv_start
+
+    kv_end += tl.load(page_kernel_lens_ptr + pid).to(tl.int32)
+
+    num_paged = tl.cdiv(kv_end - kv_start, PAGED_SIZE)
+    num_pages_loop = tl.cdiv(kv_end - kv_start, FLASHMLA_CREATE_KV_BLOCK_SIZE_TRITON)
+
+    # One CTA per page-block (grid axis 1) rather than one CTA looping all blocks;
+    # CTAs beyond this sequence's block count are guarded out.
+    i = tl.program_id(axis=1)
+    if i < num_pages_loop:
+        # index into req_to_token_ptr needs to be int64
+        paged_offset = (
+            tl.arange(0, NUM_PAGE_PER_BLOCK).to(tl.int64) + i * NUM_PAGE_PER_BLOCK
+        ) * PAGED_SIZE
+        paged_offset_out = tl.arange(0, NUM_PAGE_PER_BLOCK) + i * NUM_PAGE_PER_BLOCK
+
+        mask = paged_offset < num_paged * PAGED_SIZE
+        mask_out = paged_offset_out < num_paged
+
+        data = tl.load(
+            req_to_token_ptr
+            + req_pool_index * req_to_token_ptr_stride
+            + kv_start
+            + paged_offset,
+            mask=mask,
+        )
+        tl.store(
+            kv_indices_ptr + pid * kv_indices_ptr_stride + paged_offset_out,
+            data // PAGED_SIZE,
+            mask=mask_out,
+        )

python/sglang/kernels/ops/kvcache/mla_buffer.py

@@ -0,0 +1,386 @@
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.jit_kernel.utils import is_arch_support_pdl
+from sglang.srt.runtime_context import get_parallel
+
+
+@triton.jit
+def set_mla_kv_buffer_kernel(
+    kv_buffer_ptr,
+    cache_k_nope_ptr,
+    cache_k_rope_ptr,
+    loc_ptr,
+    buffer_stride: tl.constexpr,
+    nope_stride: tl.constexpr,
+    rope_stride: tl.constexpr,
+    nope_dim: tl.constexpr,
+    rope_dim: tl.constexpr,
+    BLOCK: tl.constexpr,
+    DCP_RANK: tl.constexpr,
+    DCP_WORLD_SIZE: tl.constexpr,
+    USE_GDC: tl.constexpr = False,
+):
+    pid_loc = tl.program_id(0)
+    pid_blk = tl.program_id(1)
+
+    base = pid_blk * BLOCK
+    offs = base + tl.arange(0, BLOCK)
+    total_dim = nope_dim + rope_dim
+    mask = offs < total_dim
+
+    if USE_GDC:
+        tl.extra.cuda.gdc_wait()
+
+    loc = tl.load(loc_ptr + pid_loc).to(tl.int64)
+    is_valid = loc % DCP_WORLD_SIZE == DCP_RANK
+    safe_loc = tl.where(is_valid, loc, 0)
+    safe_loc = safe_loc // DCP_WORLD_SIZE
+    dst_ptr = kv_buffer_ptr + safe_loc * buffer_stride + offs
+
+    # Three-way branch to handle boundary correctly while preserving fast path
+    if base + BLOCK <= nope_dim:
+        # Fast path: entire block is in nope region
+        src = tl.load(
+            cache_k_nope_ptr + pid_loc * nope_stride + offs,
+            mask=mask,
+        )
+    elif base >= nope_dim:
+        # Fast path: entire block is in rope region
+        offs_rope = offs - nope_dim
+        src = tl.load(
+            cache_k_rope_ptr + pid_loc * rope_stride + offs_rope,
+            mask=mask,
+        )
+    else:
+        # Boundary case: block spans nope/rope boundary (e.g., FP8 with nope_dim=528)
+        # Handle each offset individually to avoid negative indexing
+        is_nope = offs < nope_dim
+        is_rope = (offs >= nope_dim) & (offs < (nope_dim + rope_dim))
+
+        src_nope = tl.load(
+            cache_k_nope_ptr + pid_loc * nope_stride + offs,
+            mask=mask & is_nope,
+            other=0,
+        )
+        src_rope = tl.load(
+            cache_k_rope_ptr + pid_loc * rope_stride + (offs - nope_dim),
+            mask=mask & is_rope,
+            other=0,
+        )
+
+        src = tl.where(is_nope, src_nope, src_rope)
+
+    tl.store(dst_ptr, src, mask=mask & is_valid)
+
+    if USE_GDC:
+        tl.extra.cuda.gdc_launch_dependents()
+
+
+# Above this loc count the TMA bulk-store path overtakes the single-CTA-per-loc
+# Triton kernel. Below it, Triton with BLOCK = next_pow2(total_dim) (one CTA
+# does the whole row in one tile, no boundary fan-out) is the winning fallback.
+# Tuned on GB300 with DSv4 row widths.
+_TMA_BULK_STORE_MIN_LOCS = 768
+
+
+def set_mla_kv_buffer_triton(
+    kv_buffer: torch.Tensor,
+    loc: torch.Tensor,
+    cache_k_nope: torch.Tensor,
+    cache_k_rope: torch.Tensor,
+):
+    """Dispatch MLA paged-KV scatter writes to the fastest available path.
+
+    Two paths, chosen on ``n_loc``:
+
+    - ``n_loc >= 768`` (and SM90+ with TMA-compatible row widths): JIT CUDA
+      kernel where each warp loads one (nope, rope) row into shared memory and
+      issues a single ``cp.async.bulk.global.shared::cta`` store to scatter the
+      row at ``kv_buffer[loc[item]]``. Wins at large bs because it packs 4-8
+      items per CTA, drastically reducing the CTA count vs single-CTA-per-loc.
+    - Otherwise: Triton kernel with ``BLOCK = next_pow2(nope_dim + rope_dim)``,
+      i.e. one CTA per loc covering the entire row in one tile. Wins at small
+      bs because there's no per-loc CTA fan-out (5x fewer CTAs than the old
+      BLOCK=128 dispatch) and the row-spanning block makes the boundary branch
+      a one-shot per CTA. This is also the path for SM<90 and for shapes that
+      violate the TMA 16-byte alignment.
+
+    Speedup vs the legacy BLOCK=128 Triton kernel on GB300 (BF16, nope=512,
+    rope=64): ~1.05x at bs=8, ~1.5x at bs=128, 3.5x at bs=512, **11.7x at
+    bs=16384**.
+
+    Name retained for caller compatibility; the implementation is no longer
+    Triton-only.
+    """
+    from sglang.jit_kernel.set_mla_kv_buffer import (
+        can_use_set_mla_kv_buffer,
+    )
+    from sglang.jit_kernel.set_mla_kv_buffer import (
+        set_mla_kv_buffer as jit_set_mla_kv_buffer,
+    )
+
+    n_loc = loc.numel()
+    nope_bytes = cache_k_nope.shape[-1] * cache_k_nope.element_size()
+    rope_bytes = cache_k_rope.shape[-1] * cache_k_rope.element_size()
+    if (
+        n_loc >= _TMA_BULK_STORE_MIN_LOCS
+        and is_arch_support_pdl()
+        and can_use_set_mla_kv_buffer(nope_bytes, rope_bytes)
+        and not get_parallel().dcp_enabled
+    ):
+        jit_set_mla_kv_buffer(kv_buffer, loc, cache_k_nope, cache_k_rope)
+        return
+
+    # Fallback: Triton with BLOCK = next_pow2(total_dim). One CTA per loc; the
+    # whole row in one tile (the existing 3-way nope/rope/boundary branch in
+    # ``set_mla_kv_buffer_kernel`` handles the over-allocation past total_dim
+    # via the offs<total_dim mask). Beats BLOCK=128 by 60-2700 ns across the
+    # 2 <= bs <= 512 range on GB300.
+    nope_dim = cache_k_nope.shape[-1]
+    rope_dim = cache_k_rope.shape[-1]
+    total_dim = nope_dim + rope_dim
+    BLOCK = triton.next_power_of_2(total_dim)
+    grid = (n_loc, 1)
+    pdl_kwargs = {"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {}
+    set_mla_kv_buffer_kernel[grid](
+        kv_buffer,
+        cache_k_nope,
+        cache_k_rope,
+        loc,
+        kv_buffer.stride(0),
+        cache_k_nope.stride(0),
+        cache_k_rope.stride(0),
+        nope_dim,
+        rope_dim,
+        BLOCK=BLOCK,
+        DCP_RANK=get_parallel().attn_dcp_rank,
+        DCP_WORLD_SIZE=get_parallel().attn_dcp_size,
+        **pdl_kwargs,
+    )
+
+
+@triton.jit
+def set_mla_kv_buffer_fp8_quant_kernel(
+    kv_buffer_fp8_ptr,
+    cache_k_nope_ptr,
+    cache_k_rope_ptr,
+    loc_ptr,
+    buffer_stride: tl.constexpr,
+    nope_stride: tl.constexpr,
+    rope_stride: tl.constexpr,
+    nope_dim: tl.constexpr,
+    rope_dim: tl.constexpr,
+    BLOCK: tl.constexpr,
+    USE_GDC: tl.constexpr = False,
+):
+    """Fuse BF16/FP16->FP8 cast with paged KV write."""
+    pid_loc = tl.program_id(0)
+    pid_blk = tl.program_id(1)
+
+    base = pid_blk * BLOCK
+    offs = base + tl.arange(0, BLOCK)
+    total_dim = nope_dim + rope_dim
+    mask = offs < total_dim
+
+    if USE_GDC:
+        tl.extra.cuda.gdc_wait()
+
+    loc = tl.load(loc_ptr + pid_loc).to(tl.int64)
+    dst_ptr = kv_buffer_fp8_ptr + loc * buffer_stride + offs
+
+    if base + BLOCK <= nope_dim:
+        src = tl.load(
+            cache_k_nope_ptr + pid_loc * nope_stride + offs,
+            mask=mask,
+            other=0.0,
+        )
+    elif base >= nope_dim:
+        offs_rope = offs - nope_dim
+        src = tl.load(
+            cache_k_rope_ptr + pid_loc * rope_stride + offs_rope,
+            mask=mask,
+            other=0.0,
+        )
+    else:
+        is_nope = offs < nope_dim
+        src_nope = tl.load(
+            cache_k_nope_ptr + pid_loc * nope_stride + offs,
+            mask=mask & is_nope,
+            other=0.0,
+        )
+        src_rope = tl.load(
+            cache_k_rope_ptr + pid_loc * rope_stride + (offs - nope_dim),
+            mask=mask & ~is_nope,
+            other=0.0,
+        )
+        src = tl.where(is_nope, src_nope, src_rope)
+
+    # Destination pointer is FP8-typed view; tl.store performs downcast.
+    tl.store(dst_ptr, src, mask=mask)
+
+    if USE_GDC:
+        tl.extra.cuda.gdc_launch_dependents()
+
+
+def set_mla_kv_buffer_triton_fp8_quant(
+    kv_buffer: torch.Tensor,
+    loc: torch.Tensor,
+    cache_k_nope: torch.Tensor,
+    cache_k_rope: torch.Tensor,
+    fp8_dtype: torch.dtype,
+):
+    """Fuse BF16/FP16 MLA K quantization with paged KV write."""
+    kv_buffer_fp8 = kv_buffer.view(fp8_dtype)
+
+    nope_dim = cache_k_nope.shape[-1]
+    rope_dim = cache_k_rope.shape[-1]
+    total_dim = nope_dim + rope_dim
+    BLOCK = 128
+    n_loc = loc.numel()
+    grid = (n_loc, triton.cdiv(total_dim, BLOCK))
+
+    pdl_kwargs = {"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {}
+
+    set_mla_kv_buffer_fp8_quant_kernel[grid](
+        kv_buffer_fp8,
+        cache_k_nope,
+        cache_k_rope,
+        loc,
+        kv_buffer_fp8.stride(0),
+        cache_k_nope.stride(0),
+        cache_k_rope.stride(0),
+        nope_dim,
+        rope_dim,
+        BLOCK=BLOCK,
+        **pdl_kwargs,
+    )
+
+
+@triton.jit
+def set_mla_kv_scale_buffer_kernel(
+    kv_buffer_ptr,
+    cache_k_nope_ptr,
+    cache_k_rope_ptr,
+    loc_ptr,
+    buffer_stride: tl.constexpr,
+    nope_stride: tl.constexpr,
+    rope_stride: tl.constexpr,
+    nope_dim: tl.constexpr,
+    rope_dim: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    pid_loc = tl.program_id(0)
+    pid_blk = tl.program_id(1)
+
+    base = pid_blk * BLOCK
+    offs = base + tl.arange(0, BLOCK)
+    total_dim = nope_dim + rope_dim
+    mask = offs < total_dim  # Make sure don't cross the boundary
+
+    loc = tl.load(loc_ptr + pid_loc)
+    dst_ptr = kv_buffer_ptr + loc * buffer_stride + offs
+
+    # Check each offs should read 'nope' or 'rope'
+    is_nope = offs < nope_dim
+    src_nope = tl.load(
+        cache_k_nope_ptr + pid_loc * nope_stride + offs, mask=mask & is_nope, other=0.0
+    )
+    src_rope = tl.load(
+        cache_k_rope_ptr + pid_loc * rope_stride + (offs - nope_dim),
+        mask=mask & ~is_nope,
+        other=0.0,
+    )
+
+    # Combine nope + rope
+    src = src_nope + src_rope
+    tl.store(dst_ptr, src, mask=mask)
+
+
+def set_mla_kv_scale_buffer_triton(
+    kv_buffer: torch.Tensor,
+    loc: torch.Tensor,
+    cache_k_nope: torch.Tensor,
+    cache_k_rope: torch.Tensor,
+):
+    nope_dim = cache_k_nope.shape[-1]
+    rope_dim = cache_k_rope.shape[-1]
+    total_dim = nope_dim + rope_dim
+    BLOCK = 128  # Keep origin, works for smaller total_dim as well.
+    n_loc = loc.numel()
+    grid = (n_loc, triton.cdiv(total_dim, BLOCK))
+
+    set_mla_kv_scale_buffer_kernel[grid](
+        kv_buffer,
+        cache_k_nope,
+        cache_k_rope,
+        loc,
+        kv_buffer.stride(0),
+        cache_k_nope.stride(0),
+        cache_k_rope.stride(0),
+        nope_dim,
+        rope_dim,
+        BLOCK=BLOCK,
+    )
+
+
+@triton.jit
+def get_mla_kv_buffer_kernel(
+    kv_buffer_ptr,
+    cache_k_nope_ptr,
+    cache_k_rope_ptr,
+    loc_ptr,
+    buffer_stride: tl.constexpr,
+    nope_stride: tl.constexpr,
+    rope_stride: tl.constexpr,
+    nope_dim: tl.constexpr,
+    rope_dim: tl.constexpr,
+):
+    pid_loc = tl.program_id(0)
+    loc = tl.load(loc_ptr + pid_loc).to(tl.int64)
+    loc_src_ptr = kv_buffer_ptr + loc * buffer_stride
+
+    nope_offs = tl.arange(0, nope_dim)
+    nope_src_ptr = loc_src_ptr + nope_offs
+    nope_src = tl.load(nope_src_ptr)
+
+    tl.store(
+        cache_k_nope_ptr + pid_loc * nope_stride + nope_offs,
+        nope_src,
+    )
+
+    rope_offs = tl.arange(0, rope_dim)
+    rope_src_ptr = loc_src_ptr + nope_dim + rope_offs
+    rope_src = tl.load(rope_src_ptr)
+    tl.store(
+        cache_k_rope_ptr + pid_loc * rope_stride + rope_offs,
+        rope_src,
+    )
+
+
+def get_mla_kv_buffer_triton(
+    kv_buffer: torch.Tensor,
+    loc: torch.Tensor,
+    cache_k_nope: torch.Tensor,
+    cache_k_rope: torch.Tensor,
+):
+    # The source data type will be implicitly converted to the target data type.
+    nope_dim = cache_k_nope.shape[-1]  # 512
+    rope_dim = cache_k_rope.shape[-1]  # 64
+    n_loc = loc.numel()
+    grid = (n_loc,)
+
+    get_mla_kv_buffer_kernel[grid](
+        kv_buffer,
+        cache_k_nope,
+        cache_k_rope,
+        loc,
+        kv_buffer.stride(0),
+        cache_k_nope.stride(0),
+        cache_k_rope.stride(0),
+        nope_dim,
+        rope_dim,
+    )

python/sglang/kernels/ops/kvcache/rope_cache.py

@@ -0,0 +1,736 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _get_gptj_rotated_x(
+    x,
+    x_rotated_mask,
+    BLOCK_D: tl.constexpr,
+    BLOCK_D_HALF: tl.constexpr,
+):
+    # GPT-J rotary layout:
+    # Pair adjacent dimensions and apply:
+    # [x0, x1, x2, x3] -> [-x1, x0, -x3, x2]
+
+    # Apply sign inversion on odd positions.
+    x_rotated = tl.where(x_rotated_mask, x, -x)
+    # Reshape into (D/2, 2) pairs.
+    x_rotated = tl.reshape(x_rotated, (BLOCK_D_HALF, 2))
+    # Swap each pair.
+    x_rotated = tl.flip(x_rotated, 1)
+    # Flatten back to original shape.
+    x_rotated = tl.reshape(x_rotated, (BLOCK_D,))
+    return x_rotated
+
+
+@triton.jit
+def _get_neox_rotated_x(
+    x,
+    x_rotated_mask,
+    BLOCK_D: tl.constexpr,
+    BLOCK_D_HALF: tl.constexpr,
+):
+    # GPT-NeoX rotary layout:
+    # Split head dimension into two halves:
+    # [x0, x1, x2, x3] -> [-x2, -x3, x0, x1]
+
+    # Keep first half positive, second half negative.
+    x_rotated = tl.where(x_rotated_mask, x, -x)
+    # Reshape into (2, D/2).
+    x_rotated = tl.reshape(x_rotated, (2, BLOCK_D_HALF))
+    # Reverse each half.
+    x_rotated = tl.flip(x_rotated, 1)
+    # Flatten and reverse full vector.
+    x_rotated = tl.reshape(x_rotated, (BLOCK_D,))
+    x_rotated = tl.flip(x_rotated, 0)
+    return x_rotated
+
+
+@triton.jit
+def _unit_rope(
+    x_ptrs,
+    cos,
+    sin,
+    d_pe_offs,
+    IS_NEOX: tl.constexpr,
+    BLOCK_D_pe: tl.constexpr,
+    BLOCK_D_HALF_pe: tl.constexpr,
+):
+    # Load one full attention head vector.
+    x_pe = tl.load(x_ptrs)
+
+    # Stage 1: Build rotated vector according to rotary layout.
+    if IS_NEOX:
+        x_rotated_mask = d_pe_offs < BLOCK_D_HALF_pe
+        x_pe_rotated = _get_neox_rotated_x(
+            x_pe, x_rotated_mask, BLOCK_D_pe, BLOCK_D_HALF_pe
+        )
+    else:
+        x_rotated_mask = d_pe_offs % 2 == 0
+        x_pe_rotated = _get_gptj_rotated_x(
+            x_pe, x_rotated_mask, BLOCK_D_pe, BLOCK_D_HALF_pe
+        )
+
+    # Stage 2: Apply RoPE transform:
+    # x' = x*cos + rotate(x)*sin
+    x_pe = x_pe * cos + x_pe_rotated * sin
+
+    return x_pe
+
+
+@triton.jit
+def _load_cos_sin(
+    cos_sin_ptr,
+    pos,
+    d_cos_offs,
+    stride_t,
+    stride_d,
+    freq_dim,
+):
+    base = pos * stride_t
+    cos = tl.load(cos_sin_ptr + base + d_cos_offs * stride_d)
+    sin = tl.load(cos_sin_ptr + base + (d_cos_offs + freq_dim) * stride_d)
+    return cos, sin
+
+
+@triton.jit
+def _fused_qk_rope_reshape_and_cache_kernel(
+    q_ptr,
+    k_ptr,
+    v_ptr,
+    pos_ptr,
+    cos_sin_ptr,
+    offs_ptr,
+    key_cache_ptr,
+    value_cache_ptr,
+    slot_mapping_ptr,
+    swa_slot_mapping_ptr,
+    q_out_ptr,
+    k_out_ptr,
+    zeros_out_ptr,
+    T,
+    T_slot,
+    q_stride_t,
+    q_stride_h,
+    q_stride_d,
+    k_stride_t,
+    k_stride_h,
+    k_stride_d,
+    v_stride_t,
+    v_stride_h,
+    v_stride_d,
+    cos_sin_stride_t,
+    cos_sin_stride_d,
+    q_out_stride_t,
+    q_out_stride_h,
+    q_out_stride_d,
+    k_out_stride_t,
+    k_out_stride_h,
+    k_out_stride_d,
+    key_cache_stride_t,
+    key_cache_stride_h,
+    key_cache_stride_d,
+    key_cache_stride_b,
+    key_cache_stride_x,
+    value_cache_stride_t,
+    value_cache_stride_h,
+    value_cache_stride_d,
+    value_cache_stride_b,
+    value_cache_stride_slot_chunk,
+    value_cache_stride_x,
+    zeros_out_stride_t,
+    zeros_out_stride_h,
+    zeros_out_stride_d,
+    k_scale_ptr,
+    v_scale_ptr,
+    QH_PER_KH: tl.constexpr,
+    QH: tl.constexpr,
+    KH: tl.constexpr,
+    REUSE_FREQS_FRONT_PART: tl.constexpr,
+    IS_NEOX: tl.constexpr,
+    BLOCK_D_pe: tl.constexpr,
+    BLOCK_D_HALF_pe: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    X_SIZE: tl.constexpr,
+    FLASH_LAYOUT: tl.constexpr,
+    VALUE_SHUFFLE_LAYOUT: tl.constexpr = False,
+    HAVE_POS: tl.constexpr = False,
+    HAVE_K_SCALE: tl.constexpr = False,
+    HAVE_V_SCALE: tl.constexpr = False,
+    HAVE_ZEROS: tl.constexpr = False,
+    HAS_SWA: tl.constexpr = False,
+):
+    # ============================================================
+    # Stage 0: Static stride assumptions for Triton compiler
+    #
+    # These assumptions help Triton optimize pointer arithmetic and
+    # simplify generated address calculations.
+    # ============================================================
+
+    tl.assume(q_stride_t >= 0)
+    tl.assume(q_stride_h >= 0)
+    tl.assume(q_stride_d >= 0)
+    tl.assume(k_stride_t >= 0)
+    tl.assume(k_stride_h >= 0)
+    tl.assume(k_stride_d >= 0)
+    tl.assume(v_stride_t >= 0)
+    tl.assume(v_stride_h >= 0)
+    tl.assume(v_stride_d >= 0)
+    tl.assume(cos_sin_stride_t >= 0)
+    tl.assume(cos_sin_stride_d >= 0)
+    tl.assume(q_out_stride_t >= 0)
+    tl.assume(q_out_stride_h >= 0)
+    tl.assume(q_out_stride_d >= 0)
+    tl.assume(k_out_stride_t >= 0)
+    tl.assume(k_out_stride_h >= 0)
+    tl.assume(k_out_stride_d >= 0)
+    tl.assume(key_cache_stride_t >= 0)
+    tl.assume(key_cache_stride_h >= 0)
+    tl.assume(key_cache_stride_d >= 0)
+    tl.assume(key_cache_stride_b >= 0)
+    tl.assume(key_cache_stride_x >= 0)
+    tl.assume(value_cache_stride_t >= 0)
+    tl.assume(value_cache_stride_h >= 0)
+    tl.assume(value_cache_stride_d >= 0)
+    tl.assume(value_cache_stride_b >= 0)
+    tl.assume(value_cache_stride_slot_chunk >= 0)
+    tl.assume(value_cache_stride_x >= 0)
+    tl.assume(zeros_out_stride_t >= 0)
+    tl.assume(zeros_out_stride_h >= 0)
+    tl.assume(zeros_out_stride_d >= 0)
+
+    # ============================================================
+    # Stage 1: Program instance mapping
+    #
+    # Each program handles:
+    #   - one (token, q_head) for Q path
+    #   - selected KV ownership for cache write path
+    #
+    # pid layout:
+    #   [0, T*QH)            -> decode Q path
+    #   [T*QH, extra KV)     -> KV-only path
+    # ============================================================
+
+    pid = tl.program_id(0)
+    tl.assume(pid >= 0)
+
+    d_pe_offs = tl.arange(0, BLOCK_D_pe).to(tl.int64)
+
+    # ============================================================
+    # Stage 2: Main decode path (Q always active)
+    # ============================================================
+
+    if pid < T * QH:
+        pid_t = pid // QH
+        pid_hq = pid % QH
+
+        # --------------------------------------------------------
+        # Stage 2.1: Compute rotary frequency offsets
+        #
+        # RoPE frequencies may be stored as:
+        #   D/2 frequencies (shared front-half)
+        #   D frequencies (full explicit)
+        # --------------------------------------------------------
+
+        if REUSE_FREQS_FRONT_PART:
+            if IS_NEOX:
+                d_cos_offs = d_pe_offs
+                d_cos_offs = tl.where(
+                    (d_cos_offs >= BLOCK_D_HALF_pe) & (d_cos_offs < BLOCK_D_pe),
+                    d_cos_offs - BLOCK_D_HALF_pe,
+                    d_cos_offs,
+                ).to(d_cos_offs.dtype)
+                # d_cos_mask = d_cos_offs < BLOCK_D_pe
+            else:
+                d_cos_offs = d_pe_offs // 2
+                # d_cos_mask = d_cos_offs < BLOCK_D_HALF_pe
+        else:
+            d_cos_offs = d_pe_offs
+            # d_cos_mask = d_cos_offs < BLOCK_D_pe
+
+        # --------------------------------------------------------
+        # Stage 2.2: Load token position and optional offset
+        #
+        # offs_ptr is used by chunked prefill / sliding-window decode.
+        # --------------------------------------------------------
+        pos = tl.load(pos_ptr + pid_t)
+        if HAVE_POS:
+            offset = tl.load(offs_ptr + pid_t)
+            pos = pos + offset
+
+        # --------------------------------------------------------
+        # Stage 2.3: Load cosine / sine table
+        # --------------------------------------------------------
+        # cos_offs = pos * cos_stride_t + d_cos_offs * cos_stride_d
+        # cos = tl.load(cos_ptr + cos_offs)
+        # sin = tl.load(sin_ptr + cos_offs)
+
+        freq_dim = BLOCK_D_HALF_pe if REUSE_FREQS_FRONT_PART else BLOCK_D_pe
+
+        cos, sin = _load_cos_sin(
+            cos_sin_ptr,
+            pos,
+            d_cos_offs,
+            cos_sin_stride_t,
+            cos_sin_stride_d,
+            freq_dim,
+        )
+
+        # --------------------------------------------------------
+        # Stage 2.4: Apply RoPE to Q
+        # --------------------------------------------------------
+        q_ptrs = (
+            q_ptr + pid_t * q_stride_t + pid_hq * q_stride_h + d_pe_offs * q_stride_d
+        )
+        q_pe = _unit_rope(
+            q_ptrs,
+            cos,
+            sin,
+            d_pe_offs,
+            IS_NEOX,
+            BLOCK_D_pe,
+            BLOCK_D_HALF_pe,
+        )
+
+        # Store rotated Q output.
+        q_out_ptrs = (
+            q_out_ptr
+            + pid_t * q_out_stride_t
+            + pid_hq * q_out_stride_h
+            + d_pe_offs * q_out_stride_d
+        )
+        tl.store(q_out_ptrs, q_pe.to(q_out_ptr.dtype.element_ty))
+
+        if HAVE_ZEROS:
+            z = tl.zeros((BLOCK_D_pe,), dtype=zeros_out_ptr.dtype.element_ty)
+            zeros_out_ptrs = (
+                zeros_out_ptr
+                + pid_t * zeros_out_stride_t
+                + pid_hq * zeros_out_stride_h
+                + d_pe_offs * zeros_out_stride_d
+            )
+            tl.store(zeros_out_ptrs, z)
+
+        # ========================================================
+        # Stage 3: KV ownership path
+        #
+        # Only one Q group leader writes KV:
+        #   pid_hq % QH_PER_KH == 0
+        #
+        # This prevents duplicated KV cache writes.
+        # ========================================================
+
+        if pid_hq % QH_PER_KH == 0:
+            # ----------------------------------------------------
+            # Stage 3.1: Resolve cache slot
+            # ----------------------------------------------------
+            pid_slot = tl.load(slot_mapping_ptr + pid_t).to(tl.int64)
+            if HAS_SWA:
+                pid_slot = tl.load(swa_slot_mapping_ptr + pid_slot)
+
+            # ------------------------------------------------
+            # Stage 3.2: Apply RoPE to K
+            # ------------------------------------------------
+            if pid_slot >= 0:
+                pid_t_slot = pid_slot // BLOCK_SIZE
+                pid_b = pid_slot % BLOCK_SIZE
+                pid_hk = pid_hq // QH_PER_KH
+                if HAVE_K_SCALE:
+                    k_scale = tl.load(k_scale_ptr)
+                else:
+                    k_scale = 1
+                k_ptrs = (
+                    k_ptr
+                    + pid_t * k_stride_t
+                    + pid_hk * k_stride_h
+                    + d_pe_offs * k_stride_d
+                )
+                k_pe = _unit_rope(
+                    k_ptrs,
+                    cos,
+                    sin,
+                    d_pe_offs,
+                    IS_NEOX,
+                    BLOCK_D_pe,
+                    BLOCK_D_HALF_pe,
+                )
+
+                k_out_ptrs = (
+                    k_out_ptr
+                    + pid_t * k_out_stride_t
+                    + pid_hk * k_out_stride_h
+                    + d_pe_offs * k_out_stride_d
+                )
+                tl.store(k_out_ptrs, k_pe.to(k_out_ptr.dtype.element_ty))
+
+                # ------------------------------------------------
+                # Stage 3.3: Optional fp8 scaling before cache
+                # ------------------------------------------------
+
+                k_scale_rcprl = 1 / k_scale
+                k_pe = k_pe * k_scale_rcprl
+
+                # ------------------------------------------------
+                # Stage 3.4: Write K cache
+                #
+                # Two layouts supported:
+                #   FLASH_LAYOUT
+                #   paged KV layout
+                # ------------------------------------------------
+
+                if FLASH_LAYOUT:
+                    k_out_ptrs = (
+                        key_cache_ptr
+                        + pid_t_slot * key_cache_stride_t
+                        + pid_b * key_cache_stride_b
+                        + pid_hk * key_cache_stride_h
+                        + d_pe_offs * key_cache_stride_d
+                    )
+                else:
+                    k_pe = tl.reshape(k_pe, (BLOCK_D_pe // X_SIZE, X_SIZE))
+                    dx_offs = tl.arange(0, BLOCK_D_pe // X_SIZE).to(tl.int64)
+                    x_offs = tl.arange(0, X_SIZE).to(tl.int64)
+                    k_out_ptrs = (
+                        key_cache_ptr
+                        + pid_t_slot * key_cache_stride_t
+                        + pid_hk * key_cache_stride_h
+                        + dx_offs[:, None] * key_cache_stride_d
+                        + pid_b * key_cache_stride_b
+                        + x_offs[None, :] * key_cache_stride_x
+                    )
+
+                tl.store(k_out_ptrs, k_pe.to(key_cache_ptr.dtype.element_ty))
+
+                # ------------------------------------------------
+                # Stage 3.5: Write V cache
+                #
+                # Supports:
+                #   normal layout
+                #   shuffle layout
+                # ------------------------------------------------
+
+                v_ptrs = (
+                    v_ptr
+                    + pid_t * v_stride_t
+                    + pid_hk * v_stride_h
+                    + d_pe_offs * v_stride_d
+                )
+                if HAVE_V_SCALE:
+                    v_scale = tl.load(v_scale_ptr)
+                else:
+                    v_scale = 1
+                v_scale_rcprl = 1 / v_scale
+                v = tl.load(v_ptrs) * v_scale_rcprl
+                if VALUE_SHUFFLE_LAYOUT:
+                    slot_chunk = pid_b // X_SIZE
+                    x_off = pid_b % X_SIZE
+                    v_out_ptrs = (
+                        value_cache_ptr
+                        + pid_t_slot * value_cache_stride_t
+                        + pid_hk * value_cache_stride_h
+                        + slot_chunk * value_cache_stride_slot_chunk
+                        + d_pe_offs.to(tl.int64) * value_cache_stride_d
+                        + x_off * value_cache_stride_x
+                    )
+                else:
+                    v_out_ptrs = (
+                        value_cache_ptr
+                        + pid_t_slot * value_cache_stride_t
+                        + pid_hk * value_cache_stride_h
+                        + d_pe_offs.to(tl.int64) * value_cache_stride_d
+                        + pid_b * value_cache_stride_b
+                    )
+                tl.store(v_out_ptrs, v.to(value_cache_ptr.dtype.element_ty))
+    # ============================================================
+    # Stage 4: Extra KV-only path
+    #
+    # Handles tokens that only require cache update:
+    #   T_slot > T
+    #
+    # No Q / no RoPE on Q branch.
+    # ============================================================
+    else:
+        pid = pid - T * QH + T * KH
+        if pid < T_slot * KH:
+            pid_t = pid // KH
+            pid_hk = pid % KH
+            pid_slot = tl.load(slot_mapping_ptr + pid_t).to(tl.int64)
+            if HAS_SWA:
+                pid_slot = tl.load(swa_slot_mapping_ptr + pid_slot)
+
+            if pid_slot >= 0:
+                pid_t_slot = pid_slot // BLOCK_SIZE
+                pid_b = pid_slot % BLOCK_SIZE
+                if HAVE_K_SCALE:
+                    k_scale = tl.load(k_scale_ptr)
+                else:
+                    k_scale = 1
+                k_ptrs = (
+                    k_ptr
+                    + pid_t * k_stride_t
+                    + pid_hk * k_stride_h
+                    + d_pe_offs * k_stride_d
+                )
+
+                k_pe = tl.load(k_ptrs)
+
+                k_out_ptrs = (
+                    k_out_ptr
+                    + pid_t * k_out_stride_t
+                    + pid_hk * k_out_stride_h
+                    + d_pe_offs * k_out_stride_d
+                )
+                tl.store(k_out_ptrs, k_pe.to(k_out_ptr.dtype.element_ty))
+
+                k_scale_rcprl = 1 / k_scale
+                k_pe = k_pe * k_scale_rcprl
+
+                if FLASH_LAYOUT:
+                    k_out_ptrs = (
+                        key_cache_ptr
+                        + pid_t_slot * key_cache_stride_t
+                        + d_pe_offs * key_cache_stride_d
+                        + pid_b * key_cache_stride_b
+                        + pid_hk * key_cache_stride_h
+                    )
+                else:
+                    k_pe = tl.reshape(k_pe, (BLOCK_D_pe // X_SIZE, X_SIZE))
+                    dx_offs = tl.arange(0, BLOCK_D_pe // X_SIZE).to(tl.int64)
+                    x_offs = tl.arange(0, X_SIZE).to(tl.int64)
+                    k_out_ptrs = (
+                        key_cache_ptr
+                        + pid_t_slot * key_cache_stride_t
+                        + pid_hk * key_cache_stride_h
+                        + dx_offs[:, None] * key_cache_stride_d
+                        + pid_b * key_cache_stride_b
+                        + x_offs[None, :] * key_cache_stride_x
+                    )
+                tl.store(k_out_ptrs, k_pe.to(key_cache_ptr.dtype.element_ty))
+
+                v_ptrs = (
+                    v_ptr
+                    + pid_t * v_stride_t
+                    + pid_hk * v_stride_h
+                    + d_pe_offs * v_stride_d
+                )
+                if HAVE_V_SCALE:
+                    v_scale = tl.load(v_scale_ptr)
+                else:
+                    v_scale = 1
+                v_scale_rcprl = 1 / v_scale
+                v = tl.load(v_ptrs) * v_scale_rcprl
+                if VALUE_SHUFFLE_LAYOUT:
+                    slot_chunk = pid_b // X_SIZE
+                    x_off = pid_b % X_SIZE
+                    v_out_ptrs = (
+                        value_cache_ptr
+                        + pid_t_slot * value_cache_stride_t
+                        + pid_hk * value_cache_stride_h
+                        + slot_chunk * value_cache_stride_slot_chunk
+                        + d_pe_offs * value_cache_stride_d
+                        + x_off * value_cache_stride_x
+                    )
+                else:
+                    v_out_ptrs = (
+                        value_cache_ptr
+                        + pid_t_slot * value_cache_stride_t
+                        + pid_hk * value_cache_stride_h
+                        + d_pe_offs * value_cache_stride_d
+                        + pid_b * value_cache_stride_b
+                    )
+                tl.store(v_out_ptrs, v.to(value_cache_ptr.dtype.element_ty))
+
+
+def fused_qk_rope_reshape_and_cache(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    pos: torch.Tensor,
+    cos_sin: torch.Tensor,
+    k_scale: torch.Tensor,
+    v_scale: torch.Tensor,
+    is_neox: bool,
+    flash_layout: bool,
+    apply_scale: bool = True,
+    offs: torch.Tensor = None,
+    q_out: torch.Tensor = None,
+    k_out: torch.Tensor = None,
+    output_zeros: bool = True,
+    zeros_out: torch.Tensor = None,
+    swa_slot_mapping=None,
+):
+    """
+    Perform RoPE on q and k and along the last dimension and copy k and v in to key_cache and value_cache inplace
+
+    Key parameters:
+    - q: shape (T, QH, D).
+    - k: shape (T_slot, KH, D).
+    - v: shape (T_slot, KH, D).
+    - if flash_layout:
+    -     key_cache: shape (T_cache, block_size, KH, D).
+    -     value_cache: shape (T_cache, block_size, KH, D).
+    - else:
+    -     key_cache: shape (T_cache, KH, D // x, block_size, x).
+    -     value_cache: shape (T_cache, KH, D, block_size).
+    - slot_mapping: shape (T_slot, ).
+
+    T is the number of decode tokens, T_cahce * block_size is the max number of tokens of kv_cache
+    QH must be multiple of KH
+
+    Returns:
+    - q_out: same shape as input q.
+    - k_out: same shape as input k.
+    - key_cache: same shape as input key_cache (inplace).
+    - value_cache: same shape as input value_cache (inplace).
+    - zeros_out: same shape as input q.
+    """
+
+    t, qh, d = q.shape
+    tk, kh, dk = k.shape
+    tv, vh, dv = v.shape
+    if flash_layout:
+        t_cache, block_size, kh_cache, dk_cache = key_cache.shape
+        t_cache_v, block_size_v, vh_cache, dv_cache = value_cache.shape
+        value_shuffle_layout = False
+    else:
+        t_cache, kh_cache, dkx_cache, block_size, x_cache = key_cache.shape
+        if value_cache.ndim == 5:
+            # value_cache shuffle: (num_blocks, num_kv_heads, block_size // x, head_size, x)
+            t_cache_v, vh_cache, slot_chunk_v, dv_cache, x_v = value_cache.shape
+            value_shuffle_layout = True
+            block_size_v = slot_chunk_v * x_v
+            assert block_size_v == block_size and x_v == x_cache, (
+                f"value_cache shuffle (T,KH,block_size//x,D,x) must match key: "
+                f"{block_size_v=} {block_size=} {x_v=} {x_cache=}"
+            )
+        else:
+            t_cache_v, vh_cache, dv_cache, block_size_v = value_cache.shape
+            value_shuffle_layout = False
+    (t_slot,) = slot_mapping.shape
+
+    assert (
+        t == tk == tv and t_slot <= tk
+    ), f"Number of tokens should be identical for q, kand v. The number of tokens of slot_mapping should no more than that of q, k and v, {t=} {tk=} {tv=} {t_slot=}"
+    assert (
+        block_size == block_size_v
+    ), f"block size should be identical for key_cache, and value_cache {block_size} {block_size_v}"
+    assert (
+        kh == vh == kh_cache == vh_cache
+    ), "KV head should be identical for k, v, key_cache, and value_cache"
+    assert (
+        t_cache == t_cache_v
+    ), "Number of tokens should be identical for key_cache, and value_cache"
+    if flash_layout:
+        assert (
+            d == dk == dv == dk_cache == dv_cache
+        ), "D dimension should be identical for q, k, and v"
+    else:
+        assert (
+            d == dk == dv == dkx_cache * x_cache == dv_cache
+        ), "D dimension should be identical for q, k, and v"
+        assert x_cache == triton.next_power_of_2(x_cache), "x_size should be power of 2"
+
+    assert d == triton.next_power_of_2(d), "D dimension should be power of 2"
+    assert block_size == triton.next_power_of_2(
+        block_size
+    ), "block_size should be power of 2"
+    assert qh % kh == 0, "Q heads must be multiple of H heads"
+    d_freq = cos_sin.shape[-1] // 2
+    assert (d_freq == d // 2) or (
+        d_freq == d
+    ), "cos/sin last dim should be the same or half of the qk last dim"
+    reuse_freqs_front_part = d_freq == d // 2
+
+    if q_out is None:
+        q_out = torch.empty((t, qh, d), dtype=q.dtype, device=q.device)
+
+    if k_out is None:
+        k_out = torch.empty((tk, kh, dk), dtype=k.dtype, device=q.device)
+
+    if zeros_out is not None:
+        tz, qhz, dz = zeros_out.shape
+        assert (
+            t == tz and qh == qhz and d == dz
+        ), f"q and zeros shape mismatch {q.shape=} {zeros_out.shape=}"
+        output_zeros = True
+    elif output_zeros:
+        zeros_out = torch.empty((t, qh, d), dtype=q.dtype, device=q.device)
+    else:
+        zeros_out = None
+
+    n_pid = t * qh + (t_slot - t) * kh if t_slot >= t else t * qh
+    grid = (n_pid, 1, 1)
+    _fused_qk_rope_reshape_and_cache_kernel[grid](
+        q,
+        k,
+        v,
+        pos,
+        cos_sin,
+        offs,
+        key_cache,
+        value_cache,
+        slot_mapping,
+        swa_slot_mapping,
+        q_out,
+        k_out,
+        zeros_out,
+        t,
+        t_slot,
+        *q.stride(),
+        *k.stride(),
+        *v.stride(),
+        cos_sin.stride(0),
+        cos_sin.stride(-1),
+        *q_out.stride(),
+        *k_out.stride(),
+        key_cache.stride(0) if not flash_layout else key_cache.stride(0),
+        key_cache.stride(1) if not flash_layout else key_cache.stride(2),
+        key_cache.stride(2) if not flash_layout else key_cache.stride(3),
+        key_cache.stride(3) if not flash_layout else key_cache.stride(1),
+        key_cache.stride(4) if not flash_layout else 0,
+        value_cache.stride(0) if not flash_layout else value_cache.stride(0),
+        value_cache.stride(1) if not flash_layout else value_cache.stride(2),
+        (
+            value_cache.stride(3)
+            if (not flash_layout and value_shuffle_layout)
+            else (value_cache.stride(2) if not flash_layout else value_cache.stride(3))
+        ),
+        (
+            0
+            if (not flash_layout and value_shuffle_layout)
+            else (value_cache.stride(3) if not flash_layout else value_cache.stride(1))
+        ),
+        value_cache.stride(2) if (not flash_layout and value_shuffle_layout) else 0,
+        value_cache.stride(4) if (not flash_layout and value_shuffle_layout) else 0,
+        zeros_out.stride(0) if zeros_out is not None else 0,
+        zeros_out.stride(1) if zeros_out is not None else 0,
+        zeros_out.stride(2) if zeros_out is not None else 0,
+        k_scale_ptr=k_scale,
+        v_scale_ptr=v_scale,
+        QH_PER_KH=qh // kh,
+        QH=qh,
+        KH=kh,
+        REUSE_FREQS_FRONT_PART=reuse_freqs_front_part,
+        IS_NEOX=is_neox,
+        BLOCK_D_pe=d,
+        BLOCK_D_HALF_pe=d // 2,
+        BLOCK_SIZE=block_size,
+        X_SIZE=x_cache if not flash_layout else 0,
+        FLASH_LAYOUT=flash_layout,
+        VALUE_SHUFFLE_LAYOUT=value_shuffle_layout,
+        HAVE_POS=(offs is not None),
+        HAVE_K_SCALE=(k_scale is not None and apply_scale),
+        HAVE_V_SCALE=(v_scale is not None and apply_scale),
+        HAVE_ZEROS=output_zeros,
+        HAS_SWA=(swa_slot_mapping is not None),
+        num_warps=1,
+    )
+
+    if zeros_out is not None:
+        return q_out.view(-1, qh * d), k_out, key_cache, value_cache, zeros_out
+    return q_out.view(-1, qh * d), k_out, key_cache, value_cache

python/sglang/kernels/ops/kvcache/trtllm_fp8_kv_kernel.py

@@ -0,0 +1,504 @@
+"""
+Fused FP8 quantization + paged KV cache write kernel for TRTLLM MHA backend.
+
+This kernel fuses the following operations:
+1. FP8 quantization of K and V tensors (from BF16/FP16 to FP8)
+2. Per-token or per-page scale computation
+3. Writing quantized K/V to paged KV cache layout
+
+Performance benefits:
+- Eliminates intermediate FP8 tensors in memory
+- Reduces kernel launch overhead
+- Better memory bandwidth utilization
+"""
+
+import logging
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+logger = logging.getLogger(__name__)
+
+
+@triton.jit
+def _process_kv_tensor(
+    token_id,
+    head_block_id,
+    page_id,
+    page_offset,
+    input_ptr,
+    cache_ptr,
+    inv_scale,
+    use_provided_scale: tl.constexpr,
+    num_kv_heads: tl.constexpr,
+    head_dim: tl.constexpr,
+    input_stride_token: tl.constexpr,
+    input_stride_head: tl.constexpr,
+    input_stride_dim: tl.constexpr,
+    cache_stride_page: tl.constexpr,
+    cache_stride_offset: tl.constexpr,
+    cache_stride_head: tl.constexpr,
+    cache_stride_dim: tl.constexpr,
+    BLOCK_HEAD: tl.constexpr,
+    BLOCK_DIM: tl.constexpr,
+):
+    """Process a block of heads for a single K or V tensor."""
+    head_idx = head_block_id * BLOCK_HEAD
+    num_heads_in_block = min(BLOCK_HEAD, num_kv_heads - head_idx)
+
+    for dim_idx in range(0, head_dim, BLOCK_DIM):
+        num_dims_in_block = min(BLOCK_DIM, head_dim - dim_idx)
+
+        head_offsets = head_idx + tl.arange(0, BLOCK_HEAD)
+        dim_offsets = dim_idx + tl.arange(0, BLOCK_DIM)
+
+        head_mask = head_offsets < (head_idx + num_heads_in_block)
+        dim_mask = dim_offsets < (dim_idx + num_dims_in_block)
+
+        # Load from input using 3D strides
+        input_offsets = (
+            token_id * input_stride_token
+            + head_offsets[:, None] * input_stride_head
+            + dim_offsets[None, :] * input_stride_dim
+        )
+        mask = head_mask[:, None] & dim_mask[None, :]
+
+        block = tl.load(input_ptr + input_offsets, mask=mask, other=0.0)
+
+        # Quantize to FP8
+        if use_provided_scale:
+            block_fp8 = (block * inv_scale).to(tl.float8e4nv)
+        else:
+            block_fp8 = block.to(tl.float8e4nv)
+
+        # Write to cache at [page_id, page_offset, head, dim]
+        cache_offsets = (
+            page_id * cache_stride_page
+            + page_offset * cache_stride_offset
+            + head_offsets[:, None] * cache_stride_head
+            + dim_offsets[None, :] * cache_stride_dim
+        )
+
+        tl.store(cache_ptr + cache_offsets, block_fp8, mask=mask)
+
+
+@triton.jit
+def _fused_fp8_set_kv_buffer_kernel(
+    # Input tensors (post-RoPE K and V in FP16/BF16)
+    k_ptr,  # [num_tokens, num_kv_heads, head_dim]
+    v_ptr,  # [num_tokens, num_kv_heads, head_dim]
+    # Output KV cache buffers (FP8 paged layout)
+    k_cache_ptr,  # [total_slots, num_kv_heads, head_dim]
+    v_cache_ptr,  # [total_slots, num_kv_heads, head_dim]
+    # Cache location indices
+    cache_loc_ptr,  # [num_tokens] -> token to cache location mapping
+    # Pointers to scalar inverse scales (computed on GPU in wrapper)
+    inv_k_scale_ptr,  # pointer to 0-D tensor on GPU
+    inv_v_scale_ptr,  # pointer to 0-D tensor on GPU
+    use_provided_scale: tl.constexpr,  # whether to use provided scale
+    # Tensor dimensions
+    num_kv_heads: tl.constexpr,
+    head_dim: tl.constexpr,
+    page_size: tl.constexpr,
+    # Strides for K input [num_tokens, num_kv_heads, head_dim]
+    k_stride_token: tl.constexpr,
+    k_stride_head: tl.constexpr,
+    k_stride_dim: tl.constexpr,
+    # Strides for K cache [total_slots, num_kv_heads, head_dim] (logically paged)
+    k_cache_stride_page: tl.constexpr,
+    k_cache_stride_offset: tl.constexpr,
+    k_cache_stride_head: tl.constexpr,
+    k_cache_stride_dim: tl.constexpr,
+    # Strides for V input [num_tokens, num_kv_heads, head_dim]
+    v_stride_token: tl.constexpr,
+    v_stride_head: tl.constexpr,
+    v_stride_dim: tl.constexpr,
+    # Strides for V cache [total_slots, num_kv_heads, head_dim] (logically paged)
+    v_cache_stride_page: tl.constexpr,
+    v_cache_stride_offset: tl.constexpr,
+    v_cache_stride_head: tl.constexpr,
+    v_cache_stride_dim: tl.constexpr,
+    # Block sizes
+    BLOCK_HEAD: tl.constexpr,  # Number of heads per block
+    BLOCK_DIM: tl.constexpr,  # Head dimension block size
+):
+    """
+    Fused FP8 quantization + paged KV cache write kernel.
+
+    Each program processes one token-head_block-kv combination, quantizing and writing
+    to the appropriate page in the KV cache.
+
+    Grid: (num_tokens, num_head_blocks, 2) where dim2: 0=K, 1=V
+    """
+    # Get program IDs
+    token_id = tl.program_id(0)
+    head_block_id = tl.program_id(1)
+    kv_idx = tl.program_id(2)  # 0 for K, 1 for V
+
+    # Get cache location for this token
+    cache_loc = tl.load(cache_loc_ptr + token_id)
+
+    # Compute page_id and offset within page
+    page_id = cache_loc // page_size
+    page_offset = cache_loc % page_size
+
+    # Select K or V based on kv_idx
+    if kv_idx == 0:
+        # Process K tensor
+        if use_provided_scale:
+            inv_scale = tl.load(inv_k_scale_ptr)
+        else:
+            inv_scale = 1.0
+        _process_kv_tensor(
+            token_id,
+            head_block_id,
+            page_id,
+            page_offset,
+            k_ptr,
+            k_cache_ptr,
+            inv_scale,
+            use_provided_scale,
+            num_kv_heads,
+            head_dim,
+            k_stride_token,
+            k_stride_head,
+            k_stride_dim,
+            k_cache_stride_page,
+            k_cache_stride_offset,
+            k_cache_stride_head,
+            k_cache_stride_dim,
+            BLOCK_HEAD,
+            BLOCK_DIM,
+        )
+    else:
+        # Process V tensor
+        if use_provided_scale:
+            inv_scale = tl.load(inv_v_scale_ptr)
+        else:
+            inv_scale = 1.0
+        _process_kv_tensor(
+            token_id,
+            head_block_id,
+            page_id,
+            page_offset,
+            v_ptr,
+            v_cache_ptr,
+            inv_scale,
+            use_provided_scale,
+            num_kv_heads,
+            head_dim,
+            v_stride_token,
+            v_stride_head,
+            v_stride_dim,
+            v_cache_stride_page,
+            v_cache_stride_offset,
+            v_cache_stride_head,
+            v_cache_stride_dim,
+            BLOCK_HEAD,
+            BLOCK_DIM,
+        )
+
+
+def fused_fp8_set_kv_buffer(
+    k: torch.Tensor,  # [num_tokens, num_kv_heads, head_dim] or [num_tokens, num_kv_heads * head_dim]
+    v: torch.Tensor,  # [num_tokens, num_kv_heads, head_dim] or [num_tokens, num_kv_heads * head_dim]
+    k_cache: torch.Tensor,  # [total_slots, num_kv_heads, head_dim] or [num_pages, page_size, num_kv_heads, head_dim]
+    v_cache: torch.Tensor,  # [total_slots, num_kv_heads, head_dim] or [num_pages, page_size, num_kv_heads, head_dim]
+    cache_loc: torch.Tensor,  # [num_tokens], dtype=int32
+    k_scale: Optional[
+        float
+    ] = None,  # Scalar scale (matching original set_kv_buffer signature)
+    v_scale: Optional[float] = None,
+    page_size: int = 16,
+    use_triton: bool = True,  # Whether to use Triton kernel (set to False to force naive fallback)
+) -> None:
+    """
+    Python wrapper for the fused FP8 quantization + paged KV cache write kernel.
+
+    This function replicates the exact behavior of the original set_kv_buffer but with
+    a fused kernel that combines FP8 quantization and cache write.
+
+    Args:
+        k: Key tensor after RoPE, can be 2D or 3D
+        v: Value tensor, can be 2D or 3D
+        k_cache: Paged K cache buffer in FP8
+        v_cache: Paged V cache buffer in FP8
+        cache_loc: Cache location for each token, shape [num_tokens]
+        k_scale: Optional scalar scale for K (matching original set_kv_buffer)
+        v_scale: Optional scalar scale for V (matching original set_kv_buffer)
+        page_size: Number of tokens per page
+        use_triton: Whether to use optimized Triton kernel
+    """
+    num_tokens = k.shape[0]
+
+    # Step 1: Infer num_kv_heads and head_dim from cache shape
+    if k_cache.ndim == 3:
+        # 3D cache layout: [total_slots, num_kv_heads, head_dim]
+        total_slots, num_kv_heads, head_dim = k_cache.shape
+        assert (
+            total_slots % page_size == 0
+        ), f"total_slots ({total_slots}) must be divisible by page_size ({page_size})"
+        num_pages = total_slots // page_size
+    elif k_cache.ndim == 4:
+        # 4D cache layout: [num_pages, page_size, num_kv_heads, head_dim]
+        num_pages, ps, num_kv_heads, head_dim = k_cache.shape
+        assert (
+            ps == page_size
+        ), f"page_size mismatch: cache has {ps}, expected {page_size}"
+        total_slots = num_pages * page_size
+    else:
+        raise ValueError(f"Unsupported k_cache.ndim={k_cache.ndim}, expected 3 or 4")
+
+    # Step 2: Validate k, v shapes and normalize
+    # Store original 3D shape for Triton path
+    k_3d = None
+    v_3d = None
+
+    if k.ndim == 3:
+        # Input is [num_tokens, num_kv_heads, head_dim]
+        assert (
+            k.shape[1] == num_kv_heads
+        ), f"num_kv_heads mismatch: k.shape[1]={k.shape[1]} vs cache={num_kv_heads}"
+        assert (
+            k.shape[2] == head_dim
+        ), f"head_dim mismatch: k.shape[2]={k.shape[2]} vs cache={head_dim}"
+        assert v.shape[1] == num_kv_heads and v.shape[2] == head_dim, "v shape mismatch"
+
+        # Keep 3D for Triton kernel
+        k_3d = k
+        v_3d = v
+        # Create 2D view for naive fallback (will be used only if use_triton=False)
+        k_2d = k.reshape(num_tokens, num_kv_heads * head_dim)
+        v_2d = v.reshape(num_tokens, num_kv_heads * head_dim)
+    elif k.ndim == 2:
+        # Input is already [num_tokens, num_kv_heads * head_dim]
+        assert (
+            k.shape[1] == num_kv_heads * head_dim
+        ), f"k.shape[1]={k.shape[1]} != {num_kv_heads * head_dim}"
+        assert (
+            v.shape[1] == num_kv_heads * head_dim
+        ), f"v.shape[1]={v.shape[1]} != {num_kv_heads * head_dim}"
+
+        # Create 3D view for Triton kernel
+        k_3d = k.view(num_tokens, num_kv_heads, head_dim)
+        v_3d = v.view(num_tokens, num_kv_heads, head_dim)
+        # Keep 2D for naive
+        k_2d = k
+        v_2d = v
+    else:
+        raise ValueError(f"Unsupported k.ndim={k.ndim}, expected 2 or 3")
+
+    # Step 3: Compute cache strides based on layout
+    if k_cache.ndim == 3:
+        # 3D cache: [total_slots, num_kv_heads, head_dim]
+        stride_slot = k_cache.stride(0)
+        stride_head = k_cache.stride(1)
+        stride_dim = k_cache.stride(2)
+
+        k_cache_stride_page = stride_slot * page_size
+        k_cache_stride_offset = stride_slot
+        k_cache_stride_head = stride_head
+        k_cache_stride_dim = stride_dim
+
+        v_stride_slot = v_cache.stride(0)
+        v_stride_head = v_cache.stride(1)
+        v_stride_dim = v_cache.stride(2)
+
+        v_cache_stride_page = v_stride_slot * page_size
+        v_cache_stride_offset = v_stride_slot
+        v_cache_stride_head = v_stride_head
+        v_cache_stride_dim = v_stride_dim
+    else:
+        # 4D cache: [num_pages, page_size, num_kv_heads, head_dim]
+        k_cache_stride_page = k_cache.stride(0)
+        k_cache_stride_offset = k_cache.stride(1)
+        k_cache_stride_head = k_cache.stride(2)
+        k_cache_stride_dim = k_cache.stride(3)
+
+        v_cache_stride_page = v_cache.stride(0)
+        v_cache_stride_offset = v_cache.stride(1)
+        v_cache_stride_head = v_cache.stride(2)
+        v_cache_stride_dim = v_cache.stride(3)
+
+    # Decide whether to use provided scale
+    use_provided_scale = k_scale is not None and v_scale is not None
+
+    if use_triton and num_tokens > 0:
+        # Use optimized Triton kernel
+        # Compute input strides for 3D k, v: [num_tokens, num_kv_heads, head_dim]
+        k_stride_token = k_3d.stride(0)
+        k_stride_head = k_3d.stride(1)
+        k_stride_dim = k_3d.stride(2)
+
+        v_stride_token = v_3d.stride(0)
+        v_stride_head = v_3d.stride(1)
+        v_stride_dim = v_3d.stride(2)
+
+        # Block sizes for tiling (tunable)
+        BLOCK_HEAD = min(num_kv_heads, 8)  # Process up to 8 heads at once
+        BLOCK_DIM = min(head_dim, 128)  # Process up to 128 dims at once
+
+        # Compute number of head blocks
+        num_head_blocks = (num_kv_heads + BLOCK_HEAD - 1) // BLOCK_HEAD
+
+        # Grid: (num_tokens, num_head_blocks, 2)
+        # - dim 0: tokens
+        # - dim 1: head blocks
+        # - dim 2: K/V (0=K, 1=V)
+        grid = (num_tokens, num_head_blocks, 2)
+
+        device = k_3d.device
+
+        def _to_tensor_scale(scale):
+            """Convert scale to 0-D CUDA tensor (accepts Python float or Tensor)."""
+            if isinstance(scale, torch.Tensor):
+                return scale.to(device=device, dtype=torch.float32)
+            else:
+                # Python float / np scalar
+                return torch.tensor(float(scale), device=device, dtype=torch.float32)
+
+        # Compute inverse scales on GPU to avoid GPU→CPU sync in CUDA graph capture.
+        # Previously we used float(k_scale) which triggers synchronization and fails
+        # during CUDA graph capture with cudaErrorStreamCaptureUnsupported.
+        if use_provided_scale:
+            k_scale_tensor = _to_tensor_scale(k_scale)
+            v_scale_tensor = _to_tensor_scale(v_scale)
+
+            # Pure GPU scalar operation, safe for CUDA graph
+            inv_k_scale = (1.0 / k_scale_tensor).to(device=device, dtype=torch.float32)
+            inv_v_scale = (1.0 / v_scale_tensor).to(device=device, dtype=torch.float32)
+
+            inv_k_scale_ptr = inv_k_scale
+            inv_v_scale_ptr = inv_v_scale
+        else:
+            # When use_provided_scale=False, kernel uses constant 1.0 for inv_scale.
+            # Triton will optimize away the tl.load() calls via constant folding.
+            # We pass dummy pointers (k_3d) which won't be accessed in the kernel.
+            # This avoids creating new GPU tensors during CUDA graph capture.
+            inv_k_scale_ptr = k_3d
+            inv_v_scale_ptr = k_3d
+
+        # Launch Triton kernel
+        _fused_fp8_set_kv_buffer_kernel[grid](
+            k_3d,
+            v_3d,
+            k_cache,
+            v_cache,
+            cache_loc,
+            inv_k_scale_ptr,
+            inv_v_scale_ptr,
+            use_provided_scale,
+            num_kv_heads,
+            head_dim,
+            page_size,
+            k_stride_token,
+            k_stride_head,
+            k_stride_dim,
+            k_cache_stride_page,
+            k_cache_stride_offset,
+            k_cache_stride_head,
+            k_cache_stride_dim,
+            v_stride_token,
+            v_stride_head,
+            v_stride_dim,
+            v_cache_stride_page,
+            v_cache_stride_offset,
+            v_cache_stride_head,
+            v_cache_stride_dim,
+            BLOCK_HEAD=BLOCK_HEAD,
+            BLOCK_DIM=BLOCK_DIM,
+        )
+    else:
+        # Fallback to naive implementation
+        _naive_fp8_set_kv_buffer(
+            k_2d, v_2d, k_cache, v_cache, cache_loc, k_scale, v_scale, page_size
+        )
+
+
+def _naive_fp8_set_kv_buffer(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    cache_loc: torch.Tensor,
+    k_scale: Optional[float],
+    v_scale: Optional[float],
+    page_size: int,
+) -> None:
+    """
+    Naive fallback implementation that mimics the original set_kv_buffer logic.
+
+    This directly replicates the behavior of MHATokenToKVPool.set_kv_buffer:
+    1. Apply scale (if k.dtype != cache.dtype and scale is provided)
+    2. Convert to FP8
+    3. Write to cache at cache_loc
+
+    Args:
+        k: [num_tokens, num_kv_heads * head_dim], already reshaped to 2D
+        v: [num_tokens, num_kv_heads * head_dim], already reshaped to 2D
+        k_cache: [total_slots, num_kv_heads, head_dim] or [num_pages, page_size, num_kv_heads, head_dim]
+        v_cache: Same shape as k_cache
+        cache_loc: [num_tokens]
+        k_scale: Optional scale for K
+        v_scale: Optional scale for V
+        page_size: Tokens per page
+    """
+    num_tokens = k.shape[0]
+
+    # Infer dimensions from cache
+    if k_cache.ndim == 3:
+        num_kv_heads = k_cache.shape[1]
+        head_dim = k_cache.shape[2]
+    elif k_cache.ndim == 4:
+        num_kv_heads = k_cache.shape[2]
+        head_dim = k_cache.shape[3]
+    else:
+        raise ValueError(f"Unsupported k_cache.ndim={k_cache.ndim}")
+
+    # Determine target dtype and storage dtype
+    # See: python/sglang/srt/mem_cache/memory_pool.py:445-449
+    store_dtype = k_cache.dtype
+    if store_dtype == torch.uint8:
+        # Cache is stored as uint8 for FP8 (due to index_put limitation)
+        dtype = torch.float8_e4m3fn  # Logical dtype
+    else:
+        dtype = store_dtype  # Cache dtype is the logical dtype
+
+    # Replicate the original set_kv_buffer behavior
+    # See: python/sglang/srt/mem_cache/memory_pool.py:777-799
+    if k.dtype != dtype:
+        # Need quantization - clone first to avoid modifying input
+        k = k.clone()
+        v = v.clone()
+
+        if k_scale is not None:
+            k.div_(k_scale)  # In-place division
+        if v_scale is not None:
+            v.div_(v_scale)  # In-place division
+
+        k = k.to(dtype)
+        v = v.to(dtype)
+
+    # View FP8 as uint8 if needed (for index_put compatibility)
+    if store_dtype == torch.uint8 and dtype in (torch.float8_e5m2, torch.float8_e4m3fn):
+        k = k.view(torch.uint8)
+        v = v.view(torch.uint8)
+
+    # Reshape from [T, H*D] to [T, H, D]
+    k = k.view(num_tokens, num_kv_heads, head_dim)
+    v = v.view(num_tokens, num_kv_heads, head_dim)
+
+    # Write to cache using advanced indexing (same as original)
+    if k_cache.ndim == 3:
+        # 3D cache: [total_slots, H, D]
+        k_cache[cache_loc] = k
+        v_cache[cache_loc] = v
+    else:
+        # 4D cache: [num_pages, page_size, H, D]
+        # Decompose loc into page_id and page_offset (vectorized)
+        page_ids = cache_loc // page_size
+        page_offsets = cache_loc % page_size
+        k_cache[page_ids, page_offsets] = k
+        v_cache[page_ids, page_offsets] = v

python/sglang/kernels/ops/kvcache/trtllm_mha_graph_metadata.py

@@ -0,0 +1,198 @@
+"""Fused CUDA-graph metadata update for the TRTLLM MHA backend.
+
+`TRTLLMHAAttnBackend._apply_cuda_graph_metadata` used to rebuild the
+page table(s) and seqlen buffers with ~25 small aten ops per graph
+replay (index gathers, floor_divide, cumsum, dtype casts, copies).
+On some CPUs that is ~0.7-1.0 ms of pure host dispatch, repeated 4x
+per decode step (2 draft-decode steps + target-verify + draft-extend)
+on every TP rank. The resulting per-rank CPU jitter skews the
+cudaGraphLaunch across ranks and is paid as spin time inside the first
+custom all-reduce of every replayed graph.
+
+This kernel performs the whole update in ONE launch:
+  - cache_seqlens[i]  = seq_lens[i] + seqlen_offset            (int32)
+  - cu_seqlens_k[1:]  = cumsum(cache_seqlens)                  (int32)
+  - cu_seqlens_q[1:]  = cumsum(qlens) or arange*q_stride       (optional)
+  - page_table[i, p]  = req_to_token[req_pool_indices[i],
+                                     p * page_size] // page_size
+  - swa_page_table    = full_to_swa_mapping[token] // page_size (optional)
+  - swa_out_cache_loc = full_to_swa_mapping[out_cache_loc], zero padded
+                                                               (optional)
+"""
+
+import triton
+import triton.language as tl
+
+# cu_seqlens_q handling inside the fused kernel
+Q_MODE_NONE = 0  # cu_seqlens_q is preset (decode / target-verify)
+Q_MODE_CUMSUM = 1  # cu_seqlens_q[1:] = cumsum(qlens)   (draft-extend)
+Q_MODE_STRIDED = 2  # cu_seqlens_q[1:] = arange*q_stride (draft-extend v2)
+
+
+@triton.jit
+def update_trtllm_mha_graph_metadata_kernel(
+    # inputs
+    req_pool_indices_ptr,  # [bs] int
+    seq_lens_ptr,  # [bs] int
+    req_to_token_ptr,  # [pool_size, req_to_token_stride] int32
+    swa_mapping_ptr,  # [full_size + page_size + 1] int64, or None
+    out_cache_loc_ptr,  # [num_out_tokens] int64, or None
+    qlens_ptr,  # [bs] int, or None (Q_MODE_CUMSUM only)
+    # outputs
+    cache_seqlens_ptr,  # [bs] int32
+    cu_seqlens_k_ptr,  # [bs + 1] int32
+    cu_seqlens_q_ptr,  # [bs + 1] int32, or None
+    page_table_ptr,  # [bs, page_table_stride] int32
+    swa_page_table_ptr,  # [bs, swa_page_table_stride] int32, or None
+    swa_out_cache_loc_ptr,  # [swa_out_len] int64, or None
+    # scalars
+    bs,
+    seqlen_offset,  # added to seq_lens for cache_seqlens / cu_seqlens_k
+    max_seq_pages,  # page-table columns to (re)write per row
+    q_stride,  # Q_MODE_STRIDED stride
+    num_out_tokens,  # valid prefix of out_cache_loc
+    swa_out_len,  # full swa_out_cache_loc length (zero-padded tail)
+    req_to_token_stride,
+    page_table_stride,
+    swa_page_table_stride,
+    # constexpr
+    PAGE_SIZE: tl.constexpr,
+    HAS_SWA: tl.constexpr,
+    HAS_SWA_OUT: tl.constexpr,
+    Q_MODE: tl.constexpr,
+    PAGE_BLOCK: tl.constexpr,
+    BS_BLOCK: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+
+    if pid < bs:
+        # One program per batch row: cache_seqlens + page table row(s).
+        req_pool_index = tl.load(req_pool_indices_ptr + pid).to(tl.int64)
+        seqlen = (tl.load(seq_lens_ptr + pid) + seqlen_offset).to(tl.int32)
+        tl.store(cache_seqlens_ptr + pid, seqlen)
+
+        row_in = req_to_token_ptr + req_pool_index * req_to_token_stride
+        row_out = page_table_ptr + pid.to(tl.int64) * page_table_stride
+        if HAS_SWA:
+            swa_row_out = swa_page_table_ptr + pid.to(tl.int64) * swa_page_table_stride
+        for i in range(tl.cdiv(max_seq_pages, PAGE_BLOCK)):
+            page_idx = i * PAGE_BLOCK + tl.arange(0, PAGE_BLOCK)
+            mask = page_idx < max_seq_pages
+            token = tl.load(
+                row_in + page_idx.to(tl.int64) * PAGE_SIZE, mask=mask, other=0
+            )
+            tl.store(row_out + page_idx, token // PAGE_SIZE, mask=mask)
+            if HAS_SWA:
+                token64 = token.to(tl.int64)
+                # Real req_to_token slots are >=0; the token>=0 guard + other=-1 mirror
+                # the swa_out_cache_loc -1 sentinel (uniform handling, no wrap).
+                swa_token = tl.load(
+                    swa_mapping_ptr + token64, mask=mask & (token64 >= 0), other=-1
+                )
+                swa_page = tl.where(swa_token < 0, -1, swa_token // PAGE_SIZE)
+                tl.store(swa_row_out + page_idx, swa_page.to(tl.int32), mask=mask)
+    elif pid == bs:
+        # Single program: cu_seqlens_k (+ optional cu_seqlens_q) cumsum.
+        offs = tl.arange(0, BS_BLOCK)
+        mask = offs < bs
+        seqlens = (tl.load(seq_lens_ptr + offs, mask=mask, other=0)).to(tl.int32)
+        seqlens = tl.where(mask, seqlens + seqlen_offset, 0)
+        tl.store(cu_seqlens_k_ptr + 1 + offs, tl.cumsum(seqlens, axis=0), mask=mask)
+        if Q_MODE == 1:  # Q_MODE_CUMSUM
+            qlens = tl.load(qlens_ptr + offs, mask=mask, other=0).to(tl.int32)
+            qlens = tl.where(mask, qlens, 0)
+            tl.store(cu_seqlens_q_ptr + 1 + offs, tl.cumsum(qlens, axis=0), mask=mask)
+        if Q_MODE == 2:  # Q_MODE_STRIDED
+            tl.store(
+                cu_seqlens_q_ptr + 1 + offs,
+                ((offs + 1) * q_stride).to(tl.int32),
+                mask=mask,
+            )
+    else:
+        # Remaining programs: swa_out_cache_loc translate + zero padding.
+        if HAS_SWA_OUT:
+            out_idx = (pid - bs - 1) * PAGE_BLOCK + tl.arange(0, PAGE_BLOCK)
+            in_range = out_idx < swa_out_len
+            is_real = in_range & (out_idx < num_out_tokens)
+            loc = tl.load(out_cache_loc_ptr + out_idx, mask=is_real, other=0)
+            translated = tl.load(
+                swa_mapping_ptr + loc, mask=is_real & (loc >= 0), other=0
+            )
+            translated = tl.where(is_real & (loc < 0), -1, translated)
+            tl.store(swa_out_cache_loc_ptr + out_idx, translated, mask=in_range)
+
+
+def update_trtllm_mha_graph_metadata(
+    *,
+    req_pool_indices,
+    seq_lens,
+    req_to_token,
+    cache_seqlens,
+    cu_seqlens_k,
+    page_table,
+    bs: int,
+    seqlen_offset: int,
+    max_seq_pages: int,
+    page_size: int,
+    swa_mapping=None,
+    swa_page_table=None,
+    out_cache_loc=None,
+    swa_out_cache_loc=None,
+    cu_seqlens_q=None,
+    qlens=None,
+    q_stride: int = 0,
+    q_mode: int = Q_MODE_NONE,
+):
+    """Launch the fused metadata update (one kernel for the whole replay init)."""
+    if bs == 0:
+        return
+
+    # Launch-block width: page-table columns each program writes per iteration
+    # (also the swa_out_cache_loc tile width). 512 keeps the per-program working
+    # set small enough to stay off the register-pressure / occupancy cliff while
+    # being wide enough to cover the static page-table width in few iterations.
+    PAGE_BLOCK = 512
+    has_swa = swa_page_table is not None
+    has_swa_out = swa_out_cache_loc is not None
+
+    swa_out_len = swa_out_cache_loc.shape[0] if has_swa_out else 0
+    if has_swa_out and out_cache_loc is not None:
+        num_out_tokens = min(swa_out_len, out_cache_loc.shape[0])
+    else:
+        num_out_tokens = 0
+    if num_out_tokens == 0:
+        # All loads are masked out; pass a valid dummy pointer for codegen.
+        out_cache_loc = swa_out_cache_loc
+
+    grid_extra = triton.cdiv(swa_out_len, PAGE_BLOCK) if has_swa_out else 0
+    grid = (bs + 1 + grid_extra,)
+
+    update_trtllm_mha_graph_metadata_kernel[grid](
+        req_pool_indices,
+        seq_lens,
+        req_to_token,
+        swa_mapping,
+        out_cache_loc,
+        qlens,
+        cache_seqlens,
+        cu_seqlens_k,
+        cu_seqlens_q,
+        page_table,
+        swa_page_table,
+        swa_out_cache_loc,
+        bs,
+        seqlen_offset,
+        max_seq_pages,
+        q_stride,
+        num_out_tokens,
+        swa_out_len,
+        req_to_token.stride(0),
+        page_table.stride(0),
+        swa_page_table.stride(0) if has_swa else 0,
+        PAGE_SIZE=page_size,
+        HAS_SWA=has_swa,
+        HAS_SWA_OUT=has_swa_out,
+        Q_MODE=q_mode,
+        PAGE_BLOCK=PAGE_BLOCK,
+        BS_BLOCK=triton.next_power_of_2(bs),
+    )

python/sglang/kernels/ops/kvcache/trtllm_mha_page_table.py

@@ -0,0 +1,138 @@
+"""Device-side page-table builder for the trtllm_mha attention backend.
+
+trtllm_mha builds its block (page) table from the global ``req_to_token`` pool.
+Doing it with a host-max PyTorch gather forces a ``seq_lens.max().item()`` D2H
+sync (the CPU must know the page-table width before launching). This kernel
+instead derives the per-request page count from the device-side ``seq_lens``
+tensor, so the build is sync-free: the grid/buffer use the static
+``max_num_pages`` upper bound, while each program self-guards on the real length.
+
+The kernel is MHA-owned (no dependency on the MLA kv-index kernels) and also
+emits the SWA-translated block table in the same pass via the full->SWA lookup
+table, so SWA hybrid models stay sync-free too.
+"""
+
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+# Tokens covered per CTA along the page-block (grid axis-1) dimension.
+# Must be a multiple of page_size (asserted in build_trtllm_mha_page_table).
+_MHA_KV_INDEX_BLOCK_TOKENS = 4096
+# Triton kernels can only read module globals that are tl.constexpr instances.
+_MHA_KV_INDEX_BLOCK_TOKENS_TL = tl.constexpr(_MHA_KV_INDEX_BLOCK_TOKENS)
+
+
+def get_num_mha_kv_index_blocks(num_pages: int, page_size: int) -> int:
+    """Grid axis-1 size: number of page-block CTAs spanning the widest sequence.
+
+    ``num_pages`` is the per-row width of the page-table buffer (the static
+    ``max_num_pages`` upper bound). One CTA handles ``_MHA_KV_INDEX_BLOCK_TOKENS
+    // page_size`` pages.
+    """
+    pages_per_block = _MHA_KV_INDEX_BLOCK_TOKENS // page_size
+    return (num_pages + pages_per_block - 1) // pages_per_block
+
+
+@triton.jit
+def create_trtllm_mha_kv_indices_triton(
+    req_to_token_ptr,  # [max_reqs, max_context_len], int32
+    req_pool_indices_ptr,  # [bs]
+    seq_lens_ptr,  # [bs], per-request KV length in tokens
+    full_to_swa_ptr,  # full->SWA token-slot lookup table, or dummy when not SWA
+    page_table_ptr,  # [bs, num_pages] int32 block ids (output)
+    swa_page_table_ptr,  # [bs, num_pages] int32 SWA block ids (output), or dummy
+    full_to_swa_numel,
+    req_to_token_stride: tl.constexpr,
+    page_table_stride: tl.constexpr,
+    PAGE_SIZE: tl.constexpr,
+    HAS_SWA: tl.constexpr,
+):
+    """Fill ``page_table_ptr`` (and ``swa_page_table_ptr`` when ``HAS_SWA``).
+
+    Program ``(pid_req, pid_blk)`` writes the block ids of request ``pid_req`` for
+    the page-block ``pid_blk``. It reads the KV token slot at each page boundary
+    from ``req_to_token`` and converts it to a block id (``slot // PAGE_SIZE``).
+    Programs past the request's page count are guarded out, so the work (and the
+    DRAM traffic) is bounded by the device-side ``seq_lens`` — no host max needed.
+
+    The SWA lookup assumes valid (``>= 0``) slots, unlike
+    ``translate_loc_from_full_to_swa``'s ``-1`` sentinel handling; page-boundary
+    reads stay within ``seq_len``, so slots are always valid here.
+    """
+    PAGES_PER_BLOCK: tl.constexpr = _MHA_KV_INDEX_BLOCK_TOKENS_TL // PAGE_SIZE
+    pid_req = tl.program_id(0)
+    pid_blk = tl.program_id(1)
+
+    seq_len = tl.load(seq_lens_ptr + pid_req)
+    num_pages = tl.cdiv(seq_len, PAGE_SIZE)
+    num_page_blocks = tl.cdiv(seq_len, _MHA_KV_INDEX_BLOCK_TOKENS_TL)
+    if pid_blk >= num_page_blocks:
+        return
+
+    req_pool_index = tl.load(req_pool_indices_ptr + pid_req)
+    page_idx = tl.arange(0, PAGES_PER_BLOCK) + pid_blk * PAGES_PER_BLOCK
+    token_pos = page_idx.to(tl.int64) * PAGE_SIZE
+    mask = page_idx < num_pages
+
+    slot = tl.load(
+        req_to_token_ptr
+        + req_pool_index.to(tl.int64) * req_to_token_stride
+        + token_pos,
+        mask=mask,
+    )
+    out_off = pid_req * page_table_stride + page_idx
+    tl.store(page_table_ptr + out_off, (slot // PAGE_SIZE).to(tl.int32), mask=mask)
+    if HAS_SWA:
+        swa_index = tl.minimum(tl.maximum(slot, 0), full_to_swa_numel - 1)
+        swa_slot = tl.load(full_to_swa_ptr + swa_index.to(tl.int64), mask=mask)
+        tl.store(
+            swa_page_table_ptr + out_off,
+            (swa_slot // PAGE_SIZE).to(tl.int32),
+            mask=mask,
+        )
+
+
+def build_trtllm_mha_page_table(
+    req_to_token: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    cache_seqlens: torch.Tensor,
+    page_table: torch.Tensor,
+    page_size: int,
+    swa_page_table: Optional[torch.Tensor] = None,
+    full_to_swa: Optional[torch.Tensor] = None,
+) -> None:
+    """Fill ``page_table`` (and ``swa_page_table`` when SWA) on-device, no D2H sync.
+
+    Computes the launch grid from the static page-table width and dispatches
+    ``create_trtllm_mha_kv_indices_triton``. ``page_table`` (and, for SWA models,
+    ``swa_page_table``) are written in place; the caller owns the buffers so the
+    cuda-graph path can reuse its pre-allocated tensors. SWA is enabled iff
+    ``full_to_swa`` is provided, which then also requires ``swa_page_table``.
+    """
+    has_swa = full_to_swa is not None
+    assert has_swa == (
+        swa_page_table is not None
+    ), "full_to_swa and swa_page_table must be provided together"
+    assert (
+        _MHA_KV_INDEX_BLOCK_TOKENS % page_size == 0
+    ), f"page_size={page_size} must divide _MHA_KV_INDEX_BLOCK_TOKENS={_MHA_KV_INDEX_BLOCK_TOKENS}"
+    bs, num_pages = page_table.shape
+    full_to_swa_numel = full_to_swa.numel() if has_swa else 0
+    create_trtllm_mha_kv_indices_triton[
+        (bs, get_num_mha_kv_index_blocks(num_pages, page_size))
+    ](
+        req_to_token,
+        req_pool_indices,
+        cache_seqlens,
+        full_to_swa,
+        page_table,
+        swa_page_table,
+        full_to_swa_numel,
+        req_to_token.stride(0),
+        page_table.stride(0),
+        PAGE_SIZE=page_size,
+        HAS_SWA=has_swa,
+    )