[AMD] Optimize Kimi-K3 Triton MLA prefill on gfx950 (#35770)

Co-authored-by: clintg6 <7388379+clintg6@users.noreply.github.com>
Co-authored-by: Thomas Wang <thomawan@amd.com>
Co-authored-by: HAI <hixiao@gmail.com>
Co-authored-by: kk <43161300+kkHuang-amd@users.noreply.github.com>

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

@@ -16,14 +16,15 @@ Memory-efficient attention for prefill.
 It supports page size = 1 and prefill with KV cache (i.e. extend).
 """
 
+import math
+from typing import Optional
+
 import torch
 import triton
 import triton.language as tl
 
 from sglang.kernels.ops.attention.decode_attention import _extract_kv_strides
-from sglang.kernels.ops.attention.prefill_attention import (
+from sglang.kernels.ops.attention.prefill_attention import context_attention_fwd
-    context_attention_fwd,
-)
 from sglang.kernels.ops.attention.score_mod import unpack_aux_tensors
 from sglang.srt.environ import envs
 from sglang.srt.utils import (
@@ -86,11 +87,12 @@ def _get_block_sizes_for_extend_attention(Lq: int, Lv: int):
             # late-prefill kernel from ~12.57 ms to ~5.24 ms.
             BLOCK_M, BLOCK_N = (64, 32)
             num_warps = 4
-        elif _is_gfx95 and Lq <= 256:
+        elif _is_gfx95 and 128 < Lq <= 256:
-            # gfx950 (CDNA4), head_dim <= 256: every workgroup streams the whole
+            # gfx950 (CDNA4), 128 < head_dim <= 256: a larger query tile halves KV bytes
-            # prefix, so a larger query tile halves the KV bytes read per call;
+            # streamed per call (each workgroup reads the whole prefix); 8 warps
-            # BLOCK_M / num_warps = 16 rows per warp is exactly one MFMA tile at
+            # hide the loads. Measured on MI350X head_dim 256: -36% kernel time,
-            # matrix_instr_nonkdim=16. Measured on MI350X at head_dim 64, 128, 256.
+            # 28% -> 44% MFU, numerically equivalent (BLOCK_N reduction order
+            # unchanged). Other AMD archs / head dims keep the default below.
             BLOCK_M, BLOCK_N = (128, 64)
             num_warps = 8
         else:
@@ -147,6 +149,21 @@ def _get_block_sizes_for_extend_attention(Lq: int, Lv: int):
     return BLOCK_DMODEL, BLOCK_DPE, BLOCK_DV, BLOCK_M, BLOCK_N, num_warps
 
 
+def _get_num_stages_for_extend_attention(
+    Lq: int, Lv: int, block_n: int | None = None
+) -> int:
+    if _is_gfx95 and Lq == 192 and Lv == 128:
+        return 2
+    if (
+        _is_gfx95
+        and Lq == 576
+        and Lv == 512
+        and (block_n == 32 or (block_n is None and _is_triton_ge_37))
+    ):
+        return 2
+    return 1
+
+
 def _compact_extend_q_tiles_per_head(
     *,
     batch_size: int,
@@ -356,6 +373,7 @@ def _fwd_kernel(
     BLOCK_DV: tl.constexpr,
     BLOCK_M: tl.constexpr,
     BLOCK_N: tl.constexpr,
+    BLOCK_N_PREFIX: tl.constexpr,
     USE_CUSTOM_MASK: tl.constexpr,
     IS_CAUSAL: tl.constexpr,
     SKIP_PREFIX_CUSTOM_MASK: tl.constexpr,
@@ -364,9 +382,14 @@ def _fwd_kernel(
     SKIP_EXTEND: tl.constexpr,
     STORE_TRANSPOSE: tl.constexpr,
     HAS_SINK: tl.constexpr,
+    USE_COMPACT_TILE_GRID: tl.constexpr,
+    USE_EXP2: tl.constexpr,
+    USE_FP8_PREFIX: tl.constexpr,
+    USE_FP8_EXTEND: tl.constexpr,
+    FP8_MAX: tl.constexpr,
     IS_GFX1250: tl.constexpr = False,
-    USE_COMPACT_TILE_GRID: tl.constexpr = False,
     PAGE_SIZE: tl.constexpr = 1,
+    IDENTITY_KV_INDICES: tl.constexpr = False,
     SCORE_MOD: tl.constexpr = None,
     Aux0=None,
     aux0_stride_t=0,
@@ -401,6 +424,8 @@ def _fwd_kernel(
         cur_head = tl.program_id(1)
         cur_block_m = tl.program_id(2)
     cur_kv_head = cur_head // kv_group_num
+    LOG2E: tl.constexpr = 1.4426950408889634
+    LN2: tl.constexpr = 0.6931471805599453
 
     cur_seq_extend_start_idx = tl.load(qo_indptr + cur_seq)
     cur_seq_len_extend = tl.load(qo_indptr + cur_seq + 1) - cur_seq_extend_start_idx
@@ -459,15 +484,17 @@ def _fwd_kernel(
 
     # stage 1: compute scores with prefix
     offs_n = tl.arange(0, BLOCK_N)
+    # The FP8 prefix sweep can use a wider tile than the current-token sweep.
+    offs_n_prefix = tl.arange(0, BLOCK_N_PREFIX)
 
     acc = tl.zeros([BLOCK_M, BLOCK_DV], dtype=tl.float32)
     deno = tl.zeros([BLOCK_M], dtype=tl.float32)
     e_max = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
 
     prefix_end = 0 if SKIP_PREFIX else cur_seq_len_prefix
-    for start_n in range(0, prefix_end, BLOCK_N):
+    for start_n in range(0, prefix_end, BLOCK_N_PREFIX):
-        start_n = tl.multiple_of(start_n, BLOCK_N)
+        start_n = tl.multiple_of(start_n, BLOCK_N_PREFIX)
-        mask_n = (start_n + offs_n) < cur_seq_len_prefix
+        mask_n = (start_n + offs_n_prefix) < cur_seq_len_prefix
 
         final_mask = mask_m[:, None] & mask_n[None, :]
         if USE_CUSTOM_MASK and not SKIP_PREFIX_CUSTOM_MASK:
@@ -478,7 +505,7 @@ def _fwd_kernel(
                 * (cur_seq_len + window_kv_offset)
                 + window_kv_offset
                 + start_n
-                + offs_n[None, :],
+                + offs_n_prefix[None, :],
                 mask=(mask_m[:, None] & mask_n[None, :]),
                 other=0,
             )
@@ -488,7 +515,7 @@ def _fwd_kernel(
             # q_id = prefix_len + cur_m, kv_id = cur_n
             window_mask = (
                 cur_seq_len_prefix + cur_block_m * BLOCK_M + offs_m[:, None]
-            ) <= (start_n + offs_n[None, :] + SLIDING_WINDOW_SIZE)
+            ) <= (start_n + offs_n_prefix[None, :] + SLIDING_WINDOW_SIZE)
             final_mask &= window_mask
 
         SKIP_TILE = False
@@ -496,11 +523,14 @@ def _fwd_kernel(
             SKIP_TILE = tl.max(tl.max(final_mask.to(tl.int32), axis=1), axis=0) == 0
 
         if not SKIP_TILE:
-            offs_kv_loc = tl.load(
+            if IDENTITY_KV_INDICES:
-                kv_indices + cur_seq_kv_start_idx + start_n + offs_n,
+                offs_kv_loc = cur_seq_kv_start_idx + start_n + offs_n_prefix
-                mask=mask_n,
+            else:
-                other=0,
+                offs_kv_loc = tl.load(
-            )
+                    kv_indices + cur_seq_kv_start_idx + start_n + offs_n_prefix,
+                    mask=mask_n,
+                    other=0,
+                )
 
             # Page-aware KV address math. At PAGE_SIZE==1
             # (legacy / non-shared / shared-at-ps=1), Triton specializes
@@ -560,7 +590,10 @@ def _fwd_kernel(
                     qk += tl.dot(qpe, kpe.to(qpe.dtype))
                 else:
                     qk += tl.dot(qpe.to(kpe.dtype), kpe)
-            qk *= sm_scale * k_scale
+            if USE_EXP2:
+                qk *= sm_scale * k_scale * LOG2E
+            else:
+                qk *= sm_scale * k_scale
 
             if logit_cap > 0:
                 qk = logit_cap * tanh(qk / logit_cap)
@@ -572,7 +605,7 @@ def _fwd_kernel(
                 qk = SCORE_MOD(
                     qk,
                     (cur_seq_len_prefix + cur_block_m * BLOCK_M + offs_m)[:, None],
-                    start_n + offs_n[None, :],
+                    start_n + offs_n_prefix[None, :],
                     (cur_seq_extend_start_idx + cur_block_m * BLOCK_M + offs_m)[
                         :, None
                     ],
@@ -590,8 +623,12 @@ def _fwd_kernel(
             row_max_fixed = tl.where(row_max == float("-inf"), -1e20, row_max)
             n_e_max = tl.maximum(row_max_fixed, e_max)
 
-            re_scale = tl.exp(e_max - n_e_max)
+            if USE_EXP2:
-            p = tl.exp(qk - n_e_max[:, None])
+                re_scale = tl.exp2(e_max - n_e_max)
+                p = tl.exp2(qk - n_e_max[:, None])
+            else:
+                re_scale = tl.exp(e_max - n_e_max)
+                p = tl.exp(qk - n_e_max[:, None])
             deno = deno * re_scale + tl.sum(p, 1)
 
             if PAGE_SIZE == 1:
@@ -612,14 +649,18 @@ def _fwd_kernel(
                 mask=mask_n[:, None] & mask_dv[None, :],
                 other=0.0,
             )
-            # keep softmax weights p in fp32 for the P·V dot (do not downcast to bf16)
+            if USE_FP8_PREFIX:
-            # on gfx1250; on other platforms restore the original p.to(v.dtype) cast.
+                p_dot = (p * FP8_MAX).to(v.dtype)
-            # TODO: remove this branch once the gfx1250 bf16 P·V issue is resolved.
+                acc = acc * re_scale[:, None] + tl.dot(p_dot, v) * (v_scale / FP8_MAX)
-            if IS_GFX1250:
-                dot = tl.dot(p, v.to(tl.float32), out_dtype=tl.float32)
             else:
-                dot = tl.dot(p.to(v.dtype), v)
+                # keep softmax weights p in fp32 for the P·V dot (do not downcast to bf16)
-            acc = acc * re_scale[:, None] + dot * v_scale
+                # on gfx1250; on other platforms restore the original p.to(v.dtype) cast.
+                # TODO: remove this branch once the gfx1250 bf16 P·V issue is resolved.
+                if IS_GFX1250:
+                    dot = tl.dot(p, v.to(tl.float32), out_dtype=tl.float32)
+                else:
+                    dot = tl.dot(p.to(v.dtype), v)
+                acc = acc * re_scale[:, None] + dot * v_scale
 
             e_max = n_e_max
 
@@ -631,16 +672,7 @@ def _fwd_kernel(
         else tl.minimum(cur_seq_len_extend, (cur_block_m + 1) * BLOCK_M)
     )
     extend_end = 0 if SKIP_EXTEND else cur_block_m_end
-    # The mask below keeps (q, kv) iff q <= kv + SLIDING_WINDOW_SIZE, so no tile
+    for start_n in range(0, extend_end, BLOCK_N):
-    # under this floor can hold an unmasked element -- tight for any BLOCK_M/BLOCK_N.
-    # SKIP_TILE already made those tiles no-ops, so bounding the loop is
-    # bit-identical and drops their cross-wave tl.max reduction.
-    extend_start = 0
-    if SLIDING_WINDOW_SIZE > 0:
-        extend_start = (
-            tl.maximum(cur_block_m * BLOCK_M - SLIDING_WINDOW_SIZE, 0) // BLOCK_N
-        ) * BLOCK_N
-    for start_n in range(extend_start, extend_end, BLOCK_N):
         start_n = tl.multiple_of(start_n, BLOCK_N)
         mask_n = (start_n + offs_n) < cur_block_m_end
 
@@ -706,7 +738,10 @@ def _fwd_kernel(
                 )
                 qk += tl.dot(qpe, kpe)
 
-            qk *= sm_scale
+            if USE_EXP2:
+                qk *= sm_scale * LOG2E
+            else:
+                qk *= sm_scale
 
             if logit_cap > 0:
                 qk = logit_cap * tanh(qk / logit_cap)
@@ -736,8 +771,12 @@ def _fwd_kernel(
             row_max_fixed = tl.where(row_max == float("-inf"), -1e20, row_max)
             n_e_max = tl.maximum(row_max_fixed, e_max)
 
-            re_scale = tl.exp(e_max - n_e_max)
+            if USE_EXP2:
-            p = tl.exp(qk - n_e_max[:, None])
+                re_scale = tl.exp2(e_max - n_e_max)
+                p = tl.exp2(qk - n_e_max[:, None])
+            else:
+                re_scale = tl.exp(e_max - n_e_max)
+                p = tl.exp(qk - n_e_max[:, None])
             deno = deno * re_scale + tl.sum(p, 1)
 
             offs_v = (
@@ -748,26 +787,41 @@ def _fwd_kernel(
             v = tl.load(
                 V_Extend + offs_v, mask=mask_n[:, None] & mask_dv[None, :], other=0.0
             )
-            # keep softmax weights p in fp32 for the P·V dot (do not downcast to bf16)
+            if USE_FP8_EXTEND:
-            # on gfx1250; on other platforms restore the original p.to(v.dtype) cast.
+                p_dot = (p * FP8_MAX).to(v.dtype)
-            # TODO: remove this branch once the gfx1250 bf16 P·V issue is resolved.
+                acc = acc * re_scale[:, None] + tl.dot(p_dot, v) * (1.0 / FP8_MAX)
-            if IS_GFX1250:
-                dot = tl.dot(p, v.to(tl.float32), out_dtype=tl.float32)
             else:
-                dot = tl.dot(p.to(v.dtype), v)
+                # keep softmax weights p in fp32 for the P·V dot (do not downcast to bf16)
-            acc = acc * re_scale[:, None] + dot
+                # on gfx1250; on other platforms restore the original p.to(v.dtype) cast.
+                # TODO: remove this branch once the gfx1250 bf16 P·V issue is resolved.
+                if IS_GFX1250:
+                    dot = tl.dot(p, v.to(tl.float32), out_dtype=tl.float32)
+                else:
+                    dot = tl.dot(p.to(v.dtype), v)
+                acc = acc * re_scale[:, None] + dot
 
             e_max = n_e_max
 
     if HAS_SINK:
         cur_sink = tl.load(sink_ptr + cur_head)
-        deno += tl.exp(cur_sink - e_max)
+        if USE_EXP2:
+            deno += tl.exp2(cur_sink * LOG2E - e_max)
+        else:
+            deno += tl.exp(cur_sink - e_max)
+
+    # A ragged prefix chunk can be empty for some requests. Represent an empty
+    # partial as output=0 and LSE=-inf so merge_state ignores it exactly.
+    no_kv = deno == 0.0
 
     if STORE_LSE:
         offs_lse = (
             cur_seq_extend_start_idx + cur_block_m * BLOCK_M + offs_m
         ) * stride_lse_bs + cur_head * stride_lse_h
-        lse = tl.log(deno) + e_max
+        if USE_EXP2:
+            lse = tl.log(deno) + e_max * LN2
+        else:
+            lse = tl.log(deno) + e_max
+        lse = tl.where(no_kv, float("-inf"), lse)
         tl.store(LSE_Extend + offs_lse, lse, mask=mask_m)
 
     offs_o = (
@@ -776,16 +830,17 @@ def _fwd_kernel(
         + cur_head * stride_oh
         + offs_dv[None, :]
     )
+    deno_safe = tl.where(no_kv, 1.0, deno)
     if STORE_TRANSPOSE:
         tl.store(
             O_Extend + offs_o.T,
-            (acc / deno[:, None]).T,
+            (acc / deno_safe[:, None]).T,
             mask=(mask_m[:, None] & mask_dv[None, :]).T,
         )
     else:
         tl.store(
             O_Extend + offs_o,
-            acc / deno[:, None],
+            acc / deno_safe[:, None],
             mask=mask_m[:, None] & mask_dv[None, :],
         )
 
@@ -820,6 +875,7 @@ def extend_attention_fwd(
     score_mod=None,
     aux_tensors=None,
     extend_seq_lens_cpu=None,
+    identity_kv_indices: bool = False,
 ):
     """
     q_extend, k_extend, v_extend, o_extend: contiguous tensors
@@ -832,6 +888,8 @@ def extend_attention_fwd(
     respectively so DCP can compute those two parts separately.
     ``score_mod`` / ``aux_tensors`` add a custom term to the attention logits;
     see triton_ops/score_mod.py for the contract.
+    ``identity_kv_indices`` promises that the prefix buffer is densely packed,
+    allowing direct addressing instead of loading an index for every token.
     """
     Lq, Lk, Lv = (
         q_extend.shape[-1],
@@ -839,20 +897,90 @@ def extend_attention_fwd(
         v_extend.shape[-1],
     )
 
-    # Get block sizes and configuration
-    BLOCK_DMODEL, BLOCK_DPE, BLOCK_DV, BLOCK_M, BLOCK_N, num_warps = (
-        _get_block_sizes_for_extend_attention(Lq, Lv)
-    )
-
     sm_scale = sm_scale or 1.0 / (Lq**0.5)
     batch_size, head_num = qo_indptr.shape[0] - 1, q_extend.shape[1]
     kv_group_num = q_extend.shape[1] // k_extend.shape[1]
+    zero_prefix_shape = (
+        head_num == 12
+        and k_extend.shape[1] == 12
+        and Lq == 192
+        and Lk == 192
+        and Lv == 128
+    )
+    absorbed_shape = (
+        head_num == 12
+        and k_extend.shape[1] == 1
+        and Lq == 576
+        and Lk == 576
+        and Lv == 512
+    )
+    kimi_k3_shape = zero_prefix_shape or absorbed_shape
+
+    # Match Aiter's opt-in behavior: cast Q, K, and V separately before the
+    # native-FP8 zero-prefix kernel at every sequence length.
+    use_fp8_zero_prefix = (
+        _is_gfx95
+        and envs.SGLANG_TRITON_FP8_PREFILL_ATTN.get()
+        and zero_prefix_shape
+        and q_extend.dtype == torch.bfloat16
+        and k_extend.dtype == torch.bfloat16
+        and v_extend.dtype == torch.bfloat16
+        and k_buffer.dtype == torch.float8_e4m3fn
+        and v_buffer.dtype == torch.float8_e4m3fn
+        and custom_mask is None
+        and is_causal
+        and sliding_window_size <= 0
+        and logit_cap <= 0
+        and xai_temperature_len <= 0
+        and sinks is None
+        and score_mod is None
+        and aux_tensors is None
+    )
+    if use_fp8_zero_prefix:
+        q_extend = q_extend.to(torch.float8_e4m3fn)
+        k_extend = k_extend.to(torch.float8_e4m3fn)
+        v_extend = v_extend.to(torch.float8_e4m3fn)
+
+    # Get block sizes and configuration for the generic fallback.
+    BLOCK_DMODEL, BLOCK_DPE, BLOCK_DV, BLOCK_M, BLOCK_N, num_warps = (
+        _get_block_sizes_for_extend_attention(Lq, Lv)
+    )
 
     USE_CUSTOM_MASK = custom_mask is not None
     # Skip custom mask for prefix part
     SKIP_PREFIX_CUSTOM_MASK = skip_prefix_custom_mask
 
     HAS_SINK = sinks is not None
+    USE_FP8_PREFIX = (
+        _is_gfx95
+        and kimi_k3_shape
+        and k_buffer.dtype == torch.float8_e4m3fn
+        and v_buffer.dtype == torch.float8_e4m3fn
+    )
+    USE_FP8_EXTEND = (
+        _is_gfx95
+        and zero_prefix_shape
+        and k_extend.dtype == torch.float8_e4m3fn
+        and v_extend.dtype == torch.float8_e4m3fn
+    )
+    FP8_MAX = (
+        torch.finfo(torch.float8_e4m3fn).max
+        if USE_FP8_PREFIX or USE_FP8_EXTEND
+        else 1.0
+    )
+    # At head_dim 192, FP8 operands allow a 128-column tile while BF16 does
+    # not fit in LDS. Widen prefix and extend sweeps independently.
+    BLOCK_N_ARCH = BLOCK_N
+    FP8_BLOCK_N = 128 if BLOCK_N_ARCH < 128 and Lq <= 192 else BLOCK_N_ARCH
+    BLOCK_N = FP8_BLOCK_N if USE_FP8_EXTEND else BLOCK_N_ARCH
+    BLOCK_N_PREFIX = FP8_BLOCK_N if USE_FP8_PREFIX else BLOCK_N_ARCH
+    USE_EXP2 = (
+        _is_gfx95
+        and kimi_k3_shape
+        and logit_cap <= 0
+        and xai_temperature_len <= 0
+        and score_mod is None
+    )
     STORE_LSE = lse_extend is not None
     stride_lse_bs = lse_extend.stride(0) if STORE_LSE else 0
     stride_lse_h = lse_extend.stride(1) if STORE_LSE else 0
@@ -877,7 +1005,9 @@ def extend_attention_fwd(
         grid = (compact_q_tiles, head_num)
     else:
         grid = (batch_size, head_num, triton.cdiv(max_len_extend, BLOCK_M))
-    num_stages = 1
+    num_stages = (
+        _get_num_stages_for_extend_attention(Lq, Lv, BLOCK_N) if kimi_k3_shape else 1
+    )
 
     extra_kargs = {}
     if _is_hip:
@@ -940,6 +1070,7 @@ def extend_attention_fwd(
         BLOCK_DV=BLOCK_DV,
         BLOCK_M=BLOCK_M,
         BLOCK_N=BLOCK_N,
+        BLOCK_N_PREFIX=BLOCK_N_PREFIX,
         Lq=Lq,
         Lv=Lv,
         USE_CUSTOM_MASK=USE_CUSTOM_MASK,
@@ -952,7 +1083,12 @@ def extend_attention_fwd(
         IS_GFX1250=_is_gfx1250,
         STORE_TRANSPOSE=_is_hip,
         USE_COMPACT_TILE_GRID=use_compact_tile_grid,
+        USE_EXP2=USE_EXP2,
+        USE_FP8_PREFIX=USE_FP8_PREFIX,
+        USE_FP8_EXTEND=USE_FP8_EXTEND,
+        FP8_MAX=FP8_MAX,
         PAGE_SIZE=page_size,
+        IDENTITY_KV_INDICES=identity_kv_indices,
         SCORE_MOD=score_mod,
         Aux0=aux0,
         aux0_stride_t=aux0_stride_t,
@@ -1470,3 +1606,536 @@ def extend_attention_fwd_unified(
         num_stages=num_stages,
         **extra_kargs,
     )
+
+
+@triton.jit
+def _dense_prefill_inner(
+    acc,
+    deno,
+    e_max,
+    q,
+    qpe,
+    K,
+    V,
+    cur_seq_kv_start,
+    cur_kv_head,
+    stride_kbs,
+    stride_kh,
+    stride_vbs,
+    stride_vh,
+    offs_d,
+    offs_dpe,
+    offs_dv,
+    q_pos,
+    mask_m,
+    mask_d,
+    mask_dv,
+    start_lo,
+    start_hi,
+    kv_end,
+    qk_scale,
+    logit_cap: tl.constexpr,
+    BLOCK_DPE: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    MASKED: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+    EVEN_D: tl.constexpr,
+    USE_FP8: tl.constexpr,
+    LOG2_FP8_MAX: tl.constexpr,
+):
+    """One online-softmax sweep over ``[start_lo, start_hi)`` of the KV axis.
+
+    Instantiated twice per kernel. ``MASKED=False`` covers the interior, where
+    every key is visible to every query row of the block: the predicates are
+    gone, so the K/V loads stay contiguous over the head dim and widen to
+    dwordx4. ``MASKED=True`` covers only the causal diagonal and the ragged
+    tail. Splitting the sweep is what keeps that per-element ``tl.where`` and
+    the narrowed loads off the long prefix, which is the bulk of the work.
+    """
+    offs_n = tl.arange(0, BLOCK_N)
+    for start_n in range(start_lo, start_hi, BLOCK_N):
+        start_n = tl.multiple_of(start_n, BLOCK_N)
+        offs_kn = start_n + offs_n
+
+        offs_k = (
+            (cur_seq_kv_start + offs_kn[None, :]) * stride_kbs
+            + cur_kv_head * stride_kh
+            + offs_d[:, None]
+        )
+        if MASKED:
+            mask_n = offs_kn < kv_end
+            k = tl.load(K + offs_k, mask=mask_n[None, :] & mask_d[:, None], other=0.0)
+        elif EVEN_D:
+            k = tl.load(K + offs_k)
+        else:
+            k = tl.load(K + offs_k, mask=mask_d[:, None], other=0.0)
+
+        qk = tl.dot(q, k, out_dtype=tl.float32)
+        if BLOCK_DPE > 0:
+            offs_kpe = (
+                (cur_seq_kv_start + offs_kn[None, :]) * stride_kbs
+                + cur_kv_head * stride_kh
+                + offs_dpe[:, None]
+            )
+            if MASKED:
+                kpe = tl.load(K + offs_kpe, mask=mask_n[None, :], other=0.0)
+            else:
+                kpe = tl.load(K + offs_kpe)
+            qk += tl.dot(qpe, kpe, out_dtype=tl.float32)
+
+        if logit_cap > 0:
+            qk *= qk_scale
+            qk = logit_cap * tanh(qk / logit_cap)
+            qk *= 1.4426950408889634
+        else:
+            qk *= qk_scale
+
+        if MASKED:
+            final_mask = mask_m[:, None] & mask_n[None, :]
+            if IS_CAUSAL:
+                final_mask &= q_pos[:, None] >= offs_kn[None, :]
+            qk = tl.where(final_mask, qk, float("-inf"))
+            row_max = tl.max(qk, 1)
+            # A fully masked row would poison e_max with -inf; -1e20 keeps the
+            # rescale finite and still contributes nothing to deno.
+            row_max = tl.where(row_max == float("-inf"), -1e20, row_max)
+        else:
+            row_max = tl.max(qk, 1)
+
+        n_e_max = tl.maximum(row_max, e_max)
+        re_scale = tl.exp2(e_max - n_e_max)
+        if USE_FP8:
+            # Bias the exponent instead of multiplying P by FP8_MAX after the
+            # fact: exp2(qk - m + log2(FP8_MAX)) == exp2(qk - m) * FP8_MAX, so
+            # the lift costs a BLOCK_M-wide subtract rather than a
+            # BLOCK_M x BLOCK_N one, and skips a rounding step. deno picks up
+            # the same constant factor, which cancels against acc at the final
+            # divide -- so out_scale drops its 1/FP8_MAX and LSE subtracts
+            # log2(FP8_MAX) from e_max.
+            p = tl.exp2(qk - (n_e_max - LOG2_FP8_MAX)[:, None])
+        else:
+            p = tl.exp2(qk - n_e_max[:, None])
+        deno = deno * re_scale + tl.sum(p, 1)
+
+        offs_v = (
+            (cur_seq_kv_start + offs_kn[:, None]) * stride_vbs
+            + cur_kv_head * stride_vh
+            + offs_dv[None, :]
+        )
+        if MASKED:
+            v = tl.load(V + offs_v, mask=mask_n[:, None] & mask_dv[None, :], other=0.0)
+        elif EVEN_D:
+            v = tl.load(V + offs_v)
+        else:
+            v = tl.load(V + offs_v, mask=mask_dv[None, :], other=0.0)
+
+        # P is already lifted off the e4m3 denormal floor by the exponent bias
+        # above when USE_FP8; the cast is all that is left.
+        #
+        # Guarding this rescale on `tl.min(re_scale) < 1.0` (it is exactly 1.0
+        # once the running max settles, which is most of a long prefix) was
+        # measured at 21% SLOWER: the branch splits the loop body and the
+        # pipeliner stops prefetching K/V across iterations. Keep it
+        # unconditional.
+        acc = acc * re_scale[:, None] + tl.dot(p.to(v.dtype), v)
+
+        e_max = n_e_max
+
+    return acc, deno, e_max
+
+
+@triton.jit
+def _fwd_kernel_dense_prefill(
+    Q,
+    K,
+    V,
+    O,
+    Lse,
+    qo_indptr,
+    kv_indptr,
+    sm_scale,
+    k_scale,
+    v_scale,
+    stride_qbs,
+    stride_qh,
+    stride_kbs,
+    stride_kh,
+    stride_vbs,
+    stride_vh,
+    stride_obs,
+    stride_oh,
+    stride_lse_bs,
+    stride_lse_h,
+    kv_group_num: tl.constexpr,
+    logit_cap: tl.constexpr,
+    Lq: tl.constexpr,
+    Lv: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_DPE: tl.constexpr,
+    BLOCK_DV: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+    STORE_LSE: tl.constexpr,
+    USE_FP8: tl.constexpr,
+    LOG2_FP8_MAX: tl.constexpr,
+    EVEN_D: tl.constexpr,
+    NUM_BLOCKS_M: tl.constexpr,
+):
+    """Single-loop dense prefill: K/V hold prefix + current chunk contiguously.
+
+    ``_fwd_kernel`` needs two stages because its prefix lives in the paged
+    latent KV cache while its suffix is contiguous. Once the prefix has been
+    up-projected into dense per-head K/V (``AttnForwardMethod.MHA_ONE_SHOT``),
+    both halves share one base pointer, one dtype and one scale, so the split
+    buys nothing and only costs pipelining and registers.
+
+    Causal masking is bottom-right aligned: query ``m`` of a sequence sits at
+    absolute position ``prefix_len + m``, where ``prefix_len = kv_len - q_len``.
+    """
+    cur_seq = tl.program_id(0)
+    cur_head = tl.program_id(1)
+    # Causal cost grows with block index: block m sweeps prefix_len + m*BLOCK_M
+    # keys. Program IDs dispatch roughly in order, so issuing the heavy blocks
+    # first lets the cheap ones backfill the tail instead of trailing it.
+    cur_block_m = NUM_BLOCKS_M - 1 - tl.program_id(2)
+    cur_kv_head = cur_head // kv_group_num
+
+    LOG2E: tl.constexpr = 1.4426950408889634
+    LN2: tl.constexpr = 0.6931471805599453
+
+    cur_seq_q_start = tl.load(qo_indptr + cur_seq)
+    cur_seq_q_len = tl.load(qo_indptr + cur_seq + 1) - cur_seq_q_start
+    cur_seq_kv_start = tl.load(kv_indptr + cur_seq)
+    cur_seq_kv_len = tl.load(kv_indptr + cur_seq + 1) - cur_seq_kv_start
+    cur_seq_prefix_len = cur_seq_kv_len - cur_seq_q_len
+
+    # Grid axis 2 spans the batch-max query length; short sequences bail early.
+    if cur_block_m * BLOCK_M >= cur_seq_q_len:
+        return
+
+    offs_m = tl.arange(0, BLOCK_M)
+    offs_n = tl.arange(0, BLOCK_N)
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    offs_dv = tl.arange(0, BLOCK_DV)
+
+    mask_m = (cur_block_m * BLOCK_M + offs_m) < cur_seq_q_len
+    mask_d = offs_d < Lq
+    mask_dv = offs_dv < Lv
+
+    offs_q = (
+        (cur_seq_q_start + cur_block_m * BLOCK_M + offs_m[:, None]) * stride_qbs
+        + cur_head * stride_qh
+        + offs_d[None, :]
+    )
+    q = tl.load(Q + offs_q, mask=mask_m[:, None] & mask_d[None, :], other=0.0)
+
+    if BLOCK_DPE > 0:
+        offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE)
+        offs_qpe = (
+            (cur_seq_q_start + cur_block_m * BLOCK_M + offs_m[:, None]) * stride_qbs
+            + cur_head * stride_qh
+            + offs_dpe[None, :]
+        )
+        qpe = tl.load(Q + offs_qpe, mask=mask_m[:, None], other=0.0)
+    else:
+        # Never read: BLOCK_DPE is constexpr, so the whole rope branch is
+        # folded away. These only keep the inner helper's signature uniform.
+        offs_dpe = offs_d
+        qpe = q
+
+    # Absolute position of each query row inside its sequence.
+    q_pos = cur_seq_prefix_len + cur_block_m * BLOCK_M + offs_m
+
+    acc = tl.zeros([BLOCK_M, BLOCK_DV], dtype=tl.float32)
+    deno = tl.zeros([BLOCK_M], dtype=tl.float32)
+    e_max = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
+
+    if IS_CAUSAL:
+        kv_end = tl.minimum(
+            cur_seq_kv_len, cur_seq_prefix_len + (cur_block_m + 1) * BLOCK_M
+        )
+        # The lowest query row of this block sees keys 0..prefix_len+m*BLOCK_M,
+        # so every whole BLOCK_N below that bound is unmasked for all rows.
+        n_full = ((cur_seq_prefix_len + cur_block_m * BLOCK_M + 1) // BLOCK_N) * BLOCK_N
+        n_full = tl.minimum(n_full, kv_end)
+    else:
+        kv_end = cur_seq_kv_len
+        n_full = (kv_end // BLOCK_N) * BLOCK_N
+
+    qk_scale = sm_scale * k_scale
+    if logit_cap <= 0:
+        qk_scale *= LOG2E
+
+    acc, deno, e_max = _dense_prefill_inner(
+        acc,
+        deno,
+        e_max,
+        q,
+        qpe,
+        K,
+        V,
+        cur_seq_kv_start,
+        cur_kv_head,
+        stride_kbs,
+        stride_kh,
+        stride_vbs,
+        stride_vh,
+        offs_d,
+        offs_dpe,
+        offs_dv,
+        q_pos,
+        mask_m,
+        mask_d,
+        mask_dv,
+        0,
+        n_full,
+        kv_end,
+        qk_scale,
+        logit_cap,
+        BLOCK_DPE=BLOCK_DPE,
+        BLOCK_N=BLOCK_N,
+        MASKED=False,
+        IS_CAUSAL=IS_CAUSAL,
+        EVEN_D=EVEN_D,
+        USE_FP8=USE_FP8,
+        LOG2_FP8_MAX=LOG2_FP8_MAX,
+    )
+    acc, deno, e_max = _dense_prefill_inner(
+        acc,
+        deno,
+        e_max,
+        q,
+        qpe,
+        K,
+        V,
+        cur_seq_kv_start,
+        cur_kv_head,
+        stride_kbs,
+        stride_kh,
+        stride_vbs,
+        stride_vh,
+        offs_d,
+        offs_dpe,
+        offs_dv,
+        q_pos,
+        mask_m,
+        mask_d,
+        mask_dv,
+        n_full,
+        kv_end,
+        kv_end,
+        qk_scale,
+        logit_cap,
+        BLOCK_DPE=BLOCK_DPE,
+        BLOCK_N=BLOCK_N,
+        MASKED=True,
+        IS_CAUSAL=IS_CAUSAL,
+        EVEN_D=EVEN_D,
+        USE_FP8=USE_FP8,
+        LOG2_FP8_MAX=LOG2_FP8_MAX,
+    )
+
+    no_kv = deno == 0.0
+
+    if STORE_LSE:
+        offs_lse = (
+            cur_seq_q_start + cur_block_m * BLOCK_M + offs_m
+        ) * stride_lse_bs + cur_head * stride_lse_h
+        # e_max is in log2 units because qk carries the folded LOG2E. Under
+        # FP8 deno also carries the exponent-bias lift; undoing it in log space
+        # is one more constant on the same term.
+        if USE_FP8:
+            lse = tl.log(deno) + (e_max - LOG2_FP8_MAX) * LN2
+        else:
+            lse = tl.log(deno) + e_max * LN2
+        lse = tl.where(no_kv, float("-inf"), lse)
+        tl.store(Lse + offs_lse, lse, mask=mask_m)
+
+    offs_o = (
+        (cur_seq_q_start + cur_block_m * BLOCK_M + offs_m[:, None]) * stride_obs
+        + cur_head * stride_oh
+        + offs_dv[None, :]
+    )
+    # The FP8 lift applied to P divides out between acc and deno, so v_scale
+    # is the only surviving factor.
+    out_scale = v_scale
+    deno_safe = tl.where(no_kv, 1.0, deno)
+    tl.store(
+        O + offs_o,
+        acc * (out_scale / deno_safe[:, None]),
+        mask=mask_m[:, None] & mask_dv[None, :],
+    )
+
+
+def can_use_dense_prefill_fp8(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    is_causal: bool,
+    logit_cap: float,
+) -> bool:
+    """Whether q/k/v may be cast to FP8 ahead of ``dense_prefill_attention_fwd``.
+
+    Deliberately as narrow as the zero-prefix gate in ``extend_attention_fwd``:
+    gfx950, BF16 dense inputs, plain causal softmax. Callers are responsible
+    for having already rejected custom masks, sinks, SWA and score mods.
+    """
+    return (
+        _is_gfx95
+        and envs.SGLANG_TRITON_FP8_PREFILL_ATTN.get()
+        and q.dtype == torch.bfloat16
+        and k.dtype == torch.bfloat16
+        and v.dtype == torch.bfloat16
+        and is_causal
+        and logit_cap <= 0
+    )
+
+
+def dense_prefill_attention_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    o: torch.Tensor,
+    qo_indptr: torch.Tensor,
+    kv_indptr: torch.Tensor,
+    max_len_q: int,
+    sm_scale: Optional[float] = None,
+    k_scale: float = 1.0,
+    v_scale: float = 1.0,
+    logit_cap: float = 0.0,
+    is_causal: bool = True,
+    lse: Optional[torch.Tensor] = None,
+) -> None:
+    """Dense varlen prefill over a fully materialized K/V.
+
+    q/o are addressed by ``qo_indptr``; k/v by ``kv_indptr``, whose
+    per-sequence length must be >= the query length. The excess leading rows
+    are the cached prefix, which every query of that sequence attends. Writes
+    ``o``, and ``lse`` (natural log) when provided.
+
+    Shapes::
+
+        q : [sum(q_len),  H,    Lq]      o : [sum(q_len), H, Lv]
+        k : [sum(kv_len), H_kv, Lq]      v : [sum(kv_len), H_kv, Lv]
+    """
+    Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
+    assert Lq == Lk, f"q/k head dims must match, got {Lq} vs {Lk}"
+    assert k.shape[0] == v.shape[0], (
+        f"k/v token counts must match, got {k.shape[0]} vs {v.shape[0]}"
+    )
+
+    sm_scale = sm_scale or 1.0 / (Lq**0.5)
+    batch_size = qo_indptr.shape[0] - 1
+    head_num = q.shape[1]
+    kv_group_num = q.shape[1] // k.shape[1]
+
+    # Shares the extend-attention tables: same head dims, same archs, and the
+    # 192/128 K3 entry is the shape this path exists to serve.
+    BLOCK_DMODEL, BLOCK_DPE, BLOCK_DV, BLOCK_M, BLOCK_N, num_warps = (
+        _get_block_sizes_for_extend_attention(Lq, Lv)
+    )
+    assert BLOCK_DMODEL + BLOCK_DPE >= Lq, (
+        f"tile {BLOCK_DMODEL}+{BLOCK_DPE} cannot cover head dim {Lq}"
+    )
+
+    # Both sides of every tl.dot must share a dtype, so a mixed bf16-q/fp8-k
+    # pair does not compile. That pair is reachable: on gfx95 with MXFP4
+    # kv_b_proj weights, forward_mha_rocm fuses the up-projection with the FP8
+    # cast (fused_gemm_afp4wfp4_split_cat) and hands back k/v already in e4m3
+    # while q is still bf16. Follow the cheap direction -- promote q rather
+    # than upcast the far larger k/v -- and keep the P scaling consistent with
+    # it. p is cast to v.dtype, so v decides the underflow lift too.
+    fp8 = torch.float8_e4m3fn
+    if fp8 in (q.dtype, k.dtype, v.dtype):
+        q, k, v = q.to(fp8), k.to(fp8), v.to(fp8)
+    else:
+        assert q.dtype == k.dtype == v.dtype, (
+            f"q/k/v dtypes must match, got {q.dtype}/{k.dtype}/{v.dtype}"
+        )
+    use_fp8 = q.dtype == fp8
+    fp8_max = torch.finfo(fp8).max if use_fp8 else 1.0
+
+    store_lse = lse is not None
+    stride_lse_bs = lse.stride(0) if store_lse else 0
+    stride_lse_h = lse.stride(1) if store_lse else 0
+
+    extra_kargs = {}
+    if _is_hip:
+        # No kpack: gfx950 overwrites it to 1 and warns on every launch.
+        # matrix_instr_nonkdim=32 was measured and is ~1.4x slower here.
+        extra_kargs = {"waves_per_eu": 1, "matrix_instr_nonkdim": 16}
+
+    num_stages = _get_num_stages_for_extend_attention(Lq, Lv, BLOCK_N)
+    if _is_gfx95 and Lq == 192 and Lv == 128 and use_fp8:
+        # The shared table is tuned for the two-stage absorbed kernel; this
+        # sweep is mask-free over its interior and wants a wider KV tile.
+        # FP8 only: at BLOCK_N=128 the BF16 K/V tiles need 168 KB of LDS
+        # against a 160 KB limit, so BF16 stays on the narrow tile below.
+        # Swept BLOCK_M x BLOCK_N x warps x stages over the three shapes this
+        # path actually sees (FP8, ms):
+        #
+        #                    q16384/p65536   q8192/p32768   q2464/p65536
+        #   128/64  w4 st3       7.89            4.00           1.28
+        #   128/128 w4 st2       7.12            3.56           1.18
+        #   256/128 w4 st2       6.91            3.50           2.08
+        #
+        # BLOCK_M=256 edges ahead on full chunks but halves the M-block count,
+        # which starves the 256 CUs on the short trailing chunk of a request --
+        # 1.6x slower there. 128/128 wins everywhere. At BLOCK_N=128 the loop
+        # body already covers the load latency, so the deeper pipeline that
+        # helped at BLOCK_N=64 no longer pays (7.12 at 2 stages vs 7.24 at 3).
+        BLOCK_N = 128
+        num_warps = 4
+        num_stages = 2
+    elif _is_gfx95 and Lq == 192 and Lv == 128:
+        # BF16, stuck on BLOCK_N=64 by LDS: there the mask-free interior does
+        # pipeline deeper than the table's two stages (10.90ms vs 11.10ms at
+        # 16K queries over a 64K prefix). 4 stages does not compile (Triton
+        # asserts in its pipeliner).
+        num_stages = 3
+
+    num_blocks_m = triton.cdiv(max_len_q, BLOCK_M)
+    grid = (batch_size, head_num, num_blocks_m)
+
+    _fwd_kernel_dense_prefill[grid](
+        q,
+        k,
+        v,
+        o,
+        lse,
+        qo_indptr,
+        kv_indptr,
+        sm_scale,
+        k_scale,
+        v_scale,
+        q.stride(0),
+        q.stride(1),
+        k.stride(0),
+        k.stride(1),
+        v.stride(0),
+        v.stride(1),
+        o.stride(0),
+        o.stride(1),
+        stride_lse_bs,
+        stride_lse_h,
+        kv_group_num=kv_group_num,
+        logit_cap=logit_cap,
+        Lq=Lq,
+        Lv=Lv,
+        BLOCK_DMODEL=BLOCK_DMODEL,
+        BLOCK_DPE=BLOCK_DPE,
+        BLOCK_DV=BLOCK_DV,
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        IS_CAUSAL=is_causal,
+        STORE_LSE=store_lse,
+        USE_FP8=use_fp8,
+        LOG2_FP8_MAX=math.log2(fp8_max),
+        EVEN_D=(BLOCK_DMODEL + BLOCK_DPE == Lq and BLOCK_DV == Lv),
+        NUM_BLOCKS_M=num_blocks_m,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        **extra_kargs,
+    )

python/sglang/srt/environ.py

@@ -1037,6 +1037,15 @@ class Envs:
     # gfx950 MLA decode stage-1: pick the launch geometry and split count per batch.
     # Reorders the fp32 accumulation, so off by default.
     SGLANG_MLA_DECODE_TUNE = EnvBool(False)
+    # Native FP8 prefill for exact gfx950 Kimi-K3 zero-prefix and absorbed
+    # cached-prefix shapes. Validated at 98% GSM8K accuracy.
+    SGLANG_TRITON_FP8_PREFILL_ATTN = EnvBool(True)
+    # Route Triton MLA prefill that carries a cached prefix through dense
+    # (non-absorbed) one-shot MHA: up-project the prefix out of the latent KV
+    # cache and run a single dense FP8 kernel instead of the absorbed 576/512
+    # prefill. Materializes K/V for the whole batch, so it only engages when
+    # the batch fits the chunk budget.
+    SGLANG_TRITON_DENSE_PREFILL_ATTN = EnvBool(True)
     SGLANG_ENABLE_TORCH_COMPILE = EnvBool(False)
     SGLANG_TRITON_PREFILL_TRUNCATION_ALIGN_SIZE = EnvInt(4096)
     SGLANG_TRITON_DECODE_SPLIT_TILE_SIZE = EnvInt(256)

python/sglang/srt/layers/attention/triton_backend.py

@@ -7,6 +7,9 @@ import torch
 import triton
 
 from sglang.kernels.ops.attention.metadata import get_num_kv_splits_triton
+from sglang.kernels.ops.attention.mla_kv_pack_quantize_fp8 import (
+    mla_kv_pack_quantize_fp8,
+)
 from sglang.srt.configs.hybrid_arch import mambaish_config
 from sglang.srt.configs.model_config import (
     AttentionArch,
@@ -36,12 +39,7 @@ from sglang.srt.model_executor.cuda_graph_config import (
     cuda_graph_fully_disabled,
 )
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
-from sglang.srt.runtime_context import (
+from sglang.srt.runtime_context import get_exec, get_parallel, get_schedule, get_spec
-    get_exec,
-    get_parallel,
-    get_schedule,
-    get_spec,
-)
 from sglang.srt.speculative.spec_utils import (
     draft_kv_indices_buffer_width,
     draft_kv_indices_used_len,
@@ -90,8 +88,6 @@ def _should_use_verify_shared_kv(model_config, topk, use_mla, use_verify_splitkv
     if use_mla:
         return is_kimi_k3(model_config.hf_config)
     if is_dspark_draft(model_config.hf_config):
-        # Added for the K3 DSpark draft model, which is qwen3 type attention,
-        # and using bidirectional (non-causal) mode.
         return use_verify_splitkv
     return (
         use_verify_splitkv
@@ -166,15 +162,13 @@ class TritonAttnBackend(AttentionBackend):
         )
         from sglang.kernels.ops.attention.extend_attention import (
             build_unified_kv_indices,
+            can_use_dense_prefill_fp8,
+            dense_prefill_attention_fwd,
             extend_attention_fwd,
             extend_attention_fwd_unified,
         )
-        from sglang.kernels.ops.attention.verify_mla import (
+        from sglang.kernels.ops.attention.verify_mla import verify_shared_kv_fwd
-            verify_shared_kv_fwd,
+        from sglang.kernels.ops.attention.verify_splitkv import verify_splitkv_fwd
-        )
-        from sglang.kernels.ops.attention.verify_splitkv import (
-            verify_splitkv_fwd,
-        )
 
         super().__init__()
 
@@ -189,6 +183,15 @@ class TritonAttnBackend(AttentionBackend):
             extend_attention_fwd_unified
         )
         self.build_unified_kv_indices = torch.compiler.disable(build_unified_kv_indices)
+        # Dense (non-absorbed) MLA prefill over a materialized prefix; see
+        # handle_attention_triton for when the dispatcher selects it.
+        self.dense_prefill_attention_fwd = torch.compiler.disable(
+            dense_prefill_attention_fwd
+        )
+        self.can_use_dense_prefill_fp8 = can_use_dense_prefill_fp8
+        # Cumulative full sequence lengths addressing the one-shot K/V; built
+        # on first use per forward and reset by init_forward_metadata.
+        self._dense_one_shot_kv_indptr = None
         # Split-KV EAGLE-verify kernel; enabled below once topk is known (valid only at topk == 1).
         self.verify_splitkv_fwd = torch.compiler.disable(verify_splitkv_fwd)
         # Grouped-head split-KV verify kernel for MLA or one shared local KV head.
@@ -235,6 +238,24 @@ class TritonAttnBackend(AttentionBackend):
         self.num_kv_head = model_runner.model_config.get_num_kv_heads(
             get_parallel().attn_tp_size, get_parallel().attn_dcp_size
         )
+        mla_config = model_runner.model_config
+        self.use_dense_fp8_chunked_prefill = (
+            self.use_mla
+            and is_gfx95_supported()
+            and envs.SGLANG_TRITON_DENSE_PREFILL_ATTN.get()
+            and envs.SGLANG_TRITON_FP8_PREFILL_ATTN.get()
+            and model_runner.kv_cache_dtype == torch.float8_e4m3fn
+            and self.num_head == 12
+            and mla_config.qk_nope_head_dim + mla_config.qk_rope_head_dim == 192
+            and mla_config.v_head_dim == 128
+            and mla_config.kv_lora_rank == 512
+        )
+        # forward_mha discovers these hooks dynamically. Hiding them when the
+        # exact Kimi-K3 FP8 configuration is absent keeps all other models on
+        # their existing code paths.
+        if not self.use_dense_fp8_chunked_prefill:
+            self.prepare_chunked_prefill_qkv = None
+            self.pack_prefix_chunk_kv = None
         # The decode kernel's "// Lv" stride trick requires attn_logits.shape[-1]
         # to exactly match the layer's v_head_dim, so hybrid SWA models with
         # differing SWA/full v_head_dim need a second buffer for SWA layers.
@@ -745,6 +766,7 @@ class TritonAttnBackend(AttentionBackend):
     def init_forward_metadata(self, forward_batch: ForwardBatch):
         """Init auxiliary variables for triton attention backend."""
 
+        self._dense_one_shot_kv_indptr = None
         bs = forward_batch.batch_size
         window_kv_indptr = self.window_kv_indptr
         window_kv_indices = None
@@ -1285,6 +1307,186 @@ class TritonAttnBackend(AttentionBackend):
     ):
         pass
 
+    @property
+    def pack_all_prefix_chunks(self) -> bool:
+        """Pack every prefix chunk into one FP8 buffer when capacity allows."""
+        return self.use_dense_fp8_chunked_prefill
+
+    @property
+    def fuse_prefix_into_extend(self) -> bool:
+        """Attend the packed prefix and current chunk in one launch."""
+        return self.use_dense_fp8_chunked_prefill
+
+    def prepare_chunked_prefill_qkv(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        forward_batch: ForwardBatch,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Convert the current dense MHA chunk once and reuse Q for prefix passes."""
+        fp8_dtype = torch.float8_e4m3fn
+        output_dtype = q.dtype
+        if output_dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
+            output_dtype = torch.bfloat16
+        forward_batch._triton_dense_fp8_output_dtype = output_dtype
+
+        if q.dtype != fp8_dtype:
+            q = q.to(fp8_dtype)
+        if k.dtype != fp8_dtype:
+            k = k.to(fp8_dtype)
+        if v.dtype != fp8_dtype:
+            v = v.to(fp8_dtype)
+        return q.contiguous(), k.contiguous(), v.contiguous()
+
+    def pack_prefix_chunk_kv(
+        self,
+        k_nope: torch.Tensor,
+        k_pe: torch.Tensor,
+        v: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Pack a materialized dense prefix directly into unit-scale FP8 K/V."""
+        return mla_kv_pack_quantize_fp8(
+            k_nope,
+            k_pe,
+            v,
+            fp8_dtype=torch.float8_e4m3fn,
+            enable_pdl=False,
+        )
+
+    def _can_run_dense_fp8_chunked_mha(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+    ) -> bool:
+        return (
+            self.use_dense_fp8_chunked_prefill
+            and forward_batch.attn_attend_prefix_cache is not None
+            and self.forward_metadata.custom_mask is None
+            and q.dtype == torch.float8_e4m3fn
+            and k.dtype == torch.float8_e4m3fn
+            and v.dtype == torch.float8_e4m3fn
+            and layer.tp_q_head_num == 12
+            and layer.tp_k_head_num == 12
+            and layer.qk_head_dim == 192
+            and layer.v_head_dim == 128
+            and (layer.sliding_window_size is None or layer.sliding_window_size <= -1)
+            and layer.logit_cap <= 0
+        )
+
+    def _forward_dense_fp8_chunked_mha(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+    ):
+        """Run current or cached dense FP8 K/V through Triton extend attention."""
+        output_dtype = getattr(
+            forward_batch, "_triton_dense_fp8_output_dtype", torch.bfloat16
+        )
+        output = torch.empty(
+            (q.shape[0], layer.tp_q_head_num, layer.v_head_dim),
+            dtype=output_dtype,
+            device=q.device,
+        )
+
+        prefix_k = getattr(forward_batch, "fused_prefix_k", None)
+        if prefix_k is not None:
+            # Prefix is non-causal while the current chunk is causal. The
+            # extend kernel already implements precisely that two-stage mask.
+            prefix_v = forward_batch.fused_prefix_v
+            self.extend_attention_fwd(
+                q,
+                k,
+                v,
+                output,
+                prefix_k,
+                prefix_v,
+                self.forward_metadata.qo_indptr,
+                forward_batch.prefix_chunk_cu_seq_lens[0],
+                forward_batch.prefix_dense_kv_indices[: prefix_k.shape[0]],
+                None,
+                True,
+                None,
+                self.forward_metadata.max_extend_len,
+                1.0,
+                1.0,
+                sm_scale=layer.scaling,
+                page_size=1,
+                extend_seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
+                identity_kv_indices=True,
+            )
+            return output
+
+        lse = torch.empty(
+            (q.shape[0], layer.tp_q_head_num),
+            dtype=torch.float32,
+            device=q.device,
+        )
+
+        if forward_batch.attn_attend_prefix_cache:
+            chunk_idx = forward_batch.prefix_chunk_idx
+            assert chunk_idx is not None and chunk_idx >= 0
+            kv_indptr = forward_batch.prefix_chunk_cu_seq_lens[chunk_idx]
+            kv_indices = forward_batch.prefix_dense_kv_indices[: k.shape[0]]
+            self.extend_attention_fwd(
+                q,
+                k[:0],
+                v[:0],
+                output,
+                k,
+                v,
+                self.forward_metadata.qo_indptr,
+                kv_indptr,
+                kv_indices,
+                None,
+                False,
+                None,
+                self.forward_metadata.max_extend_len,
+                1.0,
+                1.0,
+                sm_scale=layer.scaling,
+                lse_extend=lse,
+                skip_extend=True,
+                page_size=1,
+                extend_seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
+                identity_kv_indices=True,
+            )
+            # Empty ragged rows are returned as output=0, LSE=-inf, so the
+            # portable merge_state operation ignores them exactly.
+        else:
+            self.extend_attention_fwd(
+                q,
+                k,
+                v,
+                output,
+                k[:0],
+                v[:0],
+                self.forward_metadata.qo_indptr,
+                forward_batch.mha_empty_kv_indptr,
+                self.forward_metadata.kv_indices[:0],
+                None,
+                True,
+                None,
+                self.forward_metadata.max_extend_len,
+                1.0,
+                1.0,
+                sm_scale=layer.scaling,
+                lse_extend=lse,
+                skip_prefix=True,
+                page_size=1,
+                extend_seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
+            )
+
+        if forward_batch.mha_return_lse:
+            return output, lse
+        return output
+
     def _set_kv_buffer(
         self,
         forward_batch: ForwardBatch,
@@ -1330,6 +1532,16 @@ class TritonAttnBackend(AttentionBackend):
         score_mod=None,
         aux_tensors=None,
     ):
+        if (
+            k is not None
+            and v is not None
+            and sinks is None
+            and score_mod is None
+            and aux_tensors is None
+            and self._can_run_dense_fp8_chunked_mha(q, k, v, layer, forward_batch)
+        ):
+            return self._forward_dense_fp8_chunked_mha(q, k, v, layer, forward_batch)
+
         # TODO: reuse the buffer across layers
         attn_out = getattr(forward_batch, "_attn_output", None)
         if attn_out is not None:
@@ -1398,6 +1610,31 @@ class TritonAttnBackend(AttentionBackend):
         ):
             causal = False
 
+        # Dense one-shot MLA prefill (AttnForwardMethod.MHA_ONE_SHOT): k/v were
+        # up-projected out of the latent cache and span prefix + current chunk,
+        # so they no longer line up row-for-row with q the way
+        # extend_attention_fwd requires. Route to the single-loop dense kernel.
+        # A prefix-chunk phase (attn_attend_prefix_cache) also carries a longer
+        # k/v, but the dispatcher never hands Triton MHA_CHUNKED_KV.
+        if (
+            forward_batch.mha_one_shot
+            and not forward_batch.attn_attend_prefix_cache
+            and k is not None
+            and k.shape[0] != q.shape[0]
+        ):
+            return self._forward_extend_dense_one_shot(
+                q,
+                k,
+                v,
+                o,
+                layer,
+                forward_batch,
+                causal,
+                logits_soft_cap,
+                sinks=sinks,
+                score_mod=score_mod,
+            )
+
         if self.dcp_size > 1:
             if score_mod is not None:
                 raise NotImplementedError(
@@ -1517,6 +1754,81 @@ class TritonAttnBackend(AttentionBackend):
         )
         return o
 
+    def _dense_one_shot_kv_indptr_for(self, forward_batch: ForwardBatch):
+        """Cumulative full sequence lengths addressing the one-shot K/V rows.
+
+        The MHA one-shot K/V is gathered with fetch_mha_one_shot_kv_indices(),
+        which lays sequences out back to back at their full seq_len -- so the
+        row offsets are cumsum(seq_lens), not the prefix-only kv_indptr that
+        forward_metadata carries for the paged extend path.
+        """
+        if self._dense_one_shot_kv_indptr is None:
+            bs = forward_batch.batch_size
+            kv_indptr = torch.zeros(bs + 1, dtype=torch.int32, device=self.device)
+            kv_indptr[1:] = torch.cumsum(forward_batch.seq_lens[:bs], dim=0)
+            self._dense_one_shot_kv_indptr = kv_indptr
+        return self._dense_one_shot_kv_indptr
+
+    def _forward_extend_dense_one_shot(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        o: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        causal: bool,
+        logits_soft_cap: float,
+        sinks: Optional[torch.Tensor] = None,
+        score_mod=None,
+    ):
+        # Guarded rather than silently fallen back on: dropping to
+        # extend_attention_fwd with a longer-than-q k/v would read the wrong
+        # rows and quietly return wrong numbers.
+        if sinks is not None or score_mod is not None:
+            raise NotImplementedError(
+                "Triton dense one-shot prefill does not support sinks/score_mod"
+            )
+        if layer.sliding_window_size is not None and layer.sliding_window_size > -1:
+            raise NotImplementedError(
+                "Triton dense one-shot prefill does not support sliding windows"
+            )
+        if layer.k_scale is not None or layer.v_scale is not None:
+            raise NotImplementedError(
+                "Triton dense one-shot prefill does not support KV descales"
+            )
+        if layer.xai_temperature_len is not None and layer.xai_temperature_len > 0:
+            raise NotImplementedError(
+                "Triton dense one-shot prefill does not support xai temperature"
+            )
+
+        q = q.view(-1, layer.tp_q_head_num, layer.qk_head_dim)
+        k = k.view(-1, layer.tp_k_head_num, layer.qk_head_dim)
+        v = v.view(-1, layer.tp_k_head_num, layer.v_head_dim)
+
+        if self.can_use_dense_prefill_fp8(
+            q, k, v, is_causal=causal, logit_cap=logits_soft_cap
+        ):
+            # Cast Q, K and V separately, matching the zero-prefix FP8 gate in
+            # extend_attention_fwd (and Aiter's opt-in behavior).
+            q = q.to(torch.float8_e4m3fn)
+            k = k.to(torch.float8_e4m3fn)
+            v = v.to(torch.float8_e4m3fn)
+
+        self.dense_prefill_attention_fwd(
+            q,
+            k.contiguous(),
+            v.contiguous(),
+            o.view(-1, layer.tp_q_head_num, layer.v_head_dim),
+            self.forward_metadata.qo_indptr,
+            self._dense_one_shot_kv_indptr_for(forward_batch),
+            self.forward_metadata.max_extend_len,
+            sm_scale=layer.scaling,
+            logit_cap=logits_soft_cap,
+            is_causal=causal,
+        )
+        return o
+
     def _forward_extend_dcp(
         self,
         q: torch.Tensor,

python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py

@@ -47,9 +47,20 @@ class ForwardBatchDeepSeekMHAMixin:
     prefix_chunk_num_tokens: Optional[List[int]] = None
     # KV Indices for each chunk
     prefix_chunk_kv_indices: Optional[List[torch.Tensor]] = None
+    # All chunk KV indices in chunk-major order. The packed dense-prefix path
+    # gathers this once, then slices materialized K/V by chunk token count.
+    prefix_all_kv_indices: Optional[torch.Tensor] = None
+    # Reusable identity indices into a temporary dense prefix K/V buffer.
+    prefix_dense_kv_indices: Optional[torch.Tensor] = None
+    # Reusable empty prefix indptr for current-chunk-only attention.
+    mha_empty_kv_indptr: Optional[torch.Tensor] = None
     # For MLA chunked prefix cache used in chunked prefill
     # Tell attention backend whether lse needs to be returned
     mha_return_lse: Optional[bool] = None
+    # Packed dense prefix K/V for one fused prefix+extend attention launch.
+    # These are set only for the duration of an attention call.
+    fused_prefix_k: Optional[torch.Tensor] = None
+    fused_prefix_v: Optional[torch.Tensor] = None
     # Whether to apply MHA_ONE_SHOT forward method
     mha_one_shot: Optional[bool] = None
     # KV Indices for MHA_ONE_SHOT forward method
@@ -64,7 +75,12 @@ class ForwardBatchDeepSeekMHAMixin:
     def set_attn_attend_prefix_cache(self, attn_attend_prefix_cache: bool):
         self.attn_attend_prefix_cache = attn_attend_prefix_cache
 
-    def prepare_chunked_kv_indices(self, device: torch.device):
+    def prepare_chunked_kv_indices(
+        self,
+        device: torch.device,
+        pack_all_prefix_chunks: bool = False,
+        dense_metadata: bool = False,
+    ):
         self.prefix_chunk_kv_indices = []
         req_to_token = get_req_to_token_pool().req_to_token
         for idx in range(self.num_prefix_chunks):
@@ -95,6 +111,24 @@ class ForwardBatchDeepSeekMHAMixin:
             chunk_kv_indices = translator.translate_dcp_read_ids(chunk_kv_indices)
             self.prefix_chunk_kv_indices.append(chunk_kv_indices)
 
+        if not (pack_all_prefix_chunks or dense_metadata):
+            return
+
+        # Only the packed path gathers the whole prefix in one operation. The
+        # identity indices and empty indptr describe a dense per-chunk buffer
+        # and are required for both packed and bounded chunked paths.
+        if pack_all_prefix_chunks:
+            self.prefix_all_kv_indices = torch.cat(self.prefix_chunk_kv_indices)
+        max_dense_chunk_tokens = max(
+            chunk_indices.numel() for chunk_indices in self.prefix_chunk_kv_indices
+        )
+        self.prefix_dense_kv_indices = torch.arange(
+            max_dense_chunk_tokens, dtype=torch.int32, device=device
+        )
+        self.mha_empty_kv_indptr = torch.zeros(
+            self.batch_size + 1, dtype=torch.int32, device=device
+        )
+
     # Here we suppose the length of each chunk is equal
     # For example, if we have 4 sequences with prefix length [256, 512, 768, 1024], prefix_chunk_len = 256
     # num_prefix_chunks = cdiv(1024, 256) = 4
@@ -125,25 +159,23 @@ class ForwardBatchDeepSeekMHAMixin:
 
     # Called before each attention module if using chunked kv cache for prefill
     # Some of the codes are adapted from https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/backends/mla/common.py
-    def prepare_chunked_prefix_cache_info(self, device: torch.device):
+    def prepare_chunked_prefix_cache_info(
+        self,
+        device: torch.device,
+        pack_all_prefix_chunks: bool = False,
+        single_chunk: bool = False,
+        dense_metadata: bool = False,
+    ):
 
         from sglang.srt.mem_cache.memory_pool import (
             HybridLinearKVPool,
             MLATokenToKVPool,
         )
-        from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 
         token_to_kv_pool = get_token_to_kv_pool()
-        assert (
+        assert isinstance(token_to_kv_pool, MLATokenToKVPool) or (
-            isinstance(token_to_kv_pool, MLATokenToKVPool)
+            isinstance(token_to_kv_pool, HybridLinearKVPool)
-            or (
+            and isinstance(token_to_kv_pool.full_kv_pool, MLATokenToKVPool)
-                isinstance(token_to_kv_pool, HybridLinearKVPool)
-                and isinstance(token_to_kv_pool.full_kv_pool, MLATokenToKVPool)
-            )
-            or (
-                isinstance(token_to_kv_pool, SWAKVPool)
-                and isinstance(token_to_kv_pool.full_kv_pool, MLATokenToKVPool)
-            )
         ), "Currently chunked prefix cache can only be used by Deepseek models"
 
         if not any(self.extend_prefix_lens_cpu):
@@ -158,7 +190,13 @@ class ForwardBatchDeepSeekMHAMixin:
 
         # chunk_capacity is the maximum number of tokens in each chunk
         chunk_capacity = self.get_max_chunk_capacity()
-        self.prefix_chunk_len = chunk_capacity // self.batch_size
+        if single_chunk:
+            # The caller has already checked that the full packed prefix fits
+            # chunk_capacity. One request-major chunk gives the fused kernel the
+            # same per-request layout described by prefix_chunk_cu_seq_lens[0].
+            self.prefix_chunk_len = max(self.extend_prefix_lens_cpu)
+        else:
+            self.prefix_chunk_len = chunk_capacity // self.batch_size
 
         self.num_prefix_chunks = (
             max(self.extend_prefix_lens_cpu) + self.prefix_chunk_len - 1
@@ -210,7 +248,7 @@ class ForwardBatchDeepSeekMHAMixin:
         ]
 
         # Precompute the kv indices for each chunk
-        self.prepare_chunked_kv_indices(device)
+        self.prepare_chunked_kv_indices(device, pack_all_prefix_chunks, dense_metadata)
 
     def fetch_mha_one_shot_kv_indices(self):
         if self.mha_one_shot_kv_indices is not None:

python/sglang/srt/models/deepseek_common/attention_backend_handler.py

@@ -1,3 +1,4 @@
+from sglang.srt.environ import envs
 from sglang.srt.layers.attention.tbo_backend import TboAttnBackend
 from sglang.srt.layers.utils.cp_utils import mla_use_prefill_cp
 from sglang.srt.model_executor.forward_context import get_attn_backend
@@ -13,17 +14,21 @@ from sglang.srt.models.deepseek_common.attention_forward_methods.forward_methods
 from sglang.srt.models.deepseek_common.utils import _is_hip
 from sglang.srt.runtime_context import (
     get_exec,
+    get_parallel,
     get_platform,
 )
-from sglang.srt.utils import use_intel_amx_backend
+from sglang.srt.utils import (
+    is_gfx95_supported,
+    use_intel_amx_backend,
+)
 
 MHA_ONE_SHOT_SUPPORTED_BACKENDS = ["fa3", "flashinfer", "flashmla"]
 
 # ROCm runs dedicated MHA/MLA implementations (forward_mha_rocm.py /
 # forward_mla_rocm.py) so the shared CUDA paths carry no AMD branches. Backend
 # handlers keep returning the generic method; the platform swap happens here.
-# MHA_CHUNKED_KV has no ROCm entry because its accumulation step needs the
+# MHA_CHUNKED_KV deliberately stays generic on ROCm. Its shared implementation
-# CUDA-only merge_state_v2 kernel.
+# selects the ROCm prepare/fetch helpers and the portable merge_state wrapper.
 _ROCM_FORWARD_METHODS = {
     AttnForwardMethod.MHA: AttnForwardMethod.MHA_ROCM,
     AttnForwardMethod.MHA_ONE_SHOT: AttnForwardMethod.MHA_ONE_SHOT_ROCM,
@@ -212,6 +217,26 @@ def handle_attention_dsa(attn, forward_batch):
     return AttnForwardMethod.MLA
 
 
+def _can_use_triton_dense_fp8_prefill(attn, forward_batch) -> bool:
+    prefix_lens = forward_batch.extend_prefix_lens_cpu
+    return (
+        _is_hip
+        and is_gfx95_supported()
+        and envs.SGLANG_TRITON_FP8_PREFILL_ATTN.get()
+        and attn.kv_cache_dtype == "fp8_e4m3"
+        and attn.num_local_heads == 12
+        and attn.qk_nope_head_dim == 128
+        and attn.qk_rope_head_dim == 64
+        and attn.v_head_dim == 128
+        and attn.kv_lora_rank == 512
+        and not get_parallel().dcp_enabled
+        and not mla_use_prefill_cp(forward_batch)
+        and forward_batch.forward_mode.is_extend_without_speculative()
+        and prefix_lens is not None
+        and any(prefix_lens)
+    )
+
+
 def handle_attention_triton(attn, forward_batch):
     if is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph():
         return AttnForwardMethod.MLA
@@ -220,13 +245,18 @@ def handle_attention_triton(attn, forward_batch):
     if get_exec().deterministic.enable_deterministic_inference:
         return _dispatch_mla_subtype(attn, forward_batch)
 
+    # Kimi-K3 with an FP8 latent cache uses dense 192/128 K/V for cached
+    # prefixes. Always select chunked-KV here: its fast path packs the prefix
+    # once and fuses it with the current chunk in the normal extend kernel.
+    if _can_use_triton_dense_fp8_prefill(attn, forward_batch):
+        return AttnForwardMethod.MHA_CHUNKED_KV
+
     if (
         forward_batch.forward_mode.is_extend_without_speculative()
         and sum(forward_batch.extend_prefix_lens_cpu) == 0
     ):
         return AttnForwardMethod.MHA
-    else:
+    return _dispatch_mla_subtype(attn, forward_batch)
-        return _dispatch_mla_subtype(attn, forward_batch)
 
 
 def handle_attention_intel_xpu(attn, forward_batch):

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py

@@ -21,23 +21,18 @@ from sglang.srt.model_executor.forward_context import (
 )
 from sglang.srt.models.deepseek_common.utils import (
     _is_cuda,
+    _is_hip,
     _is_musa,
     _is_npu,
     _use_aiter_gfx95,
 )
-from sglang.srt.runtime_context import (
+from sglang.srt.runtime_context import get_exec, get_parallel, get_schedule
-    get_exec,
-    get_parallel,
-    get_schedule,
-)
 from sglang.srt.utils import BumpAllocator, next_power_of_2
 
 if TYPE_CHECKING:
     from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA
 
 if _is_cuda:
-    from sgl_kernel import merge_state_v2
-
     from sglang.kernels.ops.attention.concat_mla import concat_mla_k
 elif _is_musa:
     from sgl_kernel import concat_mla_k
@@ -45,9 +40,52 @@ elif _is_musa:
 
 def resolve_attn_backend(forward_batch: ForwardBatch):
     backend = get_attn_backend()
-    if isinstance(backend, TboAttnBackend):
+    while True:
-        backend = backend.primary
+        if isinstance(backend, TboAttnBackend):
-    return backend
+            backend = backend.primary
+            continue
+        # Hybrid KDA/MLA models route full-attention calls through an outer
+        # HybridLinearAttnBackend. Model-side MHA preparation hooks belong to
+        # its full-attention child.
+        if hasattr(backend, "full_attn_backend"):
+            backend = backend.full_attn_backend
+            continue
+        # A split prefill/decode HybridAttnBackend may itself be the full-attn
+        # child. Resolve the backend serving this forward mode as well.
+        if (
+            hasattr(backend, "prefill_backend")
+            and hasattr(backend, "decode_backend")
+            and hasattr(backend, "_select_backend")
+        ):
+            backend = backend._select_backend(forward_batch.forward_mode)
+            continue
+        return backend
+
+
+def use_dense_prefix_kv(backend) -> bool:
+    """Whether the backend substitutes dense K/V for prefix chunks."""
+    return (
+        getattr(backend, "pack_prefix_chunk_kv", None) is not None
+        and getattr(backend, "pack_all_prefix_chunks", False)
+        and not get_parallel().dcp_enabled
+    )
+
+
+def use_packed_prefix_chunks(backend, forward_batch: ForwardBatch) -> bool:
+    """Whether the complete prefix can be packed into one bounded buffer."""
+    prefix_lens = forward_batch.extend_prefix_lens_cpu
+    return (
+        use_dense_prefix_kv(backend)
+        and prefix_lens is not None
+        and sum(prefix_lens) <= forward_batch.get_max_chunk_capacity()
+    )
+
+
+def use_fused_prefix_extend(backend, forward_batch: ForwardBatch) -> bool:
+    """Whether the packed prefix and current chunk can share one launch."""
+    return getattr(
+        backend, "fuse_prefix_into_extend", False
+    ) and use_packed_prefix_chunks(backend, forward_batch)
 
 
 def forward_dsa_indexer_for_mha(
@@ -281,7 +319,14 @@ class DeepseekMHAForwardMixin:
         # The top comments in https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/backends/mla/common.py
         # will be helpful for understanding the purpose of this function.
 
-        # First do normal mha forward to get output for extended part
+        # Preserve the ROCm fused RMS/quantized projection path for the current
+        # chunk; the shared preparation path is primarily the CUDA version.
+        if _is_hip:
+            return self.forward_normal_rocm_prepare(
+                positions, hidden_states, forward_batch, zero_allocator
+            )
+
+        # First do normal mha forward to get output for extended part.
         return self.forward_normal_prepare(
             positions, hidden_states, forward_batch, zero_allocator
         )
@@ -297,12 +342,34 @@ class DeepseekMHAForwardMixin:
         has_extend_prefix = forward_batch.extend_prefix_lens_cpu is not None and any(
             forward_batch.extend_prefix_lens_cpu
         )
+        backend = resolve_attn_backend(forward_batch)
+        prepare_qkv_fn = getattr(backend, "prepare_chunked_prefill_qkv", None)
+        if has_extend_prefix and prepare_qkv_fn is not None:
+            q, k, v = prepare_qkv_fn(q, k, v, forward_batch)
+
+        fused_prefix = has_extend_prefix and use_fused_prefix_extend(
+            backend, forward_batch
+        )
+
         # Only initialize the info once
         if has_extend_prefix and forward_batch.num_prefix_chunks is None:
-            forward_batch.prepare_chunked_prefix_cache_info(q.device)
+            forward_batch.prepare_chunked_prefix_cache_info(
+                q.device,
+                pack_all_prefix_chunks=use_packed_prefix_chunks(backend, forward_batch),
+                single_chunk=fused_prefix,
+                dense_metadata=use_dense_prefix_kv(backend),
+            )
             if hasattr(get_attn_backend(), "init_mha_chunk_metadata"):
                 get_attn_backend().init_mha_chunk_metadata(forward_batch)
 
+        if fused_prefix:
+            attn_output = self._fused_prefix_extend_attn_mha(q, k, v, forward_batch)
+            attn_output = attn_output.reshape(
+                -1, self.num_local_heads * self.v_head_dim
+            )
+            output, _ = self.o_proj(attn_output)
+            return output
+
         forward_batch.mha_return_lse = has_extend_prefix
         # Do mha for extended part without prefix
         forward_batch.set_attn_attend_prefix_cache(False)
@@ -356,6 +423,48 @@ class DeepseekMHAForwardMixin:
         forward_batch.set_attn_attend_prefix_cache(False)
         return self.forward_normal_core(q, k, v, forward_batch, gate)
 
+    def _fused_prefix_extend_attn_mha(
+        self: DeepseekV2AttentionMLA,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        forward_batch: ForwardBatch,
+    ) -> torch.Tensor:
+        """Attend the complete prefix and current chunk in one kernel launch."""
+        assert forward_batch.num_prefix_chunks == 1, (
+            "fused prefix+extend needs the single-chunk layout, got "
+            f"{forward_batch.num_prefix_chunks} chunks"
+        )
+        backend = resolve_attn_backend(forward_batch)
+        get_mla_kv_buffer = (
+            self._get_mla_kv_buffer_rocm if _is_hip else self._get_mla_kv_buffer
+        )
+
+        # With one chunk, these indices are request-major and are described by
+        # prefix_chunk_cu_seq_lens[0]. Project and pack the prefix only once.
+        kv_a_normed, k_pe = get_mla_kv_buffer(
+            forward_batch.prefix_all_kv_indices, torch.bfloat16, forward_batch
+        )
+        kv = self.kv_b_proj(kv_a_normed)[0]
+        kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim)
+        prefix_k, prefix_v = backend.pack_prefix_chunk_kv(
+            kv[..., : self.qk_nope_head_dim],
+            k_pe,
+            kv[..., self.qk_nope_head_dim :],
+        )
+        del kv_a_normed, k_pe, kv
+
+        forward_batch.mha_return_lse = False
+        forward_batch.set_attn_attend_prefix_cache(False)
+        forward_batch.set_prefix_chunk_idx(0)
+        forward_batch.fused_prefix_k = prefix_k
+        forward_batch.fused_prefix_v = prefix_v
+        try:
+            return self.attn_mha(q, k, v, forward_batch, save_kv_cache=False)
+        finally:
+            forward_batch.fused_prefix_k = None
+            forward_batch.fused_prefix_v = None
+
     def _chunked_prefix_attn_mha(
         self: DeepseekV2AttentionMLA,
         q: torch.Tensor,
@@ -364,46 +473,82 @@ class DeepseekMHAForwardMixin:
         forward_batch: ForwardBatch,
     ) -> torch.Tensor:
         # kv_b_proj needs BF16 input, but legacy q.dtype was BF16 by accident.
+        from sglang.srt.layers.attention.merge_state import merge_state
+
         backend = resolve_attn_backend(forward_batch)
         pack_fn = getattr(backend, "pack_prefix_chunk_kv", None)
         kv_a_dtype = torch.bfloat16 if pack_fn is not None else q.dtype
+        get_mla_kv_buffer = (
+            self._get_mla_kv_buffer_rocm if _is_hip else self._get_mla_kv_buffer
+        )
 
-        assert forward_batch.num_prefix_chunks is not None
+        # If the complete prefix fits the capacity, gather/project/pack once.
-        for i in range(forward_batch.num_prefix_chunks):
+        # Larger prefixes retain bounded per-chunk materialization.
-            forward_batch.set_prefix_chunk_idx(i)
+        packed_prefix_k = None
-
+        packed_prefix_v = None
-            kv_indices = forward_batch.prefix_chunk_kv_indices[i]
+        pack_all_prefix_chunks = use_packed_prefix_chunks(backend, forward_batch)
-            # Fetch latent cache from memory pool with precomputed chunked kv indices
+        if pack_all_prefix_chunks:
-            kv_a_normed, k_pe = self._get_mla_kv_buffer(
+            kv_a_normed, k_pe = get_mla_kv_buffer(
-                kv_indices, kv_a_dtype, forward_batch
+                forward_batch.prefix_all_kv_indices,
-            )
+                kv_a_dtype,
-            kv_a_normed, k_pe = all_gather_kv_cache_for_mha_chunk_extend(
+                forward_batch,
-                kv_a_normed,
-                k_pe,
-                forward_batch.prefix_chunk_seq_lens_cpu[i],
-                forward_batch.prefix_chunk_starts_cpu[i],
             )
             kv = self.kv_b_proj(kv_a_normed)[0]
             kv = kv.view(
                 -1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim
             )
-            v = kv[..., self.qk_nope_head_dim :]
             k_nope = kv[..., : self.qk_nope_head_dim]
+            v_dense = kv[..., self.qk_nope_head_dim :]
+            packed_prefix_k, packed_prefix_v = pack_fn(k_nope, k_pe, v_dense)
+            del kv_a_normed, k_pe, kv, k_nope, v_dense
 
-            if pack_fn is not None:
+        assert forward_batch.num_prefix_chunks is not None
-                k, v = pack_fn(k_nope, k_pe, v)
+        packed_offset = 0
+        for i in range(forward_batch.num_prefix_chunks):
+            forward_batch.set_prefix_chunk_idx(i)
+
+            if pack_all_prefix_chunks:
+                chunk_num_tokens = forward_batch.prefix_chunk_num_tokens[i]
+                packed_end = packed_offset + chunk_num_tokens
+                k = packed_prefix_k[packed_offset:packed_end]
+                v = packed_prefix_v[packed_offset:packed_end]
+                packed_offset = packed_end
             else:
-                k = torch.empty(
+                kv_indices = forward_batch.prefix_chunk_kv_indices[i]
-                    (
+                kv_a_normed, k_pe = get_mla_kv_buffer(
-                        k_nope.shape[0],
+                    kv_indices, kv_a_dtype, forward_batch
-                        self.num_local_heads,
-                        self.qk_nope_head_dim + self.qk_rope_head_dim,
-                    ),
-                    dtype=v.dtype,
-                    device=v.device,
                 )
-                k[..., : self.qk_nope_head_dim] = k_nope
+                kv_a_normed, k_pe = all_gather_kv_cache_for_mha_chunk_extend(
-                k[..., self.qk_nope_head_dim :] = k_pe
+                    kv_a_normed,
+                    k_pe,
+                    forward_batch.prefix_chunk_seq_lens_cpu[i],
+                    forward_batch.prefix_chunk_starts_cpu[i],
+                )
+                kv = self.kv_b_proj(kv_a_normed)[0]
+                kv = kv.view(
+                    -1,
+                    self.num_local_heads,
+                    self.qk_nope_head_dim + self.v_head_dim,
+                )
+                v_dense = kv[..., self.qk_nope_head_dim :]
+                k_nope = kv[..., : self.qk_nope_head_dim]
+
+                if pack_fn is not None:
+                    k, v = pack_fn(k_nope, k_pe, v_dense)
+                else:
+                    v = v_dense
+                    k = torch.empty(
+                        (
+                            k_nope.shape[0],
+                            self.num_local_heads,
+                            self.qk_nope_head_dim + self.qk_rope_head_dim,
+                        ),
+                        dtype=v.dtype,
+                        device=v.device,
+                    )
+                    k[..., : self.qk_nope_head_dim] = k_nope
+                    k[..., self.qk_nope_head_dim :] = k_pe
+                del kv_a_normed, k_pe, kv, k_nope, v_dense
 
             output, lse = self.attn_mha(
                 q,
@@ -418,9 +563,20 @@ class DeepseekMHAForwardMixin:
             )
             tmp_output = torch.empty_like(accum_output)
             tmp_lse = torch.empty_like(accum_lse)
-            merge_state_v2(output, lse, accum_output, accum_lse, tmp_output, tmp_lse)
+            merge_state(
+                output,
+                lse,
+                accum_output,
+                accum_lse,
+                tmp_output,
+                tmp_lse,
+            )
             accum_output, accum_lse = tmp_output, tmp_lse
-            del kv, k, v, output, lse, tmp_output, tmp_lse
+            del k, v, output, lse, tmp_output, tmp_lse
+
+        if pack_all_prefix_chunks:
+            assert packed_offset == packed_prefix_k.shape[0]
+            del packed_prefix_k, packed_prefix_v
 
         return accum_output
 

test/registered/unit/layers/attention/test_triton_dense_prefill_gfx950.py

@@ -0,0 +1,188 @@
+"""Dense (non-absorbed) Triton prefill over a materialized prefix + chunk.
+
+Covers ``AttnForwardMethod.MHA_ONE_SHOT`` for Kimi-K3 on the triton backend,
+where the cached prefix is up-projected to the 192/128 MHA shape and attended
+in one pass instead of running the 576/512 absorbed kernel.
+"""
+
+import unittest
+
+import torch
+
+from sglang.kernels.ops.attention.extend_attention import (
+    can_use_dense_prefill_fp8,
+    dense_prefill_attention_fwd,
+)
+from sglang.srt.environ import envs
+from sglang.srt.utils import get_device, is_gfx95_supported, is_hip
+from sglang.test.ci.ci_register import register_amd_ci
+
+register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd-mi35x")
+
+H_Q, D_QK, D_V = 12, 192, 128
+FP8 = torch.float8_e4m3fn
+
+
+def _reference(q, k, v, qo_indptr, kv_indptr, scale, is_causal):
+    """Bottom-right aligned causal attention, one sequence at a time, in fp32."""
+    out = torch.empty(q.shape[0], H_Q, D_V, dtype=torch.float32, device=q.device)
+    lse = torch.empty(q.shape[0], H_Q, dtype=torch.float32, device=q.device)
+    for i in range(len(qo_indptr) - 1):
+        q_lo, q_hi = int(qo_indptr[i]), int(qo_indptr[i + 1])
+        k_lo, k_hi = int(kv_indptr[i]), int(kv_indptr[i + 1])
+        q_len, kv_len = q_hi - q_lo, k_hi - k_lo
+        scores = (
+            torch.matmul(
+                q[q_lo:q_hi].float().transpose(0, 1),
+                k[k_lo:k_hi].float().transpose(0, 1).transpose(1, 2),
+            )
+            * scale
+        )
+        if is_causal:
+            # Query m sits at absolute position (kv_len - q_len) + m.
+            q_pos = torch.arange(q_len, device=q.device)[:, None] + (kv_len - q_len)
+            k_pos = torch.arange(kv_len, device=q.device)[None, :]
+            scores = scores.masked_fill(q_pos < k_pos, float("-inf"))
+        probs = torch.softmax(scores, dim=-1)
+        out[q_lo:q_hi] = torch.matmul(
+            probs, v[k_lo:k_hi].float().transpose(0, 1)
+        ).transpose(0, 1)
+        lse[q_lo:q_hi] = torch.logsumexp(scores, dim=-1).transpose(0, 1)
+    return out, lse
+
+
+@unittest.skipUnless(
+    is_hip() and is_gfx95_supported(), "Kimi-K3 dense Triton prefill requires gfx950"
+)
+class TestKimiK3TritonDensePrefill(unittest.TestCase):
+    def setUp(self):
+        torch.manual_seed(42)
+        torch.cuda.manual_seed_all(42)
+        self.device = get_device()
+        self.scale = D_QK**-0.5
+
+    def _run(self, q_lens, prefix_lens, *, mode="bf16", is_causal=True, want_lse=False):
+        device = self.device
+        kv_lens = [q + p for q, p in zip(q_lens, prefix_lens)]
+        qo_indptr = torch.zeros(len(q_lens) + 1, dtype=torch.int32, device=device)
+        kv_indptr = torch.zeros(len(q_lens) + 1, dtype=torch.int32, device=device)
+        qo_indptr[1:] = torch.tensor(q_lens, device=device).cumsum(0)
+        kv_indptr[1:] = torch.tensor(kv_lens, device=device).cumsum(0)
+        total_q, total_kv = int(qo_indptr[-1]), int(kv_indptr[-1])
+
+        q = torch.randn(total_q, H_Q, D_QK, dtype=torch.bfloat16, device=device) * 0.25
+        k = torch.randn(total_kv, H_Q, D_QK, dtype=torch.bfloat16, device=device) * 0.25
+        v = torch.randn(total_kv, H_Q, D_V, dtype=torch.bfloat16, device=device) * 0.25
+
+        # Quantize before taking the reference so the comparison isolates the
+        # kernel from the cast: an FP8 run should match FP8 inputs exactly.
+        if mode == "fp8":
+            q, k, v = q.to(FP8), k.to(FP8), v.to(FP8)
+        elif mode == "mixed":
+            # What forward_mha_rocm actually hands over on gfx95 with MXFP4
+            # kv_b_proj weights: k/v already e4m3 from the fused up-projection,
+            # q still bf16.
+            k, v = k.to(FP8), v.to(FP8)
+        ref_out, ref_lse = _reference(
+            q.float() if mode == "bf16" else q.to(FP8).float(),
+            k.float() if mode == "bf16" else k.float(),
+            v.float() if mode == "bf16" else v.float(),
+            qo_indptr,
+            kv_indptr,
+            self.scale,
+            is_causal,
+        )
+
+        out = torch.empty(total_q, H_Q, D_V, dtype=torch.bfloat16, device=device)
+        lse = (
+            torch.empty(total_q, H_Q, dtype=torch.float32, device=device)
+            if want_lse
+            else None
+        )
+        dense_prefill_attention_fwd(
+            q,
+            k,
+            v,
+            out,
+            qo_indptr,
+            kv_indptr,
+            max(q_lens),
+            sm_scale=self.scale,
+            is_causal=is_causal,
+            lse=lse,
+        )
+        return out, lse, ref_out, ref_lse
+
+    def test_causal_shapes(self):
+        # Prefix lengths deliberately straddle the BLOCK_N=64 boundary: the
+        # kernel splits its KV sweep into an unmasked interior and a masked
+        # tail at a BLOCK_N multiple, so an off-by-one there is only visible
+        # when the prefix is not a clean multiple.
+        cases = [
+            ([128], [0]),  # no prefix: every block is diagonal
+            ([7], [0]),  # q shorter than BLOCK_M
+            ([1], [1000]),  # single query, long prefix
+            ([128], [63]),
+            ([128], [64]),
+            ([128], [65]),
+            ([256], [1025]),
+            ([100, 37, 256, 51], [0, 500, 300, 77]),  # ragged, mixed prefixes
+        ]
+        for q_lens, prefix_lens in cases:
+            with self.subTest(q=q_lens, prefix=prefix_lens):
+                out, _, ref, _ = self._run(q_lens, prefix_lens)
+                torch.testing.assert_close(out.float(), ref, rtol=2e-2, atol=2e-2)
+
+    def test_non_causal(self):
+        out, _, ref, _ = self._run([128, 64], [256, 130], is_causal=False)
+        torch.testing.assert_close(out.float(), ref, rtol=2e-2, atol=2e-2)
+
+    def test_lse_matches_reference(self):
+        out, lse, ref, ref_lse = self._run([192, 64], [300, 129], want_lse=True)
+        torch.testing.assert_close(out.float(), ref, rtol=2e-2, atol=2e-2)
+        # LSE is what the chunked-KV merge would consume, so it has to be a
+        # natural log in absolute terms, not just proportional.
+        torch.testing.assert_close(lse, ref_lse, rtol=1e-3, atol=1e-3)
+
+    def test_fp8_matches_quantized_reference(self):
+        out, _, ref, _ = self._run([128, 64], [512, 77], mode="fp8")
+        torch.testing.assert_close(out.float(), ref, rtol=5e-2, atol=5e-2)
+
+    def test_mixed_bf16_query_fp8_kv(self):
+        # Both operands of a tl.dot must share a dtype, so the wrapper has to
+        # promote q rather than compile a bf16 x fp8 pair. Regression guard:
+        # this shape reaches the kernel straight from forward_mha_rocm.
+        out, _, ref, _ = self._run([128, 64], [512, 77], mode="mixed")
+        torch.testing.assert_close(out.float(), ref, rtol=5e-2, atol=5e-2)
+
+    def test_fp8_gate(self):
+        device = self.device
+        q = torch.empty(1, H_Q, D_QK, dtype=torch.bfloat16, device=device)
+        k = torch.empty(1, H_Q, D_QK, dtype=torch.bfloat16, device=device)
+        v = torch.empty(1, H_Q, D_V, dtype=torch.bfloat16, device=device)
+
+        with envs.SGLANG_TRITON_FP8_PREFILL_ATTN.override(True):
+            self.assertTrue(
+                can_use_dense_prefill_fp8(q, k, v, is_causal=True, logit_cap=0.0)
+            )
+            # Non-causal and logit-capped softmax are outside the gate, and an
+            # already-quantized input must not be cast a second time.
+            self.assertFalse(
+                can_use_dense_prefill_fp8(q, k, v, is_causal=False, logit_cap=0.0)
+            )
+            self.assertFalse(
+                can_use_dense_prefill_fp8(q, k, v, is_causal=True, logit_cap=1.0)
+            )
+            self.assertFalse(
+                can_use_dense_prefill_fp8(
+                    q.to(FP8), k, v, is_causal=True, logit_cap=0.0
+                )
+            )
+        with envs.SGLANG_TRITON_FP8_PREFILL_ATTN.override(False):
+            self.assertFalse(
+                can_use_dense_prefill_fp8(q, k, v, is_causal=True, logit_cap=0.0)
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/layers/attention/test_triton_mla_prefill_gfx950.py

@@ -0,0 +1,310 @@
+import unittest
+
+import torch
+
+from sglang.kernels.ops.attention.extend_attention import extend_attention_fwd
+from sglang.kernels.ops.attention.extend_attention_split_dim import (
+    can_use_split_dim_absorbed_extend,
+)
+from sglang.srt.environ import envs
+from sglang.srt.utils import get_device, is_gfx95_supported, is_hip
+from sglang.test.ci.ci_register import register_amd_ci
+
+register_amd_ci(est_time=15, suite="stage-b-test-1-gpu-small-amd-mi35x")
+
+
+@unittest.skipUnless(
+    is_hip() and is_gfx95_supported(), "Kimi-K3 Triton prefill requires gfx950"
+)
+class TestKimiK3TritonPrefill(unittest.TestCase):
+    def setUp(self):
+        torch.manual_seed(42)
+        torch.cuda.manual_seed_all(42)
+
+    def test_split_dim_ragged_bf16(self):
+        device = get_device()
+        h_q, h_kv, d_qk, d_v = 12, 1, 576, 512
+        extend_lens = (97, 128)
+        prefix_lens = (53, 91)
+        scale = 192**-0.5
+        total_extend, total_prefix = sum(extend_lens), sum(prefix_lens)
+
+        q = torch.randn(total_extend, h_q, d_qk, dtype=torch.bfloat16, device=device)
+        k = torch.randn(total_extend, h_kv, d_qk, dtype=torch.bfloat16, device=device)
+        v = torch.randn(total_extend, h_kv, d_v, dtype=torch.bfloat16, device=device)
+        k_buffer = torch.randn(
+            total_prefix, h_kv, d_qk, dtype=torch.bfloat16, device=device
+        )
+        v_buffer = torch.randn(
+            total_prefix, h_kv, d_v, dtype=torch.bfloat16, device=device
+        )
+        qo_indptr = torch.tensor(
+            [0, extend_lens[0], total_extend], dtype=torch.int32, device=device
+        )
+        kv_indptr = torch.tensor(
+            [0, prefix_lens[0], total_prefix], dtype=torch.int32, device=device
+        )
+        kv_indices = torch.arange(total_prefix, dtype=torch.int64, device=device)
+        output = torch.empty(
+            total_extend, h_q, d_v, dtype=torch.bfloat16, device=device
+        )
+
+        extend_attention_fwd(
+            q,
+            k,
+            v,
+            output,
+            k_buffer,
+            v_buffer,
+            qo_indptr,
+            kv_indptr,
+            kv_indices,
+            custom_mask=None,
+            is_causal=True,
+            mask_indptr=None,
+            max_len_extend=max(extend_lens),
+            k_scale=1.0,
+            v_scale=1.0,
+            sm_scale=scale,
+        )
+
+        reference = torch.empty_like(output, dtype=torch.float32)
+        for batch, (extend_len, prefix_len) in enumerate(zip(extend_lens, prefix_lens)):
+            q_start = int(qo_indptr[batch])
+            prefix_start = int(kv_indptr[batch])
+            q_batch = q[q_start : q_start + extend_len].float()
+            k_prefix = k_buffer[prefix_start : prefix_start + prefix_len, 0].float()
+            v_prefix = v_buffer[prefix_start : prefix_start + prefix_len, 0].float()
+            k_current = k[q_start : q_start + extend_len, 0].float()
+            v_current = v[q_start : q_start + extend_len, 0].float()
+            causal = torch.triu(
+                torch.ones(
+                    extend_len,
+                    extend_len,
+                    dtype=torch.bool,
+                    device=device,
+                ),
+                diagonal=1,
+            )
+            for head in range(h_q):
+                prefix_scores = q_batch[:, head] @ k_prefix.T * scale
+                current_scores = q_batch[:, head] @ k_current.T * scale
+                current_scores.masked_fill_(causal, float("-inf"))
+                scores = torch.cat([prefix_scores, current_scores], dim=1)
+                values = torch.cat([v_prefix, v_current], dim=0)
+                reference[q_start : q_start + extend_len, head] = (
+                    torch.softmax(scores, dim=-1) @ values
+                )
+
+        torch.testing.assert_close(output.float(), reference, rtol=1e-2, atol=1e-2)
+
+    def test_split_dim_dispatch_gates(self):
+        device = get_device()
+        q = torch.empty(1, 12, 576, dtype=torch.bfloat16, device=device)
+        k = torch.empty(1, 1, 576, dtype=torch.bfloat16, device=device)
+        v = torch.empty(1, 1, 512, dtype=torch.bfloat16, device=device)
+        o = torch.empty(1, 12, 512, dtype=torch.bfloat16, device=device)
+        k_buffer = torch.empty(1, 1, 576, dtype=torch.bfloat16, device=device)
+        v_buffer = torch.empty(1, 1, 512, dtype=torch.bfloat16, device=device)
+        kwargs = dict(
+            lse=None,
+            sinks=None,
+            k_scale=1.0,
+            v_scale=1.0,
+            custom_mask=None,
+            is_causal=True,
+            sliding_window_size=-1,
+            logit_cap=0.0,
+            xai_temperature_len=-1,
+            skip_prefix=False,
+            skip_extend=False,
+            page_size=1,
+            score_mod=None,
+            aux_tensors=None,
+        )
+        self.assertTrue(
+            can_use_split_dim_absorbed_extend(q, k, v, o, k_buffer, v_buffer, **kwargs)
+        )
+
+        fp8_k_buffer = k_buffer.to(torch.float8_e4m3fn)
+        fp8_v_buffer = v_buffer.to(torch.float8_e4m3fn)
+        with envs.SGLANG_TRITON_FP8_PREFILL_ATTN.override(False):
+            self.assertFalse(
+                can_use_split_dim_absorbed_extend(
+                    q, k, v, o, fp8_k_buffer, fp8_v_buffer, **kwargs
+                )
+            )
+        with envs.SGLANG_TRITON_FP8_PREFILL_ATTN.override(True):
+            self.assertTrue(
+                can_use_split_dim_absorbed_extend(
+                    q, k, v, o, fp8_k_buffer, fp8_v_buffer, **kwargs
+                )
+            )
+            self.assertTrue(
+                can_use_split_dim_absorbed_extend(
+                    q,
+                    k,
+                    v,
+                    o,
+                    fp8_k_buffer,
+                    fp8_v_buffer,
+                    **{**kwargs, "k_scale": 0.5, "v_scale": 0.25},
+                )
+            )
+
+        for override in (
+            {"page_size": 2},
+            {"logit_cap": 1.0},
+            {"sliding_window_size": 128},
+            {"skip_prefix": True},
+            {"is_causal": False},
+            {"lse": torch.empty(1, 12, dtype=torch.float32, device=device)},
+            {"sinks": torch.empty(12, dtype=torch.float32, device=device)},
+            {"k_scale": 0.5},
+        ):
+            self.assertFalse(
+                can_use_split_dim_absorbed_extend(
+                    q,
+                    k,
+                    v,
+                    o,
+                    k_buffer,
+                    v_buffer,
+                    **{**kwargs, **override},
+                )
+            )
+
+    def test_zero_prefix_fp8_flag(self):
+        device = get_device()
+        tokens, heads, d_qk, d_v = 128, 12, 192, 128
+        q = torch.randn(tokens, heads, d_qk, dtype=torch.bfloat16, device=device) * 0.25
+        k = torch.randn(tokens, heads, d_qk, dtype=torch.bfloat16, device=device) * 0.25
+        v = torch.randn(tokens, heads, d_v, dtype=torch.bfloat16, device=device) * 0.25
+        k_buffer = torch.empty(1, heads, d_qk, dtype=torch.float8_e4m3fn, device=device)
+        v_buffer = torch.empty(1, heads, d_v, dtype=torch.float8_e4m3fn, device=device)
+        qo_indptr = torch.tensor([0, tokens], dtype=torch.int32, device=device)
+        kv_indptr = torch.tensor([0, 0], dtype=torch.int32, device=device)
+        kv_indices = torch.empty(0, dtype=torch.int64, device=device)
+        bf16_output = torch.empty(
+            tokens, heads, d_v, dtype=torch.bfloat16, device=device
+        )
+        fp8_output = torch.empty_like(bf16_output)
+        scale = d_qk**-0.5
+
+        def run(output):
+            extend_attention_fwd(
+                q,
+                k,
+                v,
+                output,
+                k_buffer,
+                v_buffer,
+                qo_indptr,
+                kv_indptr,
+                kv_indices,
+                custom_mask=None,
+                is_causal=True,
+                mask_indptr=None,
+                max_len_extend=tokens,
+                k_scale=1.0,
+                v_scale=1.0,
+                sm_scale=scale,
+            )
+
+        with envs.SGLANG_TRITON_FP8_PREFILL_ATTN.override(False):
+            run(bf16_output)
+        with envs.SGLANG_TRITON_FP8_PREFILL_ATTN.override(True):
+            run(fp8_output)
+
+        causal = (
+            torch.arange(tokens, device=device)[None, :]
+            <= torch.arange(tokens, device=device)[:, None]
+        )
+
+        def reference(q_ref, k_ref, v_ref):
+            scores = torch.einsum("qhd,khd->qhk", q_ref, k_ref) * scale
+            scores.masked_fill_(~causal[:, None, :], float("-inf"))
+            return torch.einsum("qhk,khd->qhd", torch.softmax(scores, dim=-1), v_ref)
+
+        bf16_reference = reference(q.float(), k.float(), v.float())
+        fp8_reference = reference(
+            q.to(torch.float8_e4m3fn).float(),
+            k.to(torch.float8_e4m3fn).float(),
+            v.to(torch.float8_e4m3fn).float(),
+        )
+        torch.testing.assert_close(
+            bf16_output.float(), bf16_reference, rtol=1e-2, atol=1e-2
+        )
+        torch.testing.assert_close(
+            fp8_output.float(), fp8_reference, rtol=2e-2, atol=2e-2
+        )
+
+    def test_absorbed_fp8_prefix(self):
+        device = get_device()
+        tokens, prefix, heads, d_qk, d_v = 64, 73, 12, 576, 512
+        q = torch.randn(tokens, heads, d_qk, dtype=torch.bfloat16, device=device) * 0.25
+        k = torch.randn(tokens, 1, d_qk, dtype=torch.bfloat16, device=device) * 0.25
+        v = torch.randn(tokens, 1, d_v, dtype=torch.bfloat16, device=device) * 0.25
+        k_buffer = (
+            torch.randn(prefix, 1, d_qk, dtype=torch.bfloat16, device=device) * 0.25
+        ).to(torch.float8_e4m3fn)
+        v_buffer = (
+            torch.randn(prefix, 1, d_v, dtype=torch.bfloat16, device=device) * 0.25
+        ).to(torch.float8_e4m3fn)
+        qo_indptr = torch.tensor([0, tokens], dtype=torch.int32, device=device)
+        kv_indptr = torch.tensor([0, prefix], dtype=torch.int32, device=device)
+        kv_indices = torch.arange(prefix, dtype=torch.int64, device=device)
+        output = torch.empty(tokens, heads, d_v, dtype=torch.bfloat16, device=device)
+        generic_output = torch.empty_like(output)
+        scale, k_scale, v_scale = 192**-0.5, 0.5, 0.25
+
+        def run(candidate):
+            extend_attention_fwd(
+                q,
+                k,
+                v,
+                candidate,
+                k_buffer,
+                v_buffer,
+                qo_indptr,
+                kv_indptr,
+                kv_indices,
+                custom_mask=None,
+                is_causal=True,
+                mask_indptr=None,
+                max_len_extend=tokens,
+                k_scale=k_scale,
+                v_scale=v_scale,
+                sm_scale=scale,
+            )
+
+        with envs.SGLANG_TRITON_FP8_PREFILL_ATTN.override(False):
+            run(generic_output)
+        with envs.SGLANG_TRITON_FP8_PREFILL_ATTN.override(True):
+            run(output)
+
+        q_fp8 = q.to(torch.float8_e4m3fn).float()
+        prefix_scores = (
+            torch.einsum("qhd,kd->qhk", q_fp8, k_buffer[:, 0].float()) * scale * k_scale
+        )
+        current_scores = torch.einsum("qhd,kd->qhk", q.float(), k[:, 0].float())
+        current_scores *= scale
+        causal = (
+            torch.arange(tokens, device=device)[None, :]
+            <= torch.arange(tokens, device=device)[:, None]
+        )
+        current_scores.masked_fill_(~causal[:, None, :], float("-inf"))
+        scores = torch.cat([prefix_scores, current_scores], dim=-1)
+        values = torch.cat([v_buffer[:, 0].float() * v_scale, v[:, 0].float()], dim=0)
+        reference = torch.einsum("qhk,kv->qhv", torch.softmax(scores, dim=-1), values)
+        torch.testing.assert_close(output.float(), reference, rtol=3e-2, atol=3e-2)
+        torch.testing.assert_close(
+            generic_output.float(), reference, rtol=3e-2, atol=3e-2
+        )
+        torch.testing.assert_close(
+            output.float(), generic_output.float(), rtol=1e-2, atol=1e-2
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()