# Copyright (c) 2025, Tri Dao. import math import operator import cutlass import cutlass.cute as cute import flash_attn.cute.utils as utils class PackGQA: def __init__( self, m_block_size: cutlass.Constexpr[int], head_dim_padded: cutlass.Constexpr[int], check_hdim_oob: cutlass.Constexpr[bool], qhead_per_kvhead: cutlass.Constexpr[bool], ): self.m_block_size = m_block_size self.head_dim_padded = head_dim_padded self.check_hdim_oob = check_hdim_oob self.qhead_per_kvhead = qhead_per_kvhead @cute.jit def compute_ptr( self, tensor: cute.Tensor, cRows: cute.Tensor, tidx: cutlass.Int32, block: cutlass.Int32, threads_per_row: cutlass.Constexpr[int], num_threads: cutlass.Constexpr[int], ): num_ptr_per_thread = cute.ceil_div(cute.size(cRows), threads_per_row) tPrPtr = cute.make_fragment(num_ptr_per_thread, cutlass.Int64) for i in cutlass.range_constexpr(num_ptr_per_thread): row = i * num_threads + cRows[tidx % threads_per_row][0] idx = block * self.m_block_size + row m_idx = idx // self.qhead_per_kvhead h_idx = idx - m_idx * self.qhead_per_kvhead tPrPtr[i] = utils.elem_pointer(tensor, ((h_idx, m_idx),)).toint() return tPrPtr @cute.jit def load_Q( self, mQ: cute.Tensor, # ((qhead_per_kvhead, seqlen_q), headdim) sQ: cute.Tensor, # (m_block_size, head_dim_padded) gmem_tiled_copy: cute.TiledCopy, tidx: cutlass.Int32, block: cutlass.Int32, seqlen: cutlass.Int32, ): gmem_thr_copy = gmem_tiled_copy.get_slice(tidx) cQ = cute.make_identity_tensor((self.m_block_size, self.head_dim_padded)) tQsQ = gmem_thr_copy.partition_D(sQ) tQcQ = gmem_thr_copy.partition_S(cQ) t0QcQ = gmem_thr_copy.get_slice(0).partition_S(cQ) tQpQ = utils.predicate_k(tQcQ, limit=mQ.shape[1]) tQcQ_row = tQcQ[0, None, 0] threads_per_row = gmem_tiled_copy.layout_tv_tiled.shape[0][0] assert cute.arch.WARP_SIZE % threads_per_row == 0, "threads_per_row must divide WARP_SIZE" num_threads = gmem_tiled_copy.size tPrQPtr = self.compute_ptr(mQ[None, 0], tQcQ_row, tidx, block, threads_per_row, num_threads) for m in cutlass.range_constexpr(cute.size(tQsQ.shape[1])): q_ptr_i64 = utils.shuffle_sync( tPrQPtr[m // threads_per_row], m % threads_per_row, width=threads_per_row ) q_gmem_ptr = cute.make_ptr( mQ.element_type, q_ptr_i64, cute.AddressSpace.gmem, assumed_align=16 ) if ( t0QcQ[0, m, 0][0] < seqlen * self.qhead_per_kvhead - block * self.m_block_size - tQcQ_row[0][0] ): mQ_cur = cute.make_tensor(q_gmem_ptr, (self.head_dim_padded,)) elems_per_load = cute.size(tQsQ.shape[0][0]) mQ_cur_copy = cute.tiled_divide(mQ_cur, (elems_per_load,)) for k in cutlass.range_constexpr(cute.size(tQsQ.shape[2])): ki = tQcQ[0, 0, k][1] // elems_per_load cute.copy( gmem_thr_copy, mQ_cur_copy[None, ki], tQsQ[None, m, k], pred=tQpQ[None, m, k] if cutlass.const_expr(self.check_hdim_oob) else None, ) # We don't need to clear the sQ smem tiles since we'll only write out the valid outputs @cute.jit def store_LSE( self, mLSE: cute.Tensor, # (qhead_per_kvhead, seqlen_q) tLSErLSE: cute.Tensor, # (m_block_size, head_dim_padded) tiled_mma: cute.TiledMma, tidx: cutlass.Int32, block: cutlass.Int32, seqlen: cutlass.Int32, ): thr_mma = tiled_mma.get_slice(tidx) caccO = cute.make_identity_tensor((self.m_block_size, self.head_dim_padded)) taccOcO = thr_mma.partition_C(caccO) taccOcO_row = utils.make_acc_tensor_mn_view(taccOcO)[None, 0] assert cute.size(tLSErLSE) == cute.size(taccOcO_row) threads_per_row = tiled_mma.tv_layout_C.shape[0][0] assert cute.arch.WARP_SIZE % threads_per_row == 0, "threads_per_row must divide WARP_SIZE" assert cute.size(tLSErLSE) <= threads_per_row num_threads = tiled_mma.size tPrLSEPtr = self.compute_ptr(mLSE, taccOcO_row, tidx, block, threads_per_row, num_threads) for m in cutlass.range_constexpr(cute.size(tLSErLSE)): lse_ptr_i64 = utils.shuffle_sync( tPrLSEPtr[m // threads_per_row], m % threads_per_row, width=threads_per_row, ) lse_gmem_ptr = cute.make_ptr( mLSE.element_type, lse_ptr_i64, cute.AddressSpace.gmem, assumed_align=4 ) row = block * self.m_block_size + taccOcO_row[m][0] # Only the thread corresponding to column 0 writes out the lse to gmem if taccOcO[0][1] == 0 and row < seqlen * self.qhead_per_kvhead: mLSE_copy = cute.make_tensor(lse_gmem_ptr, (1,)) mLSE_copy[0] = tLSErLSE[m] @cute.jit def store_O( self, mO: cute.Tensor, # ((qhead_per_kvhead, seqlen_q), headdim) tOrO: cute.Tensor, # (m_block_size, head_dim_padded) split across threads according to gmem_tiled_copy gmem_tiled_copy: cute.TiledCopy, tidx: cutlass.Int32, block: cutlass.Int32, seqlen: cutlass.Int32, ): gmem_thr_copy = gmem_tiled_copy.get_slice(tidx) cO = cute.make_identity_tensor((self.m_block_size, self.head_dim_padded)) tOcO = gmem_thr_copy.partition_S(cO) t0OcO = gmem_thr_copy.get_slice(0).partition_S(cO) tOpO = utils.predicate_k(tOcO, limit=mO.shape[1]) tOcO_row = tOcO[0, None, 0] threads_per_row = gmem_tiled_copy.layout_tv_tiled.shape[0][0] assert cute.arch.WARP_SIZE % threads_per_row == 0, "threads_per_row must divide WARP_SIZE" num_threads = gmem_tiled_copy.size tPrOPtr = self.compute_ptr(mO[None, 0], tOcO_row, tidx, block, threads_per_row, num_threads) for m in cutlass.range_constexpr(cute.size(tOrO.shape[1])): o_ptr_i64 = utils.shuffle_sync( tPrOPtr[m // threads_per_row], m % threads_per_row, width=threads_per_row ) o_gmem_ptr = cute.make_ptr( mO.element_type, o_ptr_i64, cute.AddressSpace.gmem, assumed_align=16 ) if ( t0OcO[0, m, 0][0] < seqlen * self.qhead_per_kvhead - block * self.m_block_size - tOcO_row[0][0] ): mO_cur = cute.make_tensor(o_gmem_ptr, (self.head_dim_padded,)) elems_per_load = cute.size(tOrO.shape[0][0]) mO_cur_copy = cute.tiled_divide(mO_cur, (elems_per_load,)) for k in cutlass.range_constexpr(cute.size(tOrO.shape[2])): ki = tOcO[0, 0, k][1] // elems_per_load cute.copy( gmem_thr_copy, tOrO[None, m, k], mO_cur_copy[None, ki], pred=tOpO[None, m, k] if cutlass.const_expr(self.check_hdim_oob) else None, )