chore: import upstream snapshot with attribution

sgl-kernel/python/sgl_kernel/__init__.py

@@ -0,0 +1,243 @@
+import platform
+import sys
+
+from sgl_kernel.version import __version__  # noqa: F401
+
+# On macOS only the Metal extension is shipped; skip CUDA op loading and
+# re-exports so those symbols are not exposed on Apple Silicon.
+if sys.platform == "darwin" and platform.machine() == "arm64":
+    from sgl_kernel.metal import *
+else:
+    import torch
+    from sgl_kernel.debug_utils import maybe_wrap_debug_kernel
+    from sgl_kernel.load_utils import (
+        _load_architecture_specific_ops,
+        _preload_cuda_library,
+    )
+
+    # Initialize the ops library based on current GPU
+    common_ops = _load_architecture_specific_ops()
+
+    # Preload the CUDA library to avoid the issue of libcudart.so.12 not found
+    if torch.version.cuda is not None:
+        _preload_cuda_library()
+
+    from sgl_kernel.allreduce import *
+    from sgl_kernel.attention import (
+        cutlass_mla_decode,
+        cutlass_mla_get_workspace_size,
+        merge_state_v2,
+    )
+    from sgl_kernel.cutlass_moe import (
+        cutlass_w4a8_moe_mm,
+        get_cutlass_w4a8_moe_mm_data,
+    )
+    from sgl_kernel.elementwise import (
+        concat_mla_absorb_q,
+        concat_mla_k,
+        copy_to_gpu_no_ce,
+        dsv4_fused_k_norm_rope_flashmla,
+        dsv4_fused_q_indexer_rope_hadamard_quant,
+        dsv4_fused_q_norm_rope,
+        fused_add_rmsnorm,
+        gelu_and_mul,
+        gelu_tanh_and_mul,
+        gemma_fused_add_rmsnorm,
+        gemma_rmsnorm,
+        rmsnorm,
+        rotary_embedding,
+        silu_and_mul,
+    )
+    from sgl_kernel.expert_specialization import (
+        es_fp8_blockwise_scaled_grouped_mm,
+        es_sm100_mxfp8_blockscaled_grouped_mm,
+        es_sm100_mxfp8_blockscaled_grouped_quant,
+    )
+    from sgl_kernel.gemm import (
+        awq_dequantize,
+        bmm_fp8,
+        dsv3_fused_a_gemm,
+        fp8_blockwise_scaled_mm,
+        fp8_scaled_mm,
+        gptq_gemm,
+        gptq_shuffle,
+        int8_scaled_mm,
+        qserve_w4a8_per_chn_gemm,
+        qserve_w4a8_per_group_gemm,
+        sgl_per_token_group_quant_8bit,
+        sgl_per_token_group_quant_fp8,
+        sgl_per_token_group_quant_int8,
+        sgl_per_token_quant_fp8,
+        shuffle_rows,
+    )
+    from sgl_kernel.grammar import apply_token_bitmask_inplace_cuda
+    from sgl_kernel.infllm_v2 import (
+        infllmv2_attn_stage1,
+        max_pooling_1d_varlen,
+    )
+    from sgl_kernel.kvcacheio import (
+        copy_all_layer_kv_cache_cpu,
+        transfer_kv_all_layer,
+        transfer_kv_all_layer_mla,
+        transfer_kv_per_layer,
+        transfer_kv_per_layer_mla,
+    )
+    from sgl_kernel.mamba import (
+        causal_conv1d_fn_cpu,
+        causal_conv1d_fwd,
+        causal_conv1d_update,
+        causal_conv1d_update_cpu,
+        chunk_gated_delta_rule_cpu,
+    )
+    from sgl_kernel.memory import weak_ref_tensor
+    from sgl_kernel.moe import (
+        apply_shuffle_mul_sum,
+        fp8_blockwise_scaled_grouped_mm,
+        fused_qk_norm_rope,
+        moe_align_block_size,
+        moe_sum,
+        moe_sum_reduce,
+        prepare_moe_input,
+        topk_sigmoid,
+        topk_softmax,
+    )
+    from sgl_kernel.quantization import (
+        ggml_dequantize,
+        ggml_moe_a8,
+        ggml_moe_a8_vec,
+        ggml_moe_get_block_size,
+        ggml_mul_mat_a8,
+        ggml_mul_mat_vec_a8,
+    )
+    from sgl_kernel.sampling import (
+        top_k_renorm_prob,
+        top_p_renorm_prob,
+    )
+    from sgl_kernel.speculative import (
+        assign_draft_cache_locs_contiguous_cpu,
+        assign_extend_cache_locs_cpu,
+        assign_req_to_token_pool_cpu,
+        build_draft_decode_metadata_cpu,
+        build_tree_kernel_efficient,
+        build_tree_kernel_efficient_cpu,
+        fill_accept_out_cache_loc_cpu,
+        fill_bonus_tokens_cpu,
+        reconstruct_indices_from_tree_mask,
+        rotate_input_ids_cpu,
+        segment_packbits,
+        tree_speculative_sampling_target_only,
+        verify_tree_greedy,
+        verify_tree_greedy_cpu,
+    )
+    from sgl_kernel.top_k import (
+        fast_topk,
+        fast_topk_transform_fused,
+        fast_topk_transform_ragged_fused,
+        fast_topk_v2,
+    )
+    from sgl_kernel.version import __version__
+
+    if torch.version.hip is not None:
+        from sgl_kernel.elementwise import gelu_quick
+        from sgl_kernel.top_k import deepseek_v4_topk_transform_512
+
+    if hasattr(torch.version, "musa") and torch.version.musa is not None:
+        from sgl_kernel.musa import (
+            min_p_sampling_from_probs,
+            musa_batched_rotary_embedding_contiguous,
+            musa_fused_gemv,
+            musa_fused_moe_gemv,
+            musa_fused_mul_add,
+            musa_rotary_embedding_contiguous,
+            top_k_top_p_sampling_from_probs,
+        )
+
+    _DEBUG_EXPORT_NAMES = [
+        "apply_shuffle_mul_sum",
+        "apply_token_bitmask_inplace_cuda",
+        "awq_dequantize",
+        "bmm_fp8",
+        "build_tree_kernel_efficient",
+        "causal_conv1d_fwd",
+        "causal_conv1d_update",
+        "concat_mla_absorb_q",
+        "concat_mla_k",
+        "copy_to_gpu_no_ce",
+        "cutlass_mla_decode",
+        "cutlass_mla_get_workspace_size",
+        "dsv3_fused_a_gemm",
+        "dsv3_router_gemm",
+        "dsv4_fused_k_norm_rope_flashmla",
+        "dsv4_fused_q_indexer_rope_hadamard_quant",
+        "dsv4_fused_q_norm_rope",
+        "es_fp8_blockwise_scaled_grouped_mm",
+        "es_sm100_mxfp8_blockscaled_grouped_mm",
+        "es_sm100_mxfp8_blockscaled_grouped_quant",
+        "fast_topk",
+        "fast_topk_transform_fused",
+        "fast_topk_transform_ragged_fused",
+        "fast_topk_v2",
+        "fp8_blockwise_scaled_grouped_mm",
+        "fp8_blockwise_scaled_mm",
+        "fp8_scaled_mm",
+        "fused_add_rmsnorm",
+        "fused_qk_norm_rope",
+        "gelu_and_mul",
+        "gelu_tanh_and_mul",
+        "gemma_fused_add_rmsnorm",
+        "gemma_rmsnorm",
+        "gptq_gemm",
+        "gptq_shuffle",
+        "int8_scaled_mm",
+        "merge_state_v2",
+        "moe_align_block_size",
+        "moe_sum",
+        "moe_sum_reduce",
+        "prepare_moe_input",
+        "qserve_w4a8_per_chn_gemm",
+        "qserve_w4a8_per_group_gemm",
+        "reconstruct_indices_from_tree_mask",
+        "rmsnorm",
+        "rotary_embedding",
+        "segment_packbits",
+        "sgl_per_token_group_quant_8bit",
+        "sgl_per_token_group_quant_fp8",
+        "sgl_per_token_group_quant_int8",
+        "sgl_per_token_quant_fp8",
+        "shuffle_rows",
+        "silu_and_mul",
+        "top_k_renorm_prob",
+        "top_p_renorm_prob",
+        "topk_sigmoid",
+        "topk_softmax",
+        "transfer_kv_all_layer",
+        "transfer_kv_all_layer_mla",
+        "transfer_kv_per_layer",
+        "transfer_kv_per_layer_mla",
+        "tree_speculative_sampling_target_only",
+        "verify_tree_greedy",
+        "weak_ref_tensor",
+    ]
+
+    if torch.version.hip is not None:
+        _DEBUG_EXPORT_NAMES.append("gelu_quick")
+        _DEBUG_EXPORT_NAMES.append("deepseek_v4_topk_transform_512")
+
+    for _name in _DEBUG_EXPORT_NAMES:
+        if _name in globals():
+            globals()[_name] = maybe_wrap_debug_kernel(
+                globals()[_name], f"sgl_kernel.{_name}"
+            )
+
+    del _name
+    del _DEBUG_EXPORT_NAMES
+
+    def create_greenctx_stream_by_value(*args, **kwargs):
+        from sgl_kernel.spatial import create_greenctx_stream_by_value as _impl
+
+        return _impl(*args, **kwargs)
+
+    def get_sm_available(*args, **kwargs):
+        from sgl_kernel.spatial import get_sm_available as _impl
+
+        return _impl(*args, **kwargs)

sgl-kernel/python/sgl_kernel/allreduce.py

@@ -0,0 +1,130 @@
+from typing import List, Optional, Tuple
+
+import torch
+
+if torch.version.hip is not None:
+    # ROCM custom allreduce
+    def init_custom_ar(
+        meta: torch.Tensor,
+        rank_data: torch.Tensor,
+        handles: List[str],
+        offsets: List[int],
+        rank: int,
+        full_nvlink: bool,
+    ) -> int:
+        return torch.ops.sgl_kernel.init_custom_ar.default(
+            meta, rank_data, handles, offsets, rank, full_nvlink
+        )
+
+    def all_reduce_reg(fa: int, inp: torch.Tensor, out: torch.Tensor) -> None:
+        torch.ops.sgl_kernel.all_reduce_reg.default(fa, inp, out)
+
+    def all_reduce_unreg(
+        fa: int, inp: torch.Tensor, reg_buffer: torch.Tensor, out: torch.Tensor
+    ) -> None:
+        torch.ops.sgl_kernel.all_reduce_unreg.default(fa, inp, reg_buffer, out)
+
+    def deterministic_all_reduce_reg(
+        fa: int, inp: torch.Tensor, out: torch.Tensor
+    ) -> None:
+        torch.ops.sgl_kernel.deterministic_all_reduce_reg.default(fa, inp, out)
+
+    def deterministic_all_reduce_unreg(
+        fa: int, inp: torch.Tensor, reg_buffer: torch.Tensor, out: torch.Tensor
+    ) -> None:
+        torch.ops.sgl_kernel.deterministic_all_reduce_unreg.default(
+            fa, inp, reg_buffer, out
+        )
+
+    def dispose(fa: int) -> None:
+        torch.ops.sgl_kernel.dispose.default(fa)
+
+    def meta_size() -> int:
+        return torch.ops.sgl_kernel.meta_size.default()
+
+    def register_buffer(
+        fa: int, t: torch.Tensor, handles: List[str], offsets: List[int]
+    ) -> None:
+        return torch.ops.sgl_kernel.register_buffer.default(fa, t, handles, offsets)
+
+    def get_graph_buffer_ipc_meta(fa: int) -> Tuple[torch.Tensor, List[int]]:
+        return torch.ops.sgl_kernel.get_graph_buffer_ipc_meta.default(fa)
+
+    def register_graph_buffers(
+        fa: int, handles: List[str], offsets: List[List[int]]
+    ) -> None:
+        torch.ops.sgl_kernel.register_graph_buffers.default(fa, handles, offsets)
+
+    def allocate_meta_buffer(size: int) -> torch.Tensor:
+        return torch.ops.sgl_kernel.allocate_meta_buffer.default(size)
+
+    def get_meta_buffer_ipc_handle(inp: torch.Tensor) -> torch.Tensor:
+        return torch.ops.sgl_kernel.get_meta_buffer_ipc_handle.default(inp)
+
+    # ROCM quick allreduce
+    def init_custom_qr(
+        rank: int, world_size: int, qr_max_size: Optional[int] = None
+    ) -> int:
+        return torch.ops.sgl_kernel.init_custom_qr.default(
+            world_size, rank, qr_max_size
+        )
+
+    def qr_get_handle(fa: int) -> torch.Tensor:
+        return torch.ops.sgl_kernel.qr_get_handle.default(fa)
+
+    def qr_open_handles(fa: int, handles: list[torch.Tensor]) -> None:
+        torch.ops.sgl_kernel.qr_open_handles.default(fa, handles)
+
+    def qr_all_reduce(
+        fa: int,
+        profile: int,
+        inp: torch.Tensor,
+        out: torch.Tensor,
+        cast_bf162half: bool,
+    ) -> None:
+        torch.ops.sgl_kernel.qr_all_reduce.default(
+            fa, profile, inp, out, cast_bf162half
+        )
+
+    def qr_destroy(fa: int) -> None:
+        torch.ops.sgl_kernel.qr_destroy.default(fa)
+
+    def qr_max_size() -> int:
+        return torch.ops.sgl_kernel.qr_max_size.default()
+
+else:
+
+    def init_custom_ar(
+        ipc_tensors: List[int], rank_data: torch.Tensor, rank: int, full_nvlink: bool
+    ) -> int:
+        return torch.ops.sgl_kernel.init_custom_ar.default(
+            ipc_tensors, rank_data, rank, full_nvlink
+        )
+
+    def dispose(fa: int) -> None:
+        torch.ops.sgl_kernel.dispose.default(fa)
+
+    def all_reduce(
+        fa: int,
+        inp: torch.Tensor,
+        out: torch.Tensor,
+        reg_buffer: int,
+        reg_buffer_sz_bytes: int,
+    ) -> None:
+        torch.ops.sgl_kernel.all_reduce.default(
+            fa, inp, out, reg_buffer, reg_buffer_sz_bytes
+        )
+
+    def get_graph_buffer_ipc_meta(fa) -> Tuple[List[int], List[int]]:
+        return torch.ops.sgl_kernel.get_graph_buffer_ipc_meta.default(fa)
+
+    def register_buffer(fa: int, fake_ipc_ptrs: List[int]) -> None:
+        return torch.ops.sgl_kernel.register_buffer.default(fa, fake_ipc_ptrs)
+
+    def register_graph_buffers(
+        fa: int, handles: List[List[int]], offsets: List[List[int]]
+    ) -> None:
+        torch.ops.sgl_kernel.register_graph_buffers.default(fa, handles, offsets)
+
+    def meta_size() -> int:
+        return torch.ops.sgl_kernel.meta_size.default()

sgl-kernel/python/sgl_kernel/attention.py

