quack-kernels 0.2.2__py3-none-any.whl → 0.2.3__py3-none-any.whl

quack/utils.py

@@ -1,7 +1,8 @@
 # Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao.
 
 import math
-from typing import Optional, Tuple, Type, Union
+from functools import partial
+from typing import Optional, Tuple, Union
 
 import cutlass
 import cutlass.cute as cute
@@ -9,70 +10,18 @@ import cutlass.cute as cute
 from cutlass import Float32, Int32, const_expr
 from cutlass.cutlass_dsl import T, dsl_user_op
 from cutlass._mlir.dialects import llvm, nvvm, vector
-from cutlass.cute.runtime import from_dlpack
 
 
-def convert_from_dlpack(x, leading_dim, alignment=16, divisibility=1) -> cute.Tensor:
-    return (
-        from_dlpack(x, assumed_align=alignment)
-        .mark_layout_dynamic(leading_dim=leading_dim)
-        .mark_compact_shape_dynamic(
-            mode=leading_dim, stride_order=x.dim_order(), divisibility=divisibility
-        )
-    )
-
-
-def transpose_view(a: cute.Tensor) -> cute.Tensor:
-    """Transpose the first two dimensions of a tensor on smem."""
-    shape = (a.shape[1], a.shape[0], *a.shape[2:])
-    order = (1, 0, *range(2, cute.rank(a)))
-    return cute.composition(a, cute.make_ordered_layout(shape, order=order))
-
-
-def select(a: cute.Tensor, mode: list[int]) -> cute.Tensor:
-    return cute.make_tensor(a.iterator, cute.select(a.layout, mode))
-
-
-@dsl_user_op
-def get_copy_atom(
-    dtype: Type[cutlass.Numeric], num_copy_elems: int, is_async: bool = False, *, loc=None, ip=None
-) -> cute.CopyAtom:
-    num_copy_bits = const_expr(min(128, num_copy_elems * dtype.width))
-    copy_op = cute.nvgpu.cpasync.CopyG2SOp() if is_async else cute.nvgpu.CopyUniversalOp()
-    return cute.make_copy_atom(copy_op, dtype, num_bits_per_copy=num_copy_bits)
-
-
-@dsl_user_op
-def copy(
-    src: cute.Tensor,
-    dst: cute.Tensor,
-    *,
-    pred: Optional[cute.Tensor] = None,
-    num_copy_elems: int = 1,
-    is_async: bool = False,
-    loc=None,
-    ip=None,
-    **kwargs,
-) -> None:
-    copy_atom = get_copy_atom(src.element_type, num_copy_elems, is_async)
-    cute.copy(copy_atom, src, dst, pred=pred, loc=loc, ip=ip, **kwargs)
-
-
-def tiled_copy_2d(
-    dtype: Type[cutlass.Numeric], major_mode_size: int, num_threads: int, is_async: bool = True
-) -> cute.TiledCopy:
-    num_copy_bits = math.gcd(major_mode_size, 128 // dtype.width) * dtype.width
-    copy_elems = num_copy_bits // dtype.width
-    copy_op = cute.nvgpu.cpasync.CopyG2SOp() if is_async else cute.nvgpu.CopyUniversalOp()
-    copy_atom = cute.make_copy_atom(copy_op, dtype, num_bits_per_copy=num_copy_bits)
-    gmem_threads_per_row = major_mode_size // copy_elems
-    assert num_threads % gmem_threads_per_row == 0
-    thr_layout = cute.make_ordered_layout(
-        (num_threads // gmem_threads_per_row, gmem_threads_per_row),
-        order=(1, 0),
-    )
-    val_layout = cute.make_layout((1, copy_elems))
-    return cute.make_tiled_copy_tv(copy_atom, thr_layout, val_layout)
+# cute.arch.{fma,mul,add}_packed_f32x2 uses RZ rounding mode by default
+fma_packed_f32x2 = partial(cute.arch.fma_packed_f32x2, rnd=nvvm.RoundingModeKind.RN)
+mul_packed_f32x2 = partial(cute.arch.mul_packed_f32x2, rnd=nvvm.RoundingModeKind.RN)
+add_packed_f32x2 = partial(cute.arch.add_packed_f32x2, rnd=nvvm.RoundingModeKind.RN)
+sub_packed_f32x2 = partial(
+    cute.arch.calc_packed_f32x2_op,
+    src_c=None,
+    calc_func=nvvm.sub_packed_f32x2,
+    rnd=nvvm.RoundingModeKind.RN,
+)
 
 
 @dsl_user_op
@@ -91,11 +40,11 @@ def load_scalar_or_pointer(x: Float32 | cute.Pointer) -> Float32:
 
 @dsl_user_op
 def set_block_rank(
-    smem_ptr: cute.Pointer, peer_cta_rank_in_cluster: cute.Int32, *, loc=None, ip=None
-) -> cutlass.Int32:
+    smem_ptr: cute.Pointer, peer_cta_rank_in_cluster: Int32, *, loc=None, ip=None
+) -> Int32:
     """Map the given smem pointer to the address at another CTA rank in the cluster."""
     smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value()
-    return cutlass.Int32(
+    return Int32(
         llvm.inline_asm(
             T.i32(),
             [smem_ptr_i32, peer_cta_rank_in_cluster.ir_value()],
@@ -110,7 +59,7 @@ def set_block_rank(
 
 @dsl_user_op
 def store_shared_remote(
-    val: float | Float32 | cutlass.Int64,
+    val: float | Float32 | Int32 | cutlass.Int64,
     smem_ptr: cute.Pointer,
     mbar_ptr: cute.Pointer,
     peer_cta_rank_in_cluster: cute.typing.Int,
@@ -153,6 +102,21 @@ def fmin(a: Union[float, Float32], b: Union[float, Float32], *, loc=None, ip=Non
     )
 
 
+@dsl_user_op
+def sqrt(a: float | Float32, *, loc=None, ip=None) -> Float32:
+    return Float32(
+        llvm.inline_asm(
+            T.f32(),
+            [Float32(a).ir_value(loc=loc, ip=ip)],
+            "sqrt.approx.f32 $0, $1;",
+            "=f,f",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    )
+
+
 @dsl_user_op
 def ceil(a: float | Float32, *, loc=None, ip=None) -> Int32:
     return Int32(
@@ -187,55 +151,6 @@ def prmt(a: int | Int32, b: int | Int32, c: int | Int32, *, loc=None, ip=None) -
     )
 
 
-@cute.jit
-def permute_gated_Cregs_b16(t: cute.Tensor) -> None:
-    assert t.element_type.width == 16
-    assert cute.size(t.shape) % 4 == 0, "Tensor size must be a multiple of 4 for b16 permutation"
-    t_u32 = cute.recast_tensor(t, Int32)
-
-    quad_idx = cute.arch.lane_idx() % 4
-    lane_03 = quad_idx == 0 or quad_idx == 3
-    selector_upper = Int32(0x5410) if lane_03 else Int32(0x1054)
-    selector_lower = Int32(0x7632) if lane_03 else Int32(0x3276)
-    # upper_map = [0, 3, 1, 2]
-    # lower_map = [1, 2, 0, 3]
-    # upper_idx = upper_map[quad_idx]
-    # indexing isn't supported so we have to do arithmetic
-    upper_idx = quad_idx // 2 if quad_idx % 2 == 0 else 3 - quad_idx // 2
-    lower_idx = upper_idx ^ 1
-
-    # 1 -> 0b11111, 2 -> 0b11110, 4 -> 0b11100, 8 -> 0b11000, 16 -> 0b10000, 32 -> 0b00000
-    width = 4
-    mask = cute.arch.WARP_SIZE - width
-    clamp = cute.arch.WARP_SIZE - 1
-    mask_and_clamp = mask << 8 | clamp
-
-    for i in cutlass.range(cute.size(t_u32.shape) // 2, unroll_full=True):
-        upper, lower = t_u32[i * 2 + 0], t_u32[i * 2 + 1]
-        upper0 = upper if lane_03 else lower
-        lower0 = lower if lane_03 else upper
-        upper0 = cute.arch.shuffle_sync(upper0, offset=upper_idx, mask_and_clamp=mask_and_clamp)
-        lower0 = cute.arch.shuffle_sync(lower0, offset=lower_idx, mask_and_clamp=mask_and_clamp)
-        t_u32[i * 2 + 0] = prmt(upper0, lower0, selector_upper)
-        t_u32[i * 2 + 1] = prmt(upper0, lower0, selector_lower)
-
-
-@cute.jit
-def predicate_k(tAcA: cute.Tensor, limit: cutlass.Int32) -> cute.Tensor:
-    # Only compute predicates for the "k" dimension. For the mn dimension, we will use "if"
-    tApA = cute.make_fragment(
-        cute.make_layout(
-            (cute.size(tAcA, mode=[0, 1]), cute.size(tAcA, mode=[1]), cute.size(tAcA, mode=[2])),
-            stride=(cute.size(tAcA, mode=[2]), 0, 1),
-        ),
-        cutlass.Boolean,
-    )
-    for rest_v in cutlass.range_constexpr(tApA.shape[0]):
-        for rest_k in cutlass.range_constexpr(tApA.shape[2]):
-            tApA[rest_v, 0, rest_k] = cute.elem_less(tAcA[(0, rest_v), 0, rest_k][1], limit)
-    return tApA
-
-
 @cute.jit
 def fill_oob(tXsX: cute.Tensor, tXpX: Optional[cute.Tensor], fill_value: cute.Numeric) -> None:
     """Fill out-of-bounds values in shared memory tensor.
@@ -281,43 +196,8 @@ def i64_to_f32x2(c: cutlass.Int64, *, loc=None, ip=None) -> Tuple[Float32, Float
     return res0, res1
 
 
-@dsl_user_op
-def domain_offset_i64(coord: cute.Coord, tensor: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor:
-    flat_coord_i64 = tuple(cutlass.Int64(c) for c in cute.flatten(coord))
-    flat_stride = cute.flatten_to_tuple(tensor.stride)
-    assert len(flat_coord_i64) == len(flat_stride), (
-        "Coordinate and stride must have the same length"
-    )
-    offset = sum(c * s for c, s in zip(flat_coord_i64, flat_stride))
-    assert isinstance(tensor.iterator, cute.Pointer)
-    # HACK: we assume that applying the offset does not change the pointer alignment
-    new_ptr = cute.make_ptr(
-        tensor.element_type,
-        tensor.iterator.toint() + offset * tensor.element_type.width // 8,
-        tensor.memspace,
-        assumed_align=tensor.iterator.max_alignment,
-    )
-    return cute.make_tensor(new_ptr, tensor.layout)
-
-
-@dsl_user_op
-def coord_offset_i64(
-    idx: cute.typing.Int, tensor: cute.Tensor, dim: int, *, loc=None, ip=None
-) -> cute.Tensor:
-    offset = cutlass.Int64(idx) * cute.size(tensor.stride[dim])
-    assert isinstance(tensor.iterator, cute.Pointer)
-    # HACK: we assume that applying the offset does not change the pointer alignment
-    new_ptr = cute.make_ptr(
-        tensor.element_type,
-        tensor.iterator.toint() + offset * tensor.element_type.width // 8,
-        tensor.memspace,
-        assumed_align=tensor.iterator.max_alignment,
-    )
-    return cute.make_tensor(new_ptr, tensor.layout)
-
-
 @cute.jit
-def warp_prefix_sum(val: cutlass.Int32, lane: Optional[cutlass.Int32] = None) -> cutlass.Int32:
+def warp_prefix_sum(val: Int32, lane: Optional[Int32] = None) -> Int32:
     if const_expr(lane is None):
         lane = cute.arch.lane_idx()
     for i in cutlass.range_constexpr(int(math.log2(cute.arch.WARP_SIZE))):
@@ -329,74 +209,6 @@ def warp_prefix_sum(val: cutlass.Int32, lane: Optional[cutlass.Int32] = None) ->
     return val
 
 
-def convert_layout_acc_mn(acc_layout: cute.Layout) -> cute.Layout:
-    """
-    For Sm80, convert ((2, 2), MMA_M, MMA_N, ...) to ((2, MMA_M), (2, MMA_N), ...).
-    For Sm90, convert ((2, 2, V), MMA_M, MMA_N, ...) to ((2, MMA_M), (2, V, MMA_N), ...).
-    """
-    acc_layout_col_major = cute.make_layout(acc_layout.shape)
-    acc_layout_mn = cute.make_layout(
-        (
-            (acc_layout_col_major.shape[0][1], acc_layout_col_major.shape[1]),  # MMA_M
-            (
-                acc_layout_col_major.shape[0][0],
-                *acc_layout_col_major.shape[0][2:],
-                acc_layout_col_major.shape[2],
-            ),  # MMA_N
-            *acc_layout_col_major.shape[3:],
-        ),
-        stride=(
-            (acc_layout_col_major.stride[0][1], acc_layout_col_major.stride[1]),  # MMA_M
-            (
-                acc_layout_col_major.stride[0][0],
-                *acc_layout_col_major.stride[0][2:],
-                acc_layout_col_major.stride[2],
-            ),  # MMA_N
-            *acc_layout_col_major.stride[3:],
-        ),
-    )
-    return cute.composition(acc_layout, acc_layout_mn)
-
-
-def make_acc_tensor_mn_view(acc: cute.Tensor) -> cute.Tensor:
-    return cute.make_tensor(acc.iterator, convert_layout_acc_mn(acc.layout))
-
-
-@dsl_user_op
-def sm90_get_smem_load_op(
-    layout_c: cutlass.utils.LayoutEnum,
-    elem_ty_c: Type[cutlass.Numeric],
-    *,
-    loc=None,
-    ip=None,
-) -> cute.CopyAtom:
-    """
-    Selects the largest vectorized smem load atom available subject to constraint of gmem layout.
-
-    Parameters:
-    -----------
-    layout_c : LayoutEnum
-        The layout enum of the output tensor D.
-
-    elem_ty_c : Type[Numeric]
-        The element type for output tensor D.
-
-    Returns:
-    --------
-    Either SmemLoadMatrix or SimtSyncCopy, based on the input parameters.
-    """
-
-    if not isinstance(elem_ty_c, cutlass.cutlass_dsl.NumericMeta):
-        raise TypeError(f"elem_ty_c must be a Numeric, but got {elem_ty_c}")
-    is_m_major = layout_c.is_m_major_c()
-    if elem_ty_c.width == 16:
-        return cute.make_copy_atom(
-            cute.nvgpu.warp.LdMatrix8x8x16bOp(is_m_major, 4), elem_ty_c, loc=loc, ip=ip
-        )
-    else:
-        return cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), elem_ty_c, loc=loc, ip=ip)
-
-
 @dsl_user_op
 def atomic_add_i32(a: int | Int32, gmem_ptr: cute.Pointer, *, loc=None, ip=None) -> Int32:
     return nvvm.atomicrmw(

quack/varlen_utils.py

@@ -3,9 +3,13 @@
 from typing import Optional
 from dataclasses import dataclass
 
+import cutlass
 import cutlass.cute as cute
+from cutlass import Int32, Boolean, const_expr
+from cutlass.utils import LayoutEnum
 
-from quack.cute_dsl_utils import ArgumentsBase
+from quack.cute_dsl_utils import ArgumentsBase, ParamsBase
+from quack.tensormap_manager import TensorMapManagerSm90
 
 
 # Grouping arguments together that should be passed to __call__
@@ -15,3 +19,368 @@ class VarlenArguments(ArgumentsBase):
     mCuSeqlensK: Optional[cute.Tensor] = None
     mTensormaps: Optional[cute.Tensor] = None
     mAIdx: Optional[cute.Tensor] = None
+
+
+class VarlenManager:
+    bytes_per_tensormap = 128
+
+    @dataclass
+    class Params(ParamsBase):
+        cu_seqlens_m: Optional[cute.Tensor] = None
+        cu_seqlens_k: Optional[cute.Tensor] = None
+        tensormaps: Optional[cute.Tensor] = None
+        mAIdx: Optional[cute.Tensor] = None
+
+        @staticmethod
+        @cute.jit
+        def create(args: VarlenArguments, *, loc=None, ip=None) -> "VarlenManager.Params":
+            return VarlenManager.Params(
+                cu_seqlens_m=args.mCuSeqlensM,
+                cu_seqlens_k=args.mCuSeqlensK,
+                tensormaps=args.mTensormaps,
+                mAIdx=args.mAIdx,
+            )
+
+    def __init__(
+        self,
+        params: Params,
+        tensormap_manager: Optional[cutlass.utils.TensorMapManager],
+        tensormap_a_ptr: Optional[cute.Pointer],
+        tensormap_b_ptr: Optional[cute.Pointer],
+        tensormap_d_ptr: Optional[cute.Pointer],
+        tensormap_epi_ptrs: list[Optional[cute.Pointer]],
+        len_m_static: Int32,
+        len_k_static: Int32,
+        last_batch_idx: Int32 = Int32(-1),
+        is_group_changed: Boolean = Boolean(True),
+        *,
+        loc=None,
+        ip=None,
+    ):
+        self.params = params
+        self.tensormap_manager = tensormap_manager
+        self._tensormap_a_ptr = tensormap_a_ptr
+        self._tensormap_b_ptr = tensormap_b_ptr
+        self._tensormap_d_ptr = tensormap_d_ptr
+        self._tensormap_epi_ptrs = tensormap_epi_ptrs
+        self._len_m_static = len_m_static
+        self._len_k_static = len_k_static
+        self._last_batch_idx = last_batch_idx
+        self._is_group_changed = is_group_changed
+        self.varlen_m = const_expr(params.cu_seqlens_m is not None)
+        self.varlen_k = const_expr(params.cu_seqlens_k is not None)
+        self.gather_A = const_expr(params.mAIdx is not None)
+        self._loc = loc
+        self._ip = ip
+
+    @staticmethod
+    def to_underlying_arguments(args: VarlenArguments, *, loc=None, ip=None) -> Params:
+        assert not (args.mCuSeqlensM is not None and args.mCuSeqlensK is not None), (
+            "Only support either varlen_m or varlen_k"
+        )
+        return VarlenManager.Params.create(args, loc=loc, ip=ip)
+
+    @staticmethod
+    @cute.jit
+    def create(
+        params: Params,
+        has_D: bool,
+        num_epi_tensormaps: int,
+        len_m_static: Int32,
+        len_k_static: Int32,
+        pingpong: bool = False,
+        warp_idx: int | Int32 = 0,
+        *,
+        loc=None,
+        ip=None,
+    ) -> "VarlenManager":
+        tensormap_manager = None
+        tensormap_a_ptr, tensormap_b_ptr, tensormap_d_ptr = None, None, None
+        tensormap_epi_ptrs = [None] * num_epi_tensormaps
+        varlen_m = const_expr(params.cu_seqlens_m is not None)
+        varlen_k = const_expr(params.cu_seqlens_k is not None)
+        if const_expr(varlen_m or varlen_k):
+            tensormap_manager = TensorMapManagerSm90(
+                cutlass.utils.TensorMapUpdateMode.GMEM, VarlenManager.bytes_per_tensormap
+            )
+            # equivalent to bidx + bidy * gridDim.x + bidxz * gridDim.x * gridDim.y
+            tensormap_workspace_idx = cute.make_layout(cute.arch.grid_dim())(cute.arch.block_idx())
+            if const_expr(varlen_m):
+                tensormap_d_idx = warp_idx // 4 if const_expr(pingpong) else 0
+                tensormap_epi_offset = tensormap_d_idx
+                if const_expr(has_D):
+                    tensormap_d_ptr = tensormap_manager.get_tensormap_ptr(
+                        params.tensormaps[tensormap_workspace_idx, tensormap_d_idx, None].iterator
+                    )
+                    tensormap_epi_offset += 1 if not pingpong else 2
+                tensormap_epi_ptrs = [
+                    tensormap_manager.get_tensormap_ptr(
+                        params.tensormaps[
+                            tensormap_workspace_idx,
+                            tensormap_epi_offset + i * (1 if not pingpong else 2),
+                            None,
+                        ].iterator
+                    )
+                    for i in range(num_epi_tensormaps)
+                ]
+            else:
+                assert varlen_k
+                gather_A = const_expr(params.mAIdx is not None)
+                if const_expr(not gather_A):
+                    tensormap_a_ptr = tensormap_manager.get_tensormap_ptr(
+                        params.tensormaps[tensormap_workspace_idx, 0, None].iterator
+                    )
+                tensormap_b_ptr = tensormap_manager.get_tensormap_ptr(
+                    params.tensormaps[
+                        tensormap_workspace_idx, 1 if not gather_A else 0, None
+                    ].iterator
+                )
+        return VarlenManager(
+            params,
+            tensormap_manager,
+            tensormap_a_ptr,
+            tensormap_b_ptr,
+            tensormap_d_ptr,
+            tensormap_epi_ptrs,
+            len_m_static=len_m_static,
+            len_k_static=len_k_static,
+        )
+
+    def len_m(self, batch_idx: Int32) -> Int32:
+        if const_expr(self.varlen_m):
+            return self.params.cu_seqlens_m[batch_idx + 1] - self.params.cu_seqlens_m[batch_idx]
+        else:
+            return self._len_m_static
+
+    def len_k(self, batch_idx: Int32) -> Int32:
+        if const_expr(self.varlen_k):
+            return self.params.cu_seqlens_k[batch_idx + 1] - self.params.cu_seqlens_k[batch_idx]
+        else:
+            return self._len_k_static
+
+    def offset_batch_A(self, mA_mkl: cute.Tensor, batch_idx: Int32) -> cute.Tensor:
+        params = self.params
+        if const_expr(self.varlen_m):
+            mA_mk = cute.domain_offset((params.cu_seqlens_m[batch_idx], 0), mA_mkl)
+        elif const_expr(self.varlen_k):
+            mA_mk = cute.domain_offset((0, params.cu_seqlens_k[batch_idx]), mA_mkl)
+        else:
+            mA_mk = mA_mkl[None, None, batch_idx]
+        return mA_mk
+
+    def offset_batch_AIdx(self, batch_idx: Int32) -> cute.Tensor:
+        params = self.params
+        if const_expr(self.varlen_m):
+            mAIdx_mk = cute.domain_offset((params.cu_seqlens_m[batch_idx],), params.mAIdx)
+        elif const_expr(self.varlen_k):
+            mAIdx_mk = cute.domain_offset((params.cu_seqlens_k[batch_idx],), params.mAIdx)
+        else:
+            mAIdx_mk = params.mAIdx[None, batch_idx]
+        return mAIdx_mk
+
+    def offset_batch_B(self, mB_nkl: cute.Tensor, batch_idx: Int32) -> cute.Tensor:
+        params = self.params
+        if const_expr(self.varlen_k):
+            mB_nk = cute.domain_offset((0, params.cu_seqlens_k[batch_idx]), mB_nkl)
+        else:
+            mB_nk = mB_nkl[None, None, batch_idx]
+        return mB_nk
+
+    def offset_batch_epi(self, mD_mnl: cute.Tensor, batch_idx: Int32) -> cute.Tensor:
+        params = self.params
+        if const_expr(self.varlen_m):
+            mD_mn = cute.domain_offset((params.cu_seqlens_m[batch_idx], 0), mD_mnl)
+        else:
+            mD_mn = mD_mnl[None, None, batch_idx]
+        return mD_mn
+
+    def init_tensormap_AB(
+        self,
+        tma_atom_a: Optional[cute.CopyAtom],
+        tma_atom_b: cute.CopyAtom,
+        is_manager_warp: bool | Boolean = True,
+    ) -> None:
+        if const_expr(self.varlen_k):
+            if const_expr(not self.gather_A):
+                self.tensormap_manager.init_tensormap_from_atom(
+                    tma_atom_a, self._tensormap_a_ptr, is_manager_warp
+                )
+            self.tensormap_manager.init_tensormap_from_atom(
+                tma_atom_b, self._tensormap_b_ptr, is_manager_warp
+            )
+
+    def init_tensormap_epi(
+        self,
+        tma_atom_d: Optional[cute.CopyAtom],
+        tma_atoms_epi: list[cute.CopyAtom],
+        is_manager_warp: bool | Boolean = True,
+    ) -> None:
+        if const_expr(self.varlen_m):
+            if const_expr(self._tensormap_d_ptr is not None):
+                self.tensormap_manager.init_tensormap_from_atom(
+                    tma_atom_d, self._tensormap_d_ptr, is_manager_warp
+                )
+            for tma_atom, tensormap_epi_ptr in zip(tma_atoms_epi, self._tensormap_epi_ptrs):
+                self.tensormap_manager.init_tensormap_from_atom(
+                    tma_atom, tensormap_epi_ptr, is_manager_warp
+                )
+
+    def fence_tensormap_init(self) -> None:
+        self.tensormap_manager.fence_tensormap_initialization()
+
+    @cute.jit
+    def update_tensormap_AB(
+        self,
+        batch_idx: Int32,
+        a_layout: LayoutEnum,
+        b_layout: LayoutEnum,
+        is_manager_warp: bool | Boolean = True,
+    ) -> None:
+        if const_expr(self.varlen_k):
+            self._is_group_changed = Boolean(batch_idx != self._last_batch_idx)
+            self._last_batch_idx = batch_idx
+            if self._is_group_changed:
+                # construct tensor A/B based on real address, shape and stride information
+                cu_seqlens_k = self.params.cu_seqlens_k
+                tensormap_ptrs = [self._tensormap_b_ptr]
+                shapes = [cu_seqlens_k[batch_idx + 1]]
+                orders = [0 if const_expr(b_layout == LayoutEnum.ROW_MAJOR) else 1]
+                if const_expr(not self.gather_A):
+                    tensormap_ptrs.insert(0, self._tensormap_a_ptr)
+                    shapes.insert(0, cu_seqlens_k[batch_idx + 1])
+                    orders.insert(0, 0 if const_expr(a_layout == LayoutEnum.ROW_MAJOR) else 1)
+                self.tensormap_manager.update_tensormap_shape(
+                    tensormap_ptrs,
+                    is_manager_warp=is_manager_warp,
+                    shapes=shapes,
+                    orders=orders,
+                    tensormap_smem_ptr=None,
+                )
+
+    @cute.jit
+    def update_tensormap_epi(
+        self,
+        batch_idx: Int32,
+        d_layout: LayoutEnum,
+        epi_shapes: list[Int32],
+        epi_orders: list[int],
+        is_manager_warp: bool | Boolean = True,
+    ) -> None:
+        if const_expr(self.varlen_m):
+            self._is_group_changed = Boolean(batch_idx != self._last_batch_idx)
+            self._last_batch_idx = batch_idx
+            # Cute-DSL doesn't like this under if statement
+            order_d = (
+                (0 if const_expr(d_layout.is_m_major_c()) else 1) if d_layout is not None else None
+            )
+            if self._is_group_changed:
+                # construct tensor A/B based on real address, shape and stride information
+                cu_seqlens_m = self.params.cu_seqlens_m
+                # construct tensor D based on real address, shape and stride information
+                tensormap_ptrs, shapes, orders = [], [], []
+                if const_expr(self._tensormap_d_ptr is not None):
+                    tensormap_ptrs.append(self._tensormap_d_ptr)
+                    shapes.append(cu_seqlens_m[batch_idx + 1])
+                    orders.append(order_d)
+                tensormap_ptrs.extend(self._tensormap_epi_ptrs)
+                shapes.extend(epi_shapes)
+                orders.extend(epi_orders)
+                self.tensormap_manager.update_tensormap_shape(
+                    tensormap_ptrs,
+                    is_manager_warp=is_manager_warp,
+                    shapes=shapes,
+                    orders=orders,
+                    tensormap_smem_ptr=None,
+                )
+
+    @cute.jit
+    def fence_tensormap_update_AB(self, is_manager_warp: bool | Boolean = True) -> None:
+        if const_expr(self.varlen_k):
+            if self._is_group_changed and is_manager_warp:
+                if const_expr(not self.gather_A):
+                    self.tensormap_manager.fence_tensormap_update(self._tensormap_a_ptr)
+                self.tensormap_manager.fence_tensormap_update(self._tensormap_b_ptr)
+
+    @cute.jit
+    def fence_tensormap_update_epi(self, is_manager_warp: bool | Boolean = True) -> None:
+        if const_expr(self.varlen_m):
+            if self._is_group_changed and is_manager_warp:
+                if const_expr(self._tensormap_d_ptr is not None):
+                    self.tensormap_manager.fence_tensormap_update(self._tensormap_d_ptr)
+                for tensormap_epi_ptr in self._tensormap_epi_ptrs:
+                    if const_expr(tensormap_epi_ptr is not None):
+                        self.tensormap_manager.fence_tensormap_update(tensormap_epi_ptr)
+
+    def get_tma_desc_a_ptr(self) -> Optional[cute.Pointer]:
+        tma_desc_a_ptr = None
+        if const_expr(self.varlen_k and self._tensormap_a_ptr is not None):
+            tma_desc_a_ptr = self.tensormap_manager.get_tensormap_ptr(
+                self._tensormap_a_ptr, cute.AddressSpace.generic
+            )
+        return tma_desc_a_ptr
+
+    def get_tma_desc_b_ptr(self) -> Optional[cute.Pointer]:
+        tma_desc_b_ptr = None
+        if const_expr(self.varlen_k):
+            tma_desc_b_ptr = self.tensormap_manager.get_tensormap_ptr(
+                self._tensormap_b_ptr, cute.AddressSpace.generic
+            )
+        return tma_desc_b_ptr
+
+    def get_tma_desc_d_ptr(self) -> Optional[cute.Pointer]:
+        tma_desc_d_ptr = None
+        if const_expr(self.varlen_m and self._tensormap_d_ptr is not None):
+            tma_desc_d_ptr = self.tensormap_manager.get_tensormap_ptr(
+                self._tensormap_d_ptr, cute.AddressSpace.generic
+            )
+        return tma_desc_d_ptr
+
+    def get_tma_desc_epi_ptrs(self) -> list[Optional[cute.Pointer]]:
+        tma_desc_epi_ptrs = [None] * len(self._tensormap_epi_ptrs)
+        if const_expr(self.varlen_m):
+            for i, tensormap_epi_ptr in enumerate(self._tensormap_epi_ptrs):
+                if const_expr(tensormap_epi_ptr is not None):
+                    tma_desc_epi_ptrs[i] = self.tensormap_manager.get_tensormap_ptr(
+                        tensormap_epi_ptr, cute.AddressSpace.generic
+                    )
+        return tma_desc_epi_ptrs
+
+    def __extract_mlir_values__(self):
+        values, self._values_pos = [], []
+        for obj in [
+            self.params,
+            self.tensormap_manager,
+            self._tensormap_a_ptr,
+            self._tensormap_b_ptr,
+            self._tensormap_d_ptr,
+            self._tensormap_epi_ptrs,
+            self._len_m_static,
+            self._len_k_static,
+            self._last_batch_idx,
+            self._is_group_changed,
+        ]:
+            obj_values = cutlass.extract_mlir_values(obj)
+            values += obj_values
+            self._values_pos.append(len(obj_values))
+        return values
+
+    def __new_from_mlir_values__(self, values):
+        obj_list = []
+        for obj, n_items in zip(
+            [
+                self.params,
+                self.tensormap_manager,
+                self._tensormap_a_ptr,
+                self._tensormap_b_ptr,
+                self._tensormap_d_ptr,
+                self._tensormap_epi_ptrs,
+                self._len_m_static,
+                self._len_k_static,
+                self._last_batch_idx,
+                self._is_group_changed,
+            ],
+            self._values_pos,
+        ):
+            obj_list.append(cutlass.new_from_mlir_values(obj, values[:n_items]))
+            values = values[n_items:]
+        return self.__class__(*(tuple(obj_list)), loc=self._loc)