chore: import upstream snapshot with attribution

python/sglang/srt/layers/attention/dsa/dequant_k_cache.py

@@ -0,0 +1,289 @@
+import torch
+import triton
+import triton.language as tl
+
+
+def dequantize_k_cache(quant_k_cache):
+    return _dequantize_k_cache_fast_wrapped(quant_k_cache)
+
+
+def _dequantize_k_cache_ref(
+    quant_k_cache: torch.Tensor,  # (num_blocks, block_size, 1, bytes_per_token)
+    dv: int = 512,
+    tile_size: int = 128,
+    d: int = 576,
+) -> torch.Tensor:
+    """
+    De-quantize the k-cache
+    """
+    assert dv % tile_size == 0
+    original_ndim = quant_k_cache.ndim
+    if original_ndim == 3:
+        # set block_size = 1
+        quant_k_cache = quant_k_cache.unsqueeze(1)
+    num_tiles = dv // tile_size
+    num_blocks, block_size, h_k, _ = quant_k_cache.shape
+    assert h_k == 1
+    result = torch.empty(
+        (num_blocks, block_size, d), dtype=torch.bfloat16, device=quant_k_cache.device
+    )
+
+    quant_k_cache = quant_k_cache.view(num_blocks, block_size, -1)
+
+    input_nope = quant_k_cache[..., :dv]
+    input_scale = quant_k_cache[..., dv : dv + num_tiles * 4].view(torch.float32)
+    input_rope = quant_k_cache[..., dv + num_tiles * 4 :].view(torch.bfloat16)
+    result[..., dv:] = input_rope
+
+    for tile_idx in range(0, num_tiles):
+        cur_nope = input_nope[
+            ..., tile_idx * tile_size : (tile_idx + 1) * tile_size
+        ].to(torch.float32)
+        cur_scales = input_scale[..., tile_idx].unsqueeze(-1)
+        result[..., tile_idx * tile_size : (tile_idx + 1) * tile_size] = (
+            cur_nope * cur_scales
+        )
+
+    if original_ndim == 3:
+        return result.view(num_blocks, 1, -1)
+    else:
+        return result.view(num_blocks, block_size, 1, -1)
+
+
+def _dequantize_k_cache_fast_wrapped(
+    quant_k_cache: torch.Tensor,
+    dv: int = 512,
+    tile_size: int = 128,
+) -> torch.Tensor:
+    original_ndim = quant_k_cache.ndim
+    if original_ndim == 3:
+        # set block_size = 1
+        quant_k_cache = quant_k_cache.unsqueeze(1)
+    num_blocks, block_size, _, dim_quant = quant_k_cache.shape
+    assert dv == 512
+    assert dim_quant == 656
+    assert tile_size == 128
+    quant_k_cache = quant_k_cache.view((-1, dim_quant))
+
+    output = _dequantize_k_cache_fast(quant_k_cache)
+
+    if original_ndim == 3:
+        return output.view(num_blocks, 1, -1)
+    else:
+        return output.view(num_blocks, block_size, 1, -1)
+
+
+def _dequantize_k_cache_fast(quant_k_cache, group_size: int = 128):
+    num_tokens, dim_quant = quant_k_cache.shape
+
+    assert quant_k_cache.dtype == torch.float8_e4m3fn
+    dim_nope = 512
+    dim_rope = 64
+    num_tiles = dim_nope // group_size
+    assert dim_quant == 656
+
+    output = torch.empty(
+        (num_tokens, dim_nope + dim_rope),
+        dtype=torch.bfloat16,
+        device=quant_k_cache.device,
+    )
+
+    num_blocks_per_token = triton.cdiv(dim_nope + dim_rope, group_size)
+    assert num_blocks_per_token == 5
+
+    assert dim_nope % group_size == 0
+
+    input_nope_q = quant_k_cache[:, :dim_nope]
+    input_nope_s = quant_k_cache[:, dim_nope : dim_nope + num_tiles * 4].view(
+        torch.float32
+    )
+    input_rope = quant_k_cache[:, dim_nope + num_tiles * 4 :].view(torch.bfloat16)
+
+    _dequantize_k_cache_fast_kernel[(num_tokens, num_blocks_per_token)](
+        output,
+        input_nope_q,
+        input_nope_s,
+        input_rope,
+        output.stride(0),
+        input_nope_q.stride(0),
+        input_nope_s.stride(0),
+        input_rope.stride(0),
+        NUM_NOPE_BLOCKS=num_tiles,
+        GROUP_SIZE=group_size,
+        DIM_NOPE=dim_nope,
+        DIM_ROPE=dim_rope,
+    )
+
+    return output
+
+
+@triton.jit
+def _dequantize_k_cache_fast_kernel(
+    output_ptr,
+    input_nope_q_ptr,
+    input_nope_s_ptr,
+    input_rope_ptr,
+    output_stride_0: int,
+    input_nope_q_stride_0: int,
+    input_nope_s_stride_0: int,
+    input_rope_stride_0: int,
+    NUM_NOPE_BLOCKS: tl.constexpr,
+    GROUP_SIZE: tl.constexpr,
+    DIM_NOPE: tl.constexpr,
+    DIM_ROPE: tl.constexpr,
+):
+    token_id = tl.program_id(0)
+    raw_block_id = tl.program_id(1)
+
+    if raw_block_id < NUM_NOPE_BLOCKS:
+        # a. dequant nope
+        effective_block_id = raw_block_id
+
+        offs_q = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
+        mask = offs_q < DIM_NOPE
+        ptr_q = input_nope_q_ptr + token_id * input_nope_q_stride_0 + offs_q
+        ptr_s = input_nope_s_ptr + token_id * input_nope_s_stride_0 + effective_block_id
+
+        y_q = tl.load(ptr_q, mask=mask, other=0.0).to(tl.float32)
+        y_s = tl.load(ptr_s)
+
+        y = (y_q * y_s).to(output_ptr.dtype.element_ty)
+
+        dst_ptr = output_ptr + token_id * output_stride_0 + offs_q
+        tl.store(dst_ptr, y, mask=mask)
+    else:
+        # b. copy rope
+        effective_block_id = raw_block_id - NUM_NOPE_BLOCKS
+
+        offs = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
+        mask = offs < DIM_ROPE
+
+        src_ptr = input_rope_ptr + token_id * input_rope_stride_0 + offs
+        dst_ptr = output_ptr + token_id * output_stride_0 + DIM_NOPE + offs
+
+        data = tl.load(src_ptr, mask=mask).to(tl.bfloat16)
+        tl.store(dst_ptr, data, mask=mask)
+
+
+def dequantize_k_cache_paged(
+    quant_k_cache: torch.Tensor,
+    page_table_1_flattened: torch.Tensor,
+    group_size: int = 128,
+) -> torch.Tensor:
+    """
+    De-quantize the k-cache with paged layout
+    Args:
+        quant_k_cache: [total_num_tokens, 1, dim_quant] or [num_blocks, block_size, 1, dim_quant], the quantized k-cache in paged layout
+        page_table_1_flattened: [num_tokens], the flattened page_table_1 with the page indices in each requests concatenated together
+    Returns:
+        output: [num_tokens, 1, dim_nope + dim_rope], the de-quantized k-cache
+    """
+    dim_quant = quant_k_cache.shape[-1]
+    assert (
+        dim_quant == 656
+    ), f"dim_quant: {dim_quant} != 656 detected in dequantize_k_cache_paged"
+    quant_k_cache = quant_k_cache.view((-1, dim_quant))
+
+    # num_tokens can exceed kv_cache_size due to prefix sharing (multiple seqs share same KV slots)
+    # Index bounds validated in dsa_backend.init_forward_metadata
+    num_tokens = page_table_1_flattened.shape[0]
+    assert quant_k_cache.dtype == torch.float8_e4m3fn
+    dim_nope = 512
+    dim_rope = 64
+    num_tiles = dim_nope // group_size  # 512 // 128 = 4
+
+    output = torch.empty(
+        (num_tokens, 1, dim_nope + dim_rope),
+        dtype=torch.bfloat16,
+        device=quant_k_cache.device,
+    )
+
+    # cdiv(512 + 64, 128) = 5
+    num_blocks_per_token = triton.cdiv(dim_nope + dim_rope, group_size)
+    assert num_blocks_per_token == 5
+
+    assert dim_nope % group_size == 0
+
+    input_nope_q = quant_k_cache[:, :dim_nope]
+    # [:, 512:512+4*4] = [:, 512:528]
+    input_nope_s = quant_k_cache[:, dim_nope : dim_nope + num_tiles * 4].view(
+        torch.float32
+    )
+    # [:, 528:]
+    input_rope = quant_k_cache[:, dim_nope + num_tiles * 4 :].view(torch.bfloat16)
+
+    _dequantize_k_cache_paged_kernel[(num_tokens, num_blocks_per_token)](
+        output,
+        input_nope_q,
+        input_nope_s,
+        input_rope,
+        page_table_1_flattened,
+        output.stride(0),
+        input_nope_q.stride(0),
+        input_nope_s.stride(0),
+        input_rope.stride(0),
+        NUM_NOPE_BLOCKS=num_tiles,
+        GROUP_SIZE=group_size,
+        DIM_NOPE=dim_nope,
+        DIM_ROPE=dim_rope,
+    )
+
+    return output
+
+
+@triton.jit
+def _dequantize_k_cache_paged_kernel(
+    output_ptr,
+    input_nope_q_ptr,
+    input_nope_s_ptr,
+    input_rope_ptr,
+    page_table_1_ptr,
+    output_stride_0: int,
+    input_nope_q_stride_0: int,
+    input_nope_s_stride_0: int,
+    input_rope_stride_0: int,
+    NUM_NOPE_BLOCKS: tl.constexpr,
+    GROUP_SIZE: tl.constexpr,
+    DIM_NOPE: tl.constexpr,
+    DIM_ROPE: tl.constexpr,
+):
+    token_id = tl.program_id(0)
+    token_id_paged = tl.load(page_table_1_ptr + token_id).to(tl.int32)
+    raw_block_id = tl.program_id(1)
+
+    if raw_block_id < NUM_NOPE_BLOCKS:
+        # a. dequant nope
+        effective_block_id = raw_block_id
+
+        offs_q = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
+        mask = offs_q < DIM_NOPE
+        ptr_q = input_nope_q_ptr + token_id_paged * input_nope_q_stride_0 + offs_q
+        ptr_s = (
+            input_nope_s_ptr
+            + token_id_paged * input_nope_s_stride_0
+            + effective_block_id
+        )
+
+        y_q = tl.load(ptr_q, mask=mask, other=0.0).to(tl.float32)
+        y_s = tl.load(ptr_s)
+
+        y = (y_q * y_s).to(output_ptr.dtype.element_ty)
+
+        dst_ptr = output_ptr + token_id * output_stride_0 + offs_q
+        tl.store(dst_ptr, y, mask=mask)
+    else:
+        # b. copy rope
+        effective_block_id = raw_block_id - NUM_NOPE_BLOCKS
+
+        offs = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
+        mask = offs < DIM_ROPE
+
+        src_ptr = input_rope_ptr + token_id_paged * input_rope_stride_0 + offs
+        dst_ptr = output_ptr + token_id * output_stride_0 + DIM_NOPE + offs
+
+        data = tl.load(src_ptr, mask=mask).to(tl.bfloat16)
+        tl.store(dst_ptr, data, mask=mask)
+
+
+if __name__ == "__main__":
+    raise Exception("UT is in quant_k_cache.py")

python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py

