from __future__ import annotations
import math
import cudnn
import torch

from cuda.bindings import driver as cuda
from cudnn.datatypes import _torch_to_cudnn_data_type
from cudnn.api_base import APIBase, TupleDict
from typing import Optional

from ..utils import make_tensor_strided_like


class SlidingWindowAttention(APIBase):
    def __init__(
        self,
        sample_q: torch.Tensor,
        sample_k: torch.Tensor,
        sample_v: torch.Tensor,
        sample_o: torch.Tensor,
        sample_stats: Optional[torch.Tensor] = None,
        left_bound: int = 0,
        right_bound: int = 0,
        sample_seq_len_q: Optional[torch.Tensor] = None,
        sample_seq_len_kv: Optional[torch.Tensor] = None,
        sample_q_ragged_offset: Optional[torch.Tensor] = None,
        sample_k_ragged_offset: Optional[torch.Tensor] = None,
        sample_v_ragged_offset: Optional[torch.Tensor] = None,
        sample_o_ragged_offset: Optional[torch.Tensor] = None,
        sample_stats_ragged_offset: Optional[torch.Tensor] = None,
        max_seq_len_q: Optional[int] = None,
        max_seq_len_kv: Optional[int] = None,
        attn_scale: Optional[float] = None,
        intermediate_data_type: torch.dtype = torch.float32,
        compute_data_type: torch.dtype = torch.float32,
        cudnn_handle: Optional[cudnn.handle] = None,
    ):
        super().__init__()
        self._logger.debug("Entering __init__")

        self.sample_q = sample_q
        self.sample_k = sample_k
        self.sample_v = sample_v
        self.sample_o = sample_o
        self.is_infer = sample_stats is None
        self.sample_stats = self._pad_tensor_to_ndim(sample_stats, self.sample_o.ndim, "sample_stats") if sample_stats is not None else None
        self.left_bound = left_bound
        self.right_bound = right_bound

        self.sample_seq_len_q = self._pad_tensor_to_ndim(sample_seq_len_q, 4, "sample_seq_len_q")
        self.sample_seq_len_kv = self._pad_tensor_to_ndim(sample_seq_len_kv, 4, "sample_seq_len_kv")
        self.max_seq_len_q = max_seq_len_q
        self.max_seq_len_kv = max_seq_len_kv
        self.sample_q_ragged_offset = self._pad_tensor_to_ndim(sample_q_ragged_offset, 4, "sample_q_ragged_offset")
        self.sample_k_ragged_offset = self._pad_tensor_to_ndim(sample_k_ragged_offset, 4, "sample_k_ragged_offset")
        self.sample_v_ragged_offset = self._pad_tensor_to_ndim(sample_v_ragged_offset, 4, "sample_v_ragged_offset")
        self.sample_o_ragged_offset = self._pad_tensor_to_ndim(sample_o_ragged_offset, 4, "sample_o_ragged_offset")
        self.sample_stats_ragged_offset = (
            self._pad_tensor_to_ndim(sample_stats_ragged_offset, 4, "sample_stats_ragged_offset") if sample_stats_ragged_offset is not None else None
        )

        self.attn_scale = attn_scale if attn_scale is not None else 1.0 / math.sqrt(self.sample_q.shape[-1])
        self.intermediate_data_type = intermediate_data_type
        self.compute_data_type = compute_data_type

        self.dtype = None
        self.sm_version = None
        self.input_layout = None

        if cudnn_handle is None:
            self._logger.critical(
                "cudnn_handle not provided, creating new handle. This is not recommended as this is significant overhead and will occur for each SlidingWindowAttention object created."
            )
        self._cudnn_handle = cudnn_handle if cudnn_handle is not None else cudnn.create_handle()
        self._cudnn_swa_graph = None
        self._cudnn_compiled = False
        self._logger.debug(
            f"__init__ completed with args: sample_q {tuple(sample_q.shape)}, sample_k {tuple(sample_k.shape)}, sample_v {tuple(sample_v.shape)}, sample_o {tuple(sample_o.shape)}, sample_stats {None if sample_stats is None else tuple(sample_stats.shape)}, left_bound {left_bound}, right_bound {right_bound}, sample_seq_len_q {None if sample_seq_len_q is None else tuple(sample_seq_len_q.shape)}, sample_seq_len_kv {None if sample_seq_len_kv is None else tuple(sample_seq_len_kv.shape)}, max_seq_len_q {max_seq_len_q}, max_seq_len_kv {max_seq_len_kv}, is_infer {self.is_infer}, attn_scale {attn_scale}, intermediate_data_type {intermediate_data_type}, compute_data_type {compute_data_type}"
        )

    def _calculate_ragged_offsets(
        self,
        seq_len_q,
        seq_len_kv,
        sample_q,
        sample_k,
        sample_v,
        sample_o,
        sample_stats,
    ):
        """Calculate ragged offsets for fully packed THD layout."""

        def compute_exclusive_prefix_sum(tensor):
            assert tensor.shape[1:] == (
                1,
                1,
                1,
            ), f"Expected shape (b,1,1,1), got {tensor.shape}"
            return torch.cat(
                (
                    torch.zeros((1, 1, 1, 1), dtype=tensor.dtype, device=tensor.device),
                    torch.cumsum(tensor, dim=0),
                )
            )

        # Calculate ragged offsets
        q_ragged_offset = (compute_exclusive_prefix_sum(seq_len_q) * self.sample_q.stride()[0]).to(dtype=torch.int64)
        k_ragged_offset = (compute_exclusive_prefix_sum(seq_len_kv) * self.sample_k.stride()[0]).to(dtype=torch.int64)
        v_ragged_offset = (compute_exclusive_prefix_sum(seq_len_kv) * self.sample_v.stride()[0]).to(dtype=torch.int64)
        o_ragged_offset = (compute_exclusive_prefix_sum(seq_len_q) * self.sample_o.stride()[0]).to(dtype=torch.int64)
        stats_ragged_offset = (compute_exclusive_prefix_sum(seq_len_q) * self.sample_stats.stride()[0]).to(dtype=torch.int64) if not self.is_infer else None

        return (
            q_ragged_offset,
            k_ragged_offset,
            v_ragged_offset,
            o_ragged_offset,
            stats_ragged_offset,
        )

    def check_support(self) -> bool:
        self._logger.debug("Entering check_support")

        if not torch.cuda.is_available():
            raise RuntimeError("CUDA is not available")
        self.dtype = self.sample_q.dtype
        self.sm_version = torch.cuda.get_device_capability()[0] * 10 + torch.cuda.get_device_capability()[1]
        if self.sample_q.ndim == 4:
            self._logger.debug("Inferred bshd layout")
            self.input_layout = "bshd"
        elif self.sample_q.ndim == 3:
            self._logger.debug("Inferred thd layout")
            self.input_layout = "thd"
        else:
            raise ValueError(f"Invalid input layout: {self.sample_q.ndim}")

        swa_graph = cudnn.pygraph(
            io_data_type=_torch_to_cudnn_data_type(self.dtype),
            intermediate_data_type=_torch_to_cudnn_data_type(self.intermediate_data_type),
            compute_data_type=_torch_to_cudnn_data_type(self.compute_data_type),
            handle=self._cudnn_handle,
            sm_version=self.sm_version,
        )

        (
            self.q_cudnn,
            self.k_cudnn,
            self.v_cudnn,
            self.seq_len_q_cudnn,
            self.seq_len_kv_cudnn,
            self.q_ragged_offset_cudnn,
            self.k_ragged_offset_cudnn,
            self.v_ragged_offset_cudnn,
            self.o_ragged_offset_cudnn,
            self.stats_ragged_offset_cudnn,
        ) = (None, None, None, None, None, None, None, None, None, None)
        if self.input_layout == "bshd":
            b, h_q, s_q, d_qk = self.sample_q.shape
            b, h_kv, s_kv, d_qk = self.sample_k.shape
            b, h_kv, s_kv, d_v = self.sample_v.shape
            b, h_q, s_q, d_v = self.sample_o.shape

            if self.sample_q.shape != (b, h_q, s_q, d_qk):
                raise ValueError(f"Input shape mismatch: expected Q tensor shape {b, h_q, s_q, d_qk}, got {self.sample_q.shape}")
            if self.sample_k.shape != (b, h_kv, s_kv, d_qk):
                raise ValueError(f"Input shape mismatch: expected K tensor shape {b, h_kv, s_kv, d_qk}, got {self.sample_k.shape}")
            if self.sample_v.shape != (b, h_kv, s_kv, d_v):
                raise ValueError(f"Input shape mismatch: expected V tensor shape {b, h_kv, s_kv, d_v}, got {self.sample_v.shape}")
            if self.sample_o.shape != (b, h_q, s_q, d_v):
                raise ValueError(f"Output shape mismatch: expected O tensor shape {b, h_q, s_q, d_v}, got {self.sample_o.shape}")
            if not self.is_infer:
                self.sample_stats = self._pad_tensor_to_ndim(self.sample_stats, 4, "sample_stats")
                if self.sample_stats.shape != (b, h_q, s_q, 1):
                    raise ValueError(f"Output shape mismatch: expected Stats tensor shape {b, h_q, s_q, 1}, got {self.sample_stats.shape}")
            if self.sample_seq_len_q is not None or self.sample_seq_len_kv is not None:
                raise ValueError(
                    f"sample_seq_len_q and sample_seq_len_kv should be None for bshd layout, got {self.sample_seq_len_q} and {self.sample_seq_len_kv}"
                )
            if self.max_seq_len_q is not None or self.max_seq_len_kv is not None:
                raise ValueError(f"max_seq_len_q and max_seq_len_kv should be None for bshd layout, got {self.max_seq_len_q} and {self.max_seq_len_kv}")
            if (
                self.sample_q_ragged_offset is not None
                or self.sample_k_ragged_offset is not None
                or self.sample_v_ragged_offset is not None
                or self.sample_o_ragged_offset is not None
                or self.sample_stats_ragged_offset is not None
            ):
                raise ValueError(
                    f"sample_q_ragged_offset, sample_k_ragged_offset, sample_v_ragged_offset, sample_o_ragged_offset, and sample_stats_ragged_offset should be None for bshd layout, got {self.sample_q_ragged_offset}, {self.sample_k_ragged_offset}, {self.sample_v_ragged_offset}, {self.sample_o_ragged_offset}, and {self.sample_stats_ragged_offset}"
                )

            self.q_cudnn = swa_graph.tensor_like(self.sample_q)
            self.k_cudnn = swa_graph.tensor_like(self.sample_k)
            self.v_cudnn = swa_graph.tensor_like(self.sample_v)
        elif self.input_layout == "thd":
            t, h_q, d_qk = self.sample_q.shape
            t, h_kv, d_qk = self.sample_k.shape
            t, h_kv, d_v = self.sample_v.shape
            t, h_q, d_v = self.sample_o.shape

            if self.sample_q.shape != (t, h_q, d_qk):
                raise ValueError(f"Input shape mismatch: expected Q tensor shape {t, h_q, d_qk}, got {self.sample_q.shape}")
            if self.sample_k.shape != (t, h_kv, d_qk):
                raise ValueError(f"Input shape mismatch: expected K tensor shape {t, h_kv, d_qk}, got {self.sample_k.shape}")
            if self.sample_v.shape != (t, h_kv, d_v):
                raise ValueError(f"Input shape mismatch: expected V tensor shape {t, h_kv, d_v}, got {self.sample_v.shape}")
            if self.sample_o.shape != (t, h_q, d_v):
                raise ValueError(f"Output shape mismatch: expected O tensor shape {t, h_q, d_v}, got {self.sample_o.shape}")
            if not self.is_infer:
                self.sample_stats = self._pad_tensor_to_ndim(self.sample_stats, 3, "sample_stats")
                if self.sample_stats.shape != (t, h_q, 1):
                    raise ValueError(f"Output shape mismatch: expected Stats tensor shape {t, h_q, 1}, got {self.sample_stats.shape}")

            if self.sample_seq_len_q is None or self.sample_seq_len_kv is None:
                raise ValueError(
                    f"sample_seq_len_q and sample_seq_len_kv must be provided for thd layout, got {self.sample_seq_len_q} and {self.sample_seq_len_kv}"
                )
            if self.max_seq_len_q is None or self.max_seq_len_kv is None:
                raise ValueError(f"max_seq_len_q and max_seq_len_kv must be provided for thd layout, got {self.max_seq_len_q} and {self.max_seq_len_kv}")

            if (
                self.sample_q_ragged_offset is None
                or self.sample_k_ragged_offset is None
                or self.sample_v_ragged_offset is None
                or self.sample_o_ragged_offset is None
                or (self.sample_stats_ragged_offset is None and not self.is_infer)
            ):
                if (
                    self.sample_q_ragged_offset is not None
                    or self.sample_k_ragged_offset is not None
                    or self.sample_v_ragged_offset is not None
                    or self.sample_o_ragged_offset is not None
                    or (not self.is_infer and self.sample_stats_ragged_offset is not None)
                ):
                    raise ValueError(
                        f"sample_q_ragged_offset, sample_k_ragged_offset, sample_v_ragged_offset, sample_o_ragged_offset, and sample_stats_ragged_offset must be all provided or all None, got {self.sample_q_ragged_offset}, {self.sample_k_ragged_offset}, {self.sample_v_ragged_offset}, {self.sample_o_ragged_offset}, and {self.sample_stats_ragged_offset}"
                    )
                self._logger.info("Calculating ragged offsets internally assuming fully packed THD layout")
                (
                    self.sample_q_ragged_offset,
                    self.sample_k_ragged_offset,
                    self.sample_v_ragged_offset,
                    self.sample_o_ragged_offset,
                    self.sample_stats_ragged_offset,
                ) = self._calculate_ragged_offsets(
                    self.sample_seq_len_q,
                    self.sample_seq_len_kv,
                    self.sample_q,
                    self.sample_k,
                    self.sample_v,
                    self.sample_o,
                    self.sample_stats,
                )

            b = len(self.sample_seq_len_q)
            self.q_cudnn = swa_graph.tensor(
                dim=(b, h_q, self.max_seq_len_q, d_qk),
                stride=(
                    self.sample_q.stride()[0] * self.max_seq_len_q,
                    self.sample_q.stride()[1],
                    self.sample_q.stride()[0],
                    self.sample_q.stride()[2],
                ),
            )
            self.k_cudnn = swa_graph.tensor(
                dim=(b, h_kv, self.max_seq_len_kv, d_qk),
                stride=(
                    self.sample_k.stride()[0] * self.max_seq_len_kv,
                    self.sample_k.stride()[1],
                    self.sample_k.stride()[0],
                    self.sample_k.stride()[2],
                ),
            )
            self.v_cudnn = swa_graph.tensor(
                dim=(b, h_kv, self.max_seq_len_kv, d_v),
                stride=(
                    self.sample_v.stride()[0] * self.max_seq_len_kv,
                    self.sample_v.stride()[1],
                    self.sample_v.stride()[0],
                    self.sample_v.stride()[2],
                ),
            )
            self.seq_len_q_cudnn = swa_graph.tensor_like(self.sample_seq_len_q)
            self.seq_len_kv_cudnn = swa_graph.tensor_like(self.sample_seq_len_kv)
            self.q_ragged_offset_cudnn = swa_graph.tensor_like(self.sample_q_ragged_offset)
            self.k_ragged_offset_cudnn = swa_graph.tensor_like(self.sample_k_ragged_offset)
            self.v_ragged_offset_cudnn = swa_graph.tensor_like(self.sample_v_ragged_offset)
            self.o_ragged_offset_cudnn = swa_graph.tensor_like(self.sample_o_ragged_offset)
            if not self.is_infer:
                self.stats_ragged_offset_cudnn = swa_graph.tensor_like(self.sample_stats_ragged_offset)

            self.q_cudnn.set_ragged_offset(self.q_ragged_offset_cudnn)
            self.k_cudnn.set_ragged_offset(self.k_ragged_offset_cudnn)
            self.v_cudnn.set_ragged_offset(self.v_ragged_offset_cudnn)

        self.o_cudnn, self.stats_cudnn = swa_graph.sdpa(
            name="sdpa",
            q=self.q_cudnn,
            k=self.k_cudnn,
            v=self.v_cudnn,
            generate_stats=not self.is_infer,
            attn_scale=self.attn_scale,
            bias=None,
            use_alibi_mask=False,
            use_padding_mask=(self.input_layout == "thd"),
            seq_len_q=self.seq_len_q_cudnn,
            seq_len_kv=self.seq_len_kv_cudnn,
            diagonal_band_left_bound=self.left_bound,
            diagonal_band_right_bound=self.right_bound,
            diagonal_alignment=cudnn.diagonal_alignment.TOP_LEFT,
            dropout=None,
            rng_dump=None,
            paged_attention_k_table=None,
            paged_attention_v_table=None,
            paged_attention_max_seq_len_kv=None,
        )
        self.o_cudnn.set_output(True)

        if self.input_layout == "bshd":
            self.o_cudnn.set_dim(self.sample_o.shape).set_stride(self.sample_o.stride())
            if not self.is_infer:
                self.stats_cudnn.set_output(True).set_data_type(cudnn.data_type.FLOAT)
                self.stats_cudnn.set_dim(self.sample_stats.shape).set_stride(self.sample_stats.stride())
        elif self.input_layout == "thd":
            self.o_cudnn.set_dim((b, h_q, self.max_seq_len_q, d_v))
            self.o_cudnn.set_stride(
                (
                    self.sample_o.stride()[0] * self.max_seq_len_q,
                    self.sample_o.stride()[1],
                    self.sample_o.stride()[0],
                    self.sample_o.stride()[2],
                )
            )
            self.o_cudnn.set_ragged_offset(self.o_ragged_offset_cudnn)

            if not self.is_infer:
                self.stats_cudnn.set_output(True).set_data_type(cudnn.data_type.FLOAT)
                self.stats_cudnn.set_dim((b, h_q, self.max_seq_len_q, 1))
                self.stats_cudnn.set_stride(
                    (
                        self.sample_stats.stride()[0] * self.max_seq_len_q,
                        self.sample_stats.stride()[1],
                        self.sample_stats.stride()[0],
                        self.sample_stats.stride()[2],
                    )
                )
                self.stats_cudnn.set_ragged_offset(self.stats_ragged_offset_cudnn)

        try:
            swa_graph.validate()
        except cudnn.cudnnGraphNotSupportedError as e:
            self._logger.error(f"Graph not supported (cudnnGraphNotSupportedError): {e}")
            return False
        except Exception as e:
            self._logger.error(f"Graph not supported: {e}")
            return False

        self._cudnn_swa_graph = swa_graph
        self._is_supported = True
        self._logger.debug("check_support completed successfully")
        return True

    def compile(self) -> None:
        self._ensure_support_checked()

        self._cudnn_swa_graph.build_operation_graph()
        self._cudnn_swa_graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])
        self._cudnn_swa_graph.check_support()
        self._cudnn_swa_graph.build_plans()

        self._cudnn_compiled = True
        self._logger.debug("SlidingWindowAttention kernel compiled successfully")

    def execute(
        self,
        q_tensor: torch.Tensor,
        k_tensor: torch.Tensor,
        v_tensor: torch.Tensor,
        o_tensor: torch.Tensor,
        stats_tensor: Optional[torch.Tensor] = None,
        seq_len_q_tensor: Optional[torch.Tensor] = None,
        seq_len_kv_tensor: Optional[torch.Tensor] = None,
        q_ragged_offset_tensor: Optional[torch.Tensor] = None,
        k_ragged_offset_tensor: Optional[torch.Tensor] = None,
        v_ragged_offset_tensor: Optional[torch.Tensor] = None,
        o_ragged_offset_tensor: Optional[torch.Tensor] = None,
        stats_ragged_offset_tensor: Optional[torch.Tensor] = None,
        current_stream: Optional[cuda.CUstream] = None,
        cudnn_handle: Optional[cudnn.handle] = None,
        skip_compile: bool = False,
    ) -> None:
        self._logger.debug("Entering execute")
        cudnn_handle = self._cudnn_handle if cudnn_handle is None else cudnn_handle
        if current_stream is not None:
            self._logger.info("Overwriting cudnn_handle stream with provided cuda stream. Do not pass in current_stream if this is not intended.")
            cudnn.set_stream(cudnn_handle, current_stream)

        if skip_compile:
            raise NotImplementedError("cudnn sliding window attention kernel does not support skip_compile")
        if self._cudnn_swa_graph is None or not self._cudnn_compiled:
            raise ValueError("SlidingWindowAttention kernel not compiled")
        self._logger.debug("Executing with compiled kernel")

        self._logger.debug("Reshaping tensors to kernel expected format")
        stats_tensor = self._pad_tensor_to_ndim(stats_tensor, self.sample_o.ndim, "stats_tensor") if stats_tensor is not None else None
        seq_len_q_tensor = self._pad_tensor_to_ndim(seq_len_q_tensor, 4, "seq_len_q_tensor")
        seq_len_kv_tensor = self._pad_tensor_to_ndim(seq_len_kv_tensor, 4, "seq_len_kv_tensor")
        q_ragged_offset_tensor = self._pad_tensor_to_ndim(q_ragged_offset_tensor, 4, "q_ragged_offset_tensor")
        k_ragged_offset_tensor = self._pad_tensor_to_ndim(k_ragged_offset_tensor, 4, "k_ragged_offset_tensor")
        v_ragged_offset_tensor = self._pad_tensor_to_ndim(v_ragged_offset_tensor, 4, "v_ragged_offset_tensor")
        o_ragged_offset_tensor = self._pad_tensor_to_ndim(o_ragged_offset_tensor, 4, "o_ragged_offset_tensor")
        stats_ragged_offset_tensor = self._pad_tensor_to_ndim(stats_ragged_offset_tensor, 4, "stats_ragged_offset_tensor")

        if not self.is_infer and stats_tensor is None:
            raise ValueError(f"stats_tensor must be provided when compiled in non-inference mode, got {stats_tensor}")

        if self.input_layout == "thd":
            if seq_len_q_tensor is None or seq_len_kv_tensor is None:
                raise ValueError(f"seq_len_q_tensor and seq_len_kv_tensor must be provided for thd layout, got {seq_len_q_tensor} and {seq_len_kv_tensor}")
            if (
                q_ragged_offset_tensor is None
                or k_ragged_offset_tensor is None
                or v_ragged_offset_tensor is None
                or o_ragged_offset_tensor is None
                or (stats_ragged_offset_tensor is None and not self.is_infer)
            ):
                if (
                    q_ragged_offset_tensor is not None
                    or k_ragged_offset_tensor is not None
                    or v_ragged_offset_tensor is not None
                    or o_ragged_offset_tensor is not None
                    or (not self.is_infer and stats_ragged_offset_tensor is not None)
                ):
                    raise ValueError(
                        f"q_ragged_offset_tensor, k_ragged_offset_tensor, v_ragged_offset_tensor, o_ragged_offset_tensor, and stats_ragged_offset_tensor must be all provided or all None, got {q_ragged_offset_tensor}, {k_ragged_offset_tensor}, {v_ragged_offset_tensor}, {o_ragged_offset_tensor}, and {stats_ragged_offset_tensor}"
                    )
                self._logger.info("Calculating ragged offsets internally assuming fully packed THD layout")
                (
                    q_ragged_offset_tensor,
                    k_ragged_offset_tensor,
                    v_ragged_offset_tensor,
                    o_ragged_offset_tensor,
                    stats_ragged_offset_tensor,
                ) = self._calculate_ragged_offsets(
                    seq_len_q_tensor,
                    seq_len_kv_tensor,
                    self.sample_q,
                    self.sample_k,
                    self.sample_v,
                    self.sample_o,
                    self.sample_stats,
                )

        variant_pack = {
            self.q_cudnn: q_tensor,
            self.k_cudnn: k_tensor,
            self.v_cudnn: v_tensor,
            self.o_cudnn: o_tensor,
            self.seq_len_q_cudnn: seq_len_q_tensor,
            self.seq_len_kv_cudnn: seq_len_kv_tensor,
            self.q_ragged_offset_cudnn: q_ragged_offset_tensor,
            self.k_ragged_offset_cudnn: k_ragged_offset_tensor,
            self.v_ragged_offset_cudnn: v_ragged_offset_tensor,
            self.o_ragged_offset_cudnn: o_ragged_offset_tensor,
        }
        if not self.is_infer:
            variant_pack[self.stats_cudnn] = stats_tensor
            variant_pack[self.stats_ragged_offset_cudnn] = stats_ragged_offset_tensor

        workspace = torch.empty(
            self._cudnn_swa_graph.get_workspace_size(),
            device=q_tensor.device,
            dtype=torch.uint8,
        )
        self._cudnn_swa_graph.execute(variant_pack, workspace, handle=cudnn_handle)
        torch.cuda.synchronize()
        self._logger.debug("Executed successfully")

    def __call__(self, *args, **kwargs) -> None:
        self.execute(*args, skip_compile=True, **kwargs)


