sgl-project--sglang/python/sglang/jit_kernel/fused_metadata_copy.py

"""
Fused metadata copy kernel for DSA backend CUDA graph replay.

This module provides JIT-compiled CUDA kernels for fusing multiple tensor
copy operations into single kernel launches, reducing kernel launch overhead
and improving CUDA graph replay performance.

The kernels are compiled on-demand using TVM FFI and cached for subsequent use.
"""

from __future__ import annotations

import logging
from typing import Optional

import torch

from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args

logger = logging.getLogger(__name__)


# ============================================================================
# JIT Module Compilation
# ============================================================================


@cache_once
def _jit_fused_metadata_copy_module(
    forward_mode: int, has_real_page_table: bool, has_flashmla: bool
):
    """Compile JIT module for single-backend fused metadata copy.

    Args:
        forward_mode: 0=DECODE, 1=TARGET_VERIFY, 2=DRAFT_EXTEND
        has_real_page_table: Whether real_page_table tensors are used
        has_flashmla: Whether FlashMLA metadata tensors are used
    """
    args = make_cpp_args(forward_mode, has_real_page_table, has_flashmla)
    try:
        return load_jit(
            "fused_metadata_copy",
            *args,
            cuda_files=["elementwise/fused_metadata_copy.cuh"],
            cuda_wrappers=[
                (
                    "fused_metadata_copy",
                    f"FusedMetadataCopyKernel<{args}>::run",
                )
            ],
        )
    except Exception as e:
        logger.error(
            f"Failed to compile JIT fused metadata copy kernel "
            f"(forward_mode={forward_mode}, has_real_page_table={has_real_page_table}, "
            f"has_flashmla={has_flashmla}): {e}"
        )
        raise


@cache_once
def _jit_fused_metadata_copy_multi_module(
    has_real_page_table: bool, has_flashmla: bool
):
    """Compile JIT module for multi-backend fused metadata copy (DECODE mode only).

    Args:
        has_real_page_table: Whether real_page_table tensors are used
        has_flashmla: Whether FlashMLA metadata tensors are used
    """
    args = make_cpp_args(has_real_page_table, has_flashmla)
    try:
        return load_jit(
            "fused_metadata_copy_multi",
            *args,
            cuda_files=["elementwise/fused_metadata_copy.cuh"],
            cuda_wrappers=[
                (
                    "fused_metadata_copy_multi",
                    f"FusedMetadataCopyMultiKernel<{args}>::run",
                )
            ],
        )
    except Exception as e:
        logger.error(
            f"Failed to compile JIT fused metadata copy multi kernel "
            f"(has_real_page_table={has_real_page_table}, has_flashmla={has_flashmla}): {e}"
        )
        raise


# ============================================================================
# Public API
# ============================================================================


