lib/python3.12/site-packages/cudnn/native_sparse_attention/compression/api.py

from typing import Tuple, Optional
import math

from cuda.bindings import driver as cuda
import torch

import cutlass
import cutlass.cute as cute
from cutlass.cute.runtime import make_fake_stream
from cutlass.cute.typing import Int32

from cudnn.api_base import APIBase, TupleDict
from cudnn.datatypes import _convert_to_cutlass_data_type

from ..utils import make_tensor_strided_like
from .fmha import BlackwellFusedMultiHeadAttentionForward
from . import fmha_helpers as fmha_utils


class CompressionAttention(APIBase):
    def __init__(
        self,
        sample_q: torch.Tensor,
        sample_k: torch.Tensor,
        sample_v: torch.Tensor,
        sample_o: torch.Tensor,
        sample_lse: Optional[torch.Tensor] = None,
        sample_cum_seqlen_q: Optional[torch.Tensor] = None,
        sample_cum_seqlen_k: Optional[torch.Tensor] = None,
        qk_acc_dtype: torch.dtype = torch.float32,
        pv_acc_dtype: torch.dtype = torch.float32,
        mma_tiler_mn: Tuple[int, int] = (128, 128),
        is_persistent: bool = False,
        scale_q: float = 1.0,
        scale_k: float = 1.0,
        scale_v: float = 1.0,
        inv_scale_o: float = 1.0,
        scale_softmax: Optional[float] = None,
    ):
        super().__init__()
        self._kernel = BlackwellFusedMultiHeadAttentionForward

        self._logger.warning("CompressionAttention is an experimental API")
        self._logger.debug("Entering __init__")

        self.q_desc = self._make_tensor_desc(sample_q, name="sample_q")
        self.k_desc = self._make_tensor_desc(sample_k, name="sample_k")
        self.v_desc = self._make_tensor_desc(sample_v, name="sample_v")
        self.o_desc = self._make_tensor_desc(sample_o, name="sample_o")
        self.lse_desc = self._make_tensor_desc(sample_lse, name="sample_lse")
        self.enable_lse = sample_lse is not None
        self.cum_seqlen_q_desc = self._unpad_tensor_to_ndim(self._make_tensor_desc(sample_cum_seqlen_q, name="sample_cum_seqlen_q"), 1, "sample_cum_seqlen_q")
        self.cum_seqlen_k_desc = self._unpad_tensor_to_ndim(self._make_tensor_desc(sample_cum_seqlen_k, name="sample_cum_seqlen_k"), 1, "sample_cum_seqlen_k")
        self.max_cum_seqlen_q = int(sample_cum_seqlen_q.max().item()) if sample_cum_seqlen_q is not None else None
        self.max_cum_seqlen_k = int(sample_cum_seqlen_k.max().item()) if sample_cum_seqlen_k is not None else None

        # Types and kernel configuration
        self.qk_acc_dtype_torch = qk_acc_dtype
        self.pv_acc_dtype_torch = pv_acc_dtype
        self.mma_tiler_mn = mma_tiler_mn
        self.is_persistent = is_persistent

        # Scale config
        self.scale_q = scale_q
        self.scale_k = scale_k
        self.scale_v = scale_v
        self.inv_scale_o = inv_scale_o
        self.scale_softmax = scale_softmax

        # Derived attributes (populated in check_support)
        self.batch_size = None
        self.s_q = None
        self.s_k = None
        self.h_q = None
        self.h_k = None
        self.h_r = None
        self.head_dim = None
        self.problem_size = None
        self._compiled_kernel = None

        self._logger.debug(
            f"__init__ completed with args: sample_q {self.q_desc.shape}, sample_k {self.k_desc.shape}, sample_v {self.v_desc.shape}, sample_o {self.o_desc.shape}, sample_lse {self.lse_desc.shape if self.lse_desc is not None else 'None'}, sample_cum_seqlen_q {self.cum_seqlen_q_desc.shape if self.cum_seqlen_q_desc is not None else 'None'}, sample_cum_seqlen_k {self.cum_seqlen_k_desc.shape if self.cum_seqlen_k_desc is not None else 'None'}, qk_acc_dtype {qk_acc_dtype}, pv_acc_dtype {pv_acc_dtype}, mma_tiler_mn {mma_tiler_mn}, is_persistent {is_persistent}, scale_q {scale_q}, scale_k {scale_k}, scale_v {scale_v}, inv_scale_o {inv_scale_o}, scale_softmax {scale_softmax}"
        )

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

        # shape normalization and validation
        self._logger.debug("Checking shape normalization and validation")
        if self.q_desc.ndim == 4:
            self.input_layout = "B,H,S,D"

            b, h_qo, s_qo, d_qk = self.q_desc.shape
            b, h_kv, s_kv, d_qk = self.k_desc.shape
            b, h_kv, s_kv, d_v = self.v_desc.shape
            b, h_q, s_qo, d_v = self.o_desc.shape

            self._check_tensor_shape(self.q_desc, (b, h_qo, s_qo, d_qk), name="Q")
            self._check_tensor_shape(self.k_desc, (b, h_kv, s_kv, d_qk), name="K")
            self._check_tensor_shape(self.v_desc, (b, h_kv, s_kv, d_v), name="V")
            self._check_tensor_shape(self.o_desc, (b, h_q, s_qo, d_v), name="O")
            if self.enable_lse:

                self.lse_desc = self._unpad_tensor_to_ndim(self.lse_desc, 3, name="LSE")
                self._check_tensor_shape(self.lse_desc, (b, h_q, s_qo), name="LSE")
                self._check_tensor_stride(
                    self.lse_desc, stride=(h_q * s_qo, s_qo, 1), name="LSE", extra_error_msg="LSE tensor must be contiguous"
                )  # TODO @mingyangw: contiguous check
            if self.cum_seqlen_q_desc is not None or self.cum_seqlen_k_desc is not None:
                self._logger.warning("cum_seqlen_q and cum_seqlen_k are ignored for B,H,S,D layout")

            # Shapes
            self.batch_size = b
            self.s_q = s_qo
            self.s_kv = s_kv
            self.h_q = h_q
            self.h_kv = h_kv
            self.h_r = h_q // h_kv
            self.head_dim = d_qk
        elif self.q_desc.ndim == 3:
            self.input_layout = "T,H,D"

            t, h_q, d_qk = self.q_desc.shape
            t_kv, h_kv, d_qk = self.k_desc.shape  # T has been compressed for K and V
            t_kv, h_kv, d_v = self.v_desc.shape
            t, h_q, d_v = self.o_desc.shape

            self._check_tensor_shape(self.q_desc, (t, h_q, d_qk), name="Q")
            self._check_tensor_shape(self.k_desc, (t_kv, h_kv, d_qk), name="K")
            self._check_tensor_shape(self.v_desc, (t_kv, h_kv, d_v), name="V")
            self._check_tensor_shape(self.o_desc, (t, h_q, d_v), name="O")
            if self.enable_lse:
                self.lse_desc = self._unpad_tensor_to_ndim(self.lse_desc, 2, name="LSE")
                self._check_tensor_shape(self.lse_desc, (t, h_q), name="LSE")

            if self.cum_seqlen_q_desc is None or self.cum_seqlen_k_desc is None:
                raise ValueError(f"cum_seqlen_q and cum_seqlen_k must be provided for T,H,D layout, got {self.cum_seqlen_q_desc} and {self.cum_seqlen_k_desc}")

            if self.cum_seqlen_q_desc.ndim != 1 or self.cum_seqlen_k_desc.ndim != 1:
                raise ValueError(f"cum_seqlen_q and cum_seqlen_k must be 1D tensors, got {self.cum_seqlen_q_desc.ndim} and {self.cum_seqlen_k_desc.ndim}")
            self._check_dtype(self.cum_seqlen_q_desc, [torch.int32, torch.int64], name="cum_seqlen_q")
            self._check_dtype(self.cum_seqlen_k_desc, [torch.int32, torch.int64], name="cum_seqlen_k")
            if self.cum_seqlen_q_desc.shape[0] != self.cum_seqlen_k_desc.shape[0]:
                raise ValueError(
                    f"cum_seqlen_q and cum_seqlen_k must have the same length, got {self.cum_seqlen_q_desc.shape[0]} and {self.cum_seqlen_k_desc.shape[0]}"
                )

            self.batch_size = self.cum_seqlen_q_desc.shape[0] - 1
            self.s_q = None
            self.s_kv = None
            self.h_q = h_q
            self.h_kv = h_kv
            self.h_r = h_q // h_kv
            self.head_dim = d_qk

        else:
            raise ValueError(f"Invalid input layout: q must be rank-3 (T,H,D) or rank-4 (B,H,S,D), got {self.q_desc.ndim}")
        if d_qk != d_v:
            raise ValueError("D_qk must match D_v")
        if d_qk not in {32, 64, 128}:
            raise ValueError("Head dimension D_qk must be 32, 64, or 128")
        if h_q % h_kv != 0:
            raise ValueError("H_q must be divisible by H_k (GQA/MQA constraint)")

        self._logger.debug("Checking dtypes")
        in_dtype = self._check_dtype(self.q_desc, dtype=[torch.float16, torch.bfloat16, torch.float8_e4m3fn], name="Q")
        self._check_dtype(self.k_desc, dtype=in_dtype, name="K", extra_error_msg="K must have the same dtype as Q")
        self._check_dtype(self.v_desc, dtype=in_dtype, name="V", extra_error_msg="V must have the same dtype as Q")
        self._check_dtype(self.o_desc, dtype=[torch.float16, torch.bfloat16, torch.float8_e4m3fn], name="O")
        self._check_dtype(self.qk_acc_dtype_torch, dtype=torch.float32, name="qk_acc_dtype", extra_error_msg="qk_acc_dtype must be Float32")
        self._check_dtype(self.pv_acc_dtype_torch, dtype=torch.float32, name="pv_acc_dtype", extra_error_msg="pv_acc_dtype must be Float32")

        # Scale defaults
        if self.scale_softmax is None:
            self._logger.debug("No scale_softmax provided, using default 1/sqrt(d)")
            self.scale_softmax = 1.0 / math.sqrt(self.head_dim)

        # Environment checks
        self._logger.debug("Checking environment")
        if not torch.cuda.is_available():
            raise RuntimeError("CUDA is not available")
        device = torch.cuda.current_device()
        major, minor = torch.cuda.get_device_capability(device)
        compute_capability = major * 10 + minor
        if compute_capability < 100:
            raise RuntimeError(f"CompressionAttention requires SM100+ compute capability, but found SM{compute_capability} on device {device}")
        if compute_capability == 103:
            raise RuntimeError("cuteDSL is not supported on SM103")

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

    def compile(self) -> None:
        self._logger.debug("Entering compile")
        self._ensure_support_checked()
        if self._compiled_kernel is not None:
            self._logger.debug("Kernel already compiled; skipping recompilation")
            return

        fmha_kernel = self._kernel(
            _convert_to_cutlass_data_type(self.qk_acc_dtype_torch),
            _convert_to_cutlass_data_type(self.pv_acc_dtype_torch),
            (*self.mma_tiler_mn, self.head_dim),
            self.is_persistent,
            mask_type=fmha_utils.MaskType.COMPRESSED_CAUSAL_MASK,
        )

        # Scales
        log2_e = math.log2(math.exp(1.0))
        scale_softmax = self.scale_q * self.scale_k * self.scale_softmax
        scale_softmax_log2 = scale_softmax * log2_e
        scale_output = self.scale_v * self.inv_scale_o

        s_q = self.s_q if self.input_layout == "B,H,S,D" else self.max_cum_seqlen_q
        s_kv = self.s_kv if self.input_layout == "B,H,S,D" else self.max_cum_seqlen_k
        self.problem_size = (
            self.batch_size,
            s_q,
            s_q,
            s_kv,
            self.h_q,
            self.h_kv,
            self.head_dim,
        )
        fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False)

        self._logger.debug("Compiling CompressionAttention kernel with cute.compile")
        if self.input_layout == "B,H,S,D":
            _q_desc = self.q_desc.transpose(1, 2)
            _k_desc = self.k_desc.transpose(1, 2)
            _v_desc = self.v_desc.transpose(1, 2)
            _o_desc = self.o_desc.transpose(1, 2)
            _lse_desc = self.lse_desc.transpose(1, 2) if self.enable_lse else None
        elif self.input_layout == "T,H,D":
            _q_desc = self.q_desc.as_strided(size=(1, *self.q_desc.shape), stride=(self.q_desc.stride[0], *self.q_desc.stride))
            _k_desc = self.k_desc.as_strided(size=(1, *self.k_desc.shape), stride=(self.k_desc.stride[0], *self.k_desc.stride))
            _v_desc = self.v_desc.as_strided(size=(1, *self.v_desc.shape), stride=(self.v_desc.stride[0], *self.v_desc.stride))
            _o_desc = self.o_desc.as_strided(size=(1, *self.o_desc.shape), stride=(self.o_desc.stride[0], *self.o_desc.stride))
            # lse_local = self.sample_lse.as_strided(size=(1, *self.sample_lse.shape), stride=(self.sample_lse.stride()[0], *self.sample_lse.stride()))
            _lse_desc = self.lse_desc.unsqueeze(0) if self.enable_lse else None
        else:
            raise NotImplementedError(f"Invalid input layout: {self.input_layout}")

        # breakpoint()
        _compiled_kernel = cute.compile(
            fmha_kernel,
            Q=self._make_fake_cute_tensor_from_desc(_q_desc, assumed_align=16),
            K=self._make_fake_cute_tensor_from_desc(_k_desc, assumed_align=16),
            V=self._make_fake_cute_tensor_from_desc(_v_desc, assumed_align=16),
            O=self._make_fake_cute_tensor_from_desc(_o_desc, assumed_align=16),
            problem_size=self.problem_size,
            cum_seqlen_q=(self._make_fake_cute_tensor_from_desc(self.cum_seqlen_q_desc, assumed_align=16) if self.input_layout == "T,H,D" else None),
            cum_seqlen_k=(self._make_fake_cute_tensor_from_desc(self.cum_seqlen_k_desc, assumed_align=16) if self.input_layout == "T,H,D" else None),
            LSE=(self._make_fake_cute_tensor_from_desc(_lse_desc, assumed_align=16) if self.enable_lse else None),
            scale_softmax_log2=scale_softmax_log2,
            scale_softmax=scale_softmax,
            scale_output=scale_output,
            window_size_left=None,
            window_size_right=Int32(0),
            stream=fake_stream,
            options="--enable-tvm-ffi",
        )

        def tensor_api(
            q_tensor,
            k_tensor,
            v_tensor,
            o_tensor,
            problem_size,
            cum_seqlen_q,
            cum_seqlen_k,
            lse_tensor,
            scale_softmax_log2,
            scale_softmax,
            scale_output,
            window_size_left,
            window_size_right,
            stream,
        ):
            if self.enable_lse:
                lse_tensor = self._unpad_tensor_to_ndim(lse_tensor, self.o_desc.ndim - 1, "lse_tensor")
            if self.input_layout == "B,H,S,D":
                q_tensor = q_tensor.transpose(1, 2)
                k_tensor = k_tensor.transpose(1, 2)
                v_tensor = v_tensor.transpose(1, 2)
                o_tensor = o_tensor.transpose(1, 2)
                lse_tensor = lse_tensor.transpose(1, 2) if self.enable_lse else None
            elif self.input_layout == "T,H,D":
                q_tensor = q_tensor.as_strided(size=(1, *q_tensor.shape), stride=(q_tensor.stride()[0], *q_tensor.stride()))
                k_tensor = k_tensor.as_strided(size=(1, *k_tensor.shape), stride=(k_tensor.stride()[0], *k_tensor.stride()))
                v_tensor = v_tensor.as_strided(size=(1, *v_tensor.shape), stride=(v_tensor.stride()[0], *v_tensor.stride()))
                o_tensor = o_tensor.as_strided(size=(1, *o_tensor.shape), stride=(o_tensor.stride()[0], *o_tensor.stride()))
                lse_tensor = lse_tensor.unsqueeze(0) if self.enable_lse else None
                cum_seqlen_q = self._unpad_tensor_to_ndim(cum_seqlen_q, 1, "cum_seqlen_q")
                cum_seqlen_k = self._unpad_tensor_to_ndim(cum_seqlen_k, 1, "cum_seqlen_k")

            return _compiled_kernel(
                q_tensor,
                k_tensor,
                v_tensor,
                o_tensor,
                problem_size,
                cum_seqlen_q,
                cum_seqlen_k,
                lse_tensor,
                scale_softmax_log2,
                scale_softmax,
                scale_output,
                window_size_left,
                window_size_right,
                stream,
            )

        self._compiled_kernel = tensor_api

        self._logger.debug("Kernel compiled successfully")

    def execute(
        self,
        q_tensor: torch.Tensor,
        k_tensor: torch.Tensor,
        v_tensor: torch.Tensor,
        o_tensor: torch.Tensor,
        lse_tensor: Optional[torch.Tensor] = None,
        cum_seqlen_q_tensor: Optional[torch.Tensor] = None,
        cum_seqlen_k_tensor: Optional[torch.Tensor] = None,
        current_stream: Optional[cuda.CUstream] = None,
        scale_q: Optional[float] = None,
        scale_k: Optional[float] = None,
        scale_v: Optional[float] = None,
        inv_scale_o: Optional[float] = None,
        scale_softmax: Optional[float] = None,
    ) -> None:
        self._logger.debug("Entering execute")
        current_stream = self._get_default_stream(current_stream)

        if self.enable_lse:
            if lse_tensor is None:
                raise ValueError("kernel was compiled with lse_tensor provided, but lse_tensor was not provided during execute")

        if self.input_layout == "T,H,D":
            if cum_seqlen_q_tensor is None or cum_seqlen_k_tensor is None:
                raise ValueError("cum_seqlen_q_tensor and cum_seqlen_k_tensor must be provided during execute for T,H,D layout")

        # Scale values
        scale_q = self.scale_q if scale_q is None else scale_q
        scale_k = self.scale_k if scale_k is None else scale_k
        scale_v = self.scale_v if scale_v is None else scale_v
        inv_scale_o = self.inv_scale_o if inv_scale_o is None else inv_scale_o
        scale_softmax = self.scale_softmax if scale_softmax is None else scale_softmax
        log2_e = math.log2(math.e)
        scale_softmax_val = scale_q * scale_k * scale_softmax
        scale_softmax_log2_val = scale_softmax_val * log2_e
        scale_output_val = scale_v * inv_scale_o

        if self._compiled_kernel is None:
            raise ValueError("CompressionAttention kernel not compiled")
        self._logger.debug("Executing with compiled kernel")
        self._compiled_kernel(
            q_tensor=q_tensor,
            k_tensor=k_tensor,
            v_tensor=v_tensor,
            o_tensor=o_tensor,
            problem_size=self.problem_size,
            cum_seqlen_q=(cum_seqlen_q_tensor if self.input_layout == "T,H,D" else None),
            cum_seqlen_k=(cum_seqlen_k_tensor if self.input_layout == "T,H,D" else None),
            lse_tensor=lse_tensor,
            scale_softmax_log2=scale_softmax_log2_val,
            scale_softmax=scale_softmax_val,
            scale_output=scale_output_val,
            window_size_left=None,
            window_size_right=Int32(0),
            stream=current_stream,
        )
        self._logger.debug("Executed with compiled kernel successfully")