import logging

_logger = logging.getLogger(__name__)
_cache_of_SlidingWindowAttentionObjects = {}


def sliding_window_attention_wrapper(
    q_tensor: torch.Tensor,
    k_tensor: torch.Tensor,
    v_tensor: torch.Tensor,
    seq_len_q_tensor: Optional[torch.Tensor] = None,
    seq_len_kv_tensor: Optional[torch.Tensor] = None,
    q_ragged_offset_tensor: Optional[torch.Tensor] = None,
    k_ragged_offset_tensor: Optional[torch.Tensor] = None,
    v_ragged_offset_tensor: Optional[torch.Tensor] = None,
    o_ragged_offset_tensor: Optional[torch.Tensor] = None,
    stats_ragged_offset_tensor: Optional[torch.Tensor] = None,
    left_bound: int = 0,
    right_bound: int = 0,
    is_infer: bool = False,
    attn_scale: Optional[float] = None,
    o_dtype: Optional[torch.dtype] = None,
    intermediate_data_type: torch.dtype = torch.float32,
    compute_data_type: torch.dtype = torch.float32,
    cudnn_handle: Optional[cudnn.handle] = None,
    stream: Optional[cuda.CUstream] = None,
) -> TupleDict:
    o_tensor, stats_tensor = None, None
    o_dtype = o_dtype if o_dtype is not None else q_tensor.dtype
    if q_tensor.ndim == 3:  # thd
        _logger.debug("sliding_window_attention_wrapper: Creating empty output tensor o for thd layout")
        t, h_q, d = q_tensor.shape
        _, h_k, d_v = v_tensor.shape
        o_tensor = make_tensor_strided_like(q_tensor, (t, h_q, d_v), dtype=o_dtype, device=q_tensor.device)
        if not is_infer:
            _logger.debug("sliding_window_attention_wrapper: Creating empty output tensor stats for thd layout")
            stats_tensor = make_tensor_strided_like(q_tensor, (t, h_q, 1), dtype=torch.float32, device=q_tensor.device)
    else:  # bshd
        _logger.debug("sliding_window_attention_wrapper: Creating empty output tensor o for bshd layout")
        b, h_q, s_q, d = q_tensor.shape
        _, h_k, s_k, d_v = v_tensor.shape
        o_tensor = make_tensor_strided_like(q_tensor, (b, h_q, s_q, d_v), dtype=o_dtype, device=q_tensor.device)
        if not is_infer:
            _logger.debug("sliding_window_attention_wrapper: Creating empty output tensor stats for bshd layout")
            stats_tensor = make_tensor_strided_like(q_tensor, (b, h_q, s_q, 1), dtype=torch.float32, device=q_tensor.device)

    cache_key = (
        q_tensor.shape,
        k_tensor.shape,
        v_tensor.shape,
        seq_len_q_tensor.shape if seq_len_q_tensor is not None else None,
        seq_len_kv_tensor.shape if seq_len_kv_tensor is not None else None,
        q_ragged_offset_tensor.shape if q_ragged_offset_tensor is not None else None,
        k_ragged_offset_tensor.shape if k_ragged_offset_tensor is not None else None,
        v_ragged_offset_tensor.shape if v_ragged_offset_tensor is not None else None,
        o_ragged_offset_tensor.shape if o_ragged_offset_tensor is not None else None,
        (stats_ragged_offset_tensor.shape if stats_ragged_offset_tensor is not None else None),
        q_tensor.stride(),
        k_tensor.stride(),
        v_tensor.stride(),
        seq_len_q_tensor.stride() if seq_len_q_tensor is not None else None,
        seq_len_kv_tensor.stride() if seq_len_kv_tensor is not None else None,
        q_ragged_offset_tensor.stride() if q_ragged_offset_tensor is not None else None,
        k_ragged_offset_tensor.stride() if k_ragged_offset_tensor is not None else None,
        v_ragged_offset_tensor.stride() if v_ragged_offset_tensor is not None else None,
        o_ragged_offset_tensor.stride() if o_ragged_offset_tensor is not None else None,
        (stats_ragged_offset_tensor.stride() if stats_ragged_offset_tensor is not None else None),
        q_tensor.dtype,
        k_tensor.dtype,
        v_tensor.dtype,
        left_bound,
        right_bound,
        is_infer,
        attn_scale,
        intermediate_data_type,
        compute_data_type,
    )
    sliding_window_attention_object = None
    if cache_key in _cache_of_SlidingWindowAttentionObjects:
        _logger.debug("sliding_window_attention_wrapper: Using previously cached SlidingWindowAttention object")
        sliding_window_attention_object = _cache_of_SlidingWindowAttentionObjects[cache_key]

        sliding_window_attention_object.execute(
            q_tensor=q_tensor,
            k_tensor=k_tensor,
            v_tensor=v_tensor,
            o_tensor=o_tensor,
            stats_tensor=stats_tensor,
            seq_len_q_tensor=seq_len_q_tensor,
            seq_len_kv_tensor=seq_len_kv_tensor,
            q_ragged_offset_tensor=q_ragged_offset_tensor,
            k_ragged_offset_tensor=k_ragged_offset_tensor,
            v_ragged_offset_tensor=v_ragged_offset_tensor,
            o_ragged_offset_tensor=o_ragged_offset_tensor,
            stats_ragged_offset_tensor=stats_ragged_offset_tensor,
            current_stream=stream,
            cudnn_handle=cudnn_handle,
        )
    else:
        _logger.debug("sliding_window_attention_wrapper: No previously cached SlidingWindowAttention object found, creating new SlidingWindowAttention object")
        sliding_window_attention_object = SlidingWindowAttention(
            sample_q=q_tensor,
            sample_k=k_tensor,
            sample_v=v_tensor,
            sample_o=o_tensor,
            sample_stats=stats_tensor,
            sample_seq_len_q=seq_len_q_tensor,
            sample_seq_len_kv=seq_len_kv_tensor,
            sample_q_ragged_offset=q_ragged_offset_tensor,
            sample_k_ragged_offset=k_ragged_offset_tensor,
            sample_v_ragged_offset=v_ragged_offset_tensor,
            sample_o_ragged_offset=o_ragged_offset_tensor,
            sample_stats_ragged_offset=stats_ragged_offset_tensor,
            max_seq_len_q=(max(seq_len_q_tensor).item() if seq_len_q_tensor is not None else None),
            max_seq_len_kv=(max(seq_len_kv_tensor).item() if seq_len_kv_tensor is not None else None),
            left_bound=left_bound,
            right_bound=right_bound,
            attn_scale=attn_scale,
            intermediate_data_type=intermediate_data_type,
            compute_data_type=compute_data_type,
            cudnn_handle=cudnn_handle,
        )

        assert sliding_window_attention_object.check_support()
        sliding_window_attention_object.compile()
        sliding_window_attention_object.execute(
            q_tensor=q_tensor,
            k_tensor=k_tensor,
            v_tensor=v_tensor,
            o_tensor=o_tensor,
            stats_tensor=stats_tensor,
            seq_len_q_tensor=seq_len_q_tensor,
            seq_len_kv_tensor=seq_len_kv_tensor,
            q_ragged_offset_tensor=q_ragged_offset_tensor,
            k_ragged_offset_tensor=k_ragged_offset_tensor,
            v_ragged_offset_tensor=v_ragged_offset_tensor,
            o_ragged_offset_tensor=o_ragged_offset_tensor,
            stats_ragged_offset_tensor=stats_ragged_offset_tensor,
            current_stream=stream,
        )
        _cache_of_SlidingWindowAttentionObjects[cache_key] = sliding_window_attention_object

    return TupleDict(
        o_tensor=o_tensor,
        stats_tensor=stats_tensor,
    )