@@ -0,0 +1,113 @@
+from typing import Optional, Tuple
+
+import torch
+
+
+def merge_state_v2(
+    v_a: torch.Tensor,
+    s_a: torch.Tensor,
+    v_b: torch.Tensor,
+    s_b: torch.Tensor,
+    v_merged: Optional[torch.Tensor] = None,
+    s_merged: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    s_a = s_a.to(torch.float32)
+    s_b = s_b.to(torch.float32)
+    # TODO(DefTruth): Currently, the custom merge_attn_states kernel
+    # does not support the FP8 data type and non - CUDA devices.
+    # It may be necessary to fall back to using the Triton kernel.
+
+    # Avoid creating new tensors if they are already provided
+    if v_merged is None:
+        v_merged = torch.empty_like(v_a)
+    if s_merged is None:
+        s_merged = torch.empty_like(s_a)
+    torch.ops.sgl_kernel.merge_state_v2.default(v_a, s_a, v_b, s_b, v_merged, s_merged)
+    return v_merged, s_merged
+
+
+def cutlass_mla_decode(
+    q_nope: torch.Tensor,
+    q_pe: torch.Tensor,
+    kv_c_and_k_pe_cache: torch.Tensor,
+    seq_lens: torch.Tensor,
+    page_table: torch.Tensor,
+    workspace: torch.Tensor,
+    sm_scale: float,
+    num_kv_splits: int = 1,  # Set to 1 to avoid cuda_graph issue by default.
+) -> torch.Tensor:
+    assert q_nope.ndim == 3, f"q_nope must be a 3D tensor, but got {q_nope.ndim}"
+    assert q_pe.ndim == 3, f"q_pe must be a 3D tensor, but got {q_pe.ndim}"
+    assert (
+        kv_c_and_k_pe_cache.ndim == 3
+    ), f"kv_c_and_k_pe_cache must be a 3D tensor, but got {kv_c_and_k_pe_cache.ndim}"
+
+    B_q, H, D_q_nope = q_nope.shape
+    B_q_2, H_2, D_q_pe = q_pe.shape
+    assert (B_q == B_q_2) and (H == H_2)
+
+    _, PAGE_SIZE, D_ckv = kv_c_and_k_pe_cache.shape
+
+    D_latent = 512
+    D_rope = 64
+    assert D_q_nope == D_latent
+    assert D_q_pe == D_rope
+    assert D_ckv == D_latent + D_rope
+
+    MAX_HEADS = 128
+    assert H <= MAX_HEADS, f"H must be <= {MAX_HEADS}, but got {H}"
+    if H < MAX_HEADS:
+        q_nope_padded = q_nope.new_empty((B_q, MAX_HEADS, D_q_nope))
+        q_nope_padded[:, :H] = q_nope
+        q_nope = q_nope_padded
+
+        q_pe_padded = q_pe.new_empty((B_q, MAX_HEADS, D_q_pe))
+        q_pe_padded[:, :H] = q_pe
+        q_pe = q_pe_padded
+
+    assert len(page_table.shape) == 2
+    B_block_table, block_num = page_table.shape
+    assert B_block_table == B_q
+    assert block_num > 0, f"block num must be greater than 0, got {block_num}"
+    assert block_num % (128 / PAGE_SIZE) == 0
+
+    # TODO(kaixih@nvidia): support fp8
+    assert q_nope.dtype in (
+        torch.float16,
+        torch.bfloat16,
+    ), f"q_nope.dtype needs to be fp16 or bf16 but got {q_nope.dtype}."
+    assert q_nope.dtype == q_pe.dtype == kv_c_and_k_pe_cache.dtype
+    assert (
+        seq_lens.dtype == torch.int32
+    ), f"seq_lens.dtype needs to be int32 but got {seq_lens.dtype}."
+    assert (
+        page_table.dtype == torch.int32
+    ), f"page_table.dtype needs to be int32 but got {page_table.dtype}."
+
+    out = q_nope.new_empty((B_q, MAX_HEADS, D_latent))
+
+    torch.ops.sgl_kernel.cutlass_mla_decode.default(
+        out,
+        q_nope,
+        q_pe,
+        kv_c_and_k_pe_cache,
+        seq_lens,
+        page_table,
+        workspace,
+        sm_scale,
+        num_kv_splits,
+    )
+    return out[:, :H].contiguous()
+
+
+def cutlass_mla_get_workspace_size(
+    max_seq_len: int,
+    num_batches: int,
+    sm_count: int = 0,
+    num_kv_splits: int = 1,  # Set to 1 to avoid cuda_graph issue by default.
+) -> int:
+    assert max_seq_len > 0, f"max_seq_len must be greater than 0, got {max_seq_len}"
+    assert num_batches > 0, f"num_batches must be greater than 0, got {num_batches}"
+    return torch.ops.sgl_kernel.cutlass_mla_get_workspace_size.default(
+        max_seq_len, num_batches, sm_count, num_kv_splits
+    )

sgl-kernel/python/sgl_kernel/cutlass_moe.py

@@ -0,0 +1,112 @@
+import torch
+
+
+def get_cutlass_w4a8_moe_mm_data(
+    topk_ids: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    problem_sizes1: torch.Tensor,
+    problem_sizes2: torch.Tensor,
+    input_permutation: torch.Tensor,
+    output_permutation: torch.Tensor,
+    num_experts: int,
+    n: int,
+    k: int,
+):
+    """
+    Prepare data necessary to perform CUTLASS grouped matrix multiplications
+    used in CUTLASS-based fused MoE.
+
+    The function takes in topk_ids (token-expert mapping) and uses it to
+    compute:
+    - expert_offsets: Indices that mark at which token index each expert begins
+                      its computation after the input is sorted with
+                      input_permutation. The number of tokens computed with
+                      expert E is expert_offsets[E + 1] - expert_offsets[E]
+    - problem_sizes1, problem_sizes2: MxNxK sizes of each expert's
+                                      multiplication in two grouped MMs used in
+                                      the fused MoE operation.
+    - input_permutation: Permutation that must be used to shuffle the input
+                         before executing the MMs.
+    - output_permutation: Permutation that must be used to shuffle the output
+                          after executing the MMs.
+    """
+    torch.ops.sgl_kernel.get_cutlass_w4a8_moe_mm_data.default(
+        topk_ids,
+        expert_offsets,
+        problem_sizes1,
+        problem_sizes2,
+        input_permutation,
+        output_permutation,
+        num_experts,
+        n,
+        k,
+    )
+
+
+def cutlass_w4a8_moe_mm(
+    d: torch.Tensor,
+    a: torch.Tensor,
+    b: torch.Tensor,
+    a_scales: torch.Tensor,
+    b_scales: torch.Tensor,
+    experts_offsets: torch.tensor,
+    problem_sizes: torch.tensor,
+    a_strides: torch.tensor,
+    b_strides: torch.tensor,
+    d_strides: torch.tensor,
+    s_strides: torch.tensor,
+    chunk_size: int = 128,
+    topk: int = 8,
+):
+    """
+    Perform grouped matrix multiplication between int4 weights and fp8 activations.
+
+    This function executes multiple GEMM operations in parallel, which is useful for
+    scenarios like Mixture of Experts (MoE) where different inputs go through different
+    experts. The implementation leverages NVIDIA Hopper architecture features for
+    optimal performance with quantized weights.
+
+    Args:
+        d: Output matrices of shape [total_m, total_n]
+        a: Activation matrices in FP8 (float_e4m3_t) format
+           Each tensor should be of shape [total_m, K] in row-major layout
+        b: Weight matrices in packed int4 format
+           Each tensor should be of shape [E, N, K//2] in column-major layout
+           where each byte contains two 4-bit integers
+        a_scales: Scale factors for the inputs
+        b_scales: Scale factors for the quantized weights
+           Each tensor should be of shape [E, K//512, N*8]
+        experts_offsets: Tensor containing expert offsets for determining group boundaries
+        problem_sizes: with shape [num_experts, 3] (M, N, K for each group) (int32)
+        a_strides: Strides information for A matrices
+        b_strides: Strides information for B matrices
+        d_strides: Strides information for D matrices
+        s_strides: Strides information for b_scales matrices
+        chunk_size: Number of elements each scale value applies to (K//512), default to 128
+
+    Requirements:
+        - All tensors must be on a CUDA device
+        - Requires an NVIDIA Hopper GPU (H100)
+        - A tensors must be in float8_e4m3fn format
+        - B tensors must contain packed int4 values (stored as int8)
+
+    Note:
+        The function computes: D = (A * (B * scales))
+        for each group of tensors in parallel
+    """
+
+    torch.ops.sgl_kernel.cutlass_w4a8_moe_mm.default(
+        d,
+        a,
+        b,
+        a_scales,
+        b_scales,
+        experts_offsets,
+        problem_sizes,
+        a_strides,
+        b_strides,
+        d_strides,
+        s_strides,
+        chunk_size,
+        topk,
+    )

sgl-kernel/python/sgl_kernel/debug_utils.py

@@ -0,0 +1,45 @@
+import os
+from typing import Any, Callable, TypeVar, cast, overload
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def _wrap_debug_kernel(func: F, op_name: str | None = None) -> F:
+    try:
+        if int(os.environ.get("SGLANG_KERNEL_API_LOGLEVEL", "0")) == 0:
+            return func
+    except Exception:
+        return func
+
+    try:
+        from sglang.kernel_api_logging import debug_kernel_api
+    except Exception:
+        return func
+
+    if getattr(func, "_debug_kernel_wrapped", False):
+        return func
+
+    wrapped = debug_kernel_api(func, op_name=op_name)
+    setattr(wrapped, "_debug_kernel_wrapped", True)
+    return cast(F, wrapped)
+
+
+@overload
+def maybe_wrap_debug_kernel(func: F) -> F: ...
+
+
+@overload
+def maybe_wrap_debug_kernel(func: F, op_name: str) -> F: ...
+
+
+@overload
+def maybe_wrap_debug_kernel(*, op_name: str | None = None) -> Callable[[F], F]: ...
+
+
+def maybe_wrap_debug_kernel(
+    func: F | None = None, op_name: str | None = None
+) -> F | Callable[[F], F]:
+    if func is None:
+        return lambda wrapped_func: _wrap_debug_kernel(wrapped_func, op_name)
+
+    return _wrap_debug_kernel(func, op_name)

sgl-kernel/python/sgl_kernel/elementwise.py

@@ -0,0 +1,440 @@
+from typing import Optional
+
+import torch
+from sgl_kernel.utils import is_arch_support_pdl
+
+try:
+    import flashinfer.norm as _flashinfer_norm
+
+    _has_flashinfer = True
+except ImportError:
+    _has_flashinfer = False
+
+_FLASHINFER_NORM_SUPPORTED_DTYPES = {torch.float16, torch.bfloat16}
+
+
+def _rmsnorm_internal(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    out: Optional[torch.Tensor],
+    enable_pdl: Optional[bool],
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty_like(input)
+    if enable_pdl is None:
+        enable_pdl = is_arch_support_pdl()
+    torch.ops.sgl_kernel.rmsnorm.default(out, input, weight, eps, enable_pdl)
+    return out
+
+
+def _fused_add_rmsnorm_internal(
+    input: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    enable_pdl: Optional[bool],
+) -> None:
+    if enable_pdl is None:
+        enable_pdl = is_arch_support_pdl()
+    torch.ops.sgl_kernel.fused_add_rmsnorm.default(
+        input, residual, weight, eps, enable_pdl
+    )
+
+
+def _gemma_rmsnorm_internal(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    out: Optional[torch.Tensor],
+    enable_pdl: Optional[bool],
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty_like(input)
+    if enable_pdl is None:
+        enable_pdl = is_arch_support_pdl()
+    torch.ops.sgl_kernel.gemma_rmsnorm.default(out, input, weight, eps, enable_pdl)
+    return out
+
+
+def _gemma_fused_add_rmsnorm_internal(
+    input: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    enable_pdl: Optional[bool],
+) -> None:
+    if enable_pdl is None:
+        enable_pdl = is_arch_support_pdl()
+    torch.ops.sgl_kernel.gemma_fused_add_rmsnorm.default(
+        input, residual, weight, eps, enable_pdl
+    )
+
+
+# These implementations extensively draw from and build upon the FlashInfer project https://github.com/flashinfer-ai/flashinfer
+# Kudos to @yzh119
+def rmsnorm(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    out: Optional[torch.Tensor] = None,
+    enable_pdl: Optional[bool] = None,
+) -> torch.Tensor:
+    r"""Root mean square normalization.
+
+    ``out[i] = (input[i] / RMS(input)) * weight[i]``
+
+    Parameters
+    ----------
+    input: torch.Tensor
+        Input tensor, shape (batch_size, hidden_size).
+    weight: torch.Tensor
+        Weight tensor, shape (hidden_size,).
+    eps: float
+        Epsilon for numerical stability.
+    out: Optional[torch.Tensor]
+        The output tensor, if specified, the kernel will update this tensor inplace.
+    enable_pdl: Optional[bool]
+        Whether to enable `programmatic dependent launch
+        <https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#programmatic-dependent-launch-and-synchronization>`_
+        If None, will be automatically enabled on Hopper architecture.
+
+    Returns
+    -------
+    output: torch.Tensor
+        Normalized tensor, shape (batch_size, hidden_size).
+    """
+    # torch.compiler.is_dynamo_compiling(): FlashInfer norm paths are not safe under
+    # torch.compile(..., fullgraph=True). Dynamo traces into FlashInfer's JIT module
+    # loading path, which calls Path.exists() / os.stat() — both untraceable — causing
+    # the entire compilation to fail. We fall back to the internal implementation while
+    # tracing as a temporary workaround. Once the upstream fix is merged and we upgrade
+    # FlashInfer, this check can be removed.
+    # See: https://github.com/flashinfer-ai/flashinfer/issues/2734
+    #      https://github.com/flashinfer-ai/flashinfer/pull/2733
+    if (
+        _has_flashinfer
+        and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES
+        and not torch.compiler.is_dynamo_compiling()
+    ):
+        return _flashinfer_norm.rmsnorm(input, weight, eps, out, enable_pdl)
+    else:
+        return _rmsnorm_internal(input, weight, eps, out, enable_pdl)
+
+
+def fused_add_rmsnorm(
+    input: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    enable_pdl: Optional[bool] = None,
+) -> None:
+    r"""Fused add root mean square normalization.
+
+    Step 1:
+    ``residual[i] += input[i]``
+
+    Step 2:
+    ``input[i] = (residual[i] / RMS(residual)) * weight[i]``
+
+    Parameters
+    ----------
+    input: torch.Tensor
+        Input tensor, shape (batch_size, hidden_size).
+    residual: torch.Tensor
+        Residual tensor, shape (batch_size, hidden_size).
+    weight: torch.Tensor
+        Weight tensor, shape (hidden_size,).
+    eps: float
+        Epsilon for numerical stability.
+    enable_pdl: Optional[bool]
+        Whether to enable `programmatic dependent launch
+        <https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#programmatic-dependent-launch-and-synchronization>`_
+        If None, will be automatically enabled on Hopper architecture.
+    """
+    if (
+        _has_flashinfer
+        and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES
+        and not torch.compiler.is_dynamo_compiling()
+    ):
+        _flashinfer_norm.fused_add_rmsnorm(input, residual, weight, eps, enable_pdl)
+    else:
+        _fused_add_rmsnorm_internal(input, residual, weight, eps, enable_pdl)
+
+
+def gemma_rmsnorm(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    out: Optional[torch.Tensor] = None,
+    enable_pdl: Optional[bool] = None,
+) -> torch.Tensor:
+    r"""Gemma-style root mean square normalization.
+
+    ``out[i] = (input[i] / RMS(input)) * (weight[i] + 1)``
+
+    Parameters
+    ----------
+    input: torch.Tensor
+        Input tensor, shape (batch_size, hidden_size).
+    weight: torch.Tensor
+        Weight tensor, shape (hidden_size,).
+    eps: float
+        Epsilon for numerical stability.
+    out: Optional[torch.Tensor]
+        The output tensor, if specified, the kernel will update this tensor inplace.
+    enable_pdl: Optional[bool]
+        Whether to enable `programmatic dependent launch
+        <https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#programmatic-dependent-launch-and-synchronization>`_
+        If None, will be automatically enabled on Hopper architecture.
+
+    Returns
+    -------
+    output: torch.Tensor
+        Gemma Normalized tensor, shape (batch_size, hidden_size).
+    """
+    if (
+        _has_flashinfer
+        and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES
+        and not torch.compiler.is_dynamo_compiling()
+    ):
+        return _flashinfer_norm.gemma_rmsnorm(input, weight, eps, out, enable_pdl)
+    else:
+        return _gemma_rmsnorm_internal(input, weight, eps, out, enable_pdl)
+
+
+def gemma_fused_add_rmsnorm(
+    input: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    enable_pdl: Optional[bool] = None,
+) -> None:
+    r"""Gemma-style fused add root mean square normalization.
+
+    Step 1:
+    ``residual[i] += input[i]``
+
+    Step 2:
+    ``input[i] = (residual[i] / RMS(residual)) * (weight + 1)``
+
+    Parameters
+    ----------
+    input: torch.Tensor
+        Input tensor, shape (batch_size, hidden_size).
+    residual: torch.Tensor
+        Residual tensor, shape (batch_size, hidden_size).
+    weight: torch.Tensor
+        Weight tensor, shape (hidden_size,).
+    eps: float
+        Epsilon for numerical stability.
+    enable_pdl: Optional[bool]
+        Whether to enable `programmatic dependent launch
+        <https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#programmatic-dependent-launch-and-synchronization>`_
+        If None, will be automatically enabled on Hopper architecture.
+    """
+    if (
+        _has_flashinfer
+        and input.dtype in _FLASHINFER_NORM_SUPPORTED_DTYPES
+        and not torch.compiler.is_dynamo_compiling()
+    ):
+        _flashinfer_norm.gemma_fused_add_rmsnorm(
+            input, residual, weight, eps, enable_pdl
+        )
+    else:
+        _gemma_fused_add_rmsnorm_internal(input, residual, weight, eps, enable_pdl)
+
+
+def _check_shape(input: torch.Tensor, output: torch.Tensor) -> None:
+    assert input.ndim == output.ndim, f"{input.ndim} != {output.ndim}"
+    assert (
+        input.shape[:-1] == output.shape[:-1]
+    ), f"{input.shape[:-1]} != {output.shape[:-1]}"
+    assert (
+        input.shape[-1] == 2 * output.shape[-1]
+    ), f"{input.shape[-1]} != {2 * output.shape[-1]}"
+
+
+def silu_and_mul(input: torch.Tensor, out: torch.Tensor = None) -> torch.Tensor:
+    if input.shape[-1] * input.dtype.itemsize % 16 != 0:
+        raise ValueError("The pointers must be multiple of 16 bytes.")
+    if out is not None:
+        _check_shape(input, out)
+    else:
+        out = torch.empty(
+            input.shape[:-1] + (input.shape[-1] // 2,),
+            device=input.device,
+            dtype=input.dtype,
+        )
+    torch.ops.sgl_kernel.silu_and_mul.default(out, input)
+    return out
+
+
+def gelu_tanh_and_mul(input: torch.Tensor, out: torch.Tensor = None) -> torch.Tensor:
+    if input.shape[-1] * input.dtype.itemsize % 16 != 0:
+        raise ValueError("The pointers must be multiple of 16 bytes.")
+    if out is not None:
+        _check_shape(input, out)
+    else:
+        out = torch.empty(
+            input.shape[:-1] + (input.shape[-1] // 2,),
+            device=input.device,
+            dtype=input.dtype,
+        )
+    torch.ops.sgl_kernel.gelu_tanh_and_mul.default(out, input)
+    return out
+
+
+def gelu_and_mul(input: torch.Tensor, out: torch.Tensor = None) -> torch.Tensor:
+    if input.shape[-1] * input.dtype.itemsize % 16 != 0:
+        raise ValueError("The pointers must be multiple of 16 bytes.")
+    if out is not None:
+        _check_shape(input, out)
+    else:
+        out = torch.empty(
+            input.shape[:-1] + (input.shape[-1] // 2,),
+            device=input.device,
+            dtype=input.dtype,
+        )
+    torch.ops.sgl_kernel.gelu_and_mul.default(out, input)
+    return out
+
+
+if torch.version.hip is not None:
+
+    def gelu_quick(input: torch.Tensor, out: torch.Tensor = None) -> torch.Tensor:
+        """
+        Quick-GELU:  y = x * sigmoid(1.702 * x)
+
+        The CUDA/HIP kernel uses 128-bit (16-byte) vector loads & stores,
+        so the last-dimension byte length must be a multiple of 16 bytes.
+        """
+        if input.shape[-1] * input.dtype.itemsize % 16 != 0:
+            raise ValueError(
+                f"The last dimension ({input.shape[-1]}) x itemsize "
+                f"({input.dtype.itemsize}) must be a multiple of 16 bytes."
+            )
+
+        if out is not None:
+            assert input.shape == out.shape, f"{input.shape} != {out.shape}"
+        else:
+            out = torch.empty_like(input)
+
+        torch.ops.sgl_kernel.gelu_quick(out, input)
+        return out
+
+
+def dsv4_fused_q_norm_rope(
+    q_input: torch.Tensor,
+    freqs_cis: torch.Tensor,
+    positions: torch.Tensor,
+    eps: float = 1e-6,
+    q_output: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """DeepSeek-V4 fused Q RMSNorm (no weight) + RoPE.
+
+    Parameters
+    ----------
+    q_input  : (B, num_q_heads, head_dim) bfloat16
+    freqs_cis: (max_pos, rope_dim) float32, re/im interleaved
+    positions: (B,) int32
+    eps      : RMSNorm epsilon
+    q_output : optional pre-allocated output tensor
+    """
+    if q_output is None:
+        q_output = torch.empty_like(q_input)
+    torch.ops.sgl_kernel.dsv4_fused_q_norm_rope.default(
+        q_input, q_output, freqs_cis, positions, eps
+    )
+    return q_output
+
+
+def dsv4_fused_k_norm_rope_flashmla(
+    kv: torch.Tensor,
+    kv_weight: torch.Tensor,
+    freqs_cis: torch.Tensor,
+    positions: torch.Tensor,
+    out_loc: torch.Tensor,
+    kvcache: torch.Tensor,
+    eps: float = 1e-6,
+    page_size: int = 1,
+) -> None:
+    """DeepSeek-V4 fused K RMSNorm + RoPE + FlashMLA FP8 store.
+
+    Parameters
+    ----------
+    kv       : (B, 512) bfloat16
+    kv_weight: (512,) bfloat16
+    freqs_cis: (max_pos, 64) float32
+    positions: (B,) int32
+    out_loc  : (B,) int32  cache slot ids
+    kvcache  : (npages, page_bytes) uint8
+    eps      : RMSNorm epsilon
+    page_size: page size (power of 2)
+    """
+    torch.ops.sgl_kernel.dsv4_fused_k_norm_rope_flashmla.default(
+        kv, kv_weight, freqs_cis, positions, out_loc, kvcache, eps, page_size
+    )
+
+
+def dsv4_fused_q_indexer_rope_hadamard_quant(
+    q_input: torch.Tensor,
+    q_fp8: torch.Tensor,
+    weight: torch.Tensor,
+    weights_out: torch.Tensor,
+    weight_scale: float,
+    freqs_cis: torch.Tensor,
+    positions: torch.Tensor,
+) -> None:
+    """DeepSeek-V4 fused Q indexer: RoPE + Hadamard + FP8 quant.
+
+    Parameters
+    ----------
+    q_input    : (B, num_heads, 128) bfloat16
+    q_fp8      : (B, num_heads, 128) fp8_e4m3 output
+    weight     : (B, num_heads) bfloat16
+    weights_out: (B, num_heads, 1) float32 output
+    weight_scale: scalar
+    freqs_cis  : (max_pos, 64) float32
+    positions  : (B,) int32
+    """
+    torch.ops.sgl_kernel.dsv4_fused_q_indexer_rope_hadamard_quant.default(
+        q_input, q_fp8, weight, weights_out, weight_scale, freqs_cis, positions
+    )
+
+
+def rotary_embedding(
+    positions: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    head_size: int,
+    cos_sin_cache: torch.Tensor,
+    is_neox: bool = True,
+):
+    torch.ops.sgl_kernel.rotary_embedding.default(
+        positions, query, key, head_size, cos_sin_cache, is_neox
+    )
+
+
+def copy_to_gpu_no_ce(input: torch.Tensor, output: torch.Tensor):
+    torch.ops.sgl_kernel.copy_to_gpu_no_ce(input, output)
+
+
+def concat_mla_k(
+    k: torch.Tensor,
+    k_nope: torch.Tensor,
+    k_rope: torch.Tensor,
+):
+    torch.ops.sgl_kernel.concat_mla_k(k, k_nope, k_rope)
+
+
+def concat_mla_absorb_q(
+    a: torch.Tensor,
+    b: torch.Tensor,
+):
+    *batch_dims, _ = a.shape
+    out = torch.empty(
+        (*batch_dims, a.shape[-1] + b.shape[-1]), device=a.device, dtype=a.dtype
+    )
+    torch.ops.sgl_kernel.concat_mla_absorb_q(a, b, out)
+    return out

sgl-kernel/python/sgl_kernel/expert_specialization.py

@@ -0,0 +1,50 @@
+import torch
+
+
+def es_fp8_blockwise_scaled_grouped_mm(
+    output,
+    a,
+    b,
+    scales_a,
+    scales_b,
+    stride_a,
+    stride_b,
+    stride_d,
+    problem_sizes,
+    expert_offsets,
+    workspace,
+):
+    torch.ops.sgl_kernel.es_fp8_blockwise_scaled_grouped_mm.default(
+        output,
+        a,
+        b,
+        scales_a,
+        scales_b,
+        stride_a,
+        stride_b,
+        stride_d,
+        problem_sizes,
+        expert_offsets,
+        workspace,
+    )
+
+
+def es_sm100_mxfp8_blockscaled_grouped_mm(
+    output, a, b, sfa, sfb, problem_sizes, expert_offsets, blockscale_offsets
+):
+    torch.ops.sgl_kernel.es_sm100_mxfp8_blockscaled_grouped_mm.default(
+        a, b, sfa, sfb, output, problem_sizes, expert_offsets, blockscale_offsets
+    )
+
+
+def es_sm100_mxfp8_blockscaled_grouped_quant(
+    input, problem_sizes, expert_offsets, blockscale_offsets, quant_output, scale_factor
+):
+    torch.ops.sgl_kernel.es_sm100_mxfp8_blockscaled_grouped_quant.default(
+        input,
+        problem_sizes,
+        expert_offsets,
+        blockscale_offsets,
+        quant_output,
+        scale_factor,
+    )

sgl-kernel/python/sgl_kernel/flash_attn.py

@@ -0,0 +1,424 @@
+from functools import lru_cache
+from typing import Optional, Union
+
+import torch
+from sgl_kernel.debug_utils import maybe_wrap_debug_kernel
+
+try:
+    from sgl_kernel import flash_ops
+except:
+    raise ImportError(
+        "Can not import FA3 in sgl_kernel. Please check your installation."
+    )
+
+
+@lru_cache(maxsize=1)
+def is_fa3_supported(device=None) -> bool:
+    #  There some fa3 FYI
+    #  FA3 can fail without a enough shared memory for a some shapes, such as higher
+    #  hidden_dim or some special cases.
+    #  Right now, fa3 is supported for sm80/sm87 and sm86/sm89. The main different
+    #  Between sm80/sm87 and sm86/sm89 is the shared memory size. you can follow the link below for more information
+    #  https://docs.nvidia.com/cuda/cuda-c-programming-guide/#shared-memory-8-x
+    #  And for sgl-kernel right now, we can build fa3 on sm80/sm86/sm89/sm90a.
+    #  That means if you use A100/A*0/L20/L40/L40s/4090 you can use fa3.
+    return (torch.version.cuda >= "12.3") and (
+        torch.cuda.get_device_capability(device)[0] == 9
+        or torch.cuda.get_device_capability(device)[0] == 8
+    )
+
+
+def maybe_contiguous(x):
+    return x.contiguous() if x is not None and x.stride(-1) != 1 else x
+
+
+@maybe_wrap_debug_kernel
+def flash_attn_with_kvcache(
+    q,
+    k_cache,
+    v_cache,
+    k=None,
+    v=None,
+    qv=None,
+    rotary_cos=None,
+    rotary_sin=None,
+    cache_seqlens: Optional[Union[int, torch.Tensor]] = None,
+    cache_batch_idx: Optional[torch.Tensor] = None,
+    cache_leftpad: Optional[torch.Tensor] = None,
+    page_table: Optional[torch.Tensor] = None,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k_new: Optional[torch.Tensor] = None,
+    max_seqlen_q: Optional[int] = None,
+    rotary_seqlens: Optional[torch.Tensor] = None,
+    q_descale: Optional[torch.Tensor] = None,
+    k_descale: Optional[torch.Tensor] = None,
+    v_descale: Optional[torch.Tensor] = None,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    attention_chunk: Optional[int] = None,
+    softcap=0.0,  # 0.0 means deactivated
+    rotary_interleaved=True,
+    scheduler_metadata=None,
+    num_splits=0,  # Can be tuned for speed
+    pack_gqa=None,  # Can be tuned for speed
+    only_qv=False,  # Only use QV (skip K matmul); requires qv. Used when qk rope dim is 0.
+    sm_margin=0,  # Can be tuned if some SMs are used for communication
+    return_softmax_lse=False,
+    sinks=None,
+    score_mod=None,
+    aux_tensors=None,
+    ver=3,
+    out=None,
+):
+    """
+    If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from
+    k and v. This is useful for incremental decoding: you can pass in the cached keys/values from
+    the previous step, and update them with the new keys/values from the current step, and do
+    attention with the updated cache, all in 1 kernel.
+
+    If you pass in k / v, you must make sure that the cache is large enough to hold the new values.
+    For example, the KV cache could be pre-allocated with the max sequence length, and you can use
+    cache_seqlens to keep track of the current sequence lengths of each sequence in the batch.
+
+    Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be
+    rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.
+    If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos
+    and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.
+    If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at
+    indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens).
+
+    See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function.
+
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+
+    Note: Does not support backward pass.
+
+    Arguments:
+        q: (batch_size, seqlen, nheads, headdim)
+        k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no page_table,
+            or (num_blocks, page_block_size, nheads_k, headdim) if there's a page_table (i.e. paged KV cache)
+            page_block_size must be a multiple of 256.
+        v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim_v) if there's no page_table,
+            or (num_blocks, page_block_size, nheads_k, headdim_v) if there's a page_table (i.e. paged KV cache)
+        k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate
+            k with k_cache, starting at the indices specified by cache_seqlens.
+        v [optional]: (batch_size, seqlen_new, nheads_k, headdim_v). Similar to k.
+        qv [optional]: (batch_size, seqlen, nheads, headdim_v)
+        rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding
+            to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16.
+        rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos.
+        cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the
+            KV cache.
+        cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache.
+            If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1].
+            If the indices are not distinct, and k and v are provided, the values updated in the cache
+                 might come from any of the duplicate indices.
+        cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0.
+        page_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        attention_chunk: Optional[int]. If not None, splits the query into chunks of this size to save memory.
+        softcap: float. Anything > 0 activates softcapping attention.
+        rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in.
+            If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False,
+            rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1
+            (i.e. GPT-NeoX style).
+        num_splits: int. If > 1, split the key/value into this many chunks along the sequence.
+           If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic
+           to automatically determine the number of splits.
+           Don't change this unless you know what you are doing.
+        return_softmax_lse: bool. Whether to return the logsumexp of the attention scores.
+        score_mod [optional]: A callable that takes the attention scores and applies a modification.
+        aux_tensors [optional]: Some score_mods will want to read from global aux_tensors. This is how we thread them through to the inner kernel.
+
+    Return:
+        out: (batch_size, seqlen, nheads, headdim).
+        softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+    """
+
+    if v_cache is None:
+        raise ValueError("v_cache must be provided")
+    assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension"
+
+    if k_cache is None:
+        if not only_qv:
+            raise ValueError("k_cache can only be None when only_qv=True")
+        if q is not None:
+            k_head_size = q.shape[-1]
+            k_dtype = q.dtype
+            k_device = q.device
+        elif k is not None:
+            k_head_size = k.shape[-1]
+            k_dtype = k.dtype
+            k_device = k.device
+        else:
+            # Fallback: only_qv kernel ignores K values, so a tiny placeholder works.
+            k_head_size = 64
+            k_dtype = v_cache.dtype
+            k_device = v_cache.device
+        k_shape = (*v_cache.shape[:-1], k_head_size)
+        # The kernel path for only_qv ignores K values, but backend API still requires k tensor.
+        k_cache = torch.empty(k_shape, dtype=k_dtype, device=k_device)
+    assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension"
+
+    if q is None:
+        if not only_qv:
+            raise ValueError("q can only be None when only_qv=True")
+        if qv is None:
+            raise ValueError(
+                "q must be provided unless qv is provided with only_qv=True"
+            )
+        q_shape = (*qv.shape[:-1], k_cache.shape[-1])
+        # The kernel path for only_qv ignores q values, but backend API still requires q tensor.
+        q = torch.empty(q_shape, dtype=qv.dtype, device=qv.device)
+
+    if softmax_scale is None:
+        if only_qv:
+            if qv is None:
+                raise ValueError("only_qv=True requires qv to be provided")
+            softmax_scale = (qv.shape[-1]) ** (-0.5)
+        else:
+            softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (
+                -0.5
+            )
+    if cache_seqlens is not None and isinstance(cache_seqlens, int):
+        cache_seqlens = torch.full(
+            (q.shape[0],), cache_seqlens, dtype=torch.int32, device=v_cache.device
+        )
+        cache_seqlens = maybe_contiguous(cache_seqlens)
+
+    q, k_cache, k, v = [maybe_contiguous(x) for x in (q, k_cache, k, v)]
+    v_cache = (
+        v_cache.contiguous()
+        if v_cache.stride(-1) != 1 and v_cache.stride(-3) != 1
+        else v_cache
+    )
+    cu_seqlens_q, cu_seqlens_k_new = [
+        maybe_contiguous(x) for x in (cu_seqlens_q, cu_seqlens_k_new)
+    ]
+    page_table, cache_batch_idx, cache_leftpad = [
+        maybe_contiguous(x) for x in (page_table, cache_batch_idx, cache_leftpad)
+    ]
+    rotary_cos, rotary_sin = [maybe_contiguous(x) for x in (rotary_cos, rotary_sin)]
+    rotary_seqlens = maybe_contiguous(rotary_seqlens)
+    attention_chunk = 0 if attention_chunk is None else int(attention_chunk)
+
+    out, softmax_lse, *rest = torch.ops.sgl_kernel.fwd.default(
+        q,
+        k_cache,
+        v_cache,
+        k,
+        v,
+        qv,
+        out,  # out (pre-allocated output to avoid DtoD copy)
+        cu_seqlens_q,
+        None,  # cu_seqlens_k
+        cu_seqlens_k_new,
+        None,  # seqused_q
+        cache_seqlens,
+        max_seqlen_q,
+        None,  # max_seqlen_k
+        page_table,
+        cache_batch_idx,
+        cache_leftpad,
+        rotary_cos,
+        rotary_sin,
+        rotary_seqlens,
+        q_descale,
+        k_descale,
+        v_descale,
+        softmax_scale,
+        causal,
+        window_size[0],
+        window_size[1],
+        attention_chunk,
+        softcap,
+        rotary_interleaved,
+        scheduler_metadata,
+        num_splits,
+        pack_gqa,
+        sm_margin,
+        sinks,
+        None,  # sparse_mask_fine
+        only_qv,
+    )
+    # return (out, softmax_lse) if return_softmax_lse else out
+    return (out, softmax_lse, *rest) if return_softmax_lse else out
+
+
+@maybe_wrap_debug_kernel
+def flash_attn_varlen_func(
+    q,
+    k,
+    v,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q=None,
+    max_seqlen_k=None,
+    seqused_q=None,
+    seqused_k=None,
+    page_table=None,
+    softmax_scale=None,
+    causal=False,
+    qv=None,
+    q_descale=None,
+    k_descale=None,
+    v_descale=None,
+    window_size=(-1, -1),
+    attention_chunk=0,
+    softcap=0.0,
+    num_splits=1,
+    pack_gqa=None,
+    only_qv=False,
+    sm_margin=0,
+    return_softmax_lse=False,
+    sinks=None,
+    score_mod=None,
+    aux_tensors=None,
+    ver=3,
+    out=None,
+):
+
+    if not is_fa3_supported():
+        raise NotImplementedError(
+            "flash_attn at sgl-kernel is only supported on sm90 and above"
+        )
+
+    # FA3 requires max_seqlen_q and max_seqlen_k
+    if max_seqlen_q is None or max_seqlen_k is None:
+        raise ValueError("max_seqlen_q and max_seqlen_k are required for FA3")
+
+    if softmax_scale is None:
+        softmax_scale = (q.shape[-1] + (qv.shape[-1] if qv is not None else 0)) ** (
+            -0.5
+        )
+    attention_chunk = 0 if attention_chunk is None else int(attention_chunk)
+
+    out, softmax_lse, *rest = torch.ops.sgl_kernel.fwd.default(
+        q,
+        k,
+        v,
+        None,  # k_new
+        None,  # v_new
+        qv,  # qv
+        out,  # out
+        cu_seqlens_q,
+        cu_seqlens_k,
+        None,  # cu_seqlens_k_new
+        seqused_q,
+        seqused_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        None,  # page_table,
+        None,  # kv_batch_idx
+        None,  # leftpad_k
+        None,  # rotary cos
+        None,  # rotary sin
+        None,  # seqlens_rotary
+        q_descale,
+        k_descale,
+        v_descale,
+        softmax_scale,
+        causal,
+        window_size[0],
+        window_size[1],
+        attention_chunk,
+        softcap,
+        is_rotary_interleaved=False,
+        scheduler_metadata=None,
+        num_splits=num_splits,
+        pack_gqa=pack_gqa,
+        sm_margin=sm_margin,
+        sinks=sinks,
+        sparse_mask_fine=None,
+        only_qv=only_qv,
+    )
+
+    return (out, softmax_lse, *rest) if return_softmax_lse else out
+
+
+def get_scheduler_metadata(
+    batch_size: int,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    num_heads: int,
+    num_heads_k: int,
+    headdim: int,
+    cache_seqlens: torch.Tensor,
+    qkv_dtype=torch.bfloat16,
+    headdim_v: Optional[int] = None,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    cu_seqlens_k_new: Optional[torch.Tensor] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    leftpad_k: Optional[torch.Tensor] = None,
+    page_size: Optional[int] = None,
+    max_seqlen_k_new: int = 0,
+    causal: bool = False,
+    window_size=(-1, -1),
+    attention_chunk: int = 0,
+    has_softcap: bool = False,
+    num_splits: int = 0,
+    pack_gqa: Optional[bool] = None,
+    sm_margin: int = 0,
+):
+    """Precompute FA3 tile scheduling metadata.
+
+    Call this once per batch (not per layer) and pass the result as
+    scheduler_metadata to flash_attn_with_kvcache / flash_attn_varlen_func.
+    This avoids the prepare_varlen_num_blocks kernel running on every layer.
+    """
+    cache_seqlens = maybe_contiguous(cache_seqlens)
+    if headdim_v is None:
+        headdim_v = headdim
+
+    return torch.ops.sgl_kernel.get_scheduler_metadata(
+        batch_size,
+        max_seqlen_q,
+        max_seqlen_k,
+        num_heads,
+        num_heads_k,
+        headdim,
+        headdim_v,
+        qkv_dtype,
+        cache_seqlens,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        cu_seqlens_k_new,
+        seqused_q,
+        leftpad_k,
+        page_size,
+        max_seqlen_k_new,
+        causal,
+        window_size[0],
+        window_size[1],
+        attention_chunk,
+        has_softcap,
+        num_splits,
+        pack_gqa,
+        sm_margin,
+    )

sgl-kernel/python/sgl_kernel/flash_mla.py

@@ -0,0 +1,342 @@
+import dataclasses
+from typing import Optional, Tuple
+
+import torch
+
+try:
+    from sgl_kernel import flashmla_ops  # triggers TORCH extension registration
+except Exception as _e:
+    _flashmla_import_error = _e
+else:
+    _flashmla_import_error = None
+
+_IMPORT_ERROR = ImportError(
+    "Failed to load sgl_kernel.flashmla_ops extension. Ensure CUDA Driver >= 12.4"
+)
+
+
+@dataclasses.dataclass
+class FlashMLASchedMeta:
+    """Tile scheduler metadata for the newer FlashMLA Python API."""
+
+    @dataclasses.dataclass
+    class Config:
+        b: int
+        s_q: int
+        h_q: int
+        page_block_size: int
+        h_k: int
+        causal: bool
+        is_fp8_kvcache: bool
+        topk: Optional[int]
+        extra_page_block_size: Optional[int]
+        extra_topk: Optional[int]
+
+    have_initialized: bool = False
+    config: Optional[Config] = None
+    tile_scheduler_metadata: Optional[torch.Tensor] = None
+    num_splits: Optional[torch.Tensor] = None
+
+
+def get_mla_metadata(
+    cache_seqlens: Optional[torch.Tensor] = None,
+    num_q_tokens_per_head_k: Optional[int] = None,
+    num_heads_k: Optional[int] = None,
+    num_heads_q: Optional[int] = None,
+    is_fp8_kvcache: bool = False,
+    topk: Optional[int] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Arguments:
+        cache_seqlens: (batch_size), dtype torch.int32.
+        num_q_tokens_per_head_k: Equals to num_q_tokens_per_q_seq * num_heads_q // num_heads_k.
+        num_heads_k: The number of k heads.
+        num_heads_q: The number of q heads. This argument is optional when sparse attention is not enabled
+        is_fp8_kvcache: Whether the k_cache and v_cache are in fp8 format.
+        topk: If not None, sparse attention will be enabled, and only tokens in the `indices` array passed to `flash_mla_with_kvcache_sm90` will be attended to.
+
+    Returns:
+        tile_scheduler_metadata: (num_sm_parts, TileSchedulerMetaDataSize), dtype torch.int32.
+        num_splits: (batch_size + 1), dtype torch.int32.
+    """
+    if _flashmla_import_error is not None:
+        raise _IMPORT_ERROR from _flashmla_import_error
+
+    if cache_seqlens is None:
+        return FlashMLASchedMeta(), None
+
+    assert num_q_tokens_per_head_k is not None
+    assert num_heads_k is not None
+
+    if is_fp8_kvcache and topk is None:
+        return torch.ops.sgl_kernel.get_mla_decoding_metadata_dense_fp8.default(
+            cache_seqlens,
+            num_q_tokens_per_head_k,
+            num_heads_k,
+        )
+    return torch.ops.sgl_kernel.get_mla_decoding_metadata.default(
+        cache_seqlens,
+        num_q_tokens_per_head_k,
+        num_heads_k,
+        num_heads_q,
+        is_fp8_kvcache,
+        topk,
+    )
+
+
+def flash_mla_with_kvcache(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    block_table: Optional[torch.Tensor],
+    cache_seqlens: Optional[torch.Tensor],
+    head_dim_v: int,
+    tile_scheduler_metadata: torch.Tensor | FlashMLASchedMeta,
+    num_splits: Optional[torch.Tensor] = None,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    descale_q: torch.Tensor | None = None,
+    descale_k: torch.Tensor | None = None,
+    is_fp8_kvcache: bool = False,
+    indices: Optional[torch.Tensor] = None,
+    attn_sink: Optional[torch.Tensor] = None,
+    extra_k_cache: Optional[torch.Tensor] = None,
+    extra_indices_in_kvcache: Optional[torch.Tensor] = None,
+    topk_length: Optional[torch.Tensor] = None,
+    extra_topk_length: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Arguments:
+        q: (batch_size, seq_len_q, num_heads_q, head_dim).
+        k_cache: (num_blocks, page_block_size, num_heads_k, head_dim).
+        block_table: (batch_size, max_num_blocks_per_seq), torch.int32.
+        cache_seqlens: (batch_size), torch.int32.
+        head_dim_v: Head dimension of v.
+        tile_scheduler_metadata: (num_sm_parts, TileSchedulerMetaDataSize), torch.int32, returned by get_mla_metadata.
+        num_splits: (batch_size + 1), torch.int32, returned by get_mla_metadata.
+        softmax_scale: float. The scale of QK^T before applying softmax. Default to 1 / sqrt(head_dim).
+        causal: bool. Whether to apply causal attention mask.
+        descale_q: (batch_size), torch.float32. Descaling factors for Q, used for fp8 quantization.
+        descale_k: (batch_size), torch.float32. Descaling factors for K, used for fp8 quantization.
+        is_fp8_kvcache: bool. Whether the k_cache and v_cache are in fp8 format. For the format of FP8 KV cache, please refer to README.md
+        indices: (batch_size, seq_len_q, topk), torch.int32. If not None, sparse attention will be enabled, and only tokens in the `indices` array will be attended to. Invalid indices should be set to -1 or numbers >= total_seq_len_kv. For details about how to set up `indices`, please refer to README.md.
+
+    Returns:
+        out: (batch_size, seq_len_q, num_heads_q, head_dim_v).
+        softmax_lse: (batch_size, num_heads_q, seq_len_q), torch.float32.
+    """
+    if _flashmla_import_error is not None:
+        raise _IMPORT_ERROR from _flashmla_import_error
+
+    if softmax_scale is None:
+        softmax_scale = q.shape[-1] ** (-0.5)
+    if isinstance(tile_scheduler_metadata, FlashMLASchedMeta):
+        return _flash_mla_with_kvcache_sched_meta(
+            q=q,
+            k_cache=k_cache,
+            block_table=block_table,
+            cache_seqlens=cache_seqlens,
+            head_dim_v=head_dim_v,
+            sched_meta=tile_scheduler_metadata,
+            num_splits=num_splits,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            is_fp8_kvcache=is_fp8_kvcache,
+            indices=indices,
+            attn_sink=attn_sink,
+            extra_k_cache=extra_k_cache,
+            extra_indices_in_kvcache=extra_indices_in_kvcache,
+            topk_length=topk_length,
+            extra_topk_length=extra_topk_length,
+        )
+
+    assert num_splits is not None
+    assert block_table is not None
+    assert cache_seqlens is not None
+    assert attn_sink is None
+    assert extra_k_cache is None
+    assert extra_indices_in_kvcache is None
+    assert topk_length is None
+    assert extra_topk_length is None
+    if indices is not None:
+        assert causal == False, "causal must be `false` if sparse attention is enabled."
+    assert (descale_q is None) == (
+        descale_k is None
+    ), "descale_q and descale_k should be both None or both not None"
+
+    if indices is None and q.element_size() == 1:
+        out, softmax_lse = torch.ops.sgl_kernel.fwd_kvcache_mla_fp8.default(
+            q,
+            k_cache,
+            head_dim_v,
+            cache_seqlens,
+            block_table,
+            softmax_scale,
+            causal,
+            tile_scheduler_metadata,
+            num_splits,
+            descale_q,
+            descale_k,
+        )
+    else:
+        out, softmax_lse = torch.ops.sgl_kernel.fwd_kvcache_mla.default(
+            q,
+            k_cache,
+            head_dim_v,
+            cache_seqlens,
+            block_table,
+            softmax_scale,
+            causal,
+            tile_scheduler_metadata,
+            num_splits,
+            is_fp8_kvcache,
+            indices,
+            attn_sink,
+            extra_k_cache,
+            extra_indices_in_kvcache,
+            topk_length,
+            extra_topk_length,
+        )
+    return out, softmax_lse
+
+
+def _flash_mla_with_kvcache_sched_meta(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    block_table: Optional[torch.Tensor],
+    cache_seqlens: Optional[torch.Tensor],
+    head_dim_v: int,
+    sched_meta: FlashMLASchedMeta,
+    num_splits: Optional[torch.Tensor],
+    softmax_scale: float,
+    causal: bool,
+    is_fp8_kvcache: bool,
+    indices: Optional[torch.Tensor],
+    attn_sink: Optional[torch.Tensor],
+    extra_k_cache: Optional[torch.Tensor],
+    extra_indices_in_kvcache: Optional[torch.Tensor],
+    topk_length: Optional[torch.Tensor],
+    extra_topk_length: Optional[torch.Tensor],
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    assert num_splits is None, "num_splits must be None with FlashMLASchedMeta"
+
+    topk = indices.shape[-1] if indices is not None else None
+    extra_page_block_size = (
+        extra_k_cache.shape[1] if extra_k_cache is not None else None
+    )
+    extra_topk = (
+        extra_indices_in_kvcache.shape[-1]
+        if extra_indices_in_kvcache is not None
+        else None
+    )
+
+    if not sched_meta.have_initialized:
+        sched_meta.have_initialized = True
+        sched_meta.config = FlashMLASchedMeta.Config(
+            b=q.shape[0],
+            s_q=q.shape[1],
+            h_q=q.shape[2],
+            page_block_size=k_cache.shape[1],
+            h_k=k_cache.shape[2],
+            causal=causal,
+            is_fp8_kvcache=is_fp8_kvcache,
+            topk=topk,
+            extra_page_block_size=extra_page_block_size,
+            extra_topk=extra_topk,
+        )
+    else:
+        helper_msg = (
+            " Input arguments are inconsistent with FlashMLASchedMeta. Reuse a "
+            "scheduler only for matching tensor shapes and sparse settings."
+        )
+        assert sched_meta.config is not None
+        assert sched_meta.config.b == q.shape[0], helper_msg
+        assert sched_meta.config.s_q == q.shape[1], helper_msg
+        assert sched_meta.config.h_q == q.shape[2], helper_msg
+        assert sched_meta.config.page_block_size == k_cache.shape[1], helper_msg
+        assert sched_meta.config.h_k == k_cache.shape[2], helper_msg
+        assert sched_meta.config.causal == causal, helper_msg
+        assert sched_meta.config.is_fp8_kvcache == is_fp8_kvcache, helper_msg
+        assert sched_meta.config.topk == topk, helper_msg
+        assert (
+            sched_meta.config.extra_page_block_size == extra_page_block_size
+        ), helper_msg
+        assert sched_meta.config.extra_topk == extra_topk, helper_msg
+
+    if topk is not None:
+        assert not causal, "causal must be False when sparse attention is enabled"
+        assert is_fp8_kvcache, "is_fp8_kvcache must be True for sparse attention"
+        out, lse, new_tile_scheduler_metadata, new_num_splits = (
+            torch.ops.sgl_kernel.sparse_decode_fwd.default(
+                q,
+                k_cache,
+                indices,
+                topk_length,
+                attn_sink,
+                sched_meta.tile_scheduler_metadata,
+                sched_meta.num_splits,
+                extra_k_cache,
+                extra_indices_in_kvcache,
+                extra_topk_length,
+                head_dim_v,
+                softmax_scale,
+            )
+        )
+    else:
+        assert block_table is not None and cache_seqlens is not None
+        assert attn_sink is None
+        assert extra_k_cache is None
+        assert extra_indices_in_kvcache is None
+        assert topk_length is None
+        assert extra_topk_length is None
+        out, lse, new_tile_scheduler_metadata, new_num_splits = (
+            torch.ops.sgl_kernel.dense_decode_fwd.default(
+                q,
+                k_cache,
+                head_dim_v,
+                cache_seqlens,
+                block_table,
+                softmax_scale,
+                causal,
+                sched_meta.tile_scheduler_metadata,
+                sched_meta.num_splits,
+            )
+        )
+
+    sched_meta.tile_scheduler_metadata = new_tile_scheduler_metadata
+    sched_meta.num_splits = new_num_splits
+    return out, lse
+
+
+def flash_mla_sparse_fwd(
+    q: torch.Tensor,
+    kv: torch.Tensor,
+    indices: torch.Tensor,
+    sm_scale: float,
+    d_v: int = 512,
+    attn_sink: Optional[torch.Tensor] = None,
+    topk_length: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Sparse attention prefill kernel
+
+    Args:
+        q: [s_q, h_q, d_qk], bfloat16
+        kv: [s_kv, h_kv, d_qk], bfloat16
+        indices: [s_q, h_kv, topk], int32. Invalid indices should be set to -1 or numbers >= s_kv
+        sm_scale: float
+        d_v: The dimension of value vectors. Can only be 512
+
+    Returns:
+        (output, max_logits, lse)
+        About the definition of output, max_logits and lse, please refer to README.md
+        - output: [s_q, h_q, d_v], bfloat16
+        - max_logits:  [s_q, h_q], float
+        - lse: [s_q, h_q], float, 2-based log-sum-exp
+    """
+    if _flashmla_import_error is not None:
+        raise _IMPORT_ERROR from _flashmla_import_error
+
+    results = torch.ops.sgl_kernel.sparse_prefill_fwd.default(
+        q, kv, indices, sm_scale, d_v, attn_sink, topk_length
+    )
+    return results

sgl-kernel/python/sgl_kernel/gemm.py

@@ -0,0 +1,220 @@
+from typing import Optional
+
+import torch
+from sgl_kernel.utils import _get_cache_buf
+
+
+def awq_dequantize(
+    qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor
+) -> torch.ByteTensor:
+    return torch.ops.sgl_kernel.awq_dequantize.default(qweight, scales, qzeros)
+
+
+def int8_scaled_mm(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None):
+    return torch.ops.sgl_kernel.int8_scaled_mm.default(
+        mat_a,
+        mat_b,
+        scales_a,
+        scales_b,
+        out_dtype,
+        bias,
+    )
+
+
+def fp8_blockwise_scaled_mm(mat_a, mat_b, scales_a, scales_b, out_dtype):
+    return torch.ops.sgl_kernel.fp8_blockwise_scaled_mm.default(
+        mat_a,
+        mat_b,
+        scales_a,
+        scales_b,
+        out_dtype,
+    )
+
+
+def fp8_scaled_mm(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None):
+    return torch.ops.sgl_kernel.fp8_scaled_mm.default(
+        mat_a,
+        mat_b,
+        scales_a,
+        scales_b,
+        out_dtype,
+        bias,
+    )
+
+
+def _bmm_fp8_internal(
+    workspace_buffer: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    D: torch.Tensor,
+    A_scale: torch.Tensor,
+    B_scale: torch.Tensor,
+) -> None:
+    cublas_handle = torch.cuda.current_blas_handle()
+    torch.ops.sgl_kernel.bmm_fp8.default(
+        A,
+        B,
+        D,
+        A_scale,
+        B_scale,
+        workspace_buffer,
+        cublas_handle,
+    )
+
+
+def bmm_fp8(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    A_scale: torch.Tensor,
+    B_scale: torch.Tensor,
+    dtype: torch.dtype,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            (A.shape[0], A.shape[1], B.shape[2]),
+            device=A.device,
+            dtype=dtype,
+        )
+    workspace_buffer = _get_cache_buf("bmm_fp8_workspace", 32 * 1024 * 1024, A.device)
+    _bmm_fp8_internal(workspace_buffer, A, B, out, A_scale, B_scale)
+    return out
+
+
+def dsv3_fused_a_gemm(
+    mat_a: torch.Tensor,
+    mat_b: torch.Tensor,
+    output: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if output is None:
+        output = torch.empty(
+            (mat_a.shape[0], mat_b.shape[1]),
+            device=mat_a.device,
+            dtype=mat_a.dtype,
+        )
+    torch.ops.sgl_kernel.dsv3_fused_a_gemm.default(output, mat_a, mat_b)
+    return output
+
+
+def sgl_per_token_group_quant_8bit(
+    input: torch.Tensor,
+    output_q: torch.Tensor,
+    output_s: torch.Tensor,
+    group_size: int,
+    eps: float,
+    fp8_min: float,
+    fp8_max: float,
+    scale_ue8m0: bool = False,
+    fuse_silu_and_mul: bool = False,
+    masked_m: Optional[torch.Tensor] = None,
+    enable_v2: Optional[bool] = None,
+) -> None:
+    _V2_KERNEL_SUPPORTED_GROUP_SIZES = [16, 32, 64, 128]
+    if enable_v2 is None:
+        enable_v2 = group_size in _V2_KERNEL_SUPPORTED_GROUP_SIZES
+
+    if enable_v2:
+        return torch.ops.sgl_kernel.sgl_per_token_group_quant_8bit_v2.default(
+            input,
+            output_q,
+            output_s,
+            group_size,
+            eps,
+            fp8_min,
+            fp8_max,
+            scale_ue8m0,
+            fuse_silu_and_mul,
+            masked_m,
+        )
+
+    assert not fuse_silu_and_mul, "only v2 support fuse_silu_and_mul"
+    assert masked_m is None, "only v2 support masked_m"
+    torch.ops.sgl_kernel.sgl_per_token_group_quant_8bit.default(
+        input, output_q, output_s, group_size, eps, fp8_min, fp8_max, scale_ue8m0
+    )
+
+
+# For legacy usage
+sgl_per_token_group_quant_fp8 = sgl_per_token_group_quant_8bit
+sgl_per_token_group_quant_int8 = sgl_per_token_group_quant_8bit
+
+
+def sgl_per_token_quant_fp8(
+    input: torch.Tensor,
+    output_q: torch.Tensor,
+    output_s: torch.Tensor,
+) -> None:
+    torch.ops.sgl_kernel.sgl_per_token_quant_fp8.default(input, output_q, output_s)
+
+
+def qserve_w4a8_per_chn_gemm(
+    in_feats: torch.Tensor,
+    kernel: torch.Tensor,
+    wscales: torch.Tensor,
+    ascales: torch.Tensor,
+    w_szs: torch.Tensor,
+    a_ssums: torch.Tensor,
+    out_feats: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out_feats is None:
+        # NOTE(HandH1998): qserve_w4a8_per_chn_gemm only supports out dtype=torch.float16 now
+        out_feats = torch.empty(
+            (in_feats.shape[0], kernel.shape[0]),
+            device=in_feats.device,
+            dtype=torch.float16,
+        )
+    torch.ops.sgl_kernel.qserve_w4a8_per_chn_gemm.default(
+        in_feats, kernel, wscales, ascales, w_szs, a_ssums, out_feats
+    )
+    return out_feats
+
+
+def qserve_w4a8_per_group_gemm(
+    in_feats: torch.Tensor,
+    kernel: torch.Tensor,
+    zeros: torch.Tensor,
+    scales_i8: torch.Tensor,
+    wscales: torch.Tensor,
+    ascales: torch.Tensor,
+    out_feats: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out_feats is None:
+        # NOTE(HandH1998): qserve_w4a8_per_group_gemm only supports out dtype=torch.float16 now
+        out_feats = torch.empty(
+            (in_feats.shape[0], kernel.shape[0]),
+            device=in_feats.device,
+            dtype=torch.float16,
+        )
+    torch.ops.sgl_kernel.qserve_w4a8_per_group_gemm.default(
+        in_feats, kernel, zeros, scales_i8, wscales, ascales, out_feats
+    )
+    return out_feats
+
+
+def shuffle_rows(input_tensor, dst2src_map, output_tensor_shape):
+    output_tensor = torch.empty(
+        output_tensor_shape,
+        device=input_tensor.device,
+        dtype=input_tensor.dtype,
+    )
+    torch.ops.sgl_kernel.shuffle_rows.default(input_tensor, dst2src_map, output_tensor)
+    return output_tensor
+
+
+# GPTQ kernels
+def gptq_gemm(
+    a: torch.Tensor,
+    b_q_weight: torch.Tensor,
+    b_gptq_qzeros: torch.Tensor,
+    b_gptq_scales: torch.Tensor,
+    b_g_idx: torch.Tensor,
+    use_shuffle: bool,
+    bit: int,
+) -> torch.Tensor:
+    return torch.ops.sgl_kernel.gptq_gemm(
+        a, b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx, use_shuffle, bit
+    )
+
+
+def gptq_shuffle(q_weight: torch.Tensor, q_perm: torch.Tensor, bit: int) -> None:
+    torch.torch.ops.sgl_kernel.gptq_shuffle(q_weight, q_perm, bit)

sgl-kernel/python/sgl_kernel/grammar.py

@@ -0,0 +1,15 @@
+from typing import List, Optional, Union
+
+import torch
+
+
+def apply_token_bitmask_inplace_cuda(
+    logits: torch.Tensor,
+    bitmask: torch.Tensor,
+    indices: Optional[Union[List[int], torch.Tensor]] = None,
+) -> None:
+    if isinstance(indices, list):
+        indices = torch.tensor(indices, dtype=torch.int32, device=logits.device)
+    if indices is not None:
+        indices = indices.to(logits.device)
+    torch.ops.sgl_kernel.apply_token_bitmask_inplace_cuda(logits, bitmask, indices)

sgl-kernel/python/sgl_kernel/infllm_v2/__init__.py

@@ -0,0 +1,7 @@
+from sgl_kernel.infllm_v2.attention import infllmv2_attn_stage1
+from sgl_kernel.infllm_v2.max_pooling import max_pooling_1d_varlen
+
+__all__ = [
+    "infllmv2_attn_stage1",
+    "max_pooling_1d_varlen",
+]

sgl-kernel/python/sgl_kernel/infllm_v2/_loader.py

@@ -0,0 +1,98 @@
+"""Robust loader for the standalone ``infllm_ops`` pybind extension.
+
+The InfLLM-V2 FlashAttention backend is built as its own module ``infllm_ops``
+(installed into the ``sgl_kernel`` package directory). Under editable installs
+the compiled ``.so`` may live in ``site-packages/sgl_kernel`` while the imported
+``sgl_kernel`` package resolves to the source tree, so a plain ``from sgl_kernel
+import infllm_ops`` is not always sufficient. This loader searches the known
+candidate locations and loads the extension by file path.
+"""
+
+import glob
+import importlib.util
+import site
+import sys
+from pathlib import Path
+from typing import List, Optional
+
+_infllm_ops = None
+
+
+def _candidate_dirs() -> List[Path]:
+    dirs: List[Path] = []
+
+    # 1) The directory of the sgl_kernel package as currently imported.
+    try:
+        import sgl_kernel
+
+        dirs.append(Path(sgl_kernel.__file__).parent)
+    except Exception:
+        pass
+
+    # 2) This module's parent package directory (source tree).
+    dirs.append(Path(__file__).resolve().parent.parent)
+
+    # 3) Every ``sgl_kernel`` directory found on the install paths.
+    search_roots: List[str] = []
+    try:
+        search_roots.extend(site.getsitepackages())
+    except Exception:
+        pass
+    try:
+        search_roots.append(site.getusersitepackages())
+    except Exception:
+        pass
+    search_roots.extend(p for p in sys.path if p)
+    for root in search_roots:
+        dirs.append(Path(root) / "sgl_kernel")
+
+    # De-duplicate while preserving order.
+    seen = set()
+    unique: List[Path] = []
+    for d in dirs:
+        key = str(d)
+        if key not in seen:
+            seen.add(key)
+            unique.append(d)
+    return unique
+
+
+def _find_so() -> Optional[Path]:
+    for d in _candidate_dirs():
+        if not d.is_dir():
+            continue
+        matches = sorted(glob.glob(str(d / "infllm_ops*.so")))
+        if matches:
+            return Path(matches[0])
+    return None
+
+
+def load_infllm_ops():
+    """Import and return the ``infllm_ops`` extension module (cached)."""
+    global _infllm_ops
+    if _infllm_ops is not None:
+        return _infllm_ops
+
+    # Fast path: a normal import may already work.
+    try:
+        from sgl_kernel import infllm_ops as _mod  # type: ignore
+
+        _infllm_ops = _mod
+        return _infllm_ops
+    except Exception:
+        pass
+
+    so_path = _find_so()
+    if so_path is None:
+        raise ImportError(
+            "[sgl_kernel] Could not locate the 'infllm_ops' extension (infllm_ops*.so). "
+            "Ensure sgl-kernel was built with the InfLLM-V2 FlashAttention backend."
+        )
+
+    spec = importlib.util.spec_from_file_location("infllm_ops", str(so_path))
+    if spec is None or spec.loader is None:
+        raise ImportError(f"[sgl_kernel] Could not create module spec for {so_path}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    _infllm_ops = module
+    return _infllm_ops

sgl-kernel/python/sgl_kernel/infllm_v2/attention.py

@@ -0,0 +1,82 @@
+"""InfLLM-V2 sparse FlashAttention public API.
+
+Ported (drop-in) from ``3rdparty/infllmv2_cuda_impl/infllm_v2/infllmv2_sparse_attention.py``.
+The CUDA backend now lives in the standalone ``infllm_ops`` extension.
+"""
+
+import torch
+from sgl_kernel.infllm_v2._loader import load_infllm_ops
+
+
+def maybe_contiguous(x):
+    return x.contiguous() if x is not None and x.stride(-1) != 1 else x
+
+
+def infllmv2_attn_stage1(
+    q,
+    k,
+    v,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    cu_seqlens_v,
+    max_seqlen_q,
+    max_seqlen_k,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    softcap=0.0,
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=True,
+    block_table=None,
+):
+    """Neighborhood Sparse Attention (NSA) Stage 1 with varlen support.
+
+    Drop-in replacement for ``infllm_v2.infllmv2_attn_stage1``. Returns the
+    attention-score matrix with the NSA sparsity pattern, shape
+    ``(num_heads_k, total_q, max_seqlen_k)``.
+    """
+    infllm_ops = load_infllm_ops()
+    if softmax_scale is None:
+        softmax_scale = q.shape[-1] ** (-0.5)
+
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+
+    total_q, nheads, head_dim = q.shape
+    nheads_k = k.shape[1]
+    nheads_per_group = nheads // nheads_k
+
+    q = q.reshape(total_q, nheads_k, nheads_per_group, head_dim)
+    q = (
+        q.transpose(1, 2)
+        .reshape(total_q * nheads_per_group, nheads_k, head_dim)
+        .contiguous()
+    )
+
+    result = infllm_ops.varlen_fwd_stage1(
+        q,
+        k,
+        v,
+        None,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        cu_seqlens_v,
+        None,
+        None,
+        block_table,
+        alibi_slopes,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        True,
+        causal,
+        window_size[0],
+        window_size[1],
+        softcap,
+        True,
+        None,
+    )
+
+    return result[0]

sgl-kernel/python/sgl_kernel/infllm_v2/max_pooling.py

@@ -0,0 +1,61 @@
+import torch
+
+
+def max_pooling_1d_varlen(
+    input: torch.Tensor,  # num_heads x total_q x max_k
+    cu_seqlens_q: torch.Tensor,  # batch_size + 1
+    cu_seqlens_k: torch.Tensor,  # batch_size + 1
+    cache_lens: torch.Tensor,  # batch_size
+    max_seqlen_q: int,
+    max_context_len: int,
+    local_blocks: int,
+    init_blocks: int,
+    block_size: int = 64,
+    stride: int = 16,
+    total_q: int = -1,
+) -> torch.Tensor:
+    """Variable-length 1D max pooling over packed sequences.
+
+    Drop-in replacement for ``infllm_v2.max_pooling_1d_varlen``.
+    """
+    assert input.dtype in (torch.float16, torch.bfloat16)
+    assert cu_seqlens_q.dtype == torch.int32
+    assert cu_seqlens_k.dtype == torch.int32
+    assert cache_lens.dtype == torch.int32
+    assert input.dim() == 3, f"Expected 3D input, got {input.dim()}D"
+
+    input = input.contiguous()
+    cu_seqlens_q = cu_seqlens_q.contiguous()
+    cu_seqlens_k = cu_seqlens_k.contiguous()
+    cache_lens = cache_lens.contiguous()
+
+    max_seqlen_k = max_context_len // stride
+    out_len = (max_context_len + block_size - 1) // block_size
+
+    stride = block_size // stride
+    kernel_size = stride + 1
+    padding = 1
+
+    num_heads = input.shape[0]
+    total_q = input.shape[1]
+
+    output = torch.zeros(
+        num_heads, total_q, out_len, device=input.device, dtype=input.dtype
+    )
+    torch.ops.sgl_kernel.infllm_v2_max_pooling_1d_varlen.default(
+        input,
+        output,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        cache_lens,
+        max_seqlen_q,
+        max_seqlen_k,
+        kernel_size,
+        stride,
+        padding,
+        block_size,
+        local_blocks,
+        init_blocks,
+        total_q,
+    )
+    return output

sgl-kernel/python/sgl_kernel/kvcacheio.py

@@ -0,0 +1,321 @@
+from typing import List
+
+import torch
+
+
+def is_hip() -> bool:
+    return torch.version.hip is not None
+
+
+_is_hip = is_hip()
+
+
+def transfer_kv_per_layer(
+    src_k: torch.Tensor,
+    dst_k: torch.Tensor,
+    src_v: torch.Tensor,
+    dst_v: torch.Tensor,
+    src_indices: torch.Tensor,
+    dst_indices: torch.Tensor,
+    item_size: int,
+    block_quota: int = 2,
+    num_warps_per_block: int = 16 if _is_hip else 32,
+):
+    torch.ops.sgl_kernel.transfer_kv_per_layer.default(
+        src_k,
+        dst_k,
+        src_v,
+        dst_v,
+        src_indices,
+        dst_indices,
+        item_size,
+        block_quota,
+        num_warps_per_block,
+    )
+
+
+def transfer_kv_per_layer_pf_lf(
+    src_k: torch.Tensor,
+    dst_k: torch.Tensor,
+    src_v: torch.Tensor,
+    dst_v: torch.Tensor,
+    src_indices: torch.Tensor,
+    dst_indices: torch.Tensor,
+    layer_id: int,
+    item_size: int,
+    src_layout_dim: int,
+    block_quota: int = 2,
+    num_warps_per_block: int = 16 if _is_hip else 32,
+):
+    torch.ops.sgl_kernel.transfer_kv_per_layer_pf_lf.default(
+        src_k,
+        dst_k,
+        src_v,
+        dst_v,
+        src_indices,
+        dst_indices,
+        layer_id,
+        item_size,
+        src_layout_dim,
+        block_quota,
+        num_warps_per_block,
+    )
+
+
+def transfer_kv_per_layer_ph_lf(
+    src_k: torch.Tensor,
+    dst_k: torch.Tensor,
+    src_v: torch.Tensor,
+    dst_v: torch.Tensor,
+    src_indices: torch.Tensor,
+    dst_indices: torch.Tensor,
+    layer_id: int,
+    item_size: int,
+    src_layout_dim: int,
+    page_size: int,
+    head_num: int,
+    block_quota: int = 2,
+    num_warps_per_block: int = 16 if _is_hip else 32,
+):
+    torch.ops.sgl_kernel.transfer_kv_per_layer_ph_lf.default(
+        src_k,
+        dst_k,
+        src_v,
+        dst_v,
+        src_indices,
+        dst_indices,
+        layer_id,
+        item_size,
+        src_layout_dim,
+        page_size,
+        head_num,
+        block_quota,
+        num_warps_per_block,
+    )
+
+
+def transfer_kv_all_layer(
+    src_k_layers: torch.Tensor,
+    dst_k_layers: torch.Tensor,
+    src_v_layers: torch.Tensor,
+    dst_v_layers: torch.Tensor,
+    src_indices: torch.Tensor,
+    dst_indices: torch.Tensor,
+    item_size: int,
+    num_layers: int,
+    block_quota: int = 2,
+    num_warps_per_block: int = 16 if _is_hip else 32,
+):
+    torch.ops.sgl_kernel.transfer_kv_all_layer.default(
+        src_k_layers,
+        dst_k_layers,
+        src_v_layers,
+        dst_v_layers,
+        src_indices,
+        dst_indices,
+        item_size,
+        num_layers,
+        block_quota,
+        num_warps_per_block,
+    )
+
+
+def transfer_kv_all_layer_lf_pf(
+    src_k_layers: torch.Tensor,
+    dst_k: torch.Tensor,
+    src_v_layers: torch.Tensor,
+    dst_v: torch.Tensor,
+    src_indices: torch.Tensor,
+    dst_indices: torch.Tensor,
+    item_size: int,
+    dst_layout_dim: int,
+    num_layers: int,
+    block_quota: int = 2,
+    num_warps_per_block: int = 16 if _is_hip else 32,
+):
+    torch.ops.sgl_kernel.transfer_kv_all_layer_lf_pf.default(
+        src_k_layers,
+        dst_k,
+        src_v_layers,
+        dst_v,
+        src_indices,
+        dst_indices,
+        item_size,
+        dst_layout_dim,
+        num_layers,
+        block_quota,
+        num_warps_per_block,
+    )
+
+
+def transfer_kv_all_layer_lf_ph(
+    src_k_layers: torch.Tensor,
+    dst_k: torch.Tensor,
+    src_v_layers: torch.Tensor,
+    dst_v: torch.Tensor,
+    src_indices: torch.Tensor,
+    dst_indices: torch.Tensor,
+    item_size: int,
+    dst_layout_dim: int,
+    num_layers: int,
+    page_size: int,
+    head_num: int,
+    block_quota: int = 2,
+    num_warps_per_block: int = 16 if _is_hip else 32,
+):
+    torch.ops.sgl_kernel.transfer_kv_all_layer_lf_ph.default(
+        src_k_layers,
+        dst_k,
+        src_v_layers,
+        dst_v,
+        src_indices,
+        dst_indices,
+        item_size,
+        dst_layout_dim,
+        num_layers,
+        page_size,
+        head_num,
+        block_quota,
+        num_warps_per_block,
+    )
+
+
+def transfer_kv_direct(
+    src_layers: List[torch.Tensor],
+    dst_layers: List[torch.Tensor],
+    src_indices: torch.Tensor,
+    dst_indices: torch.Tensor,
+    page_size: int,
+):
+    torch.ops.sgl_kernel.transfer_kv_direct.default(
+        src_layers, dst_layers, src_indices, dst_indices, page_size
+    )
+
+
+def transfer_kv_per_layer_direct_pf_lf(
+    src_ptrs: List[torch.Tensor],
+    dst_ptrs: List[torch.Tensor],
+    src_indices: torch.Tensor,
+    dst_indices: torch.Tensor,
+    layer_id: int,
+    page_size: int,
+):
+    torch.ops.sgl_kernel.transfer_kv_per_layer_direct_pf_lf.default(
+        src_ptrs, dst_ptrs, src_indices, dst_indices, layer_id, page_size
+    )
+
+
+def transfer_kv_all_layer_direct_lf_pf(
+    src_ptrs: List[torch.Tensor],
+    dst_ptrs: List[torch.Tensor],
+    src_indices: torch.Tensor,
+    dst_indices: torch.Tensor,
+    page_size: int,
+):
+    torch.ops.sgl_kernel.transfer_kv_all_layer_direct_lf_pf.default(
+        src_ptrs, dst_ptrs, src_indices, dst_indices, page_size
+    )
+
+
+def transfer_kv_per_layer_mla(
+    src: torch.Tensor,
+    dst: torch.Tensor,
+    src_indices: torch.Tensor,
+    dst_indices: torch.Tensor,
+    item_size: int,
+    block_quota: int = 2,
+    num_warps_per_block: int = 16 if _is_hip else 32,
+):
+    torch.ops.sgl_kernel.transfer_kv_per_layer_mla.default(
+        src,
+        dst,
+        src_indices,
+        dst_indices,
+        item_size,
+        block_quota,
+        num_warps_per_block,
+    )
+
+
+def transfer_kv_per_layer_mla_pf_lf(
+    src: torch.Tensor,
+    dst: torch.Tensor,
+    src_indices: torch.Tensor,
+    dst_indices: torch.Tensor,
+    layer_id: int,
+    item_size: int,
+    src_layout_dim: int,
+    block_quota: int = 2,
+    num_warps_per_block: int = 16 if _is_hip else 32,
+):
+    torch.ops.sgl_kernel.transfer_kv_per_layer_mla_pf_lf.default(
+        src,
+        dst,
+        src_indices,
+        dst_indices,
+        layer_id,
+        item_size,
+        src_layout_dim,
+        block_quota,
+        num_warps_per_block,
+    )
+
+
+def transfer_kv_all_layer_mla(
+    src_layers: torch.Tensor,
+    dst_layers: torch.Tensor,
+    src_indices: torch.Tensor,
+    dst_indices: torch.Tensor,
+    item_size: int,
+    num_layers: int,
+    block_quota: int = 2,
+    num_warps_per_block: int = 16 if _is_hip else 32,
+):
+    torch.ops.sgl_kernel.transfer_kv_all_layer_mla.default(
+        src_layers,
+        dst_layers,
+        src_indices,
+        dst_indices,
+        item_size,
+        num_layers,
+        block_quota,
+        num_warps_per_block,
+    )
+
+
+def transfer_kv_all_layer_mla_lf_pf(
+    src_layers: torch.Tensor,
+    dst: torch.Tensor,
+    src_indices: torch.Tensor,
+    dst_indices: torch.Tensor,
+    item_size: int,
+    dst_layout_dim: int,
+    num_layers: int,
+    block_quota: int = 2,
+    num_warps_per_block: int = 16 if _is_hip else 32,
+):
+    torch.ops.sgl_kernel.transfer_kv_all_layer_mla_lf_pf.default(
+        src_layers,
+        dst,
+        src_indices,
+        dst_indices,
+        item_size,
+        dst_layout_dim,
+        num_layers,
+        block_quota,
+        num_warps_per_block,
+    )
+
+
+def copy_all_layer_kv_cache_cpu(
+    data_ptrs: torch.Tensor,
+    strides: torch.Tensor,
+    tgt_loc: torch.Tensor,
+    src_loc: torch.Tensor,
+):
+    torch.ops.sgl_kernel.copy_all_layer_kv_cache_cpu(
+        data_ptrs,
+        strides,
+        tgt_loc,
+        src_loc,
+    )

sgl-kernel/python/sgl_kernel/load_utils.py

@@ -0,0 +1,247 @@
+import ctypes
+import glob
+import importlib.util
+import logging
+import os
+import shutil
+from pathlib import Path
+from typing import List
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+def _get_compute_capability():
+    """Get the compute capability of the current GPU."""
+    if not torch.cuda.is_available():
+        return None
+
+    # Get the current device
+    device = torch.cuda.current_device()
+    properties = torch.cuda.get_device_properties(device)
+
+    # Return as integer (major * 10 + minor)
+    return properties.major * 10 + properties.minor
+
+
+def _filter_compiled_extensions(file_list):
+    """Filter and prioritize compiled extensions over Python source files."""
+    compiled_extensions = [".so", ".pyd", ".dll"]  # Common compiled extension suffixes
+    compiled_files = []
+    other_files = []
+
+    for file_path in file_list:
+        path = Path(file_path)
+        # Check if it's a compiled extension (including complex names like .abi3.so, .cpython-312.so)
+        if any(
+            str(path).endswith(ext) or ext in str(path) for ext in compiled_extensions
+        ):
+            compiled_files.append(file_path)
+        else:
+            other_files.append(file_path)
+
+    # Return compiled files first, then others
+    return compiled_files + other_files
+
+
+def _load_architecture_specific_ops():
+    """Load the appropriate common_ops library based on GPU architecture."""
+    compute_capability = _get_compute_capability()
+    logger.debug(
+        f"[sgl_kernel] GPU Detection: compute_capability = {compute_capability}"
+    )
+
+    # Get the directory where sgl_kernel is installed
+    sgl_kernel_dir = Path(__file__).parent
+    logger.debug(f"[sgl_kernel] sgl_kernel directory: {sgl_kernel_dir}")
+
+    # Determine which version to load based on GPU architecture
+    if compute_capability == 90:
+        ops_subdir = "sm90"
+        variant_name = "SM90 (Hopper/H100 with fast math optimization)"
+    elif compute_capability is not None:
+        ops_subdir = "sm100"
+        variant_name = f"SM{compute_capability} (precise math for compatibility)"
+    else:
+        ops_subdir = "sm100"
+        variant_name = "CPU/No GPU detected (using precise math)"
+
+    # Look for the compiled module with any valid extension
+
+    ops_pattern = str(sgl_kernel_dir / ops_subdir / "common_ops.*")
+    raw_matching_files = glob.glob(ops_pattern)
+    matching_files = _filter_compiled_extensions(raw_matching_files)
+
+    logger.debug(f"[sgl_kernel] Attempting to load {variant_name}")
+    logger.debug(f"[sgl_kernel] Looking for library matching pattern: {ops_pattern}")
+    logger.debug(f"[sgl_kernel] Found files: {raw_matching_files}")
+    logger.debug(f"[sgl_kernel] Prioritized files: {matching_files}")
+
+    previous_import_errors: List[Exception] = []
+
+    # Try to load from the architecture-specific directory
+    if matching_files:
+        ops_path = Path(matching_files[0])  # Use the first prioritized file
+        logger.debug(f"[sgl_kernel] Found architecture-specific library: {ops_path}")
+        try:
+            # Load the module from specific path using importlib
+            spec = importlib.util.spec_from_file_location("common_ops", str(ops_path))
+            if spec is None:
+                raise ImportError(f"Could not create module spec for {ops_path}")
+
+            common_ops = importlib.util.module_from_spec(spec)
+            if spec.loader is None:
+                raise ImportError(f"Module spec has no loader for {ops_path}")
+
+            logger.debug(f"[sgl_kernel] Loading module from {ops_path}...")
+            spec.loader.exec_module(common_ops)
+            logger.debug(f"[sgl_kernel] ✓ Successfully loaded {variant_name}")
+            logger.debug(f"[sgl_kernel] ✓ Module file: {common_ops.__file__}")
+            return common_ops
+
+        except Exception as e:
+            previous_import_errors.append(e)
+            logger.debug(
+                f"[sgl_kernel] ✗ Failed to load from {ops_path}: {type(e).__name__}: {e}"
+            )
+            # Continue to fallback
+    else:
+        logger.debug(
+            f"[sgl_kernel] ✗ Architecture-specific library not found matching pattern: {ops_pattern}"
+        )
+
+    # Try alternative directory (in case installation structure differs)
+    alt_pattern = str(sgl_kernel_dir / "common_ops.*")
+    raw_alt_files = glob.glob(alt_pattern)
+    alt_matching_files = _filter_compiled_extensions(raw_alt_files)
+    logger.debug(f"[sgl_kernel] Attempting fallback: looking for pattern {alt_pattern}")
+    logger.debug(f"[sgl_kernel] Found fallback files: {raw_alt_files}")
+    logger.debug(f"[sgl_kernel] Prioritized fallback files: {alt_matching_files}")
+
+    if alt_matching_files:
+        alt_path = Path(alt_matching_files[0])  # Use the first prioritized file
+        logger.debug(f"[sgl_kernel] Found fallback library: {alt_path}")
+        try:
+            spec = importlib.util.spec_from_file_location("common_ops", str(alt_path))
+            if spec is None:
+                raise ImportError(f"Could not create module spec for {alt_path}")
+
+            common_ops = importlib.util.module_from_spec(spec)
+            if spec.loader is None:
+                raise ImportError(f"Module spec has no loader for {alt_path}")
+
+            logger.debug(f"[sgl_kernel] Loading fallback module from {alt_path}...")
+            spec.loader.exec_module(common_ops)
+            logger.debug(f"[sgl_kernel] ✓ Successfully loaded fallback library")
+            logger.debug(f"[sgl_kernel] ✓ Module file: {common_ops.__file__}")
+            return common_ops
+
+        except Exception as e:
+            previous_import_errors.append(e)
+            logger.debug(
+                f"[sgl_kernel] ✗ Failed to load fallback from {alt_path}: {type(e).__name__}: {e}"
+            )
+    else:
+        logger.debug(
+            f"[sgl_kernel] ✗ Fallback library not found matching pattern: {alt_pattern}"
+        )
+
+    # Final attempt: try standard Python import (for backward compatibility)
+    logger.debug(
+        f"[sgl_kernel] Final attempt: trying standard Python import 'common_ops'"
+    )
+    try:
+        import common_ops
+
+        logger.debug(f"[sgl_kernel] ✓ Successfully imported via standard Python import")
+        logger.debug(f"[sgl_kernel] ✓ Module file: {common_ops.__file__}")
+        return common_ops
+    except ImportError as e:
+        previous_import_errors.append(e)
+        logger.debug(f"[sgl_kernel] ✗ Standard Python import failed: {e}")
+
+    attempt_error_msg = "\n".join(
+        f"- {type(err).__name__}: {err}" for err in previous_import_errors
+    )
+
+    # All attempts failed
+    cuda_version = torch.version.cuda
+    if cuda_version and cuda_version.startswith("12"):
+        install_hint = (
+            "pip install sglang-kernel --index-url https://docs.sglang.ai/whl/cu129/"
+        )
+    else:
+        install_hint = "pip install --upgrade sglang-kernel"
+
+    error_msg = f"""
+[sgl_kernel] CRITICAL: Could not load any common_ops library!
+
+Attempted locations:
+1. Architecture-specific pattern: {ops_pattern} - found files: {matching_files}
+2. Fallback pattern: {alt_pattern} - found files: {alt_matching_files}
+3. Standard Python import: common_ops - failed
+
+GPU Info:
+- Compute capability: {compute_capability}
+- Expected variant: {variant_name}
+- CUDA version: {cuda_version}
+
+Please ensure sgl_kernel is properly installed with:
+{install_hint}
+
+Error details from previous import attempts:
+{attempt_error_msg}
+"""
+    logger.debug(error_msg)
+    raise ImportError(error_msg)
+
+
+# copy & modify from torch/utils/cpp_extension.py
+def _find_cuda_home():
+    """Find the CUDA install path."""
+    # Guess #1
+    cuda_home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH")
+    if cuda_home is None:
+        # Guess #2
+        nvcc_path = shutil.which("nvcc")
+        if nvcc_path is not None:
+            cuda_home = os.path.dirname(os.path.dirname(nvcc_path))
+        else:
+            # Guess #3
+            cuda_home = "/usr/local/cuda"
+    return cuda_home
+
+
+def _preload_cuda_library():
+    """Preload the CUDA runtime library to help avoid 'libcudart.so not found' issues."""
+    cuda_home = Path(_find_cuda_home())
+
+    candidate_dirs = [
+        cuda_home / "lib",
+        cuda_home / "lib64",
+        Path("/usr/lib/x86_64-linux-gnu"),
+        Path("/usr/lib/aarch64-linux-gnu"),
+        Path("/usr/lib64"),
+        Path("/usr/lib"),
+    ]
+
+    # Determine CUDA major version to try the matching library first.
+    # On CUDA 13 systems (e.g., DGX Spark), only libcudart.so.13 exists.
+    cuda_major = torch.version.cuda.split(".")[0] if torch.version.cuda else "12"
+    lib_versions = list(dict.fromkeys([cuda_major, "13", "12"]))
+
+    for base in candidate_dirs:
+        for lib_version in lib_versions:
+            candidate = base / f"libcudart.so.{lib_version}"
+            if candidate.exists():
+                try:
+                    cuda_runtime_lib = candidate.resolve()
+                    ctypes.CDLL(str(cuda_runtime_lib), mode=ctypes.RTLD_GLOBAL)
+                    logger.debug(f"Preloaded CUDA runtime under {cuda_runtime_lib}")
+                    return
+                except Exception as e:
+                    logger.debug(f"Failed to load {candidate}: {e}")
+                    continue
+
+    logger.debug("[sgl_kernel] Could not preload CUDA runtime library")

sgl-kernel/python/sgl_kernel/mamba.py

@@ -0,0 +1,122 @@
+from typing import Optional
+
+import torch
+
+
+# mamba
+def causal_conv1d_fwd(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias_: Optional[torch.Tensor],
+    conv_states: Optional[torch.Tensor],
+    query_start_loc: Optional[torch.Tensor],
+    cache_indices: Optional[torch.Tensor],
+    has_initial_state: Optional[torch.Tensor],
+    silu_activation: bool,
+    pad_slot_id: int,
+):
+    torch.ops.sgl_kernel.causal_conv1d_fwd(
+        x,
+        weight,
+        bias_,
+        conv_states,
+        query_start_loc,
+        cache_indices,
+        has_initial_state,
+        silu_activation,
+        pad_slot_id,
+    )
+
+
+def causal_conv1d_update(
+    x: torch.Tensor,
+    conv_state: torch.Tensor,
+    weight: torch.Tensor,
+    bias_: Optional[torch.Tensor],
+    silu_activation: bool,
+    cache_seqlens: Optional[torch.Tensor],
+    conv_state_indices: Optional[torch.Tensor],
+    pad_slot_id: int,
+):
+    torch.ops.sgl_kernel.causal_conv1d_update(
+        x,
+        conv_state,
+        weight,
+        bias_,
+        silu_activation,
+        cache_seqlens,
+        conv_state_indices,
+        pad_slot_id,
+    )
+
+
+def causal_conv1d_fn_cpu(
+    mixed_qkv_transposed,
+    conv_weights,
+    bias,
+    activation,
+    conv_states,
+    has_initial_state,
+    cache_indices,
+    query_start_loc,
+    seq_lens_cpu,
+):
+    return torch.ops.sgl_kernel.causal_conv1d_fwd_cpu(
+        mixed_qkv_transposed,
+        conv_weights,
+        bias,
+        conv_states,
+        query_start_loc,
+        cache_indices,
+        has_initial_state,
+        activation == "silu",
+        -1,
+        True,
+    )
+
+
+def causal_conv1d_update_cpu(
+    mixed_qkv, conv_states, conv_weights, bias, activation, conv_state_indices
+):
+    return torch.ops.sgl_kernel.causal_conv1d_update_cpu(
+        mixed_qkv,
+        conv_states,
+        conv_weights,
+        bias,
+        activation == "silu",
+        None,
+        conv_state_indices,
+        -1,
+        True,
+    )
+
+
+def chunk_gated_delta_rule_cpu(
+    q,
+    k,
+    v,
+    g,
+    beta,
+    initial_state,
+    cu_seqlens,
+    head_first,
+    use_qk_l2norm_in_kernel,
+    initial_state_indices,
+):
+    core_attn_out, last_recurrent_state = (
+        torch.ops.sgl_kernel.chunk_gated_delta_rule_cpu(
+            q,
+            k,
+            v,
+            g,
+            beta,
+            initial_state,
+            True,  # output_final_state
+            cu_seqlens,
+            head_first,
+            use_qk_l2norm_in_kernel,
+            initial_state_indices,
+        )
+    )
+    h = None  # Todo: add return h support
+    return core_attn_out, last_recurrent_state, h

sgl-kernel/python/sgl_kernel/memory.py

@@ -0,0 +1,9 @@
+import torch
+
+
+def weak_ref_tensor(tensor):
+    return (
+        torch.ops.sgl_kernel.weak_ref_tensor(tensor)
+        if isinstance(tensor, torch.Tensor)
+        else tensor
+    )

sgl-kernel/python/sgl_kernel/metal.py

@@ -0,0 +1,112 @@
+"""Python entry points for the sgl_kernel Metal extension."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import mlx.core as mx
+
+_METALLIB_NAME = "sgl_metal_kernels.metallib"
+
+try:
+    from . import _metal
+
+    _metallib_path = Path(_metal.__file__).resolve().parent / _METALLIB_NAME
+    if not _metallib_path.is_file():
+        raise ImportError(
+            f"{_METALLIB_NAME} not found next to the native Metal extension "
+            f"at {_metallib_path}"
+        )
+    _metal.register_library(str(_metallib_path))
+except ImportError as _exc:  # pragma: no cover - import guarded at call time
+    _metal = None
+    _IMPORT_ERROR: Exception | None = _exc
+else:
+    _IMPORT_ERROR = None
+
+# Python wrappers for the compiled `_metal.*` entry points go below. Wrappers
+# validate input shapes/dtypes and then invoke AOT C++ entry points. They do
+# not force `mx.eval`, so MLX can keep these calls inside its lazy graph.
+
+
+def rope_pool_fused(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    positions: mx.array,
+    slots: mx.array,
+    k_pool: mx.array,
+    v_pool: mx.array,
+    *,
+    head_dim: int,
+    num_qo_heads: int,
+    num_kv_heads: int,
+    rope_base: float,
+) -> tuple[mx.array, mx.array, mx.array, mx.array]:
+    """Apply NeoX RoPE to Q/K and scatter K/V into the MLX KV pool.
+
+    Args:
+        q: Query tensor with shape `[num_tokens, num_qo_heads, head_dim]`.
+        k: Key tensor with shape `[num_tokens, num_kv_heads, head_dim]`.
+        v: Value tensor with shape `[num_tokens, num_kv_heads, head_dim]`.
+        positions: int32 positions with shape `[num_tokens]`.
+        slots: int32 KV-pool slots with shape `[num_tokens]`; values `< 0`
+            skip the pool write for that token.
+        k_pool: Existing K pool with shape `[pool_size, num_kv_heads, head_dim]`.
+        v_pool: Existing V pool with shape `[pool_size, num_kv_heads, head_dim]`.
+
+    Returns:
+        `(q_rot, k_rot, k_pool_new, v_pool_new)`.
+    """
+    if q.ndim != 3 or k.ndim != 3 or v.ndim != 3:
+        raise ValueError("rope_pool_fused expects q/k/v to be 3-D")
+    if positions.ndim != 1 or slots.ndim != 1:
+        raise ValueError("rope_pool_fused expects positions/slots to be 1-D")
+    if k_pool.ndim != 3 or v_pool.ndim != 3:
+        raise ValueError("rope_pool_fused expects pool tensors to be 3-D")
+    q_shape = tuple(q.shape)
+    k_shape = tuple(k.shape)
+    v_shape = tuple(v.shape)
+    positions_shape = tuple(positions.shape)
+    slots_shape = tuple(slots.shape)
+    k_pool_shape = tuple(k_pool.shape)
+    v_pool_shape = tuple(v_pool.shape)
+
+    if q_shape != (q_shape[0], num_qo_heads, head_dim):
+        raise ValueError(
+            "q shape must be [num_tokens, num_qo_heads, head_dim], " f"got {q.shape}"
+        )
+    if k_shape != (q_shape[0], num_kv_heads, head_dim):
+        raise ValueError(
+            "k shape must be [num_tokens, num_kv_heads, head_dim], " f"got {k.shape}"
+        )
+    if v_shape != k_shape:
+        raise ValueError(f"v shape must match k shape, got {v.shape} vs {k.shape}")
+    if positions_shape != (q_shape[0],) or slots_shape != (q_shape[0],):
+        raise ValueError("positions/slots must have one entry per token")
+    if k_pool_shape[1:] != (num_kv_heads, head_dim):
+        raise ValueError(f"k_pool has incompatible shape {k_pool.shape}")
+    if v_pool_shape != k_pool_shape:
+        raise ValueError(
+            f"v_pool shape must match k_pool shape, got {v_pool.shape} vs {k_pool.shape}"
+        )
+    if q.dtype != k.dtype or q.dtype != v.dtype:
+        raise ValueError("q/k/v dtypes must match")
+    if k_pool.dtype != q.dtype or v_pool.dtype != q.dtype:
+        raise ValueError("pool dtypes must match q/k/v dtype")
+
+    return _metal.rope_pool_fused(
+        q,
+        k,
+        v,
+        positions,
+        slots,
+        k_pool,
+        v_pool,
+        head_dim,
+        num_qo_heads,
+        num_kv_heads,
+        float(rope_base),
+    )

sgl-kernel/python/sgl_kernel/moe.py

@@ -0,0 +1,224 @@
+from typing import Optional
+
+import torch
+
+
+def moe_align_block_size(
+    topk_ids,
+    num_experts,
+    block_size,
+    sorted_token_ids,
+    experts_ids,
+    num_tokens_post_pad,
+    cumsum_buffer,
+    pad_sorted_token_ids=False,
+):
+    torch.ops.sgl_kernel.moe_align_block_size.default(
+        topk_ids,
+        num_experts,
+        block_size,
+        sorted_token_ids,
+        experts_ids,
+        num_tokens_post_pad,
+        cumsum_buffer,
+        pad_sorted_token_ids,
+    )
+
+
+def topk_softmax(
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    gating_output: torch.Tensor,
+    renormalize: bool = False,
+    moe_softcapping: float = 0.0,
+    correction_bias: Optional[torch.Tensor] = None,
+) -> None:
+    """
+    Compute top-k softmax for MoE routing.
+
+    Args:
+        topk_weights: Output tensor for top-k weights [num_tokens, topk]
+        topk_ids: Output tensor for top-k expert indices [num_tokens, topk]
+        gating_output: Gating logits [num_tokens, num_experts]
+        renormalize: Whether to renormalize the top-k weights
+        moe_softcapping: Tanh softcapping value (0.0 to disable)
+        correction_bias: Per-expert bias correction [num_experts], must be float32 if provided
+    """
+    torch.ops.sgl_kernel.topk_softmax.default(
+        topk_weights,
+        topk_ids,
+        gating_output,
+        renormalize,
+        moe_softcapping,
+        correction_bias,
+    )
+
+
+def topk_sigmoid(
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    gating_output: torch.Tensor,
+    renormalize: bool = False,
+    correction_bias: Optional[torch.Tensor] = None,
+) -> None:
+    """
+    Compute top-k sigmoid for MoE routing.
+
+    Args:
+        topk_weights: Output tensor for top-k weights [num_tokens, topk]
+        topk_ids: Output tensor for top-k expert indices [num_tokens, topk]
+        gating_output: Gating logits [num_tokens, num_experts]
+        renormalize: Whether to renormalize the top-k weights
+        correction_bias: Per-expert bias correction [num_experts], must be float32 if provided
+    """
+    torch.ops.sgl_kernel.topk_sigmoid.default(
+        topk_weights,
+        topk_ids,
+        gating_output,
+        renormalize,
+        correction_bias,
+    )
+
+
+def moe_sum_reduce(
+    input_tensor,
+    output_tensor,
+    routed_scaling_factor=0,
+):
+    torch.ops.sgl_kernel.moe_sum_reduce.default(
+        input_tensor,
+        output_tensor,
+        routed_scaling_factor,
+    )
+
+
+def moe_sum(
+    input_tensor: torch.Tensor,
+    output_tensor: torch.Tensor,
+):
+    torch.ops.sgl_kernel.moe_sum.default(
+        input_tensor,
+        output_tensor,
+    )
+
+
+# moe_fused_gate / kimi_k2_moe_fused_gate (AOT gate kernels) retired — the gate/topk
+# path is consolidated onto the unified Triton router in
+# python/sglang/jit_kernel/moe_fused_gate.py (sglang issue #26771).
+
+
+def fp8_blockwise_scaled_grouped_mm(
+    output,
+    a_ptrs,
+    b_ptrs,
+    out_ptrs,
+    a_scales_ptrs,
+    b_scales_ptrs,
+    a,
+    b,
+    scales_a,
+    scales_b,
+    stride_a,
+    stride_b,
+    stride_c,
+    layout_sfa,
+    layout_sfb,
+    problem_sizes,
+    expert_offsets,
+    workspace,
+):
+    torch.ops.sgl_kernel.fp8_blockwise_scaled_grouped_mm.default(
+        output,
+        a_ptrs,
+        b_ptrs,
+        out_ptrs,
+        a_scales_ptrs,
+        b_scales_ptrs,
+        a,
+        b,
+        scales_a,
+        scales_b,
+        stride_a,
+        stride_b,
+        stride_c,
+        layout_sfa,
+        layout_sfb,
+        problem_sizes,
+        expert_offsets,
+        workspace,
+    )
+
+
+def prepare_moe_input(
+    topk_ids,
+    expert_offsets,
+    problem_sizes1,
+    problem_sizes2,
+    input_permutation,
+    output_permutation,
+    num_experts,
+    n,
+    k,
+    blockscale_offsets: Optional[torch.Tensor] = None,
+):
+    torch.ops.sgl_kernel.prepare_moe_input.default(
+        topk_ids,
+        expert_offsets,
+        blockscale_offsets,
+        problem_sizes1,
+        problem_sizes2,
+        input_permutation,
+        output_permutation,
+        num_experts,
+        n,
+        k,
+    )
+
+
+def apply_shuffle_mul_sum(
+    input,
+    output,
+    permutation,
+    factors,
+):
+    torch.ops.sgl_kernel.apply_shuffle_mul_sum.default(
+        input, output, permutation, factors
+    )
+
+
+def fused_qk_norm_rope(
+    qkv: torch.Tensor,
+    num_heads_q: int,
+    num_heads_k: int,
+    num_heads_v: int,
+    head_dim: int,
+    eps: float,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    base: float,
+    is_neox: bool,
+    position_ids: torch.Tensor,
+    factor: float,
+    low: float,
+    high: float,
+    attention_factor: float,
+    rotary_dim: Optional[int] = None,
+) -> None:
+    torch.ops.sgl_kernel.fused_qk_norm_rope(
+        qkv,
+        num_heads_q,
+        num_heads_k,
+        num_heads_v,
+        head_dim,
+        eps,
+        q_weight,
+        k_weight,
+        base,
+        is_neox,
+        position_ids,
+        factor,
+        low,
+        high,
+        attention_factor,
+        rotary_dim if rotary_dim is not None else head_dim,
+    )

sgl-kernel/python/sgl_kernel/musa.py

@@ -0,0 +1,353 @@
+from typing import Optional, Union
+
+import torch
+from sgl_kernel.utils import _to_tensor_scalar_tuple
+
+
+def musa_batched_rotary_embedding_contiguous(
+    positions: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    head_size: int,
+    cos_sin_cache: torch.Tensor,
+    is_neox: bool,
+    rot_dim: int,
+    cos_sin_cache_offsets: torch.Tensor,
+) -> None:
+    return torch.ops.sgl_kernel.musa_batched_rotary_embedding_contiguous(
+        positions,
+        query,
+        key,
+        head_size,
+        cos_sin_cache,
+        is_neox,
+        rot_dim,
+        cos_sin_cache_offsets,
+    )
+
+
+def musa_rotary_embedding_contiguous(
+    positions: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    head_size: int,
+    cos_sin_cache: torch.Tensor,
+    is_neox: bool,
+) -> None:
+    return torch.ops.sgl_kernel.musa_rotary_embedding_contiguous(
+        positions,
+        query,
+        key,
+        head_size,
+        cos_sin_cache,
+        is_neox,
+    )
+
+
+def musa_fused_moe_gemv(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    A_scale,
+    B_scale,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    mul_routed_weight: bool,
+    topk: int,
+    use_int4_w4a16: bool,
+    use_swigelu: bool,
+) -> None:
+    return torch.ops.sgl_kernel.musa_fused_moe_gemv(
+        A,
+        B,
+        C,
+        A_scale,
+        B_scale,
+        topk_weights,
+        topk_ids,
+        mul_routed_weight,
+        topk,
+        use_int4_w4a16,
+        use_swigelu,
+    )
+
+
+def musa_fused_gemv(
+    x: torch.Tensor,
+    qweight: torch.Tensor,
+    x_scales: Optional[torch.Tensor] = None,
+    qweight_scales: Optional[torch.Tensor] = None,
+    use_swigelu: bool = False,
+    use_rms_norm: bool = False,
+    gamma: Optional[torch.Tensor] = None,
+    eps: float = 1e-6,
+):
+    use_int4_w4a16 = False
+    out_shape = x.shape[:-1] + (
+        qweight.shape[0] if not use_swigelu else qweight.shape[0] // 2,
+    )
+    assert not (
+        use_swigelu and use_rms_norm
+    ), "gemv only fused one activation (swigelu or rms_norm)!"
+
+    if use_rms_norm:
+        if gamma is None:
+            assert False, "rms_norm gamma is None!"
+
+    # fp8 grouped matmul
+    if qweight.dtype == torch.float8_e4m3fn:
+        assert qweight_scales is not None, "FP8 grouped matmul weight scales is None!"
+        output = torch.empty(out_shape, device=x.device, dtype=torch.bfloat16)
+        torch.ops.sgl_kernel.musa_fused_gemv(
+            x,
+            qweight,
+            output,
+            x_scales,
+            qweight_scales,
+            use_int4_w4a16,
+            use_swigelu,
+            use_rms_norm,
+            gamma,
+            eps,
+        )
+        return output
+    # w4a16 gemv
+    elif qweight_scales is not None:
+        assert (
+            x.dtype == torch.bfloat16 or x.dtype == torch.float16
+        ), "W4A16 gemv only support bfloat16 or float16!"
+        use_int4_w4a16 = True
+        out_shape = x.shape[:-1] + (
+            qweight.shape[0] if not use_swigelu else qweight.shape[0] // 2,
+        )
+        output = torch.empty(out_shape, device=x.device, dtype=x.dtype)
+        torch.ops.sgl_kernel.musa_fused_gemv(
+            x,
+            qweight,
+            output,
+            None,
+            qweight_scales,
+            use_int4_w4a16,
+            use_swigelu,
+            use_rms_norm,
+            gamma,
+            eps,
+        )
+        return output
+    # general gemv
+    else:
+        output = torch.empty(out_shape, device=x.device, dtype=x.dtype)
+        torch.ops.sgl_kernel.musa_fused_gemv(
+            x,
+            qweight,
+            output,
+            None,
+            None,
+            use_int4_w4a16,
+            use_swigelu,
+            use_rms_norm,
+            gamma,
+            eps,
+        )
+        return output
+
+
+def musa_fused_mul_add(
+    self: torch.Tensor,
+    bias: Optional[torch.Tensor],
+    scale: Optional[float],
+    accurate: bool = True,
+):
+    # if accurate == False, then we call inplace op: bias += (self * scale)
+    if not accurate:
+        bias.add_(self, alpha=scale)
+        return bias
+
+    # otherwise, we call custom outplace op, act: output = self * scale + bias
+    output = torch.empty_like(self)
+    torch.ops.sgl_kernel.musa_fused_mul_add(output, self, bias, scale)
+
+    return output
+
+
+def _top_k_renorm_probs_internal(
+    probs: torch.Tensor,
+    maybe_top_k_arr: Optional[torch.Tensor],
+    top_k_val: int,
+) -> torch.Tensor:
+    probs = probs.float()
+    maybe_top_k_arr = maybe_top_k_arr.int() if maybe_top_k_arr is not None else None
+    renorm_probs = torch.empty_like(probs)
+    torch.ops.sgl_kernel.top_k_renorm_probs.default(
+        probs, renorm_probs, maybe_top_k_arr, top_k_val
+    )
+    return renorm_probs
+
+
+def top_k_renorm_probs(
+    probs: torch.Tensor,
+    top_k: Union[torch.Tensor, int],
+) -> torch.Tensor:
+    return _top_k_renorm_probs_internal(probs, *_to_tensor_scalar_tuple(top_k))
+
+
+def _top_p_renorm_probs_internal(
+    probs: torch.Tensor,
+    maybe_top_p_arr: Optional[torch.Tensor],
+    top_p_val: float,
+) -> torch.Tensor:
+    probs = probs.float()
+    maybe_top_p_arr = maybe_top_p_arr.float() if maybe_top_p_arr is not None else None
+    renorm_probs = torch.empty_like(probs)
+    torch.ops.sgl_kernel.top_p_renorm_probs.default(
+        probs, renorm_probs, maybe_top_p_arr, top_p_val
+    )
+    return renorm_probs
+
+
+def top_p_renorm_probs(
+    probs: torch.Tensor,
+    top_p: Union[torch.Tensor, float],
+) -> torch.Tensor:
+    return _top_p_renorm_probs_internal(probs, *_to_tensor_scalar_tuple(top_p))
+
+
+def _top_p_sampling_from_probs_internal(
+    probs: torch.Tensor,
+    indices: Optional[torch.Tensor],
+    maybe_top_p_arr: Optional[torch.Tensor],
+    top_p_val: float,
+    deterministic: bool,
+    generator: Optional[torch.Generator],
+) -> torch.Tensor:
+    device = probs.device
+    probs = probs.float()
+    maybe_top_p_arr = maybe_top_p_arr.float() if maybe_top_p_arr is not None else None
+    samples = torch.empty(probs.size(0), dtype=torch.int32, device=device)
+    torch.ops.sgl_kernel.top_p_sampling_from_probs.default(
+        probs,
+        samples,
+        indices,
+        maybe_top_p_arr,
+        top_p_val,
+        deterministic,
+        generator,
+    )
+    return samples
+
+
+def top_p_sampling_from_probs(
+    probs: torch.Tensor,
+    top_p: Union[torch.Tensor, float],
+    indices: Optional[torch.Tensor] = None,
+    deterministic: bool = True,
+    generator: Optional[torch.Generator] = None,
+    check_nan: bool = False,
+) -> torch.Tensor:
+    if check_nan and torch.any(torch.isnan(probs)):
+        raise ValueError("Input probs contains NaN.")
+    return _top_p_sampling_from_probs_internal(
+        probs, indices, *_to_tensor_scalar_tuple(top_p), deterministic, generator
+    )
+
+
+def _top_k_top_p_sampling_from_probs_internal(
+    probs: torch.Tensor,
+    indices: Optional[torch.Tensor],
+    maybe_top_k_arr: Optional[torch.Tensor],
+    top_k_val: int,
+    maybe_top_p_arr: Optional[torch.Tensor],
+    top_p_val: float,
+    deterministic: bool,
+    generator: Optional[torch.Generator],
+) -> torch.Tensor:
+    device = probs.device
+    probs = probs.float()
+    maybe_top_k_arr = maybe_top_k_arr.int() if maybe_top_k_arr is not None else None
+    maybe_top_p_arr = maybe_top_p_arr.float() if maybe_top_p_arr is not None else None
+    samples = torch.empty(probs.size(0), dtype=torch.int32, device=device)
+    torch.ops.sgl_kernel.musa_top_k_top_p_sampling_from_probs.default(
+        probs,
+        samples,
+        indices,
+        maybe_top_k_arr,
+        top_k_val,
+        maybe_top_p_arr,
+        top_p_val,
+        deterministic,
+        generator,
+    )
+    return samples
+
+
+def top_k_top_p_sampling_from_probs(
+    probs: torch.Tensor,
+    top_k: Union[torch.Tensor, int],
+    top_p: Union[torch.Tensor, float],
+    indices: Optional[torch.Tensor] = None,
+    filter_apply_order: str = "top_k_first",
+    deterministic: bool = True,
+    generator: Optional[torch.Generator] = None,
+    check_nan: bool = False,
+) -> torch.Tensor:
+    if filter_apply_order == "top_k_first":
+        renorm_probs = top_k_renorm_probs(probs, top_k)
+        return top_p_sampling_from_probs(
+            renorm_probs,
+            top_p,
+            indices,
+            deterministic,
+            generator=generator,
+            check_nan=check_nan,
+        )
+    if filter_apply_order == "joint":
+        if check_nan and torch.any(torch.isnan(probs)):
+            raise ValueError("Input probs contains NaN.")
+        return _top_k_top_p_sampling_from_probs_internal(
+            probs,
+            indices,
+            *_to_tensor_scalar_tuple(top_k),
+            *_to_tensor_scalar_tuple(top_p),
+            deterministic,
+            generator,
+        )
+    raise ValueError(f"Invalid filter_apply_order: {filter_apply_order}")
+
+
+def _min_p_sampling_from_probs_internal(
+    probs: torch.Tensor,
+    indices: Optional[torch.Tensor],
+    maybe_min_p_arr: Optional[torch.Tensor],
+    min_p_val: float,
+    deterministic: bool,
+    generator: Optional[torch.Generator],
+) -> torch.Tensor:
+    device = probs.device
+    probs = probs.float()
+    maybe_min_p_arr = maybe_min_p_arr.float() if maybe_min_p_arr is not None else None
+    samples = torch.empty(probs.size(0), dtype=torch.int32, device=device)
+    torch.ops.sgl_kernel.min_p_sampling_from_probs.default(
+        probs,
+        samples,
+        indices,
+        maybe_min_p_arr,
+        min_p_val,
+        deterministic,
+        generator,
+    )
+    return samples
+
+
+def min_p_sampling_from_probs(
+    probs: torch.Tensor,
+    min_p: Union[torch.Tensor, float],
+    indices: Optional[torch.Tensor] = None,
+    deterministic: bool = True,
+    generator: Optional[torch.Generator] = None,
+    check_nan: bool = False,
+) -> torch.Tensor:
+    if check_nan and torch.any(torch.isnan(probs)):
+        raise ValueError("Input probs contains NaN.")
+    return _min_p_sampling_from_probs_internal(
+        probs, indices, *_to_tensor_scalar_tuple(min_p), deterministic, generator
+    )

sgl-kernel/python/sgl_kernel/quantization/__init__.py

@@ -0,0 +1,8 @@
+from .gguf import (
+    ggml_dequantize,
+    ggml_moe_a8,
+    ggml_moe_a8_vec,
+    ggml_moe_get_block_size,
+    ggml_mul_mat_a8,
+    ggml_mul_mat_vec_a8,
+)

sgl-kernel/python/sgl_kernel/quantization/gguf.py

@@ -0,0 +1,62 @@
+import torch
+
+
+def ggml_dequantize(
+    weight: torch.Tensor, quant_type: int, M: int, N: int, dtype: torch.dtype
+):
+    assert M > 0 and N > 0, "GGUF weight Input shape must be of positive dimensions"
+    return torch.ops.sgl_kernel.ggml_dequantize.default(weight, quant_type, M, N, dtype)
+
+
+def ggml_mul_mat_vec_a8(
+    weight: torch.Tensor, x: torch.Tensor, quant_type: int, row: int
+) -> torch.Tensor:
+    return torch.ops.sgl_kernel.ggml_mul_mat_vec_a8.default(weight, x, quant_type, row)
+
+
+def ggml_mul_mat_a8(
+    weight: torch.Tensor, x: torch.Tensor, quant_type: int, row: int
+) -> torch.Tensor:
+    return torch.ops.sgl_kernel.ggml_mul_mat_a8.default(weight, x, quant_type, row)
+
+
+def ggml_moe_a8(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    sorted_token_ids: torch.Tensor,
+    expert_ids: torch.Tensor,
+    num_token_post_padded: torch.Tensor,
+    type: int,
+    row: int,
+    topk: int,
+    tokens: int,
+) -> torch.Tensor:
+    return torch.ops.sgl_kernel.ggml_moe_a8.default(
+        input,
+        weight,
+        sorted_token_ids,
+        expert_ids,
+        num_token_post_padded,
+        type,
+        row,
+        topk,
+        tokens,
+    )
+
+
+def ggml_moe_a8_vec(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    topk_ids: torch.Tensor,
+    top_k: int,
+    type: int,
+    row: int,
+    tokens: int,
+) -> torch.Tensor:
+    return torch.ops.sgl_kernel.ggml_moe_a8_vec.default(
+        input, weight, topk_ids, top_k, type, row, tokens
+    )
+
+
+def ggml_moe_get_block_size(type: int) -> int:
+    return torch.ops.sgl_kernel.ggml_moe_get_block_size.default(type)

sgl-kernel/python/sgl_kernel/sampling.py

@@ -0,0 +1,115 @@
+from typing import Optional, Union
+
+import torch
+from sgl_kernel.utils import _to_tensor_scalar_tuple
+
+try:
+    import flashinfer.sampling as _flashinfer_sampling
+
+    _has_flashinfer = True
+except ImportError:
+    _has_flashinfer = False
+
+
+def _top_k_renorm_probs_internal(
+    probs: torch.Tensor,
+    maybe_top_k_arr: Optional[torch.Tensor],
+    top_k_val: int,
+) -> torch.Tensor:
+    probs = probs.float()
+    maybe_top_k_arr = maybe_top_k_arr.int() if maybe_top_k_arr is not None else None
+    renorm_probs = torch.empty_like(probs)
+    torch.ops.sgl_kernel.top_k_renorm_probs.default(
+        probs, renorm_probs, maybe_top_k_arr, top_k_val
+    )
+    return renorm_probs
+
+
+def top_k_renorm_probs(
+    probs: torch.Tensor,
+    top_k: Union[torch.Tensor, int],
+) -> torch.Tensor:
+    r"""Adapt from https://github.com/flashinfer-ai/flashinfer/flashinfer/sampling.py
+    Fused GPU kernel for renormalizing probabilities by top-k thresholding.
+
+    Parameters
+    ----------
+    probs: torch.Tensor
+        Probabilities, shape ``(batch_size, num_classes)``.
+    top_k: Union[torch.Tensor, int]
+        Either a scalar or a tensor of shape ``(batch_size,)``, representing the top-k threshold for for
+        for re-normalizing probabilities, should be in ``(0, num_classes)``.
+        If a scalar, the same threshold is used for all requests.
+        If a tensor, each request has its own threshold.
+        We keep the top-k probabilities, set the rest to zero, and renormalize the probabilities.
+
+    Returns
+    -------
+    renorm_probs: torch.Tensor
+        Renormalized probabilities, shape ``(batch_size, num_classes)``.
+
+    Note
+    ----
+    This combination of ``top_k_renorm_probs`` and ``sampling_from_probs`` should be equivalent to
+    ``top_k_sampling_from_probs``.
+    """
+    if probs.device.type == "musa" or not _has_flashinfer:
+        return _top_k_renorm_probs_internal(probs, *_to_tensor_scalar_tuple(top_k))
+    else:
+        return _flashinfer_sampling.top_k_renorm_probs(probs, top_k)
+
+
+top_k_renorm_prob = top_k_renorm_probs
+
+
+def _top_p_renorm_probs_internal(
+    probs: torch.Tensor,
+    maybe_top_p_arr: Optional[torch.Tensor],
+    top_p_val: float,
+) -> torch.Tensor:
+    probs = probs.float()
+    maybe_top_p_arr = maybe_top_p_arr.float() if maybe_top_p_arr is not None else None
+    renorm_probs = torch.empty_like(probs)
+    torch.ops.sgl_kernel.top_p_renorm_probs.default(
+        probs, renorm_probs, maybe_top_p_arr, top_p_val
+    )
+    return renorm_probs
+
+
+def top_p_renorm_probs(
+    probs: torch.Tensor,
+    top_p: Union[torch.Tensor, float],
+) -> torch.Tensor:
+    r"""Adapt from https://github.com/flashinfer-ai/flashinfer/flashinfer/sampling.py
+    Fused GPU kernel for renormalizing probabilities by top-p thresholding.
+
+    Parameters
+    ----------
+    probs: torch.Tensor
+        Probabilities, shape ``(batch_size, num_classes)``.
+    top_p: Union[torch.Tensor, float]
+        Either a scalar or a tensor of shape ``(batch_size,)``, representing the top-p threshold for for
+        re-normalizing probabilities, should be in ``(0, 1)``.
+        If a scalar, the same threshold is used for all requests.
+        If a tensor, each request has its own threshold.
+        We mask out the probabilities less than `threshold` where the cumulative sum
+        of ``probs[probs >= threshold]`` is `top_p`, and renormalize the probabilities.
+
+    Returns
+    -------
+    renorm_probs: torch.Tensor
+        Renormalized probabilities, shape ``(batch_size, num_classes)``.
+
+    Note
+    ----
+    This combination of ``top_p_renorm_probs`` and ``sampling_from_probs`` should be equivalent to
+    ``top_p_sampling_from_probs``.
+
+    """
+    if probs.device.type == "musa" or not _has_flashinfer:
+        return _top_p_renorm_probs_internal(probs, *_to_tensor_scalar_tuple(top_p))
+    else:
+        return _flashinfer_sampling.top_p_renorm_probs(probs, top_p)
+
+
+top_p_renorm_prob = top_p_renorm_probs

sgl-kernel/python/sgl_kernel/scalar_type.py

@@ -0,0 +1,352 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import functools
+import struct
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional, Union
+
+_SCALAR_TYPES_ID_MAP = {}
+
+
+# Mirrors enum in `core/scalar_type.hpp`
+class NanRepr(Enum):
+    NONE = 0  # nans are not supported
+    IEEE_754 = 1  # nans are: Exp all 1s, mantissa not all 0s
+    EXTD_RANGE_MAX_MIN = 2  # nans are: Exp all 1s, mantissa all 1s
+
+
+# This ScalarType class is a parallel implementation of the C++ ScalarType
+# class found in csrc/core/scalar_type.hpp.  These two classes should be kept
+# in sync until the inductor fully supports custom C++ classes.
+@dataclass(frozen=True)
+class ScalarType:
+    """
+    ScalarType can represent a wide range of floating point and integer
+    types, in particular it can be used to represent sub-byte data types
+    (something that torch.dtype currently does not support). It is also
+    capable of  representing types with a bias, i.e.:
+      `stored_value = value + bias`,
+    this is useful for quantized types (e.g. standard GPTQ 4bit uses a bias
+    of 8). The implementation for this class can be found in
+    csrc/core/scalar_type.hpp, these type signatures should be kept in sync
+    with that file.
+    """
+
+    exponent: int
+    """
+    Number of bits in the exponent if this is a floating point type
+    (zero if this an integer type)
+    """
+
+    mantissa: int
+    """
+    Number of bits in the mantissa if this is a floating point type,
+    or the number bits representing an integer excluding the sign bit if
+    this an integer type.
+    """
+
+    signed: bool
+    "If the type is signed (i.e. has a sign bit)"
+
+    bias: int
+    """
+    bias used to encode the values in this scalar type
+    (value = stored_value - bias, default 0) for example if we store the
+    type as an unsigned integer with a bias of 128 then the value 0 will be
+    stored as 128 and -1 will be stored as 127 and 1 will be stored as 129.
+    """
+
+    _finite_values_only: bool = False
+    """
+    Private: if infs are supported, used `has_infs()` instead.
+    """
+
+    nan_repr: NanRepr = NanRepr.IEEE_754
+    """
+    How NaNs are represent in this scalar type, returns NanRepr value.
+    (not applicable for integer types)
+    """
+
+    def _floating_point_max_int(self) -> int:
+        assert (
+            self.mantissa <= 52 and self.exponent <= 11
+        ), f"Cannot represent max/min as a double for type {self.__str__()}"
+
+        max_mantissa = (1 << self.mantissa) - 1
+        if self.nan_repr == NanRepr.EXTD_RANGE_MAX_MIN:
+            max_mantissa = max_mantissa - 1
+
+        max_exponent = (1 << self.exponent) - 2
+        if self.nan_repr == NanRepr.EXTD_RANGE_MAX_MIN or self.nan_repr == NanRepr.NONE:
+            assert (
+                self.exponent < 11
+            ), f"Cannot represent max/min as a double for type {self.__str__()}"
+            max_exponent = max_exponent + 1
+
+        # adjust the exponent to match that of a double
+        # for now we assume the exponent bias is the standard 2^(e-1) -1, (where
+        # e is the exponent bits), there is some precedent for non-standard
+        # biases, example `float8_e4m3b11fnuz` here:
+        # https://github.com/jax-ml/ml_dtypes but to avoid premature over
+        # complication we are just assuming the standard exponent bias until
+        # there is a need to support non-standard biases
+        exponent_bias = (1 << (self.exponent - 1)) - 1
+        exponent_bias_double = (1 << 10) - 1  # double e = 11
+
+        max_exponent_double = max_exponent - exponent_bias + exponent_bias_double
+
+        # shift the mantissa and exponent into the proper positions for an
+        # IEEE double and bitwise-or them together.
+        return (max_mantissa << (52 - self.mantissa)) | (max_exponent_double << 52)
+
+    def _floating_point_max(self) -> float:
+        double_raw = self._floating_point_max_int()
+        return struct.unpack("!d", struct.pack("!Q", double_raw))[0]
+
+    def _raw_max(self) -> Union[int, float]:
+        if self.is_floating_point():
+            return self._floating_point_max()
+        else:
+            assert (
+                self.size_bits < 64 or self.size_bits == 64 and self.is_signed()
+            ), "Cannot represent max as an int"
+            return (1 << self.mantissa) - 1
+
+    def _raw_min(self) -> Union[int, float]:
+        if self.is_floating_point():
+            assert (
+                self.is_signed()
+            ), "We currently assume all floating point types are signed"
+            sign_bit_double = 1 << 63
+
+            max_raw = self._floating_point_max_int()
+            min_raw = max_raw | sign_bit_double
+            return struct.unpack("!d", struct.pack("!Q", min_raw))[0]
+        else:
+            assert (
+                not self.is_signed() or self.size_bits <= 64
+            ), "Cannot represent min as a int64_t"
+
+            if self.is_signed():
+                return -(1 << (self.size_bits - 1))
+            else:
+                return 0
+
+    @functools.cached_property
+    def id(self) -> int:
+        """
+        Convert the ScalarType to an int which can be passed to pytorch custom
+        ops. This layout of the int must be kept in sync with the C++
+        ScalarType's from_id method.
+        """
+        val = 0
+        offset = 0
+
+        def or_and_advance(member, bit_width):
+            nonlocal val
+            nonlocal offset
+            bit_mask = (1 << bit_width) - 1
+            val = val | (int(member) & bit_mask) << offset
+            offset = offset + bit_width
+
+        or_and_advance(self.exponent, 8)
+        or_and_advance(self.mantissa, 8)
+        or_and_advance(self.signed, 1)
+        or_and_advance(self.bias, 32)
+        or_and_advance(self._finite_values_only, 1)
+        or_and_advance(self.nan_repr.value, 8)
+
+        assert offset <= 64, f"ScalarType fields too big {offset} to fit into an int64"
+
+        _SCALAR_TYPES_ID_MAP[val] = self
+
+        return val
+
+    @property
+    def size_bits(self) -> int:
+        return self.exponent + self.mantissa + int(self.signed)
+
+    def min(self) -> Union[int, float]:
+        """
+        Min representable value for this scalar type.
+        (accounting for bias if there is one)
+        """
+        return self._raw_min() - self.bias
+
+    def max(self) -> Union[int, float]:
+        """
+        Max representable value for this scalar type.
+        (accounting for bias if there is one)
+        """
+        return self._raw_max() - self.bias
+
+    def is_signed(self) -> bool:
+        """
+        If the type is signed (i.e. has a sign bit), same as `signed`
+        added for consistency with:
+        https://pytorch.org/docs/stable/generated/torch.Tensor.is_signed.html
+        """
+        return self.signed
+
+    def is_floating_point(self) -> bool:
+        "If the type is a floating point type"
+        return self.exponent != 0
+
+    def is_integer(self) -> bool:
+        "If the type is an integer type"
+        return self.exponent == 0
+
+    def has_bias(self) -> bool:
+        "If the type has a non-zero bias"
+        return self.bias != 0
+
+    def has_infs(self) -> bool:
+        "If the type is floating point and supports infinity"
+        return not self._finite_values_only
+
+    def has_nans(self) -> bool:
+        return self.nan_repr != NanRepr.NONE
+
+    def is_ieee_754(self) -> bool:
+        """
+        If the type is a floating point type that follows IEEE 754
+        conventions
+        """
+        return self.nan_repr == NanRepr.IEEE_754 and not self._finite_values_only
+
+    def __str__(self) -> str:
+        """
+        naming generally follows: https://github.com/jax-ml/ml_dtypes
+        for floating point types (leading f) the scheme is:
+        `float<size_bits>_e<exponent_bits>m<mantissa_bits>[flags]`
+        flags:
+          - no-flags: means it follows IEEE 754 conventions
+          - f: means finite values only (no infinities)
+          - n: means nans are supported (non-standard encoding)
+        for integer types the scheme is:
+          `[u]int<size_bits>[b<bias>]`
+          - if bias is not present it means its zero
+        """
+        if self.is_floating_point():
+            ret = (
+                "float"
+                + str(self.size_bits)
+                + "_e"
+                + str(self.exponent)
+                + "m"
+                + str(self.mantissa)
+            )
+
+            if not self.is_ieee_754():
+                if self._finite_values_only:
+                    ret = ret + "f"
+                if self.nan_repr != NanRepr.NONE:
+                    ret = ret + "n"
+
+            return ret
+        else:
+            ret = ("int" if self.is_signed() else "uint") + str(self.size_bits)
+            if self.has_bias():
+                ret = ret + "b" + str(self.bias)
+            return ret
+
+    def __repr__(self) -> str:
+        return "ScalarType." + self.__str__()
+
+    # __len__ needs to be defined (and has to throw TypeError) for pytorch's
+    # opcheck to work.
+    def __len__(self) -> int:
+        raise TypeError
+
+    #
+    # Convenience Constructors
+    #
+
+    @classmethod
+    def int_(cls, size_bits: int, bias: Optional[int]) -> "ScalarType":
+        "Create a signed integer scalar type (size_bits includes sign-bit)."
+        ret = cls(0, size_bits - 1, True, bias if bias else 0)
+        ret.id  # noqa B018: make sure the id is cached
+        return ret
+
+    @classmethod
+    def uint(cls, size_bits: int, bias: Optional[int]) -> "ScalarType":
+        """Create a unsigned integer scalar type."""
+        ret = cls(0, size_bits, False, bias if bias else 0)
+        ret.id  # noqa B018: make sure the id is cached
+        return ret
+
+    @classmethod
+    def float_IEEE754(cls, exponent: int, mantissa: int) -> "ScalarType":
+        """
+        Create a standard floating point type
+        (i.e. follows IEEE 754 conventions).
+        """
+        assert mantissa > 0 and exponent > 0
+        ret = cls(exponent, mantissa, True, 0)
+        ret.id  # noqa B018: make sure the id is cached
+        return ret
+
+    @classmethod
+    def float_(
+        cls, exponent: int, mantissa: int, finite_values_only: bool, nan_repr: NanRepr
+    ) -> "ScalarType":
+        """
+        Create a non-standard floating point type
+        (i.e. does not follow IEEE 754 conventions).
+        """
+        assert mantissa > 0 and exponent > 0
+        assert nan_repr != NanRepr.IEEE_754, (
+            "use `float_IEEE754` constructor for floating point types that "
+            "follow IEEE 754 conventions"
+        )
+        ret = cls(exponent, mantissa, True, 0, finite_values_only, nan_repr)
+        ret.id  # noqa B018: make sure the id is cached
+        return ret
+
+    @classmethod
+    def from_id(cls, scalar_type_id: int):
+        if scalar_type_id not in _SCALAR_TYPES_ID_MAP:
+            raise ValueError(f"scalar_type_id {scalar_type_id} doesn't exists.")
+        return _SCALAR_TYPES_ID_MAP[scalar_type_id]
+
+
+# naming generally follows: https://github.com/jax-ml/ml_dtypes
+# for floating point types (leading f) the scheme is:
+#  `float<size_bits>_e<exponent_bits>m<mantissa_bits>[flags]`
+#  flags:
+#  - no-flags: means it follows IEEE 754 conventions
+#  - f: means finite values only (no infinities)
+#  - n: means nans are supported (non-standard encoding)
+# for integer types the scheme is:
+#  `[u]int<size_bits>[b<bias>]`
+#  - if bias is not present it means its zero
+
+
+class scalar_types:
+    int4 = ScalarType.int_(4, None)
+    uint4 = ScalarType.uint(4, None)
+    int8 = ScalarType.int_(8, None)
+    uint8 = ScalarType.uint(8, None)
+    float8_e4m3fn = ScalarType.float_(4, 3, True, NanRepr.EXTD_RANGE_MAX_MIN)
+    float8_e5m2 = ScalarType.float_IEEE754(5, 2)
+    float16_e8m7 = ScalarType.float_IEEE754(8, 7)
+    float16_e5m10 = ScalarType.float_IEEE754(5, 10)
+
+    # fp6, https://github.com/usyd-fsalab/fp6_llm/tree/main
+    float6_e3m2f = ScalarType.float_(3, 2, True, NanRepr.NONE)
+
+    # fp4, https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf
+    float4_e2m1f = ScalarType.float_(2, 1, True, NanRepr.NONE)
+
+    # "gptq" types
+    uint2b2 = ScalarType.uint(2, 2)
+    uint3b4 = ScalarType.uint(3, 4)
+    uint4b8 = ScalarType.uint(4, 8)
+    uint8b128 = ScalarType.uint(8, 128)
+
+    # colloquial names
+    bfloat16 = float16_e8m7
+    float16 = float16_e5m10

sgl-kernel/python/sgl_kernel/sparse_flash_attn.py

@@ -0,0 +1,293 @@
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+
+def maybe_contiguous(x):
+    return x.contiguous() if x is not None and x.stride(-1) != 1 else x
+
+
+# Sparse attention utils
+def convert_vertical_slash_indexes(
+    q_seqlens: torch.Tensor,  # [BATCH, ]
+    kv_seqlens: torch.Tensor,  # [BATCH, ]
+    vertical_indexes: torch.Tensor,  # [BATCH, N_HEADS, NNZ_V]
+    slash_indexes: torch.Tensor,  # [BATCH, N_HEADS, NNZ_S]
+    context_size: int,
+    block_size_M: int,
+    block_size_N: int,
+    causal: bool = True,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    batch_size = slash_indexes.size(0)
+    num_heads = slash_indexes.size(1)
+    nnz_slash = slash_indexes.size(2)
+    nnz_vertical = vertical_indexes.size(2)
+    num_rows = (context_size + block_size_M - 1) // block_size_M
+
+    block_count = torch.zeros(
+        batch_size, num_heads, num_rows, dtype=q_seqlens.dtype, device=q_seqlens.device
+    )
+    block_offset = torch.zeros(
+        batch_size,
+        num_heads,
+        num_rows,
+        nnz_slash,
+        dtype=q_seqlens.dtype,
+        device=q_seqlens.device,
+    )
+    column_count = torch.zeros(
+        batch_size, num_heads, num_rows, dtype=q_seqlens.dtype, device=q_seqlens.device
+    )
+    column_index = torch.zeros(
+        batch_size,
+        num_heads,
+        num_rows,
+        nnz_vertical,
+        dtype=q_seqlens.dtype,
+        device=q_seqlens.device,
+    )
+
+    torch.ops.sgl_kernel.convert_vertical_slash_indexes.default(
+        block_count,
+        block_offset,
+        column_count,
+        column_index,
+        q_seqlens,
+        kv_seqlens,
+        vertical_indexes,
+        slash_indexes,
+        context_size,
+        block_size_M,
+        block_size_N,
+        causal,
+    )
+    return block_count, block_offset, column_count, column_index
+
+
+def convert_vertical_slash_indexes_mergehead(
+    q_seqlens: torch.Tensor,  # [BATCH, ]
+    kv_seqlens: torch.Tensor,  # [BATCH, ]
+    vertical_indexes: torch.Tensor,  # [BATCH, N_HEADS, NNZ_V]
+    slash_indexes: torch.Tensor,  # [BATCH, N_HEADS, NNZ_S]
+    # [N_HEADS] : different head use different number of indices
+    vertical_indices_count: torch.Tensor,
+    slash_indices_count: torch.Tensor,
+    context_size: int,
+    block_size_M: int,
+    block_size_N: int,
+    causal: bool = True,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    batch_size = slash_indexes.size(0)
+    num_heads = slash_indexes.size(1)
+    nnz_slash = slash_indexes.size(2)
+    nnz_vertical = vertical_indexes.size(2)
+    num_rows = (context_size + block_size_M - 1) // block_size_M
+
+    block_count = torch.empty(
+        batch_size, num_heads, num_rows, dtype=q_seqlens.dtype, device=q_seqlens.device
+    )
+    block_offset = torch.empty(
+        batch_size,
+        num_heads,
+        num_rows,
+        nnz_slash,
+        dtype=q_seqlens.dtype,
+        device=q_seqlens.device,
+    )
+    column_count = torch.empty(
+        batch_size, num_heads, num_rows, dtype=q_seqlens.dtype, device=q_seqlens.device
+    )
+    column_index = torch.empty(
+        batch_size,
+        num_heads,
+        num_rows,
+        nnz_vertical,
+        dtype=q_seqlens.dtype,
+        device=q_seqlens.device,
+    )
+
+    torch.ops.sgl_kernel.convert_vertical_slash_indexes_mergehead.default(
+        block_count,
+        block_offset,
+        column_count,
+        column_index,
+        q_seqlens,
+        kv_seqlens,
+        vertical_indexes,
+        slash_indexes,
+        vertical_indices_count,
+        slash_indices_count,
+        context_size,
+        block_size_M,
+        block_size_N,
+        causal,
+    )
+    return block_count, block_offset, column_count, column_index
+
+
+def sparse_attn_func(
+    q,
+    k,
+    v,
+    block_count,
+    block_offset,
+    column_count,
+    column_index,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+    *,
+    return_softmax_lse=False,
+    out=None,
+):
+    """Compute attention with vertical and slash sparsity patterns.
+    Most Arguments are the same with the flash_attn_func interface, except for 4 extra args:
+    block_count and block_offset for slash sparsity patterns, and
+    column_count and column_index for vertical sparsity patterns.
+    For more details please refer to Appendix C.4.2 of paper https://arxiv.org/abs/2407.02490.
+
+    Arguments:
+        q: (batch_size, seqlen, nheads, headdim)
+        k: (batch_size, seqlen, nheads_k, headdim)
+        v: (batch_size, seqlen, nheads_k, headdim)
+        block_count: (batch_size, nheads, cdiv(seqlen, BLOCK_M))
+        block_offset: (batch_size, nheads, cdiv(seqlen, BLOCK_M), NNZ_S)
+        column_count: (batch_size, nheads, cdiv(seqlen, BLOCK_M))
+        column_index: (batch_size, nheads, cdiv(seqlen, BLOCK_M), NNZ_V)
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (batch_size, seqlen, nheads, headdim).
+        softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+    """
+    if softmax_scale is None:
+        softmax_scale = q.shape[-1] ** (-0.5)
+
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    out, softmax_lse = torch.ops.sgl_kernel.fwd_sparse.default(
+        q,
+        k,
+        v,
+        block_count,
+        block_offset,
+        column_count,
+        column_index,
+        out,
+        alibi_slopes,
+        dropout_p,
+        softmax_scale,
+        causal,
+        softcap,
+        return_attn_probs and dropout_p > 0,
+        None,
+    )
+    return (out, softmax_lse) if return_softmax_lse else out
+
+
+def sparse_attn_varlen_func(
+    q,
+    k,
+    v,
+    block_count,
+    block_offset,
+    column_count,
+    column_index,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q,
+    max_seqlen_k,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+    *,
+    return_softmax_lse=False,
+    out=None,
+):
+    """Compute attention with vertical and slash sparsity patterns.
+    Most Arguments are the same with the flash_attn_varlen_func interface, except for 4 extra args:
+    block_count and block_offset for slash sparsity patterns, and
+    column_count and column_index for vertical sparsity patterns.
+    For more details please refer to Appendix C.4.2 of paper https://arxiv.org/abs/2407.02490.
+
+    Arguments:
+        q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.
+        k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.
+        v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.
+        block_count: (batch_size, nheads, cdiv(seqlen, BLOCK_M))
+        block_offset: (batch_size, nheads, cdiv(seqlen, BLOCK_M), NNZ_S)
+        column_count: (batch_size, nheads, cdiv(seqlen, BLOCK_M))
+        column_index: (batch_size, nheads, cdiv(seqlen, BLOCK_M), NNZ_V)
+        cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into q.
+        cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into kv.
+        max_seqlen_q: int. Maximum query sequence length in the batch.
+        max_seqlen_k: int. Maximum key sequence length in the batch.
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        softcap: float. Anything > 0 activates softcapping attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (total, nheads, headdim).
+        softmax_lse [optional, if return_softmax_lse=True]: (nheads, total_q_seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+    """
+    if softmax_scale is None:
+        softmax_scale = q.shape[-1] ** (-0.5)
+
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    out, softmax_lse = torch.ops.sgl_kernel.varlen_fwd_sparse.default(
+        q,
+        k,
+        v,
+        block_count,
+        block_offset,
+        column_count,
+        column_index,
+        out,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        None,
+        alibi_slopes,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        False,
+        causal,
+        softcap,
+        return_attn_probs and dropout_p > 0,
+        None,
+    )
+    return (out, softmax_lse) if return_softmax_lse else out

sgl-kernel/python/sgl_kernel/spatial.py

@@ -0,0 +1,63 @@
+import torch
+from torch.cuda.streams import ExternalStream
+
+try:
+    from . import spatial_ops  # triggers TORCH extension registration
+except Exception as _e:
+    _spatial_import_error = _e
+else:
+    _spatial_import_error = None
+
+_IMPORT_ERROR = ImportError(
+    "Failed to load sgl_kernel.spatial_ops extension. Ensure CUDA Driver >= 12.4"
+)
+
+
+def create_greenctx_stream_by_value(
+    SM_a: int, SM_b: int, device_id: int = None
+) -> tuple[ExternalStream, ExternalStream]:
+    """
+    Create two streams for greenctx.
+    Args:
+        sm_A (int): The SM of stream A.
+        sm_B (int): The weight of stream B.
+        device_id (int): The device id.
+    Returns:
+        tuple[ExternalStream, ExternalStream]: The two streams.
+    """
+    if _spatial_import_error is not None:
+        raise _IMPORT_ERROR from _spatial_import_error
+    if device_id is None:
+        device_id = torch.cuda.current_device()
+
+    res = torch.ops.sgl_kernel.create_greenctx_stream_by_value(SM_a, SM_b, device_id)
+
+    stream_a = ExternalStream(
+        stream_ptr=res[0], device=torch.device(f"cuda:{device_id}")
+    )
+    stream_b = ExternalStream(
+        stream_ptr=res[1], device=torch.device(f"cuda:{device_id}")
+    )
+
+    return stream_a, stream_b
+
+
+def get_sm_available(device_id: int = None) -> int:
+    """
+    Get the SMs available on the device.
+    Args:
+        device_id (int): The device id.
+    Returns:
+        int: The SMs available.
+    """
+    if _spatial_import_error is not None:
+        raise _IMPORT_ERROR from _spatial_import_error
+    if device_id is None:
+        device_id = torch.cuda.current_device()
+
+    device_props = torch.cuda.get_device_properties(device_id)
+
+    # Get the number of Streaming Multiprocessors (SMs)
+    sm_count = device_props.multi_processor_count
+
+    return sm_count

sgl-kernel/python/sgl_kernel/speculative.py

@@ -0,0 +1,308 @@
+from typing import Optional
+
+import torch
+
+
+def tree_speculative_sampling_target_only(
+    predicts: torch.Tensor,  # mutable
+    accept_index: torch.Tensor,  # mutable
+    accept_token_num: torch.Tensor,  # mutable
+    candidates: torch.Tensor,
+    retrive_index: torch.Tensor,
+    retrive_next_token: torch.Tensor,
+    retrive_next_sibling: torch.Tensor,
+    uniform_samples: torch.Tensor,
+    uniform_samples_for_final_sampling: torch.Tensor,
+    target_probs: torch.Tensor,
+    draft_probs: torch.Tensor,
+    threshold_single: float = 1.0,
+    threshold_acc: float = 1.0,
+    deterministic: bool = True,
+) -> None:
+    torch.ops.sgl_kernel.tree_speculative_sampling_target_only.default(
+        predicts,
+        accept_index,
+        accept_token_num,
+        candidates,
+        retrive_index,
+        retrive_next_token,
+        retrive_next_sibling,
+        uniform_samples,
+        uniform_samples_for_final_sampling,
+        target_probs,
+        draft_probs,
+        threshold_single,
+        threshold_acc,
+        deterministic,
+    )
+
+
+def verify_tree_greedy(
+    predicts: torch.Tensor,  # mutable
+    accept_index: torch.Tensor,  # mutable
+    accept_token_num: torch.Tensor,  # mutable
+    candidates: torch.Tensor,
+    retrive_index: torch.Tensor,
+    retrive_next_token: torch.Tensor,
+    retrive_next_sibling: torch.Tensor,
+    target_predict: torch.Tensor,
+) -> None:
+    torch.ops.sgl_kernel.verify_tree_greedy.default(
+        predicts,
+        accept_index,
+        accept_token_num,
+        candidates,
+        retrive_index,
+        retrive_next_token,
+        retrive_next_sibling,
+        target_predict,
+    )
+
+
+def build_tree_kernel_efficient(
+    parent_list: torch.Tensor,
+    selected_index: torch.Tensor,
+    verified_seq_len: torch.Tensor,
+    tree_mask: torch.Tensor,
+    positions: torch.Tensor,
+    retrive_index: torch.Tensor,
+    retrive_next_token: torch.Tensor,
+    retrive_next_sibling: torch.Tensor,
+    topk: int,
+    depth: int,
+    draft_token_num: int,
+    tree_mask_mode: int,
+) -> None:
+    torch.ops.sgl_kernel.build_tree_kernel_efficient.default(
+        parent_list,
+        selected_index,
+        verified_seq_len,
+        tree_mask,
+        positions,
+        retrive_index,
+        retrive_next_token,
+        retrive_next_sibling,
+        topk,
+        depth,
+        draft_token_num,
+        tree_mask_mode,
+    )
+
+
+def reconstruct_indices_from_tree_mask(
+    tree_mask: torch.Tensor,
+    verified_seq_len: torch.Tensor,
+    positions: torch.Tensor,
+    retrive_index: torch.Tensor,
+    retrive_next_token: torch.Tensor,
+    retrive_next_sibling: torch.Tensor,
+    batch_size: int,
+    draft_token_num: int,
+) -> None:
+    if tree_mask.is_cpu:
+        torch.ops.sgl_kernel.reconstruct_indices_from_tree_mask_cpu(
+            tree_mask,
+            verified_seq_len,
+            positions,
+            retrive_index,
+            retrive_next_token,
+            retrive_next_sibling,
+            batch_size,
+            draft_token_num,
+        )
+    else:
+        torch.ops.sgl_kernel.reconstruct_indices_from_tree_mask.default(
+            tree_mask,
+            verified_seq_len,
+            positions,
+            retrive_index,
+            retrive_next_token,
+            retrive_next_sibling,
+            batch_size,
+            draft_token_num,
+        )
+
+
+def segment_packbits(
+    x: torch.Tensor,
+    input_indptr: torch.Tensor,
+    output_indptr: torch.Tensor,
+    y: torch.Tensor,
+    batch_size: int,
+) -> None:
+    torch.ops.sgl_kernel.segment_packbits.default(
+        x,
+        input_indptr,
+        output_indptr,
+        y,
+        batch_size,
+        torch.cuda.current_stream().cuda_stream,
+    )
+
+
+def verify_tree_greedy_cpu(
+    predicts: torch.Tensor,  # mutable
+    accept_index: torch.Tensor,  # mutable
+    accept_token_num: torch.Tensor,  # mutable
+    candidates: torch.Tensor,
+    retrive_index: torch.Tensor,
+    retrive_next_token: torch.Tensor,
+    retrive_next_sibling: torch.Tensor,
+    target_predict: torch.Tensor,
+) -> None:
+    torch.ops.sgl_kernel.verify_tree_greedy_cpu(
+        predicts,
+        accept_index,
+        accept_token_num,
+        candidates,
+        retrive_index,
+        retrive_next_token,
+        retrive_next_sibling,
+        target_predict,
+    )
+
+
+def build_tree_kernel_efficient_cpu(
+    parent_list: torch.Tensor,
+    selected_index: torch.Tensor,
+    verified_seq_len: torch.Tensor,
+    tree_mask: torch.Tensor,
+    positions: torch.Tensor,
+    retrive_index: torch.Tensor,
+    retrive_next_token: torch.Tensor,
+    retrive_next_sibling: torch.Tensor,
+    topk: int,
+    depth: int,
+    draft_token_num: int,
+    tree_mask_mode: int,
+) -> None:
+    torch.ops.sgl_kernel.build_tree_kernel_efficient_cpu(
+        parent_list,
+        selected_index,
+        verified_seq_len,
+        tree_mask,
+        positions,
+        retrive_index,
+        retrive_next_token,
+        retrive_next_sibling,
+        topk,
+        depth,
+        draft_token_num,
+        tree_mask_mode,
+    )
+
+
+def assign_req_to_token_pool_cpu(
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    start_offset: torch.Tensor,
+    end_offset: torch.Tensor,
+    out_cache_loc: torch.Tensor,
+    pool_len: int,
+) -> None:
+    torch.ops.sgl_kernel.assign_req_to_token_pool_cpu(
+        req_pool_indices,
+        req_to_token,
+        start_offset,
+        end_offset,
+        out_cache_loc,
+        pool_len,
+    )
+
+
+def build_draft_decode_metadata_cpu(
+    req_to_token: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    seq_lens: torch.Tensor,
+    topk: int,
+    num_steps: int,
+    pool_len: int,
+) -> torch.Tensor:
+    return torch.ops.sgl_kernel.build_draft_decode_metadata_cpu(
+        req_to_token,
+        req_pool_indices,
+        seq_lens,
+        topk,
+        num_steps,
+        pool_len,
+    )
+
+
+def fill_bonus_tokens_cpu(
+    accept_tokens: torch.Tensor,
+    accept_lens: torch.Tensor,
+    bonus_tokens: torch.Tensor,
+    accept_stride: int,
+) -> None:
+    torch.ops.sgl_kernel.fill_bonus_tokens_cpu(
+        accept_tokens,
+        accept_lens,
+        bonus_tokens,
+        accept_stride,
+    )
+
+
+def fill_accept_out_cache_loc_cpu(
+    accept_index: torch.Tensor,
+    out_cache_loc: torch.Tensor,
+    accept_out_cache_loc: torch.Tensor,  # mutable
+) -> None:
+    torch.ops.sgl_kernel.fill_accept_out_cache_loc_cpu(
+        accept_index,
+        out_cache_loc,
+        accept_out_cache_loc,
+    )
+
+
+def assign_draft_cache_locs_contiguous_cpu(
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    seq_lens: torch.Tensor,
+    out_cache_loc: torch.Tensor,
+    pool_len: int,
+    topk: int,
+    num_steps: int,
+) -> None:
+    torch.ops.sgl_kernel.assign_draft_cache_locs_contiguous_cpu(
+        req_pool_indices,
+        req_to_token,
+        seq_lens,
+        out_cache_loc,
+        pool_len,
+        topk,
+        num_steps,
+    )
+
+
+def assign_extend_cache_locs_cpu(
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    start_offset: torch.Tensor,
+    end_offset: torch.Tensor,
+    out_cache_loc: torch.Tensor,
+    pool_len: int,
+) -> None:
+    torch.ops.sgl_kernel.assign_extend_cache_locs_cpu(
+        req_pool_indices,
+        req_to_token,
+        start_offset,
+        end_offset,
+        out_cache_loc,
+        pool_len,
+    )
+
+
+def rotate_input_ids_cpu(
+    input_ids: torch.Tensor,
+    extend_start_loc: torch.Tensor,
+    extend_seq_lens: torch.Tensor,
+    topk_index: torch.Tensor,
+    select_index: Optional[torch.Tensor] = None,
+) -> None:
+    torch.ops.sgl_kernel.rotate_input_ids_cpu(
+        input_ids,
+        extend_start_loc,
+        extend_seq_lens,
+        topk_index,
+        select_index,
+    )

sgl-kernel/python/sgl_kernel/test_utils.py

@@ -0,0 +1,125 @@
+import torch
+
+
+def create_per_token_group_quant_test_data(num_tokens, hidden_dim, num_ranks, flags):
+    device = torch.device("cuda")
+    dtype = torch.bfloat16
+
+    seed = num_tokens * 10000 + hidden_dim
+    gen_cpu = torch.Generator(device="cpu")
+    gen_cpu.manual_seed(seed)
+    gen_cuda = torch.Generator(device="cuda")
+    gen_cuda.manual_seed(seed)
+
+    if flags["fuse_silu_and_mul"]:
+        effective_hidden_dim = hidden_dim * 2
+    else:
+        effective_hidden_dim = hidden_dim
+    del hidden_dim
+
+    if (masked_layout_mode := flags["masked_layout_mode"]) is not None:
+        num_max_dispatch_tokens_per_rank = 768
+        num_global_experts = 288
+        num_local_experts, remainder = divmod(num_global_experts, num_ranks)
+        assert remainder == 0
+
+        # mimic DeepEP low_latency_dispatch output
+        x = torch.randn(
+            num_local_experts,
+            num_max_dispatch_tokens_per_rank * num_ranks,
+            effective_hidden_dim,
+            device=device,
+            dtype=dtype,
+            generator=gen_cuda,
+        )
+
+        if masked_layout_mode == "balanced":
+            masked_m = _compute_balanced_split(num_tokens, num_local_experts)
+        elif masked_layout_mode == "imbalanced":
+            masked_m = _compute_imbalanced_split(
+                num_tokens, num_local_experts, gen_cpu=gen_cpu
+            )
+        elif masked_layout_mode == "extreme":
+            masked_m = torch.tensor(
+                [num_tokens] + [0] * (num_local_experts - 1), dtype=torch.int
+            )
+        else:
+            raise NotImplementedError
+        print(f"{masked_layout_mode=} {masked_m=} {x.shape=}")
+
+        masked_m = masked_m.to(device)
+
+        return x, masked_m
+    else:
+        x = torch.randn(
+            num_tokens,
+            effective_hidden_dim,
+            device=device,
+            dtype=dtype,
+            generator=gen_cuda,
+        )
+        x[torch.randn(x.shape, device=device, generator=gen_cuda) < 0.001] *= 10
+        return x, None
+
+
+def _compute_balanced_split(total: int, arr_len: int):
+    base = total // arr_len
+    remainder = total % arr_len
+    ans = [base + 1 if i < remainder else base for i in range(arr_len)]
+    assert sum(ans) == total
+    return torch.tensor(ans, dtype=torch.int)
+
+
+def _compute_imbalanced_split(
+    total: int, arr_len: int, gen_cpu, dtype=torch.int
+) -> list[int]:
+    # can use `rand ** 2`, `rand ** 3`, etc, to change how imbalanced it is
+    noise_raw = torch.rand(arr_len, generator=gen_cpu) ** 3
+
+    noise = noise_raw / noise_raw.sum()
+    ans = (noise * total).round().to(dtype)
+
+    diff = total - ans.sum().item()
+    while diff != 0:
+        idx = torch.randint(0, arr_len, (1,), generator=gen_cpu).item()
+        if diff > 0:
+            ans[idx] += 1
+            diff -= 1
+        elif diff < 0 and ans[idx] > 0:
+            ans[idx] -= 1
+            diff += 1
+
+    assert sum(ans) == total
+    return ans
+
+
+def assert_all_close_or_tiny_diff(a: torch.Tensor, b: torch.Tensor):
+    assert (a.shape == b.shape) and (
+        a.dtype == b.dtype
+    ), f"{a.shape=} {b.shape=} {a.dtype=} {b.dtype=}"
+    numel = a.numel()
+
+    if a.dtype == torch.float8_e4m3fn:
+        a_u8 = a.view(torch.uint8)
+        b_u8 = b.view(torch.uint8)
+        diff_u8 = (a_u8.to(torch.int16) - b_u8.to(torch.int16)).abs()
+
+        count_diff_sign = ((a_u8 >= 0) & (b_u8 < 0)).sum().item()
+        count_tiny_diff = (diff_u8 == 1).sum().item()
+        count_large_diff = (diff_u8 >= 2).sum().item()
+    elif a.dtype == torch.int8:
+        diff = (a.to(torch.int16) - a.to(torch.int16)).abs()
+        count_diff_sign = ((a >= 0) & (b < 0)).sum().item()
+        count_tiny_diff = (diff == 1).sum().item()
+        count_large_diff = (diff >= 2).sum().item()
+    else:
+        raise NotImplementedError
+
+    assert (
+        (count_diff_sign == 0)
+        and (count_large_diff == 0)
+        and (
+            (count_tiny_diff / numel < 0.005)
+            or ((count_tiny_diff / numel < 0.04) and (numel <= 4096))
+        )
+    ), f"{count_diff_sign=} {count_tiny_diff=} {count_large_diff=} {numel=} {a=} {b=}"

sgl-kernel/python/sgl_kernel/testing/rotary_embedding.py

@@ -0,0 +1,271 @@
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import torch
+
+from sglang.jit_kernel.rope import FusedSetKVBufferArg as _JitFusedSetKVBufferArg
+from sglang.jit_kernel.rope import (
+    apply_rope_with_cos_sin_cache_inplace as _jit_apply_rope_with_cos_sin_cache_inplace,
+)
+
+
+@dataclass
+class FusedSetKVBufferArg:
+    value: torch.Tensor
+    k_buffer: torch.Tensor
+    v_buffer: torch.Tensor
+    cache_loc: torch.Tensor
+    # Kept for backward compatibility with old sgl_kernel test/bench callsites.
+    k_scale: Optional[float] = None
+    v_scale: Optional[float] = None
+
+    def to_jit(self) -> _JitFusedSetKVBufferArg:
+        return _JitFusedSetKVBufferArg(
+            value=self.value,
+            k_buffer=self.k_buffer,
+            v_buffer=self.v_buffer,
+            cache_loc=self.cache_loc,
+        )
+
+
+# vLLM torch native
+def _apply_rotary_emb(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    is_neox_style: bool,
+) -> torch.Tensor:
+    """
+    Args:
+        x: [num_tokens, num_heads, head_size]
+        cos: [num_tokens, head_size // 2]
+        sin: [num_tokens, head_size // 2]
+        is_neox_style: Whether to use the Neox-style or GPT-J-style rotary
+            positional embeddings.
+    """
+    cos = cos.unsqueeze(-2).to(x.dtype)
+    sin = sin.unsqueeze(-2).to(x.dtype)
+    if is_neox_style:
+        x1, x2 = torch.chunk(x, 2, dim=-1)
+    else:
+        x1 = x[..., ::2]
+        x2 = x[..., 1::2]
+    o1 = x1 * cos - x2 * sin
+    o2 = x2 * cos + x1 * sin
+    if is_neox_style:
+        return torch.cat((o1, o2), dim=-1)
+    else:
+        return torch.stack((o1, o2), dim=-1).flatten(-2)
+
+
+class RotaryEmbedding(torch.nn.Module):
+    # Reference: https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/rotary_embedding.py
+    def __init__(
+        self,
+        head_size: int,
+        rotary_dim: int,
+        max_position_embeddings: int,
+        base: int,
+        is_neox_style: bool,
+        dtype: torch.dtype,
+    ) -> None:
+        super().__init__()
+        self.head_size = head_size
+        self.rotary_dim = rotary_dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+        self.is_neox_style = is_neox_style
+        self.dtype = dtype
+
+        cache = self._compute_cos_sin_cache()
+        self.cos_sin_cache: torch.Tensor
+        self.register_buffer("cos_sin_cache", cache, persistent=False)
+
+    def _compute_inv_freq(self, base: Union[int, float]) -> torch.Tensor:
+        inv_freq = 1.0 / (
+            base
+            ** (
+                torch.arange(0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim
+            )
+        )
+        return inv_freq
+
+    def _compute_cos_sin_cache(self) -> torch.Tensor:
+        """Compute the cos and sin cache."""
+        inv_freq = self._compute_inv_freq(self.base)
+        t = torch.arange(self.max_position_embeddings, dtype=torch.float)
+
+        freqs = torch.einsum("i,j -> ij", t, inv_freq)
+        cos = freqs.cos()
+        sin = freqs.sin()
+        cache = torch.cat((cos, sin), dim=-1)
+        return cache
+
+    def forward_native(
+        self,
+        positions: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        offsets: Optional[torch.Tensor] = None,
+        fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """A PyTorch-native implementation of forward()."""
+        assert (
+            fused_set_kv_buffer_arg is None
+        ), "fused_set_kv_buffer_arg is not supported for native implementation"
+
+        if offsets is not None:
+            positions = positions + offsets
+
+        positions = positions.flatten()
+        num_tokens = positions.shape[0]
+        cos_sin = self.cos_sin_cache.index_select(0, positions)
+
+        cos, sin = cos_sin.chunk(2, dim=-1)
+
+        query_shape = query.shape
+        query = query.view(num_tokens, -1, self.head_size)
+        query_rot = query[..., : self.rotary_dim]
+        query_pass = query[..., self.rotary_dim :]
+        query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style)
+        query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
+
+        key_shape = key.shape
+        key = key.view(num_tokens, -1, self.head_size)
+        key_rot = key[..., : self.rotary_dim]
+        key_pass = key[..., self.rotary_dim :]
+        key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style)
+        key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
+
+        # Modification: convert to the correct dtype
+        query = query.to(self.dtype)
+        key = key.to(self.dtype)
+        return query, key
+
+
+class FlashInferRotaryEmbedding(RotaryEmbedding):
+    def forward_cuda(
+        self,
+        positions: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        offsets: Optional[torch.Tensor] = None,
+        fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+
+        query_view = query.view(query.shape[0], -1, self.head_size)
+        key_view = key.view(key.shape[0], -1, self.head_size)
+        _jit_apply_rope_with_cos_sin_cache_inplace(
+            q=query_view,
+            k=key_view,
+            cos_sin_cache=self.cos_sin_cache,
+            positions=positions,
+            is_neox=self.is_neox_style,
+            fused_args=(
+                fused_set_kv_buffer_arg.to_jit()
+                if fused_set_kv_buffer_arg is not None
+                else None
+            ),
+        )
+
+        return query, key
+
+
+class SglKernelRotaryEmbedding(RotaryEmbedding):
+    def forward_cuda(
+        self,
+        positions: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        offsets: Optional[torch.Tensor] = None,
+        fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        assert (
+            fused_set_kv_buffer_arg is None
+        ), "fused_set_kv_buffer_arg is not supported for sgl-kernel implementation"
+        if self.cos_sin_cache.dtype != query.dtype:
+            self.cos_sin_cache = self.cos_sin_cache.to(query.dtype)
+        torch.ops.sgl_kernel.rotary_embedding(
+            positions,
+            query,
+            key,
+            self.head_size,
+            self.cos_sin_cache,
+            self.is_neox_style,
+        )
+        return query, key
+
+
+class MHATokenToKVPool:
+    KV_POOL_SIZE = 16384
+
+    def __init__(
+        self,
+        head_num: int,
+        head_dim: int,
+    ):
+        self.head_num = head_num
+        self.head_dim = head_dim
+        self.size = MHATokenToKVPool.KV_POOL_SIZE
+        self.page_size = 1
+        self.store_dtype = torch.bfloat16
+        self.device = "cuda"
+        self.layer_num = 1
+        self.start_layer = 0
+        self._create_buffers()
+
+    def _create_buffers(self):
+        self.k_buffer = [
+            torch.zeros(
+                (self.size + self.page_size, self.head_num, self.head_dim),
+                dtype=self.store_dtype,
+                device=self.device,
+            )
+            for _ in range(self.layer_num)
+        ]
+        self.v_buffer = [
+            torch.zeros(
+                (self.size + self.page_size, self.head_num, self.head_dim),
+                dtype=self.store_dtype,
+                device=self.device,
+            )
+            for _ in range(self.layer_num)
+        ]
+
+    def set_kv_buffer(
+        self,
+        loc: torch.Tensor,
+        cache_k: torch.Tensor,
+        cache_v: torch.Tensor,
+    ):
+        layer_id = 0
+        self.k_buffer[layer_id - self.start_layer][loc] = cache_k
+        self.v_buffer[layer_id - self.start_layer][loc] = cache_v
+
+
+def create_inputs(
+    head_size: int,
+    batch_size: int,
+    seq_len: int,
+    device,
+    dtype: torch.dtype,
+    num_q_heads: int,
+    num_kv_heads: int,
+):
+    pos_ids = torch.arange(seq_len, device=device).repeat(batch_size)
+    query = torch.randn(
+        batch_size * seq_len, num_q_heads * head_size, dtype=dtype, device=device
+    )
+    key = torch.randn(
+        batch_size * seq_len, num_kv_heads * head_size, dtype=dtype, device=device
+    )
+    value = torch.randn(
+        batch_size * seq_len, num_kv_heads * head_size, dtype=dtype, device=device
+    )
+    out_cache_loc = torch.randperm(
+        MHATokenToKVPool.KV_POOL_SIZE, dtype=torch.int64, device=device
+    )[: batch_size * seq_len].clone()
+
+    return dict(
+        pos_ids=pos_ids, query=query, key=key, value=value, out_cache_loc=out_cache_loc
+    )

sgl-kernel/python/sgl_kernel/top_k.py

@@ -0,0 +1,149 @@
+from typing import Optional
+
+import torch
+
+
+def fast_topk(values, topk, dim):
+    if topk == 1:
+        # Use max along the specified dimension to get both value and index
+        return torch.max(values, dim=dim, keepdim=True)
+    else:
+        # Use topk for efficiency with larger k values
+        # TODO: implement faster cuda kernels for large vocab sizes
+        return torch.topk(values, topk, dim=dim)
+
+
+def fast_topk_v2(
+    score: torch.Tensor,
+    lengths: torch.Tensor,
+    topk: int,
+    row_starts: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """
+    Get the topk indices of the score tensor.
+    Args:
+        score: The score tensor of shape (B, L). The score tensor is the logits
+            between the query and the key whose layout is either ragged or paged.
+            row_starts is only required when the key is ragged.
+        lengths: The lengths tensor of shape (B)
+        topk: The number of topk indices to get
+        row_starts: The start index of each row in the score tensor of shape (B).
+            For each row i, topk only applies to section [row_starts[i], row_starts[i] + lengths[i]]
+            of the score tensor.
+    Returns:
+        The topk indices tensor of shape (B, topk)
+    """
+    assert (
+        topk == 2048
+    ), "fast_topk_v2 is only optimized for deepseek v3.2 model, where topk=2048"
+    assert score.dim() == 2
+    topk_indices = score.new_empty((score.size(0), topk), dtype=torch.int32)
+    torch.ops.sgl_kernel.fast_topk(score, topk_indices, lengths, row_starts)
+    return topk_indices
+
+
+def fast_topk_transform_fused(
+    score: torch.Tensor,
+    lengths: torch.Tensor,
+    page_table_size_1: torch.Tensor,  # NOTE: page size should be 1
+    cu_seqlens_q: torch.Tensor,
+    topk: int,
+    row_starts: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """
+    Get the topk indices of the score tensor and then transform the topk indices
+    to indices to the page table (page_size = 1)
+    Args:
+        score: The score tensor of shape (B, L). The score tensor is the logits
+            between the query and the key whose layout is either ragged or paged.
+            row_starts is only required when the key is ragged.
+        lengths: The lengths tensor of shape (B)
+        page_table_size_1: The page table tensor of shape (Batch, topk)
+        cu_seqlens_q: The cumulative sequence lengths tensor of shape (Batch + 1)
+        topk: The number of topk indices to get
+        row_starts: The start index of each row in the score tensor of shape (B).
+            For each row i, topk only applies to section [row_starts[i], row_starts[i] + lengths[i]]
+            of the score tensor. It's only used for cases where the key is
+            ragged, i.e. during extend and draft extend.
+    Returns:
+        The topk indices tensor of shape (B, topk)
+    """
+    assert (
+        topk == 2048
+    ), "fast_topk_transform_fused is only optimized for deepseek v3.2 model, where topk=2048"
+    assert score.dim() == 2
+    src_page_table = page_table_size_1
+    dst_page_table = score.new_empty((score.shape[0], topk), dtype=torch.int32)
+    torch.ops.sgl_kernel.fast_topk_transform_fused(
+        score, lengths, dst_page_table, src_page_table, cu_seqlens_q, row_starts
+    )
+    return dst_page_table
+
+
+def deepseek_v4_topk_transform_512(
+    scores: torch.Tensor,
+    seq_lens: torch.Tensor,
+    page_table: torch.Tensor,
+    page_indices: torch.Tensor,
+    page_size: int,
+    raw_indices: Optional[torch.Tensor] = None,
+) -> None:
+    """
+    Performs the DeepSeek-V4 indexer top-k selection and writes the paged
+    physical slot indices into ``page_indices``. Supports topk up to 1024.
+    Optionally also writes the row-relative raw token positions into
+    ``raw_indices`` for hisparse capture.
+
+    Args:
+        scores: float32 ``[B, max_seq_len]`` indexer logits, contiguous on dim 1.
+        seq_lens: int32 ``[B]``, true KV length per batch row.
+        page_table: int32 ``[B, num_pages]``, logical->physical page table,
+            contiguous on dim 1.
+        page_indices: int32 ``[B, topk]``, output buffer, contiguous. Filled
+            with paged physical slots; -1 for padding entries.
+        page_size: power-of-2 page size.
+        raw_indices: optional int32 ``[B, topk]``, contiguous. If provided,
+            filled with raw token positions within each row.
+    """
+    if raw_indices is not None:
+        assert raw_indices.dim() == 2
+    torch.ops.sgl_kernel.deepseek_v4_topk_transform_512(
+        scores, seq_lens, page_table, page_indices, page_size, raw_indices
+    )
+
+
+def fast_topk_transform_ragged_fused(
+    score: torch.Tensor,
+    lengths: torch.Tensor,
+    topk_indices_offset: torch.Tensor,  # ragged kv
+    topk: int,
+    row_starts: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """
+    Get the topk indices of the score tensor and then transform the topk indices to
+    indices to ragged kv (non-paged). This function is only used for extend,
+    not including draft extend.
+    Args:
+        score: The score tensor of shape (B, L). The score tensor is the logits
+            between the query and the key which can be ragged or paged.
+            row_starts is only required when the key is ragged.
+        lengths: The lengths tensor of shape (B)
+        topk_indices_offset: The offset of topk indices in ragged kv of shape (B)
+        topk: The number of topk indices to get
+        row_starts: The start index of each row in the score tensor of shape (B).
+            For each row i, topk only applies to section [row_starts[i], row_starts[i] + lengths[i]]
+            of the score tensor. It can be None if only the fast path is triggered,
+            in the case of all values in lengths <= topk (not checked in the kernel,
+            guaranteed by the caller).
+    Returns:
+        The topk indices tensor of shape (B, topk)
+    """
+    assert (
+        topk == 2048
+    ), "fast_topk_transform_ragged_fused is only optimized for deepseek v3.2 model, where topk=2048"
+    assert score.dim() == 2
+    topk_indices_ragged = score.new_empty((score.shape[0], topk), dtype=torch.int32)
+    torch.ops.sgl_kernel.fast_topk_transform_ragged_fused(
+        score, lengths, topk_indices_ragged, topk_indices_offset, row_starts
+    )
+    return topk_indices_ragged

sgl-kernel/python/sgl_kernel/utils.py

@@ -0,0 +1,66 @@
+# Copyright 2025 SGLang Team. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import functools
+from typing import Dict, Tuple
+
+import torch
+
+_cache_buf: Dict[Tuple[str, torch.device], torch.Tensor] = {}
+
+
+def _get_cache_buf(name: str, bytes: int, device: torch.device) -> torch.Tensor:
+    key = (name, device)
+    buf = _cache_buf.get(key)
+    if buf is None:
+        buf = torch.empty(bytes, dtype=torch.uint8, device=device)
+        _cache_buf[key] = buf
+    return buf
+
+
+def _to_tensor_scalar_tuple(x):
+    if isinstance(x, torch.Tensor):
+        return (x, 0)
+    else:
+        return (None, x)
+
+
+def cache_once(fn):
+    """
+    NOTE: `functools.lru_cache` is not compatible with `torch.compile`
+    So we manually implement a simple cache_once decorator to replace it.
+    """
+    result_map = {}
+
+    @functools.wraps(fn)
+    def wrapper(*args, **kwargs):
+        key = (args, tuple(sorted(kwargs.items())))
+        if key not in result_map:
+            result_map[key] = fn(*args, **kwargs)
+        return result_map[key]
+
+    return wrapper
+
+
+@cache_once
+def is_arch_support_pdl() -> bool:
+    if getattr(torch.version, "hip", None) or getattr(torch.version, "musa", None):
+        return False
+    try:
+        device = torch.cuda.current_device()
+        major, _ = torch.cuda.get_device_capability(device)
+    except Exception:
+        return False
+    return major >= 9

sgl-kernel/python/sgl_kernel/version.py

@@ -0,0 +1 @@
+__version__ = "0.4.4"