import logging

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


def compression_attention_wrapper(
    q_tensor: torch.Tensor,
    k_tensor: torch.Tensor,
    v_tensor: torch.Tensor,
    cum_seqlen_q_tensor: Optional[torch.Tensor] = None,
    cum_seqlen_k_tensor: Optional[torch.Tensor] = None,
    enable_lse: bool = False,
    o_dtype: Optional[torch.dtype] = None,
    qk_acc_dtype: torch.dtype = torch.float32,
    pv_acc_dtype: torch.dtype = torch.float32,
    mma_tiler_mn: Tuple[int, int] = (128, 128),
    is_persistent: bool = False,
    scale_q: float = 1.0,
    scale_k: float = 1.0,
    scale_v: float = 1.0,
    inv_scale_o: float = 1.0,
    scale_softmax: Optional[float] = None,
    stream: Optional[cuda.CUstream] = None,
) -> TupleDict:
    """
    Compression Attention Wrapper that returns output (and optionally LSE) tensors.

    Returns:
        TupleDict: (o_tensor, lse_tensor | None)
    """
    _logger.debug("compression_attention_wrapper: Creating empty output tensor o and optional lse")

    o_tensor, lse_tensor = None, None
    o_dtype = o_dtype if o_dtype is not None else q_tensor.dtype
    if q_tensor.ndim == 4:  # bshd
        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 enable_lse:
            lse_tensor = torch.empty(b, h_q, s_q, dtype=torch.float32, device=q_tensor.device).contiguous()
    elif q_tensor.ndim == 3:  # thd
        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 enable_lse:
            lse_tensor = torch.empty(1, h_q, t, dtype=torch.float32, device=q_tensor.device).contiguous().permute(2, 1, 0)
    else:
        raise ValueError(f"Invalid input layout: q_tensor must be rank-4 (B,H,S,D) or rank-3 (T,H,D), got {q_tensor.ndim}")

    cache_key = (
        q_tensor.shape,
        k_tensor.shape,
        v_tensor.shape,
        cum_seqlen_q_tensor.shape if cum_seqlen_q_tensor is not None else None,
        cum_seqlen_k_tensor.shape if cum_seqlen_k_tensor is not None else None,
        q_tensor.dtype,
        k_tensor.dtype,
        v_tensor.dtype,
        cum_seqlen_q_tensor.dtype if cum_seqlen_q_tensor is not None else None,
        cum_seqlen_k_tensor.dtype if cum_seqlen_k_tensor is not None else None,
        q_tensor.stride(),
        k_tensor.stride(),
        v_tensor.stride(),
        cum_seqlen_q_tensor.stride() if cum_seqlen_q_tensor is not None else None,
        cum_seqlen_k_tensor.stride() if cum_seqlen_k_tensor is not None else None,
        enable_lse,
        o_dtype,
        qk_acc_dtype,
        pv_acc_dtype,
        mma_tiler_mn,
        is_persistent,
        scale_q,
        scale_k,
        scale_v,
        inv_scale_o,
        scale_softmax,
    )
    if cache_key in _cache_of_CompressionAttentionObjects:
        _logger.debug("compression_attention_wrapper: Using previously cached CompressionAttention object")
        comp_attn = _cache_of_CompressionAttentionObjects[cache_key]
        comp_attn.execute(
            q_tensor=q_tensor,
            k_tensor=k_tensor,
            v_tensor=v_tensor,
            o_tensor=o_tensor,
            lse_tensor=lse_tensor,
            cum_seqlen_q_tensor=cum_seqlen_q_tensor,
            cum_seqlen_k_tensor=cum_seqlen_k_tensor,
            current_stream=stream,
        )
    else:
        _logger.debug("compression_attention_wrapper: No cached object found, creating new CompressionAttention object")
        comp_attn = CompressionAttention(
            sample_q=q_tensor,
            sample_k=k_tensor,
            sample_v=v_tensor,
            sample_o=o_tensor,
            sample_lse=lse_tensor,
            sample_cum_seqlen_q=cum_seqlen_q_tensor,
            sample_cum_seqlen_k=cum_seqlen_k_tensor,
            qk_acc_dtype=qk_acc_dtype,
            pv_acc_dtype=pv_acc_dtype,
            mma_tiler_mn=mma_tiler_mn,
            is_persistent=is_persistent,
            scale_q=scale_q,
            scale_k=scale_k,
            scale_v=scale_v,
            inv_scale_o=inv_scale_o,
            scale_softmax=scale_softmax,
        )
        assert comp_attn.check_support()
        comp_attn.compile()
        comp_attn.execute(
            q_tensor=q_tensor,
            k_tensor=k_tensor,
            v_tensor=v_tensor,
            o_tensor=o_tensor,
            lse_tensor=lse_tensor,
            cum_seqlen_q_tensor=cum_seqlen_q_tensor,
            cum_seqlen_k_tensor=cum_seqlen_k_tensor,
            current_stream=stream,
        )
        _cache_of_CompressionAttentionObjects[cache_key] = comp_attn

    return TupleDict(
        o_tensor=o_tensor,
        lse_tensor=lse_tensor,
    )