"""
Copyright (c) 2023 by FlashInfer team.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

  http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import functools
import logging
import math
from types import SimpleNamespace
from typing import Any, Dict, List, Literal, Optional, Tuple, Union, overload

import torch

from .api_logging import flashinfer_api
from .trace.templates.attention import (
    gqa_paged_prefill_trace,
    gqa_ragged_prefill_trace,
    single_prefill_with_kv_cache_trace,
    trtllm_batch_context_trace,
)
from .trace.templates.gemm import (
    fmha_v2_prefill_deepseek_trace,
    trtllm_ragged_attention_deepseek_trace,
)
from .trace.templates.page import trtllm_fmha_v2_prefill_trace
from .jit import (
    gen_batch_prefill_module,
    gen_customize_batch_prefill_module,
    gen_fmha_cutlass_sm100a_module,
    gen_fmha_v2_module,
    gen_single_prefill_module,
    get_batch_prefill_uri,
    get_single_prefill_uri,
    setup_cubin_loader,
    gen_trtllm_gen_fmha_module,
    gen_trtllm_fmha_v2_sm120_module,
)
from .cudnn import cudnn_batch_prefill_with_kv_cache
from .page import get_seq_lens
from .quantization import packbits, segment_packbits
from .utils import (
    log2e,
    FP4Tensor,
    MaskMode,
    PosEncodingMode,
    TensorLayout,
    _check_block_tables_shape,
    _check_cached_qkv_data_type,
    _check_kv_layout,
    _check_pos_encoding_mode,
    _check_workspace_buffer_alignment,
    check_shape_dtype_device,
    get_alibi_slopes,
    _get_cache_alibi_slopes_buf,
    _get_cache_buf,
    _get_trtllm_gen_multi_ctas_kv_counter_buffer,
    _resolve_trtllm_gen_multi_ctas_kv_counter_buffer,
    _unpack_paged_kv_cache,
    _should_use_fmha_v2_sm120,
    canonicalize_torch_dtype,
    determine_attention_backend,
    device_support_pdl,
    get_compute_capability,
    get_device_sm_count,
    is_float8,
    is_sm100a_supported,
    is_sm110a_supported,
    is_sm12x_supported,
    register_custom_op,
    register_fake_op,
    ceil_div,
    round_up,
    SINGLE_KERNEL_TMP_SIZE,
    prepare_jit_additional_args,
)


def _split_scale_param(scale):
    """Split scale parameter into tensor and scalar components.

    Args:
        scale: Can be a torch.Tensor (per-head), a scalar float, or None.

    Returns:
        tuple: (tensor_ptr, scalar_val) where:
            - If scale is tensor: (scale, 1.0)
            - If scale is scalar: (None, scale)
            - If scale is None: (None, 1.0)
    """
    if scale is None:
        return None, 1.0
    elif isinstance(scale, torch.Tensor):
        return scale, 1.0
    else:
        return None, float(scale)


@functools.cache
def get_fmha_module(
    dtype_q: torch.dtype,
    dtype_kv: torch.dtype,
    dtype_o: torch.dtype,
    dtype_idx: torch.dtype,
    head_dim_qk: int,
    head_dim_vo: int,
    pos_encoding_mode: int,
    use_sliding_window: bool,
    use_logits_soft_cap: bool,
    device: torch.device,
    use_fp16_qk_reduction: bool = False,
):
    if is_sm100a_supported(device) or is_sm110a_supported(device):
        return gen_fmha_cutlass_sm100a_module(
            dtype_q,
            dtype_kv,
            dtype_o,
            dtype_idx,
            head_dim_qk,
            head_dim_vo,
            pos_encoding_mode,
            use_sliding_window,
            use_logits_soft_cap,
        ).build_and_load()
    else:
        raise ValueError(
            "CUTLASS FMHA requires SM100a (B200/GB200) or SM110a. "
            "SM12x (RTX 5090/DGX Spark) lacks tcgen05 MMA required by this kernel. "
            "Use backend='fa2' instead."
        )


def make_hashable_cache(func):
    """
    Decorator that converts unhashable arguments (like lists) to hashable ones (tuples)
    before applying functools.cache.
    """

    @functools.cache
    def cached_wrapper(*args, **kwargs):
        return func(*args, **kwargs)

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        # Convert unhashable arguments to hashable ones
        hashable_args = []
        for arg in args:
            if isinstance(arg, list):
                hashable_args.append(tuple(arg))
            else:
                hashable_args.append(arg)

        hashable_kwargs = {}
        for key, value in kwargs.items():
            if isinstance(value, list):
                hashable_kwargs[key] = tuple(value)
            else:
                hashable_kwargs[key] = value

        return cached_wrapper(*hashable_args, **hashable_kwargs)

    return wrapper


@make_hashable_cache
def get_customize_batch_prefill_module(
    backend: str,
    uri: str,
    dtype_q: torch.dtype,
    dtype_kv: torch.dtype,
    dtype_o: torch.dtype,
    idtype: torch.dtype,
    head_dim_qk: int,
    head_dim_vo: int,
    additional_tensor_names: List[str],
    additional_tensor_dtypes: List[str],
    additional_scalar_names: List[str],
    additional_scalar_dtypes: List[str],
    variant_name: str,
    variant_decl: str,
    pos_encoding_mode: int = 0,
    use_sliding_window: bool = False,
    use_logits_soft_cap: bool = False,
    use_fp16_qk_reduction: bool = False,
    fp8_enabled: bool = False,
):
    return gen_customize_batch_prefill_module(
        backend,
        uri,
        dtype_q,
        dtype_kv,
        dtype_o,
        idtype,
        head_dim_qk,
        head_dim_vo,
        additional_tensor_names,
        additional_tensor_dtypes,
        additional_scalar_names,
        additional_scalar_dtypes,
        variant_name,
        variant_decl,
        pos_encoding_mode,
        use_sliding_window,
        use_logits_soft_cap,
        use_fp16_qk_reduction,
        fp8_enabled,
    ).build_and_load()


@functools.cache
def get_trtllm_gen_prefill_module():
    mod = gen_trtllm_gen_fmha_module()
    op = mod.build_and_load()
    setup_cubin_loader(mod.get_library_path())

    def _paged_run(
        query: torch.Tensor,
        k_cache: torch.Tensor,
        v_cache: torch.Tensor,
        workspace_buffer: torch.Tensor,
        block_tables: torch.Tensor,
        seq_lens: torch.Tensor,
        max_q_len: int,
        max_kv_len: int,
        bmm1_scale: Union[float, torch.Tensor],
        bmm2_scale: Union[float, torch.Tensor],
        batch_size: int,
        cum_seq_lens_q: torch.Tensor,
        cum_seq_lens_kv: torch.Tensor,
        enable_pdl: bool,
        workspace_size: int,
        window_left: int = -1,
        out: Optional[torch.Tensor] = None,
        sinks: Optional[torch.Tensor] = None,
        key_block_scales: Optional[torch.Tensor] = None,
        value_block_scales: Optional[torch.Tensor] = None,
        skip_softmax_threshold_scale_factor: Optional[float] = None,
        uses_shared_paged_kv_idx: bool = True,
        is_causal: bool = True,
        lse: Optional[torch.Tensor] = None,
        multi_ctas_kv_counter_buffer: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        sm_count = get_device_sm_count(query.device)
        if out is None:
            out = torch.empty_like(query)
        multi_ctas_kv_counter_buffer = _resolve_trtllm_gen_multi_ctas_kv_counter_buffer(
            multi_ctas_kv_counter_buffer,
            batch_size,
            query.size(1),
            sm_count,
            query.device,
        )
        if isinstance(bmm1_scale, torch.Tensor):
            assert bmm1_scale.dtype == torch.float32
            bmm1_scale = bmm1_scale * log2e
        if isinstance(bmm2_scale, torch.Tensor):
            assert bmm2_scale.dtype == torch.float32
        if lse is not None:
            check_shape_dtype_device(
                lse, (query.size(0), query.size(1)), torch.float32, query.device, "lse"
            )
            lse_stride_tokens = lse.stride(0)
            lse_stride_heads = lse.stride(1)
        else:
            lse_stride_tokens = 0
            lse_stride_heads = 0
        op.trtllm_paged_attention_context(
            out,
            None,  # fp4 output not supported in wrapper api yet.
            query,
            k_cache,
            v_cache,
            workspace_buffer,
            multi_ctas_kv_counter_buffer,
            block_tables,
            seq_lens,
            max_q_len,
            max_kv_len,
            bmm1_scale,
            bmm2_scale,
            -1,  # o_sf_scale
            -1,  # o_sf_vec_size
            0,  # o_sf_start_index
            batch_size,
            window_left,
            cum_seq_lens_q,
            cum_seq_lens_kv,
            sm_count,
            enable_pdl,
            workspace_size,
            sinks,
            key_block_scales,
            value_block_scales,
            skip_softmax_threshold_scale_factor,
            uses_shared_paged_kv_idx,
            is_causal,
            lse,
            lse_stride_tokens,
            lse_stride_heads,
        )
        return out

    def _ragged_run(*args, **kwargs):
        # TODO(Zihao): trtllm-gen backend already supports variable length attention,
        # but not integrated into flashinfer yet.
        raise NotImplementedError(
            "Variable length is not implemented for trtllm-gen backend yet."
        )

    def _plan(*args, **kwargs):
        pass

    return SimpleNamespace(
        paged_run=_paged_run,
        ragged_run=_ragged_run,
        plan=_plan,
    )


@functools.cache
def get_single_prefill_module(backend, *args):
    uri = get_single_prefill_uri(backend, *args)
    module = gen_single_prefill_module(backend, *args).build_and_load()
    run_func = module.run

    # torch library for single_prefill_with_kv_cache

    @register_custom_op(
        f"flashinfer::{uri}_run", mutates_args=("tmp", "o", "maybe_lse")
    )
    def run_single_prefill(
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        tmp: torch.Tensor,
        o: torch.Tensor,
        maybe_lse: Optional[torch.Tensor],
        mask_mode: int,
        layout: int,
        window_left: int,
        maybe_packed_custom_mask: Optional[torch.Tensor],
        maybe_alibi_slopes: Optional[torch.Tensor],
        logits_soft_cap: float,
        sm_scale: float,
        scale_q: Optional[torch.Tensor],
        scale_k: Optional[torch.Tensor],
        scale_v: Optional[torch.Tensor],
        rope_scale: float,
        rope_theta: float,
        maybe_k_cache_sf: Optional[torch.Tensor] = None,
        maybe_v_cache_sf: Optional[torch.Tensor] = None,
    ) -> None:
        if backend == "fa3":
            scale_v_tensor, scale_v_scalar = _split_scale_param(scale_v)
            if not is_float8(q):
                run_func(
                    q,
                    k,
                    v,
                    tmp,
                    o,
                    maybe_lse,
                    mask_mode,
                    layout,
                    window_left,
                    scale_v_tensor,
                    logits_soft_cap,
                    sm_scale,
                    scale_v_scalar,
                )
            else:
                # FP8 enabled
                scale_q_tensor, scale_q_scalar = _split_scale_param(scale_q)
                scale_k_tensor, scale_k_scalar = _split_scale_param(scale_k)
                run_func(
                    q,
                    k,
                    v,
                    tmp,
                    o,
                    maybe_lse,
                    mask_mode,
                    layout,
                    window_left,
                    scale_q_tensor,
                    scale_k_tensor,
                    scale_v_tensor,
                    sm_scale,
                    scale_q_scalar,
                    scale_k_scalar,
                    scale_v_scalar,
                )
        else:
            run_func(
                q,
                k,
                v,
                tmp,
                o,
                maybe_lse,
                mask_mode,
                layout,
                window_left,
                maybe_packed_custom_mask,
                maybe_alibi_slopes,
                maybe_k_cache_sf,
                maybe_v_cache_sf,
                logits_soft_cap,
                sm_scale,
                1.0 / rope_scale,  # rope_rcp_scale
                1.0 / rope_theta,  # rope_rcp_theta
            )
        return o

    @register_fake_op(f"flashinfer::{uri}_run")
    def _fake_run_single_prefill(
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        tmp: torch.Tensor,
        o: torch.Tensor,
        maybe_lse: Optional[torch.Tensor],
        mask_mode: int,
        layout: int,
        window_left: int,
        maybe_packed_custom_mask: Optional[torch.Tensor],
        maybe_alibi_slopes: Optional[torch.Tensor],
        logits_soft_cap: float,
        sm_scale: float,
        rope_scale: float,
        rope_theta: float,
        maybe_k_cache_sf: Optional[torch.Tensor] = None,
        maybe_v_cache_sf: Optional[torch.Tensor] = None,
    ) -> None:
        pass

    # Register the module
    return SimpleNamespace(run=run_single_prefill)


@functools.cache
def get_batch_prefill_module(backend, *args):
    if backend == "trtllm-gen":
        uri = "trtllm_gen_context"
        module = get_trtllm_gen_prefill_module()
        plan_func = module.plan
        workspace_size_func = None
        ragged_run_func = module.ragged_run
        paged_run_func = module.paged_run
    else:
        uri = get_batch_prefill_uri(backend, *args)
        module = gen_batch_prefill_module(backend, *args).build_and_load()
        plan_func = module.plan
        workspace_size_func = getattr(module, "workspace_size", None)
        ragged_run_func = module.ragged_run
        paged_run_func = module.paged_run

    # torch library for ragged_run

    @register_custom_op(
        f"flashinfer::{uri}_ragged_run",
        mutates_args=(
            "float_workspace_buffer",
            "int_workspace_buffer",
            "o",
            "maybe_lse",
        ),
    )
    def ragged_run(
        float_workspace_buffer: torch.Tensor,
        int_workspace_buffer: torch.Tensor,
        plan_info_vec: List[int],
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        qo_indptr: torch.Tensor,
        kv_indptr: torch.Tensor,
        o: torch.Tensor,
        maybe_lse: Optional[torch.Tensor],
        mask_mode: int,
        layout: int,
        window_left: int,
        enable_pdl: bool,
        maybe_custom_mask: Optional[torch.Tensor],
        maybe_mask_indptr: Optional[torch.Tensor],
        maybe_alibi_slopes: Optional[torch.Tensor],
        maybe_prefix_len_ptr: Optional[torch.Tensor],
        maybe_token_pos_in_items_ptr: Optional[torch.Tensor],
        maybe_max_item_len_ptr: Optional[torch.Tensor],
        logits_soft_cap: float,
        sm_scale: float,
        rope_scale: float,
        rope_theta: float,
        token_pos_in_items_len: int,
        maybe_k_cache_sf: Optional[torch.Tensor] = None,
        maybe_v_cache_sf: Optional[torch.Tensor] = None,
        scale_q: Optional[torch.Tensor] = None,
        scale_k: Optional[torch.Tensor] = None,
        scale_v: Optional[torch.Tensor] = None,
    ) -> None:
        # Check if FP8 by presence of scale tensors
        is_fp8 = scale_q is not None

        if backend == "fa2":
            ragged_run_func(
                float_workspace_buffer,
                int_workspace_buffer,
                plan_info_vec,
                q,
                k,
                v,
                qo_indptr,
                kv_indptr,
                o,
                maybe_lse,
                mask_mode,
                layout,
                window_left,
                enable_pdl,
                maybe_custom_mask,
                maybe_mask_indptr,
                maybe_alibi_slopes,
                maybe_prefix_len_ptr,
                maybe_token_pos_in_items_ptr,
                maybe_max_item_len_ptr,
                maybe_k_cache_sf,
                maybe_v_cache_sf,
                logits_soft_cap,
                sm_scale,
                1.0 / rope_scale,  # rope_rcp_scale
                1.0 / rope_theta,  # rope_rcp_theta,
                token_pos_in_items_len,
            )
        elif is_fp8:
            # FA3 FP8: scale_q, scale_k, scale_v, sm_scale, scale_q_scalar, scale_k_scalar, scale_v_scalar
            scale_q_tensor, scale_q_scalar = _split_scale_param(scale_q)
            scale_k_tensor, scale_k_scalar = _split_scale_param(scale_k)
            scale_v_tensor, scale_v_scalar = _split_scale_param(scale_v)
            ragged_run_func(
                float_workspace_buffer,
                int_workspace_buffer,
                plan_info_vec,
                q,
                k,
                v,
                qo_indptr,
                kv_indptr,
                o,
                maybe_lse,
                mask_mode,
                layout,
                window_left,
                enable_pdl,
                scale_q_tensor,
                scale_k_tensor,
                scale_v_tensor,
                sm_scale,
                scale_q_scalar,
                scale_k_scalar,
                scale_v_scalar,
            )
        else:
            # FA3 FP16: maybe_prefix_len_ptr, maybe_token_pos_in_items_ptr,
            # maybe_max_item_len_ptr, scale_v, logits_soft_cap, sm_scale, scale_v_scalar, token_pos_in_items_len
            scale_v_tensor, scale_v_scalar = _split_scale_param(scale_v)
            ragged_run_func(
                float_workspace_buffer,
                int_workspace_buffer,
                plan_info_vec,
                q,
                k,
                v,
                qo_indptr,
                kv_indptr,
                o,
                maybe_lse,
                mask_mode,
                layout,
                window_left,
                enable_pdl,
                maybe_prefix_len_ptr,
                maybe_token_pos_in_items_ptr,
                maybe_max_item_len_ptr,
                scale_v_tensor,
                logits_soft_cap,
                sm_scale,
                scale_v_scalar,
                token_pos_in_items_len,
            )

        return o

    @register_fake_op(f"flashinfer::{uri}_ragged_run")
    def _fake_ragged_run(
        float_workspace_buffer: torch.Tensor,
        int_workspace_buffer: torch.Tensor,
        plan_info_vec: List[int],
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        qo_indptr: torch.Tensor,
        kv_indptr: torch.Tensor,
        o: torch.Tensor,
        maybe_lse: Optional[torch.Tensor],
        mask_mode: int,
        layout: int,
        window_left: int,
        enable_pdl: bool,
        maybe_custom_mask: Optional[torch.Tensor],
        maybe_mask_indptr: Optional[torch.Tensor],
        maybe_alibi_slopes: Optional[torch.Tensor],
        maybe_prefix_len_ptr: Optional[torch.Tensor],
        maybe_token_pos_in_items_ptr: Optional[torch.Tensor],
        maybe_max_item_len_ptr: Optional[torch.Tensor],
        logits_soft_cap: float,
        sm_scale: float,
        rope_scale: float,
        rope_theta: float,
        token_pos_in_items_len: int,
        maybe_k_cache_sf: Optional[torch.Tensor] = None,
        maybe_v_cache_sf: Optional[torch.Tensor] = None,
    ) -> None:
        pass

    # torch library for paged_run

    @register_custom_op(
        f"flashinfer::{uri}_paged_run",
        mutates_args=(
            "float_workspace_buffer",
            "int_workspace_buffer",
            "multi_ctas_kv_counter_buffer",
            "paged_k_cache",
            "paged_v_cache",
            "o",
            "maybe_lse",
        ),
    )
    def paged_run(
        float_workspace_buffer: torch.Tensor,
        int_workspace_buffer: torch.Tensor,
        plan_info_vec: List[int],
        q: torch.Tensor,
        paged_k_cache: torch.Tensor,
        paged_v_cache: torch.Tensor,
        qo_indptr: torch.Tensor,
        paged_kv_indptr: torch.Tensor,
        paged_kv_indices: torch.Tensor,
        paged_kv_last_page_len: torch.Tensor,
        o: torch.Tensor,
        maybe_lse: Optional[torch.Tensor],
        mask_mode: int,
        layout: int,
        window_left: int,
        enable_pdl: bool,
        maybe_custom_mask: Optional[torch.Tensor],
        maybe_mask_indptr: Optional[torch.Tensor],
        maybe_alibi_slopes: Optional[torch.Tensor],
        maybe_prefix_len_ptr: Optional[torch.Tensor],
        maybe_token_pos_in_items_ptr: Optional[torch.Tensor],
        maybe_max_item_len_ptr: Optional[torch.Tensor],
        logits_soft_cap: float,
        sm_scale: float,
        scale_q: Optional[torch.Tensor],
        scale_k: Optional[torch.Tensor],
        scale_v: Optional[torch.Tensor],
        rope_scale: float,
        rope_theta: float,
        token_pos_in_items_len: int,
        workspace_size: int,
        num_qo_heads: Optional[int] = None,
        num_kv_heads: Optional[int] = None,
        block_tables: Optional[torch.Tensor] = None,
        kv_lens_buffer: Optional[torch.Tensor] = None,
        page_size: Optional[int] = None,
        max_q_len: Optional[int] = None,
        max_kv_len: Optional[int] = None,
        batch_size: Optional[int] = None,
        cum_seq_lens_q: Optional[torch.Tensor] = None,
        cum_seq_lens_kv: Optional[torch.Tensor] = None,
        sinks: Optional[torch.Tensor] = None,
        key_block_scales: Optional[torch.Tensor] = None,
        value_block_scales: Optional[torch.Tensor] = None,
        skip_softmax_threshold_scale_factor: Optional[float] = None,
        uses_shared_paged_kv_idx: bool = True,
        multi_ctas_kv_counter_buffer: Optional[torch.Tensor] = None,
    ) -> None:
        if backend == "trtllm-gen":
            assert num_qo_heads is not None
            assert num_kv_heads is not None
            assert block_tables is not None
            assert kv_lens_buffer is not None
            assert page_size is not None
            assert max_kv_len is not None
            assert batch_size is not None
            assert cum_seq_lens_q is not None
            assert cum_seq_lens_kv is not None
            assert enable_pdl is not None
            assert multi_ctas_kv_counter_buffer is not None
            assert workspace_size > 0, "workspace_size must be greater than 0"
            o = paged_run_func(
                q.contiguous(),  # NOTE(Siyuan): without contiguous, the result is incorrect
                paged_k_cache,
                paged_v_cache,
                int_workspace_buffer,
                block_tables,
                kv_lens_buffer,
                max_q_len,
                max_kv_len,
                sm_scale,
                1.0,  # NOTE(Siyuan): update this to expose bmm2 scale
                batch_size,
                cum_seq_lens_q,
                cum_seq_lens_kv,
                enable_pdl,
                workspace_size,
                window_left,
                out=o,
                sinks=sinks,
                key_block_scales=key_block_scales,
                value_block_scales=value_block_scales,
                skip_softmax_threshold_scale_factor=skip_softmax_threshold_scale_factor,
                uses_shared_paged_kv_idx=uses_shared_paged_kv_idx,
                is_causal=mask_mode != MaskMode.NON_CAUSAL.value,
                lse=maybe_lse,
                multi_ctas_kv_counter_buffer=multi_ctas_kv_counter_buffer,
            )
        elif backend == "fa2":
            assert not is_float8(q)
            paged_run_func(
                float_workspace_buffer,
                int_workspace_buffer,
                plan_info_vec,
                q,
                paged_k_cache,
                paged_v_cache,
                qo_indptr,
                paged_kv_indptr,
                paged_kv_indices,
                paged_kv_last_page_len,
                o,
                maybe_lse,
                mask_mode,
                layout,
                window_left,
                enable_pdl,
                maybe_custom_mask,
                maybe_mask_indptr,
                maybe_alibi_slopes,
                maybe_prefix_len_ptr,
                maybe_token_pos_in_items_ptr,
                maybe_max_item_len_ptr,
                key_block_scales,
                value_block_scales,
                logits_soft_cap,
                sm_scale,
                1.0 / rope_scale,  # rope_rcp_scale
                1.0 / rope_theta,  # rope_rcp_theta
                token_pos_in_items_len,
            )
        else:
            scale_v_tensor, scale_v_scalar = _split_scale_param(scale_v)
            if not is_float8(q):
                paged_run_func(
                    float_workspace_buffer,
                    int_workspace_buffer,
                    plan_info_vec,
                    q,
                    paged_k_cache,
                    paged_v_cache,
                    qo_indptr,
                    paged_kv_indptr,
                    paged_kv_indices,
                    paged_kv_last_page_len,
                    o,
                    maybe_lse,
                    mask_mode,
                    layout,
                    window_left,
                    enable_pdl,
                    maybe_prefix_len_ptr,
                    maybe_token_pos_in_items_ptr,
                    maybe_max_item_len_ptr,
                    scale_v_tensor,
                    logits_soft_cap,
                    sm_scale,
                    scale_v_scalar,
                    token_pos_in_items_len,
                )
            else:
                scale_q_tensor, scale_q_scalar = _split_scale_param(scale_q)
                scale_k_tensor, scale_k_scalar = _split_scale_param(scale_k)
                paged_run_func(
                    float_workspace_buffer,
                    int_workspace_buffer,
                    plan_info_vec,
                    q,
                    paged_k_cache,
                    paged_v_cache,
                    qo_indptr,
                    paged_kv_indptr,
                    paged_kv_indices,
                    paged_kv_last_page_len,
                    o,
                    maybe_lse,
                    mask_mode,
                    layout,
                    window_left,
                    enable_pdl,
                    scale_q_tensor,
                    scale_k_tensor,
                    scale_v_tensor,
                    sm_scale,
                    scale_q_scalar,
                    scale_k_scalar,
                    scale_v_scalar,
                )
        return o

    @register_fake_op(f"flashinfer::{uri}_paged_run")
    def _fake_paged_run(
        float_workspace_buffer: torch.Tensor,
        int_workspace_buffer: torch.Tensor,
        plan_info_vec: List[int],
        q: torch.Tensor,
        paged_k_cache: torch.Tensor,
        paged_v_cache: torch.Tensor,
        qo_indptr: torch.Tensor,
        paged_kv_indptr: torch.Tensor,
        paged_kv_indices: torch.Tensor,
        paged_kv_last_page_len: torch.Tensor,
        o: torch.Tensor,
        maybe_lse: Optional[torch.Tensor],
        mask_mode: int,
        layout: int,
        window_left: int,
        enable_pdl: bool,
        maybe_custom_mask: Optional[torch.Tensor],
        maybe_mask_indptr: Optional[torch.Tensor],
        maybe_alibi_slopes: Optional[torch.Tensor],
        maybe_prefix_len_ptr: Optional[torch.Tensor],
        maybe_token_pos_in_items_ptr: Optional[torch.Tensor],
        maybe_max_item_len_ptr: Optional[torch.Tensor],
        logits_soft_cap: float,
        sm_scale: float,
        scale_q: Optional[torch.Tensor],
        scale_k: Optional[torch.Tensor],
        scale_v: Optional[torch.Tensor],
        rope_scale: float,
        rope_theta: float,
        token_pos_in_items_len: int,
        workspace_size: int,
        num_qo_heads: Optional[int] = None,
        num_kv_heads: Optional[int] = None,
        block_tables: Optional[torch.Tensor] = None,
        kv_lens_buffer: Optional[torch.Tensor] = None,
        page_size: Optional[int] = None,
        max_q_len: Optional[int] = None,
        max_kv_len: Optional[int] = None,
        batch_size: Optional[int] = None,
        cum_seq_lens_q: Optional[torch.Tensor] = None,
        cum_seq_lens_kv: Optional[torch.Tensor] = None,
        sinks: Optional[torch.Tensor] = None,
        key_block_scales: Optional[torch.Tensor] = None,
        value_block_scales: Optional[torch.Tensor] = None,
        skip_softmax_threshold_scale_factor: Optional[float] = None,
        uses_shared_paged_kv_idx: bool = True,
        multi_ctas_kv_counter_buffer: Optional[torch.Tensor] = None,
    ) -> None:
        pass

    # Register the module.
    #
    # Note that plan is not part of model logic. It should not be included in
    # Cuda Graph or torch.compile. So, we don't provide a torch library for plan.
    return SimpleNamespace(
        plan=plan_func,
        workspace_size=workspace_size_func,
        ragged_run=ragged_run,
        paged_run=paged_run,
    )


@functools.cache
def get_batch_prefill_jit_module(module_name: str, jit_module: Any):
    plan_func = jit_module.plan
    workspace_size_func = getattr(jit_module, "workspace_size", None)
    ragged_run_func = jit_module.ragged_run
    paged_run_func = jit_module.paged_run

    # torch library for ragged_run
    @register_custom_op(
        f"flashinfer::{module_name}_ragged_run",
        mutates_args=(
            "float_workspace_buffer",
            "int_workspace_buffer",
            "o",
            "maybe_lse",
        ),
    )
    def ragged_run(
        float_workspace_buffer: torch.Tensor,
        int_workspace_buffer: torch.Tensor,
        plan_info_vec: List[int],
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        qo_indptr: torch.Tensor,
        kv_indptr: torch.Tensor,
        o: torch.Tensor,
        maybe_lse: Optional[torch.Tensor],
        mask_mode: int,
        layout: int,
        window_left: int,
        *args,
    ) -> None:
        ragged_run_func(
            float_workspace_buffer,
            int_workspace_buffer,
            plan_info_vec,
            q,
            k,
            v,
            qo_indptr,
            kv_indptr,
            o,
            maybe_lse,
            mask_mode,
            layout,
            window_left,
            *args,
        )

    @register_fake_op(f"flashinfer::{module_name}_ragged_run")
    def _fake_ragged_run(
        float_workspace_buffer: torch.Tensor,
        int_workspace_buffer: torch.Tensor,
        plan_info_vec: List[int],
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        qo_indptr: torch.Tensor,
        kv_indptr: torch.Tensor,
        o: torch.Tensor,
        maybe_lse: Optional[torch.Tensor],
        mask_mode: int,
        layout: int,
        window_left: int,
        *args,
    ) -> None:
        pass

    # torch library for paged_run
    @register_custom_op(
        f"flashinfer::{module_name}_paged_run",
        mutates_args=(
            "float_workspace_buffer",
            "int_workspace_buffer",
            "paged_k_cache",
            "paged_v_cache",
            "o",
            "maybe_lse",
        ),
    )
    def paged_run(
        float_workspace_buffer: torch.Tensor,
        int_workspace_buffer: torch.Tensor,
        plan_info_vec: List[int],
        q: torch.Tensor,
        paged_k_cache: torch.Tensor,
        paged_v_cache: torch.Tensor,
        qo_indptr: torch.Tensor,
        paged_kv_indptr: torch.Tensor,
        paged_kv_indices: torch.Tensor,
        paged_kv_last_page_len: torch.Tensor,
        o: torch.Tensor,
        maybe_lse: Optional[torch.Tensor],
        mask_mode: int,
        layout: int,
        window_left: int,
        *args,
    ) -> None:
        paged_run_func(
            float_workspace_buffer,
            int_workspace_buffer,
            plan_info_vec,
            q,
            paged_k_cache,
            paged_v_cache,
            qo_indptr,
            paged_kv_indptr,
            paged_kv_indices,
            paged_kv_last_page_len,
            o,
            maybe_lse,
            mask_mode,
            layout,
            window_left,
            *args,
        )

    @register_fake_op(f"flashinfer::{module_name}_paged_run")
    def _fake_paged_run(
        float_workspace_buffer: torch.Tensor,
        int_workspace_buffer: torch.Tensor,
        plan_info_vec: List[int],
        q: torch.Tensor,
        paged_k_cache: torch.Tensor,
        paged_v_cache: torch.Tensor,
        qo_indptr: torch.Tensor,
        paged_kv_indptr: torch.Tensor,
        paged_kv_indices: torch.Tensor,
        paged_kv_last_page_len: torch.Tensor,
        o: torch.Tensor,
        maybe_lse: Optional[torch.Tensor],
        mask_mode: int,
        layout: int,
        window_left: int,
        *args,
    ) -> None:
        pass

    # Register the module.
    #
    # Note that plan is not part of model logic. It should not be included in
    # Cuda Graph or torch.compile. So, we don't provide a torch library for plan.
    return SimpleNamespace(
        plan=plan_func,
        workspace_size=workspace_size_func,
        ragged_run=ragged_run,
        paged_run=paged_run,
    )


@flashinfer_api
def single_prefill_with_kv_cache_with_jit_module(
    jit_module: Any,
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    *args,
    kv_layout: str = "NHD",
    mask_mode: int = MaskMode.NON_CAUSAL.value,
    window_left: int = -1,
    return_lse: bool = False,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
    r"""Single-request prefill / append attention using a pre-compiled JIT module.

    Low-level entry point used by :func:`single_prefill_with_kv_cache` after backend
    dispatch; user code should normally call ``single_prefill_with_kv_cache`` directly.
    Exposed for advanced users who already hold a compiled JIT module (for example to
    customize the attention variant via extra kernel arguments).

    Parameters
    ----------
    jit_module : Any
        Compiled JIT module returned by one of the ``gen_*_module`` factories.
    q : torch.Tensor
        Query tensor, shape ``[qo_len, num_qo_heads, head_dim_qk]``.
    k : torch.Tensor
        Key tensor, shape ``[kv_len, num_kv_heads, head_dim_qk]`` (``NHD``) or
        ``[num_kv_heads, kv_len, head_dim_qk]`` (``HND``).
    v : torch.Tensor
        Value tensor, layout matches ``k``; last dimension may differ when
        ``head_dim_vo != head_dim_qk``.
    *args
        Extra positional arguments forwarded to the JIT module's ``run`` symbol
        (e.g. soft-cap, sink, or custom-mask buffers required by the chosen variant).
    kv_layout : str
        Layout of ``k`` and ``v``, either ``"NHD"`` or ``"HND"``.  Defaults to ``"NHD"``.
    mask_mode : int
        Mask mode, one of the values defined by :class:`~flashinfer.utils.MaskMode`.
        Defaults to ``MaskMode.NON_CAUSAL.value``.
    window_left : int
        Left window size for sliding-window attention; ``-1`` disables it.
    return_lse : bool
        Whether to allocate and return the log-sum-exp tensor.  Defaults to ``False``.

    Returns
    -------
    Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
        If ``return_lse`` is ``False``, the attention output tensor.  Otherwise the
        ``(output, lse)`` pair, where ``lse`` has shape
        ``[qo_len, num_qo_heads]`` and dtype ``float32``.
    """
    device = q.device
    tmp = torch.empty(SINGLE_KERNEL_TMP_SIZE, dtype=torch.uint8, device=device)
    o = torch.empty(q.shape[:-1] + v.shape[-1:], dtype=q.dtype, device=device)
    lse = None
    if return_lse:
        lse = torch.empty((q.size(0), q.size(1)), dtype=torch.float32, device=device)
    jit_module.run(
        q,
        k,
        v,
        tmp,
        o,
        lse,
        mask_mode,
        TensorLayout[kv_layout].value,
        window_left,
        *args,
    )
    return (o, lse) if return_lse else o


@overload
def single_prefill_with_kv_cache(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    scale_q: Optional[torch.Tensor] = None,
    scale_k: Optional[torch.Tensor] = None,
    scale_v: Optional[torch.Tensor] = None,
    o_dtype: Optional[torch.dtype] = None,
    custom_mask: Optional[torch.Tensor] = None,
    packed_custom_mask: Optional[torch.Tensor] = None,
    causal: bool = False,
    kv_layout: str = "NHD",
    pos_encoding_mode: str = "NONE",
    use_fp16_qk_reduction: bool = False,
    sm_scale: Optional[float] = None,
    window_left: int = -1,
    logits_soft_cap: Optional[float] = None,
    rope_scale: Optional[float] = None,
    rope_theta: Optional[float] = None,
    backend: str = "auto",
    return_lse: Literal[False] = False,
    kv_cache_sf: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
    k_scale: Optional[float] = None,
    v_scale: Optional[float] = None,
) -> torch.Tensor: ...


@overload
def single_prefill_with_kv_cache(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    scale_q: Optional[torch.Tensor] = None,
    scale_k: Optional[torch.Tensor] = None,
    scale_v: Optional[torch.Tensor] = None,
    o_dtype: Optional[torch.dtype] = None,
    custom_mask: Optional[torch.Tensor] = None,
    packed_custom_mask: Optional[torch.Tensor] = None,
    causal: bool = False,
    kv_layout: str = "NHD",
    pos_encoding_mode: str = "NONE",
    use_fp16_qk_reduction: bool = False,
    sm_scale: Optional[float] = None,
    window_left: int = -1,
    logits_soft_cap: Optional[float] = None,
    rope_scale: Optional[float] = None,
    rope_theta: Optional[float] = None,
    backend: str = "auto",
    return_lse: Literal[True] = True,
    kv_cache_sf: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
    k_scale: Optional[float] = None,
    v_scale: Optional[float] = None,
) -> Tuple[torch.Tensor, torch.Tensor]: ...


@flashinfer_api(trace=single_prefill_with_kv_cache_trace)
def single_prefill_with_kv_cache(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    scale_q: Optional[torch.Tensor] = None,
    scale_k: Optional[torch.Tensor] = None,
    scale_v: Optional[torch.Tensor] = None,
    o_dtype: Optional[torch.dtype] = None,
    custom_mask: Optional[torch.Tensor] = None,
    packed_custom_mask: Optional[torch.Tensor] = None,
    causal: bool = False,
    kv_layout: str = "NHD",
    pos_encoding_mode: str = "NONE",
    use_fp16_qk_reduction: bool = False,
    sm_scale: Optional[float] = None,
    window_left: int = -1,
    logits_soft_cap: Optional[float] = None,
    rope_scale: Optional[float] = None,
    rope_theta: Optional[float] = None,
    backend: str = "auto",
    return_lse: bool = False,
    kv_cache_sf: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
    k_scale: Optional[float] = None,
    v_scale: Optional[float] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
    r"""Prefill/Append attention with KV cache for single request, return the attention
    output.

    Parameters
    ----------
    q : torch.Tensor
        The query tensor, shape: ``[qo_len, num_qo_heads, head_dim_qk]``.
    k : torch.Tensor
        The key tensor, shape: ``[kv_len, num_kv_heads, head_dim_qk]`` if :attr:`kv_layout`
        is ``NHD``, or ``[num_kv_heads, kv_len, head_dim_qk]`` if :attr:`kv_layout` is
        ``HND``.
    v : torch.Tensor
        The key tensor, shape: ``[kv_len, num_kv_heads, head_dim_vo]`` if :attr:`kv_layout`
        is ``NHD``, ``[num_kv_heads, kv_len, head_dim_vo]`` if :attr:`kv_layout` is
        ``HND``.
    scale_q : Optional[torch.Tensor]
        The scale tensor for query, per-head quantization with shape: ``[num_qo_heads]``.
        Used with FP8 Quantization. If not provided, will be set to ``1.0``.
    scale_k : Optional[torch.Tensor]
        The scale tensor for key, per-head quantization with shape: ``[num_kv_heads]``.
        Used with FP8 Quantization. If not provided, will be set to ``1.0``.
    scale_v : Optional[torch.Tensor]
        The scale tensor for value, per-head quantization with shape: ``[num_kv_heads]``.
        Used with FP8 Quantization. If not provided, will be set to ``1.0``.
    o_dtype : Optional[torch.dtype]
        The output tensor data type, if not provided, will be set to the same as the q.
        This is necessary as output dtype cannot be automatically inferred in quant.
    custom_mask : Optional[torch.Tensor]
        The custom boolean mask tensor, shape: ``[qo_len, kv_len]``.
        The elements in the mask tensor should be either ``True`` or ``False``,
        where ``False`` means the corresponding element in the attention matrix will be
        masked out.

        When :attr:`custom_mask` is provided, and :attr:`packed_custom_mask` is not, the
        function will pack the custom mask tensor into a 1D packed mask tensor, which introduces
        additional overhead.
    packed_custom_mask : Optional[torch.Tensor]
        The 1D packed uint8 mask tensor, if provided, the :attr:`custom_mask` will be ignored.
        The packed mask tensor is generated by :func:`flashinfer.quantization.packbits`.
    causal : bool
        Whether to apply causal mask to the attention matrix.
        This is only effective when :attr:`custom_mask` is not provided.
    kv_layout : str
        The layout of the input k/v tensors, could be either ``NHD`` or ``HND``.
    pos_encoding_mode : str
        The position encoding applied inside attention kernels, could be
        ``NONE``/``ROPE_LLAMA`` (LLAMA style rotary embedding) /``ALIBI``.
        Default is ``NONE``.
    use_fp16_qk_reduction : bool
        Whether to use f16 for qk reduction (faster at the cost of slight precision
        loss).
    window_left : int
        The left (inclusive) window size for the attention window, when set to ``-1``, the window
        size will be set to the full length of the sequence. Defaults to ``-1``.
    logits_soft_cap : Optional[float]
        The attention logits soft capping value (used in Gemini, Grok and Gemma-2, etc.), if not
        provided, will be set to ``0``. If greater than 0, the logits will be capped according to
        formula:
        :math:`\texttt{logits_soft_cap} \times \mathrm{tanh}(x / \texttt{logits_soft_cap})`,
        where :math:`x` is the input logits.
    sm_scale : Optional[float]
        The scale used in softmax, if not provided, will be set to ``1.0 / sqrt(head_dim_qk)``.
    rope_scale : Optional[float]
        The scale used in RoPE interpolation, if not provided, will be set to 1.0.
    rope_theta : Optional[float]
        The theta used in RoPE, if not provided, will be set to 1e4.
    backend : str
        The implementation backend, could be ``auto``/``fa2`` or ``fa3``. Defaults to ``auto``.
        If set to ``auto``, the function will automatically choose the backend based on the
        device architecture and kernel availability.
    return_lse : bool
        Whether to return the log sum exp value of the attention logits.
    kv_cache_sf : Optional[Tuple[torch.Tensor, torch.Tensor]]
        Per-block scale factors for NVFP4 KV cache, as a tuple of ``(k_scales, v_scales)``.
        When provided, ``k`` and ``v`` are expected to be packed uint8 FP4 tensors with last
        dimension ``head_dim // 2``, and the scale factors dequantize them before attention.
        Both ``k_scales`` and ``v_scales`` use a linear (row-major) layout, and both have dtype ``torch.float8_e4m3fn``.

        Currently, NVFP4 KV only supports `fa2` backend.
    k_scale : Optional[Union[float, torch.Tensor]]
        The calibration scale of key for fp8 or nvfp4 input, if not provided, will be set to ``1.0``.
    v_scale : Optional[Union[float, torch.Tensor]]
        The calibration scale of value for fp8 or nvfp4 input, if not provided, will be set to ``1.0``.

    Returns
    -------
    Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
        If :attr:`return_lse` is ``False``, the attention output, shape: ``[qo_len, num_qo_heads, head_dim_vo]``.
        If :attr:`return_lse` is ``True``, a tuple of two tensors:

        * The attention output, shape: ``[qo_len, num_qo_heads, head_dim_vo]``.
        * The log sum exp value, shape: ``[qo_len, num_qo_heads]``.

    Examples
    --------

    >>> import torch
    >>> import flashinfer
    >>> qo_len = 128
    >>> kv_len = 4096
    >>> num_qo_heads = 32
    >>> num_kv_heads = 4
    >>> head_dim = 128
    >>> q = torch.randn(qo_len, num_qo_heads, head_dim).half().to("cuda:0")
    >>> k = torch.randn(kv_len, num_kv_heads, head_dim).half().to("cuda:0")
    >>> v = torch.randn(kv_len, num_kv_heads, head_dim).half().to("cuda:0")
    >>> o = flashinfer.single_prefill_with_kv_cache(q, k, v, causal=True,
            use_fp16_qk_reduction=True)
    >>> o.shape
    torch.Size([128, 32, 128])
    >>> mask = torch.tril(
    >>>     torch.full((qo_len, kv_len), True, device="cuda:0"),
    >>>     diagonal=(kv_len - qo_len),
    >>> )
    >>> mask
    tensor([[ True,  True,  True,  ..., False, False, False],
            [ True,  True,  True,  ..., False, False, False],
            [ True,  True,  True,  ..., False, False, False],
            ...,
            [ True,  True,  True,  ...,  True, False, False],
            [ True,  True,  True,  ...,  True,  True, False],
            [ True,  True,  True,  ...,  True,  True,  True]], device='cuda:0')
    >>> o_custom = flashinfer.single_prefill_with_kv_cache(q, k, v, custom_mask=mask)
    >>> torch.allclose(o, o_custom, rtol=1e-3, atol=1e-3)
    True

    Note
    ----
    The ``num_qo_heads`` must be a multiple of ``num_kv_heads``. If ``num_qo_heads`` is
    not equal to ``num_kv_heads``, the function will use
    `grouped query attention <https://arxiv.org/abs/2305.13245>`_.
    """
    _check_pos_encoding_mode(pos_encoding_mode)
    _check_kv_layout(kv_layout)
    tmp = torch.empty(SINGLE_KERNEL_TMP_SIZE, dtype=torch.uint8, device=q.device)
    if logits_soft_cap is None:
        logits_soft_cap = 0.0
    if sm_scale is None:
        sm_scale = 1.0 / math.sqrt(q.size(-1))
    if rope_scale is None:
        rope_scale = 1.0
    if rope_theta is None:
        rope_theta = 1e4
    if custom_mask is not None and packed_custom_mask is None:
        # create packed custom mask from custom mask
        packed_custom_mask = packbits(
            custom_mask.contiguous().view(-1), bitorder="little"
        )

    if packed_custom_mask is not None:
        mask_mode = MaskMode.CUSTOM.value
    else:
        if causal:
            mask_mode = MaskMode.CAUSAL.value
        else:
            mask_mode = MaskMode.NON_CAUSAL.value

    lse = None
    if return_lse:
        lse = torch.empty((q.size(0), q.size(1)), dtype=torch.float32, device=q.device)

    if is_float8(q):
        # FP8 quant enabled, do sanity check:
        #   1. unsupported feature
        #   2. dtype check
        assert window_left == -1
        assert q.dtype == k.dtype == v.dtype
        assert q.shape[-1] == k.shape[-1] == v.shape[-1]
        if scale_q is None:
            scale_q = torch.ones(q.shape[1], dtype=torch.float32, device=q.device)
        if scale_k is None:
            scale_k = torch.ones(k.shape[1], dtype=torch.float32, device=q.device)
        if scale_v is None:
            scale_v = torch.ones(v.shape[1], dtype=torch.float32, device=q.device)

    # For NVFP4 KV (uint8 packed), last dim is head_dim//2; output uses q head_dim.
    out_head_dim = q.shape[-1] if kv_cache_sf is not None else v.shape[-1]

    if backend == "auto":
        backend = determine_attention_backend(
            q.device,
            PosEncodingMode[pos_encoding_mode].value,
            use_fp16_qk_reduction,
            packed_custom_mask is not None,  # use_custom_mask
            q.dtype,
            k.dtype,
            head_dim_qk=q.shape[-1],
            head_dim_vo=out_head_dim,
        )

    # Unpack NVFP4 scale factors
    k_sf, v_sf = None, None
    if kv_cache_sf is not None:
        k_sf, v_sf = kv_cache_sf

    if k_scale is not None:
        sm_scale *= k_scale

    # o_dtype should be provided for FP8 attention
    if o_dtype is None:
        o_dtype = q.dtype
    out = torch.empty(q.shape[:-1] + (out_head_dim,), dtype=o_dtype, device=q.device)

    module = get_single_prefill_module(
        backend,
        q.dtype,
        k.dtype,
        out.dtype,
        q.shape[-1],  # head_dim_qk
        out_head_dim,  # head_dim_vo
        PosEncodingMode[pos_encoding_mode].value,
        window_left >= 0,  # use_sliding_window
        logits_soft_cap > 0,  # use_logits_soft_cap
        use_fp16_qk_reduction,
    )

    module.run(
        q,
        k,
        v,
        tmp,
        out,
        lse,
        mask_mode,
        TensorLayout[kv_layout].value,
        window_left,
        packed_custom_mask,
        get_alibi_slopes(q.shape[1], device=q.device)
        if pos_encoding_mode == "ALIBI"
        else None,
        logits_soft_cap,
        sm_scale,
        scale_q,
        scale_k,
        scale_v,
        rope_scale,
        rope_theta,
        k_sf,
        v_sf,
    )

    is_float_one = isinstance(v_scale, float) and v_scale == 1.0
    if v_scale is not None and not is_float_one:
        if is_float8(out):
            out = (out.to(torch.float32) * v_scale).to(out.dtype)
        else:
            out *= v_scale

    return (out, lse) if return_lse else out


single_prefill_with_kv_cache_return_lse = functools.partial(
    single_prefill_with_kv_cache, return_lse=True
)
single_prefill_with_kv_cache_return_lse.__doc__ = """Convenience wrapper for :func:`single_prefill_with_kv_cache` that always returns LSE.

    Equivalent to calling
    :func:`single_prefill_with_kv_cache` with ``return_lse=True``; accepts the
    same arguments and forwards them unchanged. See
    :func:`single_prefill_with_kv_cache` for the full parameter list (including
    FP8 / NVFP4 quantization scales such as ``scale_q``, ``scale_k``,
    ``scale_v``, ``o_dtype``, ``kv_cache_sf``, ``k_scale``, and ``v_scale``).

    Returns
    -------
    Tuple[torch.Tensor, torch.Tensor]
        A pair ``(output, lse)`` where ``output`` is the attention output and
        ``lse`` is the log-sum-exp tensor used for cascade merging.
    """


def _compute_page_mask_indptr(
    qo_indptr: torch.Tensor,
    paged_kv_indptr: torch.Tensor,
    paged_kv_last_page_len: torch.Tensor,
    page_size: int,
) -> torch.Tensor:
    if len(qo_indptr) != len(paged_kv_indptr):
        raise ValueError(
            "The length of qo_indptr and paged_kv_indptr should be the same."
        )
    mask_indptr = torch.empty_like(qo_indptr)
    mask_indptr[0] = 0
    mask_indptr[1:] = torch.cumsum(
        (qo_indptr[1:] - qo_indptr[:-1])
        * (
            (paged_kv_indptr[1:] - paged_kv_indptr[:-1] - 1) * page_size
            + paged_kv_last_page_len
        ),
        0,
    )
    return mask_indptr


class BatchPrefillWithPagedKVCacheWrapper:
    r"""Wrapper class for prefill/append attention with paged kv-cache for batch of
    requests.

    Check :ref:`our tutorial <kv-layout>` for page table layout.

    Example
    -------
    >>> import torch
    >>> import flashinfer
    >>> num_layers = 32
    >>> num_qo_heads = 64
    >>> num_kv_heads = 16
    >>> head_dim = 128
    >>> max_num_pages = 128
    >>> page_size = 16
    >>> # allocate 128MB workspace buffer
    >>> workspace_buffer = torch.zeros(128 * 1024 * 1024, dtype=torch.uint8, device="cuda:0")
    >>> prefill_wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(
    ...     workspace_buffer, "NHD"
    ... )
    >>> batch_size = 7
    >>> nnz_qo = 100
    >>> qo_indptr = torch.tensor(
    ...     [0, 33, 44, 55, 66, 77, 88, nnz_qo], dtype=torch.int32, device="cuda:0"
    ... )
    >>> paged_kv_indices = torch.arange(max_num_pages).int().to("cuda:0")
    >>> paged_kv_indptr = torch.tensor(
    ...     [0, 17, 29, 44, 48, 66, 100, 128], dtype=torch.int32, device="cuda:0"
    ... )
    >>> # 1 <= paged_kv_last_page_len <= page_size
    >>> paged_kv_last_page_len = torch.tensor(
    ...     [1, 7, 14, 4, 3, 1, 16], dtype=torch.int32, device="cuda:0"
    ... )
    >>> q_at_layer = torch.randn(num_layers, nnz_qo, num_qo_heads, head_dim).half().to("cuda:0")
    >>> kv_cache_at_layer = torch.randn(
    ...     num_layers, max_num_pages, 2, page_size, num_kv_heads, head_dim, dtype=torch.float16, device="cuda:0"
    ... )
    >>> # create auxiliary data structures for batch prefill attention
    >>> prefill_wrapper.plan(
    ...     qo_indptr,
    ...     paged_kv_indptr,
    ...     paged_kv_indices,
    ...     paged_kv_last_page_len,
    ...     num_qo_heads,
    ...     num_kv_heads,
    ...     head_dim,
    ...     page_size,
    ...     causal=True,
    ... )
    >>> outputs = []
    >>> for i in range(num_layers):
    ...     q = q_at_layer[i]
    ...     kv_cache = kv_cache_at_layer[i]
    ...     # compute batch prefill attention, reuse auxiliary data structures
    ...     o = prefill_wrapper.run(q, kv_cache)
    ...     outputs.append(o)
    ...
    >>> outputs[0].shape
    torch.Size([100, 64, 128])
    >>>
    >>> # below is another example of creating custom mask for batch prefill attention
    >>> mask_arr = []
    >>> qo_len = (qo_indptr[1:] - qo_indptr[:-1]).cpu().tolist()
    >>> kv_len = (page_size * (paged_kv_indptr[1:] - paged_kv_indptr[:-1] - 1) + paged_kv_last_page_len).cpu().tolist()
    >>> for i in range(batch_size):
    ...     mask_i = torch.tril(
    ...         torch.full((qo_len[i], kv_len[i]), True, device="cuda:0"),
    ...         diagonal=(kv_len[i] - qo_len[i]),
    ...     )
    ...     mask_arr.append(mask_i.flatten())
    ...
    >>> mask = torch.cat(mask_arr, dim=0)
    >>> prefill_wrapper.plan(
    ...     qo_indptr,
    ...     paged_kv_indptr,
    ...     paged_kv_indices,
    ...     paged_kv_last_page_len,
    ...     num_qo_heads,
    ...     num_kv_heads,
    ...     head_dim,
    ...     page_size,
    ...     custom_mask=mask,
    ... )
    >>> for i in range(num_layers):
    ...     q = q_at_layer[i]
    ...     kv_cache = kv_cache_at_layer[i]
    ...     # compute batch prefill attention, reuse auxiliary data structures
    ...     o_custom = prefill_wrapper.run(q, kv_cache)
    ...     assert torch.allclose(o_custom, outputs[i], rtol=1e-3, atol=1e-3)
    ...



    Note
    ----
    To accelerate computation, FlashInfer's batch prefill/append attention operators
    create some auxiliary data structures, these data structures can be reused across
    multiple prefill/append attention calls (e.g. different Transformer layers). This
    wrapper class manages the lifecycle of these data structures.
    """

    @flashinfer_api
    def __init__(
        self,
        float_workspace_buffer: torch.Tensor,
        kv_layout: str = "NHD",
        use_cuda_graph: bool = False,
        qo_indptr_buf: Optional[torch.Tensor] = None,
        paged_kv_indptr_buf: Optional[torch.Tensor] = None,
        paged_kv_indices_buf: Optional[torch.Tensor] = None,
        paged_kv_last_page_len_buf: Optional[torch.Tensor] = None,
        custom_mask_buf: Optional[torch.Tensor] = None,
        mask_indptr_buf: Optional[torch.Tensor] = None,
        backend: str = "auto",
        jit_args: Optional[List[Any]] = None,
        jit_kwargs: Optional[Dict[str, Any]] = None,
    ) -> None:
        r"""Constructor of :class:`BatchPrefillWithPagedKVCacheWrapper`.

        Parameters
        ----------
        float_workspace_buffer : torch.Tensor
            The user reserved workspace buffer used to store intermediate attention results in
            split-k algorithm. The recommended size is 128MB, the device of the workspace buffer
            should be the same as the device of the input tensors. The buffer must be 16-byte
            aligned; tensors created by ``torch.empty`` satisfy this on supported devices.

        kv_layout : str
            The layout of the input k/v tensors, could be either ``NHD`` or ``HND``.

        use_cuda_graph : bool
            Whether to enable CUDA graph capture for the prefill kernels, if enabled, the
            auxiliary data structures will be stored in provided buffers. The ``batch_size``
            cannot change during the lifecycle of this wrapper when CUDAGraph is enabled.

        qo_indptr_buf : Optional[torch.Tensor]
            The user reserved buffer to store the ``qo_indptr`` array, the size of the buffer
            should be ``[batch_size + 1]``.
            This argument is only effective when ``use_cuda_graph`` is ``True``.

        paged_kv_indptr_buf : Optional[torch.Tensor]
            The user reserved buffer to store the ``paged_kv_indptr`` array, the size of this
            buffer should be ``[batch_size + 1]``.
            This argument is only effective when ``use_cuda_graph`` is ``True``.

        paged_kv_indices_buf : Optional[torch.Tensor]
            The user reserved buffer to store the ``paged_kv_indices`` array, should be large
            enough to store the maximum possible size of the ``paged_kv_indices`` array during
            the lifetime of the wrapper. This argument is only effective when ``use_cuda_graph``
            is ``True``.

        paged_kv_last_page_len_buf : Optional[torch.Tensor]
            The user reserved buffer to store the ``paged_kv_last_page_len`` array, the size of
            the buffer should be ``[batch_size]``.
            This argument is only effective when ``use_cuda_graph`` is ``True``.

        custom_mask_buf : Optional[torch.Tensor]
            The user reserved buffer to store the custom mask tensor, should be large enough to
            store the maximum possible size of the packed custom mask tensor during the lifetime of
            the wrapper. This argument is only effective when ``use_cuda_graph`` is set to ``True``
            and the custom mask will be used in attention computation.

        mask_indptr_buf : Optional[torch.Tensor]
            The user reserved buffer to store the ``mask_indptr`` array, the size of the buffer
            should be ``[batch_size + 1]``.
            This argument is only effective when ``use_cuda_graph`` is ``True`` and the custom
            mask will be used in attention computation.

        backend : str
            The implementation backend, could be ``auto``/``fa2``/``fa3``/``cudnn`` or ``trtllm-gen``.
            Defaults to ``auto``.
            If set to ``auto``, the wrapper will automatically choose the backend based on the
            device architecture and kernel availability.

        jit_args : Optional[List[Any]]
            If provided, the wrapper will use the provided arguments to create the JIT module,
            otherwise, the wrapper will use default attention implementation.

        jit_kwargs : Optional[Dict[str, Any]]
            The keyword arguments to create the JIT module, defaults to None.
        """
        _check_workspace_buffer_alignment(
            float_workspace_buffer, "float_workspace_buffer"
        )
        _check_kv_layout(kv_layout)

        if backend == "cute-dsl":
            raise NotImplementedError(
                "cute-dsl backend is not yet supported for paged KV cache. "
                "Use BatchPrefillWithRaggedKVCacheWrapper instead."
            )

        if jit_args is not None and backend != "cute-dsl":
            if jit_kwargs is None:
                jit_kwargs = {}
            self._jit_module = get_batch_prefill_jit_module(
                jit_args[0],
                get_customize_batch_prefill_module(backend, *jit_args, **jit_kwargs),
            )
            # jit_args[7] is additional_tensor_names from gen_customize_batch_prefill_module
            self._jit_additional_tensor_names = list(jit_args[7])
        else:
            self._jit_module = None
            self._jit_additional_tensor_names = []

        self._kv_layout = kv_layout
        if backend == "cudnn":
            assert kv_layout == "NHD", "CUDNN backend only supports NHD layout"

        self._float_workspace_buffer = float_workspace_buffer
        self._workspace_size = (
            self._float_workspace_buffer.numel()
            * self._float_workspace_buffer.element_size()
        )
        self.device = float_workspace_buffer.device
        self._trtllm_gen_multi_ctas_kv_counter_buffer: Optional[torch.Tensor] = None

        self._kv_lens_buffer = torch.empty(
            (32768,), dtype=torch.int32, device=self.device
        )
        self._int_workspace_buffer = torch.empty(
            (8 * 1024 * 1024,), dtype=torch.uint8, device=self.device
        )
        self._pin_memory_int_workspace_buffer = torch.empty(
            self._int_workspace_buffer.shape,
            dtype=self._int_workspace_buffer.dtype,
            device="cpu",
            pin_memory=True,
        )
        self._use_cuda_graph = use_cuda_graph
        if use_cuda_graph:
            if not torch.is_tensor(qo_indptr_buf):
                raise ValueError(
                    "qo_indptr_buf should be a torch.Tensor in CUDA graph mode"
                )
            if not torch.is_tensor(paged_kv_indptr_buf):
                raise ValueError(
                    "paged_kv_indptr_buf should be a torch.Tensor in CUDA graph mode"
                )
            if not torch.is_tensor(paged_kv_indices_buf):
                raise ValueError(
                    "paged_kv_indices_buf should be a torch.Tensor in CUDA graph mode"
                )
            if not torch.is_tensor(paged_kv_last_page_len_buf):
                raise ValueError(
                    "paged_kv_last_page_len_buf should be a torch.Tensor in CUDA graph mode"
                )
            self._fixed_batch_size = len(qo_indptr_buf) - 1
            if len(paged_kv_indptr_buf) != self._fixed_batch_size + 1:
                raise ValueError(
                    "The length of paged_kv_indptr_buf should be batch_size + 1."
                )
            if len(paged_kv_last_page_len_buf) != self._fixed_batch_size:
                raise ValueError(
                    "The length of paged_kv_last_page_len_buf should be batch_size."
                )
            # NOTE(Zihao): do not check custom_mask_buf and mask_indptr_buf here, as they are optional
        else:
            self._fixed_batch_size = 0

        self._qo_indptr_buf = qo_indptr_buf
        self._paged_kv_indptr_buf = paged_kv_indptr_buf
        self._paged_kv_indices_buf = paged_kv_indices_buf
        self._paged_kv_last_page_len_buf = paged_kv_last_page_len_buf
        self._custom_mask_buf = custom_mask_buf
        self._mask_indptr_buf = mask_indptr_buf
        self._max_total_num_rows: Optional[int] = None
        self._backend = backend
        self._plan_info = None
        self._cached_module = None
        self._seq_lens_kv = None
        self._seq_lens_q = None
        self._block_tables = None

    @property
    def is_cuda_graph_enabled(self) -> bool:
        return self._use_cuda_graph

    def reset_workspace_buffer(
        self, float_workspace_buffer: torch.Tensor, int_workspace_buffer: torch.Tensor
    ) -> None:
        r"""Reset the workspace buffer.

        Parameters
        ----------
        float_workspace_buffer : torch.Tensor
            The new float workspace buffer, the device of the new float workspace buffer should
            be the same as the device of the input tensors.

        int_workspace_buffer : torch.Tensor
            The new int workspace buffer, the device of the new int workspace buffer should
            be the same as the device of the input tensors.
        """
        _check_workspace_buffer_alignment(
            float_workspace_buffer, "float_workspace_buffer"
        )
        _check_workspace_buffer_alignment(int_workspace_buffer, "int_workspace_buffer")
        self._float_workspace_buffer = float_workspace_buffer
        self._int_workspace_buffer = int_workspace_buffer
        self._pin_memory_int_workspace_buffer = torch.empty(
            self._int_workspace_buffer.shape,
            dtype=self._int_workspace_buffer.dtype,
            device="cpu",
            pin_memory=True,
        )

    @flashinfer_api
    def workspace_size(
        self,
        qo_indptr: torch.Tensor,
        paged_kv_indptr: torch.Tensor,
        paged_kv_indices: torch.Tensor,
        paged_kv_last_page_len: torch.Tensor,
        num_qo_heads: int,
        num_kv_heads: int,
        head_dim_qk: int,
        page_size: int,
        head_dim_vo: Optional[int] = None,
        custom_mask: Optional[torch.Tensor] = None,
        packed_custom_mask: Optional[torch.Tensor] = None,
        causal: bool = False,
        pos_encoding_mode: str = "NONE",
        use_fp16_qk_reduction: bool = False,
        sm_scale: Optional[float] = None,
        window_left: int = -1,
        logits_soft_cap: Optional[float] = None,
        rope_scale: Optional[float] = None,
        rope_theta: Optional[float] = None,
        q_data_type: Union[str, torch.dtype] = "float16",
        kv_data_type: Optional[Union[str, torch.dtype]] = None,
        o_data_type: Optional[Union[str, torch.dtype]] = None,
        prefix_len_ptr: Optional[torch.Tensor] = None,
        token_pos_in_items_ptr: Optional[torch.Tensor] = None,
        token_pos_in_items_len: int = 0,
        max_item_len_ptr: Optional[torch.Tensor] = None,
        seq_lens: Optional[torch.Tensor] = None,
        seq_lens_q: Optional[torch.Tensor] = None,
        block_tables: Optional[torch.Tensor] = None,
        max_token_per_sequence: Optional[int] = None,
        max_sequence_kv: Optional[int] = None,
        fixed_split_size: Optional[int] = None,
        disable_split_kv: bool = False,
    ) -> Tuple[int, int]:
        r"""Return the caller-owned workspace size required by :meth:`plan`.

        The returned tuple is ``(float_workspace_size, int_workspace_size)`` in
        bytes.  The inputs follow :meth:`plan`; host-side planning tensors such
        as ``qo_indptr`` and ``paged_kv_indptr`` are copied to CPU in the same
        way as :meth:`plan`.  The wrapper's float workspace buffer is only used
        as the device/stream selector for the underlying module query.  This
        method does not allocate buffers and does not mutate cached plan state.

        Parameters
        ----------
        qo_indptr : torch.Tensor
            The indptr of the query/output tensor, shape: ``[batch_size + 1]``.
        paged_kv_indptr : torch.Tensor
            The indptr of the paged kv-cache, shape: ``[batch_size + 1]``.
        paged_kv_indices : torch.Tensor
            The page indices of the paged kv-cache, shape: ``[paged_kv_indptr[-1]]``.
        paged_kv_last_page_len : torch.Tensor
            The number of entries in the last page of each request in the paged
            kv-cache, shape: ``[batch_size]``.
        num_qo_heads : int
            The number of query/output heads.
        num_kv_heads : int
            The number of key/value heads.
        head_dim_qk : int
            The dimension of the query/key heads.
        page_size : int
            The size of each page in the paged kv-cache.
        head_dim_vo : Optional[int]
            The dimension of the value/output heads. If not provided, defaults to
            ``head_dim_qk``.
        custom_mask : Optional[torch.Tensor]
            The flattened boolean mask tensor; when provided the packed mask is
            generated internally. See :meth:`plan` for the full description.
        packed_custom_mask : Optional[torch.Tensor]
            The 1D packed uint8 mask tensor. If provided, ``custom_mask`` is ignored.
        causal : bool
            Whether to apply a causal mask. Defaults to ``False``.
        pos_encoding_mode : str
            The position encoding applied inside attention kernels, could be
            ``NONE``/``ROPE_LLAMA`` (LLAMA style rotary embedding) /``ALIBI``.
            Defaults to ``NONE``.
        use_fp16_qk_reduction : bool
            Whether to use fp16 for qk reduction. Defaults to ``False``.
        sm_scale : Optional[float]
            Softmax scale. If ``None``, defaults to ``1.0 / sqrt(head_dim_qk)``.
        window_left : int
            The left (inclusive) window size for the attention window, when set to ``-1``, the window
            size will be set to the full length of the sequence. Defaults to ``-1``.
        logits_soft_cap : Optional[float]
            The attention logits soft capping value (used in Gemini, Grok and Gemma-2, etc.),
            if not provided, will be set to ``0``.
        rope_scale : Optional[float]
            Scale factor applied during RoPE interpolation. Defaults to ``1.0`` when ``None``.
        rope_theta : Optional[float]
            Base value for the RoPE frequencies. Defaults to ``1e4`` when ``None``.
        q_data_type : Union[str, torch.dtype]
            The data type of the query tensor. Defaults to ``torch.float16``.
        kv_data_type : Optional[Union[str, torch.dtype]]
            The data type of the key/value tensor. If ``None``, will be set to ``q_data_type``.
        o_data_type : Optional[Union[str, torch.dtype]]
            The data type of the output tensor. If ``None``, will be set to ``q_data_type``.
        prefix_len_ptr : Optional[torch.Tensor]
            A uint32 1D tensor indicating the prefix length of each prompt,
            shape: ``[batch_size]``.
        token_pos_in_items_ptr : Optional[torch.Tensor]
            A uint16 1D tensor indicating the token position of each item within its item.
        token_pos_in_items_len : int
            Zero-padding length for ``token_pos_in_items_ptr``. Defaults to ``0``.
        max_item_len_ptr : Optional[torch.Tensor]
            A uint16 vector containing the max token length of all items for each prompt.
        seq_lens : Optional[torch.Tensor]
            A uint32 1D tensor indicating the kv sequence length of each prompt,
            shape: ``[batch_size]``.
        seq_lens_q : Optional[torch.Tensor]
            A uint32 1D tensor indicating the q sequence length of each prompt,
            shape: ``[batch_size]``.
        block_tables : Optional[torch.Tensor]
            A uint32 2D tensor indicating the block table of each prompt,
            shape: ``[batch_size, max_num_blocks_per_seq]``.
        max_token_per_sequence : Optional[int]
            Required for cuDNN backend. The scalar max token length of each sequence.
        max_sequence_kv : Optional[int]
            Maximum number of KV tokens per sequence for cuDNN backend.
        fixed_split_size : Optional[int]
            The fixed split size for split-kv prefill, in pages.
        disable_split_kv : bool
            Whether to disable the split-kv. Defaults to ``False``.

        Returns
        -------
        Tuple[int, int]
            ``(float_workspace_size, int_workspace_size)`` in bytes.

        Example
        -------
        >>> float_bytes, int_bytes = wrapper.workspace_size(...)
        >>> wrapper.reset_workspace_buffer(
        ...     torch.empty(float_bytes, dtype=torch.uint8, device="cuda"),
        ...     torch.empty(int_bytes, dtype=torch.uint8, device="cuda"),
        ... )
        >>> wrapper.plan(...)
        """
        _check_workspace_buffer_alignment(
            self._float_workspace_buffer, "float_workspace_buffer"
        )
        del (
            block_tables,
            max_item_len_ptr,
            max_token_per_sequence,
            prefix_len_ptr,
            rope_scale,
            rope_theta,
            seq_lens_q,
            sm_scale,
            token_pos_in_items_len,
            token_pos_in_items_ptr,
        )
        q_data_type = canonicalize_torch_dtype(q_data_type)
        if kv_data_type is None:
            kv_data_type = q_data_type
        kv_data_type = canonicalize_torch_dtype(kv_data_type)
        if o_data_type is None:
            o_data_type = q_data_type
        o_data_type = canonicalize_torch_dtype(o_data_type)

        if logits_soft_cap is None:
            logits_soft_cap = 0.0
        if head_dim_vo is None:
            head_dim_vo = head_dim_qk
        if fixed_split_size is None:
            fixed_split_size = -1

        batch_size = len(qo_indptr) - 1
        qo_indptr_host = qo_indptr.to("cpu")
        total_num_rows = int(qo_indptr_host[-1])
        paged_kv_indptr_host = paged_kv_indptr.to("cpu")
        paged_kv_last_page_len_host = paged_kv_last_page_len.to("cpu")
        if seq_lens is None:
            kv_lens_arr_host = get_seq_lens(
                paged_kv_indptr_host, paged_kv_last_page_len_host, page_size
            )
        else:
            kv_lens_arr_host = seq_lens.cpu().flatten()

        if self.is_cuda_graph_enabled:
            if batch_size != self._fixed_batch_size:
                raise ValueError(
                    "The batch size should be fixed during the lifecycle of the wrapper in "
                    "cuda graph mode, the runtime batch size {} mismatches the batch size {} "
                    " set during initialization.".format(
                        batch_size, self._fixed_batch_size
                    )
                )
            if len(paged_kv_indices) > len(self._paged_kv_indices_buf):
                raise ValueError(
                    "The length of paged_kv_indices exceeds the allocated buffer size."
                )
            if (
                self._max_total_num_rows is not None
                and total_num_rows > self._max_total_num_rows
            ):
                raise ValueError(
                    "The total number of rows in qo_indptr {} in cuda graph mode cannot "
                    "exceed the number of rows set during initialization {}.".format(
                        total_num_rows, self._max_total_num_rows
                    )
                )

        use_custom_mask = custom_mask is not None or packed_custom_mask is not None
        backend = self._backend
        if self._jit_module is not None:
            module = self._jit_module
        else:
            if backend == "auto":
                backend = determine_attention_backend(
                    self.device,
                    PosEncodingMode[pos_encoding_mode].value,
                    use_fp16_qk_reduction,
                    use_custom_mask,
                    q_data_type,
                    kv_data_type,
                )
            if backend == "cudnn":
                raise NotImplementedError(
                    "workspace_size is not available for cudnn prefill backend"
                )
            module = get_batch_prefill_module(
                backend,
                q_data_type,
                kv_data_type,
                o_data_type,
                paged_kv_indptr.dtype,
                head_dim_qk,
                head_dim_vo,
                PosEncodingMode[pos_encoding_mode].value,
                window_left >= 0,
                logits_soft_cap > 0,
                use_fp16_qk_reduction,
            )
        if module.workspace_size is None:
            raise NotImplementedError(
                f"workspace_size is not available for prefill backend {backend!r}"
            )

        del max_sequence_kv
        total_num_rows_for_plan = self._max_total_num_rows or total_num_rows
        args = [
            self._float_workspace_buffer,
            qo_indptr_host,
            paged_kv_indptr_host,
            kv_lens_arr_host,
            total_num_rows_for_plan,
            batch_size,
            num_qo_heads,
            num_kv_heads,
            page_size,
            self.is_cuda_graph_enabled,
            head_dim_qk,
            head_dim_vo,
            causal,
            window_left,
        ]
        if backend == "fa2":
            args.append(fixed_split_size)
            args.append(disable_split_kv)
            args.append(0)  # num_colocated_ctas
            args.append(0)  # uniform_q_len
        float_workspace_size, int_workspace_size = module.workspace_size(*args)
        return int(float_workspace_size), int(int_workspace_size)

    @flashinfer_api
    def plan(
        self,
        qo_indptr: torch.Tensor,
        paged_kv_indptr: torch.Tensor,
        paged_kv_indices: torch.Tensor,
        paged_kv_last_page_len: torch.Tensor,
        num_qo_heads: int,
        num_kv_heads: int,
        head_dim_qk: int,
        page_size: int,
        head_dim_vo: Optional[int] = None,
        custom_mask: Optional[torch.Tensor] = None,
        packed_custom_mask: Optional[torch.Tensor] = None,
        causal: bool = False,
        pos_encoding_mode: str = "NONE",
        use_fp16_qk_reduction: bool = False,
        sm_scale: Optional[float] = None,
        window_left: int = -1,
        logits_soft_cap: Optional[float] = None,
        rope_scale: Optional[float] = None,
        rope_theta: Optional[float] = None,
        q_data_type: Union[str, torch.dtype] = "float16",
        kv_data_type: Optional[Union[str, torch.dtype]] = None,
        o_data_type: Optional[Union[str, torch.dtype]] = None,
        non_blocking: bool = True,
        prefix_len_ptr: Optional[torch.Tensor] = None,
        token_pos_in_items_ptr: Optional[torch.Tensor] = None,
        token_pos_in_items_len: int = 0,
        max_item_len_ptr: Optional[torch.Tensor] = None,
        seq_lens: Optional[torch.Tensor] = None,
        seq_lens_q: Optional[torch.Tensor] = None,
        block_tables: Optional[torch.Tensor] = None,
        max_token_per_sequence: Optional[int] = None,
        max_sequence_kv: Optional[int] = None,
        fixed_split_size: Optional[int] = None,
        disable_split_kv: bool = False,
    ) -> None:
        r"""Plan batch prefill/append attention on Paged KV-Cache for given problem specification.

        Parameters
        ----------
        qo_indptr : torch.Tensor
            The indptr of the query/output tensor, shape: ``[batch_size + 1]``.
        paged_kv_indptr : torch.Tensor
            The indptr of the paged kv-cache, shape: ``[batch_size + 1]``.
        paged_kv_indices : torch.Tensor
            The page indices of the paged kv-cache, shape: ``[paged_kv_indptr[-1]]``.
        paged_kv_last_page_len : torch.Tensor
            The number of entries in the last page of each request in the paged
            kv-cache, shape: ``[batch_size]``.
        num_qo_heads : int
            The number of query/output heads.
        num_kv_heads : int
            The number of key/value heads.
        head_dim_qk : int
            The dimension of the query/key heads.
        page_size : int
            The size of each page in the paged kv-cache.
        head_dim_vo : Optional[int]
            The dimension of the value/output heads, if not provided, will be set to
            ``head_dim_qk``.
        custom_mask : Optional[torch.Tensor]
            The flattened boolean mask tensor, shape: ``(sum(q_len[i] * k_len[i] for i in range(batch_size))``.
            The elements in the mask tensor should be either ``True`` or ``False``,
            where ``False`` means the corresponding element in the attention matrix will be
            masked out.

            Please refer to the :ref:`mask layout <mask-layout>` for more details about flattened
            layout of mask tensor.

            When :attr:`custom_mask` is provided, and :attr:`packed_custom_mask` is not, the
            function will pack the custom mask tensor into a 1D packed mask tensor, which introduces
            additional overhead.
        packed_custom_mask : Optional[torch.Tensor]
            The 1D packed uint8 mask tensor, if provided, the :attr:`custom_mask` will be ignored.
            The packed mask tensor is generated by :func:`flashinfer.quantization.packbits`.
        causal : bool
            Whether to apply causal mask to the attention matrix.
            This is only effective when :attr:`custom_mask` is not provided in
            :meth:`plan`.
        pos_encoding_mode : str
            The position encoding applied inside attention kernels, could be
            ``NONE``/``ROPE_LLAMA`` (LLAMA style rotary embedding) /``ALIBI``.
            Default is ``NONE``.
        use_fp16_qk_reduction : bool
            Whether to use f16 for qk reduction (faster at the cost of slight precision
            loss).
        window_left : int
            The left (inclusive) window size for the attention window, when set to ``-1``, the window
            size will be set to the full length of the sequence. Defaults to ``-1``.
        logits_soft_cap : Optional[float]
            The attention logits soft capping value (used in Gemini, Grok and Gemma-2, etc.), if not
            provided, will be set to ``0``. If greater than 0, the logits will be capped according to
            formula:
            :math:`\texttt{logits_soft_cap} \times \mathrm{tanh}(x / \texttt{logits_soft_cap})`,
            where :math:`x` is the input logits.
        sm_scale : Optional[float]
            The scale used in softmax, if not provided, will be set to
            ``1.0 / sqrt(head_dim)``.
        rope_scale : Optional[float]
            The scale used in RoPE interpolation, if not provided, will be set to
            ``1.0``.
        rope_theta : Optional[float]
            The theta used in RoPE, if not provided, will be set to ``1e4``.
        q_data_type : Union[str, torch.dtype]
            The data type of the query tensor, defaults torch.float16.
        kv_data_type : Optional[Union[str, torch.dtype]]
            The data type of the key/value tensor. If None, will be set to :attr:`q_data_type`.
        o_data_type : Optional[Union[str, torch.dtype]]
            The data type of the output tensor. If None, will be set to :attr:`q_data_type`.
            For FP8 inputs, this should typically be set to torch.float16 or torch.bfloat16.
        non_blocking : bool
            Whether to copy the input tensors to the device asynchronously, defaults to ``True``.
        prefix_len_ptr :Optional[torch.Tensor]
            prefix length. A uint32 1D tensor indicating the prefix length of each prompt. The tensor size is equal to the batch size.
        token_pos_in_items_ptr : Optional[torch.Tensor]
            A uint16 1D tensor (it will be converted to uint16 in flashinfer) indicating the token position of each item and started from 0 (delimiter)
            for each item. E.g., if we have 3 items of length 3, 2, 4 respectively for this member. This vector will be looking like
            `[0, 1, 2, 3, 0, 1, 2, 0, 1, 2, 3, 4, 0]` with 4 delimiters indexed as 0. For batch size > 1,
            we will concat them as 1D with zero paddings to make sure each has the same length, the padding length is defined by
            `token_pos_in_items_len` - length of the raw `token_pos_in_items_ptr` for each prompt.
        token_pos_in_items_len : int
            zero padding length for `token_pos_in_items_ptr` to better handle the bsz > 1 case. Still using the above 3,2,4 example.
            If we set `token_pos_in_items_len` to be 20, it will be  `[0, 1, 2, 3, 0, 1, 2, 0, 1, 2, 3, 4, 0, 0, 0, 0, 0, 0, 0, 0]`
            with 7 padded zeros. (note there're 8 zeros in the end where the first one is the delimiter token 0 in the end of the prompt)
        max_item_len_ptr : Optional[torch.Tensor]
            a uint16 vector contains the max token length of all items for each prompt
        seq_lens: Optional[torch.Tensor]
            A uint32 1D tensor indicating the kv sequence length of each prompt. shape: ``[batch_size]``.
        seq_lens_q: Optional[torch.Tensor]
            A uint32 1D tensor indicating the q sequence length of each prompt. shape: ``[batch_size]``.
            If not provided, will be set to the same value as ``seq_lens``.
        block_tables: Optional[torch.Tensor]
            A uint32 2D tensor indicating the block table of each prompt. shape: ``[batch_size, max_num_blocks_per_seq]``.
        max_token_per_sequence: Optional[int],
            Required for cudnn backend. This is the scalar max token length of each sequence.
        max_sequence_kv: Optional[int],
            Required for cudnn backend. This is the scalar max sequence length of each sequence in kv cache.
        fixed_split_size : Optional[int],
            The fixed split size for FA2 split-kv prefill/decode in pages. Recommend setting to the average sequence length of your workload.
            When enabled, will lead to deterministic softmax score reduction in the merge_states kernel, and therefore
            batch-size invariant outputs. See https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/
            Note that compatibility with CUDA graph is NOT guaranteed, as even when bs is fixed, kv seq len can change
            and lead to a varied number of launched CTAs.
        disable_split_kv : bool,
            Whether to disable the split-kv for determinism in CUDA Graph, defaults to ``False``.
        Note
        ----
        The :meth:`plan` method should be called before any :meth:`run` or
        :meth:`run_return_lse` calls, auxiliary data structures will be created
        during this call and cached for multiple kernel runs.

        The ``num_qo_heads`` must be a multiple of ``num_kv_heads``. If ``num_qo_heads``
        is not equal to ``num_kv_heads``, the function will use
        `grouped query attention <https://arxiv.org/abs/2305.13245>`_.

        The :meth:`plan` method cannot be used in Cuda Graph or in ``torch.compile``.
        """
        _check_workspace_buffer_alignment(
            self._float_workspace_buffer, "float_workspace_buffer"
        )
        _check_workspace_buffer_alignment(
            self._int_workspace_buffer, "int_workspace_buffer"
        )
        q_data_type = canonicalize_torch_dtype(q_data_type)
        if kv_data_type is None:
            kv_data_type = q_data_type
        kv_data_type = canonicalize_torch_dtype(kv_data_type)
        if o_data_type is None:
            o_data_type = q_data_type
        o_data_type = canonicalize_torch_dtype(o_data_type)

        if logits_soft_cap is None:
            logits_soft_cap = 0.0
        if head_dim_vo is None:
            head_dim_vo = head_dim_qk
        if fixed_split_size is None:
            fixed_split_size = -1

        batch_size = len(qo_indptr) - 1
        self._batch_size = batch_size
        self._num_qo_heads = num_qo_heads
        self._num_kv_heads = num_kv_heads
        if custom_mask is not None or packed_custom_mask is not None:
            mask_indptr = _compute_page_mask_indptr(
                qo_indptr,
                paged_kv_indptr,
                paged_kv_last_page_len,
                page_size,
            )
        if packed_custom_mask is None and custom_mask is not None:
            # create packed custom mask from custom mask
            packed_custom_mask, mask_indptr = segment_packbits(
                custom_mask.contiguous().view(-1),
                mask_indptr,
                bitorder="little",
            )

        self._prefix_len_ptr = prefix_len_ptr
        self._token_pos_in_items_ptr = token_pos_in_items_ptr
        self._token_pos_in_items_len = token_pos_in_items_len
        self._max_item_len_ptr = max_item_len_ptr

        # NOTE(Zihao): only required if qo_indptr/paged_kv_indptr are device tensors
        qo_indptr_host = qo_indptr.to("cpu")
        self._qo_indptr_last = int(qo_indptr_host[-1])
        total_num_rows = self._qo_indptr_last
        if max_token_per_sequence is not None:
            self._max_q_len = max_token_per_sequence
        else:
            self._max_q_len = max(qo_indptr_host[1:] - qo_indptr_host[:-1]).item()

        if max_sequence_kv is not None:
            self._max_kv_len = max_sequence_kv
        else:
            paged_kv_indptr_host = paged_kv_indptr.to("cpu")
            paged_kv_last_page_len_host = paged_kv_last_page_len.to("cpu")
            if seq_lens is None:
                kv_lens_arr_host = get_seq_lens(
                    paged_kv_indptr_host, paged_kv_last_page_len_host, page_size
                )
            else:
                kv_lens_arr_host = seq_lens.cpu().flatten()
            required_size = len(kv_lens_arr_host)
            if required_size > self._kv_lens_buffer.shape[0]:
                self._kv_lens_buffer = torch.empty(
                    (required_size,), dtype=torch.int32, device=self.device
                )
            self._kv_lens_buffer[:required_size].copy_(
                kv_lens_arr_host, non_blocking=non_blocking
            )
            self._max_kv_len = max(kv_lens_arr_host).item()

        if self.is_cuda_graph_enabled:
            if self._max_total_num_rows is None:
                self._max_total_num_rows = total_num_rows
            elif total_num_rows > self._max_total_num_rows:
                raise ValueError(
                    "The total number of rows in qo_indptr {} in cuda graph mode cannot "
                    "exceed the number of rows set during initialization {}.".format(
                        total_num_rows, self._max_total_num_rows
                    )
                )

            if batch_size != self._fixed_batch_size:
                raise ValueError(
                    "The batch size should be fixed during the lifecycle of the wrapper in "
                    "cuda graph mode, the runtime batch size {} mismatches the batch size {} "
                    " set during initialization.".format(
                        batch_size, self._fixed_batch_size
                    )
                )
            if len(paged_kv_indices) > len(self._paged_kv_indices_buf):
                raise ValueError(
                    "The length of paged_kv_indices exceeds the allocated buffer size."
                )

            self._qo_indptr_buf.copy_(qo_indptr, non_blocking=non_blocking)
            self._paged_kv_indptr_buf.copy_(paged_kv_indptr, non_blocking=non_blocking)
            self._paged_kv_last_page_len_buf.copy_(
                paged_kv_last_page_len, non_blocking=non_blocking
            )
            self._paged_kv_indices_buf[: len(paged_kv_indices)].copy_(
                paged_kv_indices,
                non_blocking=(paged_kv_indices.device == self.device) and non_blocking,
            )

            if packed_custom_mask is not None:
                if not torch.is_tensor(self._custom_mask_buf):
                    raise ValueError(
                        "custom_mask_buf must be initialized with a torch.Tensor in cuda graph mode if we use custom mask in attention computation."
                    )
                if not torch.is_tensor(self._mask_indptr_buf):
                    raise ValueError(
                        "mask_indptr_buf must be initialized with a torch.Tensor in cuda graph mode if we use custom mask in attention computation."
                    )
                self._custom_mask_buf[: len(packed_custom_mask)].copy_(
                    packed_custom_mask,
                    non_blocking=(packed_custom_mask.device == self.device)
                    and non_blocking,
                )
                # NOTE(Zihao): mask_indptr has the same length as qo_indptr
                self._mask_indptr_buf.copy_(mask_indptr, non_blocking=non_blocking)
        else:
            self._qo_indptr_buf = qo_indptr.to(self.device, non_blocking=non_blocking)
            self._paged_kv_indptr_buf = paged_kv_indptr.to(
                self.device, non_blocking=non_blocking
            )
            self._paged_kv_indices_buf = paged_kv_indices.to(
                self.device, non_blocking=non_blocking
            )
            self._paged_kv_last_page_len_buf = paged_kv_last_page_len.to(
                self.device, non_blocking=non_blocking
            )
            if packed_custom_mask is not None:
                self._custom_mask_buf = packed_custom_mask.to(
                    self.device, non_blocking=non_blocking
                )
                self._mask_indptr_buf = mask_indptr.to(
                    self.device, non_blocking=non_blocking
                )
            else:
                self._custom_mask_buf = None
                self._mask_indptr_buf = None

        self._cached_q_data_type = q_data_type
        self._cached_kv_data_type = kv_data_type
        self._cached_o_data_type = o_data_type

        if self._jit_module is not None:
            self._cached_module = self._jit_module
        else:
            if self._backend == "auto":
                self._backend = determine_attention_backend(
                    self.device,
                    PosEncodingMode[pos_encoding_mode].value,
                    use_fp16_qk_reduction,
                    self._custom_mask_buf is not None,  # use_custom_mask
                    q_data_type,
                    kv_data_type,
                    head_dim_qk=head_dim_qk,
                    head_dim_vo=head_dim_vo,
                )
            if self._backend != "cudnn":
                get_module_args = (
                    q_data_type,
                    kv_data_type,
                    o_data_type,
                    paged_kv_indptr.dtype,
                    head_dim_qk,
                    head_dim_vo,
                    PosEncodingMode[pos_encoding_mode].value,
                    window_left >= 0,  # use_sliding_window
                    logits_soft_cap > 0,  # use_logits_soft_cap
                    use_fp16_qk_reduction,
                )

                self._cached_module = get_batch_prefill_module(
                    self._backend, *get_module_args
                )

        self._block_tables = block_tables
        if self._backend == "trtllm-gen":
            # Allocated once per plan and reused across run() launches; the
            # trtllm-gen kernel self-resets the counters after each launch.
            self._trtllm_gen_multi_ctas_kv_counter_buffer = (
                _get_trtllm_gen_multi_ctas_kv_counter_buffer(
                    batch_size,
                    num_qo_heads,
                    get_device_sm_count(self.device),
                    self.device,
                )
            )
            if logits_soft_cap != 0.0:
                raise ValueError(
                    "logits_soft_cap must be 0.0 for trtllm-gen paged KV cache"
                )
            if not causal and window_left >= 0:
                raise NotImplementedError(
                    "Sliding-window non-causal attention is not supported for trtllm-gen paged KV cache. "
                    "Use window_left=-1 for dense bidirectional attention."
                )
            if self._block_tables is None:
                blocks_per_seq = [
                    (seq_len + page_size - 1) // page_size
                    for seq_len in kv_lens_arr_host
                ]
                max_num_blocks_per_seq = max(blocks_per_seq)
                self._block_tables = torch.zeros(
                    (batch_size, max_num_blocks_per_seq),
                    dtype=torch.int,
                    device=self.device,
                )
                block_id = paged_kv_indptr_host[0]
                for i in range(batch_size):
                    num_blocks_needed = blocks_per_seq[i]
                    assert self._block_tables is not None, (
                        "block_tables is not initialized"
                    )
                    self._block_tables[i, :num_blocks_needed] = paged_kv_indices[
                        block_id : block_id + num_blocks_needed
                    ]
                    block_id += num_blocks_needed

        if self._cached_module is not None:
            args = [
                self._float_workspace_buffer,
                self._int_workspace_buffer,
                self._pin_memory_int_workspace_buffer,
                qo_indptr_host,
                paged_kv_indptr_host,
                kv_lens_arr_host,
                self._max_total_num_rows or total_num_rows,
                batch_size,
                num_qo_heads,
                num_kv_heads,
                page_size,
                self.is_cuda_graph_enabled,
                head_dim_qk,
                head_dim_vo,
                causal,
                window_left,
            ]
            if self._backend == "fa2":
                args.append(fixed_split_size or -1)  # fixed_split_size
                args.append(disable_split_kv)  # disable_split_kv
                args.append(0)  # num_colocated_ctas
                args.append(0)  # uniform_q_len
            self._plan_info = self._cached_module.plan(
                *args,
            )

        self._causal = causal
        self._pos_encoding_mode = pos_encoding_mode
        self._use_fp16_qk_reduction = use_fp16_qk_reduction
        self._window_left = window_left
        self._logits_soft_cap = logits_soft_cap
        self._sm_scale = sm_scale
        self._rope_scale = rope_scale
        self._rope_theta = rope_theta
        self._seq_lens_kv = seq_lens
        self._seq_lens_q = seq_lens_q if seq_lens_q is not None else seq_lens

    begin_forward = plan

    def forward(
        self,
        q: torch.Tensor,
        paged_kv_cache: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
        causal: bool = False,
        pos_encoding_mode: str = "NONE",
        use_fp16_qk_reduction: bool = False,
        k_scale: Optional[float] = None,
        v_scale: Optional[float] = None,
        window_left: int = -1,
        logits_soft_cap: Optional[float] = None,
        sm_scale: Optional[float] = None,
        rope_scale: Optional[float] = None,
        rope_theta: Optional[float] = None,
    ) -> torch.Tensor:
        r"""Warning: This function is deprecated, please use :meth:`run` instead."""
        self._causal = causal
        self._pos_encoding_mode = pos_encoding_mode
        self._use_fp16_qk_reduction = use_fp16_qk_reduction
        self._window_left = window_left
        self._logits_soft_cap = logits_soft_cap
        self._sm_scale = sm_scale
        self._rope_scale = rope_scale
        self._rope_theta = rope_theta
        return self.run(q, paged_kv_cache, k_scale=k_scale, v_scale=v_scale)

    @overload
    def run(
        self,
        q: torch.Tensor,
        paged_kv_cache: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
        *args,
        k_scale: Optional[float] = None,
        v_scale: Optional[float] = None,
        out: Optional[torch.Tensor] = None,
        lse: Optional[torch.Tensor] = None,
        return_lse: Literal[False] = False,
        enable_pdl: Optional[bool] = None,
        window_left: Optional[int] = None,
        sinks: Optional[torch.Tensor] = None,
        kv_cache_sf: Optional[
            Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
        ] = None,
        skip_softmax_threshold_scale_factor: Optional[float] = None,
    ) -> torch.Tensor: ...

    @overload
    def run(
        self,
        q: torch.Tensor,
        paged_kv_cache: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
        *args,
        k_scale: Optional[float] = None,
        v_scale: Optional[float] = None,
        out: Optional[torch.Tensor] = None,
        lse: Optional[torch.Tensor] = None,
        return_lse: Literal[True] = True,
        enable_pdl: Optional[bool] = None,
        window_left: Optional[int] = None,
        sinks: Optional[torch.Tensor] = None,
        kv_cache_sf: Optional[
            Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
        ] = None,
        skip_softmax_threshold_scale_factor: Optional[float] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]: ...

    @flashinfer_api(trace=gqa_paged_prefill_trace)
    def run(
        self,
        q: torch.Tensor,
        paged_kv_cache: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
        *args,
        q_scale: Optional[Union[float, torch.Tensor]] = None,
        k_scale: Optional[Union[float, torch.Tensor]] = None,
        v_scale: Optional[Union[float, torch.Tensor]] = None,
        out: Optional[torch.Tensor] = None,
        lse: Optional[torch.Tensor] = None,
        return_lse: bool = False,
        enable_pdl: Optional[bool] = None,
        window_left: Optional[int] = None,
        sinks: Optional[torch.Tensor] = None,
        kv_cache_sf: Optional[
            Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
        ] = None,
        skip_softmax_threshold_scale_factor: Optional[float] = None,
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
        r"""Compute batch prefill/append attention between query and paged kv-cache.

        Parameters
        ----------
        q : torch.Tensor
            The query tensor, shape: ``[qo_indptr[-1], num_qo_heads, head_dim]``
        paged_kv_cache : Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
            The paged KV-Cache stored as a tuple of tensors or a single tensor:

            * a tuple ``(k_cache, v_cache)`` of 4-D tensors, each with shape:
              ``[max_num_pages, page_size, num_kv_heads, head_dim]`` if :attr:`kv_layout` is ``NHD``,
              and ``[max_num_pages, num_kv_heads, page_size, head_dim]`` if :attr:`kv_layout` is ``HND``.

            * a single 5-D tensor with shape:
              ``[max_num_pages, 2, page_size, num_kv_heads, head_dim]`` if
              :attr:`kv_layout` is ``NHD``, and
              ``[max_num_pages, 2, num_kv_heads, page_size, head_dim]`` if
              :attr:`kv_layout` is ``HND``. Where ``paged_kv_cache[:, 0]`` is the key-cache and
              ``paged_kv_cache[:, 1]`` is the value-cache.

        *args
            Additional arguments for custom kernels.
        q_scale : Optional[Union[float, torch.Tensor]]
            The calibration scale of query for fp8 input, if not provided, will be set to ``1.0``.
        k_scale : Optional[Union[float, torch.Tensor]]
            The calibration scale of key for fp8 or nvfp4 input, if not provided, will be set to ``1.0``.
        v_scale : Optional[Union[float, torch.Tensor]]
            The calibration scale of value for fp8 or nvfp4 input, if not provided, will be set to ``1.0``.
        out : Optional[torch.Tensor]
            The output tensor, if not provided, will be allocated internally.
        lse : Optional[torch.Tensor]
            The log-sum-exp of attention logits, if not provided, will be allocated internally.
        return_lse : bool
            Whether to return the logsumexp of attention output
        enable_pdl : bool
            Whether to enable Programmatic Dependent Launch (PDL). See https://docs.nvidia.com/cuda/cuda-c-programming-guide/#programmatic-dependent-launch-and-synchronization
            Only supported for >= sm90, and currently only for FA2 and CUDA core decode.
        window_left : Optional[int]
            Per-call override for the left (inclusive) sliding-window size.  When
            ``None``, the value supplied to :meth:`plan` is used.  Pass ``-1`` to
            disable the sliding window for this call.
        sinks : Optional[torch.Tensor]
            Per-head attention-sink logits.  When provided, the kernel applies the
            attention-with-sink variant: an additional virtual token whose logit is
            ``sinks[head_idx]`` is appended to each row of the softmax denominator.
            Shape: ``[num_qo_heads]``, dtype ``float32``.
        kv_cache_sf : Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]]
            Per-block scale factors for NVFP4 KV cache. Accepts the same formats as
            ``paged_kv_cache``:

            * a tuple ``(k_scales, v_scales)`` of 4-D tensors, each with shape:
              ``[num_pages, page_size, num_kv_heads, head_dim // 16]`` if :attr:`kv_layout` is ``NHD``,
              and ``[num_pages, num_kv_heads, page_size, head_dim // 16]`` if :attr:`kv_layout` is ``HND``.
            * a single 5-D tensor with shape:
              ``[num_pages, 2, page_size, num_kv_heads, head_dim // 16]`` if :attr:`kv_layout` is ``NHD``,
              and ``[num_pages, 2, num_kv_heads, page_size, head_dim // 16]`` if :attr:`kv_layout` is ``HND``,
              where dim 1 holds k (index 0) and v (index 1) scales.

            Both tensors have dtype ``torch.float8_e4m3fn``. ``k_scales`` uses a linear
            (row-major) layout, while ``v_scales`` must use TRT-LLM's 4-token interleaved
            layout within each ``[page_size, head_dim // 16]`` tile if backend is `trtllm-gen`. Use
            :func:`flashinfer.fp4_quantization.nvfp4_quantize_paged_kv_cache` to produce
            correctly formatted scale factors.

            For the trtllm-gen backend with ``NHD`` layout, scale tensors are transposed
            to HND internally (incurring a copy). Use ``HND`` for better performance.

            Currently, NVFP4 KV supports `fa2` and `trtllm-gen` backend.
        skip_softmax_threshold_scale_factor : Optional[float]
            Threshold scale factor for skipping softmax operations.  Providing a
            value enables skip-softmax sparsity as described in
            https://arxiv.org/abs/2512.12087.  Defaults to ``None`` (standard
            attention).  Higher values yield faster kernels at the cost of accuracy;
            the effective threshold equals the supplied factor divided by the
            context length.
        Returns
        -------
        Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
            If :attr:`return_lse` is ``False``, the attention output, shape: ``[qo_indptr[-1], num_qo_heads, head_dim]``.
            If :attr:`return_lse` is ``True``, a tuple of two tensors:

            * The attention output, shape: ``[qo_indptr[-1], num_qo_heads, head_dim]``.
            * The logsumexp of attention output, shape: ``[qo_indptr[-1], num_qo_heads]``.
        """
        if enable_pdl is None:
            enable_pdl = device_support_pdl(q.device)
        k_cache, v_cache = _unpack_paged_kv_cache(paged_kv_cache, self._kv_layout)
        _check_cached_qkv_data_type(
            q, k_cache, self._cached_q_data_type, self._cached_kv_data_type
        )
        # Validate q shape matches qo_indptr (using value cached in plan() to avoid GPU sync)
        if self._backend == "cudnn":
            if q.numel() != self._qo_indptr_last:
                raise ValueError(
                    f"q.numel() ({q.numel()}) does not match qo_indptr[-1] ({self._qo_indptr_last}). "
                    f"For cudnn paged prefill, qo_indptr uses element offsets "
                    f"(total_tokens * num_heads * head_dim)."
                )
        else:
            if q.size(0) != self._qo_indptr_last:
                raise ValueError(
                    f"q.shape[0] ({q.size(0)}) does not match qo_indptr[-1] ({self._qo_indptr_last}). "
                    f"For paged prefill, q must have shape [total_tokens, num_heads, head_dim] "
                    f"where total_tokens = qo_indptr[-1]."
                )

        if (
            k_cache.dtype == torch.uint8 or v_cache.dtype == torch.uint8
        ) and kv_cache_sf is None:
            raise ValueError("kv_cache_sf must be provided for NVFP4 KV cache.")
        key_block_scales, value_block_scales = (
            _unpack_paged_kv_cache(kv_cache_sf, self._kv_layout)
            if kv_cache_sf is not None
            else (None, None)
        )

        o_dtype = self._cached_o_data_type
        if out is not None and out.dtype != o_dtype:
            raise ValueError(
                f"The dtype of out {out.dtype} does not match the o_data_type {o_dtype} specified in plan function."
            )

        if self._kv_layout == "NHD":
            page_size = k_cache.shape[1]
        else:
            page_size = k_cache.shape[2]
        window_left = self._window_left if window_left is None else window_left
        if self._backend != "trtllm-gen":
            # NOTE(Siyuan): since window_left is appeared in the plan function, we need to make sure it is the same as the one in the plan function.
            # Remove this check if the backend supports dynamic window_left.
            assert window_left == self._window_left
        logits_soft_cap = self._logits_soft_cap
        sm_scale = self._sm_scale
        rope_scale = self._rope_scale
        rope_theta = self._rope_theta
        if logits_soft_cap is None:
            logits_soft_cap = 0.0
        if sm_scale is None:
            sm_scale = 1.0 / math.sqrt(q.size(-1))
        if self._backend != "cudnn":
            if q_scale is not None:
                sm_scale *= q_scale
            if k_scale is not None:
                sm_scale *= k_scale
        if rope_scale is None:
            rope_scale = 1.0
        if rope_theta is None:
            rope_theta = 1e4
        if return_lse:
            if lse is None:
                lse = torch.empty(
                    (q.size(0), q.size(1)), dtype=torch.float32, device=q.device
                )
            else:
                check_shape_dtype_device(
                    lse, (q.size(0), q.size(1)), torch.float32, q.device, "lse"
                )

        # For NVFP4 KV (uint8 packed), v_cache last dim is head_dim//2;
        # use q's head_dim for output instead
        out_head_dim = q.shape[-1] if kv_cache_sf is not None else v_cache.shape[-1]
        if out is None:
            # Use cached output data type if available (for FP8 attention with FP16 output)
            out_dtype = getattr(self, "_cached_o_data_type", None) or q.dtype
            out = torch.zeros(
                q.shape[:-1] + (out_head_dim,), dtype=out_dtype, device=q.device
            )
        else:
            out_dtype = getattr(self, "_cached_o_data_type", None) or q.dtype
            check_shape_dtype_device(
                out, q.shape[:-1] + (out_head_dim,), out_dtype, q.device, "out"
            )

        # Convert NHD layout to HND for trtllm-gen backend
        if self._backend == "trtllm-gen" and self._kv_layout == "NHD":
            k_cache = k_cache.transpose(-3, -2)
            v_cache = v_cache.transpose(-3, -2)
            if key_block_scales is not None:
                print(
                    "[WARNING] NVFP4 KV cache with NHD layout will be converted to HND, "
                    "incurring extra transpose and contiguous copy overhead. "
                    "Use kv_layout='HND' for better performance."
                )
                k_cache = k_cache.contiguous()
                v_cache = v_cache.contiguous()
                key_block_scales = key_block_scales.transpose(-3, -2).contiguous()
                value_block_scales = value_block_scales.transpose(-3, -2).contiguous()

        if self._custom_mask_buf is not None:
            mask_mode = MaskMode.CUSTOM.value
        else:
            if self._causal:
                mask_mode = MaskMode.CAUSAL.value
            else:
                mask_mode = MaskMode.NON_CAUSAL.value

        if self._prefix_len_ptr is not None:
            mask_mode = MaskMode.MULTIITEMSCORING.value

        if self._backend == "cudnn":
            if self._seq_lens_q is not None and self._seq_lens_q.dim() == 1:
                self._seq_lens_q = self._seq_lens_q.reshape(self._batch_size, 1, 1, 1)

            if self._seq_lens_kv is not None and self._seq_lens_kv.dim() == 1:
                self._seq_lens_kv = self._seq_lens_kv.reshape(self._batch_size, 1, 1, 1)

            cudnn_batch_prefill_with_kv_cache(
                q,
                k_cache,  # Need to be changed
                v_cache,  # Need to be changed
                self._sm_scale,
                self._float_workspace_buffer,
                actual_seq_lens_q=self._seq_lens_q,
                actual_seq_lens_kv=self._seq_lens_kv,
                max_token_per_sequence=self._max_q_len,
                max_sequence_kv=self._max_kv_len,
                block_tables=self._block_tables,
                causal=self._causal,
                return_lse=return_lse,
                q_scale=q_scale,
                k_scale=k_scale,
                v_scale=v_scale,
                batch_offsets_q=self._qo_indptr_buf,
                batch_offsets_o=self._qo_indptr_buf,
                out=out,
                lse=lse,
                o_data_type=out_dtype,
            )
        else:
            if self._backend != "trtllm-gen":
                assert self._plan_info is not None, "plan info is not initialized"
            run_args = [
                self._float_workspace_buffer,
                self._int_workspace_buffer,
                self._plan_info,
                q,
                k_cache,
                v_cache,
                self._qo_indptr_buf,
                self._paged_kv_indptr_buf,
                self._paged_kv_indices_buf,
                self._paged_kv_last_page_len_buf,
                out,
                lse,
                mask_mode,
                TensorLayout[self._kv_layout].value,
                window_left,
                enable_pdl,
            ]
            if self._jit_module is not None:
                run_args.extend(
                    prepare_jit_additional_args(
                        self._jit_additional_tensor_names,
                        {
                            "maybe_custom_mask": self._custom_mask_buf,
                            "maybe_mask_indptr": self._mask_indptr_buf,
                            "maybe_alibi_slopes": lambda: _get_cache_alibi_slopes_buf(
                                q.shape[1], q.device
                            ),
                        },
                        args,
                    )
                )
            else:
                # Extract FP8 scale tensors from *args if q is FP8
                fp8_scale_q = None
                fp8_scale_k = None
                fp8_scale_v = None
                if is_float8(q) and len(args) >= 3:
                    fp8_scale_q = args[0]
                    fp8_scale_k = args[1]
                    fp8_scale_v = args[2]
                run_args += [
                    self._custom_mask_buf,
                    self._mask_indptr_buf,
                    _get_cache_alibi_slopes_buf(q.shape[1], q.device),
                    self._prefix_len_ptr,
                    self._token_pos_in_items_ptr,
                    self._max_item_len_ptr,
                    logits_soft_cap,
                    sm_scale,
                    fp8_scale_q,
                    fp8_scale_k,
                    fp8_scale_v,
                    rope_scale,
                    rope_theta,
                    self._token_pos_in_items_len,
                    self._workspace_size,
                    self._num_qo_heads,
                    self._num_kv_heads,
                    self._block_tables,
                    self._kv_lens_buffer,
                    page_size,
                    self._max_q_len,
                    self._max_kv_len,
                    self._batch_size,
                    self._qo_indptr_buf,
                    self._paged_kv_indptr_buf,
                    sinks,
                    key_block_scales,
                    value_block_scales,
                    skip_softmax_threshold_scale_factor,
                    True,  # uses_shared_paged_kv_idx
                    self._trtllm_gen_multi_ctas_kv_counter_buffer,
                ]

            if self._backend == "trtllm-gen":
                assert self._trtllm_gen_multi_ctas_kv_counter_buffer is not None
            assert self._cached_module is not None, "cached module is not initialized"
            self._cached_module.paged_run(*run_args)

            is_float_one = isinstance(v_scale, float) and v_scale == 1.0
            if v_scale is not None and not is_float_one:
                # TODO(Zihao): fused into kernel
                if is_float8(out):
                    out = (out.to(torch.float32) * v_scale).to(out.dtype)
                else:
                    out *= v_scale

        return (out, lse) if return_lse else out

    run_return_lse = functools.partialmethod(run, return_lse=True)

    def forward_return_lse(
        self,
        q: torch.Tensor,
        paged_kv_cache: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
        causal: bool = False,
        pos_encoding_mode: str = "NONE",
        use_fp16_qk_reduction: bool = False,
        k_scale: Optional[float] = None,
        v_scale: Optional[float] = None,
        window_left: int = -1,
        logits_soft_cap: Optional[float] = None,
        sm_scale: Optional[float] = None,
        rope_scale: Optional[float] = None,
        rope_theta: Optional[float] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        r"""Warning: This function is deprecated, please use :meth:`run_return_lse` instead."""
        self._causal = causal
        self._pos_encoding_mode = pos_encoding_mode
        self._use_fp16_qk_reduction = use_fp16_qk_reduction
        self._window_left = window_left
        self._logits_soft_cap = logits_soft_cap
        self._sm_scale = sm_scale
        self._rope_scale = rope_scale
        self._rope_theta = rope_theta
        return self.run_return_lse(q, paged_kv_cache, k_scale=k_scale, v_scale=v_scale)

    def end_forward(self) -> None:
        r"""Warning: this function is deprecated and has no effect."""
        pass


def _compute_mask_indptr(
    qo_indptr: torch.Tensor, kv_indptr: torch.Tensor
) -> torch.Tensor:
    if len(qo_indptr) != len(kv_indptr):
        raise ValueError("The length of qo_indptr and kv_indptr should be the same.")
    mask_indptr = torch.empty_like(qo_indptr)
    mask_indptr[0] = 0
    mask_indptr[1:] = torch.cumsum(
        (qo_indptr[1:] - qo_indptr[:-1]) * (kv_indptr[1:] - kv_indptr[:-1]),
        0,
    )
    return mask_indptr


class BatchPrefillWithRaggedKVCacheWrapper:
    r"""Wrapper class for prefill/append attention with ragged (tensor) kv-cache for
    batch of requests.

    Check :ref:`our tutorial <kv-layout>` for ragged kv-cache layout.

    Example
    -------
    >>> import torch
    >>> import flashinfer
    >>> num_layers = 32
    >>> num_qo_heads = 64
    >>> num_kv_heads = 16
    >>> head_dim = 128
    >>> # allocate 128MB workspace buffer
    >>> workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device="cuda:0")
    >>> prefill_wrapper = flashinfer.BatchPrefillWithRaggedKVCacheWrapper(
    ...     workspace_buffer, "NHD"
    ... )
    >>> batch_size = 7
    >>> nnz_kv = 100
    >>> nnz_qo = 100
    >>> qo_indptr = torch.tensor(
    ...     [0, 33, 44, 55, 66, 77, 88, nnz_qo], dtype=torch.int32, device="cuda:0"
    ... )
    >>> kv_indptr = qo_indptr.clone()
    >>> q_at_layer = torch.randn(num_layers, nnz_qo, num_qo_heads, head_dim).half().to("cuda:0")
    >>> k_at_layer = torch.randn(num_layers, nnz_kv, num_kv_heads, head_dim).half().to("cuda:0")
    >>> v_at_layer = torch.randn(num_layers, nnz_kv, num_kv_heads, head_dim).half().to("cuda:0")
    >>> # create auxiliary data structures for batch prefill attention
    >>> prefill_wrapper.plan(
    ...     qo_indptr,
    ...     kv_indptr,
    ...     num_qo_heads,
    ...     num_kv_heads,
    ...     head_dim,
    ...     causal=True,
    ... )
    >>> outputs = []
    >>> for i in range(num_layers):
    ...     q = q_at_layer[i]
    ...     k = k_at_layer[i]
    ...     v = v_at_layer[i]
    ...     # compute batch prefill attention, reuse auxiliary data structures
    ...     o = prefill_wrapper.run(q, k, v)
    ...     outputs.append(o)
    ...
    >>> outputs[0].shape
    torch.Size([100, 64, 128])
    >>>
    >>> # below is another example of creating custom mask for batch prefill attention
    >>> mask_arr = []
    >>> qo_len = (qo_indptr[1:] - qo_indptr[:-1]).cpu().tolist()
    >>> kv_len = (kv_indptr[1:] - kv_indptr[:-1]).cpu().tolist()
    >>> for i in range(batch_size):
    ...     mask_i = torch.tril(
    ...         torch.full((qo_len[i], kv_len[i]), True, device="cuda:0"),
    ...         diagonal=(kv_len[i] - qo_len[i]),
    ...     )
    ...     mask_arr.append(mask_i.flatten())
    ...
    >>> mask = torch.cat(mask_arr, dim=0)
    >>> prefill_wrapper.plan(
    ...     qo_indptr,
    ...     kv_indptr,
    ...     num_qo_heads,
    ...     num_kv_heads,
    ...     head_dim,
    ...     custom_mask=mask
    ... )
    >>> outputs_custom_mask = []
    >>> for i in range(num_layers):
    ...     q = q_at_layer[i]
    ...     k = k_at_layer[i]
    ...     v = v_at_layer[i]
    ...     # compute batch prefill attention, reuse auxiliary data structures
    ...     o_custom = prefill_wrapper.run(q, k, v)
    ...     assert torch.allclose(o_custom, outputs[i], rtol=1e-3, atol=1e-3)
    ...
    >>> outputs_custom_mask[0].shape
    torch.Size([100, 64, 128])


    Note
    ----
    To accelerate computation, FlashInfer's batch prefill/append attention operators
    create some auxiliary data structures, these data structures can be reused across
    multiple prefill/append attention calls (e.g. different Transformer layers). This
    wrapper class manages the lifecycle of these data structures.
    """

    @flashinfer_api
    def __init__(
        self,
        float_workspace_buffer: torch.Tensor,
        kv_layout: str = "NHD",
        use_cuda_graph: bool = False,
        qo_indptr_buf: Optional[torch.Tensor] = None,
        kv_indptr_buf: Optional[torch.Tensor] = None,
        custom_mask_buf: Optional[torch.Tensor] = None,
        mask_indptr_buf: Optional[torch.Tensor] = None,
        backend: str = "auto",
        jit_args: Optional[List[Any]] = None,
        jit_kwargs: Optional[Dict[str, Any]] = None,
    ) -> None:
        r"""Constructor of :class:`BatchPrefillWithRaggedKVCacheWrapper`.

        Parameters
        ----------
        float_workspace_buffer : torch.Tensor
            The user reserved float workspace buffer used to store intermediate attention results
            in the split-k algorithm. The recommended size is 128MB, the device of the workspace
            buffer should be the same as the device of the input tensors. The buffer must be
            16-byte aligned; tensors created by ``torch.empty`` satisfy this on supported devices.

        kv_layout : str
            The layout of the input k/v tensors, could be either ``NHD`` or ``HND``.

        use_cuda_graph : bool
            Whether to enable CUDA graph capture for the prefill kernels, if enabled, the
            auxiliary data structures will be stored as the provided buffers.

        qo_indptr_buf : Optional[torch.Tensor]
            The user reserved GPU buffer to store the ``qo_indptr`` array, the size of the buffer
            should be ``[batch_size + 1]``.
            This argument is only effective when ``use_cuda_graph`` is ``True``.

        kv_indptr_buf : Optional[torch.Tensor]
            The user reserved GPU buffer to store the ``kv_indptr`` array, the size of the buffer
            should be ``[batch_size + 1]``.
            This argument is only effective when ``use_cuda_graph`` is ``True``.

        custom_mask_buf : Optional[torch.Tensor]
            The user reserved GPU buffer to store the custom mask tensor, should be large
            enough to store the maximum possible size of the packed custom mask tensor during the
            lifetime of the wrapper. This argument is only effective when ``use_cuda_graph``
            is ``True`` and custom mask will be used in attention computation.

        mask_indptr_buf : Optional[torch.Tensor]
            The user reserved GPU buffer to store the ``mask_indptr`` array, the size of the buffer
            should be ``[batch_size]``.
            This argument is only effective when ``use_cuda_graph`` is ``True`` and custom mask
            will be used in attention computation.

        backend : str
            The implementation backend, could be ``auto``/``fa2``/``fa3``/``cudnn``/``cutlass``
            or ``cute-dsl``.
            Defaults to ``auto``.
            If set to ``auto``, the wrapper will automatically choose the backend based on the
            device architecture and kernel availability.
            The ``cute-dsl`` backend uses the CuTe DSL attention kernel for Blackwell (SM100+).

        jit_args : Optional[List[Any]]
            If provided, the wrapper will use the provided arguments to create the JIT module,
            otherwise, the wrapper will use default attention implementation.

        jit_kwargs : Optional[Dict[str, Any]]
            The keyword arguments to create the JIT module, defaults to None.
        """
        _check_kv_layout(kv_layout)
        if jit_args is not None and backend != "cute-dsl":
            if jit_kwargs is None:
                jit_kwargs = {}
            self._jit_module = get_batch_prefill_jit_module(
                jit_args[0],
                get_customize_batch_prefill_module(backend, *jit_args, **jit_kwargs),
            )
            # jit_args[7] is additional_tensor_names from gen_customize_batch_prefill_module
            self._jit_additional_tensor_names = list(jit_args[7])
        else:
            self._jit_module = None
            self._jit_additional_tensor_names = []

        self._cute_dsl_wrapper = None
        if backend == "cute-dsl":
            from .cute_dsl.attention import BatchPrefillCuteDSLWrapper

            self._cute_dsl_wrapper = BatchPrefillCuteDSLWrapper(
                float_workspace_buffer, use_cuda_graph=use_cuda_graph
            )

        self._kv_layout = kv_layout
        self._float_workspace_buffer = float_workspace_buffer
        self.device = float_workspace_buffer.device
        self._int_workspace_buffer = torch.empty(
            (8 * 1024 * 1024,), dtype=torch.uint8, device=self.device
        )
        self._pin_memory_int_workspace_buffer = torch.empty(
            self._int_workspace_buffer.shape,
            dtype=torch.uint8,
            pin_memory=True,
            device="cpu",
        )
        self._use_cuda_graph = use_cuda_graph
        if use_cuda_graph:
            if not torch.is_tensor(qo_indptr_buf):
                raise ValueError(
                    "qo_indptr_buf should be a torch.Tensor in cuda graph mode"
                )
            if not torch.is_tensor(kv_indptr_buf):
                raise ValueError(
                    "kv_indptr_buf should be a torch.Tensor in cuda graph mode"
                )
            self._fixed_batch_size = len(qo_indptr_buf) - 1
            if len(kv_indptr_buf) != self._fixed_batch_size + 1:
                raise ValueError(
                    "The length of kv_indptr_buf ({}) should be the same as qo_indptr_buf ({}).".format(
                        len(kv_indptr_buf), self._fixed_batch_size
                    )
                )
            # NOTE(Zihao): do not check custom_mask_buf and mask_indptr_buf here,
            # as they may not be used.

        self._qo_indptr_buf = qo_indptr_buf
        self._kv_indptr_buf = kv_indptr_buf
        self._custom_mask_buf = custom_mask_buf
        self._mask_indptr_buf = mask_indptr_buf
        self._max_total_num_rows: Optional[int] = None
        self._backend = backend
        self._cached_module = None

    @property
    def is_cuda_graph_enabled(self) -> bool:
        return self._use_cuda_graph

    def reset_workspace_buffer(
        self, float_workspace_buffer: torch.Tensor, int_workspace_buffer
    ) -> None:
        r"""Reset the workspace buffer.

        Parameters
        ----------
        float_workspace_buffer : torch.Tensor
            The new float workspace buffer, the device of the new float workspace buffer should
            be the same as the device of the input tensors.

        int_workspace_buffer : torch.Tensor
            The new int workspace buffer, the device of the new int workspace buffer should
            be the same as the device of the input tensors.
        """
        self._float_workspace_buffer = float_workspace_buffer
        self._int_workspace_buffer = int_workspace_buffer
        self._pin_memory_int_workspace_buffer = torch.empty(
            self._int_workspace_buffer.shape,
            dtype=self._int_workspace_buffer.dtype,
            device="cpu",
            pin_memory=True,
        )

    @flashinfer_api
    def plan(
        self,
        qo_indptr: torch.Tensor,
        kv_indptr: torch.Tensor,
        num_qo_heads: int,
        num_kv_heads: int,
        head_dim_qk: int,
        head_dim_vo: Optional[int] = None,
        custom_mask: Optional[torch.Tensor] = None,
        packed_custom_mask: Optional[torch.Tensor] = None,
        causal: bool = False,
        pos_encoding_mode: str = "NONE",
        use_fp16_qk_reduction: bool = False,
        window_left: int = -1,
        logits_soft_cap: Optional[float] = None,
        sm_scale: Optional[float] = None,
        rope_scale: Optional[float] = None,
        rope_theta: Optional[float] = None,
        q_data_type: Union[str, torch.dtype] = "float16",
        kv_data_type: Optional[Union[str, torch.dtype]] = None,
        o_data_type: Optional[Union[str, torch.dtype]] = None,
        non_blocking: bool = True,
        prefix_len_ptr: Optional[torch.Tensor] = None,
        token_pos_in_items_ptr: Optional[torch.Tensor] = None,
        token_pos_in_items_len: int = 0,
        max_item_len_ptr: Optional[torch.Tensor] = None,
        fixed_split_size: Optional[int] = None,
        disable_split_kv: bool = False,
        seq_lens: Optional[torch.Tensor] = None,
        seq_lens_q: Optional[torch.Tensor] = None,
        max_token_per_sequence: Optional[int] = None,
        max_sequence_kv: Optional[int] = None,
        v_indptr: Optional[torch.Tensor] = None,
        o_indptr: Optional[torch.Tensor] = None,
    ) -> None:
        r"""Plan batch prefill/append attention on Ragged KV-Cache for given problem specification.

        Parameters
        ----------
        qo_indptr : torch.Tensor
            The indptr of the query/output tensor, shape: ``[batch_size + 1]``.
        kv_indptr : torch.Tensor
            The indptr of the key/value tensor, shape: ``[batch_size + 1]``.
        num_qo_heads : int
            The number of query/output heads.
        num_kv_heads : int
            The number of key/value heads.
        head_dim_qk : int
            The dimension of the heads on query/key tensor.
        head_dim_vo : Optional[int]
            The dimension of the heads on value/output tensor.
            If not provided, will be set to ``head_dim_qk``.
        custom_mask : Optional[torch.Tensor]
            The flattened boolean mask tensor, shape: ``(sum(q_len[i] * k_len[i] for i in range(batch_size))``.
            The elements in the mask tensor should be either ``True`` or ``False``,
            where ``False`` means the corresponding element in the attention matrix will be
            masked out.

            Please refer to the :ref:`mask layout <mask-layout>` for more details about flattened
            layout of mask tensor.

            When :attr:`custom_mask` is provided, and :attr:`packed_custom_mask` is not, the
            function will pack the custom mask tensor into a 1D packed mask tensor, which introduces
            additional overhead.
        packed_custom_mask : Optional[torch.Tensor]
            The 1D packed uint8 mask tensor, if provided, the :attr:`custom_mask` will be ignored.
            The packed mask tensor is generated by :func:`flashinfer.quantization.packbits`.

            If provided, the custom mask will be added to the attention matrix before softmax
            and after scaling. The mask tensor should be in the same device as the input tensors.
        causal : bool
            Whether to apply causal mask to the attention matrix.
            This argument is ignored if ``mask`` is provided in :meth:`plan`.
        pos_encoding_mode : str
            The position encoding applied inside attention kernels, could be
            ``NONE``/``ROPE_LLAMA`` (LLAMA style rotary embedding) /``ALIBI``.
            Default is ``NONE``.
        use_fp16_qk_reduction : bool
            Whether to use f16 for qk reduction (faster at the cost of slight precision
            loss).
        window_left : int
            The left (inclusive) window size for the attention window, when set to ``-1``, the window
            size will be set to the full length of the sequence. Defaults to ``-1``.
        logits_soft_cap : Optional[float]
            The attention logits soft capping value (used in Gemini, Grok and Gemma-2, etc.), if not
            provided, will be set to ``0``. If greater than 0, the logits will be capped according to
            formula:
            :math:`\texttt{logits_soft_cap} \times \mathrm{tanh}(x / \texttt{logits_soft_cap})`,
            where :math:`x` is the input logits.
        sm_scale : Optional[float]
            The scale used in softmax, if not provided, will be set to
            ``1.0 / sqrt(head_dim_qk)``.
        rope_scale : Optional[float]
            The scale used in RoPE interpolation, if not provided, will be set to
            ``1.0``.
        rope_theta : Optional[float]
            The theta used in RoPE, if not provided, will be set to ``1e4``.
        q_data_type : Union[str, torch.dtype]
            The data type of the query tensor, defaults to torch.float16.
        kv_data_type : Optional[Union[str, torch.dtype]]
            The data type of the key/value tensor. If None, will be set to :attr:`q_data_type`.
        o_data_type : Optional[Union[str, torch.dtype]]
            The data type of the output tensor. If None, will be set to :attr:`q_data_type`.
            For FP8 inputs, this should typically be set to torch.float16 or torch.bfloat16.
        non_blocking : bool
            Whether to copy the input tensors to the device asynchronously, defaults to ``True``.
        prefix_len_ptr :Optional[torch.Tensor]
            prefix length. A uint32 1D tensor indicating the prefix length of each prompt. The tensor size is equal to the batch size.
        token_pos_in_items_ptr : Optional[torch.Tensor]
            A uint16 1D tensor (it will be converted to uint16 in flashinfer) indicating the token position of each item and started from 0 (delimiter)
            for each item. E.g., if we have 3 items of length 3, 2, 4 respectively for this member. This vector will be looking like
            `[0, 1, 2, 3, 0, 1, 2, 0, 1, 2, 3, 4, 0]` with 4 delimiters indexed as 0. For batch size > 1,
            we will concat them as 1D with zero paddings to make sure each has the same length, the padding length is defined by
            `token_pos_in_items_len` - length of the raw `token_pos_in_items_ptr` for each prompt.
        token_pos_in_items_len : int
            zero padding length for `token_pos_in_items_ptr` to better handle the bsz > 1 case. Still using the above 3,2,4 example.
            If we set `token_pos_in_items_len` to be 20, it will be  `[0, 1, 2, 3, 0, 1, 2, 0, 1, 2, 3, 4, 0, 0, 0, 0, 0, 0, 0, 0]`
            with 7 padded zeros. (note there're 8 zeros in the end where the first one is the delimiter token 0 in the end of the prompt)
        max_item_len_ptr : Optional[torch.Tensor]
            a uint16 vector contains the max token length of all items for each prompt
        fixed_split_size : Optional[int],
            The fixed split size for split-kv FA2 prefill/decode, in pages. Recommend setting to the average sequence length of your workload.
            When enabled, will lead to deterministic softmax score reduction in the merge_states kernel, and therefore
            batch-size invariant outputs. See https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/
            Note that compatibility with CUDA graph is NOT guaranteed, as even when bs is fixed, kv seq len can change
            and lead to a varied number of launched CTAs.
        disable_split_kv : bool,
            Whether to disable the split-kv for determinism in CUDA Graph, defaults to ``False``.
        seq_lens: Optional[torch.Tensor]
            A uint32 1D tensor indicating the kv sequence length of each prompt. shape: ``[batch_size]``.
        seq_lens_q: Optional[torch.Tensor]
            A uint32 1D tensor indicating the q sequence length of each prompt. shape: ``[batch_size]``.
            If not provided, will be set to the same value as ``seq_lens``.
        max_token_per_sequence: Optional[int],
            Required for cudnn backend. This is the scalar max token length of each sequence.
        max_sequence_kv: Optional[int],
            Required for cudnn backend. This is the scalar max sequence length of each sequence in kv cache.
        v_indptr: Optional[torch.Tensor]
            Required for cudnn backend. This is the indptr of the value tensor.
        o_indptr: Optional[torch.Tensor]
            Required for cudnn backend. This is the indptr of the output tensor.
        Note
        ----
        The :meth:`plan` method should be called before any :meth:`run` or
        :meth:`run_return_lse` calls, auxiliary data structures will be created
        during this plan call and cached for multiple kernel runs.

        The ``num_qo_heads`` must be a multiple of ``num_kv_heads``. If ``num_qo_heads``
        is not equal to ``num_kv_heads``, the function will use
        `grouped query attention <https://arxiv.org/abs/2305.13245>`_.

        The :meth:`plan` method cannot be used in Cuda Graph or in ``torch.compile``.
        """
        q_data_type = canonicalize_torch_dtype(q_data_type)
        if kv_data_type is None:
            kv_data_type = q_data_type
        kv_data_type = canonicalize_torch_dtype(kv_data_type)
        if o_data_type is None:
            o_data_type = q_data_type
        o_data_type = canonicalize_torch_dtype(o_data_type)
        if head_dim_vo is None:
            head_dim_vo = head_dim_qk
        if fixed_split_size is None:
            fixed_split_size = -1
        if logits_soft_cap is None:
            logits_soft_cap = 0.0

        batch_size = len(qo_indptr) - 1
        if len(kv_indptr) != batch_size + 1:
            raise ValueError(
                "The kv_indptr length should be equal to mask_indptr length."
            )
        if custom_mask is not None or packed_custom_mask is not None:
            mask_indptr = _compute_mask_indptr(qo_indptr, kv_indptr)
        if packed_custom_mask is None and custom_mask is not None:
            # create packed custom mask from custom mask
            packed_custom_mask, mask_indptr = segment_packbits(
                custom_mask.contiguous().view(-1),
                mask_indptr,
                bitorder="little",
            )

        # NOTE(Zihao): only required if qo_indptr/paged_kv_indptr are device tensors
        qo_indptr_host = qo_indptr.to("cpu")
        kv_indptr_host = kv_indptr.to("cpu")

        self._qo_indptr_last = int(qo_indptr_host[-1])
        total_num_rows = self._qo_indptr_last

        if self.is_cuda_graph_enabled:
            if self._max_total_num_rows is None:
                self._max_total_num_rows = total_num_rows
            elif total_num_rows > self._max_total_num_rows:
                raise ValueError(
                    "The total number of rows in qo_indptr {} in cuda graph mode cannot "
                    "exceed the number of rows set during initialization {}.".format(
                        total_num_rows, self._max_total_num_rows
                    )
                )

            if batch_size != self._fixed_batch_size:
                raise ValueError(
                    "The batch size should be fixed in cudagraph mode, the runtime batch size {} "
                    " mismatches the batch size set during initialization {}.".format(
                        batch_size, self._fixed_batch_size
                    )
                )
            self._qo_indptr_buf.copy_(qo_indptr, non_blocking=non_blocking)
            self._kv_indptr_buf.copy_(kv_indptr, non_blocking=non_blocking)
            if packed_custom_mask is not None:
                if not torch.is_tensor(self._custom_mask_buf):
                    raise ValueError(
                        "custom_mask_buf must be initialized with a torch.Tensor in cuda graph mode if we use custom mask in attention computation."
                    )
                if not torch.is_tensor(self._mask_indptr_buf):
                    raise ValueError(
                        "mask_indptr_buf must be initialized with a torch.Tensor in cuda graph mode if we use custom mask in the attention computation."
                    )
                self._custom_mask_buf[: len(packed_custom_mask)] = packed_custom_mask
                self._mask_indptr_buf.copy_(mask_indptr, non_blocking=non_blocking)
        else:
            self._qo_indptr_buf = qo_indptr.to(self.device, non_blocking=non_blocking)
            self._kv_indptr_buf = kv_indptr.to(self.device, non_blocking=non_blocking)
            if packed_custom_mask is not None:
                self._custom_mask_buf = packed_custom_mask.to(
                    self.device, non_blocking=non_blocking
                )
                self._mask_indptr_buf = mask_indptr.to(
                    self.device, non_blocking=non_blocking
                )

        self._o_indptr_buf = (
            o_indptr.to(self.device, non_blocking=non_blocking)
            if o_indptr is not None
            else self._qo_indptr_buf
        )
        self._v_indptr_buf = (
            v_indptr.to(self.device, non_blocking=non_blocking)
            if v_indptr is not None
            else self._kv_indptr_buf
        )

        self._cached_q_data_type = q_data_type
        self._cached_kv_data_type = kv_data_type
        self._cached_o_data_type = o_data_type
        kv_len_arr = kv_indptr_host[1:] - kv_indptr_host[:-1]

        self._prefix_len_ptr = prefix_len_ptr
        self._token_pos_in_items_ptr = token_pos_in_items_ptr
        self._token_pos_in_items_len = token_pos_in_items_len
        self._max_item_len_ptr = max_item_len_ptr

        self._seq_lens_q = seq_lens_q
        self._seq_lens_kv = seq_lens
        self._max_token_per_sequence = max_token_per_sequence
        self._max_sequence_kv = max_sequence_kv

        # Precompute sequence length maxima (used by fmha_v2 and other backends)
        max_seq_in_batch = int(kv_len_arr.max().item())
        max_qo_len = int((qo_indptr_host[1:] - qo_indptr_host[:-1]).max().item())

        if self._backend == "cute-dsl":
            if custom_mask is not None or packed_custom_mask is not None:
                raise NotImplementedError(
                    "cute-dsl backend does not support custom_mask"
                )
            if head_dim_vo is not None and head_dim_vo != head_dim_qk:
                raise NotImplementedError(
                    "cute-dsl backend requires head_dim_vo == head_dim_qk"
                )
            if self._kv_layout == "HND":
                raise NotImplementedError("cute-dsl backend only supports NHD layout")
            if pos_encoding_mode not in ("NONE", "ALIBI"):
                raise NotImplementedError(
                    f"cute-dsl backend does not support pos_encoding_mode={pos_encoding_mode!r}. "
                    "For RoPE, apply rotary embeddings to Q/K before calling the kernel."
                )

            variant: Optional[Any] = None
            if pos_encoding_mode == "ALIBI":
                from .cute_dsl.attention import ALiBiAttention

                slopes = ALiBiAttention.get_slopes(num_qo_heads)
                variant = ALiBiAttention(
                    torch.tensor(slopes, dtype=torch.float32, device=qo_indptr.device)
                )
            if logits_soft_cap is not None and logits_soft_cap > 0:
                from .cute_dsl.attention import SoftCappingAttention

                if variant is not None:
                    raise NotImplementedError(
                        "cute-dsl backend does not support combining ALiBi with logits_soft_cap"
                    )
                variant = SoftCappingAttention(cap=logits_soft_cap)

            _sm_scale = (
                sm_scale if sm_scale is not None else 1.0 / math.sqrt(head_dim_qk)
            )
            # Variant-less head-128 dense/causal plans route to the trtllm
            # CuTe DSL FMHA kernel (prebuilt artifacts).  Sliding-window
            # plans use the modular cute-dsl path, which measures faster on
            # windows (banded masks skip dead work) for every V dtype —
            # mixed V (e.g. fp8 V with bf16 Q/K) forces the FMHA path onto
            # JIT compilation, where the modular kernel wins.
            self._cute_dsl_use_fmha = (
                variant is None and head_dim_qk == 128 and window_left < 0
            )
            if self._cute_dsl_use_fmha:
                q_lens = qo_indptr[1:] - qo_indptr[:-1]
                k_lens = kv_indptr[1:] - kv_indptr[:-1]
                self._cute_dsl_fmha_plan = {
                    "qo_indptr": qo_indptr,
                    "kv_indptr": kv_indptr,
                    "seq_lens": k_lens.to(torch.int32),
                    "batch_size": qo_indptr.shape[0] - 1,
                    "max_q_len": int(q_lens.max().item()),
                    "max_kv_len": int(k_lens.max().item()),
                    "sm_scale": _sm_scale,
                    "causal": causal,
                    "window_left": window_left,
                }
            else:
                # The modular kernel requires q and k to share one dtype
                # (kv_data_type covers both K and V; a V-only mismatch is a
                # run()-time property served natively).  Reject mismatched
                # plan dtypes here with a clear error — they would otherwise
                # surface as an opaque FFI signature mismatch at run().
                if kv_data_type is not None and q_data_type != kv_data_type:
                    raise ValueError(
                        "cute-dsl modular prefill requires q_data_type == "
                        f"kv_data_type, got {q_data_type} and {kv_data_type}"
                    )
                self._cute_dsl_wrapper.plan(
                    qo_indptr,
                    kv_indptr,
                    num_qo_heads,
                    num_kv_heads,
                    head_dim_qk,
                    head_dim_vo=head_dim_qk,
                    causal=causal,
                    sm_scale=_sm_scale,
                    q_data_type=q_data_type,
                    kv_data_type=kv_data_type
                    if kv_data_type is not None
                    else q_data_type,
                    window_left=window_left,
                    variant=variant,
                )
        elif self._jit_module is not None:
            self._cached_module = self._jit_module
        else:
            if self._backend == "auto":
                self._backend = determine_attention_backend(
                    self.device,
                    PosEncodingMode[pos_encoding_mode].value,
                    use_fp16_qk_reduction,
                    self._custom_mask_buf is not None,  # use_custom_mask
                    q_data_type,
                    kv_data_type,
                    head_dim_qk=head_dim_qk,
                    head_dim_vo=head_dim_vo,
                )
                # SM12x: use fmha_v2 HMMA kernels for large-sequence MHA.
                # fmha_v2 outperforms FA2 at longer sequences (S>=256) on SM120
                # due to optimized tiled flash attention, but has higher CTA
                # launch overhead for small sequences or high head counts.
                # Only supports: MHA, standard symmetric head dims, matching
                # q/kv dtype, no sliding window, no softcap, equal q/kv
                # segments, NHD layout. The kernel reduces QK in fp16 on bf16,
                # so only auto-select it when the caller allows fp16 QK
                # reduction (mirrors is_fa3_backend_supported's convention).
                same_qk_segments = torch.equal(qo_indptr_host, kv_indptr_host)
                is_standard_shape = head_dim_qk == head_dim_vo and head_dim_qk in (
                    64,
                    128,
                    256,
                )
                if (
                    self._backend == "fa2"
                    and self._kv_layout == "NHD"
                    and num_qo_heads == num_kv_heads
                    and same_qk_segments
                    and is_standard_shape
                    and q_data_type == kv_data_type
                    and use_fp16_qk_reduction
                    and max_seq_in_batch >= 256
                    and window_left < 0
                    and logits_soft_cap == 0.0
                    and prefix_len_ptr is None
                    and token_pos_in_items_ptr is None
                    and max_item_len_ptr is None
                    and _should_use_fmha_v2_sm120(
                        self.device,
                        PosEncodingMode[pos_encoding_mode].value,
                        self._custom_mask_buf is not None,
                        q_data_type,
                        kv_data_type,
                    )
                ):
                    self._backend = "fmha_v2"

            get_module_args = (
                q_data_type,
                kv_data_type,
                o_data_type,
                kv_indptr.dtype,
                head_dim_qk,
                head_dim_vo,
                PosEncodingMode[pos_encoding_mode].value,
                window_left >= 0,  # use_sliding_window
                logits_soft_cap > 0,  # use_logits_soft_cap
                use_fp16_qk_reduction,
            )
            if self._backend == "fmha_v2":
                # Cache plan data for run() — avoid GPU-CPU sync in hot path
                self._fmha_v2_seq_lens = kv_len_arr.to(torch.int32).to(
                    self.device, non_blocking=non_blocking
                )
                self._fmha_v2_batch_size = batch_size
                self._fmha_v2_max_q_len = max_qo_len
                self._fmha_v2_max_kv_len = max_seq_in_batch
                self._fmha_v2_qo_indptr = self._qo_indptr_buf.to(torch.int32)
                self._fmha_v2_kv_indptr = self._kv_indptr_buf.to(torch.int32)
            elif self._backend == "cutlass":
                # insert qo_indptr.device to 9th position (0-indexed) of get_module_args
                new_get_module_args = (
                    get_module_args[:9] + (qo_indptr.device,) + get_module_args[9:]
                )
                self._cached_module = get_fmha_module(*new_get_module_args)
            elif self._backend != "cudnn":
                self._cached_module = get_batch_prefill_module(
                    self._backend, *get_module_args
                )

        if self._backend == "cutlass":
            self._plan_info = fmha_varlen_plan(
                self._cached_module, qo_indptr, kv_indptr, num_qo_heads, causal
            )
            self._max_qo_len = torch.max(qo_indptr[1:] - qo_indptr[:-1]).item()
        elif self._backend == "fmha_v2":
            # fmha_v2 handles planning internally — no JIT module plan needed
            pass
        elif self._backend not in ("cudnn", "cute-dsl"):
            assert self._cached_module is not None, "cached module is not initialized"
            args = [
                self._float_workspace_buffer,
                self._int_workspace_buffer,
                self._pin_memory_int_workspace_buffer,
                qo_indptr_host,
                kv_indptr_host,
                kv_len_arr,
                self._max_total_num_rows or total_num_rows,
                batch_size,
                num_qo_heads,
                num_kv_heads,
                1,  # page_size
                self.is_cuda_graph_enabled,
                head_dim_qk,
                head_dim_vo,
                causal,
                window_left,
            ]
            if self._backend == "fa2":
                args.append(fixed_split_size or -1)  # fixed_split_size
                args.append(disable_split_kv)  # disable_split_kv
                args.append(0)  # num_colocated_ctas
                args.append(0)  # uniform_q_len
            self._plan_info = self._cached_module.plan(
                *args,
            )

        self._causal = causal
        self._pos_encoding_mode = pos_encoding_mode
        self._use_fp16_qk_reduction = use_fp16_qk_reduction
        self._window_left = window_left
        self._logits_soft_cap = logits_soft_cap
        self._sm_scale = sm_scale
        self._rope_scale = rope_scale
        self._rope_theta = rope_theta

    begin_forward = plan

    def forward(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        causal: bool = False,
        pos_encoding_mode: str = "NONE",
        use_fp16_qk_reduction: bool = False,
        window_left: int = -1,
        logits_soft_cap: Optional[float] = None,
        sm_scale: Optional[float] = None,
        rope_scale: Optional[float] = None,
        rope_theta: Optional[float] = None,
    ) -> torch.Tensor:
        r"""Warning: This function is deprecated, please use :meth:`run` instead."""
        self._causal = causal
        self._pos_encoding_mode = pos_encoding_mode
        self._use_fp16_qk_reduction = use_fp16_qk_reduction
        self._window_left = window_left
        self._logits_soft_cap = logits_soft_cap
        self._sm_scale = sm_scale
        self._rope_scale = rope_scale
        self._rope_theta = rope_theta
        return self.run(q, k, v)

    @overload
    def run(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        *args,
        out: Optional[torch.Tensor] = None,
        lse: Optional[torch.Tensor] = None,
        return_lse: Literal[False] = False,
        enable_pdl: Optional[bool] = None,
        kv_cache_sf: Optional[
            Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
        ] = None,
    ) -> torch.Tensor: ...

    @overload
    def run(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        *args,
        out: Optional[torch.Tensor] = None,
        lse: Optional[torch.Tensor] = None,
        return_lse: Literal[True] = True,
        enable_pdl: Optional[bool] = None,
        kv_cache_sf: Optional[
            Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
        ] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]: ...

    @flashinfer_api(trace=gqa_ragged_prefill_trace)
    def run(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        *args,
        q_scale: Optional[float] = None,
        k_scale: Optional[float] = None,
        v_scale: Optional[float] = None,
        o_scale: Optional[float] = None,
        out: Optional[torch.Tensor] = None,
        lse: Optional[torch.Tensor] = None,
        return_lse: bool = False,
        enable_pdl: Optional[bool] = None,
        kv_cache_sf: Optional[
            Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
        ] = None,
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
        r"""Compute batch prefill/append attention between query and kv-cache stored as
        ragged tensor.

        Parameters
        ----------
        q : torch.Tensor
            The query tensor, shape: ``[qo_indptr[-1], num_qo_heads, head_dim_qk]``
        k : torch.Tensor
            The key tensor, shape: ``[kv_indptr[-1], num_kv_heads, head_dim_qk]``
        v : torch.Tensor
            The value tensor, shape: ``[kv_indptr[-1], num_kv_heads, head_dim_vo]``
        *args
            Additional arguments for the custom kernel.
        q_scale: Optional[float]
            The calibration scale of fp8 query, if not provided, will be set to ``1.0``.
        k_scale: Optional[float]
            The calibration scale of fp8 or nvfp4 key, if not provided, will be set to ``1.0``.
        v_scale: Optional[float]
            The calibration scale of fp8 or nvfp4 value, if not provided, will be set to ``1.0``.
        o_scale: Optional[float]
            The calibration scale of output, if not provided, will be set to ``1.0``.
        out : Optional[torch.Tensor]
            The output tensor, if not provided, will be allocated internally.
        lse : Optional[torch.Tensor]
            The log-sum-exp of attention logits, if not provided, will be allocated internally.
        return_lse : bool
            Whether to return the logsumexp of attention output
        enable_pdl : bool
            Whether to enable Programmatic Dependent Launch (PDL). See https://docs.nvidia.com/cuda/cuda-c-programming-guide/#programmatic-dependent-launch-and-synchronization
            Only supported for >= sm90, and currently only for FA2 and CUDA core decode.
        kv_cache_sf : Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]]
            Per-block scale factors for NVFP4 KV input.  Accepts either a single
            packed scale tensor or a ``(k_scales, v_scales)`` tuple matching the
            structure expected by the chosen backend.  When ``None`` (default), the
            kernel runs without NVFP4 KV scaling.  See
            :func:`flashinfer.fp4_quantization.nvfp4_quantize` for layout details.
        Returns
        -------
        Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
            If :attr:`return_lse` is ``False``, the attention output, shape: ``[qo_indptr[-1], num_qo_heads, head_dim_vo]``.
            If :attr:`return_lse` is ``True``, a tuple of two tensors:

            * The attention output, shape: ``[qo_indptr[-1], num_qo_heads, head_dim_vo]``.
            * The logsumexp of attention output, shape: ``[qo_indptr[-1], num_qo_heads]``.
        """
        if enable_pdl is None:
            enable_pdl = device_support_pdl(q.device)
        _check_cached_qkv_data_type(
            q, k, self._cached_q_data_type, self._cached_kv_data_type
        )
        # Validate q shape matches qo_indptr (using value cached in plan() to avoid GPU sync)
        if self._backend == "cudnn":
            if q.numel() != self._qo_indptr_last:
                raise ValueError(
                    f"q.numel() ({q.numel()}) does not match qo_indptr[-1] ({self._qo_indptr_last}). "
                    f"For cudnn ragged prefill, qo_indptr uses element offsets "
                    f"(total_tokens * num_heads * head_dim)."
                )
        else:
            if q.size(0) != self._qo_indptr_last:
                raise ValueError(
                    f"q.shape[0] ({q.size(0)}) does not match qo_indptr[-1] ({self._qo_indptr_last}). "
                    f"For ragged prefill, q must have shape [total_tokens, num_heads, head_dim] "
                    f"where total_tokens = qo_indptr[-1]."
                )

        window_left = self._window_left
        logits_soft_cap = self._logits_soft_cap
        sm_scale = self._sm_scale
        rope_scale = self._rope_scale
        rope_theta = self._rope_theta
        if logits_soft_cap is None:
            logits_soft_cap = 0.0
        if sm_scale is None:
            sm_scale = 1.0 / math.sqrt(q.size(-1))
        # For NVFP4 KV, fuse q_scale and k_scale into sm_scale
        if kv_cache_sf is not None:
            if q_scale is not None:
                sm_scale *= q_scale
            if k_scale is not None:
                sm_scale *= k_scale
        if rope_scale is None:
            rope_scale = 1.0
        if rope_theta is None:
            rope_theta = 1e4
        if return_lse:
            if lse is None:
                lse = torch.empty(
                    (q.size(0), q.size(1)), dtype=torch.float32, device=q.device
                )
            else:
                check_shape_dtype_device(
                    lse, (q.size(0), q.size(1)), torch.float32, q.device, "lse"
                )
        # Unpack kv_cache_sf for NVFP4 ragged KV
        k_sf, v_sf = None, None
        if kv_cache_sf is not None:
            if isinstance(kv_cache_sf, tuple):
                k_sf, v_sf = kv_cache_sf
            else:
                k_sf, v_sf = kv_cache_sf.unbind(dim=1)

        # For NVFP4 KV (uint8 packed), v last dim is head_dim//2; use q head_dim for output
        out_head_dim = q.shape[-1] if kv_cache_sf is not None else v.shape[-1]
        if out is None:
            # when input dtype is fp8, we need to use bf16 output
            out_dtype = torch.bfloat16 if q.dtype.itemsize == 1 else q.dtype
            out = torch.empty(
                q.shape[:-1] + (out_head_dim,),
                dtype=out_dtype,
                device=q.device,
            )
        else:
            check_shape_dtype_device(
                out,
                q.shape[:-1] + (out_head_dim,),
                self._cached_o_data_type,
                q.device,
                "out",
            )
        if self._backend == "fmha_v2":
            if return_lse:
                raise NotImplementedError(
                    "return_lse is not yet supported for backend='fmha_v2'. "
                    "TRT-LLM fmha_v2 softmax stats have a different format "
                    "than the expected [total_q_tokens, num_heads] LSE tensor."
                )
            fmha_v2_mask_mode = "causal" if self._causal else "padding"
            # Use CONTIGUOUS_Q_KV layout: pass Q separately and stack only
            # K+V into [num_tokens, 2, num_kv_heads, head_dim]. This avoids
            # the 3-way torch.stack copy that PACKED_QKV requires.
            # (fmha_v2 auto-selection already requires MHA, so Q/K head counts match)
            #
            # CONTIGUOUS_Q_KV requires Q to be contiguous. torch.stack() below
            # materializes a contiguous kv_packed, but Q is forwarded directly, so a
            # non-contiguous (e.g. packed/strided) Q would be read with the wrong
            # strides and yield garbage. .contiguous() is a no-op when Q is already
            # contiguous (the common hot path), so it restores correctness for
            # non-contiguous Q with no cost otherwise.
            q = q.contiguous()
            kv_packed = torch.stack([k, v], dim=1)
            # Pass window_left and sinks for sliding-window + attention sink
            # support. sinks is stored on the wrapper by vLLM's metadata builder.
            fmha_v2_sinks = getattr(self, "_sinks", None)
            fmha_v2_kwargs = {}
            if self._window_left >= 0:
                fmha_v2_kwargs["window_left"] = self._window_left
            if fmha_v2_sinks is not None:
                fmha_v2_kwargs["sinks"] = fmha_v2_sinks
            out = trtllm_fmha_v2_prefill(
                (q, kv_packed),
                input_layout="CONTIGUOUS_Q_KV",
                workspace_buffer=self._float_workspace_buffer,
                seq_lens=self._fmha_v2_seq_lens,
                max_q_len=self._fmha_v2_max_q_len,
                max_kv_len=self._fmha_v2_max_kv_len,
                bmm1_scale=sm_scale,
                bmm2_scale=1.0,
                batch_size=self._fmha_v2_batch_size,
                cum_seq_lens_q=self._fmha_v2_qo_indptr,
                cum_seq_lens_kv=self._fmha_v2_kv_indptr,
                out=out,
                mask_mode=fmha_v2_mask_mode,
                save_softmax_stats=False,
                **fmha_v2_kwargs,
            )
            return out
        elif self._backend == "cute-dsl":
            if any(s is not None for s in (q_scale, k_scale, v_scale, o_scale)):
                raise NotImplementedError(
                    "cute-dsl backend does not support FP8 scale parameters"
                )
            if kv_cache_sf is not None:
                raise NotImplementedError(
                    "cute-dsl backend does not support NVFP4 packed KV cache "
                    "(kv_cache_sf)"
                )
            # enable_pdl is a launch-latency hint: forwarded on the FMHA
            # route below; the modular kernel has no PDL support.
            if self._cute_dsl_use_fmha:
                # Delegate dense/causal plans to the trtllm CuTe DSL FMHA
                # kernel (mixed V dtype included: it JIT-compiles
                # separate-V-dtype variants).  Windowed plans stay on the
                # modular path for every V dtype.
                p = self._cute_dsl_fmha_plan
                return trtllm_ragged_attention_deepseek(
                    query=q,
                    key=k,
                    value=v,
                    workspace_buffer=self._float_workspace_buffer,
                    seq_lens=p["seq_lens"],
                    max_q_len=p["max_q_len"],
                    max_kv_len=p["max_kv_len"],
                    bmm1_scale=p["sm_scale"],
                    bmm2_scale=1.0,
                    o_sf_scale=1.0,
                    batch_size=p["batch_size"],
                    window_left=p["window_left"],
                    cum_seq_lens_q=p["qo_indptr"],
                    cum_seq_lens_kv=p["kv_indptr"],
                    enable_pdl=enable_pdl,
                    is_causal=p["causal"],
                    return_lse=return_lse,
                    out=out,
                    lse=lse,
                    backend="cute-dsl",
                )
            if return_lse:
                # Standard-path modular kernel computes LSE natively; the
                # wrapper raises NotImplementedError for attention variants.
                return self._cute_dsl_wrapper.run(
                    q, k, v, out=out, return_lse=True, lse=lse
                )
            out = self._cute_dsl_wrapper.run(q, k, v, out=out)
            return out
        elif self._backend == "cutlass":
            out, lse = fmha_varlen(
                q,
                k,
                v,
                self._qo_indptr_buf,
                self._kv_indptr_buf,
                plan_info=self._plan_info,
                causal=self._causal,
                sm_scale=sm_scale,
                q_scale=q_scale,
                k_scale=k_scale,
                v_scale=v_scale,
                o_scale=o_scale,
                max_qo_len=self._max_qo_len,
                out=out,
                lse=lse,
            )
            return (out, lse) if return_lse else out
        elif self._backend == "cudnn":
            if self._seq_lens_q.dim() == 1:
                batch_size = self._seq_lens_q.shape[0]
            if self._seq_lens_q is not None and self._seq_lens_q.dim() == 1:
                self._seq_lens_q = self._seq_lens_q.reshape(batch_size, 1, 1, 1)

            if self._seq_lens_kv is not None and self._seq_lens_kv.dim() == 1:
                self._seq_lens_kv = self._seq_lens_kv.reshape(batch_size, 1, 1, 1)

            cudnn_batch_prefill_with_kv_cache(
                q,
                k,
                v,
                sm_scale,
                self._float_workspace_buffer,
                max_token_per_sequence=self._max_token_per_sequence,
                max_sequence_kv=self._max_sequence_kv,
                actual_seq_lens_q=self._seq_lens_q,
                actual_seq_lens_kv=self._seq_lens_kv,
                return_lse=return_lse,
                causal=self._causal,
                q_scale=q_scale,
                k_scale=k_scale,
                v_scale=v_scale,
                batch_offsets_q=self._qo_indptr_buf,
                batch_offsets_k=self._kv_indptr_buf,
                batch_offsets_v=self._v_indptr_buf,
                batch_offsets_o=self._o_indptr_buf,
                is_cuda_graph_compatible=self._use_cuda_graph,
                out=out,
                lse=lse,
            )

            return (out, lse) if return_lse else out

        # Skip FP8->FP16 conversion for FA3 backend with FP8 support
        # The JIT module will handle FP8 natively
        if is_float8(q) and self._backend != "fa3":
            logging.warning(
                "Our current prefill kernel implementation needs f16 input, the f8 inputs "
                " are casted to f16, which could result in performance degradation."
            )
            q = q.to(torch.float16)
            k = k.to(torch.float16)
            v = v.to(torch.float16)

        if self._custom_mask_buf is not None:
            mask_mode = MaskMode.CUSTOM.value
        else:
            if self._causal:
                mask_mode = MaskMode.CAUSAL.value
            else:
                mask_mode = MaskMode.NON_CAUSAL.value

        run_args = [
            self._float_workspace_buffer,
            self._int_workspace_buffer,
            self._plan_info,
            q,
            k,
            v,
            self._qo_indptr_buf,
            self._kv_indptr_buf,
            out,
            lse,
            mask_mode,
            TensorLayout[self._kv_layout].value,
            window_left,
            enable_pdl,
        ]
        if self._jit_module is not None:
            run_args.extend(
                prepare_jit_additional_args(
                    self._jit_additional_tensor_names,
                    {
                        "maybe_custom_mask": self._custom_mask_buf,
                        "maybe_mask_indptr": self._mask_indptr_buf,
                        "maybe_alibi_slopes": lambda: _get_cache_alibi_slopes_buf(
                            q.shape[1], self.device
                        ),
                    },
                    args,
                )
            )
        else:
            run_args += [
                self._custom_mask_buf,
                self._mask_indptr_buf,
                _get_cache_alibi_slopes_buf(q.shape[1], self.device),
                self._prefix_len_ptr,
                self._token_pos_in_items_ptr,
                self._max_item_len_ptr,
                logits_soft_cap,
                sm_scale,
                rope_scale,
                rope_theta,
                self._token_pos_in_items_len,
                k_sf,
                v_sf,
            ]
            # For FP8, append scale tensors
            if is_float8(q):
                run_args.extend(list(args))  # scale_q, scale_k, scale_v

        assert self._cached_module is not None, "cached module is not initialized"
        self._cached_module.ragged_run(*run_args)

        # Apply V scaling for NVFP4 ragged KV if v_scale is provided and not equal to 1.0
        is_float_one = isinstance(v_scale, float) and v_scale == 1.0
        if kv_cache_sf is not None and v_scale is not None and not is_float_one:
            out *= v_scale

        return (out, lse) if return_lse else out

    run_return_lse = functools.partialmethod(run, return_lse=True)

    def forward_return_lse(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        causal: bool = False,
        pos_encoding_mode: str = "NONE",
        use_fp16_qk_reduction: bool = False,
        window_left: int = -1,
        logits_soft_cap: Optional[float] = None,
        sm_scale: Optional[float] = None,
        rope_scale: Optional[float] = None,
        rope_theta: Optional[float] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        r"""Warning: This function is deprecated, please use :meth:`run_return_lse` instead."""
        self._causal = causal
        self._pos_encoding_mode = pos_encoding_mode
        self._use_fp16_qk_reduction = use_fp16_qk_reduction
        self._window_left = window_left
        self._logits_soft_cap = logits_soft_cap
        self._sm_scale = sm_scale
        self._rope_scale = rope_scale
        self._rope_theta = rope_theta
        return self.run_return_lse(q, k, v)

    def end_forward(self) -> None:
        r"""Warning: this function is deprecated and has no effect."""
        pass


def fmha_varlen_plan(
    module,
    qo_segment_offsets: torch.Tensor,
    kv_segment_offsets: torch.Tensor,
    num_qo_heads: int,
    causal: bool,
):
    num_ctas = torch.cuda.get_device_properties(
        qo_segment_offsets.device
    ).multi_processor_count
    work_indptr = torch.empty(
        num_ctas + 1, device=qo_segment_offsets.device, dtype=torch.int32
    )
    qo_tile_indices = torch.empty(
        131072, device=qo_segment_offsets.device, dtype=torch.int32
    )
    head_indices = torch.empty(
        131072, device=qo_segment_offsets.device, dtype=torch.int32
    )
    batch_indices = torch.empty(
        131072, device=qo_segment_offsets.device, dtype=torch.int32
    )
    module.plan(
        qo_segment_offsets,
        kv_segment_offsets,
        work_indptr,
        qo_tile_indices,
        head_indices,
        batch_indices,
        256,  # qo_tile_size
        num_qo_heads,
        num_ctas,
        causal,
    )
    return (
        work_indptr,
        qo_tile_indices,
        head_indices,
        batch_indices,
    )


@overload
def fmha_varlen(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    qo_segment_offsets: torch.Tensor,
    kv_segment_offsets: torch.Tensor,
    plan_info: Optional[List[torch.Tensor]] = None,
    max_qo_len: Optional[int] = None,
    out: Optional[torch.Tensor] = None,
    lse: Optional[torch.Tensor] = None,
    causal: bool = False,
    sm_scale: Optional[float] = None,
    q_scale: Optional[float] = None,
    k_scale: Optional[float] = None,
    v_scale: Optional[float] = None,
    o_scale: Optional[float] = None,
    return_lse: Literal[False] = False,
) -> torch.Tensor: ...


@overload
def fmha_varlen(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    qo_segment_offsets: torch.Tensor,
    kv_segment_offsets: torch.Tensor,
    plan_info: Optional[List[torch.Tensor]] = None,
    max_qo_len: Optional[int] = None,
    out: Optional[torch.Tensor] = None,
    lse: Optional[torch.Tensor] = None,
    causal: bool = False,
    sm_scale: Optional[float] = None,
    q_scale: Optional[float] = None,
    k_scale: Optional[float] = None,
    v_scale: Optional[float] = None,
    o_scale: Optional[float] = None,
    return_lse: Literal[True] = True,
) -> Tuple[torch.Tensor, torch.Tensor]: ...


def fmha_varlen(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    qo_segment_offsets: torch.Tensor,
    kv_segment_offsets: torch.Tensor,
    plan_info: Optional[List[torch.Tensor]] = None,
    max_qo_len: Optional[int] = None,
    out: Optional[torch.Tensor] = None,
    lse: Optional[torch.Tensor] = None,
    causal: bool = False,
    sm_scale: Optional[float] = None,
    q_scale: Optional[float] = None,
    k_scale: Optional[float] = None,
    v_scale: Optional[float] = None,
    o_scale: Optional[float] = None,
    return_lse: bool = False,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
    workspace_buffer = _get_cache_buf(
        "fmha_varlen_cutlass_workspace", 32 * 1024 * 1024, q.device
    )
    module = get_fmha_module(
        q.dtype,
        k.dtype,
        v.dtype,
        torch.int32,
        q.shape[2],
        v.shape[2],
        PosEncodingMode.NONE.value,
        False,  # use_sliding_window
        False,  # use_logits_soft_cap
        q.device,
    )

    nnz_qo, num_qo_heads, head_dim_qk = q.shape
    nnz_kv, num_kv_heads, head_dim_vo = v.shape

    mask_mode_code = 1 if causal else 0
    if sm_scale is None:
        sm_scale = 1.0 / math.sqrt(head_dim_qk)
    if q_scale is None:
        q_scale = 1.0
    if k_scale is None:
        k_scale = 1.0
    if v_scale is None:
        v_scale = 1.0
    if o_scale is None:
        o_scale = 1.0

    qo_total_len = nnz_qo
    if max_qo_len is None:
        max_qo_len = torch.max(qo_segment_offsets[1:] - qo_segment_offsets[:-1]).item()

    if plan_info is None:
        plan_info = fmha_varlen_plan(
            module, qo_segment_offsets, kv_segment_offsets, num_qo_heads, causal
        )

    (
        work_indptr,
        qo_tile_indices,
        head_indices,
        batch_indices,
    ) = plan_info

    if out is None:
        # when input dtype is fp8, we need to use bf16 output
        out_dtype = torch.bfloat16 if q.dtype.itemsize == 1 else q.dtype
        out = torch.empty(
            qo_total_len + max(max_qo_len, 128),
            num_qo_heads,
            head_dim_vo,
            device=q.device,
            dtype=out_dtype,
        )[max(max_qo_len, 128) :]

    if lse is None and return_lse:
        lse = torch.empty(
            qo_total_len, num_qo_heads, device=q.device, dtype=torch.float32
        )

    module.run(
        workspace_buffer,
        q,
        k,
        v,
        qo_segment_offsets,
        kv_segment_offsets,
        work_indptr,
        qo_tile_indices,
        head_indices,
        batch_indices,
        out,
        lse,
        mask_mode_code,
        sm_scale,
        q_scale,
        k_scale,
        v_scale,
        o_scale,
        max_qo_len,
    )

    return out, lse


@functools.cache
def get_trtllm_gen_fmha_module():
    mod = gen_trtllm_gen_fmha_module()
    op = mod.build_and_load()
    setup_cubin_loader(mod.get_library_path())
    return op


def trtllm_sage_attention_quantize(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    q_block_size: int = 1,
    k_block_size: int = 16,
    qk_quant_dtype: torch.dtype = torch.int8,
) -> Tuple[
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
]:
    """Quantize Q/K/V into the layout consumed by TRTLLM SageAttention.

    Q and K use per-token-block scaling and can be quantized to INT8 or FP8
    E4M3. V is always quantized to FP8 E4M3 with one scale per head channel.
    All inputs use contiguous ``[tokens, heads, head_dim]`` layout.

    Parameters
    ----------
    query, key, value : torch.Tensor
        FP16 or BF16 CUDA tensors. Key and value must have matching token and
        head counts, and all three tensors must have the same head dimension.
    q_block_size, k_block_size : int
        Number of tokens sharing each Q or K scale. Supported values are 1, 4,
        and 16.
    qk_quant_dtype : torch.dtype
        ``torch.int8`` (the SageAttention INT8 path) or
        ``torch.float8_e4m3fn``.

    Returns
    -------
    (q_quant, k_quant, v_quant, q_scale, k_scale, v_scale)
        Quantized Q/K/V followed by their FP32 dequantization scales. Q and K
        scales have shapes ``[q_heads, ceil(q_tokens / q_block_size)]`` and
        ``[kv_heads, ceil(kv_tokens / k_block_size)]``. V scales have shape
        ``[kv_heads, head_dim]``. These outputs can be passed directly to
        :func:`trtllm_ragged_attention_deepseek`.
    """
    if query.dtype not in (torch.float16, torch.bfloat16):
        raise ValueError(f"query must be float16 or bfloat16, got {query.dtype}")
    if key.dtype != query.dtype or value.dtype != query.dtype:
        raise ValueError("query, key, and value must have the same dtype")
    if query.device.type != "cuda":
        raise ValueError("query, key, and value must be CUDA tensors")
    if key.device != query.device or value.device != query.device:
        raise ValueError("query, key, and value must be on the same CUDA device")
    if query.ndim != 3 or key.ndim != 3 or value.ndim != 3:
        raise ValueError("query, key, and value must be 3D [tokens, heads, head_dim]")
    if query.shape[2] != key.shape[2] or key.shape[2] != value.shape[2]:
        raise ValueError("query, key, and value must have the same head dimension")
    if key.shape[:2] != value.shape[:2]:
        raise ValueError("key and value must have matching token and head counts")
    if query.shape[2] not in (64, 128, 256):
        raise ValueError("head_dim must be 64, 128, or 256")
    if q_block_size not in (1, 4, 16) or k_block_size not in (1, 4, 16):
        raise ValueError("q_block_size and k_block_size must be 1, 4, or 16")
    if qk_quant_dtype not in (torch.int8, torch.float8_e4m3fn):
        raise ValueError("qk_quant_dtype must be torch.int8 or torch.float8_e4m3fn")

    query = query.contiguous()
    key = key.contiguous()
    value = value.contiguous()
    q_quant = torch.empty_like(query, dtype=qk_quant_dtype)
    k_quant = torch.empty_like(key, dtype=qk_quant_dtype)
    v_quant = torch.empty_like(value, dtype=torch.float8_e4m3fn)
    q_scale = torch.empty(
        (query.shape[1], ceil_div(query.shape[0], q_block_size)),
        dtype=torch.float32,
        device=query.device,
    )
    k_scale = torch.empty(
        (key.shape[1], ceil_div(key.shape[0], k_block_size)),
        dtype=torch.float32,
        device=query.device,
    )
    v_scale = torch.empty(
        (value.shape[1], value.shape[2]),
        dtype=torch.float32,
        device=query.device,
    )

    get_trtllm_gen_fmha_module().trtllm_sage_attention_quantize(
        q_quant,
        k_quant,
        v_quant,
        q_scale,
        k_scale,
        v_scale,
        query,
        key,
        value,
        q_block_size,
        k_block_size,
        get_device_sm_count(query.device),
    )
    return q_quant, k_quant, v_quant, q_scale, k_scale, v_scale


@flashinfer_api(trace=trtllm_ragged_attention_deepseek_trace)
def trtllm_ragged_attention_deepseek(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    workspace_buffer: torch.Tensor,
    seq_lens: torch.Tensor,
    max_q_len: int,
    max_kv_len: int,
    bmm1_scale: Union[float, torch.Tensor],
    bmm2_scale: Union[float, torch.Tensor],
    o_sf_scale: float,
    batch_size: int,
    window_left: int,
    cum_seq_lens_q: torch.Tensor,
    cum_seq_lens_kv: torch.Tensor,
    enable_pdl: bool | None,
    is_causal: bool,
    return_lse: bool,
    attention_sinks: Optional[torch.Tensor] = None,
    skip_softmax_threshold_scale_factor: Optional[float] = None,
    out: Optional[torch.Tensor] = None,
    lse: Optional[torch.Tensor] = None,
    sage_attn_sfs: Tuple[
        Optional[torch.Tensor],
        Optional[torch.Tensor],
        Optional[torch.Tensor],
        Optional[torch.Tensor],
    ] = (None, None, None, None),
    num_elts_per_sage_attn_blk: Tuple[int, int, int, int] = (0, 0, 0, 0),
    backend: str = "trtllm-gen",
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
    """
    Parameters
    ----------
    query : torch.Tensor
        query tensor with shape [num_tokens, num_heads, head_dim]
    key : torch.Tensor
        key tensor with shape [num_tokens, num_heads, head_dim]
    value : torch.Tensor
        value tensor with shape [num_tokens, num_heads, head_dim]
    workspace_buffer : torch.Tensor
        workspace buffer
    seq_lens : torch.Tensor
        sequence lengths
    max_q_len : int
        max query length
    max_kv_len : int
        max key/value length
    bmm1_scale : Union[float, torch.Tensor]
        scale for bmm1, scale_q * scale_k * 1.0 / (head_dim_qk ** 0.5)
        when using trtllm-gen backend, it can be a torch.Tensor with dtype torch.float32.
    bmm2_scale : Union[float, torch.Tensor]
        scale for bmm2, scale_v
        when using trtllm-gen backend, it can be a torch.Tensor with dtype torch.float32.
    o_sf_scale : float
        scale for output
    batch_size : int
        batch size
    window_left : int
        window left
    cum_seq_lens_q : torch.Tensor
        cumulative sequence lengths for query
    cum_seq_lens_kv : torch.Tensor
        cumulative sequence lengths for key/value
    enable_pdl : bool
        enable pdl
    is_causal : bool
        is causal
    return_lse : bool
        Whether to allocate and return the log-sum-exp tensor in addition to the
        attention output.  When ``True`` the function returns
        ``(out, lse)``; when ``False`` only ``out`` is returned.
    attention_sinks : Optional[torch.Tensor]
        attention sinks
    skip_softmax_threshold_scale_factor : Optional[float]
        threshold scale factor for skipping softmax operations.
        Providing a value for this parameter enables skip-softmax sparsity as described in: https://arxiv.org/abs/2512.12087
        If no value is provided, then standard attention is used.
        Setting the threshold to a higher value generally increases kernel performance at the cost of accuracy degradation.
        The actual threshold value equals the provided threshold_scale_factor divided by the context length.
    out : Optional[torch.Tensor]
        output tensor, if not provided, will be allocated with shape [query.shape[0], query.shape[1], value.shape[2]]
    lse : Optional[torch.Tensor]
        lse tensor, if not provided, will be allocated with shape [query.shape[0], query.shape[1]]
    sage_attn_sfs : Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]
        SageAttention scale-factor tensors for the four sub-blocks ``(q_sf, k_sf, p_sf, v_sf)``.
        The outputs from :func:`trtllm_sage_attention_quantize` can be supplied as Q, K, and V scales (with ``p_sf=None``).
        Defaults to ``(None, None, None, None)`` which disables SageAttention.
        Currently only consulted by the ``trtllm-gen`` backend.
    num_elts_per_sage_attn_blk : Tuple[int, int, int, int]
        Per-block element counts for the SageAttention scale-factor tensors,
        matching the order of ``sage_attn_sfs``.  Defaults to ``(0, 0, 0, 0)``,
        which disables SageAttention.  Only consulted when ``sage_attn_sfs``
        contains non-``None`` tensors.
    backend : str
        Attention backend to use. "trtllm-gen" (default) or "cute-dsl".

    Returns
    -------
    out: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
        output torch.Tensor or Tuple[torch.Tensor, torch.Tensor].
        If return_lse is True, the output will be a tuple of two tensors, the first is the output tensor, the second is the lse tensor.
        If return_lse is False, the output will be a single tensor.
    """
    is_dsr1 = query.shape[2] == 192 and key.shape[2] == 192 and value.shape[2] == 128
    is_smaller_dimensions = (
        query.shape[2] == 128 and key.shape[2] == 128 and value.shape[2] == 128
    )
    is_glm5 = query.shape[2] == 256 and key.shape[2] == 256 and value.shape[2] == 256
    assert is_dsr1 or is_smaller_dimensions or is_glm5, (
        f"Unsupported head dimensions: query={query.shape[2]}, key={key.shape[2]}, value={value.shape[2]}. "
        "Supported: (192, 192, 128) for DeepSeek-R1, (128, 128, 128), or (256, 256, 256) for GLM-5."
    )

    if enable_pdl is None:
        enable_pdl = device_support_pdl(query.device)

    # --- Output allocation (shared by all backends) ---
    if out is None:
        # FP8 inputs produce bfloat16 output by default (TRT-LLM kernels
        # do not support FP8 output for ragged attention)
        if query.dtype in (torch.float8_e4m3fn, torch.float8_e5m2, torch.int8):
            out_dtype = torch.bfloat16
        else:
            out_dtype = query.dtype
        out = torch.empty(
            query.shape[0],
            query.shape[1],
            value.shape[2],
            device=query.device,
            dtype=out_dtype,
        )
    else:
        expected_shape = (query.shape[0], query.shape[1], value.shape[2])
        if tuple(out.shape) != expected_shape:
            raise ValueError(
                f"out must have shape {expected_shape}, got {tuple(out.shape)}"
            )
        if query.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) and out.dtype in (
            torch.float8_e4m3fn,
            torch.float8_e5m2,
        ):
            raise ValueError(
                "FP8 output is not supported for trtllm_ragged_attention_deepseek; "
                "use bfloat16 or float16 for out."
            )

    # --- LSE allocation (shared by all backends) ---
    if return_lse and lse is None:
        lse = torch.empty(
            query.shape[0],
            query.shape[1],
            device=query.device,
            dtype=torch.float32,
        )

    if backend == "cute-dsl":
        from .attention.cute_dsl.fmha import cute_dsl_fmha_ragged_prefill

        _SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float8_e4m3fn)
        assert query.dtype in _SUPPORTED_DTYPES, (
            f"cute-dsl backend only supports {_SUPPORTED_DTYPES}, got {query.dtype}"
        )
        # TODO: support device tensor scales to avoid D2H sync overhead
        assert not isinstance(bmm1_scale, torch.Tensor), (
            "cute-dsl backend does not support device tensor bmm1_scale"
        )
        assert not isinstance(bmm2_scale, torch.Tensor), (
            "cute-dsl backend does not support device tensor bmm2_scale"
        )
        assert sum(num_elts_per_sage_attn_blk) == 0, (
            "cute-dsl backend does not support sage attention scale factors"
        )
        _bmm1 = bmm1_scale
        _bmm2 = bmm2_scale

        # bmm1_scale = scale_q * scale_k * sm_scale (already fused by caller)
        # bmm2_scale = scale_v
        # Pass the fused value as sm_scale with scale_q=scale_k=1.0
        cute_dsl_fmha_ragged_prefill(
            q=query,
            k=key,
            v=value,
            o=out,
            qo_indptr=cum_seq_lens_q,
            kv_indptr=cum_seq_lens_kv,
            is_causal=is_causal,
            sm_scale=_bmm1,
            window_left=window_left,
            lse=lse if return_lse else None,
            attention_sinks=attention_sinks,
            scale_q=1.0,
            scale_k=1.0,
            scale_v=_bmm2,
            scale_o=1.0,
            max_qo_len=max_q_len,
            max_kv_len=max_kv_len,
            skip_softmax_threshold_scale_factor=skip_softmax_threshold_scale_factor,
            enable_pdl=enable_pdl,
        )
    else:
        # --- trtllm-gen backend ---
        run_func = get_trtllm_gen_fmha_module().trtllm_ragged_attention
        sm_count = get_device_sm_count(query.device)

        if isinstance(bmm1_scale, torch.Tensor):
            assert bmm1_scale.dtype == torch.float32
            bmm1_scale = bmm1_scale * log2e
        if isinstance(bmm2_scale, torch.Tensor):
            assert bmm2_scale.dtype == torch.float32

        workspace_size = workspace_buffer.numel() * workspace_buffer.element_size()
        sage_attn_sfs_q, sage_attn_sfs_k, sage_attn_sfs_p, sage_attn_sfs_v = (
            sage_attn_sfs
        )
        num_elts_sage_q, num_elts_sage_k, num_elts_sage_p, num_elts_sage_v = (
            num_elts_per_sage_attn_blk
        )
        run_func(
            out,
            query,
            key,
            value,
            workspace_buffer,
            seq_lens,
            max_q_len,
            max_kv_len,
            bmm1_scale,
            bmm2_scale,
            o_sf_scale,
            batch_size,
            window_left,
            cum_seq_lens_q,
            cum_seq_lens_kv,
            sm_count,
            enable_pdl,
            is_causal,
            workspace_size,
            attention_sinks,
            skip_softmax_threshold_scale_factor,
            lse,
            sage_attn_sfs_q,
            sage_attn_sfs_k,
            sage_attn_sfs_p,
            sage_attn_sfs_v,
            num_elts_sage_q,
            num_elts_sage_k,
            num_elts_sage_p,
            num_elts_sage_v,
        )

    if return_lse:
        assert lse is not None, (
            "lse assumed not None beyond this point when return_lse is True"
        )
        return out, lse
    else:
        return out


@flashinfer_api(trace=trtllm_batch_context_trace)
def trtllm_batch_context_with_kv_cache(
    query: torch.Tensor,
    kv_cache: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
    workspace_buffer: torch.Tensor,
    block_tables: torch.Tensor,
    seq_lens: torch.Tensor,
    max_q_len: int,
    max_kv_len: int,
    bmm1_scale: Union[float, torch.Tensor],
    bmm2_scale: Union[float, torch.Tensor],
    batch_size: int,
    cum_seq_lens_q: torch.Tensor,
    cum_seq_lens_kv: torch.Tensor,
    window_left: int = -1,
    out: Optional[Union[torch.Tensor, FP4Tensor]] = None,
    out_dtype: Optional[Union[torch.dtype, str]] = None,
    o_sf_scale: Optional[float] = None,
    o_sf_vec_size: Optional[int] = None,
    kv_layout: str = "HND",
    enable_pdl: Optional[bool] = None,
    sinks: Optional[List[torch.Tensor]] = None,
    kv_cache_sf: Optional[
        Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
    ] = None,
    skip_softmax_threshold_scale_factor: Optional[float] = None,
    uses_shared_paged_kv_idx: bool = True,
    causal: bool = True,
    lse: Optional[torch.Tensor] = None,
    return_lse: bool = False,
    multi_ctas_kv_counter_buffer: Optional[torch.Tensor] = None,
) -> Union[
    torch.Tensor, FP4Tensor, Tuple[Union[torch.Tensor, FP4Tensor], torch.Tensor]
]:
    """
    Parameters
    ----------
    query : torch.Tensor
        query tensor with shape [num_tokens, num_heads, head_dim]
    kv_cache : Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
        If kv_cache is a single tensor, it should be a tensor with shape [num_pages, 1 or 2, num_kv_heads, page_size, head_dim] if :attr:`kv_layout` is "HND",
        or [num_pages, 1 or 2, page_size, num_kv_heads, head_dim] if :attr:`kv_layout` is "NHD".
        If kv_cache is a tuple of two tensors, it should be a tuple of two tensors with shape [num_pages, num_kv_heads, page_size, head_dim] if :attr:`kv_layout` is "HND",
        or [num_pages, page_size, num_kv_heads, head_dim] if :attr:`kv_layout` is "NHD".
        The first tensor is the key cache, the second tensor is the value cache.

        **Contiguity requirements (trtllm-gen backend):**

        - The ``head_dim`` (last dim) **must** have stride 1. This is a TMA hardware constraint
        - The head and batch/page dims can have arbitrary strides.
    workspace_buffer : torch.Tensor
        Workspace used for trtllm-gen softmax stats/scratch.
    block_tables : torch.Tensor
        Page table of kv cache.
        When ``uses_shared_paged_kv_idx`` is True (default): shape ``[batch_size, max_num_pages_per_seq]``.
        When ``uses_shared_paged_kv_idx`` is False: shape ``[batch_size, 2, max_num_pages_per_seq]``
        where dim 1 distinguishes K (0) and V (1) page indices.
    seq_lens : torch.Tensor
        A uint32 1D tensor indicating the kv sequence length of each prompt. shape: ``[batch_size]``
    max_q_len : int
        max sequence length for query
    max_kv_len : int
        max sequence length for kv_cache
    bmm1_scale : Union[float, torch.Tensor]
        fused scale for bmm1 input.
        when using trtllm-gen backend, it can be a torch.Tensor with dtype torch.float32.
    bmm2_scale : Union[float, torch.Tensor]
        fused scale for bmm2 input.
        when using trtllm-gen backend, it can be a torch.Tensor with dtype torch.float32.
    batch_size : int
        batch size
    cum_seq_lens_q : torch.Tensor
        cumulative sequence length for query. shape: ``[batch_size + 1]``
    cum_seq_lens_kv : torch.Tensor
        cumulative sequence length for kv_cache. shape: ``[batch_size + 1]``
    window_left : int = -1
        The left (inclusive) window size for the attention window, when set to ``-1``, the window
        size will be set to the full length of the sequence. Defaults to ``-1``.
    out : Optional[Union[torch.Tensor, FP4Tensor]] = None
        output tensor, if not provided, will be allocated with ``out_dtype``, if ``out_dtype`` is not provided, will use the type of ``query``.
    out_dtype : Optional[Union[torch.dtype, str]] = None
        output dtype, if not provided, will use the type of ``out``. For nvfp4, use string ``nvfp4``.
    o_sf_scale : Optional[float] = None
        scale for nvfp4 output tensor scale factor.
    o_sf_vec_size : Optional[int] = None
        vector size for nvfp4 output tensor scale factor.
    enable_pdl : Optional[bool] = None
        Whether to enable Programmatic Dependent Launch (PDL). See https://docs.nvidia.com/cuda/cuda-c-programming-guide/#programmatic-dependent-launch-and-synchronization
        Defaults to ``None``, which means it will be enabled if the device supports PDL.
    kv_layout : str = "HND"
        Layout of kv-cache, can be "HND" or "NHD", default is "HND".
        For the trtllm-gen backend with NVFP4 KV cache, using ``NHD`` will trigger an
        automatic transpose and ``.contiguous()`` copy of both the KV data and block scale
        tensors to convert them to HND layout. This incurs extra memory allocation and
        data copy overhead. Use ``HND`` for better performance.
    sinks : Optional[List[torch.Tensor]] = None
        additional value per head in the denominator of the softmax.
    kv_cache_sf : Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None
        Per-block scale factors for NVFP4 KV cache. Accepts the same formats as
        ``kv_cache``:

        * a tuple ``(k_scales, v_scales)`` of 4-D tensors, each following
          :attr:`kv_layout`: ``[num_pages, num_kv_heads, page_size, head_dim // 16]`` for
          HND, or ``[num_pages, page_size, num_kv_heads, head_dim // 16]`` for NHD.
        * a single 5-D tensor with shape ``[num_pages, 2, ...]`` matching the layout of
          ``kv_cache``, split on dim 1 to yield k (index 0) and v (index 1) scales.

        Both tensors have dtype ``torch.float8_e4m3fn``. ``k_scales`` uses a linear
        (row-major) layout, while ``v_scales`` must use TRT-LLM's 4-token interleaved
        layout within each ``[page_size, head_dim // 16]`` tile. Use
        :func:`flashinfer.fp4_quantization.nvfp4_quantize_paged_kv_cache` to produce
        correctly formatted scale factors.

        When :attr:`kv_layout` is ``NHD``, scale tensors are transposed to HND internally
        (incurring a ``.contiguous()`` copy). Use ``HND`` for better performance.

        **Contiguity requirements (trtllm-gen backend):**

        - The last two dims (``page_size``, ``head_dim // 16``) **must** be contiguous
          (i.e., ``stride[-1] == 1`` and ``stride[-2] == head_dim // 16``). This is because
          the kernel reshapes them into ``(16, page_size * head_dim / 16 / 16)`` to satisfy
          TMA's 16-byte box width minimum.
        - The head and batch/page dims can have arbitrary strides.
    skip_softmax_threshold_scale_factor: Optional[float] = None
        threshold scale factor for skipping softmax operations.
        Providing a value for this parameter enables skip-softmax sparsity as described in: https://arxiv.org/abs/2512.12087
        If no value is provided, then standard attention is used.
        Setting the threshold to a higher value generally increases kernel performance at the cost of accuracy degradation.
        The actual threshold value equals the provided threshold_scale_factor divided by the context length.
    uses_shared_paged_kv_idx : bool = True
        Whether the K and V page indices are shared as a unified index.
        True (default) uses vLLM/FlashInfer layout with a 2D page table.
        False uses TRT-LLM layout with a 3D page table ``[batch_size, 2, max_num_pages_per_seq]``.
    causal : bool = True
        Whether to apply a causal mask. Set to ``False`` to request dense / bidirectional
        attention. For the TRTLLM-gen paged context path, non-causal currently requires
        ``window_left == -1``.
    lse : Optional[torch.Tensor] = None
        Optional pre-allocated buffer for Log-Sum-Exp values. Must have shape
        ``[num_tokens, num_qo_heads]`` with dtype ``torch.float32``. If provided,
        the tensor is filled regardless of :attr:`return_lse`. If ``return_lse`` is
        True and this is None, a buffer will be allocated.
    return_lse : bool = False
        Whether to return Log-Sum-Exp values. When True, returns ``(out, lse)``.
    multi_ctas_kv_counter_buffer : Optional[torch.Tensor] = None
        Optional separate trtllm-gen multi-CTA KV counter buffer. If not provided,
        FlashInfer creates a fresh zeroed internal buffer. When provided, it must be
        zero-initialized at allocation (e.g. via ``torch.zeros``); the kernel
        self-resets the counters after each launch, so it does not need to be
        re-zeroed between calls.
    Returns
    -------
    out: Union[torch.Tensor, FP4Tensor]
        output torch.Tensor or FP4Tensor.
    lse : torch.Tensor, optional
        Only returned when ``return_lse`` is True. Shape
        ``[num_tokens, num_qo_heads]`` with dtype ``torch.float32``.
    """

    if enable_pdl is None:
        enable_pdl = device_support_pdl(query.device)
    if not causal and window_left >= 0:
        raise NotImplementedError(
            "Sliding-window non-causal attention is not supported for trtllm-gen paged KV cache. "
            "Use window_left=-1 for dense bidirectional attention."
        )

    if isinstance(kv_cache, tuple):
        k_cache, v_cache = kv_cache
    else:
        if kv_cache.shape[1] == 1:
            k_cache, v_cache = kv_cache, kv_cache
        else:
            assert kv_cache.shape[1] == 2, (
                "When kv_cache is a single tensor, the second dimension must be 1 or 2"
            )
            # NOTE(Zihao): unbind transforms [num_pages, 2, ...] to ([num_pages, ...], [num_pages, ...])
            # it doesn't change underlying storage
            k_cache, v_cache = kv_cache.unbind(dim=1)

    if k_cache.dtype == torch.uint8 or v_cache.dtype == torch.uint8:
        if get_compute_capability(query.device) == (10, 7):
            raise ValueError("KV Cache NVFP4 is not supported on SM107")
        if kv_cache_sf is None:
            raise ValueError("kv_cache_sf must be provided for NVFP4 KV cache.")
    key_block_scales, value_block_scales = (
        _unpack_paged_kv_cache(kv_cache_sf, kv_layout)
        if kv_cache_sf is not None
        else (None, None)
    )

    # Convert NHD layout to HND if necessary
    if kv_layout == "NHD":
        k_cache = k_cache.transpose(-3, -2)
        v_cache = v_cache.transpose(-3, -2)
        if key_block_scales is not None:
            print(
                "[WARNING] NVFP4 KV cache with NHD layout will be converted to HND, "
                "incurring extra transpose and contiguous copy overhead. "
                "Use kv_layout='HND' for better performance."
            )
            k_cache = k_cache.contiguous()
            v_cache = v_cache.contiguous()
            key_block_scales = key_block_scales.transpose(-3, -2).contiguous()
            value_block_scales = value_block_scales.transpose(-3, -2).contiguous()

    run_func = get_trtllm_gen_fmha_module().trtllm_paged_attention_context
    sm_count = get_device_sm_count(query.device)

    if out_dtype == "nvfp4" or (out_dtype is None and isinstance(out, FP4Tensor)):
        assert query.dtype == torch.float8_e4m3fn, (
            "query must be fp8 when out_dtype is nvfp4."
        )
        assert o_sf_scale is not None
        assert o_sf_vec_size in [None, 16], "only o_sf_vec_size = 16 is supported"
        o_sf_vec_size = o_sf_vec_size or 16

        fp4_out_shape = query.shape[:-1] + (ceil_div(query.shape[-1], 2),)

        if isinstance(out, FP4Tensor):
            fp4_out_scale_shape = (
                out.scale.shape[0],
                round_up(query.shape[1] * query.shape[2] // o_sf_vec_size, 4),
            )
            out_scale_factor = out.scale
            o_sf_start_index = out.scale_start_index
            out = out.data
            # out_dtype may be None
            out_dtype = out_dtype or "nvfp4"
        elif out is None:
            fp4_out_scale_shape = (
                round_up(query.shape[0], 128),
                round_up(query.shape[1] * query.shape[2] // o_sf_vec_size, 4),
            )
            out_scale_factor = torch.empty(
                fp4_out_scale_shape, dtype=torch.float8_e4m3fn, device=query.device
            )
            o_sf_start_index = 0
            out = torch.empty(fp4_out_shape, dtype=torch.uint8, device=query.device)
        else:
            raise ValueError(f"Invalid out: {out}")

        assert out_dtype == "nvfp4"
        assert isinstance(out, torch.Tensor)

        # Use uint8 as the container dtype to compliant with next fp4 gemm.
        check_shape_dtype_device(out, fp4_out_shape, torch.uint8, query.device, "out")

        check_shape_dtype_device(
            out_scale_factor,
            fp4_out_scale_shape,
            torch.float8_e4m3fn,
            query.device,
            "out_scale_factor",
        )

        # Check o_sf_start_index is valid
        if (
            o_sf_start_index < 0
            or o_sf_start_index + out.shape[0] > out_scale_factor.shape[0]
        ):
            raise ValueError(
                f"o_sf_start_index is out of the valid range of out_scale_factor. "
                f"o_sf_start_index={o_sf_start_index}, out.shape[0]={out.shape[0]}, "
                f"out_scale_factor.shape[0]={out_scale_factor.shape[0]}"
            )

    elif isinstance(out_dtype, torch.dtype) or out_dtype is None:
        assert o_sf_scale is None
        assert o_sf_vec_size is None
        out_scale_factor = None
        o_sf_start_index = 0
        if out_dtype is None:
            out_dtype = out.dtype if out is not None else query.dtype
        out = out if out is not None else torch.empty_like(query, dtype=out_dtype)
        if out_dtype not in (query.dtype, torch.float16, torch.bfloat16):
            raise ValueError(f"Unsupported out_dtype: {out_dtype}")
        check_shape_dtype_device(out, query.shape, out_dtype, query.device, "out")
    else:
        raise ValueError(f"Invalid out_dtype: {out_dtype}")

    if isinstance(bmm1_scale, torch.Tensor):
        assert bmm1_scale.dtype == torch.float32
        bmm1_scale = bmm1_scale * log2e
    if isinstance(bmm2_scale, torch.Tensor):
        assert bmm2_scale.dtype == torch.float32
    _check_block_tables_shape(block_tables, uses_shared_paged_kv_idx)
    workspace_size = workspace_buffer.numel() * workspace_buffer.element_size()

    num_qo_heads = query.size(1)
    multi_ctas_kv_counter_buffer = _resolve_trtllm_gen_multi_ctas_kv_counter_buffer(
        multi_ctas_kv_counter_buffer,
        batch_size,
        num_qo_heads,
        sm_count,
        query.device,
    )
    lse_shape = (query.size(0), num_qo_heads)
    if lse is not None:
        check_shape_dtype_device(lse, lse_shape, torch.float32, query.device, "lse")
        lse_stride_tokens = lse.stride(0)
        lse_stride_heads = lse.stride(1)
    elif return_lse:
        lse = torch.empty(lse_shape, dtype=torch.float32, device=query.device)
        lse_stride_tokens = lse.stride(0)
        lse_stride_heads = lse.stride(1)
    else:
        lse_stride_tokens = 0
        lse_stride_heads = 0

    run_func(
        out,
        out_scale_factor,
        query,
        k_cache,
        v_cache,
        workspace_buffer,
        multi_ctas_kv_counter_buffer,
        block_tables,
        seq_lens,
        max_q_len,
        max_kv_len,
        bmm1_scale,
        bmm2_scale,
        o_sf_scale or -1.0,
        o_sf_vec_size or -1,
        o_sf_start_index,
        batch_size,
        window_left,
        cum_seq_lens_q,
        cum_seq_lens_kv,
        sm_count,
        enable_pdl,
        workspace_size,
        sinks,
        key_block_scales,
        value_block_scales,
        skip_softmax_threshold_scale_factor,
        uses_shared_paged_kv_idx,
        causal,
        lse,
        lse_stride_tokens,
        lse_stride_heads,
    )
    result_out = (
        out
        if out_dtype != "nvfp4"
        else FP4Tensor(out, out_scale_factor, o_sf_start_index, query.shape)
    )
    if return_lse:
        return result_out, lse
    return result_out


@functools.cache
def get_trtllm_fmha_v2_sm120_module():
    return gen_trtllm_fmha_v2_sm120_module().build_and_load()


@flashinfer_api(trace=fmha_v2_prefill_deepseek_trace)
def fmha_v2_prefill_deepseek(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    out: torch.Tensor,
    num_heads: int,
    head_dim: int,
    seq_len: int,
    scale_softmax: float,
    scale_bmm1: Optional[float] = None,
    scale_bmm2: Optional[float] = None,
    return_lse: bool = False,
    lse: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
    """
    Parameters
    ----------
    query : torch.Tensor
        query tensor with shape [batch_size, seq_len, num_heads, head_dim]
    key : torch.Tensor
        key tensor with shape [batch_size, seq_len, num_heads, head_dim]
    value : torch.Tensor
        value tensor with shape [batch_size, seq_len, num_heads, head_dim]
    out : torch.Tensor
        output tensor with shape [batch_size, seq_len, num_heads, head_dim]
    return_lse : bool
        whether to return the log-sum-exp of attention output
    num_heads : int
        number of heads
    head_dim : int
        head dimension
    seq_len : int
        sequence length
    scale_softmax : float
        scale for softmax
    scale_bmm1 : Optional[float]
        scale for bmm1
    scale_bmm2 : Optional[float]
        scale for bmm2
    lse : Optional[torch.Tensor]
        log-sum-exp of attention output
    Returns
    -------
    out: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
        output torch.Tensor or Tuple[torch.Tensor, torch.Tensor].
        If return_lse is True, the output will be a tuple of two tensors, the first is the output tensor, the second is the lse tensor.
        If return_lse is False, the output will be a single tensor.
    """
    if not is_sm12x_supported(query.device):
        raise ValueError("fmha_v2_prefill_deepseek is only supported on SM12x GPUs.")
    assert query.shape[3] == 192 and key.shape[3] == 192 and value.shape[3] == 128, (
        "currently only support deepseek r1 192 query and 128 value"
    )
    module = get_trtllm_fmha_v2_sm120_module()
    is_e4m3 = query.dtype == torch.float8_e4m3fn
    is_bf16_output = out.dtype == torch.bfloat16
    scale_softmax = (
        scale_softmax if scale_softmax is not None else 1.0 if is_e4m3 else 0.0
    )
    scale_bmm1 = scale_bmm1 if scale_bmm1 is not None else 1.0
    scale_bmm2 = scale_bmm2 if scale_bmm2 is not None else 1.0
    module.run(
        query,
        key,
        value,
        out,
        lse,
        num_heads,
        head_dim,
        seq_len,
        scale_softmax,
        scale_bmm1,
        scale_bmm2,
        is_e4m3,
        is_bf16_output,
    )
    if return_lse:
        return out, lse
    else:
        return out


@functools.cache
def get_trtllm_fmha_v2_module(
    input_layout: str, input_dtype: torch.dtype, output_dtype: torch.dtype = None
):
    return gen_fmha_v2_module(input_layout, input_dtype, output_dtype).build_and_load()


@flashinfer_api(trace=trtllm_fmha_v2_prefill_trace)
def trtllm_fmha_v2_prefill(
    qkv: Union[
        torch.Tensor,
        Tuple[torch.Tensor, torch.Tensor],
        Tuple[torch.Tensor, torch.Tensor, torch.Tensor],
    ],
    input_layout: str,
    workspace_buffer: torch.Tensor,
    seq_lens: torch.Tensor,
    max_q_len: int,
    max_kv_len: int,
    bmm1_scale: float,
    bmm2_scale: float,
    batch_size: int,
    cum_seq_lens_q: torch.Tensor,
    cum_seq_lens_kv: torch.Tensor,
    block_tables: Optional[torch.Tensor] = None,
    out: Optional[torch.Tensor] = None,
    out_dtype: Optional[Union[torch.dtype, str]] = None,
    sinks: Optional[List[torch.Tensor]] = None,
    pos_encoding_mode: str = None,
    logits_soft_cap_scale: Optional[float] = None,
    mask_mode: str = "causal",
    window_left: int = -1,
    chunked_attention_size: int = 0,
    save_softmax_stats: bool = False,
    skip_softmax_threshold_scale_factor: float = 0,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
    r"""TRT-LLM FMHAv2 prefill attention.

    Parameters
    ----------
    qkv
        Query/key/value input; expected format is determined by :attr:`input_layout`.
    input_layout
        Specifies the layout of the query/key/value tensors:
        - ``PACKED_QKV``: ``qkv`` is a single tensor of shape
          ``[num_tokens, 3, num_heads, head_dim]``.
        - ``CONTIGUOUS_Q_KV``: ``qkv`` is ``(Q, KV)`` where Q has shape
          ``[num_tokens, num_heads, head_dim]`` and KV has shape
          ``[num_tokens, 2, num_kv_heads, head_dim]``
          (``KV[:, 0, ...]`` is key, ``KV[:, 1, ...]`` is value).
        - ``Q_PAGED_KV_HND``: ``qkv`` is ``(Q, paged_KV)`` where Q has shape
          ``[num_tokens, num_heads, head_dim]`` and paged_KV has shape
          ``[num_pages, 2, num_kv_heads, page_size, head_dim]``
          (``paged_KV[:, 0, ...]`` is key cache, ``paged_KV[:, 1, ...]`` is value cache).
        - ``Q_PAGED_KV_NHD``: same as ``Q_PAGED_KV_HND`` but paged_KV shape is
          ``[num_pages, 2, page_size, num_kv_heads, head_dim]``.
        - ``SEPARATE_Q_K_V``: ``qkv`` is ``(Q, K, V)`` where Q has shape
          ``[num_tokens, num_heads, head_dim]`` and K, V have shape
          ``[num_tokens, num_kv_heads, head_dim]``.
    workspace_buffer
        The workspace buffer used for trtllm-gen softmax stats/scratch.
    seq_lens
        The KV sequence length of each request, shape: ``[batch_size]``.
    max_q_len
        The maximum sequence length for query.
    max_kv_len
        The maximum sequence length for KV cache.
    bmm1_scale
        The fused scale for BMM1 (QK^T) computation.
    bmm2_scale
        The fused scale for BMM2 (softmax(QK^T) * V) computation.
    batch_size
        The batch size.
    cum_seq_lens_q
        The cumulative sequence lengths for query, shape: ``[batch_size + 1]``.
    cum_seq_lens_kv
        The cumulative sequence lengths for KV cache, shape: ``[batch_size + 1]``.
    block_tables
        The page table for KV cache, shape: ``[batch_size, max_num_pages_per_seq]``.
        Required when using paged KV cache format.
    out
        The output tensor. If not provided, will be allocated with ``out_dtype``.
        If ``out_dtype`` is also not provided, will use the dtype of query.
    out_dtype
        The output dtype. If not provided, will use the dtype of ``out`` or query.
    sinks
        Additional value per head in the denominator of the softmax.
    pos_encoding_mode
        The position encoding mode, could be ``alibi``. Defaults to ``None``.
    logits_soft_cap_scale
        The logits soft cap scale. Defaults to ``None``, which means no soft cap.
    mask_mode
        The mask mode, could be ``causal``, ``sliding_window``, or ``chunked``.
        Defaults to ``causal``.
    window_left
        The left (inclusive) window size for the attention window, when set to ``-1``,
        the window size will be set to the full length of the sequence. Defaults to ``-1``.
        Only effective when :attr:`mask_mode` is ``sliding_window``.
    chunked_attention_size
        The chunked attention size. Defaults to ``0``, which means no chunked attention.
        Only effective when :attr:`mask_mode` is ``chunked``. Must be a power of 2.
        Chunked-causal masking is implemented only by the warp-specialized SM90
        (Hopper) kernels; it is **not supported on SM120 (Blackwell)**, where the
        tiled kernels honor sliding-window causal (``window_left``) only.
    save_softmax_stats
        Whether to save the softmax statistics. Defaults to ``False``.
    skip_softmax_threshold_scale_factor
        The factor of skip-softmax (Sparse Attention),
        Skip softmax and BMM2 when exp(local_max - global_max) < threshold,
        where threshold = skip_softmax_threshold_scale_factor / seqlen.
        Defaults to ``0`` (disabled).
    Returns
    -------
        If :attr:`save_softmax_stats` is ``False``, the attention output tensor.
        If :attr:`save_softmax_stats` is ``True``, a tuple of two tensors:
        * The attention output tensor.
        * The softmax statistics tensor (LSE).
    """

    if input_layout == "PACKED_QKV":
        assert isinstance(qkv, torch.Tensor)
        if qkv.dim() != 4 or qkv.shape[1] != 3:
            raise ValueError(
                f"PACKED_QKV expects shape [tokens, 3, num_heads, head_dim], got {tuple(qkv.shape)}"
            )
        query = qkv
        k_cache, v_cache = qkv, qkv  # placeholders
    elif input_layout == "CONTIGUOUS_Q_KV":
        assert isinstance(qkv, tuple)
        query, kv_cache = qkv[0], qkv[1]
        if kv_cache.dim() != 4 or kv_cache.shape[1] != 2:
            raise ValueError(
                f"CONTIGUOUS_Q_KV expects KV shape [tokens, 2, num_kv_heads, head_dim], got {tuple(kv_cache.shape)}"
            )
        k_cache = kv_cache
        v_cache = kv_cache  # placeholder (not used for this layout)
    elif input_layout == "Q_PAGED_KV_NHD":
        assert isinstance(qkv, tuple)
        query, paged_kv = qkv[0], qkv[1]
        if paged_kv.dim() != 5 or paged_kv.shape[1] != 2:
            raise ValueError(
                f"Q_PAGED_KV_NHD expects paged_KV shape [pages, 2, page_size, num_kv_heads, head_dim], got {tuple(paged_kv.shape)}"
            )
        k_cache, v_cache = paged_kv.unbind(dim=1)
    elif input_layout == "Q_PAGED_KV_HND":
        assert isinstance(qkv, tuple)
        query, paged_kv = qkv[0], qkv[1]
        if paged_kv.dim() != 5 or paged_kv.shape[1] != 2:
            raise ValueError(
                f"Q_PAGED_KV_HND expects paged_KV shape [pages, 2, num_kv_heads, page_size, head_dim], got {tuple(paged_kv.shape)}"
            )
        k_cache, v_cache = paged_kv.unbind(dim=1)
        if block_tables is None:
            raise ValueError(
                "block_tables is required for Q_PAGED_KV_HND input layout."
            )
    elif input_layout == "SEPARATE_Q_K_V":
        assert isinstance(qkv, tuple)
        query, k_cache, v_cache = qkv[0], qkv[1], qkv[2]
        if hasattr(torch, "float8_e4m3fn") and query.dtype == torch.float8_e4m3fn:
            raise ValueError(
                "FP8 (e4m3) is not supported for the SEPARATE_Q_K_V input layout. "
                "Use PACKED_QKV, CONTIGUOUS_Q_KV, or Q_PAGED_KV layout instead."
            )
        if logits_soft_cap_scale is not None and logits_soft_cap_scale > 0:
            raise ValueError(
                "Logits soft capping is not supported for the SEPARATE_Q_K_V input layout. "
                "Use PACKED_QKV, CONTIGUOUS_Q_KV, or Q_PAGED_KV layout instead."
            )
    else:
        raise ValueError(
            f"Unsupported input_layout: {input_layout!r}. Expected one of: "
            "PACKED_QKV, CONTIGUOUS_Q_KV, Q_PAGED_KV_HND, Q_PAGED_KV_NHD, SEPARATE_Q_K_V."
        )

    if input_layout == "PACKED_QKV":
        # Packed QKV: query is [tokens, 3, H, D]
        num_qo_heads = query.shape[2]
        page_size = 0  # Not applicable for packed layouts
        head_dim_v = query.shape[3]  # Assume same as head_dim_qk
    elif input_layout in ("Q_PAGED_KV_NHD", "Q_PAGED_KV_HND"):
        # Q is 3D: [tokens, H, D]
        num_qo_heads = query.shape[1]
        # Paged KV NHD: [num_pages, page_size, H_kv, D]
        # Paged KV HND: [num_pages, H_kv, page_size, D]
        page_size = k_cache.shape[1 if "NHD" in input_layout else 2]
        head_dim_v = v_cache.shape[3]
    elif input_layout == "CONTIGUOUS_Q_KV":
        # Q is 3D: [tokens, H, D], KV is 4D: [tokens, 2, H_kv, D]
        # k_cache holds the combined KV tensor
        num_qo_heads = query.shape[1]
        page_size = 0  # Not applicable for non-paged layouts
        head_dim_v = k_cache.shape[3]  # D from KV tensor
    else:
        # SEPARATE_Q_K_V: all 3D ragged [tokens, H, D]
        num_qo_heads = query.shape[1]
        page_size = 0  # Not applicable for non-paged layouts
        head_dim_v = v_cache.shape[2]

    uses_sliding_window = window_left is not None and window_left >= 0
    uses_chunked = chunked_attention_size is not None and chunked_attention_size > 0
    is_non_causal = mask_mode is not None and mask_mode.lower() == "padding"

    if (uses_sliding_window or uses_chunked) and is_non_causal:
        feature = "Sliding window" if uses_sliding_window else "Chunked"
        raise ValueError(
            f"{feature} attention requires causal masking. "
            f"The underlying kernel only supports SLIDING_OR_CHUNKED_CAUSAL mode. "
        )

    # Chunked-causal masking is only implemented by the warp-specialized (SM90)
    # kernels. The non-warp-specialized "tiled" kernels used on SM120 honor
    # ``sliding_window_size`` only and silently ignore ``chunked_attention_size``,
    # so reject chunked attention here rather than returning wrong results.
    # Sliding-window causal (``window_left``) IS supported on SM120.
    if uses_chunked and is_sm12x_supported(query.device):
        raise ValueError(
            "Chunked-causal attention (chunked_attention_size > 0) is not "
            "supported for FMHAv2 on SM120 (Blackwell). The SM120 tiled kernels "
            "implement sliding-window causal only. Use window_left for "
            "sliding-window masking, or run chunked attention on SM90 (Hopper)."
        )

    # Determine output dtype
    if out is not None:
        o_dtype = out.dtype
    elif out_dtype is not None:
        o_dtype = canonicalize_torch_dtype(out_dtype)
    else:
        o_dtype = query.dtype

    # Allocate output tensor if not provided
    # Use head_dim_v (actual value) not head_dim_vo (0 means "same as head_dim_qk")
    if out is None:
        out = torch.empty(
            (query.shape[0], num_qo_heads, head_dim_v),
            dtype=o_dtype,
            device=query.device,
        )

    is_e4m3 = (
        query.dtype == torch.float8_e4m3fn if hasattr(torch, "float8_e4m3fn") else False
    )
    if is_e4m3:
        if is_sm12x_supported(query.device):
            raise ValueError(
                "FP8 (e4m3) is not yet supported for FMHAv2 on SM120 (Blackwell). "
                "Use fp16 or bf16 instead."
            )
    # Always pass 1.0: the C++ auto-detect (scale_softmax == 0.0) handles FP16/INT8/E4M3
    # but has no branch for BF16, where 0.0 would zero out the softmax output.
    scale_softmax = 1.0
    softcapping_scale = (
        logits_soft_cap_scale if logits_soft_cap_scale is not None else 0.0
    )

    module = get_trtllm_fmha_v2_module(
        input_layout,
        query.dtype,
        o_dtype if query.dtype == torch.float8_e4m3fn else None,
    )

    # Allocate LSE tensor if saving softmax stats
    # Kernel writes in ragged (flat
    # ) format: [total_q_tokens, num_qo_heads, 2]
    # total_q_tokens == query.shape[0] for all ragged layouts
    lse = None
    if save_softmax_stats:
        lse = torch.empty(
            (query.shape[0], num_qo_heads, 2),
            dtype=torch.float32,
            device=query.device,
        )

    module.run(
        query,  # Q tensor
        k_cache,  # K tensor
        v_cache,  # V tensor
        out,  # Output tensor
        workspace_buffer,  # Workspace buffer
        workspace_buffer.numel()
        * workspace_buffer.element_size(),  # Workspace buffer size in bytes
        block_tables,
        page_size,
        seq_lens,  # Sequence length for kv_cache
        cum_seq_lens_q,  # Cumulative sequence length for query
        cum_seq_lens_kv,  # Cumulative sequence length for kv_cache
        input_layout.lower(),  # Input layout
        max_q_len,  # Max sequence length for query
        max_kv_len,  # Max sequence length for kv_cache
        batch_size,  # Batch size
        mask_mode.lower(),  # Attention mask type
        scale_softmax,  # Softmax scale
        bmm1_scale,  # BMM1 scale
        bmm2_scale,  # BMM2 scale (host float; encoded by C++ set_alpha)
        window_left,  # Window left
        chunked_attention_size,  # Chunked attention size
        pos_encoding_mode is not None
        and pos_encoding_mode.lower() == "alibi",  # Alibi mode
        softcapping_scale,  # Softcapping scale (0.0 = disabled)
        skip_softmax_threshold_scale_factor,  # threshold_scale_factor for skip-softmax (0.0 = disable)
        lse,  # Optional LSE tensor (None if not saving softmax stats)
        sinks,  # Optional sinks tensor
    )

    if save_softmax_stats:
        return out, lse
    else:
        return out
