"""
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 math
from types import SimpleNamespace
from typing import Any, List, Optional, Tuple, Union

import torch

from .api_logging import flashinfer_api
from .trace.templates.attention import (
    batch_pod_with_paged_kv_cache_run_trace,
    pod_with_paged_kv_cache_run_trace,
)
from .jit import gen_pod_module, gen_batch_pod_module
from .page import get_seq_lens
from .prefill import get_batch_prefill_module
from .quantization import packbits
from .utils import (
    MaskMode,
    PosEncodingMode,
    TensorLayout,
    _check_cached_qkv_data_type,
    _check_kv_layout,
    _check_pos_encoding_mode,
    _get_cache_alibi_slopes_buf,
    _get_cache_buf,
    _get_range_buf,
    _unpack_paged_kv_cache,
    canonicalize_torch_dtype,
    device_support_pdl,
)


@functools.cache
def get_pod_module(*args):
    module = gen_pod_module(*args).build_and_load()
    return SimpleNamespace(run_tensor=module.pod_with_kv_cache_tensor)


@functools.cache
def get_batch_pod_module(*args):
    module = gen_batch_pod_module(*args).build_and_load()
    return SimpleNamespace(run_tensor=module.batch_pod_with_kv_cache_tensor)


class PODWithPagedKVCacheWrapper:
    r"""Wrapper class for POD-Attention with paged kv-cache (first proposed in
    `<https://arxiv.org/abs/2410.18038>`_) for batch of requests.

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

    Examples
    --------
    >>> import torch
    >>> import flashinfer
    >>> num_layers = 32
    >>> num_qo_heads = 64
    >>> num_kv_heads = 8
    >>> head_dim = 128
    >>> max_num_pages = 128
    >>> page_size = 16
    >>> # allocate 128MB workspace buffer
    >>> workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device="cuda:0")
    >>> decode_wrapper = flashinfer.PODWithPagedKVCacheWrapper(
    ...     workspace_buffer, "NHD"
    ... )
    >>> batch_size = 7
    >>> kv_page_indices = torch.arange(max_num_pages).int().to("cuda:0")
    >>> kv_page_indptr = torch.tensor(
    ...     [0, 17, 29, 44, 48, 66, 100, 128], dtype=torch.int32, device="cuda:0"
    ... )
    >>> # 1 <= kv_last_page_len <= page_size
    >>> kv_last_page_len = torch.tensor(
    ...     [1, 7, 14, 4, 3, 1, 16], dtype=torch.int32, device="cuda:0"
    ... )
    >>> kv_cache_at_layer = [
    ...     torch.randn(
    ...         max_num_pages, 2, page_size, num_kv_heads, head_dim, dtype=torch.float16, device="cuda:0"
    ...     ) for _ in range(num_layers)
    ... ]
    >>> # create auxiliary data structures for batch decode attention
    >>> decode_wrapper.plan(
    ...     kv_page_indptr,
    ...     kv_page_indices,
    ...     kv_last_page_len,
    ...     num_qo_heads,
    ...     num_kv_heads,
    ...     head_dim,
    ...     page_size,
    ...     pos_encoding_mode="NONE",
    ...     data_type=torch.float16
    ... )
    >>> outputs = []
    >>> for i in range(num_layers):
    ...     q = torch.randn(batch_size, num_qo_heads, head_dim).half().to("cuda:0")
    ...     kv_cache = kv_cache_at_layer[i]
    ...     # compute batch decode attention, reuse auxiliary data structures for all layers
    ...     # TODO_AK: DEMONSTRATE USAGE OF POD
    ...     outputs.append(o)
    ...
    >>> outputs[0].shape
    torch.Size([7, 64, 128])

    Note
    ----
    To accelerate computation, FlashInfer's POD-Attention creates some
    auxiliary data structures, these data structures can be reused across multiple
    batch decode 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,
        paged_kv_indptr_buffer: Optional[torch.Tensor] = None,
        paged_kv_indices_buffer: Optional[torch.Tensor] = None,
        paged_kv_last_page_len_buffer: Optional[torch.Tensor] = None,
        jit_args: Optional[List[Any]] = None,
    ) -> None:
        r"""Constructor of :class:`PODWithPagedKVCacheWrapper`.

        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.

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

        use_cuda_graph : bool
            Whether to enable CUDAGraph for batch decode attention, if enabled, the
            auxiliary data structures will be stored as the provided buffers. The ``batch_size``
            cannot change during the lifecycle of this wrapper when CUDAGraph is enabled.

        paged_kv_indptr_buffer : Optional[torch.Tensor]
            The user reserved buffer on GPU to store the indptr of the paged kv cache, the size
            of the buffer should be ``[batch_size + 1]``.
            Only needed when ``use_cuda_graph`` is ``True``.

        paged_kv_indices_buffer : Optional[torch.Tensor]
            The user reserved buffer on GPU to store the page indices of the paged kv cache,
            should be large enough to store the maximum number of page indices
            (``max_num_pages``) during the lifecycle of this wrapper.
            Only needed when ``use_cuda_graph`` is ``True``.

        paged_kv_last_page_len_buffer : Optional[torch.Tensor]
            The user reserved buffer on GPU to store the number of entries in the last page, the
            size of the buffer should be ``[batch_size]``.
            Only needed when ``use_cuda_graph`` is ``True``.

        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.
        """
        _check_kv_layout(kv_layout)
        """
        if jit_args is not None:
            if use_tensor_cores:
                self._jit_module = get_batch_prefill_jit_module(
                    jit_args[0], gen_customize_batch_prefill_module("fa2", *jit_args)
                )
            else:
                self._jit_module = get_batch_decode_jit_module(
                    jit_args[0], gen_customize_batch_decode_module(*jit_args)
                )
        else:
        """
        # Override options. Only tensor core version is performant.
        use_tensor_cores = True
        self._jit_module: SimpleNamespace = None

        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(
            (8 * 1024 * 1024,),
            dtype=torch.uint8,
            pin_memory=True,
            device="cpu",
        )

        if use_cuda_graph:
            if not torch.is_tensor(paged_kv_indptr_buffer):
                raise ValueError(
                    "paged_kv_indptr_buffer should be a torch.Tensor in cudagraph mode"
                )
            if not torch.is_tensor(paged_kv_indices_buffer):
                raise ValueError(
                    "paged_kv_indices_buffer should be a torch.Tensor in cudagraph mode"
                )
            if not torch.is_tensor(paged_kv_last_page_len_buffer):
                raise ValueError(
                    "paged_kv_last_page_len_buffer should be a torch.Tensor in cudagraph mode"
                )
            self._fixed_batch_size = len(paged_kv_last_page_len_buffer)
            if len(paged_kv_indptr_buffer) != self._fixed_batch_size + 1:
                raise ValueError(
                    "The size of paged_kv_indptr_buffer should be batch_size + 1"
                )
        else:
            self._fixed_batch_size = 0

        self._paged_kv_indptr_buf = paged_kv_indptr_buffer
        self._paged_kv_indices_buf = paged_kv_indices_buffer
        self._paged_kv_last_page_len_buf = paged_kv_last_page_len_buffer
        self._use_tensor_cores = use_tensor_cores
        self._use_cuda_graph = use_cuda_graph

        if use_cuda_graph:
            # NOTE(Zihao): if once created, no need to update it in plan/run
            self._qo_indptr_buf = torch.arange(
                self._fixed_batch_size + 1,
                dtype=torch.int32,
                device=float_workspace_buffer.device,
            )

    @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.
        """
        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,
        indptr: torch.Tensor,
        indices: torch.Tensor,
        last_page_len: torch.Tensor,
        num_qo_heads: int,
        num_kv_heads: int,
        head_dim: int,
        page_size: int,
        pos_encoding_mode: str = "NONE",
        window_left: int = -1,
        q_data_type: Optional[Union[str, torch.dtype]] = "float16",
        kv_data_type: Optional[Union[str, torch.dtype]] = None,
        data_type: Optional[Union[str, torch.dtype]] = None,
        sm_scale: Optional[float] = None,
        rope_scale: Optional[float] = None,
        rope_theta: Optional[float] = None,
        non_blocking: bool = True,
    ) -> None:
        r"""Plan POD's batch decode for given problem specification.

        Parameters
        ----------
        indptr : torch.Tensor
            The indptr of the paged kv cache, shape: ``[batch_size + 1]``
        indices : torch.Tensor
            The page indices of the paged kv cache, shape: ``[qo_indptr[-1]]``
        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 : int
            The dimension of the heads
        page_size : int
            The page size of the paged kv cache
        pos_encoding_mode : str
            The position encoding applied inside attention kernels, could be
            ``NONE``/``ROPE_LLAMA`` (LLAMA style rotary embedding) /``ALIBI``.
            Defaults to ``NONE``.
        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``.
        q_data_type : Optional[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
            ``q_data_type``. Defaults to ``None``.
        data_type: Optional[Union[str, torch.dtype]]
            The data type of both the query and key/value tensors. Defaults to torch.float16.
            data_type is deprecated, please use q_data_type and kv_data_type instead.
        sm_scale : Optional[float]
            Softmax scale.  If ``None``, defaults to ``1 / sqrt(head_dim)``.  Cached on
            the wrapper and reused at :meth:`run` time.
        rope_scale : Optional[float]
            Scale factor applied during RoPE interpolation.  Only consulted when
            ``pos_encoding_mode != "NONE"``.  Defaults to ``1.0`` when ``None``.
        rope_theta : Optional[float]
            Base value for the RoPE frequencies.  Only consulted when
            ``pos_encoding_mode != "NONE"``.  Defaults to ``1e4`` when ``None``.
        non_blocking : bool
            Whether to copy the input tensors to the device asynchronously, defaults to ``True``.


        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 run calls.

        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``.
        """
        # Logits soft cap is not supported currently
        batch_size = len(last_page_len)
        logits_soft_cap = 0.0

        qo_indptr_host = _get_range_buf(batch_size + 1, "cpu")
        if self.is_cuda_graph_enabled:
            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
                    )
                )
            if len(indices) > len(self._paged_kv_indices_buf):
                raise ValueError(
                    "The size of indices should be less than or equal to the allocated buffer"
                )
            self._paged_kv_indptr_buf.copy_(indptr, non_blocking=non_blocking)
            self._paged_kv_last_page_len_buf.copy_(
                last_page_len, non_blocking=non_blocking
            )
            self._paged_kv_indices_buf[: len(indices)].copy_(
                indices, non_blocking=(indices.device == self.device) and non_blocking
            )
        else:
            self._paged_kv_indptr_buf = indptr.to(
                self.device, non_blocking=non_blocking
            )
            self._paged_kv_indices_buf = indices.to(
                self.device, non_blocking=non_blocking
            )
            self._paged_kv_last_page_len_buf = last_page_len.to(
                self.device, non_blocking=non_blocking
            )
            self._qo_indptr_buf = qo_indptr_host.to(
                self.device, non_blocking=non_blocking
            )

        indptr_host = indptr.to("cpu")
        last_page_len_host = last_page_len.to("cpu")

        if data_type is not None:
            if q_data_type is None:
                q_data_type = data_type
            if kv_data_type is None:
                kv_data_type = data_type

        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)

        self._cached_q_data_type = q_data_type
        self._cached_kv_data_type = kv_data_type
        kv_lens_arr_host = get_seq_lens(indptr_host, last_page_len_host, page_size)
        if self._jit_module is not None:
            self._cached_module = self._jit_module
        else:
            self._cached_module = get_batch_prefill_module(
                "fa2",
                q_data_type,
                kv_data_type,
                q_data_type,
                indptr.dtype,
                head_dim,  # head_dim_qk
                head_dim,  # head_dim_vo
                PosEncodingMode[pos_encoding_mode].value,
                window_left != -1,  # use_sliding_window
                logits_soft_cap > 0,  # use_logits_soft_cap
                False,  # use_fp16_qk_reduction
            )

        self._plan_info = self._cached_module.plan(
            self._float_workspace_buffer,
            self._int_workspace_buffer,
            self._pin_memory_int_workspace_buffer,
            qo_indptr_host,
            indptr_host,
            kv_lens_arr_host,
            batch_size,  # total_num_rows
            batch_size,
            num_qo_heads,
            num_kv_heads,
            page_size,
            self.is_cuda_graph_enabled,
            head_dim,
            head_dim,
            False,  # causal
            window_left,
            -1,  # fixed_split_size
            False,  # disable_split_kv
            0,  # num_colocated_ctas
            0,  # uniform_q_len
        )

        self._indptr_type = indptr.dtype
        self._pos_encoding_mode = pos_encoding_mode
        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

    @flashinfer_api(trace=pod_with_paged_kv_cache_run_trace)
    def run(
        self,
        # Main params (prefill and decode)
        q_p: torch.Tensor,
        k_p: torch.Tensor,
        v_p: torch.Tensor,
        q_d: torch.Tensor,
        paged_kv_cache_d: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
        # Prefill options
        custom_mask_p: Optional[torch.Tensor] = None,
        packed_custom_mask_p: Optional[torch.Tensor] = None,
        causal_p: bool = False,
        kv_layout_p: str = "NHD",
        pos_encoding_mode_p: str = "NONE",
        sm_scale_p: Optional[float] = None,
        window_left_p: int = -1,
        rope_scale_p: Optional[float] = None,
        rope_theta_p: Optional[float] = None,
        return_lse_p: bool = False,
        # Decode options
        custom_mask_d: Optional[torch.Tensor] = None,
        packed_custom_mask_d: Optional[torch.Tensor] = None,
        causal_d: bool = False,
        kv_layout_d: str = "NHD",
        pos_encoding_mode_d: str = "NONE",
        sm_scale_d: Optional[float] = None,
        window_left_d: int = -1,
        rope_scale_d: Optional[float] = None,
        rope_theta_d: Optional[float] = None,
        q_scale: Optional[float] = None,
        k_scale: Optional[float] = None,
        v_scale: Optional[float] = None,
        return_lse_d: bool = False,
        use_fp16_qk_reduction: bool = False,
        enable_pdl: Optional[bool] = None,
        *args,
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
        r"""Compute POD (Prefill-On-Decode) fused attention for a batch of requests.

        Single-shot fused attention that simultaneously runs single-request
        prefill (``q_p`` against ``k_p``/``v_p``) and batch-decode
        (``q_d`` against ``paged_kv_cache_d``) so prefill and decode tokens of a
        chunked-prefill / continuous-batching iteration share kernel launch and
        scheduling resources.

        Parameters
        ----------
        q_p : torch.Tensor
            Prefill query tensor, shape ``[qo_len, num_qo_heads, head_dim]``.
        k_p : torch.Tensor
            Prefill key tensor.  Layout matches ``kv_layout_p``.
        v_p : torch.Tensor
            Prefill value tensor.  Layout matches ``kv_layout_p``.
        q_d : torch.Tensor
            Decode query tensor, shape ``[batch_size, num_qo_heads, head_dim]``.
        paged_kv_cache_d : Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
            Paged KV cache for the decode requests.  Layout matches the wrapper's
            ``kv_layout`` set in :meth:`__init__`.
        custom_mask_p, packed_custom_mask_p : Optional[torch.Tensor]
            Optional dense / bit-packed custom mask for the prefill side.  See
            :func:`flashinfer.single_prefill_with_kv_cache` for the layout.
        causal_p : bool
            Whether to apply a causal mask to the prefill side.  Defaults to
            ``False``.
        kv_layout_p : str
            Layout of ``k_p`` and ``v_p``, either ``"NHD"`` or ``"HND"``.
        pos_encoding_mode_p : str
            Position-encoding mode for the prefill side.  Defaults to ``"NONE"``.
        sm_scale_p : Optional[float]
            Softmax scale for the prefill side.  Defaults to
            ``1 / sqrt(head_dim)``.
        window_left_p : int
            Left window size for sliding-window prefill; ``-1`` disables it.
        rope_scale_p, rope_theta_p : Optional[float]
            RoPE scale / theta for the prefill side.  Defaults to ``1.0`` / ``1e4``.
        return_lse_p : bool
            If ``True``, allocate an LSE buffer for the prefill kernel to write
            into.  Defaults to ``False``.  **Note: the buffer is allocated and
            filled by the kernel but is not currently returned to the caller**
            -- ``run`` always returns just ``(out_p, out_d)``.  This is a known
            limitation of the current POD wrapper (kernel API exists, Python
            wrapper does not yet plumb LSE through the return value).
        custom_mask_d, packed_custom_mask_d : Optional[torch.Tensor]
            Optional dense / bit-packed custom mask for the decode side.
        causal_d : bool
            Whether to apply a causal mask to the decode side.  Defaults to
            ``False``.
        kv_layout_d : str
            **Currently ignored**: the decode KV layout is always taken from
            the wrapper's ``kv_layout`` (set in :meth:`__init__`); this
            argument is accepted for signature symmetry with ``kv_layout_p``
            but the value is not consulted by the kernel.
        pos_encoding_mode_d : str
            **Currently ignored**: overridden by ``self._pos_encoding_mode``
            from :meth:`plan`.
        sm_scale_d : Optional[float]
            **Currently ignored**: overridden by ``self._sm_scale`` from
            :meth:`plan` (which itself defaults to ``1 / sqrt(head_dim)``).
        window_left_d : int
            **Currently ignored**: overridden by ``self._window_left`` from
            :meth:`plan`.
        rope_scale_d, rope_theta_d : Optional[float]
            **Currently ignored**: overridden by ``self._rope_scale`` /
            ``self._rope_theta`` from :meth:`plan`.
        q_scale, k_scale, v_scale : Optional[float]
            FP8 calibration scales applied to the decode side.  Folded into the
            decode ``sm_scale`` (``q_scale``, ``k_scale``) or the kernel output
            (``v_scale``).
        return_lse_d : bool
            If ``True``, allocate an LSE buffer for the decode kernel to write
            into.  Defaults to ``False``.  See ``return_lse_p`` -- same
            caveat: allocated but not returned.
        use_fp16_qk_reduction : bool
            Whether to accumulate ``QK`` in FP16 (lower precision, higher
            throughput).  Defaults to ``False``.
        enable_pdl : Optional[bool]
            Programmatic Dependent Launch toggle.  When ``None`` (default), the
            wrapper auto-detects support from the query device.
        *args
            Reserved for forward-compat with future kernel parameters.

        Returns
        -------
        Tuple[torch.Tensor, torch.Tensor]
            ``(out_p, out_d)``: the prefill output (shape
            ``[qo_len, num_qo_heads, head_dim]``) and the decode output (shape
            ``[batch_size, num_qo_heads, head_dim]``).  LSE tensors are
            **not** part of the return value even when ``return_lse_p`` /
            ``return_lse_d`` is ``True`` (see those parameter notes).
        """
        if enable_pdl is None:
            enable_pdl = device_support_pdl(q_p.device)

        # Currently unsupported
        logits_soft_cap_p = None
        logits_soft_cap_d = None
        # Prefill setup
        _check_pos_encoding_mode(pos_encoding_mode_p)
        _check_kv_layout(kv_layout_p)
        tmp_p = _get_cache_buf("pod_with_kv_cache_tmp", 32 * 1024 * 1024, q_p.device)
        if logits_soft_cap_p is None:
            logits_soft_cap_p = 0.0
        if sm_scale_p is None:
            sm_scale_p = 1.0 / math.sqrt(q_p.size(-1))
        if rope_scale_p is None:
            rope_scale_p = 1.0
        if rope_theta_p is None:
            rope_theta_p = 1e4
        if custom_mask_p is not None and packed_custom_mask_p is None:
            # create packed custom mask from custom mask
            packed_custom_mask_p = packbits(
                custom_mask_p.contiguous().view(-1), bitorder="little"
            )

        if packed_custom_mask_p is not None:
            mask_mode_p = MaskMode.CUSTOM.value
        else:
            if causal_p:
                mask_mode_p = MaskMode.CAUSAL.value
            else:
                mask_mode_p = MaskMode.NON_CAUSAL.value

        lse_p = None
        if return_lse_p:
            lse_p = torch.empty(
                (q_p.size(0), q_p.size(1)), dtype=torch.float32, device=q_p.device
            )

        out_p = torch.empty_like(q_p)

        # Decode setup
        k_cache_d, v_cache_d = _unpack_paged_kv_cache(paged_kv_cache_d, self._kv_layout)
        _check_cached_qkv_data_type(
            q_d, k_cache_d, self._cached_q_data_type, self._cached_kv_data_type
        )
        # TODO_AK: Where are these coming from?
        pos_encoding_mode_d = self._pos_encoding_mode
        window_left_d = self._window_left
        logits_soft_cap_d = self._logits_soft_cap
        sm_scale_d = self._sm_scale
        rope_scale_d = self._rope_scale
        rope_theta_d = self._rope_theta
        _check_pos_encoding_mode(pos_encoding_mode_d)
        # What are the above for and what are the below?
        if logits_soft_cap_d is None:
            logits_soft_cap_d = 0.0
        if sm_scale_d is None:
            head_dim = q_d.shape[-1]
            sm_scale_d = 1.0 / math.sqrt(head_dim)
        if q_scale is not None:
            sm_scale_d *= q_scale
        if k_scale is not None:
            sm_scale_d *= k_scale
        if rope_scale_d is None:
            rope_scale_d = 1.0
        if rope_theta_d is None:
            rope_theta_d = 1e4

        lse_d = None
        if return_lse_d:
            lse_d = torch.empty(
                (q_d.size(0), q_d.size(1)), dtype=torch.float32, device=q_d.device
            )
        out_d = torch.empty_like(q_d)

        module_getter = get_pod_module(
            # Prefill params
            q_p.dtype,
            k_p.dtype,
            q_p.dtype,
            q_p.shape[-1],
            PosEncodingMode[pos_encoding_mode_p].value,
            window_left_p >= 0,  # use_sliding_window
            logits_soft_cap_p > 0,  # use_logits_soft_cap
            use_fp16_qk_reduction,
            # Decode params
            # q_d.dtype,
            # self._cached_kv_data_type,
            # self._cached_q_data_type,
            self._indptr_type,
            # head_dim,  # head_dim_qk
            # head_dim,  # head_dim_vo
            PosEncodingMode[pos_encoding_mode_d].value,
            window_left_d != -1,  # use_sliding_window
            logits_soft_cap_d > 0,  # use_logits_soft_cap
        )
        module_getter.run_tensor(
            # Prefill params
            q_p,
            k_p,
            v_p,
            tmp_p,
            out_p,
            lse_p,
            mask_mode_p,
            TensorLayout[kv_layout_p].value,
            window_left_p,
            packed_custom_mask_p,
            _get_cache_alibi_slopes_buf(q_p.shape[1], q_p.device),
            logits_soft_cap_p,
            sm_scale_p,
            1.0 / rope_scale_p,
            1.0 / rope_theta_p,
            # Decode params
            self._float_workspace_buffer,
            self._int_workspace_buffer,
            self._plan_info,
            q_d,
            k_cache_d,
            v_cache_d,
            self._qo_indptr_buf,
            self._paged_kv_indptr_buf,
            self._paged_kv_indices_buf,
            self._paged_kv_last_page_len_buf,
            out_d,
            lse_d,
            MaskMode.NON_CAUSAL.value,
            TensorLayout[self._kv_layout].value,
            window_left_d,
            None,  # packed_custom_mask
            None,  # mask_indptr_buf
            _get_cache_alibi_slopes_buf(q_d.shape[1], q_d.device),
            logits_soft_cap_d,
            sm_scale_d,
            1.0 / rope_scale_d,
            1.0 / rope_theta_d,
            enable_pdl,
        )

        if v_scale is not None:
            out_d *= v_scale

        return (out_p, out_d)

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


class BatchPODWithPagedKVCacheWrapper:
    r"""Wrapper class for POD-Attention with paged kv-cache (first proposed in
    `<https://arxiv.org/abs/2410.18038>`_) for batch of requests.

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

    Examples
    --------
    >>> import torch
    >>> import flashinfer
    >>> num_layers = 8
    >>> num_qo_heads = 64
    >>> num_kv_heads = 8
    >>> head_dim = 128
    >>> max_num_pages = 128
    >>> device = 0
    >>> page_block_size = 1
    >>> causal = True
    >>> # allocate 128MB workspace buffer
    >>> workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device="cuda:0")
    >>> wrapper = flashinfer.BatchPODWithPagedKVCacheWrapper(
    ...     workspace_buffer, "NHD"
    ... )
    >>> # Prefill and decode parameters
    >>> p_qo_lens = [2048] * 2
    >>> d_qo_lens = [1] * 128
    >>> p_kv_lens = [2048] * 2
    >>> d_kv_lens = [2048] * 128
    >>> # Prefill plan inputs
    >>> p_seq_lens_blocks = torch.ceil(
    ...     torch.tensor(p_kv_lens, dtype=torch.int32) / page_block_size
    ... ).int()
    >>> p_q_indptr = torch.cat(
    ...     [torch.tensor([0]), torch.cumsum(torch.tensor(p_qo_lens), 0)], dim=0
    ... ).int()
    >>> p_kv_indptr = torch.cat(
    ...     [torch.tensor([0]), torch.cumsum(p_seq_lens_blocks, 0)], dim=0
    ... ).int()
    >>> kv_indices_p = torch.arange(0, p_kv_indptr[-1], device=device, dtype=torch.int32)
    >>> last_page_len_p = (p_seq_lens_blocks - 1) % page_block_size + 1
    >>> # Decode plan inputs
    >>> d_seq_lens_blocks = torch.ceil(
    ...     torch.tensor(d_kv_lens, dtype=torch.int32) / page_block_size
    ... ).int()
    >>> d_q_indptr = torch.cat(
    ...     [torch.tensor([0]), torch.cumsum(torch.tensor(d_qo_lens), 0)], dim=0
    ... ).int()
    >>> d_kv_indptr = torch.cat(
    ...     [torch.tensor([0]), torch.cumsum(d_seq_lens_blocks, 0)], dim=0
    ... ).int()
    >>> kv_indices_d = torch.arange(0, d_kv_indptr[-1], device=device, dtype=torch.int32)
    >>> last_page_len_d = (d_seq_lens_blocks - 1) % page_block_size + 1
    >>> # create auxiliary data structures for batch decode attention
    >>> wrapper.plan(
    ...     # Prefill params
    ...     p_q_indptr.to(device),
    ...     p_kv_indptr.to(device),
    ...     kv_indices_p.to(device),
    ...     last_page_len_p,
    ...     # Decode params
    ...     d_q_indptr.to(device),
    ...     d_kv_indptr.to(device),
    ...     kv_indices_d.to(device),
    ...     last_page_len_d,
    ...     # Common params
    ...     num_qo_heads=num_qo_heads,
    ...     num_kv_heads=num_kv_heads,
    ...     head_dim=head_dim,
    ...     page_size=page_block_size,
    ...     q_data_type=torch.bfloat16,
    ...     kv_data_type=torch.bfloat16,
    ... )
    >>> # Prefill input tensors
    >>> q_p = torch.rand(p_q_indptr[-1].item(), num_qo_heads, head_dim).to(
    ...     device, dtype=torch.bfloat16
    ... )
    >>> kv_p = torch.randn(p_kv_indptr[-1], 2, page_block_size, num_kv_heads, head_dim).to(
    ...     device, dtype=torch.bfloat16
    ... ).unbind(1)
    >>> # Decode input tensors
    >>> q_d = torch.rand(d_q_indptr[-1].item(), num_qo_heads, head_dim).to(
    ...     device, dtype=torch.bfloat16
    ... )
    >>> kv_d = torch.randn(d_kv_indptr[-1], 2, page_block_size, num_kv_heads, head_dim).to(
    ...     device, dtype=torch.bfloat16
    ... ).unbind(1)
    >>> for i in range(num_layers):
    ...     o_p_batch, o_d_batch = wrapper.run(
    ...         q_p,
    ...         kv_p,
    ...         q_d,
    ...         kv_d,
    ...         causal_p=causal,
    ...     )
    >>> print(o_p_batch.shape, o_d_batch.shape)
    torch.Size([4096, 64, 128]) torch.Size([128, 64, 128])

    Note
    ----
    To accelerate computation, FlashInfer's POD-Attention creates some
    auxiliary data structures, these data structures can be reused across multiple
    batch decode 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",
    ) -> None:
        r"""Constructor of :class:`BatchPODWithPagedKVCacheWrapper`.

        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.

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

        """
        _check_kv_layout(kv_layout)
        # Override options. Only tensor core version is performant.
        use_tensor_cores = True
        self._jit_module: SimpleNamespace = None

        self._kv_layout = kv_layout
        float_workspace_buffer_p, float_workspace_buffer_d = torch.chunk(
            float_workspace_buffer, 2, dim=0
        )
        self._float_workspace_buffer_p = float_workspace_buffer_p
        self._float_workspace_buffer_d = float_workspace_buffer_d
        self.device = float_workspace_buffer_p.device
        self._int_workspace_buffer_p = torch.empty(
            (8 * 1024 * 1024,), dtype=torch.uint8, device=self.device
        )
        self._int_workspace_buffer_d = torch.empty(
            (8 * 1024 * 1024,), dtype=torch.uint8, device=self.device
        )
        self._pin_memory_int_workspace_buffer_p = torch.empty(
            (8 * 1024 * 1024,),
            dtype=torch.uint8,
            pin_memory=True,
            device="cpu",
        )
        self._pin_memory_int_workspace_buffer_d = torch.empty(
            (8 * 1024 * 1024,),
            dtype=torch.uint8,
            pin_memory=True,
            device="cpu",
        )

        # SM aware scheduling buffer, requires SMs count + 2 entries
        dev_prop = torch.cuda.get_device_properties(self.device)
        self._sm_aware_sched = torch.empty(
            (dev_prop.multi_processor_count + 2), dtype=torch.int, device=self.device
        )

        self._fixed_batch_size = 0

        self._paged_kv_indptr_buf = None
        self._paged_kv_indices_buf = None
        self._paged_kv_last_page_len_buf = None
        self._use_tensor_cores = use_tensor_cores
        self._use_cuda_graph = False

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

    @flashinfer_api
    def plan(
        self,
        qo_indptr_p: torch.Tensor,
        kv_indptr_p: torch.Tensor,
        kv_indices_p: torch.Tensor,
        last_page_len_p: torch.Tensor,
        qo_indptr_d: torch.Tensor,
        kv_indptr_d: torch.Tensor,
        kv_indices_d: torch.Tensor,
        last_page_len_d: torch.Tensor,
        num_qo_heads: int,
        num_kv_heads: int,
        head_dim: int,
        page_size: int,
        pos_encoding_mode: str = "NONE",
        window_left: int = -1,
        q_data_type: Optional[Union[str, torch.dtype]] = "float16",
        kv_data_type: Optional[Union[str, torch.dtype]] = None,
        data_type: Optional[Union[str, torch.dtype]] = None,
        sm_scale: Optional[float] = None,
        rope_scale: Optional[float] = None,
        rope_theta: Optional[float] = None,
        non_blocking: bool = True,
    ) -> None:
        r"""Plan POD's batch prefill and decode for given problem specification.

        Parameters
        ----------
        qo_indptr_p : torch.Tensor
            The prefill indptr of the query/output tensor, shape: ``[batch_size + 1]``.
        kv_indptr_p : torch.Tensor
            The prefill indptr of the paged kv-cache, shape: ``[batch_size + 1]``.
        kv_indices_p : torch.Tensor
            The prefill page indices of the paged kv-cache, shape: ``[kv_indptr[-1]]``.
        last_page_len_p : torch.Tensor
            The number of entries in the last page of each prefill request in the paged
            kv-cache, shape: ``[batch_size]``.
        qo_indptr_d : torch.Tensor
            The decode indptr of the query/output tensor, shape: ``[batch_size + 1]``.
        kv_indptr_d : torch.Tensor
            The decode indptr of the paged kv-cache, shape: ``[batch_size + 1]``.
        kv_indices_d : torch.Tensor
            The decode page indices of the paged kv-cache, shape: ``[kv_indptr[-1]]``.
        last_page_len_d : torch.Tensor
            The number of entries in the last page of each decode 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 : int
            The dimension of the heads
        page_size : int
            The page size of the paged kv cache
        pos_encoding_mode : str
            The position encoding applied inside attention kernels, could be
            ``NONE``/``ROPE_LLAMA`` (LLAMA style rotary embedding) /``ALIBI``.
            Defaults to ``NONE``.
        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``.
        q_data_type : Optional[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
            ``q_data_type``. Defaults to ``None``.
        data_type: Optional[Union[str, torch.dtype]]
            The data type of both the query and key/value tensors. Defaults to torch.float16.
            data_type is deprecated, please use q_data_type and kv_data_type instead.
        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``.
        non_blocking : bool
            Whether to copy the input tensors to the device asynchronously, defaults to ``True``.

        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 run calls.

        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``.
        """
        # Logits soft cap is not supported currently
        logits_soft_cap = 0.0

        # Setup prefill params
        batch_size_p = len(last_page_len_p)
        qo_indptr_host_p = qo_indptr_p.to("cpu")
        total_num_rows_p = int(qo_indptr_host_p[-1])
        self._kv_indptr_buf_p = kv_indptr_p.to(self.device, non_blocking=non_blocking)
        self._kv_indices_buf_p = kv_indices_p.to(self.device, non_blocking=non_blocking)
        self._kv_last_page_len_buf_p = last_page_len_p.to(
            self.device, non_blocking=non_blocking
        )
        self._qo_indptr_buf_p = qo_indptr_host_p.to(
            self.device, non_blocking=non_blocking
        )
        kv_indptr_host_p = kv_indptr_p.to("cpu")
        last_page_len_host_p = last_page_len_p.to("cpu")
        kv_lens_arr_host_p = get_seq_lens(
            kv_indptr_host_p, last_page_len_host_p, page_size
        )

        if data_type is not None:
            if q_data_type is None:
                q_data_type = data_type
            if kv_data_type is None:
                kv_data_type = data_type

        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)

        self._cached_q_data_type = q_data_type
        self._cached_kv_data_type = kv_data_type
        if self._jit_module is not None:
            self._cached_module = self._jit_module
        else:
            self._cached_module = get_batch_prefill_module(
                "fa2",
                q_data_type,
                kv_data_type,
                q_data_type,
                kv_indptr_p.dtype,
                head_dim,  # head_dim_qk
                head_dim,  # head_dim_vo
                PosEncodingMode[pos_encoding_mode].value,
                window_left != -1,  # use_sliding_window
                logits_soft_cap > 0,  # use_logits_soft_cap
                False,  # use_fp16_qk_reduction
            )

        # Setup decode params
        batch_size_d = len(last_page_len_d)
        qo_indptr_host_d = qo_indptr_d.to("cpu")
        total_num_rows_d = int(qo_indptr_host_d[-1])
        self._kv_indptr_buf_d = kv_indptr_d.to(self.device, non_blocking=non_blocking)
        self._kv_indices_buf_d = kv_indices_d.to(self.device, non_blocking=non_blocking)
        self._kv_last_page_len_buf_d = last_page_len_d.to(
            self.device, non_blocking=non_blocking
        )
        self._qo_indptr_buf_d = qo_indptr_host_d.to(
            self.device, non_blocking=non_blocking
        )
        kv_indptr_host_d = kv_indptr_d.to("cpu")
        last_page_len_host_d = last_page_len_d.to("cpu")
        kv_lens_arr_host_d = get_seq_lens(
            kv_indptr_host_d, last_page_len_host_d, page_size
        )

        self._plan_info_d = self._cached_module.plan(
            self._float_workspace_buffer_d,
            self._int_workspace_buffer_d,
            self._pin_memory_int_workspace_buffer_d,
            qo_indptr_host_d,
            kv_indptr_host_d,
            kv_lens_arr_host_d,
            total_num_rows_d,  # total_num_rows
            batch_size_d,
            num_qo_heads,
            num_kv_heads,
            page_size,
            self.is_cuda_graph_enabled,
            head_dim,
            head_dim,
            False,  # causal
            window_left,
            -1,  # fixed_split_size
            False,  # disable_split_kv
            0,  # num_colocated_ctas
            0,  # uniform_q_len
        )

        num_colocated_ctas = self._plan_info_d[0]
        # Splitting small prefill causes unecessary bandwidth contention
        if total_num_rows_p > 1536:
            num_colocated_ctas = 0
        self._plan_info_p = self._cached_module.plan(
            self._float_workspace_buffer_p,
            self._int_workspace_buffer_p,
            self._pin_memory_int_workspace_buffer_p,
            qo_indptr_host_p,
            kv_indptr_host_p,
            kv_lens_arr_host_p,
            total_num_rows_p,  # total_num_rows
            batch_size_p,
            num_qo_heads,
            num_kv_heads,
            page_size,
            self.is_cuda_graph_enabled,
            head_dim,
            head_dim,
            False,  # causal
            window_left,
            -1,  # fixed_split_size
            False,  # disable_split_kv
            num_colocated_ctas,
            0,  # uniform_q_len
        )
        self._indptr_type = kv_indptr_p.dtype
        self._pos_encoding_mode = pos_encoding_mode
        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

    @flashinfer_api(trace=batch_pod_with_paged_kv_cache_run_trace)
    def run(
        self,
        # Main params (prefill and decode)
        q_p: torch.Tensor,
        paged_kv_cache_p: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
        q_d: torch.Tensor,
        paged_kv_cache_d: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
        # Prefill options
        custom_mask_p: Optional[torch.Tensor] = None,
        packed_custom_mask_p: Optional[torch.Tensor] = None,
        causal_p: bool = False,
        # Decode options
        q_scale: Optional[float] = None,
        k_scale: Optional[float] = None,
        v_scale: Optional[float] = None,
        # Common options
        return_lse: bool = False,
        use_fp16_qk_reduction: bool = False,
        enable_pdl: Optional[bool] = None,
    ) -> Union[
        Tuple[torch.Tensor, torch.Tensor],
        Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]],
    ]:
        r"""Compute batched POD (Prefill-On-Decode) fused attention.

        Fused single-shot attention that runs batched paged prefill (``q_p`` against
        ``paged_kv_cache_p``) and batched paged decode (``q_d`` against
        ``paged_kv_cache_d``) in the same kernel launch, sharing scheduling and
        execution resources.  All prefill / decode shape and policy parameters
        (pos-encoding, sliding window, sm_scale, RoPE, etc.) are taken from the
        cached values supplied to :meth:`plan`.

        Parameters
        ----------
        q_p : torch.Tensor
            Prefill query tensor, shape ``[qo_indptr_p[-1], num_qo_heads, head_dim]``.
        paged_kv_cache_p : Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
            Paged KV cache for the prefill requests.  Layout matches
            ``kv_layout`` set in :meth:`__init__`.
        q_d : torch.Tensor
            Decode query tensor, shape ``[batch_size_d, num_qo_heads, head_dim]``.
        paged_kv_cache_d : Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
            Paged KV cache for the decode requests.  Layout matches
            ``kv_layout`` set in :meth:`__init__`.
        custom_mask_p : Optional[torch.Tensor]
            Optional dense custom mask for the prefill side (auto-packed when
            ``packed_custom_mask_p`` is ``None``).
        packed_custom_mask_p : Optional[torch.Tensor]
            Optional bit-packed custom mask for the prefill side.
        causal_p : bool
            Whether to apply a causal mask to the prefill side.  Defaults to
            ``False``.
        q_scale, k_scale, v_scale : Optional[float]
            FP8 calibration scales applied to the decode side.  Folded into the
            decode ``sm_scale`` (``q_scale``, ``k_scale``) or the kernel output
            (``v_scale``).
        return_lse : bool
            Whether to return the LSE tensors for both prefill and decode.
            Defaults to ``False``.
        use_fp16_qk_reduction : bool
            Whether to accumulate ``QK`` in FP16 (lower precision, higher
            throughput).  Defaults to ``False``.
        enable_pdl : Optional[bool]
            Programmatic Dependent Launch toggle.  When ``None`` (default), the
            wrapper auto-detects support from the query device.

        Returns
        -------
        Union[Tuple[torch.Tensor, torch.Tensor], Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]]
            By default ``(out_p, out_d)``: the prefill output and the decode
            output.  When ``return_lse`` is ``True`` the return becomes
            ``((out_p, lse_p), (out_d, lse_d))``.
        """
        if enable_pdl is None:
            enable_pdl = device_support_pdl(q_p.device)

        # Currently unsupported
        logits_soft_cap_p = None
        logits_soft_cap_d = None
        # Prefill setup
        k_cache_p, v_cache_p = _unpack_paged_kv_cache(paged_kv_cache_p, self._kv_layout)
        _check_cached_qkv_data_type(
            q_p, k_cache_p, self._cached_q_data_type, self._cached_kv_data_type
        )
        # Get params from plan
        pos_encoding_mode_p = self._pos_encoding_mode
        window_left_p = self._window_left
        logits_soft_cap_p = self._logits_soft_cap
        sm_scale_p = self._sm_scale
        rope_scale_p = self._rope_scale
        rope_theta_p = self._rope_theta
        _check_pos_encoding_mode(pos_encoding_mode_p)
        if logits_soft_cap_p is None:
            logits_soft_cap_p = 0.0
        if sm_scale_p is None:
            head_dim = q_p.shape[-1]
            sm_scale_p = 1.0 / math.sqrt(head_dim)
        if rope_scale_p is None:
            rope_scale_p = 1.0
        if rope_theta_p is None:
            rope_theta_p = 1e4

        if custom_mask_p is not None and packed_custom_mask_p is None:
            # create packed custom mask from custom mask
            packed_custom_mask_p = packbits(
                custom_mask_p.contiguous().view(-1), bitorder="little"
            )

        if packed_custom_mask_p is not None:
            mask_mode_p = MaskMode.CUSTOM.value
        else:
            if causal_p:
                mask_mode_p = MaskMode.CAUSAL.value
            else:
                mask_mode_p = MaskMode.NON_CAUSAL.value

        lse_p = None
        if return_lse:
            lse_p = torch.empty(
                (q_p.size(0), q_p.size(1)), dtype=torch.float32, device=q_p.device
            )
        out_p = torch.empty_like(q_p)

        # Decode setup
        k_cache_d, v_cache_d = _unpack_paged_kv_cache(paged_kv_cache_d, self._kv_layout)
        _check_cached_qkv_data_type(
            q_d, k_cache_d, self._cached_q_data_type, self._cached_kv_data_type
        )
        # Get params from plan
        pos_encoding_mode_d = self._pos_encoding_mode
        window_left_d = self._window_left
        logits_soft_cap_d = self._logits_soft_cap
        sm_scale_d = self._sm_scale
        rope_scale_d = self._rope_scale
        rope_theta_d = self._rope_theta
        _check_pos_encoding_mode(pos_encoding_mode_d)
        if logits_soft_cap_d is None:
            logits_soft_cap_d = 0.0
        if sm_scale_d is None:
            head_dim = q_d.shape[-1]
            sm_scale_d = 1.0 / math.sqrt(head_dim)
        if q_scale is not None:
            sm_scale_d *= q_scale
        if k_scale is not None:
            sm_scale_d *= k_scale
        if rope_scale_d is None:
            rope_scale_d = 1.0
        if rope_theta_d is None:
            rope_theta_d = 1e4

        lse_d = None
        if return_lse:
            lse_d = torch.empty(
                (q_d.size(0), q_d.size(1)), dtype=torch.float32, device=q_d.device
            )
        out_d = torch.empty_like(q_d)

        module_getter = get_batch_pod_module(
            # Prefill params
            q_p.dtype,
            k_cache_p.dtype,
            q_p.dtype,
            q_p.shape[-1],
            PosEncodingMode[pos_encoding_mode_p].value,
            window_left_p >= 0,  # use_sliding_window
            logits_soft_cap_p > 0,  # use_logits_soft_cap
            use_fp16_qk_reduction,
            # Decode params
            self._indptr_type,
            PosEncodingMode[pos_encoding_mode_d].value,
            window_left_d != -1,  # use_sliding_window
            logits_soft_cap_d > 0,  # use_logits_soft_cap
        )
        module_getter.run_tensor(
            # Prefill params
            self._float_workspace_buffer_p,
            self._int_workspace_buffer_p,
            self._plan_info_p,
            q_p,
            k_cache_p,
            v_cache_p,
            self._qo_indptr_buf_p,
            self._kv_indptr_buf_p,
            self._kv_indices_buf_p,
            self._kv_last_page_len_buf_p,
            out_p,
            lse_p,
            mask_mode_p,
            TensorLayout[self._kv_layout].value,
            window_left_p,
            packed_custom_mask_p,  # packed_custom_mask
            None,  # mask_indptr_buf
            _get_cache_alibi_slopes_buf(q_p.shape[1], q_p.device),
            logits_soft_cap_p,
            sm_scale_p,
            1.0 / rope_scale_p,
            1.0 / rope_theta_p,
            # Decode params
            self._float_workspace_buffer_d,
            self._int_workspace_buffer_d,
            self._plan_info_d,
            q_d,
            k_cache_d,
            v_cache_d,
            self._qo_indptr_buf_d,
            self._kv_indptr_buf_d,
            self._kv_indices_buf_d,
            self._kv_last_page_len_buf_d,
            out_d,
            lse_d,
            MaskMode.NON_CAUSAL.value,
            TensorLayout[self._kv_layout].value,
            window_left_d,
            None,  # packed_custom_mask
            None,  # mask_indptr_buf
            _get_cache_alibi_slopes_buf(q_d.shape[1], q_d.device),
            logits_soft_cap_d,
            sm_scale_d,
            1.0 / rope_scale_d,
            1.0 / rope_theta_d,
            enable_pdl,
            self._sm_aware_sched,
        )

        if v_scale is not None:
            out_d *= v_scale

        return ((out_p, out_d), (lse_p, lse_d)) if return_lse else (out_p, out_d)

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