@@ -0,0 +1,387 @@
+"""Multi-step precompute utilities for Native Sparse Attention backend.
+
+This module provides optimization utilities for multi-step speculative decoding
+by precomputing shared metadata once and copying it to multiple backend instances.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.layers.attention.dsa.utils import compute_dsa_seqlens
+from sglang.srt.layers.attention.utils import seqlens_expand_triton
+from sglang.srt.utils import is_cuda, is_hip
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardMode
+
+_is_cuda = is_cuda()
+_is_hip = is_hip()
+_USE_FUSED_METADATA_GENERATION = (
+    envs.SGLANG_DSA_USE_FUSED_METADATA_GENERATION.get() and not _is_hip
+)
+
+
+@dataclass
+class PrecomputedMetadata:
+    """Precomputed metadata shared across multiple backend instances.
+
+    Used for multi-step speculative decoding where multiple backends
+    need identical metadata. Precomputing once and copying N times
+    is much faster than computing N times.
+
+    """
+
+    # Basic seqlens
+    cache_seqlens: torch.Tensor  # int32, [bs]
+    cu_seqlens_k: torch.Tensor  # int32, [bs+1]
+
+    # Page table
+    page_indices: torch.Tensor  # int32, [bs, max_len] or [expanded_bs, max_len]
+    real_page_table: Optional[torch.Tensor]  # int32, transformed version
+
+    # DSA seqlens
+    seqlens_expanded: torch.Tensor  # int32, [expanded_size]
+    dsa_cache_seqlens: torch.Tensor  # int32, [expanded_size]
+    dsa_cu_seqlens_k: torch.Tensor  # int32, [expanded_size+1]
+    seqlens_expanded_size: int
+
+    # Dimensions
+    max_len: int  # for decode/draft_extend
+    max_seqlen_k: int  # for target_verify
+
+    # FlashMLA (optional)
+    flashmla_metadata: Optional[torch.Tensor] = None
+
+
+def compute_cu_seqlens(seqlens: torch.Tensor) -> torch.Tensor:
+    """Compute cumulative sequence lengths with padding."""
+    assert seqlens.dtype == torch.int32
+    return torch.nn.functional.pad(
+        torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0)
+    )
+
+
+class DeepseekSparseAttnBackendMTPPrecomputeMixin:
+    """Mixin class providing metadata precomputation for multi-step speculative decoding.
+
+    This mixin provides the _precompute_replay_metadata method and its helpers,
+    which are used to optimize CUDA graph replay in multi-step scenarios.
+    """
+
+    def _precompute_replay_metadata(
+        self,
+        bs: int,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        forward_mode: ForwardMode,
+    ) -> PrecomputedMetadata:
+        """Precompute all shared metadata for multi-step backends.
+
+        This function extracts and computes all operations that are
+        identical across different backend instances in multi-step
+        speculative decoding.
+
+        Args:
+            bs: Batch size
+            req_pool_indices: Request pool indices [bs]
+            seq_lens: Sequence lengths [bs]
+            seq_lens_cpu: Sequence lengths on CPU [bs]
+            forward_mode: Forward mode (decode/target_verify)
+
+        Returns:
+            PrecomputedMetadata containing all shared intermediate results
+        """
+        # Slice inputs to batch size
+        seq_lens = seq_lens[:bs]
+        if seq_lens_cpu is not None:
+            seq_lens_cpu = seq_lens_cpu[:bs]
+        req_pool_indices = req_pool_indices[:bs]
+
+        # Dispatch to mode-specific precomputation
+        if forward_mode.is_decode_or_idle():
+            return self._precompute_decode_mode(
+                bs, req_pool_indices, seq_lens, seq_lens_cpu
+            )
+        elif forward_mode.is_target_verify():
+            return self._precompute_target_verify_mode(
+                bs, req_pool_indices, seq_lens, seq_lens_cpu
+            )
+        else:
+            raise ValueError(f"Unsupported forward mode: {forward_mode}")
+
+    def _precompute_decode_mode(
+        self,
+        bs: int,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+    ) -> PrecomputedMetadata:
+        """Precompute metadata for normal decode mode."""
+        max_len = self.decode_cuda_graph_metadata[bs].page_table_1.shape[1]
+
+        if _USE_FUSED_METADATA_GENERATION and _is_cuda and not _is_hip:
+            from sglang.kernels.ops.attention.dsa_metadata import (
+                fused_dsa_decode_metadata,
+            )
+
+            cache_seqlens = torch.empty(bs, dtype=torch.int32, device=self.device)
+            cu_seqlens_k = torch.empty(bs + 1, dtype=torch.int32, device=self.device)
+            page_indices = torch.empty(
+                (bs, max_len), dtype=torch.int32, device=self.device
+            )
+            dsa_cache_seqlens = torch.empty(bs, dtype=torch.int32, device=self.device)
+            dsa_cu_seqlens_k = torch.empty(
+                bs + 1, dtype=torch.int32, device=self.device
+            )
+            if self.real_page_size > 1:
+                real_cols = (max_len + self.real_page_size - 1) // self.real_page_size
+                real_page_table = torch.empty(
+                    (bs, real_cols), dtype=torch.int32, device=self.device
+                )
+                real_page_table_arg = real_page_table
+            else:
+                real_page_table = None
+                real_page_table_arg = page_indices
+
+            fused_dsa_decode_metadata(
+                seq_lens=seq_lens,
+                req_pool_indices=req_pool_indices,
+                req_to_token=self.req_to_token,
+                cache_seqlens=cache_seqlens,
+                cu_seqlens_k=cu_seqlens_k,
+                page_table_1=page_indices,
+                dsa_cache_seqlens=dsa_cache_seqlens,
+                dsa_cu_seqlens_k=dsa_cu_seqlens_k,
+                real_page_table=real_page_table_arg,
+                bs=bs,
+                max_len=max_len,
+                dsa_index_topk=self.dsa_index_topk,
+                real_page_size=self.real_page_size,
+            )
+            seqlens_expanded = cache_seqlens
+            seqlens_expanded_size = bs
+
+            flashmla_metadata = None
+            if self.dsa_decode_impl == "flashmla_kv":
+                flashmla_metadata = self._compute_flashmla_metadata(
+                    cache_seqlens=dsa_cache_seqlens,
+                    seq_len_q=1,
+                )
+
+            return PrecomputedMetadata(
+                cache_seqlens=cache_seqlens,
+                cu_seqlens_k=cu_seqlens_k,
+                page_indices=page_indices,
+                real_page_table=real_page_table,
+                seqlens_expanded=seqlens_expanded,
+                dsa_cache_seqlens=dsa_cache_seqlens,
+                dsa_cu_seqlens_k=dsa_cu_seqlens_k,
+                seqlens_expanded_size=seqlens_expanded_size,
+                max_len=max_len,
+                max_seqlen_k=max_len,
+                flashmla_metadata=flashmla_metadata,
+            )
+
+        # Convert to int32 and compute cumsum
+        cache_seqlens = seq_lens.to(torch.int32)
+        cu_seqlens_k = compute_cu_seqlens(cache_seqlens)
+
+        # Get page indices from cache
+        page_indices = self.req_to_token[req_pool_indices, :max_len].contiguous()
+
+        # Compute DSA seqlens
+        dsa_cache_seqlens = compute_dsa_seqlens(
+            cache_seqlens, dsa_index_topk=self.dsa_index_topk
+        )
+        seqlens_expanded = cache_seqlens
+        seqlens_expanded_size = seqlens_expanded.shape[0]
+
+        # Compute DSA cumsum
+        dsa_cu_seqlens_k = compute_cu_seqlens(dsa_cache_seqlens)
+
+        # Transform page table if needed
+        if self.real_page_size > 1:
+            real_page_table = self._transform_table_1_to_real(page_indices)
+        else:
+            real_page_table = None  # Will use page_indices directly
+
+        # Compute FlashMLA metadata if needed
+        flashmla_metadata = None
+        if self.dsa_decode_impl == "flashmla_kv":
+            flashmla_metadata = self._compute_flashmla_metadata(
+                cache_seqlens=dsa_cache_seqlens,
+                seq_len_q=1,
+            )
+
+        return PrecomputedMetadata(
+            cache_seqlens=cache_seqlens,
+            cu_seqlens_k=cu_seqlens_k,
+            page_indices=page_indices,
+            real_page_table=real_page_table,
+            seqlens_expanded=seqlens_expanded,
+            dsa_cache_seqlens=dsa_cache_seqlens,
+            dsa_cu_seqlens_k=dsa_cu_seqlens_k,
+            seqlens_expanded_size=seqlens_expanded_size,
+            max_len=max_len,
+            max_seqlen_k=max_len,
+            flashmla_metadata=flashmla_metadata,
+        )
+
+    def _precompute_target_verify_mode(
+        self,
+        bs: int,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+    ) -> PrecomputedMetadata:
+        """Precompute metadata for target verify mode."""
+        max_seqlen_k = self.decode_cuda_graph_metadata[bs].page_table_1.shape[1]
+        seqlens_expanded_size = bs * self.speculative_num_draft_tokens
+
+        if _USE_FUSED_METADATA_GENERATION and _is_cuda and not _is_hip:
+            from sglang.kernels.ops.attention.dsa_metadata import (
+                fused_dsa_target_verify_metadata,
+            )
+
+            cache_seqlens = torch.empty(bs, dtype=torch.int32, device=self.device)
+            cu_seqlens_k = torch.empty(bs + 1, dtype=torch.int32, device=self.device)
+            page_indices = torch.empty(
+                (seqlens_expanded_size, max_seqlen_k),
+                dtype=torch.int32,
+                device=self.device,
+            )
+            seqlens_expanded = torch.empty(
+                seqlens_expanded_size, dtype=torch.int32, device=self.device
+            )
+            dsa_cache_seqlens = torch.empty(
+                seqlens_expanded_size, dtype=torch.int32, device=self.device
+            )
+            dsa_cu_seqlens_k = torch.empty(
+                seqlens_expanded_size + 1,
+                dtype=torch.int32,
+                device=self.device,
+            )
+            if self.real_page_size > 1:
+                real_cols = (
+                    max_seqlen_k + self.real_page_size - 1
+                ) // self.real_page_size
+                real_page_table = torch.empty(
+                    (seqlens_expanded_size, real_cols),
+                    dtype=torch.int32,
+                    device=self.device,
+                )
+                real_page_table_arg = real_page_table
+            else:
+                real_page_table = None
+                real_page_table_arg = page_indices
+
+            fused_dsa_target_verify_metadata(
+                seq_lens=seq_lens,
+                req_pool_indices=req_pool_indices,
+                req_to_token=self.req_to_token,
+                cache_seqlens=cache_seqlens,
+                cu_seqlens_k=cu_seqlens_k,
+                page_table_1=page_indices,
+                seqlens_expanded=seqlens_expanded,
+                dsa_cache_seqlens=dsa_cache_seqlens,
+                dsa_cu_seqlens_k=dsa_cu_seqlens_k,
+                real_page_table=real_page_table_arg,
+                bs=bs,
+                max_seqlen_k=max_seqlen_k,
+                dsa_index_topk=self.dsa_index_topk,
+                real_page_size=self.real_page_size,
+                next_n=self.speculative_num_draft_tokens,
+            )
+
+            flashmla_metadata = None
+            if self.dsa_decode_impl == "flashmla_kv":
+                flashmla_metadata = self._compute_flashmla_metadata(
+                    cache_seqlens=dsa_cache_seqlens,
+                    seq_len_q=1,
+                )
+
+            return PrecomputedMetadata(
+                cache_seqlens=cache_seqlens,
+                cu_seqlens_k=cu_seqlens_k,
+                page_indices=page_indices,
+                real_page_table=real_page_table,
+                seqlens_expanded=seqlens_expanded,
+                dsa_cache_seqlens=dsa_cache_seqlens,
+                dsa_cu_seqlens_k=dsa_cu_seqlens_k,
+                seqlens_expanded_size=seqlens_expanded_size,
+                max_len=-1,
+                max_seqlen_k=max_seqlen_k,
+                flashmla_metadata=flashmla_metadata,
+            )
+
+        # Cache seqlens with draft tokens
+        cache_seqlens = (seq_lens + self.speculative_num_draft_tokens).to(torch.int32)
+        cu_seqlens_k = compute_cu_seqlens(cache_seqlens)
+
+        # Page indices (repeated for each draft token)
+        page_indices = self.req_to_token[req_pool_indices, :max_seqlen_k]
+        page_indices = torch.repeat_interleave(
+            page_indices, repeats=self.speculative_num_draft_tokens, dim=0
+        ).contiguous()
+
+        # Generate expanded seqlens on device. seq_lens_cpu is optional for DSA
+        # CUDA graph replay, so this fallback must not require a host mirror.
+        extend_seq_lens = torch.full(
+            (bs,),
+            self.speculative_num_draft_tokens,
+            dtype=torch.int32,
+            device=self.device,
+        )
+        seqlens_expanded = seqlens_expand_triton(
+            extend_seq_lens,
+            cache_seqlens,
+            bs * self.speculative_num_draft_tokens,
+            self.speculative_num_draft_tokens,
+        )
+
+        # Compute DSA seqlens
+        dsa_cache_seqlens = compute_dsa_seqlens(seqlens_expanded, self.dsa_index_topk)
+        seqlens_expanded_size = seqlens_expanded.shape[0]
+
+        # DSA cumsum
+        dsa_cu_seqlens_k = compute_cu_seqlens(dsa_cache_seqlens)
+
+        # Transform page table
+        if self.real_page_size > 1:
+            real_page_table = self._transform_table_1_to_real(page_indices)
+        else:
+            real_page_table = None
+
+        # FlashMLA metadata
+        flashmla_metadata = None
+        if self.dsa_decode_impl == "flashmla_kv":
+            flashmla_metadata = self._compute_flashmla_metadata(
+                cache_seqlens=dsa_cache_seqlens,
+                seq_len_q=1,
+            )
+
+        return PrecomputedMetadata(
+            cache_seqlens=cache_seqlens,
+            cu_seqlens_k=cu_seqlens_k,
+            page_indices=page_indices,
+            real_page_table=real_page_table,
+            seqlens_expanded=seqlens_expanded,
+            dsa_cache_seqlens=dsa_cache_seqlens,
+            dsa_cu_seqlens_k=dsa_cu_seqlens_k,
+            seqlens_expanded_size=seqlens_expanded_size,
+            max_len=-1,  # Not used in this mode
+            max_seqlen_k=max_seqlen_k,
+            flashmla_metadata=flashmla_metadata,
+        )
+
+
+# Backward-compat alias
+DeepseekSparseAttnBackendMTPPrecomputeMixin = (
+    DeepseekSparseAttnBackendMTPPrecomputeMixin
+)

python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py

@@ -0,0 +1,367 @@
+from __future__ import annotations
+
+from enum import Enum, IntEnum, auto
+from typing import Callable, Dict, List, Optional, Tuple
+
+import torch
+
+from sglang.srt.environ import envs
+
+_FLASHINFER_TIE_BREAK_VALUES = {
+    "small": 1,
+    "large": 2,
+}
+
+
+class TopkTransformMethod(IntEnum):
+    # Transform topk indices to indices to the page table (page_size = 1)
+    PAGED = auto()
+    # Transform topk indices to indices to ragged kv (non-paged)
+    RAGGED = auto()
+
+
+class DSATopKBackend(Enum):
+    SGL_KERNEL = "sgl-kernel"
+    TORCH = "torch"
+    FLASHINFER = "flashinfer"
+
+    def is_sgl_kernel(self) -> bool:
+        return self == DSATopKBackend.SGL_KERNEL
+
+    def is_torch(self) -> bool:
+        return self == DSATopKBackend.TORCH
+
+    def is_flashinfer(self) -> bool:
+        return self == DSATopKBackend.FLASHINFER
+
+    def topk_func(
+        self,
+        score: torch.Tensor,
+        lengths: torch.Tensor,
+        topk: int,
+        row_starts: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if self.is_sgl_kernel():
+            from sgl_kernel import fast_topk_v2
+
+            return fast_topk_v2(score, lengths, topk, row_starts=row_starts)
+        if self.is_torch():
+            return _topk_unfused(
+                score,
+                lengths,
+                topk,
+                row_starts=row_starts,
+                topk_op=torch.topk,
+                topk_op_kwargs={"dim": -1},
+            )
+        if self.is_flashinfer():
+            import flashinfer
+
+            return _topk_unfused(
+                score,
+                lengths,
+                topk,
+                row_starts=row_starts,
+                topk_op=flashinfer.top_k,
+                topk_op_kwargs={
+                    "sorted": False,
+                    "deterministic": envs.SGLANG_DSA_TOPK_FLASHINFER_DETERMINISTIC.get(),
+                    "tie_break": _flashinfer_tie_break_value(),
+                    "dsa_graph_safe": True,
+                },
+            )
+        raise RuntimeError(f"Unsupported {self = }.")
+
+    def topk_transform(
+        self,
+        logits: torch.Tensor,
+        lengths: torch.Tensor,
+        topk: int,
+        topk_transform_method: TopkTransformMethod,
+        attn_metadata,
+        cu_seqlens_q_topk: Optional[torch.Tensor] = None,
+        topk_indices_offset: Optional[torch.Tensor] = None,
+        row_starts: Optional[torch.Tensor] = None,
+        batch_idx_list: Optional[List[int]] = None,
+        force_unfused_topk: bool = False,
+    ) -> torch.Tensor:
+        if not envs.SGLANG_DSA_FUSE_TOPK.get() or force_unfused_topk:
+            return self.topk_func(logits, lengths, topk, row_starts=row_starts)
+
+        # Decode-shaped PAGED top-k (plain decode AND spec verify / draft-extend,
+        # whose expanded rows match the same shape) routes to the DeepSeek-V4 top-k
+        # v2 JIT kernel, which fuses top-k selection and the page-table transform in
+        # one launch and consumes the indexer's own page_size>=1 table directly, so
+        # no page_size=1 table is materialized. Shared by DeepSeek-V3.2 and GLM DSA.
+        # This is a deterministic dispatch on the work shape, not a best-effort
+        # attempt: the fused-decode CUDA graph drops the page_size=1 table for
+        # exactly this case (see dsa_drop_wide_page_table), so once the shape
+        # matches we commit to v2 and never silently fall back to the legacy
+        # page_size=1 path from here.
+        if (
+            envs.SGLANG_OPT_USE_TOPK_V2.get()
+            and topk_transform_method == TopkTransformMethod.PAGED
+            and row_starts is None
+            and batch_idx_list is None
+            and 0 < topk <= 2048
+            and lengths.shape[0]
+            == logits.shape[0]
+            == attn_metadata.real_page_table.shape[0]
+        ):
+            return _topk_transform_v2_paged(logits, lengths, topk, attn_metadata)
+
+        # The legacy transforms below read attn_metadata.page_table_1 (page_size=1),
+        # which is always present here: the fold only drops it for the decode case
+        # dispatched to v2 above.
+        assert attn_metadata.page_table_1 is not None
+
+        if self.is_sgl_kernel():
+            from sgl_kernel import (
+                fast_topk_transform_fused,
+                fast_topk_transform_ragged_fused,
+            )
+
+            if topk_transform_method == TopkTransformMethod.PAGED:
+                page_table_size_1 = (
+                    attn_metadata.page_table_1[batch_idx_list]
+                    if batch_idx_list is not None
+                    else attn_metadata.page_table_1
+                )
+                return fast_topk_transform_fused(
+                    score=logits,
+                    lengths=lengths,
+                    page_table_size_1=page_table_size_1,
+                    cu_seqlens_q=cu_seqlens_q_topk,
+                    topk=topk,
+                    row_starts=row_starts,
+                )
+            if topk_transform_method == TopkTransformMethod.RAGGED:
+                if topk_indices_offset is None:
+                    raise RuntimeError(
+                        "RAGGED topk_transform requires topk_indices_offset; "
+                        "expected extend-without-speculative metadata."
+                    )
+                return fast_topk_transform_ragged_fused(
+                    score=logits,
+                    lengths=lengths,
+                    topk_indices_offset=topk_indices_offset,
+                    topk=topk,
+                    row_starts=row_starts,
+                )
+            raise RuntimeError(f"Unsupported {topk_transform_method = }.")
+
+        if self.is_flashinfer():
+            import flashinfer
+
+            if topk_transform_method == TopkTransformMethod.PAGED:
+                row_to_batch, local_row_starts = _build_flashinfer_paged_args(
+                    attn_metadata=attn_metadata,
+                    row_starts=row_starts,
+                    cu_seqlens_q_topk=cu_seqlens_q_topk,
+                    batch_idx_list=batch_idx_list,
+                    device=logits.device,
+                    num_rows=logits.shape[0],
+                )
+                return flashinfer.top_k_page_table_transform(
+                    logits.contiguous(),
+                    attn_metadata.page_table_1.contiguous(),
+                    lengths.contiguous(),
+                    topk,
+                    row_to_batch=row_to_batch,
+                    deterministic=envs.SGLANG_DSA_TOPK_FLASHINFER_DETERMINISTIC.get(),
+                    tie_break=_flashinfer_tie_break_value(),
+                    dsa_graph_safe=True,
+                    row_starts=local_row_starts,
+                )
+            if topk_transform_method == TopkTransformMethod.RAGGED:
+                if topk_indices_offset is None:
+                    raise RuntimeError(
+                        "RAGGED topk_transform requires topk_indices_offset; "
+                        "expected extend-without-speculative metadata."
+                    )
+                return flashinfer.top_k_ragged_transform(
+                    logits.contiguous(),
+                    topk_indices_offset.contiguous(),
+                    lengths.contiguous(),
+                    topk,
+                    deterministic=envs.SGLANG_DSA_TOPK_FLASHINFER_DETERMINISTIC.get(),
+                    tie_break=_flashinfer_tie_break_value(),
+                    dsa_graph_safe=True,
+                    row_starts=row_starts,
+                )
+            raise RuntimeError(f"Unsupported {topk_transform_method = }.")
+
+        raise RuntimeError(f"Unsupported {self = } for SGLANG_DSA_FUSE_TOPK.")
+
+
+def _topk_unfused(
+    score: torch.Tensor,
+    lengths: torch.Tensor,
+    topk: int,
+    row_starts: Optional[torch.Tensor] = None,
+    topk_op: Callable[..., Tuple[torch.Tensor, torch.Tensor]] = torch.topk,
+    topk_op_kwargs: Optional[Dict[str, object]] = None,
+) -> torch.Tensor:
+    batch_size, max_score_len = score.shape
+    topk_indices = score.new_full((batch_size, topk), -1, dtype=torch.int32)
+    if batch_size == 0 or topk == 0 or max_score_len == 0:
+        return topk_indices
+
+    if row_starts is None:
+        row_starts = torch.zeros_like(lengths, dtype=torch.int32, device=score.device)
+    else:
+        row_starts = row_starts.to(dtype=torch.int32, device=score.device)
+    lengths = lengths.to(dtype=torch.int32, device=score.device)
+
+    col_indices = torch.arange(max_score_len, dtype=torch.int32, device=score.device)
+    col_indices = col_indices.unsqueeze(0)
+    row_starts_unsqueezed = row_starts.unsqueeze(1)
+    row_ends_unsqueezed = (row_starts + lengths).unsqueeze(1)
+    valid_mask = (col_indices >= row_starts_unsqueezed) & (
+        col_indices < row_ends_unsqueezed
+    )
+
+    masked_logits = score.masked_fill(~valid_mask, float("-inf"))
+    valid_topk = min(topk, max_score_len)
+    topk_kwargs = topk_op_kwargs or {}
+    topk_scores, topk_col_indices = topk_op(masked_logits, valid_topk, **topk_kwargs)
+    topk_local_indices = topk_col_indices.to(torch.int32) - row_starts_unsqueezed
+    topk_local_indices = topk_local_indices.masked_fill(
+        topk_scores == float("-inf"), -1
+    )
+    topk_indices[:, :valid_topk] = topk_local_indices
+
+    return topk_indices
+
+
+def _topk_transform_v2_paged(
+    logits: torch.Tensor,
+    lengths: torch.Tensor,
+    topk: int,
+    attn_metadata,
+) -> torch.Tensor:
+    """Fused top-k + page-table transform via the DeepSeek-V4 v2 JIT kernel.
+
+    Returns the transformed page indices ``(num_rows, topk)`` int32 (physical
+    page_size=1 KV slots, ``-1`` padded) -- identical in meaning to
+    ``fast_topk_transform_fused`` / ``flashinfer.top_k_page_table_transform``.
+    The kernel selects, per row, the top-k of ``logits[row, :lengths[row]]`` and
+    maps each selected position ``p`` through the page table as
+    ``real_page_table[row, p // page_size] * page_size + (p % page_size)``. Feeding
+    it the indexer's compact ``real_page_table`` (page_size = pool page size,
+    typically 64) yields the same physical slots as gathering the page_size=1
+    table, without materializing that wide table.
+
+    This is a committed contract, not a best-effort path: ``topk_transform`` routes
+    here only for the decode-shaped PAGED case, and the fused-decode CUDA graph
+    drops the page_size=1 table for exactly this case (see
+    ``dsa_drop_wide_page_table``). The preconditions below are therefore
+    invariants the caller must uphold -- they assert (raise) on violation rather
+    than fall back to the slow legacy path (which may not even have a page_size=1
+    table to fall back to) or silently paper over bad input (padding, recomputing
+    the plan) at the cost of the performance this path exists to deliver.
+
+    ``lengths`` entries must be NON-NEGATIVE: the kernel reads them as
+    ``uint32_t``, so a negative row length (DP-padded / idle-companion rows)
+    reinterprets as ~4e9 tokens and illegal-addresses. Metadata producers clamp
+    padded rows to 0 (see ``fused_dsa_draft_extend_metadata`` /
+    ``seqlens_expand_kernel``); 0 takes the trivial all-(-1) output path.
+    """
+    from sglang.jit_kernel.dsv4.topk import topk_transform_512_v2
+    from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
+
+    num_rows = logits.shape[0]
+
+    # The indexer (DeepGEMM) emits fp32 scores with unit row stride and a 16B-aligned
+    # row stride (a multiple of 4), which is exactly the kernel's ABI (it checks
+    # score_stride % 4 == 0 with strides {S, 1}). This holds even though the scores
+    # may be a padded view (stride(0) > width, so not `is_contiguous()`); assert the
+    # real requirement rather than force a contiguous copy of the wide score buffer.
+    assert (
+        logits.dtype == torch.float32
+        and logits.stride(1) == 1
+        and logits.stride(0) % 4 == 0
+    ), f"v2 top-k expects fp32 scores with unit row stride and 16B-aligned score_stride, got {logits.dtype=} {logits.stride()=}"
+    assert 0 < topk <= 2048, f"v2 top-k supports 0 < topk <= 2048, got {topk=}"
+
+    page_table = attn_metadata.real_page_table
+    assert page_table.dtype == torch.int32
+    lengths_i32 = lengths.to(torch.int32)
+
+    # The plan is preprocessed once per forward (DSAMetadata.topk_v2_plan,
+    # refreshed in-place under CUDA graph) and reused across layers. A missing or
+    # mismatched plan means the caller skipped that preprocessing -- fail loudly
+    # rather than silently recompute it per layer.
+    plan = attn_metadata.topk_v2_plan
+    assert (
+        plan is not None and plan.shape[0] == num_rows + 1
+    ), "topk_v2_plan must be preprocessed per forward (see DSAMetadata.topk_v2_plan)"
+
+    page_size = get_token_to_kv_pool().page_size
+    out = logits.new_full((num_rows, topk), -1, dtype=torch.int32)
+    topk_transform_512_v2(logits, lengths_i32, page_table, out, page_size, plan)
+    return out
+
+
+def _build_flashinfer_paged_args(
+    attn_metadata,
+    row_starts: Optional[torch.Tensor],
+    cu_seqlens_q_topk: Optional[torch.Tensor],
+    batch_idx_list: Optional[List[int]],
+    device: torch.device,
+    num_rows: int,
+) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
+    row_to_batch = (
+        torch.as_tensor(batch_idx_list, dtype=torch.int32, device=device)
+        if batch_idx_list is not None
+        else None
+    )
+
+    if (
+        row_to_batch is not None
+        and cu_seqlens_q_topk is not None
+        and row_to_batch.shape[0] != num_rows
+    ):
+        q_lens = (cu_seqlens_q_topk[1:] - cu_seqlens_q_topk[:-1]).to(
+            dtype=torch.int32, device=device
+        )
+        row_to_batch = torch.repeat_interleave(row_to_batch, q_lens)
+
+    if row_to_batch is None and cu_seqlens_q_topk is not None:
+        # Decode-like case (one query row per batch) does not need an explicit mapping.
+        # Avoid dynamic tensor construction in this branch to keep CUDA graph capture safe.
+        num_batches = cu_seqlens_q_topk.shape[0] - 1
+        if not (row_starts is None and num_rows == num_batches):
+            q_lens = (cu_seqlens_q_topk[1:] - cu_seqlens_q_topk[:-1]).to(
+                dtype=torch.int32, device=device
+            )
+            row_to_batch = torch.repeat_interleave(
+                torch.arange(q_lens.shape[0], dtype=torch.int32, device=device),
+                q_lens,
+            )
+
+    if row_starts is not None and row_to_batch is None:
+        raise RuntimeError(
+            "PAGED topk_transform with row_starts requires cu_seqlens_q metadata."
+        )
+
+    local_row_starts = row_starts
+    if local_row_starts is not None and row_to_batch is not None:
+        local_row_starts = (
+            local_row_starts - attn_metadata.cu_seqlens_k[:-1][row_to_batch]
+        )
+
+    return row_to_batch, local_row_starts
+
+
+def _flashinfer_tie_break_value() -> int:
+    mode = envs.SGLANG_DSA_TOPK_FLASHINFER_TIE_BREAK.get()
+    if mode is None:
+        return 0
+    mode = mode.lower()
+    if mode not in _FLASHINFER_TIE_BREAK_VALUES:
+        raise RuntimeError(
+            "SGLANG_DSA_TOPK_FLASHINFER_TIE_BREAK must be one of "
+            f"{tuple(_FLASHINFER_TIE_BREAK_VALUES)} or unset, got {mode!r}."
+        )
+    return _FLASHINFER_TIE_BREAK_VALUES[mode]

python/sglang/srt/layers/attention/dsa/index_buf_accessor.py

@@ -0,0 +1,690 @@
+from typing import TYPE_CHECKING
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.layers.attention.dsa.utils import aiter_can_use_preshuffle_paged_mqa
+from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
+from sglang.srt.utils import get_bool_env_var, is_hip
+
+_is_hip = is_hip()
+_is_fp8_fnuz = is_fp8_fnuz()
+_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
+# aiter cp_gather kernel with preshuffle=True is only valid when the indexer
+# uses the page_size=64 preshuffle layout (i.e. when the matching MQA gluon path
+# is also enabled).
+_use_aiter_preshuffle = aiter_can_use_preshuffle_paged_mqa()
+
+if _use_aiter_preshuffle:
+    from aiter.ops.cache import cp_gather_indexer_k_quant_cache
+
+if TYPE_CHECKING:
+    from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool
+
+"""
+k: data, 128 item per token, fp8
+s: scale, 1 item per token, fp32
+"""
+
+
+class GetK:
+    @classmethod
+    def execute(cls, *args, **kwargs):
+        return cls.triton(*args, **kwargs)
+
+    @classmethod
+    def slow(
+        cls, pool: "DSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
+    ):
+        num_pages = (seq_len + pool.page_size - 1) // pool.page_size
+        seq_len_ = num_pages * pool.page_size
+        index_k_fp8 = torch.empty(
+            (seq_len_, pool.index_head_dim),
+            dtype=torch.uint8,
+            device=pool.device,
+        )
+        for i in range(num_pages):
+            page_index = page_indices[i]
+            index_k_fp8[i * pool.page_size : (i + 1) * pool.page_size] = buf[
+                page_index
+            ][: pool.page_size * pool.index_head_dim].view(-1, pool.index_head_dim)
+
+        return index_k_fp8[:seq_len]
+
+    @classmethod
+    def torch_fast(
+        cls, pool: "DSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
+    ):
+        """
+        :param page_indices: (num_pages,), int32
+        :return: (seq_len, index_head_dim), uint8
+        """
+
+        # can handle per 128B instead of per element
+
+        # page_indices: (num_pages,), element := a page index
+        buf_numel_per_page = buf.shape[1]
+
+        num_k_bytes_per_page = pool.page_size * pool.index_head_dim
+        num_k_bytes_per_token = pool.index_head_dim
+
+        # buf: (num_pages, page_size 64 * head_dim 128 + page_size 64 * fp32_nbytes 4), uint8
+        # flat_buf: (whatever,), uint8
+        flat_buf = buf.flatten()
+
+        # flat_indices: (num_pages, num_k_bytes_per_page), int32, element := an index into flat_buf that we want to access
+        flat_indices = (page_indices * buf_numel_per_page)[:, None] + torch.arange(
+            num_k_bytes_per_page, dtype=torch.int32, device="cuda"
+        )[None, :]
+        flat_indices = flat_indices.flatten()[: seq_len * num_k_bytes_per_token]
+
+        out = flat_buf[flat_indices]
+        return out.view(-1, 128)
+
+    @classmethod
+    def triton(
+        cls, pool: "DSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
+    ):
+        """
+        Triton implementation for gathering K data from paged buffer.
+        :param page_indices: (num_pages,), int32/int64
+        :return: (seq_len, index_head_dim), uint8
+        """
+        return _get_k_triton(
+            buf=buf,
+            page_indices=page_indices,
+            seq_len=seq_len,
+            page_size=pool.page_size,
+            index_head_dim=pool.index_head_dim,
+        )
+
+
+class GetS:
+    @classmethod
+    def execute(cls, *args, **kwargs):
+        return cls.triton(*args, **kwargs)
+
+    @classmethod
+    def slow(
+        cls, pool: "DSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
+    ):
+        num_pages = (seq_len + pool.page_size - 1) // pool.page_size
+        seq_len_ = num_pages * pool.page_size
+        assert pool.index_head_dim // pool.quant_block_size == 1
+        index_k_scale_fp8 = torch.empty(
+            (seq_len_, 4),
+            dtype=torch.uint8,
+            device=pool.device,
+        )
+        for i in range(num_pages):
+            page_index = page_indices[i]
+            index_k_scale_fp8[i * pool.page_size : (i + 1) * pool.page_size] = buf[
+                page_index
+            ][pool.page_size * pool.index_head_dim :].view(-1, 4)
+        return index_k_scale_fp8[:seq_len]
+
+    @classmethod
+    def torch_fast(
+        cls, pool: "DSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
+    ):
+        """
+        :param page_indices: (num_pages,), int32
+        :return: (seq_len, index_head_dim // quant_block_size), uint8
+        """
+        buf_numel_per_page = buf.shape[1]
+
+        num_s_bytes_per_page = buf.shape[1] - pool.page_size * pool.index_head_dim
+        num_s_bytes_per_token = pool.index_head_dim // pool.quant_block_size * 4
+        s_offset_in_page = pool.page_size * pool.index_head_dim
+
+        flat_buf = buf.flatten()
+        flat_indices = (
+            (page_indices * buf_numel_per_page)[:, None]
+            + torch.arange(num_s_bytes_per_page, dtype=torch.int32, device="cuda")[
+                None, :
+            ]
+            + s_offset_in_page
+        )
+        flat_indices = flat_indices.flatten()[: seq_len * num_s_bytes_per_token]
+
+        out = flat_buf[flat_indices]
+        return out.view(-1, 4)
+
+    @classmethod
+    def triton(
+        cls, pool: "DSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
+    ):
+        """
+        Triton implementation for gathering S (scale) data from paged buffer.
+        :param page_indices: (num_pages,), int32/int64
+        :return: (seq_len, 4), uint8
+        """
+        return _get_s_triton(
+            buf=buf,
+            page_indices=page_indices,
+            seq_len=seq_len,
+            page_size=pool.page_size,
+            index_head_dim=pool.index_head_dim,
+        )
+
+
+class GetKAndS:
+    @classmethod
+    def execute(cls, *args, **kwargs):
+        # The aiter path uses cp_gather_indexer_k_quant_cache(preshuffle=True),
+        # which only matches the layout produced when the rest of the indexer
+        # is on the page_size=64 preshuffle path. Otherwise fall back to the
+        # triton implementation (which works on the page_size=1 legacy layout).
+        if _use_aiter_preshuffle:
+            return cls.aiter(*args, **kwargs)
+        return cls.triton(*args, **kwargs)
+
+    @classmethod
+    def aiter(
+        cls,
+        pool: "DSATokenToKVPool",
+        buf: torch.Tensor,
+        page_indices: torch.Tensor,
+        seq_len_tensor: torch.Tensor,
+        seq_len_sum: int,
+        max_seq_len: int,
+    ):
+        from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype
+
+        page_size = pool.page_size
+        index_head_dim = pool.index_head_dim
+        quant_block_size = pool.quant_block_size
+        scale_elems = index_head_dim // quant_block_size
+
+        kv_cache = buf.view(-1, page_size, index_head_dim + scale_elems * 4).view(
+            fp8_dtype
+        )
+        dst_k = torch.empty(
+            (seq_len_sum, index_head_dim), dtype=torch.uint8, device=buf.device
+        )
+        dst_scale = torch.empty(
+            (seq_len_sum, scale_elems * 4), dtype=torch.uint8, device=buf.device
+        )
+
+        cu_seq_lens = torch.zeros(
+            seq_len_tensor.shape[0] + 1, dtype=torch.int32, device=buf.device
+        )
+        torch.cumsum(seq_len_tensor.to(torch.int32), dim=0, out=cu_seq_lens[1:])
+
+        cp_gather_indexer_k_quant_cache(
+            kv_cache,
+            dst_k.view(fp8_dtype),
+            dst_scale,
+            page_indices.to(torch.int32),
+            cu_seq_lens,
+            preshuffle=True,
+        )
+        return dst_k, dst_scale
+
+    @classmethod
+    def triton(
+        cls,
+        pool: "DSATokenToKVPool",
+        buf: torch.Tensor,
+        page_indices: torch.Tensor,
+        seq_len_tensor: torch.Tensor,
+        seq_len_sum: int,
+        max_seq_len: int,
+    ):
+        """
+        Triton implementation for gathering both K and S data from paged buffer in a single call.
+        :param page_indices: (num_pages,), int32/int64
+        :param seq_len_tensor: (num_pages,), int32/int64
+        :param seq_len_sum: sum of all sequence len, int32
+        :param max_seq_len: max of all sequence len, int32
+        :return: tuple of (k_fp8, k_scale) where
+                 k_fp8: (seq_len, index_head_dim), uint8
+                 k_scale: (seq_len, 4), uint8
+        """
+        return _get_k_and_s_triton(
+            buf=buf,
+            page_indices=page_indices,
+            seq_lens=seq_len_tensor,
+            seq_len_sum=seq_len_sum,
+            max_seq_len=max_seq_len,
+            page_size=pool.page_size,
+            index_head_dim=pool.index_head_dim,
+        )
+
+
+class SetKAndS:
+    @classmethod
+    def execute(cls, *args, buf, **kwargs):
+        cls.triton(*args, **kwargs, buf=buf)
+
+    @classmethod
+    def triton(cls, pool, buf, loc, index_k, index_k_scale):
+        loc = loc.to(torch.int64)
+
+        _set_k_and_s_triton(
+            buf=buf,
+            loc=loc,
+            index_k=index_k,
+            index_k_scale=index_k_scale,
+            page_size=pool.page_size,
+        )
+
+
+def _set_k_and_s_triton(
+    buf: torch.Tensor,
+    loc: torch.Tensor,
+    index_k: torch.Tensor,
+    index_k_scale: torch.Tensor,
+    page_size: int,
+):
+    """
+    :param buf: (num_pages, page_size 64 * (128B data + 4B scale)), uint8
+    :param loc: (num_tokens_to_write,), int, element := the token index to write to
+    :param index_k: (num_tokens_to_write, 128 elem), fp8
+    :param index_k_scale: (num_tokens_to_write, 1 elem), fp32
+    :return:
+    """
+    num_pages, buf_numel_per_page = buf.shape
+    (num_tokens_to_write,) = loc.shape
+    num_tokens_to_write_, index_head_dim = index_k.shape
+
+    # Handle both 1D (num_tokens,) and 2D (num_tokens, 1) shapes for index_k_scale
+    if index_k_scale.ndim == 1:
+        num_tokens_to_write__ = index_k_scale.shape[0]
+        scale_dim = 1
+    elif index_k_scale.ndim == 2:
+        num_tokens_to_write__, scale_dim = index_k_scale.shape
+    else:
+        raise ValueError(
+            f"index_k_scale must be 1D or 2D, got shape {index_k_scale.shape}"
+        )
+    assert buf_numel_per_page == page_size * (128 + 4)
+    assert num_tokens_to_write == num_tokens_to_write_ == num_tokens_to_write__
+    assert index_head_dim == 128
+    assert scale_dim == 1
+    if _is_hip:
+        if _use_aiter_preshuffle:
+            assert (
+                page_size % 16 == 0
+            ), f"HIP preshuffle requires page_size to be a multiple of 16, got {page_size}"
+    else:
+        assert page_size == 64
+
+    assert buf.dtype == torch.uint8
+    assert loc.dtype == torch.int64, f"{loc.dtype=}"  # can be int32
+    if _is_fp8_fnuz:
+        assert index_k.dtype == torch.float8_e4m3fnuz
+    else:
+        assert index_k.dtype == torch.float8_e4m3fn
+    assert index_k_scale.dtype == torch.float32
+
+    assert buf.is_contiguous()
+    assert loc.is_contiguous()
+    assert index_k.is_contiguous()
+    assert index_k_scale.is_contiguous()
+
+    if _is_fp8_fnuz:
+        buf_fp8 = buf.view(torch.float8_e4m3fnuz)
+    else:
+        buf_fp8 = buf.view(torch.float8_e4m3fn)
+    buf_fp32 = buf.view(torch.float32)
+
+    _set_k_and_s_triton_kernel[(num_tokens_to_write,)](
+        buf_fp8,
+        buf_fp32,
+        loc,
+        index_k,
+        index_k_scale,
+        index_k.stride(0),
+        PAGE_SIZE=page_size,
+        BUF_NUMEL_PER_PAGE=buf_numel_per_page,
+        NUM_K_ELEMS_PER_TOKEN=index_head_dim,
+        S_OFFSET_NBYTES_IN_PAGE=page_size * index_head_dim,
+    )
+
+
+@triton.jit
+def _set_k_and_s_triton_kernel(
+    buf_fp8_ptr,
+    buf_fp32_ptr,
+    loc_ptr,
+    index_k_ptr,
+    index_k_scale_ptr,
+    index_k_ptr_stride_0,
+    PAGE_SIZE: tl.constexpr,
+    BUF_NUMEL_PER_PAGE: tl.constexpr,
+    NUM_K_ELEMS_PER_TOKEN: tl.constexpr,
+    S_OFFSET_NBYTES_IN_PAGE: tl.constexpr,
+):
+    token_id = tl.program_id(0)
+
+    loc = tl.load(loc_ptr + token_id)
+
+    in_k_offsets = token_id * index_k_ptr_stride_0 + tl.arange(0, NUM_K_ELEMS_PER_TOKEN)
+
+    # no need for `mask`, since we read 128B for k and 4B for scale, both pow of 2
+    k = tl.load(index_k_ptr + in_k_offsets)
+    k_scale = tl.load(index_k_scale_ptr + token_id)
+
+    loc_page_index = loc // PAGE_SIZE
+    loc_token_offset_in_page = loc % PAGE_SIZE
+
+    out_k_offsets = (
+        loc_page_index * BUF_NUMEL_PER_PAGE
+        + loc_token_offset_in_page * NUM_K_ELEMS_PER_TOKEN
+        + tl.arange(0, NUM_K_ELEMS_PER_TOKEN)
+    )
+
+    # "//4" b/c it is fp32 instead of uint8
+    out_s_offset = (
+        loc_page_index * BUF_NUMEL_PER_PAGE // 4
+        + S_OFFSET_NBYTES_IN_PAGE // 4
+        + loc_token_offset_in_page
+    )
+
+    tl.store(buf_fp8_ptr + out_k_offsets, k)
+    tl.store(buf_fp32_ptr + out_s_offset, k_scale)
+
+
+def _get_k_triton(
+    buf: torch.Tensor,
+    page_indices: torch.Tensor,
+    seq_len: int,
+    page_size: int,
+    index_head_dim: int,
+):
+    """
+    Gather K (key) data from paged buffer using Triton.
+
+    :param buf: (num_pages, page_size * 128 + page_size * 4), uint8
+    :param page_indices: (num_pages,), int32/int64
+    :param seq_len: int, number of tokens to gather
+    :param page_size: int, typically 64
+    :param index_head_dim: int, typically 128
+    :return: (seq_len, index_head_dim), uint8
+    """
+    num_pages, buf_numel_per_page = buf.shape
+
+    # Allocate output
+    out = torch.empty((seq_len, index_head_dim), dtype=torch.uint8, device=buf.device)
+
+    # Launch kernel with one thread per token
+    grid = (seq_len,)
+    _get_k_triton_kernel[grid](
+        buf,
+        page_indices,
+        out,
+        seq_len,
+        page_size,
+        buf_numel_per_page,
+        index_head_dim,
+        BLOCK_SIZE=128,
+    )
+
+    return out
+
+
+@triton.jit
+def _get_k_triton_kernel(
+    buf_ptr,
+    page_indices_ptr,
+    out_ptr,
+    seq_len: tl.constexpr,
+    page_size: tl.constexpr,
+    buf_numel_per_page: tl.constexpr,
+    index_head_dim: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """
+    Each program handles one token (seq_len tokens total).
+    Loads 128 bytes from the appropriate page.
+    """
+    token_id = tl.program_id(0)
+
+    # Calculate which page and offset within page
+    page_idx = token_id // page_size
+    token_offset_in_page = token_id % page_size
+
+    # Load the page index from page_indices
+    page_index = tl.load(page_indices_ptr + page_idx)
+
+    # Calculate source offset in buf
+    # buf[page_index, token_offset_in_page * index_head_dim : ...]
+    src_base_offset = (
+        page_index * buf_numel_per_page + token_offset_in_page * index_head_dim
+    )
+
+    # Load 128 bytes (index_head_dim elements)
+    offsets = tl.arange(0, BLOCK_SIZE)
+    mask = offsets < index_head_dim
+    data = tl.load(buf_ptr + src_base_offset + offsets, mask=mask)
+
+    # Store to output
+    dst_offset = token_id * index_head_dim
+    tl.store(out_ptr + dst_offset + offsets, data, mask=mask)
+
+
+def _get_s_triton(
+    buf: torch.Tensor,
+    page_indices: torch.Tensor,
+    seq_len: int,
+    page_size: int,
+    index_head_dim: int,
+):
+    """
+    Gather S (scale) data from paged buffer using Triton.
+
+    :param buf: (num_pages, page_size * 128 + page_size * 4), uint8
+    :param page_indices: (num_pages,), int32/int64
+    :param seq_len: int, number of tokens to gather
+    :param page_size: int, typically 64
+    :param index_head_dim: int, typically 128
+    :return: (seq_len, 4), uint8 (representing fp32 scale)
+    """
+    num_pages, buf_numel_per_page = buf.shape
+    s_offset_in_page = page_size * index_head_dim  # Scales start after K data
+
+    # Allocate output
+    out = torch.empty((seq_len, 4), dtype=torch.uint8, device=buf.device)
+
+    # Launch kernel with one thread per token
+    grid = (seq_len,)
+    _get_s_triton_kernel[grid](
+        buf,
+        page_indices,
+        out,
+        seq_len,
+        page_size,
+        buf_numel_per_page,
+        s_offset_in_page,
+    )
+
+    return out
+
+
+@triton.jit
+def _get_s_triton_kernel(
+    buf_ptr,
+    page_indices_ptr,
+    out_ptr,
+    seq_len: tl.constexpr,
+    page_size: tl.constexpr,
+    buf_numel_per_page: tl.constexpr,
+    s_offset_in_page: tl.constexpr,
+):
+    """
+    Each program handles one token (seq_len tokens total).
+    Loads 4 bytes (fp32 scale) from the appropriate page.
+    """
+    token_id = tl.program_id(0)
+
+    # Calculate which page and offset within page
+    page_idx = token_id // page_size
+    token_offset_in_page = token_id % page_size
+
+    # Load the page index from page_indices
+    page_index = tl.load(page_indices_ptr + page_idx)
+
+    # Calculate source offset in buf
+    # Scales are stored after K data: page_size * index_head_dim offset
+    # buf[page_index, s_offset_in_page + token_offset_in_page * 4 : ...]
+    src_base_offset = (
+        page_index * buf_numel_per_page + s_offset_in_page + token_offset_in_page * 4
+    )
+
+    # Load 4 bytes (fp32 scale)
+    offsets = tl.arange(0, 4)
+    data = tl.load(buf_ptr + src_base_offset + offsets)
+
+    # Store to output
+    dst_offset = token_id * 4
+    tl.store(out_ptr + dst_offset + offsets, data)
+
+
+def _get_k_and_s_triton(
+    buf: torch.Tensor,
+    page_indices: torch.Tensor,
+    seq_lens: torch.Tensor,
+    seq_len_sum: int,
+    max_seq_len: int,
+    page_size: int,
+    index_head_dim: int,
+):
+    """
+    Fused gather of both K (key) and S (scale) data from paged buffer using Triton.
+    This is more efficient than calling GetK and GetS separately.
+
+    :param buf: (num_pages, page_size * 128 + page_size * 4), uint8
+    :param page_indices: (num_pages,), int32/int64
+    :param seq_lens: tensor of sequence lens, int64
+    :param seq_len_sum: sum of all sequence len, int32
+    :param max_seq_len: max of sequence len, int32
+    :param page_size: int, typically 64
+    :param index_head_dim: int, typically 128
+    :return: tuple of (k_out, s_out) where
+             k_out: (seq_len, index_head_dim), uint8
+             s_out: (seq_len, 4), uint8
+    """
+    # Allocate outputs
+    k_out = torch.empty(
+        (seq_len_sum, index_head_dim), dtype=torch.uint8, device=buf.device
+    )
+    s_out = torch.empty((seq_len_sum, 4), dtype=torch.uint8, device=buf.device)
+
+    _, buf_numel_per_page = buf.shape
+    _, page_indice_batch_offset = page_indices.shape
+    s_offset_in_page = page_size * index_head_dim
+
+    # Launch kernel with one thread per token
+    BLOCK_SIZE = 256
+    BLOCK_SIZE_K = 128
+
+    num_token_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
+    num_k_threads = (index_head_dim + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K
+
+    seq_num = seq_lens.shape[0]
+    grid = (seq_num, num_token_blocks, num_k_threads)
+    seq_num_pow2 = 1
+    while seq_num_pow2 < seq_num:
+        seq_num_pow2 *= 2
+
+    _get_k_and_s_triton_kernel[grid](
+        buf_ptr=buf,
+        page_indices_ptr=page_indices,
+        k_out_ptr=k_out,
+        s_out_ptr=s_out,
+        seq_len_ptr=seq_lens,
+        seq_len_num_pow=seq_num_pow2,
+        page_size=page_size,
+        buf_numel_per_page=buf_numel_per_page,
+        index_head_dim=index_head_dim,
+        s_offset_in_page=s_offset_in_page,
+        page_indice_batch_offset=page_indice_batch_offset,
+        BLOCK_SIZE=BLOCK_SIZE,
+        BLOCK_SIZE_K=BLOCK_SIZE_K,
+    )
+
+    return k_out, s_out
+
+
+@triton.jit
+def _get_k_and_s_triton_kernel(
+    buf_ptr,
+    page_indices_ptr,
+    k_out_ptr,
+    s_out_ptr,
+    seq_len_ptr,
+    seq_len_num_pow: tl.constexpr,
+    page_size: tl.constexpr,
+    buf_numel_per_page: tl.constexpr,
+    index_head_dim: tl.constexpr,
+    s_offset_in_page: tl.constexpr,
+    page_indice_batch_offset,
+    BLOCK_SIZE: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+):
+    """
+    Fused kernel that gathers both K and S data in a single pass.
+    Each program handles one token (seq_len tokens total).
+    Loads 128 bytes (K) + 4 bytes (S) from the appropriate page.
+    """
+    batch_id = tl.program_id(0)
+    block_token_start = tl.program_id(1) * BLOCK_SIZE
+    thread_idx = tl.program_id(2)
+
+    # Define the token range within the block and the K dimension range handled by the thread.
+    token_ids_in_block = tl.arange(0, BLOCK_SIZE)
+    token_ids = block_token_start + token_ids_in_block
+    k_offsets = thread_idx * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
+
+    seq_len = tl.load(seq_len_ptr + batch_id)
+    token_valid_mask = token_ids < seq_len
+
+    pre_batch_idx = tl.arange(0, seq_len_num_pow)
+    mask_pre_batch_idx = pre_batch_idx < batch_id
+    prev_seq_lens = tl.load(seq_len_ptr + pre_batch_idx, mask=mask_pre_batch_idx)
+    batch_token_offset = tl.sum(prev_seq_lens)
+
+    # Batch calculate the page index and in-page offset of each token.
+    page_idx = token_ids // page_size
+    token_offset_in_page = token_ids % page_size
+    page_indices_base = batch_id * page_indice_batch_offset
+    page_idx_valid_mask = page_idx < page_indice_batch_offset
+    page_index = tl.load(
+        page_indices_ptr + page_idx + page_indices_base,
+        mask=token_valid_mask & page_idx_valid_mask,
+    )
+
+    # ===== Load K data =====
+    # The address calculation logic for K: page_index * total number of elements in a single page + K offset of the token within the page.
+    k_src_token_offset = token_offset_in_page * index_head_dim
+    k_src_base_offset = page_index * buf_numel_per_page + k_src_token_offset
+
+    k_load_addr = buf_ptr + k_src_base_offset[:, None] + k_offsets[None, :]
+    k_dim_mask = k_offsets[None, :] < index_head_dim
+    k_mask = token_valid_mask[:, None] & k_dim_mask
+
+    k_data = tl.load(k_load_addr, mask=k_mask, other=0)
+
+    # Store K to output
+    k_dst_token_offset = batch_token_offset + token_ids
+    k_dst_base_offset = k_dst_token_offset * index_head_dim
+    k_store_addr = k_out_ptr + k_dst_base_offset[:, None] + k_offsets[None, :]
+    tl.store(k_store_addr, k_data, mask=k_mask)
+
+    # ===== Load S data =====
+    # The address calculation logic for S: page_index * total number of elements in a single page + starting offset of S within the page + offset of token within S in the page
+    s_src_token_offset = s_offset_in_page + token_offset_in_page * 4
+    s_src_base_offset = page_index * buf_numel_per_page + s_src_token_offset
+
+    s_offsets = tl.arange(0, 4)
+    s_load_addr = buf_ptr + s_src_base_offset[:, None] + s_offsets[None, :]
+    s_mask = token_valid_mask[:, None] & (s_offsets[None, :] < 4)
+    s_data = tl.load(s_load_addr, mask=s_mask, other=0)
+
+    # Store S to output
+    s_dst_token_offset = batch_token_offset + token_ids
+    s_dst_base_offset = s_dst_token_offset * 4
+    s_store_addr = s_out_ptr + s_dst_base_offset[:, None] + s_offsets[None, :]
+    tl.store(s_store_addr, s_data, mask=s_mask)

python/sglang/srt/layers/attention/dsa/paged_mqa_logits_backend.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from enum import Enum
+
+from sglang.srt.utils import is_hip, is_sm100_supported
+
+
+class DSAPagedMQALogitsBackend(Enum):
+    DEEPGEMM = "deepgemm"
+    CUTEDSL = "cutedsl"
+    AITER = "aiter"
+
+    def is_deepgemm(self) -> bool:
+        return self == DSAPagedMQALogitsBackend.DEEPGEMM
+
+    def is_cutedsl(self) -> bool:
+        return self == DSAPagedMQALogitsBackend.CUTEDSL
+
+    def is_aiter(self) -> bool:
+        return self == DSAPagedMQALogitsBackend.AITER
+
+    @staticmethod
+    def resolve(value: str) -> DSAPagedMQALogitsBackend:
+        if is_hip():
+            if value not in ("auto", "aiter"):
+                raise ValueError(
+                    f"dsa_paged_mqa_logits_backend={value!r} is not supported on "
+                    "ROCm; only 'aiter' is implemented."
+                )
+            return DSAPagedMQALogitsBackend.AITER
+
+        if value == "auto" or value == "deepgemm":
+            return DSAPagedMQALogitsBackend.DEEPGEMM
+        if value == "aiter":
+            raise ValueError("dsa_paged_mqa_logits_backend='aiter' requires ROCm.")
+        if value == "cutedsl":
+            if not is_sm100_supported():
+                raise ValueError(
+                    "dsa_paged_mqa_logits_backend='cutedsl' requires SM100 (Blackwell)."
+                )
+            return DSAPagedMQALogitsBackend.CUTEDSL
+        raise ValueError(f"Unknown dsa_paged_mqa_logits_backend: {value!r}")

python/sglang/srt/layers/attention/dsa/quant_k_cache.py

@@ -0,0 +1,449 @@
+import torch
+import triton
+import triton.language as tl
+
+
+def quantize_k_cache(cache_k):
+    return _quantize_k_cache_fast_wrapped(cache_k)
+
+
+def quantize_k_cache_separate(
+    k_nope: torch.Tensor,
+    k_rope: torch.Tensor,
+    tile_size: int = 128,
+):
+    """
+    Quantize k_nope and k_rope separately without concat, returns two tensors.
+
+    This avoids the concat operation and enables direct reuse of set_mla_kv_buffer_triton
+    by returning two separate byte tensors for the nope and rope parts.
+
+    Args:
+        k_nope: (num_tokens, dim_nope) or (num_tokens, 1, dim_nope)
+                Must have dim_nope=512 for FP8 MLA quantization
+        k_rope: (num_tokens, dim_rope) or (num_tokens, 1, dim_rope)
+                Must have dim_rope=64 for FP8 MLA quantization
+        tile_size: quantization tile size (default 128)
+
+    Returns:
+        Tuple of (nope_part, rope_part) where:
+        - nope_part: (num_tokens, 1, 528) as uint8 view, contains [nope_fp8(512) | scales(16)]
+        - rope_part: (num_tokens, 1, 128) as uint8 view, contains [rope_bf16_bytes(128)]
+
+        These two tensors can be directly passed to set_mla_kv_buffer_triton(kv_buffer, loc, nope_part, rope_part)
+    """
+    # Squeeze middle dimension if present
+    k_nope_2d = k_nope.squeeze(1) if k_nope.ndim == 3 else k_nope
+    k_rope_2d = k_rope.squeeze(1) if k_rope.ndim == 3 else k_rope
+
+    num_tokens = k_nope_2d.shape[0]
+    dim_nope = k_nope_2d.shape[1]
+    dim_rope = k_rope_2d.shape[1]
+
+    # Validate dimensions for FP8 MLA
+    if dim_nope != 512:
+        raise ValueError(f"Expected dim_nope=512 for FP8 MLA, got {dim_nope}")
+    if dim_rope != 64:
+        raise ValueError(f"Expected dim_rope=64 for FP8 MLA, got {dim_rope}")
+    if k_rope_2d.shape[0] != num_tokens:
+        raise ValueError(
+            f"k_nope and k_rope must have same num_tokens, got {num_tokens} vs {k_rope_2d.shape[0]}"
+        )
+
+    return _quantize_k_cache_fast_separate(
+        k_nope=k_nope_2d, k_rope=k_rope_2d, group_size=tile_size
+    )
+
+
+# Copied from original
+def _quantize_k_cache_ref(
+    input_k_cache: torch.Tensor,  # (num_blocks, block_size, h_k, d)
+    dv: int = 512,
+    tile_size: int = 128,
+) -> torch.Tensor:
+    """
+    Quantize the k-cache
+    Return a tensor with shape (num_blocks, block_size, h_k, dv + 4(dv/tile_size) + t(d-dv)) of dtype uint8_t, where t = input_k_cache.element_size()
+    For more detail about the layout of K/V, please refer to comments in flash_mla_interface.py or README.md
+    """
+    assert dv % tile_size == 0
+    num_tiles = dv // tile_size
+    num_blocks, block_size, h_k, d = input_k_cache.shape
+    assert h_k == 1
+    input_k_cache = input_k_cache.squeeze(2)  # [num_blocks, block_size, d]
+    input_elem_size = input_k_cache.element_size()
+
+    result = torch.empty(
+        (num_blocks, block_size, dv + num_tiles * 4 + input_elem_size * (d - dv)),
+        dtype=torch.float8_e4m3fn,
+        device=input_k_cache.device,
+    )
+    result_k_nope_part = result[..., :dv]
+    result_k_scale_factor = result[..., dv : dv + num_tiles * 4].view(torch.float32)
+    result_k_rope_part = result[..., dv + num_tiles * 4 :].view(input_k_cache.dtype)
+    result_k_rope_part[:] = input_k_cache[..., dv:]
+
+    for tile_idx in range(0, num_tiles):
+        cur_scale_factors_inv = (
+            torch.abs(
+                input_k_cache[..., tile_idx * tile_size : (tile_idx + 1) * tile_size]
+            )
+            .max(dim=-1)
+            .values
+            / 448.0
+        )  # [num_blocks, block_size]
+        result_k_scale_factor[:, :, tile_idx] = cur_scale_factors_inv
+
+        cur_scale_factors_inv.unsqueeze_(-1)  # [num_blocks, block_size, 1]
+        cur_quantized_nope = (
+            input_k_cache[
+                ..., tile_idx * tile_size : (tile_idx + 1) * tile_size
+            ].float()
+            / cur_scale_factors_inv.float()
+        ).to(torch.float8_e4m3fn)
+        result_k_nope_part[..., tile_idx * tile_size : (tile_idx + 1) * tile_size] = (
+            cur_quantized_nope
+        )
+
+    result = result.view(num_blocks, block_size, 1, -1)
+    return result
+
+
+def _quantize_k_cache_fast_wrapped(
+    input_k_cache: torch.Tensor,
+    dv: int = 512,
+    tile_size: int = 128,
+) -> torch.Tensor:
+    # TODO the final API may be 2D instead of 4D, thus we convert them here
+    num_blocks, block_size, _, dim_nope_and_rope = input_k_cache.shape
+    assert dv == 512
+    assert dim_nope_and_rope == 512 + 64
+    assert tile_size == 128
+    input_k_cache = input_k_cache.view((-1, dim_nope_and_rope))
+
+    # TODO deliberately split into two tensors, then upstream can provide the two tensors instead of concat into one
+    k_nope = input_k_cache[:, :dv]
+    k_rope = input_k_cache[:, dv:]
+
+    output = _quantize_k_cache_fast(k_nope=k_nope, k_rope=k_rope)
+
+    return output.view(num_blocks, block_size, 1, -1)
+
+
+def _quantize_k_cache_fast(k_nope, k_rope, group_size: int = 128):
+    """
+    :param k_nope: (num_tokens, dim_nope 512)
+    :param k_rope: (num_tokens, dim_rope 64)
+    """
+
+    assert k_nope.dtype == torch.bfloat16
+    assert k_rope.dtype == torch.bfloat16
+
+    num_tokens, dim_nope = k_nope.shape
+    num_tokens_, dim_rope = k_rope.shape
+    assert num_tokens == num_tokens_
+    assert dim_nope == 512
+    assert dim_rope == 64
+    assert k_nope.dtype == k_rope.dtype
+    num_tiles = dim_nope // group_size
+
+    assert k_nope.stride(1) == 1
+    assert k_rope.stride(1) == 1
+
+    output = torch.empty(
+        (num_tokens, dim_nope + num_tiles * 4 + k_rope.element_size() * dim_rope),
+        dtype=torch.float8_e4m3fn,
+        device=k_nope.device,
+    )
+    output_nope_q = output[..., :dim_nope]
+    output_nope_s = output[..., dim_nope : dim_nope + num_tiles * 4].view(torch.float32)
+    output_rope = output[..., dim_nope + num_tiles * 4 :].view(torch.bfloat16)
+
+    num_blocks_per_token = triton.cdiv(dim_nope + dim_rope, group_size)
+    assert num_blocks_per_token == 5
+
+    assert dim_nope % group_size == 0
+    NUM_NOPE_BLOCKS = dim_nope // group_size
+
+    _quantize_k_cache_fast_kernel[(num_tokens, num_blocks_per_token)](
+        output_nope_q,
+        output_nope_s,
+        output_rope,
+        k_nope,
+        k_rope,
+        output_nope_q.stride(0),
+        output_nope_s.stride(0),
+        output_rope.stride(0),
+        k_nope.stride(0),
+        k_rope.stride(0),
+        NUM_NOPE_BLOCKS=NUM_NOPE_BLOCKS,
+        GROUP_SIZE=group_size,
+        DIM_NOPE=dim_nope,
+        DIM_ROPE=dim_rope,
+        FP8_MIN=torch.finfo(torch.float8_e4m3fn).min,
+        FP8_MAX=torch.finfo(torch.float8_e4m3fn).max,
+    )
+
+    return output
+
+
+def _quantize_k_cache_fast_separate(k_nope, k_rope, group_size: int = 128):
+    """
+    Quantize k_nope and k_rope in a single Triton kernel, directly outputting two separate tensors.
+
+    This avoids packing/unpacking and enables direct use with set_mla_kv_buffer_triton.
+
+    :param k_nope: (num_tokens, dim_nope 512) bfloat16
+    :param k_rope: (num_tokens, dim_rope 64) bfloat16
+    :param group_size: quantization tile size (default 128, kernel is tuned for this value)
+    :return: Tuple of (nope_part_u8, rope_part_u8)
+        - nope_part_u8: (num_tokens, 1, nope_part_bytes) uint8, layout [nope_fp8(dim_nope) | scales(num_tiles*4)]
+        - rope_part_u8: (num_tokens, 1, rope_part_bytes) uint8, layout [rope_bf16_bytes(dim_rope*2)]
+    """
+    num_tokens, dim_nope = k_nope.shape
+    num_tokens_, dim_rope = k_rope.shape
+
+    assert num_tokens == num_tokens_, f"k_nope and k_rope must have same num_tokens"
+
+    # Ensure contiguous tensors for kernel
+    k_nope = k_nope.contiguous()
+    k_rope = k_rope.contiguous()
+
+    num_tiles = dim_nope // group_size
+
+    # Calculate byte sizes based on validated dimensions
+    # nope_part: [FP8 quantized data (dim_nope bytes)] + [FP32 scales (num_tiles * 4 bytes)]
+    # rope_part: [BF16 raw data (dim_rope * 2 bytes)]
+    nope_part_bytes = (
+        dim_nope + num_tiles * 4
+    )  # e.g., 512 + 4*4 = 528 for dim_nope=512, group_size=128
+    rope_part_bytes = (
+        dim_rope * k_rope.element_size()
+    )  # e.g., 64 * 2 = 128 for dim_rope=64, BF16
+
+    # Allocate two separate output buffers (as uint8 for direct byte-level access)
+    nope_part_u8 = torch.empty(
+        (num_tokens, nope_part_bytes), dtype=torch.uint8, device=k_nope.device
+    )
+    rope_part_u8 = torch.empty(
+        (num_tokens, rope_part_bytes), dtype=torch.uint8, device=k_rope.device
+    )
+
+    # Create typed views for the kernel to write into
+    # Fixed byte layout for nope_part: [nope_fp8 (dim_nope bytes) | scales_fp32 (num_tiles*4 bytes)]
+    # Fixed byte layout for rope_part: [rope_bf16 (dim_rope*2 bytes)]
+    nope_q_view = nope_part_u8[:, :dim_nope].view(torch.float8_e4m3fn)
+    nope_s_view = nope_part_u8[:, dim_nope:].view(torch.float32)
+    rope_view = rope_part_u8.view(torch.bfloat16)
+
+    # Kernel launch parameters
+    num_blocks_per_token = triton.cdiv(dim_nope + dim_rope, group_size)
+    NUM_NOPE_BLOCKS = dim_nope // group_size
+
+    # Use the same kernel as _quantize_k_cache_fast (reuse existing implementation)
+    _quantize_k_cache_fast_kernel[(num_tokens, num_blocks_per_token)](
+        nope_q_view,
+        nope_s_view,
+        rope_view,
+        k_nope,
+        k_rope,
+        nope_q_view.stride(0),
+        nope_s_view.stride(0),
+        rope_view.stride(0),
+        k_nope.stride(0),
+        k_rope.stride(0),
+        NUM_NOPE_BLOCKS=NUM_NOPE_BLOCKS,
+        GROUP_SIZE=group_size,
+        DIM_NOPE=dim_nope,
+        DIM_ROPE=dim_rope,
+        FP8_MIN=torch.finfo(torch.float8_e4m3fn).min,
+        FP8_MAX=torch.finfo(torch.float8_e4m3fn).max,
+    )
+
+    # Add middle dimension for compatibility with set_mla_kv_buffer_triton
+    return nope_part_u8.unsqueeze(1), rope_part_u8.unsqueeze(1)
+
+
+@triton.jit
+def _quantize_k_cache_fast_kernel(
+    output_nope_q_ptr,
+    output_nope_s_ptr,
+    output_rope_ptr,
+    k_nope_ptr,
+    k_rope_ptr,
+    output_nope_q_stride_0: int,
+    output_nope_s_stride_0: int,
+    output_rope_stride_0: int,
+    k_nope_stride_0: int,
+    k_rope_stride_0: int,
+    NUM_NOPE_BLOCKS: tl.constexpr,
+    GROUP_SIZE: tl.constexpr,
+    DIM_NOPE: tl.constexpr,
+    DIM_ROPE: tl.constexpr,
+    FP8_MIN: tl.constexpr,
+    FP8_MAX: tl.constexpr,
+):
+    token_id = tl.program_id(0)
+    raw_block_id = tl.program_id(1)
+
+    if raw_block_id < NUM_NOPE_BLOCKS:
+        # a. quant nope
+        effective_block_id = raw_block_id
+
+        offs = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
+        mask = offs < DIM_NOPE
+        ptr = k_nope_ptr + token_id * k_nope_stride_0 + offs
+
+        y = tl.load(ptr, mask=mask, other=0.0).to(tl.float32)
+
+        # the ref impl do not have a `tl.maximum(... eps)`, so we remove it here
+        y_s = tl.max(tl.abs(y)) / FP8_MAX
+        y_s_inv = 1.0 / y_s
+        y_q = tl.clamp(y * y_s_inv, FP8_MIN, FP8_MAX).to(
+            output_nope_q_ptr.dtype.element_ty
+        )
+
+        dst_q_ptr = output_nope_q_ptr + token_id * output_nope_q_stride_0 + offs
+        dst_s_ptr = (
+            output_nope_s_ptr + token_id * output_nope_s_stride_0 + effective_block_id
+        )
+
+        tl.store(dst_q_ptr, y_q, mask=mask)
+        tl.store(dst_s_ptr, y_s)
+    else:
+        # b. copy rope
+        effective_block_id = raw_block_id - NUM_NOPE_BLOCKS
+
+        offs = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
+        mask = offs < DIM_ROPE
+
+        src_ptr = k_rope_ptr + token_id * k_rope_stride_0 + offs
+        dst_ptr = output_rope_ptr + token_id * output_rope_stride_0 + offs
+
+        data = tl.load(src_ptr, mask=mask)
+        tl.store(dst_ptr, data, mask=mask)
+
+
+if __name__ == "__main__":
+    import dequant_k_cache
+
+    for num_blocks, block_size in [
+        (1, 1),
+        (10, 64),
+    ]:
+        dim_nope_and_rope = 512 + 64
+
+        input_k_cache = torch.randn(
+            (num_blocks, block_size, 1, dim_nope_and_rope),
+            dtype=torch.bfloat16,
+            device="cuda",
+        )
+
+        ref_quant = _quantize_k_cache_ref(input_k_cache)
+        actual_quant = _quantize_k_cache_fast_wrapped(input_k_cache)
+
+        ref_ref_dequant = dequant_k_cache._dequantize_k_cache_slow(ref_quant)
+        ref_actual_dequant = dequant_k_cache._dequantize_k_cache_fast_wrapped(ref_quant)
+        actual_actual_dequant = dequant_k_cache._dequantize_k_cache_fast_wrapped(
+            actual_quant
+        )
+
+        print(f"{ref_ref_dequant=}")
+        print(f"{actual_actual_dequant=}")
+        print(f"{actual_actual_dequant - ref_ref_dequant=}")
+        print(f"{torch.mean(ref_ref_dequant - actual_actual_dequant)=}")
+
+        # TODO too different?
+        torch.testing.assert_close(
+            ref_ref_dequant, ref_actual_dequant, atol=0.2, rtol=0.2
+        )
+        torch.testing.assert_close(
+            ref_ref_dequant, actual_actual_dequant, atol=0.2, rtol=0.2
+        )
+
+        # test dequant_k_cache_paged
+        page_table_1 = torch.arange(
+            num_blocks * block_size, dtype=torch.int32, device="cuda"
+        )
+        actual_dequant_paged = dequant_k_cache.dequantize_k_cache_paged(
+            actual_quant, page_table_1
+        ).reshape(actual_actual_dequant.shape)
+        print(f"{torch.mean(actual_actual_dequant - actual_dequant_paged)=}")
+        torch.testing.assert_close(
+            ref_ref_dequant, actual_dequant_paged, atol=0.2, rtol=0.2
+        )
+
+    print("Passed")
+
+    # Test quantize_k_cache_separate: verify output matches concat path
+    print("\nTesting quantize_k_cache_separate...")
+    for num_tokens in [64, 100]:
+        dim_nope = 512
+        dim_rope = 64
+
+        k_nope = torch.randn(
+            num_tokens, 1, dim_nope, dtype=torch.bfloat16, device="cuda"
+        )
+        k_rope = torch.randn(
+            num_tokens, 1, dim_rope, dtype=torch.bfloat16, device="cuda"
+        )
+
+        # Old path: concat then quantize
+        k_concat = torch.cat([k_nope, k_rope], dim=-1).squeeze(1)  # (num_tokens, 576)
+        old_output = quantize_k_cache(k_concat.unsqueeze(1).unsqueeze(1))  # 4D input
+        old_output = old_output.squeeze(1).squeeze(1)  # Back to (num_tokens, 656)
+
+        # New path: quantize separately
+        nope_part, rope_part = quantize_k_cache_separate(k_nope, k_rope)
+        new_bytes = torch.cat([nope_part.squeeze(1), rope_part.squeeze(1)], dim=-1)
+
+        # Compare byte-level equality
+        old_bytes = old_output.view(torch.uint8)
+
+        if old_bytes.shape != new_bytes.shape:
+            raise RuntimeError(
+                f"Shape mismatch: {old_bytes.shape} vs {new_bytes.shape}"
+            )
+
+        diff_bytes = (old_bytes != new_bytes).sum().item()
+        if diff_bytes > 0:
+            max_diff = (old_bytes.float() - new_bytes.float()).abs().max().item()
+            raise RuntimeError(
+                f"quantize_k_cache_separate output doesn't match concat path: "
+                f"{diff_bytes} differing bytes, max_diff={max_diff}"
+            )
+
+        print(f"  num_tokens={num_tokens}: PASSED (outputs match byte-wise)")
+
+    print("quantize_k_cache_separate tests passed!")
+
+    print("\nDo benchmark...")
+
+    for num_blocks, block_size in [
+        (1, 64),
+        (64, 64),
+        (128, 64),
+        (256, 64),
+        (512, 64),
+        (1024, 64),
+        (2048, 64),
+    ]:
+        dim_nope_and_rope = 512 + 64
+
+        input_k_cache = torch.randn(
+            (num_blocks, block_size, 1, dim_nope_and_rope),
+            dtype=torch.bfloat16,
+            device="cuda",
+        )
+
+        actual_quant = _quantize_k_cache_fast_wrapped(input_k_cache)
+
+        page_table_1 = torch.arange(
+            num_blocks * block_size, dtype=torch.int32, device="cuda"
+        )
+
+        def run_ans():
+            return dequant_k_cache.dequantize_k_cache_paged(actual_quant, page_table_1)
+
+        ans_time: float = triton.testing.do_bench(run_ans, warmup=10, rep=20) / 1000  # type: ignore
+        print(f"seq_kv: {num_blocks * block_size}, time: {ans_time * 1e6: 4.0f} us")

python/sglang/srt/layers/attention/dsa/transform_index.py

@@ -0,0 +1,144 @@
+from typing import List, Optional
+
+import torch
+import triton
+import triton.language as tl
+
+
+def transform_index_page_table_prefill(**kwargs):
+    return transform_index_page_table_prefill_ref(**kwargs)
+
+
+def transform_index_page_table_decode(**kwargs):
+    return transform_index_page_table_decode_ref(**kwargs)
+
+
+@triton.jit
+def transform_index_page_table_decode_kernel(
+    page_table_ptr: torch.Tensor,
+    topk_indices_ptr: torch.Tensor,
+    result_ptr: torch.Tensor,
+    page_size: tl.constexpr,
+    max_seqlen_k: tl.constexpr,
+):
+    TOPK: tl.constexpr = 2048
+    req_id = tl.program_id(0)
+    page_table_ptr = page_table_ptr + req_id * max_seqlen_k
+    topk_indices_ptr = topk_indices_ptr + req_id * TOPK
+    result_ptr = result_ptr + req_id * TOPK
+
+    offset = tl.arange(0, TOPK)  # topk should be 2048
+    loaded_topk_indices = tl.load(topk_indices_ptr + offset)
+    mask = loaded_topk_indices >= 0
+    loaded_kv_indices = tl.load(page_table_ptr + loaded_topk_indices, mask=mask)
+    tl.store(result_ptr + offset, loaded_kv_indices, mask=mask)
+    tl.store(result_ptr + offset, -1, mask=~mask)
+
+
+def transform_index_page_table_decode_fast(
+    page_table: torch.Tensor,
+    topk_indices: torch.Tensor,
+    result: Optional[torch.Tensor] = None,
+    page_size: int = 1,
+) -> torch.Tensor:
+    """
+    Transform the page table according to topk indices for sparse topk attention.
+    Args:
+        page_table: [qo_len, max_seqlen_k], the original page table
+        topk_indices: [qo_len, topk], the topk indices for each query position
+    Returns:
+        transformed_page_table: [qo_len, topk], the transformed page table
+        For out-of-bound indices in topk_indices, this should be filled with -1.
+    """
+    assert page_size == 1
+    assert page_table.shape[0] == topk_indices.shape[0]
+    assert topk_indices.shape[1] == 2048
+    qo_len = topk_indices.shape[0]
+    max_seqlen_k = page_table.shape[1]
+    if result is None:
+        result = torch.empty_like(topk_indices, dtype=torch.int32)
+    # Launch triton kernel
+    grid = (qo_len,)
+    transform_index_page_table_decode_kernel[grid](
+        page_table,
+        topk_indices,
+        result,
+        page_size,
+        max_seqlen_k=max_seqlen_k,
+    )
+    return result
+
+
+def transform_index_page_table_prefill_fast(
+    page_table: torch.Tensor,
+    topk_indices: torch.Tensor,
+    extend_lens_cpu: List[int],
+    page_size: int = 1,
+) -> torch.Tensor:
+    # TODO(baizhou): can be implemented with another triton kernel
+    assert page_size == 1
+    result = torch.empty_like(topk_indices, dtype=torch.int32)
+    assert len(extend_lens_cpu) == page_table.shape[0]
+    offset = 0
+    for i, l in enumerate(extend_lens_cpu):
+        transform_index_page_table_decode_fast(
+            page_table[i].unsqueeze(0).expand(l, -1),
+            topk_indices[offset : offset + l],
+            result=result[offset : offset + l],
+        )
+        offset += l
+    assert offset == topk_indices.shape[0]
+    return result
+
+
+def transform_index_page_table_decode_ref(
+    page_table: torch.Tensor,
+    topk_indices: torch.Tensor,
+    result: Optional[torch.Tensor] = None,
+    page_size: int = 1,
+) -> torch.Tensor:
+    assert page_size == 1
+    assert page_table.shape[0] == topk_indices.shape[0]
+    if result is None:
+        result = torch.empty_like(topk_indices, dtype=torch.int32)
+    assert result.shape == topk_indices.shape
+    torch.gather(
+        page_table.to(result.dtype),
+        dim=1,
+        index=topk_indices.clamp(min=0),
+        out=result,
+    )
+    result[topk_indices < 0] = -1
+    return result
+
+
+def transform_index_page_table_prefill_ref(
+    page_table: torch.Tensor,
+    topk_indices: torch.Tensor,
+    extend_lens_cpu: List[int],
+    page_size: int = 1,
+) -> torch.Tensor:
+    assert page_size == 1
+    result = torch.empty_like(topk_indices, dtype=torch.int32)
+    assert len(extend_lens_cpu) == page_table.shape[0]
+    offset = 0
+    for i, l in enumerate(extend_lens_cpu):
+        transform_index_page_table_decode_ref(
+            page_table[i].unsqueeze(0).expand(l, -1),
+            topk_indices[offset : offset + l],
+            result=result[offset : offset + l],
+        )
+        offset += l
+    assert offset == topk_indices.shape[0]
+    return result
+
+
+if __name__ == "__main__":
+    bs, topk, max_seqlen = 10, 2048, 3000
+    page_table = torch.randint(0, 100, (bs, max_seqlen), device="cuda")
+    topk_indices = torch.full((bs, topk), -1, device="cuda")
+    topk_indices[:, :1600] = torch.arange(1600).unsqueeze(0).repeat(bs, 1)
+    ref_result = transform_index_page_table_decode_ref(page_table, topk_indices)
+    result = transform_index_page_table_decode_fast(page_table, topk_indices)
+    assert torch.all(result == ref_result)
+    print("Passed")

python/sglang/srt/layers/attention/dsa/triton_kernel.py

@@ -0,0 +1,196 @@
+from typing import Optional, Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+
+# Triton implementation
+@triton.jit
+def _act_quant_kernel(
+    X_ptr,
+    Y_ptr,
+    S_ptr,
+    M,
+    N,
+    group_size: tl.constexpr,
+    round_scale: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """
+    Triton kernel for activation quantization.
+
+    Each block processes BLOCK_M rows and group_size columns.
+    """
+    # Get block IDs
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+
+    # FP8 constants
+    fp8_min = -448.0
+    fp8_max = 448.0
+    fp8_max_inv = 1.0 / fp8_max
+
+    # Calculate row and column offsets
+    row_start = pid_m * BLOCK_M
+    col_start = pid_n * group_size
+
+    # Create offset arrays
+    rows = row_start + tl.arange(0, BLOCK_M)
+    cols = col_start + tl.arange(0, BLOCK_N)
+
+    # Mask for valid rows and columns
+    row_mask = rows < M
+    col_mask = cols < N
+    mask = row_mask[:, None] & col_mask[None, :]
+
+    # Load input data
+    x_ptrs = X_ptr + rows[:, None] * N + cols[None, :]
+    x = tl.load(x_ptrs, mask=mask, other=0.0).to(tl.float32)
+
+    # Compute absolute max along columns (group_size dimension) for each row
+    x_abs = tl.abs(x)
+    amax = tl.max(x_abs, axis=1)  # Shape: (BLOCK_M,)
+
+    # Clamp amax to avoid division by zero
+    amax = tl.maximum(amax, 1e-4)
+
+    # Compute scale
+    if round_scale:
+        # Fast round scale using bit manipulation approximation
+        # This is a simplified version - the exact bit manipulation is harder in Triton
+        # Using log2 + ceil + pow2 as approximation
+        log_val = tl.log2(amax * fp8_max_inv)
+        log_ceil = tl.ceil(log_val)
+        scale = tl.exp2(log_ceil)
+    else:
+        scale = amax * fp8_max_inv
+
+    # Quantize: y = clamp(x / scale, fp8_min, fp8_max)
+    scale_broadcast = scale[:, None]
+    y = x / scale_broadcast
+    y = tl.minimum(tl.maximum(y, fp8_min), fp8_max)
+
+    # Store quantized output
+    y_ptrs = Y_ptr + rows[:, None] * N + cols[None, :]
+    tl.store(y_ptrs, y, mask=mask)
+
+    # Store scales
+    s_cols = pid_n
+    s_ptrs = S_ptr + rows * (N // group_size) + s_cols
+    s_mask = row_mask
+    tl.store(s_ptrs, scale, mask=s_mask)
+
+
+def act_quant(
+    x: torch.Tensor, block_size: int = 128, scale_fmt: Optional[str] = None
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Quantizes the input tensor `x` using block-wise quantization with Triton.
+
+    Args:
+        x (torch.Tensor): The input tensor to be quantized. Must be contiguous and its last dimension size must be divisible by `block_size`.
+        block_size (int, optional): The size of the blocks to be used for quantization. Default is 128.
+        scale_fmt (Optional[str], optional): The format of the scale. Default is None.
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: A tuple containing:
+            - The quantized tensor with dtype `torch.float8_e4m3fn`.
+            - A tensor of scaling factors with dtype `torch.float32`.
+    """
+    assert x.is_contiguous(), "Input tensor must be contiguous"
+    assert (
+        x.size(-1) % block_size == 0
+    ), f"Last dimension size must be divisible by block_size (block_size={block_size})"
+
+    # Flatten all dims except last
+    N = x.size(-1)
+    x_flat = x.view(-1, N)
+    M = x_flat.size(0)
+
+    # Allocate output tensors
+    y = torch.empty_like(x, dtype=torch.float8_e4m3fn)
+    y_flat = y.view(-1, N)
+    s = x.new_empty(*x.size()[:-1], N // block_size, dtype=torch.float32)
+    s_flat = s.view(-1, N // block_size)
+
+    # Launch kernel
+    BLOCK_M = 32
+    BLOCK_N = block_size
+    grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, block_size))
+    round_scale = scale_fmt is not None
+
+    _act_quant_kernel[grid](
+        x_flat,
+        y_flat,
+        s_flat,
+        M,
+        N,
+        group_size=block_size,
+        round_scale=round_scale,
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        num_stages=0 if round_scale else 2,
+    )
+
+    return y, s
+
+
+@triton.jit
+def _get_valid_kv_indices_kernel(
+    page_table_ptr,  # [bs, topk]
+    kv_indptr_ptr,  # [bs + 1]
+    kv_indices_ptr,  # [bs * topk] output buffer
+    bs: tl.constexpr,
+    topk: tl.constexpr,
+):
+    """
+    Extract valid indices (non -1) from page_table into kv_indices.
+    Each program handles one batch.
+    """
+    batch_id = tl.program_id(0)
+
+    # Get the start position for this batch in kv_indices
+    dst_start = tl.load(kv_indptr_ptr + batch_id)
+
+    # Load all topk indices for this batch
+    src_offset = batch_id * topk
+    offsets = tl.arange(0, topk)
+    indices = tl.load(page_table_ptr + src_offset + offsets)
+
+    # Count valid indices and compact them
+    mask = indices != -1
+
+    # Use prefix sum to compute destination positions for valid elements
+    # For each position, count how many valid elements are before it
+    prefix_sum = tl.cumsum(mask.to(tl.int32), axis=0) - 1
+
+    # Store valid indices to their compacted positions
+    dst_positions = dst_start + prefix_sum
+    tl.store(kv_indices_ptr + dst_positions, indices, mask=mask)
+
+
+def get_valid_kv_indices(
+    page_table_1: torch.Tensor,
+    kv_indptr: torch.Tensor,
+    kv_indices: torch.Tensor,
+    bs: int,
+):
+    """
+    Extract valid indices from page_table_1 into kv_indices buffer.
+
+    Args:
+        page_table_1: [bs, topk] page table with -1 as invalid
+        kv_indptr: [bs + 1] cumulative count of valid indices per batch
+        kv_indices: [bs * topk] pre-allocated output buffer
+        bs: batch size
+    """
+    topk = page_table_1.shape[1]
+    grid = (bs,)
+    _get_valid_kv_indices_kernel[grid](
+        page_table_1,
+        kv_indptr,
+        kv_indices,
+        bs,
+        topk,
+    )

python/sglang/srt/layers/attention/dsa/triton_sparse_mla.py

@@ -0,0 +1,160 @@
+"""Triton sparse-MLA forward for the DSA fp8 prefill path.
+
+A per-query flash-attention kernel over the indexer-selected topk KV. On
+gfx950 this is ~1.6x faster than the TileLang partial+combine kernel for the
+prefill regime (n_groups=1): the attention tile is tiny (M=16 heads = one
+16x16 MFMA), so a small-warp per-program kernel avoids the intra-block
+coordination overhead of the 256-thread TileLang block.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
+
+_IS_FNUZ = is_fp8_fnuz()
+_FP8_MAX = 240.0 if _IS_FNUZ else 448.0
+
+
+def _prune_configs(configs, named_args, **kwargs):
+    """Drop configs whose KV tile exceeds topk (pure waste)."""
+    topk = named_args["topk"]
+    keep = [c for c in configs if c.kwargs["BLOCK_N"] <= topk]
+    return keep or [configs[0]]
+
+
+# The best (BLOCK_N, num_warps, num_stages) is shape- and arch-sensitive, so
+# autotune over a grid keyed on the attention shape. Benchmarked once per key
+# (a one-time stall on the first prefill of each new shape), then cached.
+_AUTOTUNE_CONFIGS = [
+    triton.Config({"BLOCK_N": bn}, num_warps=w, num_stages=ns)
+    for bn in (32, 64, 128)
+    for w in (1, 2, 4)
+    for ns in (1, 2)
+]
+
+
+@triton.autotune(
+    configs=_AUTOTUNE_CONFIGS,
+    key=["topk", "H", "DIM"],
+    prune_configs_by={"early_config_prune": _prune_configs},
+)
+@triton.jit
+def _sparse_mla_fwd_kernel(
+    q_nope_ptr,
+    q_rope_ptr,
+    kv_ptr,
+    idx_ptr,
+    o_ptr,
+    sm_scale,
+    fp8_max,
+    topk,
+    H: tl.constexpr,
+    DIM: tl.constexpr,
+    D_V: tl.constexpr,
+    D_TAIL: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    s_i = tl.program_id(0)
+
+    h = tl.arange(0, H)
+    dv = tl.arange(0, D_V)
+    dt = tl.arange(0, D_TAIL)
+    # q is read as two separate tensors (q_nope width D_V, q_rope width D_TAIL):
+    # the upstream concat into a single [.., DIM] tensor is skipped since this
+    # kernel splits q into main/tail anyway.
+    q_main = tl.load(q_nope_ptr + s_i * H * D_V + h[:, None] * D_V + dv[None, :]).to(
+        q_nope_ptr.dtype.element_ty
+    )  # [H, D_V]
+    q_tail = tl.load(
+        q_rope_ptr + s_i * H * D_TAIL + h[:, None] * D_TAIL + dt[None, :]
+    ).to(
+        q_nope_ptr.dtype.element_ty
+    )  # [H, D_TAIL]
+
+    m_i = tl.full([H], -float("inf"), tl.float32)
+    l_i = tl.zeros([H], tl.float32)
+    acc = tl.zeros([H, D_V], tl.float32)
+
+    n = tl.arange(0, BLOCK_N)
+    for k0 in range(0, topk, BLOCK_N):
+        kmask = (k0 + n) < topk
+        idx = tl.load(idx_ptr + s_i * topk + k0 + n, mask=kmask, other=-1)
+        valid = (idx >= 0) & kmask
+        page = tl.where(valid, idx, 0)
+        kbase = kv_ptr + page[:, None] * DIM
+        kv_main = tl.load(kbase + dv[None, :], mask=valid[:, None], other=0.0).to(
+            q_nope_ptr.dtype.element_ty
+        )  # [BLOCK_N, D_V] -- reused as V
+        kv_tail = tl.load(
+            kbase + (D_V + dt)[None, :], mask=valid[:, None], other=0.0
+        ).to(
+            q_nope_ptr.dtype.element_ty
+        )  # [BLOCK_N, D_TAIL]
+
+        qk = tl.dot(q_main, tl.trans(kv_main)).to(tl.float32)
+        qk += tl.dot(q_tail, tl.trans(kv_tail)).to(tl.float32)
+        qk = qk * sm_scale
+        qk = tl.where(valid[None, :], qk, -float("inf"))
+
+        m_new = tl.maximum(m_i, tl.max(qk, axis=1))
+        # Guard an all-masked row (m_new == -inf): shift by 0 instead so that
+        # exp(-inf - 0) = 0 rather than exp(-inf + inf) = NaN. Identical to
+        # m_new whenever the row has >=1 valid key.
+        m_safe = tl.where(m_new == -float("inf"), 0.0, m_new)
+        alpha = tl.exp(m_i - m_safe)
+        p = tl.exp(qk - m_safe[:, None])  # [H, BLOCK_N]
+        l_i = l_i * alpha + tl.sum(p, axis=1)
+
+        p_fp8 = (p * fp8_max).to(q_nope_ptr.dtype.element_ty)
+        pv = tl.dot(p_fp8, kv_main).to(tl.float32) * (1.0 / fp8_max)
+        acc = acc * alpha[:, None] + pv
+        m_i = m_new
+
+    l_safe = tl.where(l_i == 0.0, 1.0, l_i)
+    acc = acc / l_safe[:, None]
+    tl.store(
+        o_ptr + s_i * H * D_V + h[:, None] * D_V + dv[None, :],
+        acc.to(o_ptr.dtype.element_ty),
+    )
+
+
+def triton_sparse_mla_fwd(
+    q_nope: torch.Tensor,
+    q_rope: torch.Tensor,
+    kv: torch.Tensor,
+    indices: torch.Tensor,
+    sm_scale: float,
+    d_v: int = 512,
+) -> torch.Tensor:
+    """q_nope: [seq, H, d_v] fp8, q_rope: [seq, H, dim-d_v] fp8,
+    kv: [num_pages, 1, dim] fp8, indices: [seq, 1, topk].
+
+    Reads q from the two un-concatenated tensors directly (no q_nope/q_rope
+    concat). Returns [1, seq, H, d_v] bf16 to match tilelang_sparse_fwd.
+    """
+    seq, H, d_v_in = q_nope.shape
+    assert d_v_in == d_v
+    d_tail = q_rope.shape[-1]
+    dim = kv.shape[-1]
+    topk = indices.shape[-1]
+    q_nope = q_nope.contiguous()
+    q_rope = q_rope.contiguous()
+    out = torch.empty(seq, H, d_v, device=q_nope.device, dtype=torch.bfloat16)
+    # BLOCK_N / num_warps / num_stages are chosen by @triton.autotune.
+    _sparse_mla_fwd_kernel[(seq,)](
+        q_nope,
+        q_rope,
+        kv,
+        indices,
+        out,
+        sm_scale,
+        _FP8_MAX,
+        topk,
+        H=H,
+        DIM=dim,
+        D_V=d_v,
+        D_TAIL=d_tail,
+    )
+    return out.unsqueeze(0)

python/sglang/srt/layers/attention/dsa/utils.py

@@ -0,0 +1,317 @@
+from functools import lru_cache
+from typing import TYPE_CHECKING, List, Tuple, Union
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.environ import envs
+from sglang.srt.layers.dp_attention import (
+    DpPaddingMode,
+)
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
+    is_in_breakable_cuda_graph,
+)
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    is_in_tc_piecewise_cuda_graph,
+)
+from sglang.srt.runtime_context import get_parallel, get_server_args
+from sglang.srt.utils import get_bool_env_var, is_cuda, is_hip
+from sglang.srt.utils.common import ceil_align, ceil_div
+
+
+@lru_cache(maxsize=1)
+def aiter_can_use_preshuffle_paged_mqa() -> bool:
+    """Whether aiter's preshuffle paged MQA / cache kernels can be used on this runtime.
+
+    aiter's ``deepgemm_fp8_paged_mqa_logits`` only supports ``KVBlockSize > 1`` and
+    ``Preshuffle=True`` on its gluon kernel path. The gluon path is enabled when
+    Triton >= 3.5.0, OR when ``AITER_ENABLE_AOT_GLUON_PA_MQA_LOGITS=1`` is set
+    (which additionally requires that the AOT gluon kernel artifacts ship inside
+    the aiter wheel/image). Otherwise aiter asserts ``KVBlockSize == 1`` and
+    refuses ``Preshuffle=True``.
+
+    sglang's DSA indexer uses this single decision to pick:
+      * ``page_size``: 64 (preshuffle) vs 1 (legacy) on ROCm
+      * ``Preshuffle`` / ``preshuffle`` flags on the aiter MQA + cache kernels
+      * ``get_page_table_64`` vs ``get_page_table_1`` on the metadata
+      * whether ``GetKAndS.execute`` uses the aiter or the triton implementation
+
+    The result is cached so the cost is paid once per process.
+
+    Set ``SGLANG_DSA_HIP_DISABLE_PRESHUFFLE=1`` to force the legacy path even when
+    the gluon kernel would otherwise be available (useful for CI bisection).
+    ``SGLANG_NSA_HIP_DISABLE_PRESHUFFLE`` is a deprecated alias.
+    """
+    if not is_hip():
+        return False
+    if not get_bool_env_var("SGLANG_USE_AITER"):
+        return False
+    if envs.SGLANG_DSA_HIP_DISABLE_PRESHUFFLE.get():
+        return False
+    if get_bool_env_var("AITER_ENABLE_AOT_GLUON_PA_MQA_LOGITS"):
+        return True
+    try:
+        from packaging.version import Version
+
+        return Version(Version(triton.__version__).base_version) >= Version("3.5.0")
+    except Exception:
+        return False
+
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+def compute_dsa_seqlens(original_seq_lens, dsa_index_topk: int):
+    return original_seq_lens.clamp(max=dsa_index_topk)
+
+
+def is_dsa_enable_prefill_cp():
+    return get_server_args().enable_dsa_prefill_context_parallel
+
+
+def is_dsa_prefill_cp_in_seq_split():
+    return (
+        is_dsa_enable_prefill_cp()
+        and get_server_args().dsa_prefill_cp_mode == "in-seq-split"
+    )
+
+
+def is_dsa_prefill_cp_round_robin_split():
+    return (
+        is_dsa_enable_prefill_cp()
+        and get_server_args().dsa_prefill_cp_mode == "round-robin-split"
+    )
+
+
+# Structural surface where the graph DSA split-op dispatch (DSA indexer) and the
+# MLA BMM-into-attention fusion apply: a non-speculative extend (prefill) running
+# inside a piecewise/breakable CUDA graph. Both fusions are now on by default on
+# this surface (no feature flag); each adds its own extra carve-outs at its call
+# site (e.g. the indexer also excludes DSA prefill context parallelism).
+def is_graph_dsa_split_op_surface(forward_batch: "ForwardBatch") -> bool:
+    return (
+        is_cuda()
+        and (is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph())
+        and forward_batch.forward_mode.is_extend_without_speculative()
+    )
+
+
+def can_dsa_prefill_cp_round_robin_split(forward_batch: "ForwardBatch"):
+    if not forward_batch.forward_mode.is_context_parallel_extend():
+        return False
+    cp_size = get_parallel().attn_cp_size
+    seq_len = sum(forward_batch.extend_seq_lens_cpu)
+    return (
+        is_dsa_prefill_cp_round_robin_split()
+        and seq_len > 0
+        and seq_len >= cp_size
+        and cp_size > 1
+    )
+
+
+def dsa_cp_round_robin_split_data(input_: Union[torch.Tensor, List]):
+    """
+    # for round-robin-split, split the tokens evenly according to the rule of token_idx % cp_size.
+    |   +-----------before split------------+|
+    | token0, token1, token2, token3, token4, token5, token6, token7, ...
+    |
+    |   +--------------result-------------------+
+    | dp_atten_tp0: token0, token4, token8, token12, token16, ... |
+    | dp_atten_tp1: token1, token5, token9, token13, token17, ... |
+    | dp_atten_tp2: token2, token6, token10, token14, token18, ... |
+    | dp_atten_tp3: token3, token7, token11, token15, token19, ... |
+    |   +-------------------------+
+    """
+    cp_size = get_parallel().attn_cp_size
+    cp_rank = get_parallel().attn_cp_rank
+    if isinstance(input_, (tuple, list)):
+        indices = range(cp_rank, len(input_), cp_size)
+        return input_[indices]
+
+    tokens = len(input_)
+    if tokens % cp_size != 0:
+        cur_len = tokens // cp_size + (tokens % cp_size > cp_rank)
+        if cur_len == 0:
+            return input_.new_empty(0, *input_.shape[1:])
+        indices = torch.arange(cp_rank, tokens, cp_size, device=input_.device)
+        return input_[indices]
+
+    # for torch device tensor
+    return input_.view(-1, cp_size, *input_.shape[1:])[:, cp_rank].contiguous()
+
+
+def cal_padded_tokens(forward_batch: "ForwardBatch"):
+    # Consistent with the padding calculation logic in ForwardBatch.prepare_mlp_sync_batch,
+    # calculate the actual token length after padding when attn_tp_size > 1 or in the MAX_LEN padding mode.
+    from sglang.srt.layers.utils.cp_utils import get_cp_padding_align_size
+
+    global_num_tokens = forward_batch.global_num_tokens_cpu.copy()
+    sync_group_size = len(global_num_tokens)
+    attn_cp_size = get_parallel().attn_cp_size
+    # Must match the CP padding in ForwardBatch.prepare_mlp_sync_batch.
+    cp_align_size = get_cp_padding_align_size()
+    for i in range(sync_group_size):
+        global_num_tokens[i] = ceil_align(global_num_tokens[i], cp_align_size)
+    dp_padding_mode = DpPaddingMode.get_dp_padding_mode(
+        forward_batch.is_extend_in_batch, global_num_tokens
+    )
+    if dp_padding_mode.is_max_len():
+        tokens = max(global_num_tokens)
+    elif len(global_num_tokens) > 1:
+        tokens = global_num_tokens[get_parallel().attn_dp_rank]
+    else:
+        tokens = global_num_tokens[0]
+    if can_dsa_prefill_cp_round_robin_split(forward_batch):
+        tokens = ceil_div(tokens, attn_cp_size)
+    return tokens
+
+
+def pad_dsa_cache_seqlens(forward_batch: "ForwardBatch", dsa_cache_seqlens):
+    attn_cp_size = get_parallel().attn_cp_size
+    needs_cp_pad = attn_cp_size > 1 and can_dsa_prefill_cp_round_robin_split(
+        forward_batch
+    )
+    needs_dp_pad = forward_batch.global_num_tokens_cpu is not None
+    if not needs_cp_pad and not needs_dp_pad:
+        return dsa_cache_seqlens
+    tokens = cal_padded_tokens(forward_batch)
+    pad_len = tokens - dsa_cache_seqlens.shape[0]
+    if pad_len > 0:
+        dsa_cache_seqlens = torch.cat(
+            [
+                dsa_cache_seqlens,
+                dsa_cache_seqlens.new_zeros(pad_len, *dsa_cache_seqlens.shape[1:]),
+            ]
+        )
+    return dsa_cache_seqlens
+
+
+def can_dsa_cp_split(seq_len: int, cp_size: int, use_dsa: bool, forward_batch):
+    if is_dsa_prefill_cp_round_robin_split():
+        cur_cp_seq_len = seq_len // cp_size
+        assert (
+            seq_len % cp_size == 0
+        ), f"seq_len {seq_len} is not divisible by cp_size {cp_size} when dsa_prefill_cp_mode is round-robin-split"
+    else:
+        # TODO current just support prefill batch=1 and len(input_ids) > self.cp_size * 2
+        # Note: (self.cp_size * 2) To achieve load balancing for seq computation,
+        # the seq data needs to be divided and recombined at twice the size of cp_size.
+        cur_cp_seq_len = seq_len // (cp_size * 2)
+    if (
+        cur_cp_seq_len != 0
+        and cp_size > 1
+        and use_dsa
+        and forward_batch.forward_mode.is_context_parallel_extend()
+        and is_dsa_enable_prefill_cp()
+        and sum(forward_batch.extend_seq_lens_cpu) >= cp_size
+    ):
+        return True
+    else:
+        return False
+
+
+@triton.jit
+def dsa_cp_round_robin_split_q_seqs_kernel(
+    in_seqs_ptr,
+    out_seqs_ptr,
+    bs_idx_ptr,
+    tokens: tl.constexpr,
+    cp_size: tl.constexpr,
+    cp_rank: tl.constexpr,
+):
+    extra_seq = 0
+    bs_idx = 0
+    for bs in range(tokens):
+        cur_len = tl.load(in_seqs_ptr + bs)
+        cur_len += extra_seq
+        cur_seq = cur_len // cp_size + (cur_len % cp_size > cp_rank)
+        if cur_seq > 0:
+            tl.store(bs_idx_ptr + bs_idx, bs)
+            tl.store(out_seqs_ptr + bs_idx, cur_seq)
+            bs_idx += 1
+        extra_seq = cur_len - cur_seq * cp_size
+
+
+def dsa_cp_round_robin_split_q_seqs_cpu(extend_seqs):
+    cp_size = get_parallel().attn_cp_size
+    cp_rank = get_parallel().attn_cp_rank
+    extra_seq = 0
+    q_seqs = []
+    for bs, cur_len in enumerate(extend_seqs):
+        cur_len += extra_seq
+        cur_seq = cur_len // cp_size + int(cur_len % cp_size > cp_rank)
+        q_seqs.append(cur_seq)
+        extra_seq = cur_len - cur_seq * cp_size
+    bs_idx = list([i for i, x in enumerate(q_seqs) if x > 0])
+    q_seqs = [q_len for q_len in q_seqs if q_len > 0]
+    return q_seqs, bs_idx
+
+
+def dsa_cp_round_robin_split_q_seqs(
+    extend_seqs_cpu, extend_seqs
+) -> Tuple[List, torch.Tensor, List, torch.Tensor]:
+    """
+    round-robin-split distributes tokens across ranks based on token_idx % cp_size.
+
+    Return:
+    ret_q_lens_cpu(List) and ret_q_lens(torch.Tensor): the partitioned length (excluding zeros) on the current cp rank
+        for each sequence after distribution across cp ranks.
+    bs_idx_cpu(List) and bs_idx(torch.Tensor): marks which sequences are ultimately selected,
+        i.e., those with a partitioned length greater than zero.
+    """
+    cp_size = get_parallel().attn_cp_size
+    cp_rank = get_parallel().attn_cp_rank
+    # len(ret_q_lens_cpu) == len(bs_idx_cpu)
+    ret_q_lens_cpu, bs_idx_cpu = dsa_cp_round_robin_split_q_seqs_cpu(extend_seqs_cpu)
+    ret_q_lens = torch.empty(
+        (len(bs_idx_cpu),), device=extend_seqs.device, dtype=extend_seqs.dtype
+    )
+    bs_idx = torch.empty(
+        (len(bs_idx_cpu),), device=extend_seqs.device, dtype=torch.int32
+    )
+    grid = (1,)
+    dsa_cp_round_robin_split_q_seqs_kernel[grid](
+        extend_seqs, ret_q_lens, bs_idx, len(extend_seqs), cp_size, cp_rank
+    )
+    return ret_q_lens_cpu, ret_q_lens, bs_idx_cpu, bs_idx
+
+
+def dsa_use_prefill_cp(forward_batch, dsa_enable_prefill_cp=None):
+    if dsa_enable_prefill_cp is None:
+        dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
+    if (
+        forward_batch.attn_cp_metadata is not None
+        and dsa_enable_prefill_cp
+        and forward_batch.forward_mode.is_context_parallel_extend()
+    ):
+        return True
+    else:
+        return False
+
+
+def fp8_mqa_logits_ceil_to_ue8m0(x: torch.Tensor) -> torch.Tensor:
+    return torch.pow(2.0, torch.ceil(torch.log2(x.abs())))
+
+
+def fp8_mqa_logits_make_fused_kv(
+    kv_fp8: torch.Tensor,
+    kv_scales: torch.Tensor,
+    block_kv: int,
+    head_dim: int,
+) -> torch.Tensor:
+    num_phys_blocks = kv_fp8.shape[0]
+    per_token_size = head_dim + 4
+    block_bytes = block_kv * per_token_size
+    scale_offset = block_kv * head_dim
+
+    fused = torch.zeros(
+        num_phys_blocks, block_bytes, dtype=torch.uint8, device=kv_fp8.device
+    )
+    for blk in range(num_phys_blocks):
+        fused[blk, :scale_offset] = kv_fp8[blk].view(torch.uint8).reshape(-1)
+        fused[blk, scale_offset:] = (
+            kv_scales[blk].float().contiguous().view(torch.uint8).reshape(-1)
+        )
+    return fused.view(num_phys_blocks, block_kv, 1, per_token_size)