# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. # A reimplementation of https://github.com/Dao-AILab/flash-attention/blob/main/hopper/flash_fwd_combine_kernel.h # from Cutlass C++ to Cute-DSL. # Vendored from cuDNN BSA v126; only helper imports are redirected to the # equivalent private FlashAttention/CuTe source closure in this repository. import math from typing import Type, Optional from functools import partial import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute from cutlass.cute.nvgpu import cpasync from cutlass import Float32, Int32, Boolean, const_expr from .._vendor.flash_attn.cute import utils from .._vendor.flash_attn.cute.cute_dsl_utils import assume_tensor_aligned from .._vendor.flash_attn.cute.seqlen_info import SeqlenInfo from cutlass.cute import FastDivmodDivisor class BlockSparseAttnForwardCombine: def __init__( self, dtype: Type[cutlass.Numeric], head_dim: int, tile_m: int = 8, k_block_size: int = 64, log_max_splits: int = 4, num_threads: int = 256, stages: int = 4, partial_dtype: Type[cutlass.Numeric] = Float32, ): """ Forward combine kernel for split attention computation. :param dtype: output data type :param head_dim: head dimension :param tile_m: m block size :param k_block_size: k block size :param log_max_splits: log2 of maximum splits :param num_threads: number of threads :param varlen: whether using variable length sequences :param stages: number of pipeline stages """ self.dtype = dtype if partial_dtype not in (Float32, cutlass.BFloat16): raise TypeError("O partial dtype must be Float32 or BFloat16") self.partial_dtype = partial_dtype self.head_dim = head_dim self.tile_m = tile_m self.k_block_size = k_block_size self.max_splits = 1 << log_max_splits self.num_threads = num_threads self.is_even_k = head_dim % k_block_size == 0 self.stages = stages def _setup_attributes(self): # GMEM copy setup for O partial universal_copy_bits = 128 async_copy_elems = universal_copy_bits // self.partial_dtype.width assert self.k_block_size % async_copy_elems == 0 k_block_gmem = 128 if self.k_block_size % 128 == 0 else (64 if self.k_block_size % 64 == 0 else 32) gmem_threads_per_row = k_block_gmem // async_copy_elems assert self.num_threads % gmem_threads_per_row == 0 # Async copy atom for O partial load atom_async_copy_partial = cute.make_copy_atom( cpasync.CopyG2SOp(cache_mode=cute.nvgpu.LoadCacheMode.GLOBAL), self.partial_dtype, num_bits_per_copy=universal_copy_bits, ) tOpartial_layout = cute.make_ordered_layout( (self.num_threads // gmem_threads_per_row, gmem_threads_per_row), order=(1, 0), ) vOpartial_layout = cute.make_layout((1, async_copy_elems)) # 4 vals per load self.gmem_tiled_copy_O_partial = cute.make_tiled_copy_tv(atom_async_copy_partial, tOpartial_layout, vOpartial_layout) # GMEM copy setup for final O (use universal copy for store) atom_universal_copy = cute.make_copy_atom( cute.nvgpu.CopyUniversalOp(), self.dtype, num_bits_per_copy=async_copy_elems * self.dtype.width, ) self.gmem_tiled_copy_O = cute.make_tiled_copy_tv( atom_universal_copy, tOpartial_layout, vOpartial_layout, # 4 vals per store ) # LSE copy setup with async copy (alignment = 1) lse_copy_bits = Float32.width # 1 element per copy, width is in bits m_block_smem = ( 128 if self.tile_m % 128 == 0 else (64 if self.tile_m % 64 == 0 else (32 if self.tile_m % 32 == 0 else (16 if self.tile_m % 16 == 0 else 8))) ) gmem_threads_per_row_lse = m_block_smem assert self.num_threads % gmem_threads_per_row_lse == 0 # Async copy atom for LSE load atom_async_copy_lse = cute.make_copy_atom( cpasync.CopyG2SOp(cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS), Float32, num_bits_per_copy=lse_copy_bits, ) tLSE_layout = cute.make_ordered_layout( (self.num_threads // gmem_threads_per_row_lse, gmem_threads_per_row_lse), order=(1, 0), ) vLSE_layout = cute.make_layout(1) self.gmem_tiled_copy_LSE = cute.make_tiled_copy_tv(atom_async_copy_lse, tLSE_layout, vLSE_layout) # /////////////////////////////////////////////////////////////////////////////// # Shared memory # /////////////////////////////////////////////////////////////////////////////// # Shared memory to register copy for LSE self.smem_threads_per_col_lse = self.num_threads // m_block_smem assert 32 % self.smem_threads_per_col_lse == 0 # Must divide warp size s2r_layout_atom_lse = cute.make_ordered_layout( (self.smem_threads_per_col_lse, self.num_threads // self.smem_threads_per_col_lse), order=(0, 1), ) self.s2r_tiled_copy_LSE = cute.make_tiled_copy_tv( cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), Float32), s2r_layout_atom_lse, cute.make_layout(1), ) # LSE shared memory layout with swizzling to avoid bank conflicts # This works for kBlockMSmem = 8, 16, 32, 64, 128, no bank conflicts if const_expr(m_block_smem == 8): smem_lse_swizzle = cute.make_swizzle(5, 0, 5) elif const_expr(m_block_smem == 16): smem_lse_swizzle = cute.make_swizzle(4, 0, 4) else: smem_lse_swizzle = cute.make_swizzle(3, 2, 3) smem_layout_atom_lse = cute.make_composed_layout(smem_lse_swizzle, 0, cute.make_ordered_layout((8, m_block_smem), order=(1, 0))) self.smem_layout_lse = cute.tile_to_shape(smem_layout_atom_lse, (self.max_splits, self.tile_m), (0, 1)) # O partial shared memory layout (simple layout for pipeline stages) self.smem_layout_o = cute.make_ordered_layout((self.tile_m, self.k_block_size, self.stages), order=(1, 0, 2)) @cute.jit def __call__( self, mO_partial: cute.Tensor, mLSE_partial: cute.Tensor, mO: cute.Tensor, mLSE: Optional[cute.Tensor] = None, cu_seqlens: Optional[cute.Tensor] = None, seqused: Optional[cute.Tensor] = None, num_splits_dynamic_ptr: Optional[cute.Tensor] = None, varlen_batch_idx: Optional[cute.Tensor] = None, semaphore_to_reset: Optional[cute.Tensor] = None, # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI). stream: cuda.CUstream = None, ): # Type checking if const_expr(not (mO_partial.element_type == self.partial_dtype)): raise TypeError( "O partial tensor must match the configured partial dtype" ) if const_expr(not (mO.element_type == self.dtype)): raise TypeError("O tensor must match dtype") if const_expr(mLSE_partial.element_type not in [Float32]): raise TypeError("LSE partial tensor must be Float32") if const_expr(mLSE is not None and mLSE.element_type not in [Float32]): raise TypeError("LSE tensor must be Float32") # Shape validation - input tensors are in user format, need to be converted to kernel format if const_expr(len(mO_partial.shape) not in [4, 5]): raise ValueError( "O partial tensor must have 4 or 5 dimensions: (num_splits, batch, seqlen, nheads, headdim) or (num_splits, total_q, nheads, headdim)" ) if const_expr(len(mLSE_partial.shape) not in [3, 4]): raise ValueError("LSE partial tensor must have 3 or 4 dimensions: (num_splits, batch, seqlen, nheads) or (num_splits, total_q, nheads)") if const_expr(len(mO.shape) not in [3, 4]): raise ValueError("O tensor must have 3 or 4 dimensions: (batch, seqlen, nheads, headdim) or (total_q, nheads, headdim)") if const_expr(mLSE is not None and len(mLSE.shape) not in [2, 3]): raise ValueError("LSE tensor must have 2 or 3 dimensions: (batch, seqlen, nheads) or (total_q, nheads)") mO_partial, mO = [assume_tensor_aligned(t) for t in (mO_partial, mO)] # (num_splits, b, seqlen, h, d) -> (seqlen, d, num_splits, h, b) # or (num_splits, total_q, h, d) -> (total_q, d, num_splits, h) O_partial_layout_transpose = [2, 4, 0, 3, 1] if const_expr(cu_seqlens is None) else [1, 3, 0, 2] # (b, seqlen, h, d) -> (seqlen, d, h, b) or (total_q, h, d) -> (total_q, d, h) mO_partial = cute.make_tensor(mO_partial.iterator, cute.select(mO_partial.layout, mode=O_partial_layout_transpose)) O_layout_transpose = [1, 3, 2, 0] if const_expr(cu_seqlens is None) else [0, 2, 1] mO = cute.make_tensor(mO.iterator, cute.select(mO.layout, mode=O_layout_transpose)) # (num_splits, b, seqlen, h) -> (seqlen, num_splits, h, b) # or (num_splits, total_q, h) -> (total_q, num_splits, h) LSE_partial_layout_transpose = [2, 0, 3, 1] if const_expr(cu_seqlens is None) else [1, 0, 2] mLSE_partial = cute.make_tensor( mLSE_partial.iterator, cute.select(mLSE_partial.layout, mode=LSE_partial_layout_transpose), ) # (b, seqlen, h) -> (seqlen, h, b) or (total_q, h) -> (total_q, h) LSE_layout_transpose = [1, 2, 0] if const_expr(cu_seqlens is None) else [0, 1] mLSE = cute.make_tensor(mLSE.iterator, cute.select(mLSE.layout, mode=LSE_layout_transpose)) if mLSE is not None else None # Determine if we have variable length sequences varlen = const_expr(cu_seqlens is not None or seqused is not None) self._setup_attributes() @cute.struct class SharedStorage: sLSE: cute.struct.Align[cute.struct.MemRange[Float32, cute.cosize(self.smem_layout_lse)], 128] sMaxValidSplit: cute.struct.Align[cute.struct.MemRange[Int32, self.tile_m], 128] sO: cute.struct.Align[cute.struct.MemRange[self.partial_dtype, cute.cosize(self.smem_layout_o)], 128] smem_size = SharedStorage.size_in_bytes() # Grid dimensions: (ceil_div(seqlen, m_block), ceil_div(head_dim, k_block), num_head * batch) seqlen = mO_partial.shape[0] num_head = mO_partial.shape[3] batch_size = mO_partial.shape[4] if const_expr(cu_seqlens is None) else Int32(cu_seqlens.shape[0] - 1) # Create FastDivmodDivisor objects for efficient division seqlen_divmod = FastDivmodDivisor(seqlen) grid_dim = ( cute.ceil_div(seqlen * num_head, self.tile_m), cute.ceil_div(self.head_dim, self.k_block_size), batch_size, ) self.kernel( mO_partial, mLSE_partial, mO, mLSE, cu_seqlens, seqused, num_splits_dynamic_ptr, varlen_batch_idx, semaphore_to_reset, SharedStorage, self.smem_layout_lse, self.smem_layout_o, self.gmem_tiled_copy_O_partial, self.gmem_tiled_copy_O, self.gmem_tiled_copy_LSE, self.s2r_tiled_copy_LSE, seqlen_divmod, varlen, ).launch( grid=grid_dim, block=[self.num_threads, 1, 1], smem=smem_size, stream=stream, ) @cute.kernel def kernel( self, mO_partial: cute.Tensor, mLSE_partial: cute.Tensor, mO: cute.Tensor, mLSE: Optional[cute.Tensor], cu_seqlens: Optional[cute.Tensor], seqused: Optional[cute.Tensor], num_splits_dynamic_ptr: Optional[cute.Tensor], varlen_batch_idx: Optional[cute.Tensor], semaphore_to_reset: Optional[cute.Tensor], SharedStorage: cutlass.Constexpr, smem_layout_lse: cute.Layout | cute.ComposedLayout, smem_layout_o: cute.Layout, gmem_tiled_copy_O_partial: cute.TiledCopy, gmem_tiled_copy_O: cute.TiledCopy, gmem_tiled_copy_LSE: cute.TiledCopy, s2r_tiled_copy_LSE: cute.TiledCopy, seqlen_divmod: FastDivmodDivisor, varlen: cutlass.Constexpr[bool], ): # Thread and block indices tidx, _, _ = cute.arch.thread_idx() m_block, k_block, maybe_virtual_batch = cute.arch.block_idx() # Map virtual batch index to real batch index (for persistent tile schedulers) batch_idx = varlen_batch_idx[maybe_virtual_batch] if const_expr(varlen_batch_idx is not None) else maybe_virtual_batch # /////////////////////////////////////////////////////////////////////////////// # Get shared memory buffer # /////////////////////////////////////////////////////////////////////////////// smem = cutlass.utils.SmemAllocator() storage = smem.allocate(SharedStorage) sLSE = storage.sLSE.get_tensor(smem_layout_lse) sMaxValidSplit = storage.sMaxValidSplit.get_tensor((self.tile_m,)) sO = storage.sO.get_tensor(smem_layout_o) # Handle semaphore reset after dependent grids complete. if const_expr(semaphore_to_reset is not None): if ( tidx == 0 and m_block == cute.arch.grid_dim()[0] - 1 and k_block == cute.arch.grid_dim()[1] - 1 and maybe_virtual_batch == cute.arch.grid_dim()[2] - 1 ): cute.arch.griddepcontrol_wait() semaphore_to_reset[0] = 0 # Get number of splits (use maybe_virtual_batch for per-batch-slot splits) num_splits = num_splits_dynamic_ptr[maybe_virtual_batch] if const_expr(num_splits_dynamic_ptr is not None) else mLSE_partial.shape[1] if const_expr(cu_seqlens is None): seqlen = mO_partial.shape[0] offset = Int32(0) else: offset = cu_seqlens[batch_idx] seqlen = seqused[batch_idx] if const_expr(seqused is not None) else cu_seqlens[batch_idx + 1] - offset # Extract number of heads (head index will be determined dynamically) num_head = mO_partial.shape[3] max_idx = seqlen * num_head # Early exit for single split if dynamic if (const_expr(num_splits_dynamic_ptr is None) or num_splits > 1) and (const_expr(not varlen) or m_block * self.tile_m < max_idx): # The BSA wrapper launches pre-schedule, partial attention, and # combine on the same stream without CUDA dependent-grid launch # attributes. Stream ordering already guarantees producer # completion, so do not use griddepcontrol_wait here. # =============================== # Step 1: Load LSE_partial from gmem to shared memory # =============================== if const_expr(cu_seqlens is None): mLSE_partial_cur = mLSE_partial[None, None, None, batch_idx] else: mLSE_partial_cur = cute.domain_offset((offset, 0, 0), mLSE_partial) mLSE_partial_copy = cute.tiled_divide(mLSE_partial_cur, (1,)) gmem_thr_copy_LSE = gmem_tiled_copy_LSE.get_slice(tidx) tLSEsLSE = gmem_thr_copy_LSE.partition_D(sLSE) # Create identity tensor for coordinate tracking cLSE = cute.make_identity_tensor((self.max_splits, self.tile_m)) tLSEcLSE = gmem_thr_copy_LSE.partition_S(cLSE) # Load LSE partial values for m in cutlass.range(cute.size(tLSEcLSE, mode=[2]), unroll_full=True): mi = tLSEcLSE[0, 0, m][1] # Get m coordinate idx = m_block * self.tile_m + mi if idx < max_idx: # Calculate actual sequence position and head using FastDivmodDivisor if const_expr(not varlen): head_idx, m_idx = divmod(idx, seqlen_divmod) else: head_idx = idx // seqlen m_idx = idx - head_idx * seqlen mLSE_partial_cur_copy = mLSE_partial_copy[None, m_idx, None, head_idx] for s in cutlass.range(cute.size(tLSEcLSE, mode=[1]), unroll_full=True): si = tLSEcLSE[0, s, 0][0] # Get split coordinate if si < num_splits: cute.copy( gmem_thr_copy_LSE, mLSE_partial_cur_copy[None, si], tLSEsLSE[None, s, m], ) else: tLSEsLSE[None, s, m].fill(-Float32.inf) else: # Rows past max_idx never write O/LSE to gmem, but their sLSE slots # still feed the max-valid-split reduction, which bounds the # partial-O accumulation loop. Fill them with -inf so such a row # yields max_valid_split == -1 instead of reading garbage. for s in cutlass.range(cute.size(tLSEcLSE, mode=[1]), unroll_full=True): tLSEsLSE[None, s, m].fill(-Float32.inf) cute.arch.cp_async_commit_group() # =============================== # Step 2: Load O_partial for pipeline stages # =============================== gmem_thr_copy_O_partial = gmem_tiled_copy_O_partial.get_slice(tidx) cO = cute.make_identity_tensor((self.tile_m, self.k_block_size)) tOcO = gmem_thr_copy_O_partial.partition_D(cO) tOsO_partial = gmem_thr_copy_O_partial.partition_D(sO) if const_expr(cu_seqlens is None): mO_partial_cur = mO_partial[None, None, None, None, batch_idx] else: mO_partial_cur = cute.domain_offset((offset, 0, 0, 0), mO_partial) # Precompute these values to avoid recomputing them in the loop num_rows = const_expr(cute.size(tOcO, mode=[1])) tOmidx = cute.make_rmem_tensor(num_rows, cutlass.Int32) tOhidx = cute.make_rmem_tensor(num_rows, cutlass.Int32) tOrOptr = cute.make_rmem_tensor(num_rows, cutlass.Int64) for m in cutlass.range(num_rows, unroll_full=True): mi = tOcO[0, m, 0][0] # m coordinate idx = m_block * self.tile_m + mi if const_expr(not varlen): tOhidx[m], tOmidx[m] = divmod(idx, seqlen_divmod) else: tOhidx[m] = idx // seqlen tOmidx[m] = idx - tOhidx[m] * seqlen tOrOptr[m] = utils.elem_pointer(mO_partial_cur, (tOmidx[m], k_block * self.k_block_size, 0, tOhidx[m])).toint() if idx >= max_idx: tOhidx[m] = -1 tOpO = None if const_expr(not self.is_even_k): tOpO = cute.make_rmem_tensor(cute.size(tOcO, mode=[2]), Boolean) for k in cutlass.range(cute.size(tOpO), unroll_full=True): tOpO[k] = tOcO[0, 0, k][1] < mO_partial.shape[1] - k_block * self.k_block_size load_O_partial = partial( self.load_O_partial, gmem_tiled_copy_O_partial, tOrOptr, tOsO_partial, tOhidx, tOpO, tOcO, mO_partial_cur.layout, ) # Load first few stages of O_partial for stage in cutlass.range(self.stages - 1, unroll_full=True): if stage < num_splits: load_O_partial(stage, stage) cute.arch.cp_async_commit_group() # =============================== # Step 3: Load and transpose LSE from smem to registers # =============================== # Wait for LSE and initial O partial stages to complete cute.arch.cp_async_wait_group(self.stages - 1) cute.arch.sync_threads() s2r_thr_copy_LSE = s2r_tiled_copy_LSE.get_slice(tidx) ts2rsLSE = s2r_thr_copy_LSE.partition_S(sLSE) ts2rrLSE = cute.make_rmem_tensor_like(ts2rsLSE) cute.copy(s2r_tiled_copy_LSE, ts2rsLSE, ts2rrLSE) # =============================== # Step 4: Compute final LSE along split dimension # =============================== lse_sum = cute.make_rmem_tensor(cute.size(ts2rrLSE, mode=[2]), Float32) ts2rcLSE = s2r_thr_copy_LSE.partition_D(cLSE) # We compute the max valid split for each row to short-circuit the computation later max_valid_split = cute.make_rmem_tensor(cute.size(ts2rrLSE, mode=[2]), Int32) assert cute.size(ts2rrLSE, mode=[0]) == 1 # Compute max, scales, and final LSE for each row for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True): # Find max LSE value across splits threads_per_col = const_expr(self.smem_threads_per_col_lse) lse_max = cute.arch.warp_reduction_max( ts2rrLSE[None, None, m].load().reduce(cute.ReductionOp.MAX, init_val=-Float32.inf, reduction_profile=0), threads_in_group=threads_per_col, ) # Find max valid split index max_valid_idx = -1 for s in cutlass.range(cute.size(ts2rrLSE, mode=[1]), unroll_full=True): if ts2rrLSE[0, s, m] != -Float32.inf: max_valid_idx = ts2rcLSE[0, s, 0][0] # Get split coordinate max_valid_split[m] = cute.arch.warp_reduction_max(max_valid_idx, threads_in_group=threads_per_col) # Compute exp scales and sum lse_max_cur = 0.0 if lse_max == -Float32.inf else lse_max # In case all local LSEs are -inf LOG2_E = math.log2(math.e) lse_sum_cur = 0.0 for s in cutlass.range(cute.size(ts2rrLSE, mode=[1]), unroll_full=True): scale = cute.math.exp2(ts2rrLSE[0, s, m] * LOG2_E - (lse_max_cur * LOG2_E), fastmath=True) lse_sum_cur += scale ts2rrLSE[0, s, m] = scale # Store scale for later use lse_sum_cur = cute.arch.warp_reduction_sum(lse_sum_cur, threads_in_group=threads_per_col) lse_sum[m] = cute.math.log(lse_sum_cur, fastmath=True) + lse_max # Normalize scales inv_sum = 0.0 if (lse_sum_cur == 0.0 or lse_sum_cur != lse_sum_cur) else 1.0 / lse_sum_cur ts2rrLSE[None, None, m].store(ts2rrLSE[None, None, m].load() * inv_sum) # Store the scales exp(lse - lse_logsum) back to smem cute.copy(s2r_tiled_copy_LSE, ts2rrLSE, ts2rsLSE) # Store max valid split to smem for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True): if ts2rcLSE[0, 0, m][0] == 0: # Only thread responsible for s=0 writes mi = ts2rcLSE[0, 0, m][1] if mi < self.tile_m: sMaxValidSplit[mi] = max_valid_split[m] # =============================== # Step 5: Store final LSE to gmem # =============================== if const_expr(mLSE is not None): if const_expr(cu_seqlens is None): mLSE_cur = mLSE[None, None, batch_idx] else: mLSE_cur = cute.domain_offset((offset, 0), mLSE) if k_block == 0: # Only first k_block writes LSE when mLSE is provided for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True): if ts2rcLSE[0, 0, m][0] == 0: # Only thread responsible for s=0 writes mi = ts2rcLSE[0, 0, m][1] idx = m_block * self.tile_m + mi if idx < max_idx: if const_expr(not varlen): head_idx, m_idx = divmod(idx, seqlen_divmod) else: head_idx = idx // seqlen m_idx = idx - head_idx * seqlen mLSE_cur[m_idx, head_idx] = lse_sum[m] # =============================== # Step 6: Read O_partial and accumulate final O # =============================== cute.arch.sync_threads() # Get max valid split for this thread thr_max_valid_split = sMaxValidSplit[tOcO[0, 0, 0][0]] for m in cutlass.range(1, cute.size(tOcO, mode=[1]), unroll_full=True): thr_max_valid_split = max(thr_max_valid_split, sMaxValidSplit[tOcO[0, m, 0][0]]) tOrO_partial = cute.make_rmem_tensor_like(tOsO_partial[None, None, None, 0]) tOrO = cute.make_rmem_tensor_like(tOrO_partial, Float32) tOrO.fill(0.0) stage_load = self.stages - 1 stage_compute = 0 # Main accumulation loop for s in cutlass.range(thr_max_valid_split + 1, unroll=4): # Get scales for this split scale = cute.make_rmem_tensor(num_rows, Float32) for m in cutlass.range(num_rows, unroll_full=True): scale[m] = sLSE[s, tOcO[0, m, 0][0]] # Get scale from smem # Load next stage if needed split_to_load = s + self.stages - 1 if split_to_load <= thr_max_valid_split: load_O_partial(split_to_load, stage_load) cute.arch.cp_async_commit_group() stage_load = 0 if stage_load == self.stages - 1 else stage_load + 1 # Wait for the current stage to be ready cute.arch.cp_async_wait_group(self.stages - 1) # We don't need __syncthreads() because each thread is just reading its own data from smem # Copy from smem to registers cute.autovec_copy(tOsO_partial[None, None, None, stage_compute], tOrO_partial) stage_compute = 0 if stage_compute == self.stages - 1 else stage_compute + 1 # Accumulate scaled partial results for m in cutlass.range(num_rows, unroll_full=True): if tOhidx[m] >= 0 and scale[m] > 0.0: tOrO[None, m, None].store(tOrO[None, m, None].load() + scale[m] * tOrO_partial[None, m, None].load().to(Float32)) # =============================== # Step 7: Write final O to gmem # =============================== rO = cute.make_rmem_tensor_like(tOrO, self.dtype) rO.store(tOrO.load().to(self.dtype)) if const_expr(cu_seqlens is None): mO_cur = mO[None, None, None, batch_idx] else: mO_cur = cute.domain_offset((offset, 0, 0), mO) mO_cur = utils.domain_offset_aligned((0, k_block * self.k_block_size, 0), mO_cur) elems_per_store = const_expr(cute.size(gmem_tiled_copy_O.layout_tv_tiled[1])) gmem_thr_copy_O = gmem_tiled_copy_O.get_slice(tidx) # Write final results for m in cutlass.range(num_rows, unroll_full=True): if tOhidx[m] >= 0: mO_cur_copy = cute.tiled_divide(mO_cur[tOmidx[m], None, tOhidx[m]], (elems_per_store,)) for k in cutlass.range(cute.size(tOcO, mode=[2]), unroll_full=True): k_idx = tOcO[0, 0, k][1] // elems_per_store if const_expr(self.is_even_k) or tOpO[k]: cute.copy(gmem_thr_copy_O, rO[None, m, k], mO_cur_copy[None, k_idx]) @cute.jit def load_O_partial( self, gmem_tiled_copy_O_partial: cute.TiledCopy, tOrOptr: cute.Tensor, tOsO_partial: cute.Tensor, tOhidx: cute.Tensor, tOpO: Optional[cute.Tensor], tOcO: cute.Tensor, mO_cur_partial_layout: cute.Layout, split: Int32, stage: Int32, ) -> None: elems_per_load = const_expr(cute.size(gmem_tiled_copy_O_partial.layout_tv_tiled[1])) tOsO_partial_cur = tOsO_partial[None, None, None, stage] for m in cutlass.range(cute.size(tOcO, [1]), unroll_full=True): if tOhidx[m] >= 0: o_gmem_ptr = cute.make_ptr(tOsO_partial.element_type, tOrOptr[m], cute.AddressSpace.gmem, assumed_align=16) mO_partial_cur = cute.make_tensor(o_gmem_ptr, cute.slice_(mO_cur_partial_layout, (0, None, None, 0))) mO_partial_cur_copy = cute.tiled_divide(mO_partial_cur, (elems_per_load,)) for k in cutlass.range(cute.size(tOcO, mode=[2]), unroll_full=True): k_idx = tOcO[0, 0, k][1] // elems_per_load if const_expr(tOpO is None) or tOpO[k]: cute.copy( gmem_tiled_copy_O_partial, mO_partial_cur_copy[None, k_idx, split], tOsO_partial_cur[None, m, k], )