def fused_metadata_copy_cuda(
    cache_seqlens_src: torch.Tensor,
    cu_seqlens_k_src: torch.Tensor,
    page_indices_src: torch.Tensor,
    dsa_cache_seqlens_src: torch.Tensor,
    seqlens_expanded_src: Optional[torch.Tensor],
    dsa_cu_seqlens_k_src: torch.Tensor,
    real_page_table_src: Optional[torch.Tensor],
    flashmla_num_splits_src: Optional[torch.Tensor],
    flashmla_metadata_src: Optional[torch.Tensor],
    cache_seqlens_dst: torch.Tensor,
    cu_seqlens_k_dst: torch.Tensor,
    page_table_1_dst: torch.Tensor,
    dsa_cache_seqlens_dst: torch.Tensor,
    seqlens_expanded_dst: Optional[torch.Tensor],
    dsa_cu_seqlens_k_dst: torch.Tensor,
    real_page_table_dst: Optional[torch.Tensor],
    flashmla_num_splits_dst: Optional[torch.Tensor],
    flashmla_metadata_dst: Optional[torch.Tensor],
    forward_mode: int,
    bs: int,
    max_len: int,
    max_seqlen_k: int,
    seqlens_expanded_size: int,
) -> None:
    """
    Fused metadata copy kernel for DSA backend CUDA graph replay.

    This function fuses multiple tensor copy operations into a single kernel launch,
    reducing kernel launch overhead and improving performance.

    Args:
        cache_seqlens_src: Source cache sequence lengths [bs]
        cu_seqlens_k_src: Source cumulative sequence lengths [bs+1]
        page_indices_src: Source page indices [rows, max_len]
        dsa_cache_seqlens_src: Source DSA cache sequence lengths [size]
        seqlens_expanded_src: Optional source expanded sequence lengths [size] (required for TARGET_VERIFY/DRAFT_EXTEND)
        dsa_cu_seqlens_k_src: Source DSA cumulative sequence lengths [size+1]
        real_page_table_src: Optional source real page table [rows, cols]
        flashmla_num_splits_src: Optional source FlashMLA num_splits [size+1]
        flashmla_metadata_src: Optional source FlashMLA metadata tensor
        cache_seqlens_dst: Destination cache sequence lengths [bs]
        cu_seqlens_k_dst: Destination cumulative sequence lengths [bs+1]
        page_table_1_dst: Destination page table [rows, stride]
        dsa_cache_seqlens_dst: Destination DSA cache sequence lengths [size]
        seqlens_expanded_dst: Optional destination expanded sequence lengths [size] (required for TARGET_VERIFY/DRAFT_EXTEND)
        dsa_cu_seqlens_k_dst: Destination DSA cumulative sequence lengths [size+1]
        real_page_table_dst: Optional destination real page table [rows, cols]
        flashmla_num_splits_dst: Optional destination FlashMLA num_splits [size+1]
        flashmla_metadata_dst: Optional destination FlashMLA metadata tensor
        forward_mode: Forward mode (0=DECODE, 1=TARGET_VERIFY, 2=DRAFT_EXTEND)
        bs: Batch size
        max_len: Maximum length for decode/draft_extend mode
        max_seqlen_k: Maximum sequence length for target_verify mode
        seqlens_expanded_size: Size of expanded sequence lengths
    """
    # Determine template parameters for kernel specialization
    has_real_page_table = real_page_table_src is not None
    has_flashmla = flashmla_num_splits_src is not None

    # Get JIT-compiled module for this configuration (cached after first use)
    module = _jit_fused_metadata_copy_module(
        forward_mode, has_real_page_table, has_flashmla
    )

    # Ensure all required source tensors are contiguous (required for kernel's linear indexing)
    # This matches the CHECK_INPUT checks in the verified sgl-kernel implementation
    cache_seqlens_src = cache_seqlens_src.contiguous()
    cu_seqlens_k_src = cu_seqlens_k_src.contiguous()
    page_indices_src = page_indices_src.contiguous()
    dsa_cache_seqlens_src = dsa_cache_seqlens_src.contiguous()
    if seqlens_expanded_src is not None:
        seqlens_expanded_src = seqlens_expanded_src.contiguous()
    dsa_cu_seqlens_k_src = dsa_cu_seqlens_k_src.contiguous()

    # Call JIT-compiled kernel (None values are passed as Optional with no value)
    module.fused_metadata_copy(
        cache_seqlens_src,
        cu_seqlens_k_src,
        page_indices_src,
        dsa_cache_seqlens_src,
        seqlens_expanded_src,
        dsa_cu_seqlens_k_src,
        real_page_table_src,
        flashmla_num_splits_src,
        flashmla_metadata_src,
        cache_seqlens_dst,
        cu_seqlens_k_dst,
        page_table_1_dst,
        dsa_cache_seqlens_dst,
        seqlens_expanded_dst,
        dsa_cu_seqlens_k_dst,
        real_page_table_dst,
        flashmla_num_splits_dst,
        flashmla_metadata_dst,
        bs,
        max_len,
        max_seqlen_k,
        seqlens_expanded_size,
    )


