chore: import upstream snapshot with attribution

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

@@ -0,0 +1,43 @@
+"""Attention compute kernels (Triton): decode / extend / prefill / metadata.
+
+The Triton kernels migrated here live in this package
+(``sglang.kernels.ops.attention.<module>``); import them from there. Their
+``KernelSpec`` metadata is registered below for inventory (backend = Triton).
+KV-cache index/write kernels went to the ``kvcache`` group instead.
+"""
+
+from sglang.kernels.registry import register_kernel
+from sglang.kernels.spec import KernelBackend, KernelSpec
+
+# (module, public_fn) migrated from layers/attention/triton_ops + model_executor.
+_TRITON_KERNELS = [
+    ("decode_attention", "decode_attention_fwd"),
+    ("extend_attention", "extend_attention_fwd"),
+    ("extend_attention", "build_unified_kv_indices"),
+    ("prefill_attention", "context_attention_fwd"),
+    ("merge_state", "merge_state_triton"),
+    ("metadata", "get_num_kv_splits_triton"),
+    ("metadata", "prepare_swa_spec_page_table_triton"),
+    ("metadata", "normal_decode_set_metadata"),
+    ("dsa_metadata", "fused_dsa_decode_metadata"),
+    ("dsa_metadata", "fused_dsa_target_verify_metadata"),
+    ("dsa_metadata", "fused_dsa_draft_extend_metadata"),
+    ("rocm_mla_decode_rope", "decode_attention_fwd_grouped_rope"),
+    ("verify_splitkv", "verify_splitkv_fwd"),
+    ("pad", "pad_sequence_with_mask"),
+    ("pad", "pad_draft_extend_query"),
+    ("pad", "unpad_draft_extend_output"),
+    ("pad", "seqlens_expand_triton"),
+    ("position", "compute_position_triton"),
+]
+for _mod, _fn in _TRITON_KERNELS:
+    register_kernel(
+        KernelSpec(
+            op=f"attention.{_fn}",
+            backend=KernelBackend.TRITON,
+            target=f"sglang.kernels.ops.attention.{_mod}:{_fn}",
+        )
+    )
+del _mod, _fn
+
+__all__ = []

python/sglang/kernels/ops/attention/decode_attention.py