def fused_metadata_copy_multi_cuda(
    cache_seqlens_src: torch.Tensor,
    cu_seqlens_k_src: torch.Tensor,
    page_indices_src: torch.Tensor,
    dsa_cache_seqlens_src: torch.Tensor,
    dsa_cu_seqlens_k_src: torch.Tensor,
    real_page_table_src: Optional[torch.Tensor],
    flashmla_num_splits_src: Optional[torch.Tensor],
    flashmla_metadata_src: Optional[torch.Tensor],
    cache_seqlens_dst0: torch.Tensor,
    cu_seqlens_k_dst0: torch.Tensor,
    page_table_1_dst0: torch.Tensor,
    dsa_cache_seqlens_dst0: torch.Tensor,
    dsa_cu_seqlens_k_dst0: torch.Tensor,
    real_page_table_dst0: Optional[torch.Tensor],
    flashmla_num_splits_dst0: Optional[torch.Tensor],
    flashmla_metadata_dst0: Optional[torch.Tensor],
    cache_seqlens_dst1: torch.Tensor,
    cu_seqlens_k_dst1: torch.Tensor,
    page_table_1_dst1: torch.Tensor,
    dsa_cache_seqlens_dst1: torch.Tensor,
    dsa_cu_seqlens_k_dst1: torch.Tensor,
    real_page_table_dst1: Optional[torch.Tensor],
    flashmla_num_splits_dst1: Optional[torch.Tensor],
    flashmla_metadata_dst1: Optional[torch.Tensor],
    cache_seqlens_dst2: torch.Tensor,
    cu_seqlens_k_dst2: torch.Tensor,
    page_table_1_dst2: torch.Tensor,
    dsa_cache_seqlens_dst2: torch.Tensor,
    dsa_cu_seqlens_k_dst2: torch.Tensor,
    real_page_table_dst2: Optional[torch.Tensor],
    flashmla_num_splits_dst2: Optional[torch.Tensor],
    flashmla_metadata_dst2: Optional[torch.Tensor],
    bs: int,
    max_len: int,
    seqlens_expanded_size: int,
) -> None:
    """
    Multi-backend fused metadata copy kernel for DSA backend CUDA graph replay.

    This function copies metadata from one source to THREE destinations in a single
    kernel launch, eliminating the overhead of 3 separate kernel calls. Currently
    only supports DECODE mode, which is the most common case.

    Args:
        cache_seqlens_src: Source cache sequence lengths [bs]
        cu_seqlens_k_src: Source cumulative sequence lengths [bs+1]
        page_indices_src: Source page indices [bs, max_len]
        dsa_cache_seqlens_src: Source DSA cache sequence lengths [bs]
        dsa_cu_seqlens_k_src: Source DSA cumulative sequence lengths [bs+1]
        real_page_table_src: Optional source real page table [bs, cols]
        flashmla_num_splits_src: Optional source FlashMLA num_splits [bs+1]
        flashmla_metadata_src: Optional source FlashMLA metadata tensor
        cache_seqlens_dst0-2: Destination cache sequence lengths for backends 0-2
        cu_seqlens_k_dst0-2: Destination cumulative sequence lengths for backends 0-2
        page_table_1_dst0-2: Destination page tables for backends 0-2
        dsa_cache_seqlens_dst0-2: Destination DSA cache sequence lengths for backends 0-2
        dsa_cu_seqlens_k_dst0-2: Destination DSA cumulative sequence lengths for backends 0-2
        real_page_table_dst0-2: Optional destination real page tables for backends 0-2
        flashmla_num_splits_dst0-2: Optional destination FlashMLA num_splits for backends 0-2
        flashmla_metadata_dst0-2: Optional destination FlashMLA metadata tensors for backends 0-2
        bs: Batch size
        max_len: Maximum length for decode mode
        seqlens_expanded_size: Size of expanded sequence lengths
    """
    # Determine template parameters for kernel specialization
    has_real_page_table = real_page_table_src is not None
    has_flashmla = flashmla_num_splits_src is not None

    # Get JIT-compiled module for this configuration (cached after first use)
    module = _jit_fused_metadata_copy_multi_module(has_real_page_table, has_flashmla)

    # Ensure all source tensors are contiguous (required for kernel's linear indexing)
    # This matches the CHECK_INPUT checks in the verified sgl-kernel implementation
    cache_seqlens_src = cache_seqlens_src.contiguous()
    cu_seqlens_k_src = cu_seqlens_k_src.contiguous()
    page_indices_src = page_indices_src.contiguous()
    dsa_cache_seqlens_src = dsa_cache_seqlens_src.contiguous()
    dsa_cu_seqlens_k_src = dsa_cu_seqlens_k_src.contiguous()

    # Call JIT-compiled kernel (None values are passed as Optional with no value)
    module.fused_metadata_copy_multi(
        cache_seqlens_src,
        cu_seqlens_k_src,
        page_indices_src,
        dsa_cache_seqlens_src,
        dsa_cu_seqlens_k_src,
        real_page_table_src,
        flashmla_num_splits_src,
        flashmla_metadata_src,
        cache_seqlens_dst0,
        cu_seqlens_k_dst0,
        page_table_1_dst0,
        dsa_cache_seqlens_dst0,
        dsa_cu_seqlens_k_dst0,
        real_page_table_dst0,
        flashmla_num_splits_dst0,
        flashmla_metadata_dst0,
        cache_seqlens_dst1,
        cu_seqlens_k_dst1,
        page_table_1_dst1,
        dsa_cache_seqlens_dst1,
        dsa_cu_seqlens_k_dst1,
        real_page_table_dst1,
        flashmla_num_splits_dst1,
        flashmla_metadata_dst1,
        cache_seqlens_dst2,
        cu_seqlens_k_dst2,
        page_table_1_dst2,
        dsa_cache_seqlens_dst2,
        dsa_cu_seqlens_k_dst2,
        real_page_table_dst2,
        flashmla_num_splits_dst2,
        flashmla_metadata_dst2,
        bs,
        max_len,
        seqlens_expanded_size,
    )