@@ -0,0 +1,969 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""
+Memory-efficient attention for decoding.
+It supports page size = 1.
+"""
+
+# Adapted from
+# https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage1.py
+# https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage2.py
+
+import logging
+
+import triton
+import triton.language as tl
+
+from sglang.srt.utils import is_hip
+
+_is_hip = is_hip()
+
+logger = logging.getLogger(__name__)
+
+
+_MIN_BLOCK_KV = 32
+
+
+def _extract_kv_strides(buf, page_size: int):
+    """Extract (slot_stride, head_stride, page_stride, tok_stride) for a
+    KV buffer that may be:
+      - 3-D ``[max_slots, head_num, head_dim]`` (legacy / non-shared) — the
+        contiguous layout most callers use. page/tok strides are synthesized
+        so the kernel's PAGE_SIZE>1 math collapses to ``kv_loc * stride(0)``.
+      - 4-D ``[num_pages, page_size, head_num, head_dim]`` (shared
+        pool). page/tok strides come from stride(0)/stride(1) directly;
+        legacy ``stride_bs`` is set to 0 (unused at PAGE_SIZE>1).
+
+    Returns a 4-tuple of ints suitable for passing as ``stride_buf_*bs``,
+    ``stride_buf_*h``, ``stride_buf_*page``, ``stride_buf_*tok``.
+    """
+    if buf.ndim == 4:
+        # 4-D view ``[num_pages, page_size, head_num, head_dim]``.
+        #   stride(0) = per-PAGE stride (page_bytes/itemsize)
+        #   stride(1) = within-page per-TOKEN stride (k_row/v_row bytes/itemsize)
+        # The PAGE_SIZE>1 kernel branch uses page_stride/tok_stride and does
+        # NOT read slot_stride. slot_stride is consumed ONLY by the
+        # PAGE_SIZE==1 branch (``offs = kv_loc * stride_buf_*bs``), where one
+        # page holds exactly one slot, so the per-slot stride is the per-page
+        # stride — NOT the within-page token stride. Concretely the per-slot
+        # stride is ``page_stride // page_size`` (= entry_bytes/itemsize),
+        # which at ps=1 equals page_stride. Using ``tok_stride`` here (one
+        # layer's k_row) would make the ps=1 read address ``kv_loc * k_row``
+        # instead of ``kv_loc * entry_bytes`` and read the wrong slot.
+        page_stride = buf.stride(0)
+        tok_stride = buf.stride(1)
+        head_stride = buf.stride(2)
+        slot_stride = (
+            page_stride // page_size
+        )  # per-slot stride; == page_stride at ps=1
+        assert buf.shape[1] == page_size, (
+            f"4-D KV buffer's dim-1 must equal page_size; got "
+            f"shape[1]={buf.shape[1]}, page_size={page_size}"
+        )
+    elif buf.ndim == 3:
+        # Legacy 3-D ``[N, head, dim]``. Synthesize page/tok strides such
+        # that ``(kv_loc // ps) * page_stride + (kv_loc % ps) * tok_stride
+        # == kv_loc * slot_stride`` for the page-aware branch — this lets
+        # the same kernel handle non-shared paged-allocator buffers without
+        # any caller adjustment.
+        slot_stride = buf.stride(0)
+        head_stride = buf.stride(1)
+        page_stride = slot_stride * page_size
+        tok_stride = slot_stride
+    else:  # pragma: no cover
+        raise ValueError(f"unexpected KV buffer ndim={buf.ndim}, shape={buf.shape}")
+    return slot_stride, head_stride, page_stride, tok_stride
+
+
+@triton.jit
+def tanh(x):
+    # Tanh is just a scaled sigmoid
+    return 2 * tl.sigmoid(2 * x) - 1
+
+
+@triton.jit
+def _fwd_kernel_stage1(
+    Q,
+    K_Buffer,
+    V_Buffer,
+    sm_scale_withk,
+    kv_indptr,
+    kv_indices,
+    Att_Out,
+    Att_Lse,
+    num_kv_splits,
+    stride_qbs,
+    stride_qh,
+    stride_buf_kbs,
+    stride_buf_kh,
+    stride_buf_vbs,
+    stride_buf_vh,
+    # Page-aware strides (used when PAGE_SIZE > 1). For
+    # PAGE_SIZE == 1 the address math degenerates and these are unused
+    # (Triton specializes the dead branch away at compile time).
+    stride_buf_kpage,
+    stride_buf_ktok,
+    stride_buf_vpage,
+    stride_buf_vtok,
+    stride_mid_ob,
+    stride_mid_oh,
+    stride_mid_os,
+    kv_group_num: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_DV: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    MIN_BLOCK_KV: tl.constexpr,
+    logit_cap: tl.constexpr,
+    Lk: tl.constexpr,
+    Lv: tl.constexpr,
+    xai_temperature_len: tl.constexpr,
+    PAGE_SIZE: tl.constexpr,
+):
+    # int64 to avoid overflow of flat offsets into Mid_O when
+    # batch * num_head * max_kv_splits * head_dim exceeds 2**31.
+    cur_batch = tl.program_id(0).to(tl.int64)
+    cur_head = tl.program_id(1)
+    split_kv_id = tl.program_id(2)
+
+    cur_kv_head = cur_head // kv_group_num
+
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    offs_dv = tl.arange(0, BLOCK_DV)
+    mask_d = offs_d < Lk
+    mask_dv = offs_dv < Lv
+
+    cur_batch_kv_start_idx = tl.load(kv_indptr + cur_batch)
+    cur_batch_seq_len = tl.load(kv_indptr + cur_batch + 1) - cur_batch_kv_start_idx
+    kv_splits = tl.load(num_kv_splits + cur_batch)
+
+    if xai_temperature_len > 0:
+        offs_qidx = cur_batch_seq_len - 1
+        xai_temperature_scale = 1.0 / tl.log2(float(xai_temperature_len))
+        _qtemp = tl.log2(offs_qidx.to(tl.float32)) * xai_temperature_scale
+        xai_temperature_reg = tl.where(offs_qidx > xai_temperature_len, _qtemp, 1.0)
+
+    off_q = cur_batch * stride_qbs + cur_head * stride_qh + offs_d
+
+    kv_len_per_split = (
+        tl.cdiv(tl.cdiv(cur_batch_seq_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV
+    )
+    split_kv_start = kv_len_per_split * split_kv_id
+    split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len)
+
+    e_max = -float("inf")
+    e_sum = 0.0
+    acc = tl.zeros([BLOCK_DV], dtype=tl.float32)
+
+    if split_kv_end > split_kv_start:
+        q = tl.load(Q + off_q, mask=mask_d, other=0.0)
+        for start_n in range(split_kv_start, split_kv_end, BLOCK_N):
+            offs_n = start_n + tl.arange(0, BLOCK_N)
+            kv_loc = tl.load(
+                kv_indices + cur_batch_kv_start_idx + offs_n,
+                mask=offs_n < split_kv_end,
+                other=0,
+            )
+            # Page-aware KV address math. At PAGE_SIZE==1 (legacy
+            # / non-shared / shared-at-ps=1), Triton specializes the
+            # else-branch away and the SASS is byte-identical to today.
+            if PAGE_SIZE == 1:
+                offs_buf_k = (
+                    kv_loc[:, None] * stride_buf_kbs
+                    + cur_kv_head * stride_buf_kh
+                    + offs_d[None, :]
+                )
+            else:
+                page_id = kv_loc // PAGE_SIZE
+                tok_in_p = kv_loc % PAGE_SIZE
+                offs_buf_k = (
+                    page_id[:, None] * stride_buf_kpage
+                    + tok_in_p[:, None] * stride_buf_ktok
+                    + cur_kv_head * stride_buf_kh
+                    + offs_d[None, :]
+                )
+            k = tl.load(
+                K_Buffer + offs_buf_k,
+                mask=(offs_n[:, None] < split_kv_end) & (mask_d[None, :]),
+                other=0.0,
+            )
+            qk = tl.sum(q[None, :] * k, 1)
+            qk *= sm_scale_withk
+
+            if logit_cap > 0:
+                qk = logit_cap * tanh(qk / logit_cap)
+
+            if xai_temperature_len > 0:
+                qk *= xai_temperature_reg
+
+            qk = tl.where(offs_n < split_kv_end, qk, float("-inf"))
+
+            if PAGE_SIZE == 1:
+                offs_buf_v = (
+                    kv_loc[:, None] * stride_buf_vbs
+                    + cur_kv_head * stride_buf_vh
+                    + offs_dv[None, :]
+                )
+            else:
+                offs_buf_v = (
+                    page_id[:, None] * stride_buf_vpage
+                    + tok_in_p[:, None] * stride_buf_vtok
+                    + cur_kv_head * stride_buf_vh
+                    + offs_dv[None, :]
+                )
+            v = tl.load(
+                V_Buffer + offs_buf_v,
+                mask=(offs_n[:, None] < split_kv_end) & (mask_dv[None, :]),
+                other=0.0,
+            )
+
+            n_e_max = tl.maximum(tl.max(qk, 0), e_max)
+            re_scale = tl.exp(e_max - n_e_max)
+            p = tl.exp(qk - n_e_max)
+            acc *= re_scale
+            acc += tl.sum(p[:, None] * v, 0)
+
+            e_sum = e_sum * re_scale + tl.sum(p, 0)
+            e_max = n_e_max
+
+        offs_mid_o = (
+            cur_batch * stride_mid_ob
+            + cur_head * stride_mid_oh
+            + split_kv_id * stride_mid_os
+            + offs_dv
+        )
+
+        tl.store(
+            Att_Out + offs_mid_o,
+            acc / e_sum,
+            mask=(mask_dv),
+        )
+
+        offs_mid_o_1 = (
+            cur_batch * stride_mid_ob
+            + cur_head * stride_mid_oh
+            + split_kv_id * stride_mid_os
+        ) // Lv
+
+        tl.store(
+            Att_Lse + offs_mid_o_1,
+            e_max + tl.log(e_sum),
+        )
+
+
+def _decode_att_m_fwd(
+    q,
+    k_buffer,
+    v_buffer,
+    att_out,
+    att_lse,
+    kv_indptr,
+    kv_indices,
+    num_kv_splits,
+    max_kv_splits,
+    sm_scale_withk,
+    logit_cap,
+    xai_temperature_len=-1,
+    page_size: int = 1,
+):
+    BLOCK = 64
+    # [TODO] work around SGPR limit on MI3xx
+    if _is_hip:
+        BLOCK = 8
+    MAX_KV_SPLITS = max_kv_splits
+    Lk = k_buffer.shape[-1]
+    Lv = v_buffer.shape[-1]
+
+    # head_num lives in the dim immediately before the head_dim. For 3-D
+    # ``[N, head_num, head_dim]`` that's dim 1; for 4-D
+    # ``[num_pages, page_size, head_num, head_dim]`` that's dim 2.
+    kv_head_num = k_buffer.shape[-2]
+
+    batch, head_num = q.shape[0], q.shape[1]
+
+    grid = (batch, head_num, MAX_KV_SPLITS)
+    kv_group_num = q.shape[1] // kv_head_num
+
+    if kv_group_num == 1:
+        num_warps = 4
+    else:
+        num_warps = 2
+        if _is_hip:
+            num_warps = 1
+
+    BLOCK_DMODEL = triton.next_power_of_2(Lk)
+    BLOCK_DV = triton.next_power_of_2(Lv)
+
+    k_slot_stride, k_head_stride, k_page_stride, k_tok_stride = _extract_kv_strides(
+        k_buffer, page_size
+    )
+    v_slot_stride, v_head_stride, v_page_stride, v_tok_stride = _extract_kv_strides(
+        v_buffer, page_size
+    )
+
+    _fwd_kernel_stage1[grid](
+        q,
+        k_buffer,
+        v_buffer,
+        sm_scale_withk,
+        kv_indptr,
+        kv_indices,
+        att_out,
+        att_lse,
+        num_kv_splits,
+        q.stride(0),
+        q.stride(1),
+        k_slot_stride,
+        k_head_stride,
+        v_slot_stride,
+        v_head_stride,
+        k_page_stride,
+        k_tok_stride,
+        v_page_stride,
+        v_tok_stride,
+        att_out.stride(0),
+        att_out.stride(1),
+        att_out.stride(2),
+        kv_group_num=kv_group_num,
+        BLOCK_DMODEL=BLOCK_DMODEL,
+        BLOCK_DV=BLOCK_DV,
+        BLOCK_N=BLOCK,
+        MIN_BLOCK_KV=_MIN_BLOCK_KV,
+        logit_cap=logit_cap,
+        xai_temperature_len=xai_temperature_len,
+        num_warps=num_warps,
+        num_stages=2,
+        Lk=Lk,
+        Lv=Lv,
+        PAGE_SIZE=page_size,
+    )
+
+
+@triton.jit
+def _fwd_grouped_kernel_stage1(
+    Q,
+    K_Buffer,
+    V_Buffer,
+    sm_scale_withk,
+    kv_indptr,
+    kv_indices,
+    Att_Out,
+    Att_Lse,
+    num_kv_splits,
+    stride_qbs,
+    stride_qh,
+    stride_buf_kbs,
+    stride_buf_kh,
+    stride_buf_vbs,
+    stride_buf_vh,
+    # Page-aware strides (used when PAGE_SIZE > 1).
+    stride_buf_kpage,
+    stride_buf_ktok,
+    stride_buf_vpage,
+    stride_buf_vtok,
+    stride_mid_ob,
+    stride_mid_oh,
+    stride_mid_os,
+    kv_group_num: tl.constexpr,
+    q_head_num: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_DPE: tl.constexpr,
+    BLOCK_DV: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    MIN_BLOCK_KV: tl.constexpr,
+    logit_cap: tl.constexpr,
+    xai_temperature_len: tl.constexpr,
+    Lk: tl.constexpr,
+    Lv: tl.constexpr,
+    HAS_MLA: tl.constexpr = False,
+    USE_PDL: tl.constexpr = False,
+    PAGE_SIZE: tl.constexpr = 1,
+):
+    # int64 to avoid overflow of flat offsets into Mid_O when
+    # batch * num_head * max_kv_splits * head_dim exceeds 2**31.
+    cur_batch = tl.program_id(0).to(tl.int64)
+    cur_head_id = tl.program_id(1)
+    cur_kv_head = cur_head_id // tl.cdiv(kv_group_num, BLOCK_H)
+    split_kv_id = tl.program_id(2)
+
+    if BLOCK_H < kv_group_num:
+        VALID_BLOCK_H: tl.constexpr = BLOCK_H
+    else:
+        VALID_BLOCK_H: tl.constexpr = kv_group_num
+    cur_head = cur_head_id * VALID_BLOCK_H + tl.arange(0, BLOCK_H)
+    mask_h = cur_head < (cur_head_id + 1) * VALID_BLOCK_H
+    mask_h = mask_h & (cur_head < q_head_num)
+
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    offs_dv = tl.arange(0, BLOCK_DV)
+    mask_d = offs_d < Lk
+    mask_dv = offs_dv < Lv
+
+    cur_batch_kv_start_idx = tl.load(kv_indptr + cur_batch)
+    cur_batch_seq_len = tl.load(kv_indptr + cur_batch + 1) - cur_batch_kv_start_idx
+    kv_splits = tl.load(num_kv_splits + cur_batch)
+
+    if xai_temperature_len > 0:
+        offs_qidx = cur_batch_seq_len - 1
+        xai_temperature_scale = 1.0 / tl.log2(float(xai_temperature_len))
+        _qtemp = tl.log2(offs_qidx.to(tl.float32)) * xai_temperature_scale
+        xai_temperature_reg = tl.where(offs_qidx > xai_temperature_len, _qtemp, 1.0)
+
+    offs_q = cur_batch * stride_qbs + cur_head[:, None] * stride_qh + offs_d[None, :]
+
+    if BLOCK_DPE > 0:
+        offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE)
+        mask_dpe = offs_dpe < Lk
+        off_qpe = (
+            cur_batch * stride_qbs + cur_head[:, None] * stride_qh + offs_dpe[None, :]
+        )
+
+    kv_len_per_split = (
+        tl.cdiv(tl.cdiv(cur_batch_seq_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV
+    )
+    split_kv_start = kv_len_per_split * split_kv_id
+    split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len)
+
+    e_max = tl.zeros([BLOCK_H], dtype=tl.float32) - float("inf")
+    e_sum = tl.zeros([BLOCK_H], dtype=tl.float32)
+    acc = tl.zeros([BLOCK_H, BLOCK_DV], dtype=tl.float32)
+
+    # Hoist loop-invariant base offsets
+    base_offs_k = cur_kv_head * stride_buf_kh + offs_d[:, None]
+    if BLOCK_DPE > 0:
+        base_offs_kpe = cur_kv_head * stride_buf_kh + offs_dpe[:, None]
+    if not HAS_MLA:
+        base_offs_v = cur_kv_head * stride_buf_vh + offs_dv[None, :]
+
+    if split_kv_end > split_kv_start:
+        q = tl.load(Q + offs_q, mask=(mask_h[:, None]) & (mask_d[None, :]), other=0.0)
+        q_k = q.to(K_Buffer.dtype.element_ty)
+        if BLOCK_DPE > 0:
+            qpe = tl.load(
+                Q + off_qpe, mask=(mask_h[:, None]) & (mask_dpe[None, :]), other=0.0
+            )
+        for start_n in tl.range(split_kv_start, split_kv_end, BLOCK_N):
+            offs_n = start_n + tl.arange(0, BLOCK_N)
+            kv_loc = tl.load(
+                kv_indices + cur_batch_kv_start_idx + offs_n,
+                mask=offs_n < split_kv_end,
+                other=0,
+            )
+            # Page-aware KV address math (see _fwd_kernel_stage1).
+            if PAGE_SIZE == 1:
+                offs_buf_k = kv_loc[None, :] * stride_buf_kbs + base_offs_k
+            else:
+                page_id = kv_loc // PAGE_SIZE
+                tok_in_p = kv_loc % PAGE_SIZE
+                offs_buf_k = (
+                    page_id[None, :] * stride_buf_kpage
+                    + tok_in_p[None, :] * stride_buf_ktok
+                    + base_offs_k
+                )
+            k = tl.load(
+                K_Buffer + offs_buf_k,
+                mask=(offs_n[None, :] < split_kv_end) & (mask_d[:, None]),
+                other=0.0,
+            )
+            qk = tl.dot(q_k, k)
+            if BLOCK_DPE > 0:
+                if PAGE_SIZE == 1:
+                    offs_buf_kpe = kv_loc[None, :] * stride_buf_kbs + base_offs_kpe
+                else:
+                    offs_buf_kpe = (
+                        page_id[None, :] * stride_buf_kpage
+                        + tok_in_p[None, :] * stride_buf_ktok
+                        + base_offs_kpe
+                    )
+                kpe = tl.load(
+                    K_Buffer + offs_buf_kpe,
+                    mask=(offs_n[None, :] < split_kv_end) & (mask_dpe[:, None]),
+                    other=0.0,
+                )
+                qk += tl.dot(qpe, kpe.to(qpe.dtype))
+            qk *= sm_scale_withk
+
+            if logit_cap > 0:
+                qk = logit_cap * tanh(qk / logit_cap)
+
+            if xai_temperature_len > 0:
+                qk *= xai_temperature_reg[:, None]
+
+            qk = tl.where(
+                mask_h[:, None] & (offs_n[None, :] < split_kv_end), qk, float("-inf")
+            )
+            if HAS_MLA:
+                v = tl.trans(k)
+            else:
+                if PAGE_SIZE == 1:
+                    offs_buf_v = kv_loc[:, None] * stride_buf_vbs + base_offs_v
+                else:
+                    offs_buf_v = (
+                        page_id[:, None] * stride_buf_vpage
+                        + tok_in_p[:, None] * stride_buf_vtok
+                        + base_offs_v
+                    )
+                v = tl.load(
+                    V_Buffer + offs_buf_v,
+                    mask=(offs_n[:, None] < split_kv_end) & (mask_dv[None, :]),
+                    other=0.0,
+                )
+
+            n_e_max = tl.maximum(tl.max(qk, 1), e_max)
+            re_scale = tl.exp(e_max - n_e_max)
+            p = tl.exp(qk - n_e_max[:, None])
+            acc *= re_scale[:, None]
+            acc += tl.dot(p.to(v.dtype), v)
+
+            e_sum = e_sum * re_scale + tl.sum(p, 1)
+            e_max = n_e_max
+
+        offs_mid_o = (
+            cur_batch * stride_mid_ob
+            + cur_head[:, None] * stride_mid_oh
+            + split_kv_id * stride_mid_os
+            + offs_dv[None, :]
+        )
+
+        tl.store(
+            Att_Out + offs_mid_o,
+            acc / e_sum[:, None],
+            mask=(mask_h[:, None]) & (mask_dv[None, :]),
+        )
+
+        offs_mid_o_1 = (
+            cur_batch * stride_mid_ob
+            + cur_head * stride_mid_oh
+            + split_kv_id * stride_mid_os
+        ) // Lv
+
+        tl.store(
+            Att_Lse + offs_mid_o_1,
+            e_max + tl.log(e_sum),
+            mask=mask_h,
+        )
+
+    if USE_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
+
+def _decode_grouped_att_m_fwd(
+    q,
+    k_buffer,
+    v_buffer,
+    att_out,
+    att_lse,
+    kv_indptr,
+    kv_indices,
+    num_kv_splits,
+    max_kv_splits,
+    sm_scale_withk,
+    logit_cap,
+    xai_temperature_len=-1,
+    has_mla=False,
+    use_pdl=False,
+    page_size: int = 1,
+):
+    BLOCK = 32
+    Lk = k_buffer.shape[-1]
+    Lv = v_buffer.shape[-1]
+
+    # [TODO] work around shmem limit on MI3xx
+    if _is_hip and Lk >= 576:
+        BLOCK = 16
+
+    if Lk == 576:
+        BLOCK_DMODEL = 512
+        BLOCK_DPE = 64
+    elif Lk == 288:
+        BLOCK_DMODEL = 256
+        BLOCK_DPE = 32
+    else:
+        BLOCK_DMODEL = triton.next_power_of_2(Lk)
+        BLOCK_DPE = 0
+    BLOCK_DV = triton.next_power_of_2(Lv)
+
+    # 4-D view exposes head_num at dim 2; legacy 3-D exposes
+    # it at dim 1.
+    kv_head_num = k_buffer.shape[-2]
+    batch, head_num = q.shape[0], q.shape[1]
+    kv_group_num = q.shape[1] // kv_head_num
+
+    BLOCK_H = 16
+    MAX_KV_SPLITS = max_kv_splits
+    grid = (
+        batch,
+        triton.cdiv(head_num, min(BLOCK_H, kv_group_num)),
+        MAX_KV_SPLITS,
+    )
+
+    extra_kargs = {}
+    num_stages = 2
+    if _is_hip:
+        # https://rocm.docs.amd.com/en/docs-6.2.0/how-to/llm-fine-tuning-optimization/optimizing-triton-kernel.html
+        # https://github.com/triton-lang/triton/blob/main/third_party/amd/backend/compiler.py
+        extra_kargs = {"waves_per_eu": 1, "matrix_instr_nonkdim": 16, "kpack": 2}
+        num_stages = 1
+
+    k_slot_stride, k_head_stride, k_page_stride, k_tok_stride = _extract_kv_strides(
+        k_buffer, page_size
+    )
+    v_slot_stride, v_head_stride, v_page_stride, v_tok_stride = _extract_kv_strides(
+        v_buffer, page_size
+    )
+
+    _fwd_grouped_kernel_stage1[grid](
+        q,
+        k_buffer,
+        v_buffer,
+        sm_scale_withk,
+        kv_indptr,
+        kv_indices,
+        att_out,
+        att_lse,
+        num_kv_splits,
+        q.stride(0),
+        q.stride(1),
+        k_slot_stride,
+        k_head_stride,
+        v_slot_stride,
+        v_head_stride,
+        k_page_stride,
+        k_tok_stride,
+        v_page_stride,
+        v_tok_stride,
+        att_out.stride(0),
+        att_out.stride(1),
+        att_out.stride(2),
+        kv_group_num=kv_group_num,
+        q_head_num=head_num,
+        BLOCK_DMODEL=BLOCK_DMODEL,
+        BLOCK_DPE=BLOCK_DPE,
+        BLOCK_DV=BLOCK_DV,
+        BLOCK_N=BLOCK,
+        BLOCK_H=BLOCK_H,
+        MIN_BLOCK_KV=_MIN_BLOCK_KV,
+        logit_cap=logit_cap,
+        xai_temperature_len=xai_temperature_len,
+        num_warps=4,
+        num_stages=num_stages,
+        Lk=Lk,
+        Lv=Lv,
+        HAS_MLA=has_mla,
+        USE_PDL=use_pdl,
+        PAGE_SIZE=page_size,
+        **extra_kargs,
+    )
+
+
+@triton.jit
+def _fwd_kernel_stage2(
+    Mid_O,
+    Mid_O_1,
+    O,
+    v_scale,
+    kv_indptr,
+    num_kv_splits,
+    sink_ptr,
+    stride_mid_ob,
+    stride_mid_oh,
+    stride_mid_os,
+    stride_obs,
+    stride_oh,
+    MAX_KV_SPLITS: tl.constexpr,
+    MIN_BLOCK_KV: tl.constexpr,
+    BLOCK_DV: tl.constexpr,
+    Lv: tl.constexpr,
+    HAS_SINK: tl.constexpr,
+    USE_PDL: tl.constexpr = False,
+):
+    # int64 to avoid overflow of flat offsets into Mid_O when
+    # batch * num_head * max_kv_splits * head_dim exceeds 2**31.
+    cur_batch = tl.program_id(0).to(tl.int64)
+    cur_head = tl.program_id(1)
+
+    if USE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    cur_batch_seq_len = tl.load(kv_indptr + cur_batch + 1) - tl.load(
+        kv_indptr + cur_batch
+    )
+    kv_splits = tl.load(num_kv_splits + cur_batch)
+
+    offs_d = tl.arange(0, BLOCK_DV)
+    mask_d = offs_d < Lv
+
+    e_sum = 0.0
+    e_max = -float("inf")
+    acc = tl.zeros([BLOCK_DV], dtype=tl.float32)
+
+    offs_v = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + offs_d
+    offs_logic = (cur_batch * stride_mid_ob + cur_head * stride_mid_oh) // Lv
+    kv_len_per_split = (
+        tl.cdiv(tl.cdiv(cur_batch_seq_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV
+    )
+
+    for split_kv_id in tl.range(0, MAX_KV_SPLITS, num_stages=2):
+        split_kv_start = kv_len_per_split * split_kv_id
+        split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len)
+
+        if split_kv_end > split_kv_start:
+            tv = tl.load(
+                Mid_O + offs_v + split_kv_id * stride_mid_os, mask=mask_d, other=0.0
+            )
+            tlogic = tl.load(Mid_O_1 + offs_logic + split_kv_id * stride_mid_os // Lv)
+            n_e_max = tl.maximum(tlogic, e_max)
+
+            old_scale = tl.exp(e_max - n_e_max)
+            acc *= old_scale
+            exp_logic = tl.exp(tlogic - n_e_max)
+            acc += exp_logic * tv
+
+            e_sum = e_sum * old_scale + exp_logic
+            e_max = n_e_max
+
+    if HAS_SINK:
+        cur_sink = tl.load(sink_ptr + cur_head)
+        e_sum += tl.exp(cur_sink - e_max)
+
+    tl.store(
+        O + cur_batch * stride_obs + cur_head * stride_oh + offs_d,
+        acc / e_sum * v_scale,
+        mask=mask_d,
+    )
+
+
+def _decode_softmax_reducev_fwd(
+    logits,
+    lse,
+    q,
+    o,
+    v_scale,
+    v_buffer,
+    kv_indptr,
+    num_kv_splits,
+    max_kv_splits,
+    sinks=None,
+    use_pdl=False,
+):
+    batch, head_num = q.shape[0], q.shape[1]
+    Lv = v_buffer.shape[-1]
+    BLOCK_DV = triton.next_power_of_2(Lv)
+
+    MAX_KV_SPLITS = max_kv_splits
+    HAS_SINK = sinks is not None
+
+    extra_kargs = {}
+    if _is_hip:
+        # https://rocm.docs.amd.com/en/docs-6.2.0/how-to/llm-fine-tuning-optimization/optimizing-triton-kernel.html
+        # https://github.com/triton-lang/triton/blob/main/third_party/amd/backend/compiler.py
+        extra_kargs = {"waves_per_eu": 4, "matrix_instr_nonkdim": 16, "kpack": 2}
+
+    grid = (batch, head_num)
+    _fwd_kernel_stage2[grid](
+        logits,
+        lse,
+        o,
+        v_scale,
+        kv_indptr,
+        num_kv_splits,
+        sinks,
+        logits.stride(0),
+        logits.stride(1),
+        logits.stride(2),
+        o.stride(0),
+        o.stride(1),
+        MAX_KV_SPLITS=MAX_KV_SPLITS,
+        MIN_BLOCK_KV=_MIN_BLOCK_KV,
+        BLOCK_DV=BLOCK_DV,
+        Lv=Lv,
+        HAS_SINK=HAS_SINK,
+        USE_PDL=use_pdl,
+        num_warps=4,
+        num_stages=2,
+        **({"launch_pdl": True} if use_pdl else {}),
+        **extra_kargs,
+    )
+
+
+def decode_attention_fwd_normal(
+    q,
+    k_buffer,
+    v_buffer,
+    o,
+    kv_indptr,
+    kv_indices,
+    attn_logits,
+    attn_lse,
+    num_kv_splits,
+    max_kv_splits,
+    sm_scale_withk,
+    v_scale,
+    logit_cap=0.0,
+    sinks=None,
+    xai_temperature_len=-1,
+    page_size: int = 1,
+):
+    _decode_att_m_fwd(
+        q,
+        k_buffer,
+        v_buffer,
+        attn_logits,
+        attn_lse,
+        kv_indptr,
+        kv_indices,
+        num_kv_splits,
+        max_kv_splits,
+        sm_scale_withk,
+        logit_cap,
+        xai_temperature_len,
+        page_size=page_size,
+    )
+    _decode_softmax_reducev_fwd(
+        attn_logits,
+        attn_lse,
+        q,
+        o,
+        v_scale,
+        v_buffer,
+        kv_indptr,
+        num_kv_splits,
+        max_kv_splits,
+        sinks,
+    )
+
+
+def decode_attention_fwd_grouped(
+    q,
+    k_buffer,
+    v_buffer,
+    o,
+    kv_indptr,
+    kv_indices,
+    attn_logits,
+    attn_lse,
+    num_kv_splits,
+    max_kv_splits,
+    sm_scale_withk,
+    v_scale,
+    logit_cap=0.0,
+    sinks=None,
+    xai_temperature_len=-1,
+    has_mla=False,
+    use_pdl=False,
+    page_size: int = 1,
+):
+    _decode_grouped_att_m_fwd(
+        q,
+        k_buffer,
+        v_buffer,
+        attn_logits,
+        attn_lse,
+        kv_indptr,
+        kv_indices,
+        num_kv_splits,
+        max_kv_splits,
+        sm_scale_withk,
+        logit_cap,
+        xai_temperature_len,
+        has_mla=has_mla,
+        use_pdl=use_pdl,
+        page_size=page_size,
+    )
+    _decode_softmax_reducev_fwd(
+        attn_logits,
+        attn_lse,
+        q,
+        o,
+        v_scale,
+        v_buffer,
+        kv_indptr,
+        num_kv_splits,
+        max_kv_splits,
+        sinks,
+        use_pdl=use_pdl,
+    )
+
+
+def decode_attention_fwd(
+    q,
+    k_buffer,
+    v_buffer,
+    o,
+    kv_indptr,
+    kv_indices,
+    attn_logits,
+    attn_lse,
+    num_kv_splits,
+    max_kv_splits,
+    sm_scale,
+    k_scale,
+    v_scale,
+    logit_cap=0.0,
+    sinks=None,
+    xai_temperature_len=-1,
+    has_mla=False,
+    use_pdl=False,
+    page_size: int = 1,
+):
+    assert max_kv_splits == attn_logits.shape[2]
+    assert q.shape[0] <= kv_indptr.shape[0] - 1
+    assert q.shape[0] <= attn_logits.shape[0]
+
+    # head_num lives at dim 1 (3-D) or dim 2 (4-D shared view).
+    kv_head_num = v_buffer.shape[-2]
+    kv_group_num = q.shape[1] // kv_head_num
+
+    if kv_group_num == 1:
+        # MHA
+        decode_attention_fwd_normal(
+            q,
+            k_buffer,
+            v_buffer,
+            o,
+            kv_indptr,
+            kv_indices,
+            attn_logits,
+            attn_lse,
+            num_kv_splits,
+            max_kv_splits,
+            sm_scale * k_scale,
+            v_scale,
+            logit_cap=logit_cap,
+            sinks=sinks,
+            xai_temperature_len=xai_temperature_len,
+            page_size=page_size,
+        )
+    else:
+        # GQA/MQA/MLA
+        decode_attention_fwd_grouped(
+            q,
+            k_buffer,
+            v_buffer,
+            o,
+            kv_indptr,
+            kv_indices,
+            attn_logits,
+            attn_lse,
+            num_kv_splits,
+            max_kv_splits,
+            sm_scale * k_scale,
+            v_scale,
+            logit_cap=logit_cap,
+            sinks=sinks,
+            xai_temperature_len=xai_temperature_len,
+            has_mla=has_mla,
+            use_pdl=use_pdl,
+            page_size=page_size,
+        )

python/sglang/kernels/ops/attention/dsa_metadata.py

@@ -0,0 +1,682 @@
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit(
+    do_not_specialize=[
+        "page_table_stride_0",
+        "real_page_table_stride_0",
+        "max_len",
+    ]
+)
+def _fused_dsa_decode_metadata_kernel(
+    seq_lens,
+    req_pool_indices,
+    req_to_token,
+    cache_seqlens,
+    cu_seqlens_k,
+    page_table_1,
+    dsa_cache_seqlens,
+    dsa_cu_seqlens_k,
+    real_page_table,
+    seq_lens_stride: tl.constexpr,
+    req_pool_indices_stride: tl.constexpr,
+    req_to_token_stride_0: tl.constexpr,
+    req_to_token_stride_1: tl.constexpr,
+    page_table_stride_0,
+    page_table_stride_1: tl.constexpr,
+    real_page_table_stride_0,
+    real_page_table_stride_1: tl.constexpr,
+    bs: tl.constexpr,
+    max_len,
+    dsa_index_topk: tl.constexpr,
+    real_page_size: tl.constexpr,
+    HAS_REAL_PAGE_TABLE: tl.constexpr,
+    HAS_PAGE_TABLE_1: tl.constexpr,
+    BLOCK_BS: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    pid = tl.program_id(0)
+
+    if pid == 0:
+        offs_b = tl.arange(0, BLOCK_BS)
+        mask_b = offs_b < bs
+        seq = tl.load(seq_lens + offs_b * seq_lens_stride, mask=mask_b, other=0)
+        seq_i32 = seq.to(tl.int32)
+        dsa_seq = tl.minimum(seq_i32, dsa_index_topk)
+
+        cu = tl.cumsum(seq_i32, 0)
+        dsa_cu = tl.cumsum(dsa_seq, 0)
+
+        tl.store(cache_seqlens + offs_b, seq_i32, mask=mask_b)
+        tl.store(cu_seqlens_k, tl.full((), 0, tl.int32))
+        tl.store(cu_seqlens_k + 1 + offs_b, cu, mask=mask_b)
+        tl.store(dsa_cache_seqlens + offs_b, dsa_seq, mask=mask_b)
+        tl.store(dsa_cu_seqlens_k, tl.full((), 0, tl.int32))
+        tl.store(dsa_cu_seqlens_k + 1 + offs_b, dsa_cu, mask=mask_b)
+        return
+
+    num_col_blocks = tl.cdiv(max_len, BLOCK_N)
+    page_pid = pid - 1
+    row = page_pid // num_col_blocks
+    col_block = page_pid - row * num_col_blocks
+    offs_n = col_block * BLOCK_N + tl.arange(0, BLOCK_N)
+    mask = (row < bs) & (offs_n < max_len)
+
+    req_idx = tl.load(
+        req_pool_indices + row * req_pool_indices_stride,
+        mask=row < bs,
+        other=0,
+    )
+    vals = tl.load(
+        req_to_token + req_idx * req_to_token_stride_0 + offs_n * req_to_token_stride_1,
+        mask=mask,
+        other=0,
+    ).to(tl.int32)
+    # Write the wide page_size=1 table only when the caller provides it; the
+    # fused decode CUDA graph drops it and consumes real_page_table alone.
+    if HAS_PAGE_TABLE_1:
+        tl.store(
+            page_table_1 + row * page_table_stride_0 + offs_n * page_table_stride_1,
+            vals,
+            mask=mask,
+        )
+
+    if HAS_REAL_PAGE_TABLE:
+        real_mask = mask & ((offs_n % real_page_size) == 0)
+        real_cols = offs_n // real_page_size
+        tl.store(
+            real_page_table
+            + row * real_page_table_stride_0
+            + real_cols * real_page_table_stride_1,
+            vals // real_page_size,
+            mask=real_mask,
+        )
+
+
+def fused_dsa_decode_metadata(
+    seq_lens: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    cache_seqlens: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    page_table_1: Optional[torch.Tensor],
+    dsa_cache_seqlens: torch.Tensor,
+    dsa_cu_seqlens_k: torch.Tensor,
+    real_page_table: torch.Tensor,
+    bs: int,
+    max_len: int,
+    dsa_index_topk: int,
+    real_page_size: int,
+) -> None:
+    """Fill decode-graph DSA metadata (seqlens + page tables) from req_to_token.
+
+    ``page_table_1`` (the wide page_size=1 table) is optional: pass ``None`` to
+    skip materializing it and write only the compact ``real_page_table``
+    (page_size=``real_page_size``). This is used by the fused decode CUDA graph,
+    where the wide table is never read (attention uses topk_indices, the indexer
+    uses real_page_table); ``real_page_size`` must be >1 in that case. When a
+    tensor is passed, behavior is unchanged (both tables are written).
+    """
+    assert seq_lens.is_cuda
+    assert req_pool_indices.is_cuda
+    assert req_to_token.is_cuda
+    assert cache_seqlens.is_cuda
+    assert cu_seqlens_k.is_cuda
+    assert dsa_cache_seqlens.is_cuda
+    assert dsa_cu_seqlens_k.is_cuda
+
+    if bs == 0:
+        cu_seqlens_k[:1].zero_()
+        dsa_cu_seqlens_k[:1].zero_()
+        return
+
+    has_real_page_table = real_page_size > 1
+    if has_real_page_table:
+        assert real_page_table is not None
+        assert real_page_table.is_cuda
+    else:
+        # page_size==1: real IS page_table_1, so page_table_1 must be present.
+        assert page_table_1 is not None
+        real_page_table = page_table_1
+
+    # page_table_1 (the wide page_size=1 table) may be dropped for the fused
+    # decode CUDA graph; the kernel then writes only real_page_table.
+    has_page_table_1 = page_table_1 is not None
+    if not has_page_table_1:
+        assert has_real_page_table
+        page_table_1 = real_page_table  # dummy pointer for stride args
+    else:
+        assert page_table_1.is_cuda
+
+    block_bs = triton.next_power_of_2(bs)
+    block_n = 128
+    num_col_blocks = triton.cdiv(max_len, block_n)
+    grid = (1 + bs * num_col_blocks,)
+
+    _fused_dsa_decode_metadata_kernel[grid](
+        seq_lens,
+        req_pool_indices,
+        req_to_token,
+        cache_seqlens,
+        cu_seqlens_k,
+        page_table_1,
+        dsa_cache_seqlens,
+        dsa_cu_seqlens_k,
+        real_page_table,
+        seq_lens.stride(0),
+        req_pool_indices.stride(0),
+        req_to_token.stride(0),
+        req_to_token.stride(1),
+        page_table_1.stride(0),
+        page_table_1.stride(1),
+        real_page_table.stride(0) if has_real_page_table else 0,
+        real_page_table.stride(1) if has_real_page_table else 0,
+        bs,
+        max_len,
+        dsa_index_topk,
+        real_page_size,
+        has_real_page_table,
+        has_page_table_1,
+        BLOCK_BS=block_bs,
+        BLOCK_N=block_n,
+    )
+
+
+@triton.jit(
+    do_not_specialize=[
+        "page_table_stride_0",
+        "real_page_table_stride_0",
+        "max_seqlen_k",
+    ]
+)
+def _fused_dsa_target_verify_metadata_kernel(
+    seq_lens,
+    req_pool_indices,
+    req_to_token,
+    cache_seqlens,
+    cu_seqlens_k,
+    page_table_1,
+    seqlens_expanded,
+    dsa_cache_seqlens,
+    dsa_cu_seqlens_k,
+    real_page_table,
+    paged_mqa_ctx_lens_2d,
+    seq_lens_stride: tl.constexpr,
+    req_pool_indices_stride: tl.constexpr,
+    req_to_token_stride_0: tl.constexpr,
+    req_to_token_stride_1: tl.constexpr,
+    page_table_stride_0,
+    page_table_stride_1: tl.constexpr,
+    real_page_table_stride_0,
+    real_page_table_stride_1: tl.constexpr,
+    paged_mqa_ctx_lens_stride_0: tl.constexpr,
+    paged_mqa_ctx_lens_stride_1: tl.constexpr,
+    bs: tl.constexpr,
+    max_seqlen_k,
+    dsa_index_topk: tl.constexpr,
+    real_page_size: tl.constexpr,
+    next_n: tl.constexpr,
+    HAS_REAL_PAGE_TABLE: tl.constexpr,
+    HAS_PAGED_MQA_CTX_LENS: tl.constexpr,
+    HAS_PAGE_TABLE_1: tl.constexpr,
+    BLOCK_BS: tl.constexpr,
+    BLOCK_EXPANDED: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    expanded_size: tl.constexpr = bs * next_n
+
+    if pid == 0:
+        offs_b = tl.arange(0, BLOCK_BS)
+        mask_b = offs_b < bs
+        seq = tl.load(seq_lens + offs_b * seq_lens_stride, mask=mask_b, other=0)
+        cache_seq = seq.to(tl.int32) + next_n
+        cu = tl.cumsum(cache_seq, 0)
+
+        tl.store(cache_seqlens + offs_b, cache_seq, mask=mask_b)
+        tl.store(cu_seqlens_k, tl.full((), 0, tl.int32))
+        tl.store(cu_seqlens_k + 1 + offs_b, cu, mask=mask_b)
+
+        offs_e = tl.arange(0, BLOCK_EXPANDED)
+        mask_e = offs_e < expanded_size
+        req_row = offs_e // next_n
+        draft_off = offs_e - req_row * next_n
+        base_seq = tl.load(
+            seq_lens + req_row * seq_lens_stride,
+            mask=mask_e,
+            other=0,
+        ).to(tl.int32)
+        expanded_seq = base_seq + draft_off + 1
+        expanded_seq = tl.where(mask_e, expanded_seq, 0)
+        dsa_seq = tl.minimum(expanded_seq, dsa_index_topk)
+        dsa_cu = tl.cumsum(dsa_seq, 0)
+
+        tl.store(seqlens_expanded + offs_e, expanded_seq, mask=mask_e)
+        tl.store(dsa_cache_seqlens + offs_e, dsa_seq, mask=mask_e)
+        tl.store(dsa_cu_seqlens_k, tl.full((), 0, tl.int32))
+        tl.store(dsa_cu_seqlens_k + 1 + offs_e, dsa_cu, mask=mask_e)
+
+        if HAS_PAGED_MQA_CTX_LENS:
+            tl.store(
+                paged_mqa_ctx_lens_2d
+                + req_row * paged_mqa_ctx_lens_stride_0
+                + draft_off * paged_mqa_ctx_lens_stride_1,
+                base_seq + next_n,
+                mask=mask_e,
+            )
+        return
+
+    num_col_blocks = tl.cdiv(max_seqlen_k, BLOCK_N)
+    page_pid = pid - 1
+    out_row = page_pid // num_col_blocks
+    col_block = page_pid - out_row * num_col_blocks
+    offs_n = col_block * BLOCK_N + tl.arange(0, BLOCK_N)
+    mask = (out_row < expanded_size) & (offs_n < max_seqlen_k)
+
+    req_row = out_row // next_n
+    req_idx = tl.load(
+        req_pool_indices + req_row * req_pool_indices_stride,
+        mask=out_row < expanded_size,
+        other=0,
+    )
+    vals = tl.load(
+        req_to_token + req_idx * req_to_token_stride_0 + offs_n * req_to_token_stride_1,
+        mask=mask,
+        other=0,
+    ).to(tl.int32)
+    # Write the wide page_size=1 table only when the caller provides it (see
+    # fused_dsa_decode_metadata for the optional-page_table_1 contract).
+    if HAS_PAGE_TABLE_1:
+        tl.store(
+            page_table_1 + out_row * page_table_stride_0 + offs_n * page_table_stride_1,
+            vals,
+            mask=mask,
+        )
+
+    if HAS_REAL_PAGE_TABLE:
+        real_mask = mask & ((offs_n % real_page_size) == 0)
+        real_cols = offs_n // real_page_size
+        tl.store(
+            real_page_table
+            + out_row * real_page_table_stride_0
+            + real_cols * real_page_table_stride_1,
+            vals // real_page_size,
+            mask=real_mask,
+        )
+
+
+def fused_dsa_target_verify_metadata(
+    seq_lens: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    cache_seqlens: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    page_table_1: Optional[torch.Tensor],
+    seqlens_expanded: torch.Tensor,
+    dsa_cache_seqlens: torch.Tensor,
+    dsa_cu_seqlens_k: torch.Tensor,
+    real_page_table: torch.Tensor,
+    bs: int,
+    max_seqlen_k: int,
+    dsa_index_topk: int,
+    real_page_size: int,
+    next_n: int,
+    paged_mqa_ctx_lens_2d: torch.Tensor = None,
+) -> None:
+    assert seq_lens.is_cuda
+    assert req_pool_indices.is_cuda
+    assert req_to_token.is_cuda
+    assert cache_seqlens.is_cuda
+    assert cu_seqlens_k.is_cuda
+    assert seqlens_expanded.is_cuda
+    assert dsa_cache_seqlens.is_cuda
+    assert dsa_cu_seqlens_k.is_cuda
+
+    if bs == 0:
+        cu_seqlens_k[:1].zero_()
+        dsa_cu_seqlens_k[:1].zero_()
+        return
+    assert next_n > 0
+
+    has_real_page_table = real_page_size > 1
+    if has_real_page_table:
+        assert real_page_table is not None
+        assert real_page_table.is_cuda
+    else:
+        assert page_table_1 is not None
+        real_page_table = page_table_1
+
+    # page_table_1 (the wide page_size=1 table) may be dropped for the fused
+    # decode CUDA graph; the kernel then writes only real_page_table.
+    has_page_table_1 = page_table_1 is not None
+    if not has_page_table_1:
+        assert has_real_page_table
+        page_table_1 = real_page_table  # dummy pointer for stride args
+    else:
+        assert page_table_1.is_cuda
+
+    has_paged_mqa_ctx_lens = paged_mqa_ctx_lens_2d is not None
+    if has_paged_mqa_ctx_lens:
+        assert paged_mqa_ctx_lens_2d.is_cuda
+        assert paged_mqa_ctx_lens_2d.dtype == torch.int32
+        assert paged_mqa_ctx_lens_2d.dim() == 2
+        assert paged_mqa_ctx_lens_2d.size(0) == bs
+        assert paged_mqa_ctx_lens_2d.size(1) == next_n
+    else:
+        paged_mqa_ctx_lens_2d = page_table_1
+
+    expanded_size = bs * next_n
+    block_bs = triton.next_power_of_2(bs)
+    block_expanded = triton.next_power_of_2(expanded_size)
+    block_n = 128
+    num_col_blocks = triton.cdiv(max_seqlen_k, block_n)
+    grid = (1 + expanded_size * num_col_blocks,)
+
+    _fused_dsa_target_verify_metadata_kernel[grid](
+        seq_lens,
+        req_pool_indices,
+        req_to_token,
+        cache_seqlens,
+        cu_seqlens_k,
+        page_table_1,
+        seqlens_expanded,
+        dsa_cache_seqlens,
+        dsa_cu_seqlens_k,
+        real_page_table,
+        paged_mqa_ctx_lens_2d,
+        seq_lens.stride(0),
+        req_pool_indices.stride(0),
+        req_to_token.stride(0),
+        req_to_token.stride(1),
+        page_table_1.stride(0),
+        page_table_1.stride(1),
+        real_page_table.stride(0) if has_real_page_table else 0,
+        real_page_table.stride(1) if has_real_page_table else 0,
+        paged_mqa_ctx_lens_2d.stride(0) if has_paged_mqa_ctx_lens else 0,
+        paged_mqa_ctx_lens_2d.stride(1) if has_paged_mqa_ctx_lens else 0,
+        bs,
+        max_seqlen_k,
+        dsa_index_topk,
+        real_page_size,
+        next_n,
+        has_real_page_table,
+        has_paged_mqa_ctx_lens,
+        has_page_table_1,
+        BLOCK_BS=block_bs,
+        BLOCK_EXPANDED=block_expanded,
+        BLOCK_N=block_n,
+    )
+
+
+@triton.jit(
+    do_not_specialize=[
+        "page_table_stride_0",
+        "real_page_table_stride_0",
+        "total_len",
+        "max_seqlen_k",
+    ]
+)
+def _fused_dsa_draft_extend_metadata_kernel(
+    seq_lens,
+    extend_seq_lens,
+    req_pool_indices,
+    req_to_token,
+    cache_seqlens,
+    cu_seqlens_k,
+    page_table_1,
+    seqlens_expanded,
+    dsa_cache_seqlens,
+    dsa_cu_seqlens_k,
+    real_page_table,
+    seq_lens_stride: tl.constexpr,
+    extend_seq_lens_stride: tl.constexpr,
+    req_pool_indices_stride: tl.constexpr,
+    req_to_token_stride_0: tl.constexpr,
+    req_to_token_stride_1: tl.constexpr,
+    page_table_stride_0,
+    page_table_stride_1: tl.constexpr,
+    real_page_table_stride_0,
+    real_page_table_stride_1: tl.constexpr,
+    bs: tl.constexpr,
+    total_len,
+    max_seqlen_k,
+    dsa_index_topk: tl.constexpr,
+    real_page_size: tl.constexpr,
+    HAS_REAL_PAGE_TABLE: tl.constexpr,
+    HAS_PAGE_TABLE_1: tl.constexpr,
+    STATIC_EXTEND_LEN: tl.constexpr,
+    BLOCK_BS: tl.constexpr,
+    BLOCK_EXPANDED: tl.constexpr,
+    BLOCK_ROWS: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    pid = tl.program_id(0)
+
+    if pid == 0:
+        offs_b = tl.arange(0, BLOCK_BS)
+        mask_b = offs_b < bs
+        seq = tl.load(seq_lens + offs_b * seq_lens_stride, mask=mask_b, other=0)
+        cache_seq = seq.to(tl.int32)
+        cu = tl.cumsum(cache_seq, 0)
+
+        tl.store(cache_seqlens + offs_b, cache_seq, mask=mask_b)
+        tl.store(cu_seqlens_k, tl.full((), 0, tl.int32))
+        tl.store(cu_seqlens_k + 1 + offs_b, cu, mask=mask_b)
+
+        offs_e = tl.arange(0, BLOCK_EXPANDED)
+        mask_e = offs_e < total_len
+        if STATIC_EXTEND_LEN:
+            static_qo_len = tl.load(extend_seq_lens).to(tl.int32)
+            req_row = offs_e // static_qo_len
+            local_off = offs_e - req_row * static_qo_len
+            qo_len_for_row = tl.zeros((BLOCK_EXPANDED,), tl.int32) + static_qo_len
+        else:
+            req_row = tl.full((BLOCK_EXPANDED,), 0, tl.int32)
+            local_off = tl.full((BLOCK_EXPANDED,), 0, tl.int32)
+            qo_len_for_row = tl.full((BLOCK_EXPANDED,), 1, tl.int32)
+            prefix = tl.full((), 0, tl.int32)
+
+            for i in tl.range(0, bs):
+                qo_len = tl.load(extend_seq_lens + i * extend_seq_lens_stride).to(
+                    tl.int32
+                )
+                in_row = (offs_e >= prefix) & (offs_e < prefix + qo_len)
+                req_row = tl.where(in_row, i, req_row)
+                local_off = tl.where(in_row, offs_e - prefix, local_off)
+                qo_len_for_row = tl.where(in_row, qo_len, qo_len_for_row)
+                prefix += qo_len
+
+        base_seq = tl.load(
+            seq_lens + req_row * seq_lens_stride,
+            mask=mask_e,
+            other=0,
+        ).to(tl.int32)
+        # Clamp to >= 0: DP-padded / idle-companion rows carry the CUDA-graph
+        # seq_len fill value (1), which is smaller than qo_len, so the raw
+        # per-row visible kv length goes negative. Consumers treat these
+        # lengths as unsigned (the top-k v2 kernel reads them as uint32), so a
+        # negative row becomes a ~4e9-token length and an illegal memory
+        # access. 0 keeps padded rows on the trivial all-(-1) output path.
+        expanded_seq = base_seq - qo_len_for_row + local_off + 1
+        expanded_seq = tl.maximum(expanded_seq, 0)
+        expanded_seq = tl.where(mask_e, expanded_seq, 0)
+        dsa_seq = tl.minimum(expanded_seq, dsa_index_topk)
+        dsa_cu = tl.cumsum(dsa_seq, 0)
+
+        tl.store(seqlens_expanded + offs_e, expanded_seq, mask=mask_e)
+        tl.store(dsa_cache_seqlens + offs_e, dsa_seq, mask=mask_e)
+        tl.store(dsa_cu_seqlens_k, tl.full((), 0, tl.int32))
+        tl.store(dsa_cu_seqlens_k + 1 + offs_e, dsa_cu, mask=mask_e)
+        return
+
+    num_col_blocks = tl.cdiv(max_seqlen_k, BLOCK_N)
+    page_pid = pid - 1
+    req_row = page_pid // num_col_blocks
+    col_block = page_pid - req_row * num_col_blocks
+    offs_n = col_block * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    qo_len = tl.load(
+        extend_seq_lens + req_row * extend_seq_lens_stride,
+        mask=req_row < bs,
+        other=0,
+    ).to(tl.int32)
+    if STATIC_EXTEND_LEN:
+        prefix = req_row * qo_len
+    else:
+        prefix = tl.full((), 0, tl.int32)
+        for i in tl.range(0, bs):
+            prev_qo_len = tl.load(extend_seq_lens + i * extend_seq_lens_stride).to(
+                tl.int32
+            )
+            prefix += tl.where(i < req_row, prev_qo_len, 0)
+    offs_r = tl.arange(0, BLOCK_ROWS)
+    out_rows = prefix + offs_r
+    row_mask = (req_row < bs) & (offs_r < qo_len) & (out_rows < total_len)
+    col_mask = offs_n < max_seqlen_k
+    has_rows = (req_row < bs) & (qo_len > 0)
+    mask = row_mask[:, None] & col_mask[None, :]
+
+    req_idx = tl.load(
+        req_pool_indices + req_row * req_pool_indices_stride,
+        mask=has_rows,
+        other=0,
+    )
+    vals = tl.load(
+        req_to_token + req_idx * req_to_token_stride_0 + offs_n * req_to_token_stride_1,
+        mask=col_mask & has_rows,
+        other=0,
+    ).to(tl.int32)
+    # Write the wide page_size=1 table only when the caller provides it (see
+    # fused_dsa_decode_metadata for the optional-page_table_1 contract).
+    if HAS_PAGE_TABLE_1:
+        tl.store(
+            page_table_1
+            + out_rows[:, None] * page_table_stride_0
+            + offs_n[None, :] * page_table_stride_1,
+            vals[None, :],
+            mask=mask,
+        )
+
+    if HAS_REAL_PAGE_TABLE:
+        real_mask = mask & ((offs_n[None, :] % real_page_size) == 0)
+        real_cols = offs_n // real_page_size
+        tl.store(
+            real_page_table
+            + out_rows[:, None] * real_page_table_stride_0
+            + real_cols[None, :] * real_page_table_stride_1,
+            (vals // real_page_size)[None, :],
+            mask=real_mask,
+        )
+
+
+def fused_dsa_draft_extend_metadata(
+    seq_lens: torch.Tensor,
+    extend_seq_lens: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    cache_seqlens: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    page_table_1: Optional[torch.Tensor],
+    seqlens_expanded: torch.Tensor,
+    dsa_cache_seqlens: torch.Tensor,
+    dsa_cu_seqlens_k: torch.Tensor,
+    real_page_table: torch.Tensor,
+    bs: int,
+    total_len: int,
+    max_seqlen_k: int,
+    dsa_index_topk: int,
+    real_page_size: int,
+    max_extend_len: int,
+    max_total_len: int,
+    static_extend_len: bool = False,
+) -> None:
+    assert seq_lens.is_cuda
+    assert extend_seq_lens.is_cuda
+    assert req_pool_indices.is_cuda
+    assert req_to_token.is_cuda
+    assert cache_seqlens.is_cuda
+    assert cu_seqlens_k.is_cuda
+    assert seqlens_expanded.is_cuda
+    assert dsa_cache_seqlens.is_cuda
+    assert dsa_cu_seqlens_k.is_cuda
+
+    if bs == 0:
+        cu_seqlens_k[:1].zero_()
+        dsa_cu_seqlens_k[:1].zero_()
+        return
+    if total_len == 0:
+        cache = seq_lens.to(torch.int32)
+        cache_seqlens.copy_(cache)
+        cu_seqlens_k[:1].zero_()
+        cu_seqlens_k[1 : bs + 1].copy_(torch.cumsum(cache, dim=0, dtype=torch.int32))
+        dsa_cu_seqlens_k[:1].zero_()
+        return
+    assert total_len <= max_total_len
+    # Caller-owned graph metadata guarantees each request accepts at most
+    # max_extend_len tokens. Avoid checking extend_seq_lens.max() here because
+    # that would sync in the replay hot path.
+    assert max_extend_len > 0
+    assert total_len <= bs * max_extend_len
+
+    has_real_page_table = real_page_size > 1
+    if has_real_page_table:
+        assert real_page_table is not None
+        assert real_page_table.is_cuda
+    else:
+        assert page_table_1 is not None
+        real_page_table = page_table_1
+
+    # page_table_1 (the wide page_size=1 table) may be dropped for the fused
+    # decode CUDA graph; the kernel then writes only real_page_table.
+    has_page_table_1 = page_table_1 is not None
+    if not has_page_table_1:
+        assert has_real_page_table
+        page_table_1 = real_page_table  # dummy pointer for stride args
+    else:
+        assert page_table_1.is_cuda
+
+    block_bs = triton.next_power_of_2(bs)
+    block_expanded = triton.next_power_of_2(max_total_len)
+    block_rows = triton.next_power_of_2(max_extend_len)
+    block_n = 128
+    num_col_blocks = triton.cdiv(max_seqlen_k, block_n)
+    grid = (1 + bs * num_col_blocks,)
+
+    _fused_dsa_draft_extend_metadata_kernel[grid](
+        seq_lens,
+        extend_seq_lens,
+        req_pool_indices,
+        req_to_token,
+        cache_seqlens,
+        cu_seqlens_k,
+        page_table_1,
+        seqlens_expanded,
+        dsa_cache_seqlens,
+        dsa_cu_seqlens_k,
+        real_page_table,
+        seq_lens.stride(0),
+        extend_seq_lens.stride(0),
+        req_pool_indices.stride(0),
+        req_to_token.stride(0),
+        req_to_token.stride(1),
+        page_table_1.stride(0),
+        page_table_1.stride(1),
+        real_page_table.stride(0) if has_real_page_table else 0,
+        real_page_table.stride(1) if has_real_page_table else 0,
+        bs,
+        total_len,
+        max_seqlen_k,
+        dsa_index_topk,
+        real_page_size,
+        has_real_page_table,
+        has_page_table_1,
+        static_extend_len,
+        BLOCK_BS=block_bs,
+        BLOCK_EXPANDED=block_expanded,
+        BLOCK_ROWS=block_rows,
+        BLOCK_N=block_n,
+    )

python/sglang/kernels/ops/attention/merge_state.py

@@ -0,0 +1,96 @@
+from typing import Optional, Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def merge_state_kernel(
+    output,  # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE] v_merged
+    output_lse,  # [NUM_TOKENS, NUM_HEADS] s_merged
+    prefix_output,  # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE] v_a
+    prefix_lse,  # [NUM_TOKENS, NUM_HEADS] s_a
+    suffix_output,  # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE] v_b
+    suffix_lse,  # [NUM_TOKENS, NUM_HEADS] s_b
+    HEAD_SIZE: tl.constexpr,
+    PADDED_HEAD_SIZE: tl.constexpr,
+    OUTPUT_LSE: tl.constexpr,
+):
+    token_idx = tl.program_id(0)
+    num_tokens = tl.num_programs(0)
+    head_idx = tl.program_id(1)
+    num_heads = tl.num_programs(1)
+
+    p_lse = tl.load(prefix_lse + token_idx * num_heads + head_idx)
+    s_lse = tl.load(suffix_lse + token_idx * num_heads + head_idx)
+    p_lse = float("-inf") if p_lse == float("inf") else p_lse
+    s_lse = float("-inf") if s_lse == float("inf") else s_lse
+
+    max_lse = tl.maximum(p_lse, s_lse)
+    p_lse = p_lse - max_lse
+    s_lse = s_lse - max_lse
+    out_se = tl.exp(p_lse) + tl.exp(s_lse)
+
+    if OUTPUT_LSE:
+        out_lse = tl.log(out_se) + max_lse
+        tl.store(output_lse + token_idx * num_heads + head_idx, out_lse)
+
+    head_arange = tl.arange(0, PADDED_HEAD_SIZE)
+    head_mask = head_arange < HEAD_SIZE
+    p_out = tl.load(
+        prefix_output
+        + token_idx * num_heads * HEAD_SIZE
+        + head_idx * HEAD_SIZE
+        + head_arange,
+        mask=head_mask,
+    )
+    s_out = tl.load(
+        suffix_output
+        + token_idx * num_heads * HEAD_SIZE
+        + head_idx * HEAD_SIZE
+        + head_arange,
+        mask=head_mask,
+    )
+
+    p_scale = tl.exp(p_lse) / out_se
+    s_scale = tl.exp(s_lse) / out_se
+    out = p_out * p_scale + s_out * s_scale
+    tl.store(
+        output + token_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE + head_arange,
+        out,
+        mask=head_mask,
+    )
+
+
+def merge_state_triton(
+    prefix_output: torch.Tensor,
+    prefix_lse: torch.Tensor,
+    suffix_output: torch.Tensor,
+    suffix_lse: torch.Tensor,
+    output: Optional[torch.Tensor] = None,
+    output_lse: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    # Avoid creating new tensors if they are already provided
+    if output is None:
+        output = torch.empty_like(prefix_output)
+    if output_lse is None:
+        output_lse = torch.empty_like(prefix_lse)
+
+    num_tokens = output.shape[0]
+    num_query_heads = output.shape[1]
+    head_size = output.shape[2]
+    padded_head_size = triton.next_power_of_2(head_size)
+
+    merge_state_kernel[(num_tokens, num_query_heads)](
+        output,
+        output_lse,
+        prefix_output,
+        prefix_lse,
+        suffix_output,
+        suffix_lse,
+        head_size,
+        padded_head_size,
+        output_lse is not None,
+    )
+    return output, output_lse

python/sglang/kernels/ops/attention/metadata.py

@@ -0,0 +1,467 @@
+from typing import TYPE_CHECKING, Optional
+
+import torch
+import triton
+import triton.language as tl
+
+if TYPE_CHECKING:
+    from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
+
+
+@triton.jit
+def get_num_kv_splits_triton(
+    num_kv_splits_ptr,
+    seq_lens_ptr,
+    num_seq,
+    num_group,
+    num_head,
+    num_kv_head,
+    max_kv_splits,
+    device_core_count,
+    MAX_NUM_SEQ: tl.constexpr,
+):
+    # TODO: this method is tunable, we need more online serving data to tune it
+    offs_seq = tl.arange(0, MAX_NUM_SEQ)
+    mask_seq = offs_seq < num_seq
+
+    seq_lens = tl.load(seq_lens_ptr + offs_seq, mask=mask_seq, other=0)
+    max_seq_len = tl.max(seq_lens)
+    seq_lens = tl.load(seq_lens_ptr + offs_seq, mask=mask_seq, other=max_seq_len)
+    min_seq_len = tl.min(seq_lens)
+    if max_seq_len * 8 < min_seq_len * 10:
+        min_seq_len = max_seq_len
+    max_kv_splits_1 = tl.minimum(tl.cdiv(max_seq_len, min_seq_len), max_kv_splits)
+    kv_chunk_size_1 = tl.cdiv(max_seq_len, max_kv_splits_1)
+
+    # NOTE: this is a hack to let num_kv_split grows up with seqlen gradually
+    ext_seq_len = tl.cast(max_seq_len, tl.float32) / 64.0
+    ext_device_core_count = tl.cast(
+        device_core_count * tl.maximum(tl.log2(ext_seq_len), 1.0), tl.int32
+    )
+    block_h, num_kv_group = 16, num_head // num_kv_head
+    if num_kv_group == 1:
+        token_grid = num_seq * num_group * num_head
+    else:
+        # from triton_ops/decode_attention.py:_decode_grouped_att_m_fwd
+        block_h = tl.minimum(block_h, num_kv_group)
+        token_grid = num_seq * num_group * tl.cdiv(num_head, block_h)
+    max_kv_splits_2 = tl.minimum(
+        tl.cdiv(ext_device_core_count, token_grid), max_kv_splits
+    )
+    kv_chunk_size_2 = tl.cdiv(max_seq_len, max_kv_splits_2)
+
+    num_kv_splits = tl.maximum(
+        tl.cdiv(seq_lens, kv_chunk_size_1), tl.cdiv(seq_lens, kv_chunk_size_2)
+    )
+
+    offs_token = offs_seq * num_group
+    mask_token = offs_token < num_seq * num_group
+    for i in range(0, num_group):
+        tl.store(num_kv_splits_ptr + i + offs_token, num_kv_splits, mask=mask_token)
+
+
+@triton.jit
+def _prepare_swa_spec_page_table_kernel(
+    dst_ptr,
+    src_a_ptr,
+    src_b_ptr,
+    seq_len_a_ptr,
+    seq_len_b_ptr,
+    dst_stride_m,
+    dst_stride_n,
+    a_stride_m,
+    a_stride_n,
+    b_stride_m,
+    b_stride_n,
+    LEN_A: tl.constexpr,
+    LEN_B: tl.constexpr,
+    REPEAT_STEP: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+
+    idx_a = pid_m // REPEAT_STEP
+    idx_b = pid_m
+    seq_len_a = tl.load(seq_len_a_ptr + idx_a)
+    seq_len_b = tl.load(seq_len_b_ptr + idx_b)
+
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    total_len = seq_len_a + seq_len_b
+
+    if pid_n * BLOCK_N >= total_len:
+        return
+
+    mask = offs_n < total_len
+    dst = dst_ptr + pid_m * dst_stride_m + offs_n * dst_stride_n
+
+    if (pid_n + 1) * BLOCK_N < seq_len_a:
+        a_ptr = src_a_ptr + idx_a * a_stride_m + offs_n * a_stride_n
+        a_mask = mask & (offs_n < LEN_A)
+        val = tl.load(a_ptr, mask=a_mask, other=0)
+        tl.store(dst, val, mask=mask)
+    elif pid_n * BLOCK_N >= seq_len_a:
+        offs_b = offs_n - seq_len_a
+        b_ptr = src_b_ptr + idx_b * b_stride_m + offs_b * b_stride_n
+        b_mask = mask & (offs_b < LEN_B)
+        val = tl.load(b_ptr, mask=b_mask, other=0)
+        tl.store(dst, val, mask=mask)
+    else:
+        # mixed part
+        a_offs = offs_n
+        a_mask = (a_offs < seq_len_a) & (a_offs < LEN_A)
+        a_ptr = src_a_ptr + idx_a * a_stride_m + a_offs * a_stride_n
+        a_val = tl.load(a_ptr, mask=a_mask, other=0)
+
+        b_offs = offs_n - seq_len_a
+        b_mask = (b_offs >= 0) & (b_offs < seq_len_b) & (b_offs < LEN_B)
+        b_ptr = src_b_ptr + idx_b * b_stride_m + b_offs * b_stride_n
+        b_val = tl.load(b_ptr, mask=b_mask, other=0)
+
+        result = tl.where(offs_n < seq_len_a, a_val, b_val)
+        tl.store(dst, result, mask=mask)
+
+
+def prepare_swa_spec_page_table_triton(
+    page_table_dst: torch.Tensor,
+    page_table_a: torch.Tensor,
+    page_table_b: torch.Tensor,  # expand page table
+    seq_len_a: torch.Tensor,
+    seq_len_b: torch.Tensor,  # expand seq lens
+    speculative_num_draft_tokens: int,
+):
+    # concat page_table and expand page_table by kv seq length
+    bs = seq_len_a.numel()
+    bs_expand = seq_len_b.numel()
+    assert bs_expand == bs * speculative_num_draft_tokens
+
+    LEN_A = page_table_a.shape[1]
+    LEN_B = page_table_b.shape[1]
+    LEN_OUT = LEN_A + LEN_B
+    REPEAT_STEP = speculative_num_draft_tokens
+    BLOCK_N = 256
+
+    grid = (bs_expand, triton.cdiv(LEN_OUT, BLOCK_N))
+    _prepare_swa_spec_page_table_kernel[grid](
+        page_table_dst,
+        page_table_a,
+        page_table_b,
+        seq_len_a,
+        seq_len_b,
+        page_table_dst.stride(0),
+        page_table_dst.stride(1),
+        page_table_a.stride(0),
+        page_table_a.stride(1),
+        page_table_b.stride(0),
+        page_table_b.stride(1),
+        LEN_A=LEN_A,
+        LEN_B=LEN_B,
+        REPEAT_STEP=REPEAT_STEP,
+        BLOCK_N=BLOCK_N,
+        num_warps=4,
+    )
+
+
+@triton.jit
+def _fused_metadata_kernel_general(
+    # Input tensors
+    seq_lens,
+    seq_lens_stride_0,
+    req_to_token,
+    req_to_token_stride_0,
+    req_to_token_stride_1,
+    req_pool_indices,
+    req_pool_indices_stride_0,
+    # Output buffers
+    cache_seqlens_int32,
+    cache_seqlens_int32_stride_0,
+    cu_seqlens_k,
+    cu_seqlens_k_stride_0,
+    page_table,
+    page_table_stride_0,
+    page_table_stride_1,
+    swa_page_table,
+    swa_page_table_stride_0,
+    swa_page_table_stride_1,
+    full_to_swa_mapping,
+    full_to_swa_mapping_stride_0,
+    # Scalar parameters
+    B,
+    max_seq_pages,
+    page_size: tl.constexpr,
+    seq_len_delta: tl.constexpr,
+    use_swa: tl.constexpr,
+    SHIFT: tl.constexpr,
+    BLOCK_COLS: tl.constexpr,
+):
+    pid_b = tl.program_id(0)  # batch index
+    pid_c = tl.program_id(1)  # column chunk index
+
+    # 1. Prefix sum (only one block does it)
+    if pid_b == 0 and pid_c == 0:
+        acc = 0
+        for idx in range(B):
+            seq = tl.load(seq_lens + idx * seq_lens_stride_0)
+            val = (seq + seq_len_delta).to(tl.int32)
+            tl.store(cache_seqlens_int32 + idx * cache_seqlens_int32_stride_0, val)
+            tl.store(cu_seqlens_k + idx * cu_seqlens_k_stride_0, acc)
+            acc += val
+        tl.store(cu_seqlens_k + B * cu_seqlens_k_stride_0, acc)
+
+    # 2. Gather for this batch and column chunk
+    if max_seq_pages == 0:
+        return
+
+    i = pid_b
+    # Load row index for this batch (all threads in block have same i)
+    row_idx = tl.load(req_pool_indices + i * req_pool_indices_stride_0)
+    row_offset = row_idx * req_to_token_stride_0
+
+    col_start = pid_c * BLOCK_COLS
+    col_offsets = col_start + tl.arange(0, BLOCK_COLS)
+    mask = col_offsets < max_seq_pages
+
+    # Compute column indices in the source tensor (token offset)
+    if page_size == 1:
+        col_idx = col_offsets
+    else:
+        col_idx = col_offsets << SHIFT  # faster than multiplication for power-of-two
+
+    # Load page indices from req_to_token
+    rt_offsets = row_offset + col_idx * req_to_token_stride_1
+    page_index = tl.load(
+        req_to_token + rt_offsets, mask=mask, other=0, cache_modifier=".cg"
+    )
+
+    # Compute page_table
+    if page_size == 1:
+        page_table_val = page_index
+    else:
+        page_table_val = page_index >> SHIFT
+
+    # Store to page_table
+    pt_offsets = i * page_table_stride_0 + col_offsets * page_table_stride_1
+    tl.store(page_table + pt_offsets, page_table_val, mask=mask, cache_modifier=".cg")
+
+    if use_swa:
+        swa_slot = tl.load(
+            full_to_swa_mapping + page_index * full_to_swa_mapping_stride_0,
+            mask=mask,
+            other=0,
+            cache_modifier=".cg",
+        )
+        if page_size == 1:
+            swa_val = swa_slot
+        else:
+            swa_val = swa_slot >> SHIFT
+        swa_offsets = (
+            i * swa_page_table_stride_0 + col_offsets * swa_page_table_stride_1
+        )
+        tl.store(swa_page_table + swa_offsets, swa_val, mask=mask, cache_modifier=".cg")
+
+
+@triton.jit
+def _fused_metadata_kernel_ps1_no_swa(
+    # Input tensors
+    seq_lens,
+    seq_lens_stride_0,
+    req_to_token,
+    req_to_token_stride_0,
+    req_to_token_stride_1,
+    req_pool_indices,
+    req_pool_indices_stride_0,
+    # Output buffers
+    cache_seqlens_int32,
+    cache_seqlens_int32_stride_0,
+    cu_seqlens_k,
+    cu_seqlens_k_stride_0,
+    page_table,
+    page_table_stride_0,
+    page_table_stride_1,
+    # Scalar parameters
+    B,
+    max_seq_pages,
+    seq_len_delta: tl.constexpr,
+    BLOCK_COLS: tl.constexpr,
+):
+    pid_b = tl.program_id(0)  # batch index
+    pid_c = tl.program_id(1)  # column chunk index
+
+    # 1. Prefix sum (only one block does it)
+    if pid_b == 0 and pid_c == 0:
+        acc = 0
+        for idx in range(B):
+            seq = tl.load(seq_lens + idx * seq_lens_stride_0)
+            val = (seq + seq_len_delta).to(tl.int32)
+            tl.store(cache_seqlens_int32 + idx * cache_seqlens_int32_stride_0, val)
+            tl.store(cu_seqlens_k + idx * cu_seqlens_k_stride_0, acc)
+            acc += val
+        tl.store(cu_seqlens_k + B * cu_seqlens_k_stride_0, acc)
+
+    # 2. Gather for this batch and column chunk
+    if max_seq_pages == 0:
+        return
+
+    i = pid_b
+    # Load row index for this batch (all threads in block have same i)
+    row_idx = tl.load(req_pool_indices + i * req_pool_indices_stride_0)
+    row_offset = row_idx * req_to_token_stride_0
+
+    col_start = pid_c * BLOCK_COLS
+    col_offsets = col_start + tl.arange(0, BLOCK_COLS)
+    mask = col_offsets < max_seq_pages
+
+    # page_size = 1: col_idx = col_offsets
+    rt_offsets = row_offset + col_offsets * req_to_token_stride_1
+    page_index = tl.load(
+        req_to_token + rt_offsets, mask=mask, other=0, cache_modifier=".cg"
+    )
+
+    # page_table = page_index // 1 = page_index
+    pt_offsets = i * page_table_stride_0 + col_offsets * page_table_stride_1
+    tl.store(page_table + pt_offsets, page_index, mask=mask, cache_modifier=".cg")
+
+
+def normal_decode_set_metadata(
+    cache_seqlens_int32: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    page_table: torch.Tensor,
+    req_to_token: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    strided_indices: torch.Tensor,
+    max_seq_pages: torch.Tensor,
+    seq_lens: torch.Tensor,
+    seq_len_delta: int,
+    page_size: int,
+    swa_page_table: Optional[torch.Tensor] = None,
+    token_to_kv_pool: Optional["SWAKVPool"] = None,
+):
+    """
+    Fused Triton implementation that replaces 4-5 sequential CUDA kernels with 1-2 kernels:
+      1. cache_seqlens = seq_lens + seq_len_delta (int64->int32 cast)
+      2. cu_seqlens_k = cumsum(cache_seqlens) (prefix-sum)
+      3. page_indices = req_to_token[pool_idx, stride_idx] (2-D gather)
+      4. page_table = page_indices // page_size (floor-divide)
+      5. (optional) swa_page_table for sliding window attention
+
+    Achieves ~5.2x speedup on H200 hardware for typical decode workloads.
+    """
+    assert (
+        page_size > 0 and (page_size & (page_size - 1)) == 0
+    ), f"page_size must be a power of two, got {page_size}"
+
+    batch_size = cache_seqlens_int32.shape[0]
+    device = seq_lens.device
+
+    # Ensure contiguous memory layout for efficient Triton access
+    seq_lens = seq_lens.contiguous()
+    req_to_token = req_to_token.contiguous()
+    req_pool_indices = req_pool_indices.contiguous()
+
+    # Prepare tensor strides
+    seq_lens_stride_0 = seq_lens.stride(0)
+    req_to_token_stride_0 = req_to_token.stride(0)
+    req_to_token_stride_1 = req_to_token.stride(1)
+    req_pool_indices_stride_0 = req_pool_indices.stride(0)
+    cache_seqlens_int32_stride_0 = cache_seqlens_int32.stride(0)
+    cu_seqlens_k_stride_0 = cu_seqlens_k.stride(0)
+    page_table_stride_0 = page_table.stride(0)
+    page_table_stride_1 = page_table.stride(1)
+
+    # Check if we should use the specialized fast path for page_size=1, no SWA
+    use_swa = swa_page_table is not None and token_to_kv_pool is not None
+
+    if page_size == 1 and not use_swa:
+        # Specialized kernel for the common case (page_size=1, no SWA)
+        BLOCK_COLS = 256
+        if max_seq_pages == 0:
+            grid = (1, 1)
+        else:
+            num_blocks_j = triton.cdiv(max_seq_pages, BLOCK_COLS)
+            grid = (batch_size, num_blocks_j)
+
+        _fused_metadata_kernel_ps1_no_swa[grid](
+            seq_lens,
+            seq_lens_stride_0,
+            req_to_token,
+            req_to_token_stride_0,
+            req_to_token_stride_1,
+            req_pool_indices,
+            req_pool_indices_stride_0,
+            cache_seqlens_int32,
+            cache_seqlens_int32_stride_0,
+            cu_seqlens_k,
+            cu_seqlens_k_stride_0,
+            page_table,
+            page_table_stride_0,
+            page_table_stride_1,
+            batch_size,
+            max_seq_pages,
+            seq_len_delta,
+            BLOCK_COLS=BLOCK_COLS,
+            num_warps=8,
+            num_stages=3,
+        )
+    else:
+        # General kernel for page_size > 1 or SWA cases
+        # SWA parameters
+        if use_swa:
+            from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
+
+            assert isinstance(token_to_kv_pool, SWAKVPool)
+            swa_page_table = swa_page_table.contiguous()
+            swa_page_table_stride_0 = swa_page_table.stride(0)
+            swa_page_table_stride_1 = swa_page_table.stride(1)
+            # Extract the full_to_swa_index_mapping from token_to_kv_pool
+            full_to_swa_mapping = (
+                token_to_kv_pool.full_to_swa_index_mapping.contiguous()
+            )
+            full_to_swa_mapping_stride_0 = full_to_swa_mapping.stride(0)
+        else:
+            # Dummy tensors (not used)
+            swa_page_table = torch.empty(0, dtype=torch.int32, device=device)
+            swa_page_table_stride_0 = 0
+            swa_page_table_stride_1 = 0
+            full_to_swa_mapping = torch.empty(0, dtype=torch.int32, device=device)
+            full_to_swa_mapping_stride_0 = 0
+
+        # Kernel configuration
+        BLOCK_COLS = 128
+        shift = (page_size).bit_length() - 1 if page_size > 1 else 0
+
+        if max_seq_pages == 0:
+            grid = (1, 1)
+        else:
+            num_blocks_j = triton.cdiv(max_seq_pages, BLOCK_COLS)
+            grid = (batch_size, num_blocks_j)
+
+        _fused_metadata_kernel_general[grid](
+            seq_lens,
+            seq_lens_stride_0,
+            req_to_token,
+            req_to_token_stride_0,
+            req_to_token_stride_1,
+            req_pool_indices,
+            req_pool_indices_stride_0,
+            cache_seqlens_int32,
+            cache_seqlens_int32_stride_0,
+            cu_seqlens_k,
+            cu_seqlens_k_stride_0,
+            page_table,
+            page_table_stride_0,
+            page_table_stride_1,
+            swa_page_table,
+            swa_page_table_stride_0,
+            swa_page_table_stride_1,
+            full_to_swa_mapping,
+            full_to_swa_mapping_stride_0,
+            batch_size,
+            max_seq_pages,
+            page_size,
+            seq_len_delta,
+            use_swa,
+            shift,
+            BLOCK_COLS=BLOCK_COLS,
+            num_warps=4,
+            num_stages=3,
+        )

python/sglang/kernels/ops/attention/pad.py

@@ -0,0 +1,390 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def pad_sequence_with_mask_kernel(
+    input_ptr,  # (total_tokens, hidden)
+    offsets_ptr,  # (B,)
+    lengths_ptr,  # (B,)
+    output_ptr,  # (B, max_len, hidden)
+    mask_ptr,  # (B, max_len)
+    max_len,
+    hidden_dim,
+    BLOCK_M: tl.constexpr,  # seq block
+    BLOCK_D: tl.constexpr,  # hidden block
+):
+    b = tl.program_id(0)  # batch index
+    m = tl.program_id(1)  # seq block index
+
+    offset = tl.load(offsets_ptr + b)
+    length = tl.load(lengths_ptr + b)
+
+    seq_ids = m * BLOCK_M + tl.arange(0, BLOCK_M)
+    hid_ids = tl.arange(0, BLOCK_D)
+
+    seq_mask = seq_ids < max_len
+    valid_token = seq_ids < length
+
+    # input index
+    in_token = offset + seq_ids
+    in_ptr = input_ptr + in_token[:, None] * hidden_dim + hid_ids[None, :]
+
+    # output index
+    out_ptr = (
+        output_ptr
+        + b * max_len * hidden_dim
+        + seq_ids[:, None] * hidden_dim
+        + hid_ids[None, :]
+    )
+
+    values = tl.load(
+        in_ptr,
+        mask=valid_token[:, None] & (hid_ids[None, :] < hidden_dim),
+        other=0.0,
+    )
+
+    tl.store(
+        out_ptr,
+        values,
+        mask=seq_mask[:, None] & (hid_ids[None, :] < hidden_dim),
+    )
+
+    # attention mask
+    if tl.program_id(2) == 0:
+        mask_out_ptr = mask_ptr + b * max_len + seq_ids
+        tl.store(mask_out_ptr, valid_token, mask=seq_mask)
+
+
+def pad_sequence_with_mask(
+    input_emb,  # (total_tokens, hidden)
+    offsets,  # (B,)
+    lengths,  # (B,)
+    max_len,
+):
+    B = offsets.shape[0]
+    hidden_dim = input_emb.shape[1]
+
+    output = torch.zeros(
+        (B, max_len, hidden_dim),
+        device=input_emb.device,
+        dtype=input_emb.dtype,
+    )
+    attn_mask = torch.empty(
+        (B * max_len),
+        device=input_emb.device,
+        dtype=torch.bool,
+    )
+
+    BLOCK_D = triton.next_power_of_2(hidden_dim)
+    BLOCK_M = triton.next_power_of_2(max_len)
+
+    grid = (
+        B,
+        triton.cdiv(max_len, BLOCK_M),
+        1,
+    )
+
+    pad_sequence_with_mask_kernel[grid](
+        input_emb,
+        offsets,
+        lengths,
+        output,
+        attn_mask,
+        max_len,
+        hidden_dim,
+        BLOCK_M=BLOCK_M,
+        BLOCK_D=BLOCK_D,
+    )
+
+    return B, output, attn_mask
+
+
+@triton.jit
+def pad_draft_extend_query_kernel(
+    q_ptr,  # Input query tensor [total_seq_len, num_heads, head_dim]
+    padded_q_ptr,  # Output padded query tensor [batch_size, max_seq_len, num_heads, head_dim]
+    seq_lens_q_ptr,  # Sequence lengths for each sequence [batch_size]
+    cumsum_ptr,  # Cumulative sum of sequence lengths [batch_size + 1]
+    batch_size,
+    max_seq_len,
+    num_heads,
+    head_dim,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Triton kernel for padding draft extended query tensor with parallelized head and dim processing."""
+    # Use 3D program IDs: (batch_seq, head_block, dim_block)
+    batch_seq_pid = tl.program_id(0)
+    head_pid = tl.program_id(1)
+    dim_pid = tl.program_id(2)
+
+    batch_id = batch_seq_pid // max_seq_len
+    seq_pos = batch_seq_pid % max_seq_len
+
+    if batch_id >= batch_size:
+        return
+
+    # Load sequence length for this batch
+    seq_len = tl.load(seq_lens_q_ptr + batch_id)
+
+    if seq_pos >= seq_len:
+        return
+
+    # Load cumulative sum to get start position in input tensor
+    input_start = tl.load(cumsum_ptr + batch_id)
+    input_pos = input_start + seq_pos
+
+    # Calculate head and dim block ranges
+    head_start = head_pid * BLOCK_SIZE
+    head_end = tl.minimum(head_start + BLOCK_SIZE, num_heads)
+    head_mask = tl.arange(0, BLOCK_SIZE) < (head_end - head_start)
+
+    dim_start = dim_pid * BLOCK_SIZE
+    dim_end = tl.minimum(dim_start + BLOCK_SIZE, head_dim)
+    dim_mask = tl.arange(0, BLOCK_SIZE) < (dim_end - dim_start)
+
+    # Calculate input offset
+    input_offset = (
+        input_pos * num_heads * head_dim
+        + (head_start + tl.arange(0, BLOCK_SIZE))[:, None] * head_dim
+        + (dim_start + tl.arange(0, BLOCK_SIZE))[None, :]
+    )
+
+    # Load data
+    data = tl.load(
+        q_ptr + input_offset,
+        mask=head_mask[:, None] & dim_mask[None, :],
+        other=0.0,
+    )
+
+    # Calculate output offset
+    output_offset = (
+        batch_id * max_seq_len * num_heads * head_dim
+        + seq_pos * num_heads * head_dim
+        + (head_start + tl.arange(0, BLOCK_SIZE))[:, None] * head_dim
+        + (dim_start + tl.arange(0, BLOCK_SIZE))[None, :]
+    )
+
+    # Store data
+    tl.store(
+        padded_q_ptr + output_offset,
+        data,
+        mask=head_mask[:, None] & dim_mask[None, :],
+    )
+
+
+def pad_draft_extend_query(
+    q: torch.Tensor,
+    padded_q: torch.Tensor,
+    seq_lens_q: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+) -> torch.Tensor:
+    """Pad draft extended query using Triton kernel."""
+    batch_size = cu_seqlens_q.shape[0] - 1
+    max_seq_len_q = padded_q.shape[1]
+    num_heads = padded_q.shape[2]
+    head_dim = padded_q.shape[3]
+
+    # Launch Triton kernel with 3D grid for parallelized head and dim processing
+    BLOCK_SIZE = 64
+    num_head_blocks = triton.cdiv(num_heads, BLOCK_SIZE)
+    num_dim_blocks = triton.cdiv(head_dim, BLOCK_SIZE)
+    grid = (batch_size * max_seq_len_q, num_head_blocks, num_dim_blocks)
+
+    pad_draft_extend_query_kernel[grid](
+        q_ptr=q,
+        padded_q_ptr=padded_q,
+        seq_lens_q_ptr=seq_lens_q,
+        cumsum_ptr=cu_seqlens_q,
+        batch_size=batch_size,
+        max_seq_len=max_seq_len_q,
+        num_heads=num_heads,
+        head_dim=head_dim,
+        BLOCK_SIZE=BLOCK_SIZE,
+    )
+    return padded_q
+
+
+@triton.jit
+def unpad_draft_extend_output_kernel(
+    raw_out_ptr,  # Input raw output tensor (batch_size, token_per_batch, tp_q_head_num, v_head_dim)
+    output_ptr,  # Output tensor (-1, tp_q_head_num, v_head_dim)
+    num_accept_tokens_ptr,  # Accept lengths for each sequence [batch_size]
+    cumsum_ptr,  # Cumulative sum of accept lengths [batch_size + 1]
+    batch_size,
+    token_per_batch,
+    tp_q_head_num,
+    v_head_dim,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Triton kernel for unpadding draft extended output tensor with parallelized head and dim processing."""
+    batch_seq_pid = tl.program_id(0)
+    head_pid = tl.program_id(1)
+    dim_pid = tl.program_id(2)
+
+    batch_id = batch_seq_pid // token_per_batch
+    seq_pos = batch_seq_pid % token_per_batch
+
+    if batch_id >= batch_size:
+        return
+
+    # Load accept length for this batch
+    accept_len = tl.load(num_accept_tokens_ptr + batch_id)
+
+    if seq_pos >= accept_len:
+        return
+
+    # Load cumulative sum to get start position in output tensor
+    output_start = tl.load(cumsum_ptr + batch_id)
+    output_pos = output_start + seq_pos
+
+    # Calculate head and dim block ranges
+    head_start = head_pid * BLOCK_SIZE
+    head_end = tl.minimum(head_start + BLOCK_SIZE, tp_q_head_num)
+    head_mask = tl.arange(0, BLOCK_SIZE) < (head_end - head_start)
+
+    dim_start = dim_pid * BLOCK_SIZE
+    dim_end = tl.minimum(dim_start + BLOCK_SIZE, v_head_dim)
+    dim_mask = tl.arange(0, BLOCK_SIZE) < (dim_end - dim_start)
+
+    # Calculate input offset: (batch_id, seq_pos, head_id, dim_id)
+    input_offset = (
+        batch_id * token_per_batch * tp_q_head_num * v_head_dim
+        + seq_pos * tp_q_head_num * v_head_dim
+        + (head_start + tl.arange(0, BLOCK_SIZE))[:, None] * v_head_dim
+        + (dim_start + tl.arange(0, BLOCK_SIZE))[None, :]
+    )
+
+    # Load data
+    data = tl.load(
+        raw_out_ptr + input_offset,
+        mask=head_mask[:, None] & dim_mask[None, :],
+        other=0.0,
+    )
+
+    output_offset = (
+        output_pos * tp_q_head_num * v_head_dim
+        + (head_start + tl.arange(0, BLOCK_SIZE))[:, None] * v_head_dim
+        + (dim_start + tl.arange(0, BLOCK_SIZE))[None, :]
+    )
+
+    # Store data
+    tl.store(
+        output_ptr + output_offset,
+        data,
+        mask=head_mask[:, None] & dim_mask[None, :],
+    )
+
+
+def unpad_draft_extend_output(
+    raw_out: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    seq_lens_q: torch.Tensor,
+    sum_seq_lens_q: int,
+    unpad_output_buffer: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Unpad draft extended output using Triton kernel."""
+    # raw_out: (batch_size, token_per_batch, layer.tp_q_head_num, layer.v_head_dim)
+    batch_size = seq_lens_q.shape[0]
+    token_per_batch = raw_out.shape[1]  # max_seq_len
+    tp_q_head_num = raw_out.shape[2]  # num_heads
+    v_head_dim = raw_out.shape[3]  # head_dim
+    total_tokens = sum_seq_lens_q
+
+    # Check if we're in CUDA graph mode (buffers are pre-allocated)
+    if unpad_output_buffer is not None:
+        # Use pre-allocated buffer for CUDA graph compatibility
+        output = unpad_output_buffer[:total_tokens, :, :].to(dtype=raw_out.dtype)
+    else:
+        # Dynamic allocation for non-CUDA graph mode
+        output = torch.empty(
+            (total_tokens, tp_q_head_num, v_head_dim),
+            dtype=raw_out.dtype,
+            device=raw_out.device,
+        )
+
+    # Launch Triton kernel with 3D grid for parallelized head and dim processing
+    BLOCK_SIZE = 64
+    num_head_blocks = triton.cdiv(tp_q_head_num, BLOCK_SIZE)
+    num_dim_blocks = triton.cdiv(v_head_dim, BLOCK_SIZE)
+    grid = (batch_size * token_per_batch, num_head_blocks, num_dim_blocks)
+
+    unpad_draft_extend_output_kernel[grid](
+        raw_out_ptr=raw_out,
+        output_ptr=output,
+        num_accept_tokens_ptr=seq_lens_q,
+        cumsum_ptr=cu_seqlens_q,
+        batch_size=batch_size,
+        token_per_batch=token_per_batch,
+        tp_q_head_num=tp_q_head_num,
+        v_head_dim=v_head_dim,
+        BLOCK_SIZE=BLOCK_SIZE,
+    )
+    return output[:total_tokens, :, :]
+
+
+@triton.jit
+def seqlens_expand_kernel(
+    extend_seq_lens_ptr,  # [N]
+    seq_lens_ptr,  # [N]
+    offsets_ptr,  # [N+1]
+    output_ptr,  # [sum(extend_seq_lens)]
+    N,
+    BLOCK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+
+    if pid >= N:
+        return
+
+    qo_len = tl.load(extend_seq_lens_ptr + pid)
+    kv_len = tl.load(seq_lens_ptr + pid)
+
+    start = kv_len - qo_len + 1
+    out_offset = tl.load(offsets_ptr + pid)
+
+    offs = tl.arange(0, BLOCK)
+    mask = offs < qo_len
+
+    # Clamp to >= 0: rows with kv_len < qo_len (DP-padded / idle-companion
+    # rows whose kv is the CUDA-graph fill value) would otherwise produce
+    # negative lengths, which unsigned consumers (e.g. the top-k v2 kernel,
+    # which reads lengths as uint32) turn into ~4e9-token lengths and an
+    # illegal memory access.
+    values = tl.maximum(start + offs, 0)
+    tl.store(output_ptr + out_offset + offs, values, mask=mask)
+
+
+def seqlens_expand_triton(
+    extend_seq_lens: torch.Tensor,
+    seq_lens: torch.Tensor,
+    total_len: int,
+    max_q_len: int,
+):
+    """
+    extend_seq_lens: [N], int32, CUDA
+    seq_lens:        [N], int32, CUDA
+    """
+    assert extend_seq_lens.is_cuda
+    assert seq_lens.is_cuda
+
+    N = extend_seq_lens.numel()
+
+    offsets = torch.zeros(N + 1, device=extend_seq_lens.device, dtype=torch.int32)
+    offsets[1:] = torch.cumsum(extend_seq_lens, dim=0)
+    output = torch.empty(total_len, device=extend_seq_lens.device, dtype=torch.int32)
+
+    BLOCK = triton.next_power_of_2(max_q_len)
+    grid = (N,)
+
+    seqlens_expand_kernel[grid](
+        extend_seq_lens,
+        seq_lens,
+        offsets,
+        output,
+        N,
+        BLOCK=BLOCK,
+    )
+
+    return output

python/sglang/kernels/ops/attention/position.py

@@ -0,0 +1,59 @@
+import torch
+import triton
+import triton.language as tl
+
+
+def compute_position_triton(
+    extend_prefix_lens: torch.Tensor, extend_seq_lens: torch.Tensor, extend_seq_lens_sum
+):
+    """Compute positions. It is a fused version of `compute_position_torch`."""
+    batch_size = extend_seq_lens.shape[0]
+    has_prefix = extend_prefix_lens.shape[0] == batch_size
+
+    positions = torch.empty(
+        extend_seq_lens_sum, dtype=torch.int64, device=extend_seq_lens.device
+    )
+    extend_start_loc = torch.empty(
+        batch_size, dtype=torch.int32, device=extend_seq_lens.device
+    )
+
+    # Launch kernel
+    compute_position_kernel[(batch_size,)](
+        positions,
+        extend_start_loc,
+        extend_prefix_lens,
+        extend_seq_lens,
+        has_prefix,
+    )
+
+    return positions, extend_start_loc
+
+
+@triton.jit
+def compute_position_kernel(
+    positions,
+    extend_start_loc,
+    extend_prefix_lens,
+    extend_seq_lens,
+    has_prefix: tl.constexpr,
+):
+    BLOCK_SIZE: tl.constexpr = 512
+    pid = tl.program_id(0).to(tl.int64)
+
+    prefix_len = tl.load(extend_prefix_lens + pid) if has_prefix else 0
+    seq_len = tl.load(extend_seq_lens + pid)
+
+    # NOTE: This can be slow for large bs
+    cumsum_start = tl.cast(0, tl.int64)
+    for i in range(pid):
+        cumsum_start += tl.load(extend_seq_lens + i)
+
+    num_loop = tl.cdiv(seq_len, BLOCK_SIZE)
+    for i in range(num_loop):
+        offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
+        tl.store(
+            positions + cumsum_start + offset,
+            prefix_len + offset,
+            mask=offset < seq_len,
+        )
+    tl.store(extend_start_loc + pid, cumsum_start)

python/sglang/kernels/ops/attention/prefill_attention.py

@@ -0,0 +1,219 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""
+Memory-efficient attention for prefill.
+It supporst page size = 1.
+"""
+
+# Adapted from
+# https://github.com/ModelTC/lightllm/blob/f2a54f0912293f683bf1d1695fd12c4098a5bf82/lightllm/models/llama/triton_kernel/context_flashattention_nopad.py#L1
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.utils import is_cuda, is_hip
+
+_is_cuda = is_cuda()
+_is_hip = is_hip()
+
+if _is_cuda or _is_hip:
+    CUDA_CAPABILITY = torch.cuda.get_device_capability()
+
+
+@triton.jit
+def _fwd_kernel(
+    Q,
+    K,
+    V,
+    sm_scale,
+    B_Start_Loc,
+    B_Seqlen,
+    Out,
+    stride_qbs,
+    stride_qh,
+    stride_kbs,
+    stride_kh,
+    stride_vbs,
+    stride_vh,
+    stride_obs,
+    stride_oh,
+    kv_group_num: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+    Lk: tl.constexpr,
+):
+    cur_batch = tl.program_id(0)
+    cur_head = tl.program_id(1)
+    start_m = tl.program_id(2)
+
+    cur_kv_head = cur_head // kv_group_num
+
+    cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
+    cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
+
+    block_start_loc = BLOCK_M * start_m
+
+    # initialize offsets
+    offs_n = tl.arange(0, BLOCK_N)
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    off_q = (
+        (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs
+        + cur_head * stride_qh
+        + offs_d[None, :]
+    )
+    off_k = offs_n[None, :] * stride_kbs + cur_kv_head * stride_kh + offs_d[:, None]
+    off_v = offs_n[:, None] * stride_vbs + cur_kv_head * stride_vh + offs_d[None, :]
+
+    mask_d = offs_d < Lk
+
+    q = tl.load(
+        Q + off_q,
+        mask=(offs_m[:, None] < cur_batch_seq_len) & (mask_d[None, :]),
+        other=0.0,
+    )
+
+    k_ptrs = K + off_k
+    v_ptrs = V + off_v
+
+    # initialize pointer to m and l
+    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
+    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
+
+    block_mask = tl.where(block_start_loc < cur_batch_seq_len, 1, 0)
+
+    end_n = (
+        cur_batch_seq_len
+        if not IS_CAUSAL
+        else tl.minimum((start_m + 1) * BLOCK_M, cur_batch_seq_len)
+    )
+    for start_n in range(0, block_mask * end_n, BLOCK_N):
+        start_n = tl.multiple_of(start_n, BLOCK_N)
+        # -- compute qk ----
+        k = tl.load(
+            k_ptrs + (cur_batch_in_all_start_index + start_n) * stride_kbs,
+            mask=((start_n + offs_n[None, :]) < cur_batch_seq_len) & (mask_d[:, None]),
+            other=0.0,
+        )
+        # mask = tl.load(mask_ptrs + start_n, mask=start_n + offs_n < cur_batch_end_loc, other=0.0)
+
+        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+        qk += tl.dot(q, k)
+        qk *= sm_scale
+
+        if IS_CAUSAL:
+            qk += tl.where(
+                (start_n + offs_n[None, :] < cur_batch_seq_len)
+                & (offs_m[:, None] >= (start_n + offs_n[None, :])),
+                0,
+                float("-inf"),
+            )
+        else:
+            qk += tl.where(
+                (start_n + offs_n[None, :]) < cur_batch_seq_len, 0, float("-inf")
+            )
+
+        # -- compute m_ij, p, l_ij
+        m_ij = tl.max(qk, 1)
+        p = tl.exp(qk - m_ij[:, None])
+        l_ij = tl.sum(p, 1)
+        # -- update m_i and l_i
+        m_i_new = tl.maximum(m_i, m_ij)
+        alpha = tl.exp(m_i - m_i_new)
+        beta = tl.exp(m_ij - m_i_new)
+        l_i_new = alpha * l_i + beta * l_ij
+        # -- update output accumulator --
+        # scale p
+        p_scale = beta / l_i_new
+        p = p * p_scale[:, None]
+        # scale acc
+        acc_scale = l_i / l_i_new * alpha
+        acc = acc * acc_scale[:, None]
+        # update acc
+        v = tl.load(
+            v_ptrs + (cur_batch_in_all_start_index + start_n) * stride_vbs,
+            mask=((start_n + offs_n[:, None]) < cur_batch_seq_len) & (mask_d[None, :]),
+            other=0.0,
+        )
+
+        p = p.to(v.dtype)
+        acc += tl.dot(p, v)
+        # update m_i and l_i
+        l_i = l_i_new
+        m_i = m_i_new
+    # initialize pointers to output
+    off_o = (
+        (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs
+        + cur_head * stride_oh
+        + offs_d[None, :]
+    )
+    out_ptrs = Out + off_o
+    tl.store(
+        out_ptrs, acc, mask=(offs_m[:, None] < cur_batch_seq_len) & (mask_d[None, :])
+    )
+
+
+def context_attention_fwd(
+    q, k, v, o, b_start_loc, b_seq_len, max_input_len, is_causal=True, sm_scale=None
+):
+    """
+    q, k, v: [b * s, head, head_dim]
+    b_start_loc: [b]
+    b_seq_len: [b]
+    out: [b * s, head, head_dim]
+    sm_scale: softmax scale, defaults to 1/sqrt(head_dim)
+    """
+    if (_is_cuda or _is_hip) and CUDA_CAPABILITY[0] > 8:
+        BLOCK = 128
+    else:
+        BLOCK = 64
+
+    Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
+
+    if sm_scale is None:
+        sm_scale = 1.0 / (Lq**0.5)
+    batch, head = b_seq_len.shape[0], q.shape[1]
+    kv_group_num = q.shape[1] // k.shape[1]
+
+    grid = (batch, head, triton.cdiv(max_input_len, BLOCK))
+    num_warps = 4 if Lk <= 64 else 8
+
+    _fwd_kernel[grid](
+        q,
+        k,
+        v,
+        sm_scale,
+        b_start_loc,
+        b_seq_len,
+        o,
+        q.stride(0),
+        q.stride(1),
+        k.stride(0),
+        k.stride(1),
+        v.stride(0),
+        v.stride(1),
+        o.stride(0),
+        o.stride(1),
+        kv_group_num=kv_group_num,
+        BLOCK_M=BLOCK,
+        BLOCK_DMODEL=triton.next_power_of_2(Lk),
+        BLOCK_N=BLOCK,
+        IS_CAUSAL=is_causal,
+        num_warps=num_warps,
+        num_stages=1,
+        Lk=Lk,
+    )

python/sglang/kernels/ops/attention/rocm_mla_decode_rope.py

@@ -0,0 +1,439 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""
+Memory-efficient attention for decoding.
+It supports page size = 1.
+"""
+
+# Adapted from
+# https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage1.py
+# https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage2.py
+
+import triton
+import triton.language as tl
+
+from sglang.kernels.ops.attention.decode_attention import (
+    _decode_softmax_reducev_fwd,
+)
+
+
+def is_hip():
+    return triton.runtime.driver.active.get_current_target().backend == "hip"
+
+
+_is_hip = is_hip()
+
+
+@triton.jit
+def tanh(x):
+    # Tanh is just a scaled sigmoid
+    return 2 * tl.sigmoid(2 * x) - 1
+
+
+@triton.jit
+def _fwd_grouped_kernel_stage1_rope(
+    Q,  # Holds [Q_NOPE; Q_PE], b x h x (d+r)
+    K_Buffer,  # Holds [KV; K_PE], b*s x (c+r)
+    V_buffer,  # Holds [KV], b*s x (c)
+    cos_sin_cache,  # max_seq_len x (rotary_dim * 2)
+    positions,  # sequence positions
+    sm_scale,
+    kv_indptr,
+    kv_indices,
+    Att_Out,  # b x h x NUM_KV_SPLITS x (kv_lora_rank + 1)
+    k_pe_t_out,
+    stride_qb,
+    stride_qh,
+    stride_buf_kbs,
+    stride_buf_vbs,
+    stride_mid_ob,
+    stride_mid_oh,
+    stride_mid_os,
+    stride_kpe_tokens_out_b,
+    stride_cos_sin_cache_s,
+    stride_positions_b,
+    rotary_dim: tl.constexpr,
+    kv_lora_rank: tl.constexpr,
+    qk_rope_head_dim: tl.constexpr,
+    kv_group_num: tl.constexpr,
+    q_head_num: tl.constexpr,
+    BLOCK_C: tl.constexpr,
+    BLOCK_R: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    NUM_KV_SPLITS: tl.constexpr,
+    logit_cap: tl.constexpr,
+    USE_ROPE: tl.constexpr,
+    IS_NEOX_STYLE: tl.constexpr,
+):
+
+    cur_batch = tl.program_id(0)
+    cur_head_id = tl.program_id(1)
+    split_kv_id = tl.program_id(2)
+
+    if BLOCK_H < kv_group_num:
+        VALID_BLOCK_H: tl.constexpr = BLOCK_H
+    else:
+        VALID_BLOCK_H: tl.constexpr = kv_group_num
+    cur_head = cur_head_id * VALID_BLOCK_H + tl.arange(0, BLOCK_H)
+    mask_h = cur_head < (cur_head_id + 1) * VALID_BLOCK_H
+    mask_h = mask_h & (cur_head < q_head_num)
+
+    offs_c = tl.arange(0, BLOCK_C)
+    offs_qk_r = tl.arange(kv_lora_rank, kv_lora_rank + BLOCK_R)  # to get the k_pe
+
+    off_q_pe = (
+        cur_batch * stride_qb + cur_head[:, None] * stride_qh + offs_qk_r[None, :]
+    )
+    offs_q = cur_batch * stride_qb + cur_head[:, None] * stride_qh + offs_c[None, :]
+
+    mask_c = offs_c < kv_lora_rank
+    mask_qk_r = offs_qk_r < (kv_lora_rank + qk_rope_head_dim)
+
+    cur_batch_kv_start_idx = tl.load(kv_indptr + cur_batch)
+    cur_batch_seq_len = tl.load(kv_indptr + cur_batch + 1) - cur_batch_kv_start_idx
+
+    q = tl.load(Q + offs_q, mask=(mask_h[:, None]) & (mask_c[None, :]), other=0.0)
+    q_pe = tl.load(
+        Q + off_q_pe, mask=(mask_h[:, None]) & (mask_qk_r[None, :]), other=0.0
+    )
+
+    kv_len_per_split = tl.cdiv(cur_batch_seq_len, NUM_KV_SPLITS)
+    split_kv_start = kv_len_per_split * split_kv_id
+    split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len)
+
+    # apply rotary embedding for q_pe, and k_pe (last token per batch of K_PE)
+    LAST_SPLIT = split_kv_end == cur_batch_seq_len
+    k_pe_last_token = tl.zeros([BLOCK_R], dtype=q.dtype)
+
+    if USE_ROPE:
+        if IS_NEOX_STYLE:
+            # [BLOCK_ROTARY // 2, BLOCK_ROTARY // 2 + 1, BLOCK_ROTARY // 2 + 2, ..., 0, 1, 2, ..., BLOCK_ROTARY // 2 - 1, pass:]
+            offs_qk_rot_r = kv_lora_rank + (
+                (tl.arange(0, BLOCK_R) + (rotary_dim // 2)) % rotary_dim
+            )
+            # Which elements to flip
+            mask_rotate = tl.arange(0, BLOCK_R) < (rotary_dim // 2)
+            # [0 , 1, 2, ..., rotary_dim // 2 - 1, 0 , 1, 2, ..., rotary_dim // 2 - 1]
+            offs_rotary = tl.arange(0, BLOCK_R) % (rotary_dim // 2)
+        else:
+            # [1, 0, 3, 2, 5, 4, ..., BLOCK_R, BLOCK_R - 1]
+            offs_qk_rot_r = (
+                kv_lora_rank
+                + (((tl.arange(0, BLOCK_R) + 1) % 2) * 2)
+                - 1
+                + tl.arange(0, BLOCK_R)
+            )
+            mask_rotate = tl.arange(0, BLOCK_R) % 2 < 1
+            # [0, 0, 1, 1, ..., rotary_dim // 2 - 1, rotary_dim // 2 - 1]
+            offs_rotary = tl.arange(0, BLOCK_R) // 2
+
+        if qk_rope_head_dim > rotary_dim:
+            offs_qk_rot_r = tl.where(
+                tl.arange(0, BLOCK_R) < rotary_dim, offs_qk_rot_r, tl.arange(0, BLOCK_R)
+            )
+            offs_rotary = tl.where(
+                tl.arange(0, BLOCK_R) < rotary_dim, offs_rotary, tl.arange(0, BLOCK_R)
+            )
+
+        mask_rotary = tl.arange(0, BLOCK_R) < rotary_dim
+
+        pos = tl.load(positions + cur_batch * stride_positions_b)
+        cos = tl.load(
+            cos_sin_cache + pos * stride_cos_sin_cache_s + offs_rotary,
+            mask=mask_rotary,
+            other=1.0,
+        )
+        sin = tl.load(
+            cos_sin_cache
+            + pos * stride_cos_sin_cache_s
+            + offs_rotary
+            + rotary_dim // 2,
+            mask_rotary,
+            other=0.0,
+        )
+
+        off_q_pe_rot = (
+            cur_batch * stride_qb
+            + cur_head[:, None] * stride_qh
+            + offs_qk_rot_r[None, :]
+        )
+        mask_qk_rot_r = offs_qk_rot_r < (kv_lora_rank + qk_rope_head_dim)
+
+        # 0, 2, 4,.... 1, 3, 5...
+        q_pe_rot = tl.load(
+            Q + off_q_pe_rot,
+            mask=(mask_h[:, None]) & (mask_qk_rot_r[None, :]),
+            other=0.0,
+        )
+        q_pe_rot = tl.where(mask_rotate[None, :], -q_pe_rot, q_pe_rot)
+
+        q_pe = q_pe * cos + q_pe_rot * sin
+
+        # we only apply to the last token in the K_PE
+        if LAST_SPLIT:
+            # debug assert
+            if (cur_batch == 0 and cur_head == 0) and split_kv_id < NUM_KV_SPLITS - 1:
+                tl.device_assert(False, "Only last split should compute k_pe")
+
+            kv_loc = tl.load(
+                kv_indices + cur_batch_kv_start_idx + cur_batch_seq_len - 1
+            )
+            offs_buf_k_pe_last_token = kv_loc * stride_buf_kbs + offs_qk_r
+            offs_buf_k_pe_rot_last_token = kv_loc * stride_buf_kbs + offs_qk_rot_r
+            k_pe_last_token = tl.load(K_Buffer + offs_buf_k_pe_last_token)
+
+            k_pe_rot_last_token = tl.load(K_Buffer + offs_buf_k_pe_rot_last_token)
+            k_pe_rot_last_token = tl.where(
+                mask_rotate, -k_pe_rot_last_token, k_pe_rot_last_token
+            )
+
+            k_pe_last_token = k_pe_last_token * cos + k_pe_rot_last_token * sin
+
+    e_max = tl.zeros([BLOCK_H], dtype=tl.float32) - float("inf")
+    e_sum = tl.zeros([BLOCK_H], dtype=tl.float32)
+    acc = tl.zeros([BLOCK_H, BLOCK_C], dtype=tl.float32)
+
+    if split_kv_end > split_kv_start:
+        for start_n in range(split_kv_start, split_kv_end, BLOCK_N):
+            offs_n = start_n + tl.arange(0, BLOCK_N)
+            kv_loc = tl.load(
+                kv_indices + cur_batch_kv_start_idx + offs_n,
+                mask=offs_n < split_kv_end,
+                other=0,
+            )
+
+            offs_buf_kv = kv_loc[None, :] * stride_buf_kbs + offs_c[:, None]
+            offs_buf_k_pe = kv_loc[None, :] * stride_buf_kbs + offs_qk_r[:, None]
+
+            k_pe = tl.load(
+                K_Buffer + offs_buf_k_pe,
+                mask=(offs_n[None, :] < split_kv_end) & (mask_qk_r[:, None]),
+                other=0.0,
+            )  # positional embedding part of keys
+
+            if (USE_ROPE and LAST_SPLIT) and start_n >= cur_batch_seq_len - BLOCK_N:
+                k_pe = tl.where(
+                    offs_n[None, :] != (split_kv_end - 1),
+                    k_pe,
+                    k_pe_last_token[:, None],
+                )
+
+            # (16, 64) x (64, 32)
+            # dot product of rope parts
+            qk = tl.dot(q_pe, k_pe.to(q_pe.dtype))
+
+            kv = tl.load(
+                K_Buffer + offs_buf_kv,
+                mask=(offs_n[None, :] < split_kv_end) & (mask_c[:, None]),
+                other=0.0,
+            )  # the shared latent tensor for keys and values
+
+            # (16, 512) x (512, 32)
+            # dot product of nope parts
+            qk += tl.dot(q, kv)
+
+            qk *= sm_scale
+
+            if logit_cap > 0:
+                qk = logit_cap * tanh(qk / logit_cap)
+
+            qk = tl.where(
+                mask_h[:, None] & (offs_n[None, :] < split_kv_end), qk, float("-inf")
+            )
+
+            offs_buf_v = kv_loc[:, None] * stride_buf_vbs + offs_c[None, :]
+            v = tl.load(
+                V_buffer + offs_buf_v,
+                mask=(offs_n[:, None] < split_kv_end) & (mask_c[None, :]),
+                other=0.0,
+            )
+
+            n_e_max = tl.maximum(tl.max(qk, 1), e_max)
+            re_scale = tl.exp(e_max - n_e_max)
+            p = tl.exp(qk - n_e_max[:, None])
+            acc *= re_scale[:, None]
+            # (16, 32) x (32, 512)
+            acc += tl.dot(p.to(v.dtype), v)
+
+            e_sum = e_sum * re_scale + tl.sum(p, 1)
+            e_max = n_e_max
+
+        offs_mid_o = (
+            cur_batch * stride_mid_ob
+            + cur_head[:, None] * stride_mid_oh
+            + split_kv_id * stride_mid_os
+            + offs_c[None, :]
+        )
+
+        if USE_ROPE:
+            if LAST_SPLIT:
+                k_pe_last_token_ptrs = (
+                    k_pe_t_out
+                    + cur_batch * stride_kpe_tokens_out_b
+                    + tl.arange(0, BLOCK_R)
+                )
+                tl.store(k_pe_last_token_ptrs, k_pe_last_token, mask=mask_qk_r)
+
+        tl.store(
+            Att_Out + offs_mid_o,
+            acc / e_sum[:, None],
+            mask=(mask_h[:, None]) & (mask_c[None, :]),
+        )
+
+        offs_mid_o_1 = (
+            cur_batch * stride_mid_ob
+            + cur_head * stride_mid_oh
+            + split_kv_id * stride_mid_os
+            + kv_lora_rank
+        )
+
+        tl.store(
+            Att_Out + offs_mid_o_1,
+            e_max + tl.log(e_sum),
+            mask=mask_h,
+        )
+
+
+# TODO rope offset
+def _decode_grouped_att_m_fwd_rope(
+    q,
+    k_buffer,
+    v_buffer,
+    att_out,
+    k_pe_tokens_out,
+    kv_lora_rank,  # c
+    cos_sin_cache,
+    positions,
+    rotary_dim,
+    kv_indptr,
+    kv_indices,
+    num_kv_splits,
+    sm_scale,
+    logit_cap,
+    use_rope,
+    is_neox_style=True,
+):
+    if use_rope:
+        assert (
+            k_pe_tokens_out is not None
+        ), "We must output the k_pe tokens with rope applied if rope fusion enabled."
+
+    BLOCK = 32
+
+    # # [TODO] work around shmem limit on MI3xx
+    # if _is_hip and kv_lora_rank >= 576:
+    #     BLOCK = 16
+
+    qk_rope_head_dim = k_buffer.shape[-1] - kv_lora_rank
+    batch, head_num = kv_indptr.shape[0] - 1, q.shape[1]
+    kv_group_num = q.shape[1] // k_buffer.shape[1]
+
+    BLOCK_C = triton.next_power_of_2(kv_lora_rank)
+    BLOCK_R = triton.next_power_of_2(qk_rope_head_dim)
+
+    BLOCK_H = 16
+    NUM_KV_SPLITS = num_kv_splits
+    grid = (
+        batch,
+        triton.cdiv(head_num, min(BLOCK_H, kv_group_num)),
+        NUM_KV_SPLITS,
+    )
+
+    extra_kargs = {}
+    num_stages = 2
+    if _is_hip:
+        # https://rocm.docs.amd.com/en/docs-6.2.0/how-to/llm-fine-tuning-optimization/optimizing-triton-kernel.html
+        # https://github.com/triton-lang/triton/blob/main/third_party/amd/backend/compiler.py
+        extra_kargs = {"waves_per_eu": 1, "matrix_instr_nonkdim": 16, "kpack": 2}
+        num_stages = 1
+
+    _fwd_grouped_kernel_stage1_rope[grid](
+        q,
+        k_buffer,
+        v_buffer,
+        cos_sin_cache,
+        positions,
+        sm_scale,
+        kv_indptr,
+        kv_indices,
+        att_out,
+        k_pe_tokens_out,
+        q.stride(0),
+        q.stride(1),
+        k_buffer.stride(0),
+        v_buffer.stride(0),
+        att_out.stride(0),
+        att_out.stride(1),
+        att_out.stride(2),
+        k_pe_tokens_out.stride(0) if use_rope else 0,
+        cos_sin_cache.stride(0) if use_rope else 0,
+        positions.stride(0) if use_rope else 0,
+        rotary_dim,
+        kv_lora_rank,
+        qk_rope_head_dim,
+        kv_group_num=kv_group_num,
+        q_head_num=head_num,
+        BLOCK_C=BLOCK_C,
+        BLOCK_R=BLOCK_R,
+        BLOCK_N=BLOCK,
+        BLOCK_H=BLOCK_H,
+        NUM_KV_SPLITS=NUM_KV_SPLITS,
+        logit_cap=logit_cap,
+        USE_ROPE=use_rope,
+        IS_NEOX_STYLE=is_neox_style,
+        num_warps=4,
+        num_stages=num_stages,
+        **extra_kargs,
+    )
+
+
+def decode_attention_fwd_grouped_rope(
+    q,
+    k_buffer,
+    v_buffer,
+    o,
+    kv_indptr,
+    kv_indices,
+    k_pe_tokens,
+    kv_lora_rank,
+    rotary_dim,
+    cos_sin_cache,
+    positions,
+    attn_logits,
+    num_kv_splits,
+    sm_scale,
+    logit_cap=0.0,
+    use_rope=False,
+    is_neox_style=False,
+):
+    _decode_grouped_att_m_fwd_rope(
+        q,
+        k_buffer,
+        v_buffer,
+        attn_logits,
+        k_pe_tokens,
+        kv_lora_rank,
+        cos_sin_cache,
+        positions,
+        rotary_dim,
+        kv_indptr,
+        kv_indices,
+        num_kv_splits,
+        sm_scale,
+        logit_cap,
+        use_rope,
+        is_neox_style,
+    )
+    _decode_softmax_reducev_fwd(attn_logits, q, o, v_buffer, kv_indptr, num_kv_splits)

python/sglang/kernels/ops/attention/verify_splitkv.py

@@ -0,0 +1,803 @@
+"""Split-KV (flash-decode) attention for EAGLE speculative *verify*.
+
+Only valid when speculative ``topk == 1`` (the EAGLE tree reduces to a pure
+causal chain); the caller gates on that. ``topk > 1`` trees fall back to
+``extend_attention_fwd``.
+
+On the Triton backend, EAGLE target-verify runs through the prefill
+``extend_attention_fwd``, which loops the (long) prefix KV serially per
+(sequence, head). With only a few draft-token queries, that leaves the GPU
+memory system far under-utilized at long context. This kernel instead splits
+the prefix KV across parallel programs (flash-decode style) and combines the
+partials with a log-sum-exp merge, then handles the small causal draft-draft
+block -- recovering memory bandwidth on the verify path.
+
+Two Triton kernels:
+  * ``_verify_prefix_stage1``: split-KV over the shared prefix. Applies the fp8
+    dequant multipliers ``k_scale`` (on the QK score) and ``v_scale`` (on the
+    prefix output), matching ``extend_attention_fwd``'s ``_fwd_kernel``
+    (qk *= sm_scale * k_scale; acc += dot(p, v) * v_scale on the prefix loop;
+    NO scaling on the draft-draft loop, whose K/V are the fresh bf16 draft
+    tensors, not the fp8 pool). fp8 K/V buffers are handled by casting q to the
+    buffer dtype before the dot (mirrors ``q.to(k.dtype)`` in the baseline).
+  * ``_verify_combine_stage2``: combines the prefix splits (LSE merge) with the
+    small causal draft-draft block and writes the output.
+
+``verify_splitkv_fwd(...)`` takes the SAME positional args as
+``extend_attention_fwd``; it runs the split-KV path when it can serve the case
+bit-equivalently and returns True, otherwise returns False (doing nothing) so
+the caller falls back to ``extend_attention_fwd``. Supported case: causal
+(topk=1) verify with a constant per-sequence extend length, no sinks /
+sliding-window / logit-cap / xai-temperature. Correctness is never violated.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.utils import is_hip
+
+_MIN_BLOCK_KV = 32
+
+# AMD/CDNA-only Triton launch hints (waves_per_eu, matrix_instr_nonkdim); NVIDIA's
+# Triton rejects these kwargs, so only pass them on ROCm. In production this kernel
+# is dispatched only on AMD (see TritonAttnBackend); keeping it NV-safe lets the
+# numerics test run on the CUDA CI lane.
+_IS_HIP = is_hip()
+_AMD_LAUNCH_KWARGS = {"waves_per_eu": 4, "matrix_instr_nonkdim": 16} if _IS_HIP else {}
+
+# Block-size config keyed on head_dim. The (BLOCK_N, num_warps) tile that best
+# hides latency depends on head_dim: at head_dim=256 (Qwen3 family) a narrower
+# BLOCK_N with more warps wins, since the 256-wide QK/PV tiles are register
+# heavy. head_dim=256 is the value validated on MI350X; other head dims use a
+# conservative default. Block size affects PERFORMANCE only, never correctness
+# (any valid block size produces the same result).
+DEFAULT_N_SPLITS = 8
+DEFAULT_BLOCK_N = 32
+DEFAULT_NUM_WARPS = 4
+_BLOCK_CONFIG = {
+    # head_dim: (BLOCK_N, num_warps)
+    256: (32, 4),
+}
+
+
+def block_config(head_dim):
+    """Return (BLOCK_N, num_warps) for a head_dim; default for untuned dims."""
+    return _BLOCK_CONFIG.get(head_dim, (DEFAULT_BLOCK_N, DEFAULT_NUM_WARPS))
+
+
+# ---------------------------------------------------------------------------
+# Adaptive N_SPLITS.
+# ---------------------------------------------------------------------------
+# The prefix split-KV stage launches a (bs, h_q, N_SPLITS) grid; each (b,h,s)
+# program handles kv_len_per_split = cdiv(cdiv(seqlen, N_SPLITS), MIN)*MIN keys.
+# A fixed N_SPLITS=16 over-splits short/mid contexts (each split does too little
+# work -> launch + reduction overhead dominates) and under-splits very long ones
+# (too few parallel waves to saturate the device, raising tail latency on the
+# slow split). Mirror the decode kernel's intent (decode_attention.py
+# get_num_kv_splits): pick the split count per-dispatch from the representative
+# sequence length, growing gradually with seqlen and capped at MAX.
+#
+# CRITICAL: this must be computed from STATIC shapes only (no .item()/.cpu()
+# sync), because the verify/draft-extend step runs inside a captured HIP graph
+# where a device->host copy raises hipErrorStreamCaptureUnsupported. We use the
+# average prefix length = kv_indices.shape[0] / bs, which is a pure python int
+# from tensor shapes -- no device read. N_SPLITS is then a power of two so the
+# stage2 reduction tile (tl.arange(0, N_SPLITS)) stays cheap.
+#
+# Split-count bounds (internal constants). MAX=16 is the MI350X cap: 32
+# oversubscribes the device and regresses, per tuning.
+ADAPTIVE_SPLITS = True
+MAX_N_SPLITS = 16
+MIN_N_SPLITS = 4
+
+
+def choose_n_splits(avg_seqlen):
+    """Pick N_SPLITS (power of two, in [MIN_N_SPLITS, MAX_N_SPLITS]) from the
+    average prefix length. Tuned by the real-shape sweep (head_dim=256, BS*H_Q
+    =128 base programs on ~132 CUs):
+
+        ctx  <  4k -> 4   (short: extra splits add launch/reduction overhead)
+        4k <= ctx < 8k -> 8   (sweet spot: best across 1k-16k in the sweep)
+        ctx >= 8k      -> 16  (long: a few more splits help latency-bound tail)
+
+    Never 32 (4096 grid blocks oversubscribes the device and regresses, per the
+    sweep). Computed from a static shape (avg prefix = kv_indices.shape[0]/bs),
+    so it is HIP-graph-capture safe (no device->host sync)."""
+    if not ADAPTIVE_SPLITS:
+        return DEFAULT_N_SPLITS
+    s = int(avg_seqlen)
+    if s < 4096:
+        n = 4
+    elif s < 8192:
+        n = 8
+    else:
+        n = 16
+    if n < MIN_N_SPLITS:
+        n = MIN_N_SPLITS
+    if n > MAX_N_SPLITS:
+        n = MAX_N_SPLITS
+    return n
+
+
+@triton.jit
+def _verify_prefix_stage1(
+    Q,  # [extend_tokens, H_Q, D]
+    K_Buffer,  # [pool_tokens, H_KV, D]
+    V_Buffer,  # [pool_tokens, H_KV, Dv]
+    sm_scale,
+    k_scale,  # fp8 dequant multiplier for prefix K (1.0 if bf16)
+    v_scale,  # fp8 dequant multiplier for prefix V (1.0 if bf16)
+    qo_indptr,  # [BS+1] int32  -> rows of Q (draft queries)
+    kv_indptr,  # [BS+1] int32  -> rows of kv_indices (prefix)
+    kv_indices,  # [sum prefix] int64
+    Att_Out,  # [BS, H_Q, N_SPLITS, L_EXT, Dv]  fp32
+    Att_Lse,  # [BS, H_Q, N_SPLITS, L_EXT]      fp32
+    stride_qbs,
+    stride_qh,
+    stride_buf_kbs,
+    stride_buf_kh,
+    stride_buf_vbs,
+    stride_buf_vh,
+    stride_ob,
+    stride_oh,
+    stride_os,
+    stride_ol,
+    stride_lb,
+    stride_lh,
+    stride_ls,
+    kv_group_num: tl.constexpr,
+    N_SPLITS: tl.constexpr,
+    L_EXT: tl.constexpr,  # padded power-of-2 row tile (>= real l_ext)
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_DV: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    MIN_BLOCK_KV: tl.constexpr,
+):
+    cur_batch = tl.program_id(0)
+    cur_head = tl.program_id(1)
+    split_kv_id = tl.program_id(2)
+
+    cur_kv_head = cur_head // kv_group_num
+
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    offs_dv = tl.arange(0, BLOCK_DV)
+    offs_l = tl.arange(0, L_EXT)
+
+    # real number of draft query tokens for this seq
+    cur_q_start = tl.load(qo_indptr + cur_batch)
+    l_ext = tl.load(qo_indptr + cur_batch + 1) - cur_q_start
+    mask_l = offs_l < l_ext
+
+    cur_batch_kv_start_idx = tl.load(kv_indptr + cur_batch)
+    cur_batch_seq_len = tl.load(kv_indptr + cur_batch + 1) - cur_batch_kv_start_idx
+
+    # split sizing identical to the decode kernel
+    kv_len_per_split = (
+        tl.cdiv(tl.cdiv(cur_batch_seq_len, N_SPLITS), MIN_BLOCK_KV) * MIN_BLOCK_KV
+    )
+    split_kv_start = kv_len_per_split * split_kv_id
+    split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len)
+
+    e_max = tl.zeros([L_EXT], dtype=tl.float32) - float("inf")
+    e_sum = tl.zeros([L_EXT], dtype=tl.float32)
+    acc = tl.zeros([L_EXT, BLOCK_DV], dtype=tl.float32)
+
+    if split_kv_end > split_kv_start:
+        # q tile: [L_EXT, D]
+        offs_q = (
+            (cur_q_start + offs_l)[:, None] * stride_qbs
+            + cur_head * stride_qh
+            + offs_d[None, :]
+        )
+        q = tl.load(Q + offs_q, mask=mask_l[:, None], other=0.0)
+        q_k = q.to(K_Buffer.dtype.element_ty)
+
+        base_offs_k = cur_kv_head * stride_buf_kh + offs_d[:, None]
+        base_offs_v = cur_kv_head * stride_buf_vh + offs_dv[None, :]
+
+        for start_n in tl.range(split_kv_start, split_kv_end, BLOCK_N):
+            offs_n = start_n + tl.arange(0, BLOCK_N)
+            n_mask = offs_n < split_kv_end
+            kv_loc = tl.load(
+                kv_indices + cur_batch_kv_start_idx + offs_n,
+                mask=n_mask,
+                other=0,
+            )
+            # K block: [D, BLOCK_N]
+            offs_buf_k = kv_loc[None, :] * stride_buf_kbs + base_offs_k
+            k = tl.load(K_Buffer + offs_buf_k, mask=n_mask[None, :], other=0.0)
+            qk = tl.dot(q_k, k)  # [L_EXT, BLOCK_N]
+            qk *= sm_scale * k_scale  # fp8 dequant of prefix K (k_scale==1 if bf16)
+            # NO causal mask: full prefix is visible to all draft tokens.
+            qk = tl.where(n_mask[None, :], qk, float("-inf"))
+
+            # V block: [BLOCK_N, Dv]
+            offs_buf_v = kv_loc[:, None] * stride_buf_vbs + base_offs_v
+            v = tl.load(V_Buffer + offs_buf_v, mask=n_mask[:, None], other=0.0)
+
+            n_e_max = tl.maximum(tl.max(qk, 1), e_max)
+            re_scale = tl.exp(e_max - n_e_max)
+            p = tl.exp(qk - n_e_max[:, None])
+            acc *= re_scale[:, None]
+            acc += tl.dot(p.to(v.dtype), v)
+            e_sum = e_sum * re_scale + tl.sum(p, 1)
+            e_max = n_e_max
+
+        # fp8 dequant of prefix V: scale the accumulated (pre-normalised) output.
+        acc *= v_scale
+
+        offs_o = (
+            cur_batch * stride_ob
+            + cur_head * stride_oh
+            + split_kv_id * stride_os
+            + offs_l[:, None] * stride_ol
+            + offs_dv[None, :]
+        )
+        tl.store(Att_Out + offs_o, acc / e_sum[:, None], mask=mask_l[:, None])
+
+        offs_lse = (
+            cur_batch * stride_lb
+            + cur_head * stride_lh
+            + split_kv_id * stride_ls
+            + offs_l
+        )
+        tl.store(Att_Lse + offs_lse, e_max + tl.log(e_sum), mask=mask_l)
+    else:
+        # split did not run: write a sentinel lse so stage2 can ignore it.
+        offs_lse = (
+            cur_batch * stride_lb
+            + cur_head * stride_lh
+            + split_kv_id * stride_ls
+            + offs_l
+        )
+        tl.store(
+            Att_Lse + offs_lse,
+            tl.zeros([L_EXT], tl.float32) - float("inf"),
+            mask=mask_l,
+        )
+
+
+@triton.jit
+def _verify_combine_stage2(
+    Att_Out,  # [BS, H_Q, N_SPLITS, L_EXT, Dv]  fp32
+    Att_Lse,  # [BS, H_Q, N_SPLITS, L_EXT]      fp32
+    Q,  # [extend_tokens, H_Q, D]   (draft queries)
+    K_Extend,  # [extend_tokens, H_KV, D]
+    V_Extend,  # [extend_tokens, H_KV, Dv]
+    O_Out,  # [extend_tokens, H_Q, Dv]  (final, written)
+    sm_scale,
+    qo_indptr,  # [BS+1] int32
+    stride_ob,
+    stride_oh,
+    stride_os,
+    stride_ol,
+    stride_lb,
+    stride_lh,
+    stride_ls,
+    stride_qbs,
+    stride_qh,
+    stride_kebs,
+    stride_keh,
+    stride_vebs,
+    stride_veh,
+    stride_oobs,
+    stride_ooh,
+    kv_group_num: tl.constexpr,
+    N_SPLITS: tl.constexpr,
+    L_EXT: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_DV: tl.constexpr,
+):
+    cur_batch = tl.program_id(0)
+    cur_head = tl.program_id(1)
+    cur_kv_head = cur_head // kv_group_num
+
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    offs_dv = tl.arange(0, BLOCK_DV)
+    offs_l = tl.arange(0, L_EXT)
+    offs_s = tl.arange(0, N_SPLITS)
+
+    cur_q_start = tl.load(qo_indptr + cur_batch)
+    l_ext = tl.load(qo_indptr + cur_batch + 1) - cur_q_start
+    mask_l = offs_l < l_ext
+
+    # ---- (a) combine prefix splits (logsumexp) ----------------------------
+    # lse: [N_SPLITS, L_EXT]
+    offs_lse = (
+        cur_batch * stride_lb
+        + cur_head * stride_lh
+        + offs_s[:, None] * stride_ls
+        + offs_l[None, :]
+    )
+    lse = tl.load(offs_lse + Att_Lse)  # [N_SPLITS, L_EXT]
+    m_p = tl.max(lse, 0)  # [L_EXT]
+    w = tl.exp(lse - m_p[None, :])  # [N_SPLITS, L_EXT]; -inf->0
+    denom_p = tl.sum(w, 0)  # [L_EXT]
+
+    # weighted-sum of partial outputs: o_prefix[L_EXT, Dv]
+    # Att_Out[b,h,s,l,dv]
+    offs_ao = (
+        cur_batch * stride_ob
+        + cur_head * stride_oh
+        + offs_s[:, None, None] * stride_os
+        + offs_l[None, :, None] * stride_ol
+        + offs_dv[None, None, :]
+    )
+    ao = tl.load(offs_ao + Att_Out)  # [N_SPLITS, L_EXT, Dv]
+    o_prefix = tl.sum(ao * w[:, :, None], 0)  # [L_EXT, Dv]
+    o_prefix = o_prefix / denom_p[:, None]
+    lse_prefix = m_p + tl.log(denom_p)  # [L_EXT]
+
+    # ---- (b) draft-draft causal attention (L_EXT x L_EXT) -----------------
+    # load draft queries [L_EXT, D], draft K/V [L_EXT, D]/[L_EXT, Dv]
+    offs_q = (
+        (cur_q_start + offs_l)[:, None] * stride_qbs
+        + cur_head * stride_qh
+        + offs_d[None, :]
+    )
+    q = tl.load(Q + offs_q, mask=mask_l[:, None], other=0.0).to(tl.float32)
+
+    offs_ke = (
+        (cur_q_start + offs_l)[:, None] * stride_kebs
+        + cur_kv_head * stride_keh
+        + offs_d[None, :]
+    )
+    ke = tl.load(K_Extend + offs_ke, mask=mask_l[:, None], other=0.0).to(tl.float32)
+    offs_ve = (
+        (cur_q_start + offs_l)[:, None] * stride_vebs
+        + cur_kv_head * stride_veh
+        + offs_dv[None, :]
+    )
+    ve = tl.load(V_Extend + offs_ve, mask=mask_l[:, None], other=0.0).to(tl.float32)
+
+    # scores[i,j] = q_i . k_j  (i query, j key)  -> [L_EXT, L_EXT]
+    qk = tl.sum(q[:, None, :] * ke[None, :, :], 2) * sm_scale
+    # causal among drafts: query i sees key j iff j <= i, and both valid
+    causal = (offs_l[None, :] <= offs_l[:, None]) & mask_l[None, :] & mask_l[:, None]
+    qk = tl.where(causal, qk, float("-inf"))
+    m_d = tl.max(qk, 1)  # [L_EXT]
+    pd = tl.exp(qk - m_d[:, None])  # [L_EXT, L_EXT]
+    denom_d = tl.sum(pd, 1)  # [L_EXT]
+    o_draft = tl.sum(pd[:, :, None] * ve[None, :, :], 1)  # [L_EXT, Dv]
+    o_draft = o_draft / denom_d[:, None]
+    lse_draft = m_d + tl.log(denom_d)  # [L_EXT]
+
+    # ---- (c) final LSE merge (prefix vs draft) ----------------------------
+    m = tl.maximum(lse_prefix, lse_draft)
+    wp = tl.exp(lse_prefix - m)
+    wd = tl.exp(lse_draft - m)
+    o = (o_prefix * wp[:, None] + o_draft * wd[:, None]) / (wp + wd)[:, None]
+
+    offs_oo = (
+        (cur_q_start + offs_l)[:, None] * stride_oobs
+        + cur_head * stride_ooh
+        + offs_dv[None, :]
+    )
+    tl.store(O_Out + offs_oo, o.to(O_Out.dtype.element_ty), mask=mask_l[:, None])
+
+
+class VerifySplitKV:
+    """Pre-allocates scratch buffers for a problem shape and runs the split-KV
+    verify attention end to end (two Triton launches: prefix split-KV + fused
+    combine/draft/merge). Buffers are sized by ``max_bs`` (constant for the
+    server lifetime) and reused for every batch size <= max_bs, so their
+    addresses stay fixed (CUDA/HIP-graph safe) and GPU memory does not grow per
+    batch size. The kernel grid uses the actual per-call bs (<= max_bs)."""
+
+    def __init__(
+        self,
+        max_bs,
+        h_q,
+        h_kv,
+        head_dim,
+        v_head_dim,
+        l_ext,
+        device="cuda",
+        n_splits=DEFAULT_N_SPLITS,
+        block_n=DEFAULT_BLOCK_N,
+        num_warps=DEFAULT_NUM_WARPS,
+    ):
+        self.h_q = h_q
+        self.h_kv = h_kv
+        self.group = h_q // h_kv
+        self.head_dim = head_dim
+        self.v_head_dim = v_head_dim
+        self.l_ext = l_ext  # real draft tokens per seq (fixed == 4)
+        self.l_pad = triton.next_power_of_2(l_ext)
+        self.device = device
+        self.n_splits = n_splits
+        self.block_n = block_n
+        self.num_warps = num_warps
+        self._alloc(max_bs)
+
+    def _alloc(self, max_bs):
+        # prefix split partials (fp32), sized for the maximum batch size.
+        self.max_bs = max_bs
+        self.att_out = torch.empty(
+            (max_bs, self.h_q, self.n_splits, self.l_pad, self.v_head_dim),
+            dtype=torch.float32,
+            device=self.device,
+        )
+        self.att_lse = torch.empty(
+            (max_bs, self.h_q, self.n_splits, self.l_pad),
+            dtype=torch.float32,
+            device=self.device,
+        )
+
+    def grow_buffers(self, max_bs):
+        if max_bs > self.max_bs:
+            self._alloc(max_bs)
+
+    def _run_prefix_kernel(
+        self,
+        bs,
+        q_extend,
+        k_buffer,
+        v_buffer,
+        qo_indptr,
+        kv_indptr,
+        kv_indices,
+        sm_scale,
+        k_scale,
+        v_scale,
+    ):
+        grid = (bs, self.h_q, self.n_splits)
+        _verify_prefix_stage1[grid](
+            q_extend,
+            k_buffer,
+            v_buffer,
+            sm_scale,
+            k_scale,
+            v_scale,
+            qo_indptr,
+            kv_indptr,
+            kv_indices,
+            self.att_out,
+            self.att_lse,
+            q_extend.stride(0),
+            q_extend.stride(1),
+            k_buffer.stride(0),
+            k_buffer.stride(1),
+            v_buffer.stride(0),
+            v_buffer.stride(1),
+            self.att_out.stride(0),
+            self.att_out.stride(1),
+            self.att_out.stride(2),
+            self.att_out.stride(3),
+            self.att_lse.stride(0),
+            self.att_lse.stride(1),
+            self.att_lse.stride(2),
+            kv_group_num=self.group,
+            N_SPLITS=self.n_splits,
+            L_EXT=self.l_pad,
+            BLOCK_DMODEL=triton.next_power_of_2(self.head_dim),
+            BLOCK_DV=triton.next_power_of_2(self.v_head_dim),
+            BLOCK_N=self.block_n,
+            MIN_BLOCK_KV=_MIN_BLOCK_KV,
+            num_warps=self.num_warps,
+            num_stages=1,
+            **_AMD_LAUNCH_KWARGS,
+        )
+
+    def _run_combine_kernel(
+        self, bs, q_extend, k_extend, v_extend, o_out, qo_indptr, sm_scale
+    ):
+        grid = (bs, self.h_q)
+        _verify_combine_stage2[grid](
+            self.att_out,
+            self.att_lse,
+            q_extend,
+            k_extend,
+            v_extend,
+            o_out,
+            sm_scale,
+            qo_indptr,
+            self.att_out.stride(0),
+            self.att_out.stride(1),
+            self.att_out.stride(2),
+            self.att_out.stride(3),
+            self.att_lse.stride(0),
+            self.att_lse.stride(1),
+            self.att_lse.stride(2),
+            q_extend.stride(0),
+            q_extend.stride(1),
+            k_extend.stride(0),
+            k_extend.stride(1),
+            v_extend.stride(0),
+            v_extend.stride(1),
+            o_out.stride(0),
+            o_out.stride(1),
+            kv_group_num=self.group,
+            N_SPLITS=self.n_splits,
+            L_EXT=self.l_pad,
+            BLOCK_DMODEL=triton.next_power_of_2(self.head_dim),
+            BLOCK_DV=triton.next_power_of_2(self.v_head_dim),
+            num_warps=1,
+            num_stages=1,
+        )
+
+    def __call__(
+        self,
+        q_extend,
+        k_extend,
+        v_extend,
+        k_buffer,
+        v_buffer,
+        qo_indptr,
+        kv_indptr,
+        kv_indices,
+        sm_scale,
+        o_out=None,
+        k_scale=1.0,
+        v_scale=1.0,
+    ):
+        if o_out is None:
+            o_out = torch.empty(
+                (q_extend.shape[0], self.h_q, self.v_head_dim),
+                dtype=q_extend.dtype,
+                device=q_extend.device,
+            )
+        # actual batch size for this call (<= max_bs); the grid uses it while the
+        # scratch buffers stay max_bs-sized (only the first bs slices are touched).
+        bs = qo_indptr.shape[0] - 1
+        # 1. prefix split-KV
+        self._run_prefix_kernel(
+            bs,
+            q_extend,
+            k_buffer,
+            v_buffer,
+            qo_indptr,
+            kv_indptr,
+            kv_indices,
+            sm_scale,
+            k_scale,
+            v_scale,
+        )
+        # 2+3+4. fused combine + draft-draft + merge
+        self._run_combine_kernel(
+            bs,
+            q_extend,
+            k_extend,
+            v_extend,
+            o_out,
+            qo_indptr,
+            sm_scale,
+        )
+        return o_out
+
+
+# ---------------------------------------------------------------------------
+# Live-server dispatch entry.
+# ---------------------------------------------------------------------------
+# Cache one VerifySplitKV instance per (h_q, h_kv, head_dim, v_head_dim, l_ext,
+# device, n_splits) shape -- NOT keyed on the dynamic batch size. Buffers are
+# sized by the stable max_bs (grown only if a larger one is ever requested), so
+# a single instance serves every batch size: addresses stay fixed (graph-safe)
+# and GPU memory does not grow per batch size.
+_VK_CACHE = {}
+
+
+def _get_vk(
+    max_bs, h_q, h_kv, head_dim, v_head_dim, l_ext, device, n_splits=DEFAULT_N_SPLITS
+):
+    key = (h_q, h_kv, head_dim, v_head_dim, l_ext, str(device), n_splits)
+    vk = _VK_CACHE.get(key)
+    if vk is None:
+        block_n, num_warps = block_config(head_dim)
+        vk = VerifySplitKV(
+            max_bs,
+            h_q,
+            h_kv,
+            head_dim,
+            v_head_dim,
+            l_ext,
+            device=device,
+            n_splits=n_splits,
+            block_n=block_n,
+            num_warps=num_warps,
+        )
+        _VK_CACHE[key] = vk
+    else:
+        vk.grow_buffers(max_bs)
+    return vk
+
+
+def can_handle(
+    q_extend,
+    k_extend,
+    v_extend,
+    k_buffer,
+    v_buffer,
+    qo_indptr,
+    kv_indptr,
+    kv_indices,
+    custom_mask,
+    is_causal,
+    mask_indptr,
+    max_len_extend,
+    sliding_window_size=-1,
+    sinks=None,
+    logit_cap=0.0,
+    xai_temperature_len=-1,
+):
+    """Return True iff the split-KV verify path can serve this exact problem
+    with the same result as extend_attention_fwd. Conservative: anything not
+    explicitly handled -> False -> caller falls back to the baseline.
+
+    IMPORTANT: ``custom_mask`` is intentionally NOT inspected (its values can't
+    be read inside a captured HIP graph without a host sync). The kernel always
+    computes pure-causal attention, which equals the tree mask ONLY at
+    speculative topk == 1. The caller therefore MUST gate enablement on topk == 1
+    (TritonAttnBackend does: ``use_verify_splitkv = ... and self.topk == 1``).
+    At topk > 1 the tree is not causal and this path must stay disabled."""
+    # No exotic features.
+    if sinks is not None:
+        return False
+    if sliding_window_size is not None and sliding_window_size > 0:
+        return False
+    if logit_cap and logit_cap > 0:
+        return False
+    if xai_temperature_len is not None and xai_temperature_len > 0:
+        return False
+    if not is_causal:
+        return False
+    # q layout must be [tokens, H_Q, D]; head dims handled by power-of-2 pad.
+    if q_extend.dim() != 3 or k_extend.dim() != 3 or v_extend.dim() != 3:
+        return False
+    # GQA group must divide evenly.
+    h_q = q_extend.shape[1]
+    h_kv = k_extend.shape[1]
+    if h_kv == 0 or h_q % h_kv != 0:
+        return False
+    # head dims must match buffers.
+    if k_buffer.shape[1] != h_kv or v_buffer.shape[1] != h_kv:
+        return False
+    if q_extend.shape[2] != k_extend.shape[2]:
+        return False
+    if q_extend.shape[2] != k_buffer.shape[2]:
+        return False
+    if v_extend.shape[2] != v_buffer.shape[2]:
+        return False
+    # NOTE: must NOT read any tensor *values* here (no .item()/.cpu()): the
+    # target-verify step runs inside a captured CUDA/HIP graph, where a
+    # device->host sync raises hipErrorStreamCaptureUnsupported. We therefore
+    # gate purely on static shapes/dtypes/python scalars.
+    bs = qo_indptr.shape[0] - 1
+    if bs < 1:
+        return False
+    # max_len_extend must be a known positive python int (it is the static
+    # server_args.speculative_num_draft_tokens for the verify path). For
+    # topk=1 the per-seq extend len is constant == num_draft_tokens ==
+    # max_len_extend by construction of qo_indptr (arange with that step), so
+    # the L_EXT row-tile mask is exactly right and the tree custom_mask equals
+    # causal -- no value inspection required.
+    try:
+        mle = int(max_len_extend)
+    except (TypeError, ValueError):
+        return False
+    if mle < 1:
+        return False
+    # The packed extend tensor must hold exactly bs * max_len_extend rows
+    # (constant extend len). This is a pure shape check (no sync) and rejects
+    # any ragged/variable-extend batch -> falls back to the baseline.
+    if q_extend.shape[0] != bs * mle:
+        return False
+    return True
+
+
+def verify_splitkv_fwd(
+    q_extend,
+    k_extend,
+    v_extend,
+    o_extend,
+    k_buffer,
+    v_buffer,
+    qo_indptr,
+    kv_indptr,
+    kv_indices,
+    custom_mask,
+    is_causal,
+    mask_indptr,
+    max_len_extend,
+    k_scale,
+    v_scale,
+    sm_scale=None,
+    logit_cap=0.0,
+    skip_prefix_custom_mask=True,
+    sliding_window_size=-1,
+    sinks=None,
+    window_kv_offsets=None,
+    xai_temperature_len=-1,
+    max_bs=None,
+):
+    """Drop-in for extend_attention_fwd on the EAGLE target-verify (topk=1)
+    shape. Returns True if it ran (o_extend written), False if the case is
+    unsupported and the caller must fall back to extend_attention_fwd.
+
+    ``max_bs`` (optional) is the stable maximum batch size used to size the
+    cached scratch buffers; the backend passes its req_to_token_pool size. If
+    omitted it defaults to this call's bs.
+
+    Arg order mirrors extend_attention_fwd exactly so the call site is a
+    one-line swap.
+    """
+    if not can_handle(
+        q_extend,
+        k_extend,
+        v_extend,
+        k_buffer,
+        v_buffer,
+        qo_indptr,
+        kv_indptr,
+        kv_indices,
+        custom_mask,
+        is_causal,
+        mask_indptr,
+        max_len_extend,
+        sliding_window_size=sliding_window_size,
+        sinks=sinks,
+        logit_cap=logit_cap,
+        xai_temperature_len=xai_temperature_len,
+    ):
+        return False
+
+    bs = qo_indptr.shape[0] - 1
+    h_q = q_extend.shape[1]
+    h_kv = k_extend.shape[1]
+    head_dim = q_extend.shape[2]
+    v_head_dim = v_extend.shape[2]
+    l_ext = int(max_len_extend)
+
+    if sm_scale is None:
+        sm_scale = 1.0 / (head_dim**0.5)
+    # k_scale/v_scale may be float or 0-d tensor; coerce to python float.
+    try:
+        k_scale = float(k_scale)
+    except (TypeError, ValueError):
+        k_scale = 1.0
+    try:
+        v_scale = float(v_scale)
+    except (TypeError, ValueError):
+        v_scale = 1.0
+
+    # Adaptive split count from the average prefix length. This is a
+    # pure-shape derivation (kv_indices.shape[0] / bs) -- no device->host sync,
+    # so it is safe inside a captured HIP graph. The whole batch shares one
+    # N_SPLITS (the grid dim must be a launch constexpr); the per-split kernel
+    # logic still clamps each split's [start,end) to that seq's real length, so
+    # mixed-length batches stay correct -- shorter seqs simply write fewer
+    # active splits (the rest emit the -inf lse sentinel, ignored in stage2).
+    avg_seqlen = kv_indices.shape[0] / max(1, bs)
+    n_splits = choose_n_splits(avg_seqlen)
+
+    # Size scratch by the stable max_bs (backend passes req_to_token_pool size);
+    # fall back to this call's bs if not provided / smaller.
+    if max_bs is None or max_bs < bs:
+        max_bs = bs
+    vk = _get_vk(
+        max_bs,
+        h_q,
+        h_kv,
+        head_dim,
+        v_head_dim,
+        l_ext,
+        q_extend.device,
+        n_splits=n_splits,
+    )
+    vk(
+        q_extend,
+        k_extend.contiguous(),
+        v_extend.contiguous(),
+        k_buffer,
+        v_buffer,
+        qo_indptr,
+        kv_indptr,
+        kv_indices,
+        sm_scale,
+        o_out=o_extend,
+        k_scale=k_scale,
+        v_scale=v_scale,
+    )
+
+    return True