# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. # SM90 (Hopper) forward pass for flash attention, extracted from flash_fwd.py. from types import SimpleNamespace from typing import Callable, Optional from functools import partial import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute from cutlass import Float32, Int32, const_expr from cutlass.cute.nvgpu import cpasync, warpgroup from cutlass.utils import LayoutEnum import cutlass.utils.hopper_helpers as sm90_utils_basic from cutlass import pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait from cutlass.base_dsl.arch import Arch from ._compat import copy_utils from ._compat import layout_utils from ._compat import sm90_utils from .._vendor.flash_attn.cute.cute_dsl_utils import assume_tensor_aligned from .._vendor.flash_attn.cute import utils from .._vendor.flash_attn.cute.mask import AttentionMask from .._vendor.flash_attn.cute.softmax import Softmax, apply_score_mod_inner from .._vendor.flash_attn.cute.seqlen_info import SeqlenInfoQK from .._vendor.flash_attn.cute.block_info import BlockInfo from .._vendor.flash_attn.cute.block_sparsity import BlockSparseTensors from .._vendor.flash_attn.cute import pipeline as pipeline_custom from .._vendor.flash_attn.cute.pack_gqa import PackGQA, pack_gqa_layout, make_packgqa_tiled_tma_atom from .._vendor.flash_attn.cute.named_barrier import NamedBarrierFwd from ._compat.cute_dsl_utils import ParamsBase from .._vendor.flash_attn.cute.tile_scheduler import ( TileSchedulerArguments, SingleTileScheduler, SingleTileLPTScheduler, SingleTileVarlenScheduler, ) from .._vendor.flash_attn.cute.flash_fwd import FlashAttentionForwardBase from . import atoms as sol_attn_atoms from . import exact as exact_stream from ..common import selector as sol_attn_selector SOL_ATTN_ROUTE_MASK_BARRIER_ID = 7 SOL_ATTN_ROUTE_SUM_BARRIER_ID = 8 class SolAttnMainloopSm90(FlashAttentionForwardBase): def __init__( self, *args, sol_attn_assume_lane_group_route_reduce: bool = False, sol_attn_assume_full_k_exact_blocks: bool = False, sol_attn_tail_exact_words1: bool = False, sol_attn_assume_full_route_groups: bool = False, sol_attn_static_num_full_route_groups: int = -1, sol_attn_static_tail_valid_count: int = -1, sol_attn_tail_physical_tile16: bool = False, sol_attn_exact_mask_seqlen_last_only: bool = False, sol_attn_tail16_lane_group_route_reduce: bool = False, sol_attn_num_splits: int = 1, **kwargs, ): super().__init__(*args, **kwargs) self.qk_dtype = cutlass.BFloat16 self.pv_dtype = self.dtype self.sol_attn_group_size = 64 self.sol_attn_group_words = 2 self.mma_pv_is_rs = True self.sol_attn_mma_regs_override = 128 self.sol_attn_warp_route_mask = True self.sol_attn_fast_route_lens = True self.sol_attn_early_route_mask_publish = True self.sol_attn_lane_group_route_reduce = True self.sol_attn_assume_lane_group_route_reduce = sol_attn_assume_lane_group_route_reduce self.sol_attn_assume_full_k_exact_blocks = sol_attn_assume_full_k_exact_blocks self.sol_attn_route_sum_arrive_overlap = True self.sol_attn_route_mask_after_scale = True self.sol_attn_assume_full_route_groups = sol_attn_assume_full_route_groups self.sol_attn_static_num_full_route_groups = sol_attn_static_num_full_route_groups self.sol_attn_static_tail_valid_count = sol_attn_static_tail_valid_count self.sol_attn_tail_exact_words1 = ( sol_attn_tail_exact_words1 and 0 < self.sol_attn_static_tail_valid_count <= 32 ) self.sol_attn_tail_route_mask_words1 = False self.sol_attn_tail_physical_tile16 = ( sol_attn_tail_physical_tile16 and 0 < self.sol_attn_static_tail_valid_count <= 16 ) self.sol_attn_exact_mask_seqlen_last_only = sol_attn_exact_mask_seqlen_last_only self.sol_attn_full_route_mask_seqlen_false = True self.sol_attn_tail16_lane_group_route_reduce = ( sol_attn_tail16_lane_group_route_reduce and self.sol_attn_tail_physical_tile16 ) self.sol_attn_full_block_row_sum_prescale = False self.sol_attn_neutral_softmax_state = True self.sol_attn_assume_nonempty_rows = False self.sol_attn_ballot_mask = True self.sol_attn_approx_colmask = False self.sol_attn_packed_route_reduction = False self.sol_attn_num_splits = sol_attn_num_splits self.buffer_align_bytes = 1024 self.use_tma_KV = True self.cluster_shape_mn = (1, 1) if not (self.arch >= Arch.sm_90 and self.arch <= Arch.sm_90a): raise AssertionError("The Hopper backend requires SM90") def _get_smem_layout_atom(self): sQ_layout_atom = warpgroup.make_smem_layout_atom( sm90_utils_basic.get_smem_layout_atom( LayoutEnum.ROW_MAJOR, self.qk_dtype, self.tile_hdim ), self.qk_dtype, ) sK_layout_atom = sQ_layout_atom sV_layout_atom = warpgroup.make_smem_layout_atom( sm90_utils_basic.get_smem_layout_atom( LayoutEnum.ROW_MAJOR, self.pv_dtype, self.tile_hdimv ), self.pv_dtype, ) sO_layout_atom = sV_layout_atom if not self.mma_pv_is_rs: sP_layout_atom = warpgroup.make_smem_layout_atom( sm90_utils_basic.get_smem_layout_atom( LayoutEnum.ROW_MAJOR, self.pv_dtype, self.tile_n ), self.pv_dtype, ) else: sP_layout_atom = None return sQ_layout_atom, sK_layout_atom, sV_layout_atom, sO_layout_atom, sP_layout_atom def _get_tiled_mma(self): tiled_mma_qk = sm90_utils.make_tiled_mma( cutlass.BFloat16, "K", "K", self.tile_n, source="SS", atom_layout_mnk=(self.tile_m // 64, 1, 1), b_dtype=cutlass.BFloat16, acc_dtype=Float32, ) tiled_mma_pv = sol_attn_atoms.make_pv_mma( tile_m=self.tile_m, tile_v=self.tile_hdimv, ) return tiled_mma_qk, tiled_mma_pv @cute.jit def sol_attn_qk_gemm_zero_init( self, tiled_mma: cute.TiledMma, shape: cute.Shape, tCrA: cute.Tensor, tCrB: cute.Tensor, A_idx: Optional[Int32] = None, B_idx: Optional[Int32] = None, wg_wait: int = -1, swap_AB: bool = False, ) -> cute.Tensor: """Run BF16 QK WGMMA directly into an FP32 accumulator.""" return sm90_utils.gemm_zero_init( tiled_mma, shape, tCrA, tCrB, A_idx, B_idx, wg_wait, swap_AB, ) def _get_shared_storage_cls(self): sQ_struct, sK_struct = [ cute.struct.Align[ cute.struct.MemRange[self.qk_dtype, cute.cosize(layout)], self.buffer_align_bytes ] for layout in (self.sQ_layout, self.sK_layout) ] sV_struct = cute.struct.Align[ cute.struct.MemRange[self.pv_dtype, cute.cosize(self.sV_layout)], self.buffer_align_bytes, ] cosize_sQV = max(cute.cosize(self.sQ_layout), cute.cosize(self.sV_layout)) sQV_struct = cute.struct.Align[cute.struct.MemRange[self.pv_dtype, cosize_sQV], 1024] cosize_sP = cute.cosize(self.sP_layout) if const_expr(self.sP_layout is not None) else 0 sP_struct = cute.struct.Align[cute.struct.MemRange[self.pv_dtype, cosize_sP], 1024] route_mask_struct = cute.struct.Align[ cute.struct.MemRange[Int32, 4], 16 ] route_sums_struct = cute.struct.Align[ cute.struct.MemRange[Float32, 4 * self.tile_n], 16 ] # 1 stage * 2 for Q pipeline (full + empty), self.num_stages*2 for K, self.num_stages*2 for V, mbar_ptr_Q_struct = cute.struct.MemRange[cutlass.Int64, 1 * 2] mbar_ptr_K_struct = cute.struct.MemRange[cutlass.Int64, self.num_stages * 2] mbar_ptr_V_struct = cute.struct.MemRange[cutlass.Int64, self.num_stages * 2] @cute.struct class SharedStorageQKV: mbar_ptr_Q: mbar_ptr_Q_struct mbar_ptr_K: mbar_ptr_K_struct mbar_ptr_V: mbar_ptr_V_struct sV: sV_struct sQ: sQ_struct sK: sK_struct sP: sP_struct route_mask: route_mask_struct route_sums: route_sums_struct @cute.struct class SharedStorageSharedQV: mbar_ptr_Q: mbar_ptr_Q_struct mbar_ptr_K: mbar_ptr_K_struct mbar_ptr_V: mbar_ptr_V_struct sQ: sQV_struct sK: sK_struct sP: sP_struct route_mask: route_mask_struct route_sums: route_sums_struct return SharedStorageQKV if const_expr(not self.Q_in_regs) else SharedStorageSharedQV @cute.jit def sol_attn_reduce_route_sums_lane_group( self, acc_S_mn: cute.Tensor, route_sums: cute.Tensor, warp_in_mma: Int32, lane: Int32, ): """Reduce route columns using the observed SM90 accumulator lane layout.""" for col_group in cutlass.range_constexpr(self.tile_n // 8): base = col_group * 4 partial0 = Float32(acc_S_mn[base]) + Float32(acc_S_mn[base + 1]) partial1 = Float32(acc_S_mn[base + 2]) + Float32(acc_S_mn[base + 3]) partial0 += cute.arch.shuffle_sync_down(partial0, 16) partial1 += cute.arch.shuffle_sync_down(partial1, 16) partial0 += cute.arch.shuffle_sync_down(partial0, 8) partial1 += cute.arch.shuffle_sync_down(partial1, 8) partial0 += cute.arch.shuffle_sync_down(partial0, 4) partial1 += cute.arch.shuffle_sync_down(partial1, 4) if lane < Int32(4): col = Int32(8 * col_group) + lane * Int32(2) route_sums[warp_in_mma, col] = partial0 route_sums[warp_in_mma, col + Int32(1)] = partial1 @cute.jit def sol_attn_reduce_route_sums_guarded( self, acc_S_mn: cute.Tensor, route_sums: cute.Tensor, tScS_mn: cute.Tensor, q_start: Int32, seqlen: SeqlenInfoQK, warp_in_mma: Int32, lane: Int32, ): """Fallback route-column reduction that ignores invalid q rows.""" for off in cutlass.range_constexpr(self.tile_n): partial = Float32(0.0) for i in cutlass.range(cute.size(acc_S_mn), unroll_full=True): row = tScS_mn[i][0] col = tScS_mn[i][1] valid_row = q_start + row < seqlen.seqlen_q if col == Int32(off) and valid_row: partial += Float32(acc_S_mn[i]) warp_sum = cute.arch.warp_reduction_sum(partial) if lane == Int32(0): route_sums[warp_in_mma, off] = warp_sum @cute.jit def sol_attn_reduce_route_sums_lane_group_tail16( self, acc_S_mn: cute.Tensor, route_sums: cute.Tensor, route_col_offset: Int32, warp_in_mma: Int32, lane: Int32, ): """Reduce a physical 16-column tail route tile using the accumulator lane layout.""" for col_group in cutlass.range_constexpr(2): base = col_group * 4 partial0 = Float32(acc_S_mn[base]) + Float32(acc_S_mn[base + 1]) partial1 = Float32(acc_S_mn[base + 2]) + Float32(acc_S_mn[base + 3]) partial0 += cute.arch.shuffle_sync_down(partial0, 16) partial1 += cute.arch.shuffle_sync_down(partial1, 16) partial0 += cute.arch.shuffle_sync_down(partial0, 8) partial1 += cute.arch.shuffle_sync_down(partial1, 8) partial0 += cute.arch.shuffle_sync_down(partial0, 4) partial1 += cute.arch.shuffle_sync_down(partial1, 4) if lane < Int32(4): col = route_col_offset + Int32(8 * col_group) + lane * Int32(2) route_sums[warp_in_mma, col] = partial0 route_sums[warp_in_mma, col + Int32(1)] = partial1 @cute.jit def sol_attn_reduce_route_sums_static_tail( self, acc_S_mn: cute.Tensor, route_sums: cute.Tensor, tScS_mn: cute.Tensor, q_start: Int32, route_col_offset: Int32, seqlen: SeqlenInfoQK, warp_in_mma: Int32, lane: Int32, full_q_tile: bool, ): """Reduce only the compile-time-known valid columns of a static tail route group.""" for off in cutlass.range_constexpr(self.sol_attn_static_tail_valid_count): route_col = route_col_offset + Int32(off) partial = Float32(0.0) for i in cutlass.range(cute.size(acc_S_mn), unroll_full=True): row = tScS_mn[i][0] col = tScS_mn[i][1] valid_row = True if not full_q_tile: valid_row = q_start + row < seqlen.seqlen_q if col == route_col and valid_row: partial += Float32(acc_S_mn[i]) warp_sum = cute.arch.warp_reduction_sum(partial) if lane == Int32(0): route_sums[warp_in_mma, route_col] = warp_sum @cute.jit def sol_attn_reduce_route_sums_physical16( self, acc_S_mn: cute.Tensor, route_sums: cute.Tensor, tScS_mn: cute.Tensor, q_start: Int32, route_col_offset: Int32, seqlen: SeqlenInfoQK, warp_in_mma: Int32, lane: Int32, full_q_tile: bool, ): """Reduce a physically 16-column route tile into the 64-column route_sums buffer.""" for off in cutlass.range_constexpr(16): route_col = route_col_offset + Int32(off) partial = Float32(0.0) for i in cutlass.range(cute.size(acc_S_mn), unroll_full=True): row = tScS_mn[i][0] col = tScS_mn[i][1] valid_row = True if not full_q_tile: valid_row = q_start + row < seqlen.seqlen_q if col == route_col and valid_row: partial += Float32(acc_S_mn[i]) warp_sum = cute.arch.warp_reduction_sum(partial) if lane == Int32(0): route_sums[warp_in_mma, route_col] = warp_sum @cute.jit def sol_attn_build_route_mask_from_acc( self, acc_S: cute.Tensor, route_sums: cute.Tensor, tScS_mn: cute.Tensor, m_block: Int32, group_start_n_block: Int32, valid_count: Int32, route_col_offset: Int32, seqlen: SeqlenInfoQK, batch_idx: Int32, head_idx: Int32, mGlobalThresh: cute.Tensor, softmax_scale_log2: Float32, sink_range: Int32, assume_full_route_group: cutlass.Constexpr[bool] = False, physical_route_tile_n: cutlass.Constexpr[int] = 64, route_mask_words_override: cutlass.Constexpr[int] = 0, ): """Build the exact mask from the distributed route QK accumulator. WGMMA accumulators are distributed across the 128 consumer threads. Each consumer warp first reduces its local contribution per route column, then one consumer thread combines the four warp partials into the CTA-local bitmask. """ tidx, _, _ = cute.arch.thread_idx() consumer_tidx = tidx warp_in_mma = consumer_tidx // cute.arch.WARP_SIZE lane = cute.arch.lane_idx() acc_S_mn = layout_utils.reshape_acc_to_mn(acc_S) q_start = m_block * self.tile_m q_len_i32 = seqlen.seqlen_q - q_start if q_len_i32 > Int32(self.tile_m): q_len_i32 = Int32(self.tile_m) q_len = Float32(q_len_i32) full_q_tile = q_len_i32 == Int32(self.tile_m) if const_expr(self.sol_attn_tail16_lane_group_route_reduce and physical_route_tile_n == 16): if full_q_tile: self.sol_attn_reduce_route_sums_lane_group_tail16( acc_S_mn, route_sums, route_col_offset, warp_in_mma, lane, ) else: self.sol_attn_reduce_route_sums_static_tail( acc_S_mn, route_sums, tScS_mn, q_start, route_col_offset, seqlen, warp_in_mma, lane, full_q_tile, ) elif const_expr(physical_route_tile_n == 16): self.sol_attn_reduce_route_sums_physical16( acc_S_mn, route_sums, tScS_mn, q_start, route_col_offset, seqlen, warp_in_mma, lane, full_q_tile, ) elif const_expr(self.sol_attn_assume_lane_group_route_reduce): self.sol_attn_reduce_route_sums_lane_group(acc_S_mn, route_sums, warp_in_mma, lane) elif const_expr(self.sol_attn_lane_group_route_reduce): if full_q_tile: self.sol_attn_reduce_route_sums_lane_group(acc_S_mn, route_sums, warp_in_mma, lane) else: self.sol_attn_reduce_route_sums_guarded( acc_S_mn, route_sums, tScS_mn, q_start, seqlen, warp_in_mma, lane, ) else: for off in cutlass.range_constexpr(self.tile_n): partial = Float32(0.0) for i in cutlass.range(cute.size(acc_S_mn), unroll_full=True): row = tScS_mn[i][0] col = tScS_mn[i][1] valid_row = q_start + row < seqlen.seqlen_q if col == Int32(off) and valid_row: partial += Float32(acc_S_mn[i]) warp_sum = cute.arch.warp_reduction_sum(partial) if lane == Int32(0): route_sums[warp_in_mma, off] = warp_sum if const_expr(self.sol_attn_route_sum_arrive_overlap and self.sol_attn_warp_route_mask): if warp_in_mma == Int32(0): cute.arch.barrier( barrier_id=SOL_ATTN_ROUTE_SUM_BARRIER_ID, number_of_threads=self.num_mma_threads, ) else: cute.arch.barrier_arrive( barrier_id=SOL_ATTN_ROUTE_SUM_BARRIER_ID, number_of_threads=self.num_mma_threads, ) else: cute.arch.barrier( barrier_id=SOL_ATTN_ROUTE_SUM_BARRIER_ID, number_of_threads=self.num_mma_threads, ) mask0 = Int32(0) mask1 = Int32(0) mask2 = Int32(0) mask3 = Int32(0) thresh = Float32(mGlobalThresh[m_block, head_idx, batch_idx]) if const_expr(self.sol_attn_warp_route_mask): route_mask_words = self.sol_attn_group_words if const_expr(route_mask_words_override != 0): route_mask_words = route_mask_words_override if const_expr(self.sol_attn_tail_route_mask_words1 and not assume_full_route_group): route_mask_words = 1 build_mask = warp_in_mma == Int32(0) if build_mask: sink_enabled = sink_range != Int32(0) sink_start_block = sink_range & Int32(0xFFFF) sink_end_block = (sink_range >> Int32(16)) & Int32(0xFFFF) if const_expr(self.sol_attn_packed_route_reduction): off0 = lane off1 = Int32(32) + lane route_col0 = route_col_offset + off0 route_col1 = route_col_offset + off1 col_sum0 = Float32(route_sums[0, route_col0]) + Float32( route_sums[1, route_col0] ) col_sum1 = Float32(route_sums[0, route_col1]) + Float32( route_sums[1, route_col1] ) col_sum0 += Float32(route_sums[2, route_col0]) col_sum1 += Float32(route_sums[2, route_col1]) col_sum0 += Float32(route_sums[3, route_col0]) col_sum1 += Float32(route_sums[3, route_col1]) col_mean0 = col_sum0 * softmax_scale_log2 / q_len col_mean1 = col_sum1 * softmax_scale_log2 / q_len exact0 = sol_attn_selector.sol_attn_route_is_exact( m_block, group_start_n_block + off0, col_mean0, thresh, True, ) exact1 = sol_attn_selector.sol_attn_route_is_exact( m_block, group_start_n_block + off1, col_mean1, thresh, True, ) if sink_enabled: exact0 = exact0 or ( group_start_n_block + off0 >= sink_start_block and group_start_n_block + off0 < sink_end_block ) exact1 = exact1 or ( group_start_n_block + off1 >= sink_start_block and group_start_n_block + off1 < sink_end_block ) word_bits0 = Int32(cute.arch.vote_ballot_sync(exact0)) word_bits1 = Int32(cute.arch.vote_ballot_sync(exact1)) if lane == Int32(0): mask0 = word_bits0 mask1 = word_bits1 else: for word in cutlass.range_constexpr(route_mask_words): off = Int32(word * 32) + lane route_col = route_col_offset + off valid = True if const_expr( not ( self.sol_attn_assume_full_route_groups or assume_full_route_group ) ): valid = off < valid_count exact = False if valid: col_sum = ( Float32(route_sums[0, route_col]) + Float32(route_sums[1, route_col]) + Float32(route_sums[2, route_col]) + Float32(route_sums[3, route_col]) ) col_mean = col_sum * softmax_scale_log2 / q_len exact = sol_attn_selector.sol_attn_route_is_exact( m_block, group_start_n_block + off, col_mean, thresh, valid, ) if sink_enabled: exact = exact or ( group_start_n_block + off >= sink_start_block and group_start_n_block + off < sink_end_block ) if const_expr(self.sol_attn_approx_colmask): column_mask = -Float32.inf if valid and not exact: column_mask = Float32(0.0) route_sums[0, route_col] = column_mask if const_expr(self.sol_attn_ballot_mask): word_bits = Int32(cute.arch.vote_ballot_sync(exact)) else: word_bits = Int32(0) if exact: word_bits = Int32(1) << lane word_bits = word_bits | cute.arch.shuffle_sync_down( word_bits, 16 ) word_bits = word_bits | cute.arch.shuffle_sync_down( word_bits, 8 ) word_bits = word_bits | cute.arch.shuffle_sync_down( word_bits, 4 ) word_bits = word_bits | cute.arch.shuffle_sync_down( word_bits, 2 ) word_bits = word_bits | cute.arch.shuffle_sync_down( word_bits, 1 ) if lane == Int32(0): if const_expr(word == 0): mask0 = word_bits elif const_expr(word == 1): mask1 = word_bits elif const_expr(word == 2): mask2 = word_bits else: mask3 = word_bits else: sink_enabled = sink_range != Int32(0) sink_start_block = sink_range & Int32(0xFFFF) sink_end_block = (sink_range >> Int32(16)) & Int32(0xFFFF) for off in cutlass.range_constexpr(self.sol_attn_group_size): route_col = route_col_offset + Int32(off) valid = True if const_expr( not (self.sol_attn_assume_full_route_groups or assume_full_route_group) ): valid = Int32(off) < valid_count col_sum = ( Float32(route_sums[0, route_col]) + Float32(route_sums[1, route_col]) + Float32(route_sums[2, route_col]) + Float32(route_sums[3, route_col]) ) if valid: col_mean = col_sum * softmax_scale_log2 / q_len exact = sol_attn_selector.sol_attn_route_is_exact( m_block, group_start_n_block + Int32(off), col_mean, thresh, valid, ) if sink_enabled: exact = exact or ( group_start_n_block + Int32(off) >= sink_start_block and group_start_n_block + Int32(off) < sink_end_block ) if exact: mask0, mask1, mask2, mask3 = ( sol_attn_selector.sol_attn_set_exact_bit( mask0, mask1, mask2, mask3, Int32(off) ) ) return mask0, mask1, mask2, mask3 @cute.jit def sol_attn_mask_route_approx_columns( self, acc_S: cute.Tensor, route_sums: cute.Tensor, tScS_mn: cute.Tensor, valid_count: Int32, route_col_offset: Int32, mask0: Int32, mask1: Int32, mask2: Int32, mask3: Int32, assume_full_route_group: cutlass.Constexpr[bool] = False, route_mask_words_override: cutlass.Constexpr[int] = 0, ): """Keep only approximate route columns in the route score tile.""" acc_S_mn = layout_utils.reshape_acc_to_mn(acc_S) if const_expr(self.sol_attn_approx_colmask): for i in cutlass.range(cute.size(acc_S_mn), unroll_full=True): col = tScS_mn[i][1] acc_S_mn[i] = Float32(acc_S_mn[i]) + Float32(route_sums[0, col]) else: for i in cutlass.range(cute.size(acc_S_mn), unroll_full=True): col = tScS_mn[i][1] group_col = col - route_col_offset valid = True if const_expr(self.sol_attn_group_size != self.tile_n): valid = group_col >= Int32(0) if valid: valid = group_col < valid_count elif const_expr( not (self.sol_attn_assume_full_route_groups or assume_full_route_group) ): valid = col < valid_count exact = False if valid: route_mask_words = self.sol_attn_group_words if const_expr(route_mask_words_override != 0): route_mask_words = route_mask_words_override if const_expr( self.sol_attn_tail_route_mask_words1 and not assume_full_route_group ): route_mask_words = 1 exact = sol_attn_selector.sol_attn_test_exact_bit_limited_words( mask0, mask1, mask2, mask3, group_col, route_mask_words ) if (not valid) or exact: acc_S_mn[i] = -Float32.inf @cute.jit def sol_attn_expand_route_acc_to_full_tile( self, acc_S: cute.Tensor, acc_S_full_ref: cute.Tensor, tScS_mn: cute.Tensor, tScS_full_mn: cute.Tensor, ) -> cute.Tensor: """Expand a narrow physical route accumulator into a full 64-column P tile.""" acc_S_mn = layout_utils.reshape_acc_to_mn(acc_S) acc_S_full = cute.make_rmem_tensor_like(acc_S_full_ref, Float32) acc_S_full_mn = layout_utils.reshape_acc_to_mn(acc_S_full) for i in cutlass.range(cute.size(acc_S_full_mn), unroll_full=True): row = tScS_full_mn[i][0] col = tScS_full_mn[i][1] value = -Float32.inf if col < Int32(16): for j in cutlass.range(cute.size(acc_S_mn), unroll_full=True): row16 = tScS_mn[j][0] col16 = tScS_mn[j][1] if row == row16 and col == col16: value = Float32(acc_S_mn[j]) acc_S_full_mn[i] = value return acc_S_full @cute.jit def sol_attn_expand_route_prob_to_full_tile( self, acc_P: cute.Tensor, acc_S_full_ref: cute.Tensor, tScS_mn: cute.Tensor, tScS_full_mn: cute.Tensor, ) -> cute.Tensor: """Expand a compact route probability tile into the full PV A fragment.""" acc_P_mn = layout_utils.reshape_acc_to_mn(acc_P) acc_P_full = cute.make_rmem_tensor_like(acc_S_full_ref, Float32) acc_P_full_mn = layout_utils.reshape_acc_to_mn(acc_P_full) for i in cutlass.range(cute.size(acc_P_full_mn), unroll_full=True): row = tScS_full_mn[i][0] col = tScS_full_mn[i][1] value = Float32(0.0) if col < Int32(16): for j in cutlass.range(cute.size(acc_P_mn), unroll_full=True): row16 = tScS_mn[j][0] col16 = tScS_mn[j][1] if row == row16 and col == col16: value = Float32(acc_P_mn[j]) acc_P_full_mn[i] = value return acc_P_full @cute.jit def sol_attn_apply_route_current_lens_to_row_sum( self, acc_S: cute.Tensor, tScS_mn: cute.Tensor, group_start_n_block: Int32, valid_count: Int32, route_col_offset: Int32, seqlen: SeqlenInfoQK, softmax: Softmax, ): """Correct route approx denominator for VC tiles that are block sums.""" acc_S_mn = layout_utils.reshape_acc_to_mn(acc_S) last_n_block = ( (seqlen.seqlen_k + Int32(self.tile_n - 1)) // Int32(self.tile_n) ) - Int32(1) tail_len = seqlen.seqlen_k - last_n_block * Int32(self.tile_n) for r in cutlass.range(cute.size(softmax.row_sum), unroll_full=True): extra = Float32(0.0) for c in cutlass.range(cute.size(acc_S_mn.shape[1]), unroll_full=True): col = tScS_mn[r, c][1] group_col = col - route_col_offset valid = group_col >= Int32(0) if valid: valid = group_col < valid_count if valid: kv_block_idx = group_start_n_block + group_col current_len = Int32(self.tile_n) if kv_block_idx == last_n_block: current_len = tail_len extra += Float32(acc_S_mn[r, c]) * (Float32(current_len) - Float32(1.0)) softmax.row_sum[r] += extra @cute.jit def sol_attn_apply_route_current_lens_to_row_sum_fast( self, acc_S: cute.Tensor, tScS_mn: cute.Tensor, row_sum_prev: cute.Tensor, row_scale: cute.Tensor, group_start_n_block: Int32, valid_count: Int32, route_col_offset: Int32, seqlen: SeqlenInfoQK, softmax: Softmax, is_first_block: cutlass.Constexpr[bool], ): """Fast denominator correction for full-length route groups.""" last_n_block = ( (seqlen.seqlen_k + Int32(self.tile_n - 1)) // Int32(self.tile_n) ) - Int32(1) tail_len = seqlen.seqlen_k - last_n_block * Int32(self.tile_n) group_end = group_start_n_block + valid_count full_len_group = (tail_len == Int32(self.tile_n)) or (group_end <= last_n_block) if full_len_group: block_extra = Float32(self.tile_n - 1) for r in cutlass.range(cute.size(softmax.row_sum), unroll_full=True): prev_scaled = Float32(0.0) if const_expr(not is_first_block): prev_scaled = Float32(row_sum_prev[r]) * Float32(row_scale[r]) route_row_sum = Float32(softmax.row_sum[r]) - prev_scaled softmax.row_sum[r] += route_row_sum * block_extra else: # Tail blocks need per-column current_len because the last route # column may represent fewer than tile_n values. self.sol_attn_apply_route_current_lens_to_row_sum( acc_S, tScS_mn, group_start_n_block, valid_count, route_col_offset, seqlen, softmax, ) @cute.jit def __call__( self, mQ: cute.Tensor, mK: cute.Tensor, mV: cute.Tensor, mO: cute.Tensor, mKC: cute.Tensor, mVC: cute.Tensor, mGlobalThresh: cute.Tensor, mLSE: Optional[cute.Tensor], softmax_scale: Float32, sink_range: Int32, stream: cuda.CUstream = None, ): """Configure and launch the Hopper Sol-Attn kernel.""" mCuSeqlensQ = None mCuSeqlensK = None mSeqUsedQ = None mSeqUsedK = None mPageTable = None window_size_left = None window_size_right = None learnable_sink = None blocksparse_tensors = None piecewise_k = None piecewise_v = None aux_tensors = None self.varlen_q = mCuSeqlensQ is not None or mSeqUsedQ is not None mQ, mK, mV, mO, mKC, mVC, mGlobalThresh = [ assume_tensor_aligned(t) for t in (mQ, mK, mV, mO, mKC, mVC, mGlobalThresh) ] if const_expr(piecewise_k is not None): piecewise_k, piecewise_v = [ assume_tensor_aligned(t) for t in (piecewise_k, piecewise_v) ] SOL_ATTN_BTHD_TRANSPOSE = [1, 3, 2, 0] SOL_ATTN_BNH_TRANSPOSE = [1, 2, 0] mQ, mK, mV, mO, mKC, mVC = [ layout_utils.select(t, SOL_ATTN_BTHD_TRANSPOSE) for t in (mQ, mK, mV, mO, mKC, mVC) ] mGlobalThresh = layout_utils.select( mGlobalThresh, SOL_ATTN_BNH_TRANSPOSE ) if const_expr(piecewise_k is not None): piecewise_k, piecewise_v = [ layout_utils.select(t, SOL_ATTN_BTHD_TRANSPOSE) for t in (piecewise_k, piecewise_v) ] LSE_layout_transpose = [1, 2, 0] mLSE = ( layout_utils.select(mLSE, LSE_layout_transpose) if const_expr(mLSE is not None) else None ) tiled_mma_qk, tiled_mma_pv = self._get_tiled_mma() self.num_mma_threads = tiled_mma_qk.size self.num_threads_per_warp_group = 128 self.num_wg_mma = self.num_mma_threads // self.num_threads_per_warp_group assert self.num_wg_mma in [1, 2, 3] if const_expr(self.num_wg_mma != 1): raise NotImplementedError("SOL_ATTN SM90 path requires exactly one MMA warpgroup") self.num_threads = self.num_threads_per_warp_group self.num_producer_threads = 32 self.num_Q_load_threads = self.num_threads_per_warp_group # If not TMA_Q self.num_epilogue_threads = self.num_mma_threads self.num_mma_regs, self.num_producer_regs = {1: (256, 56), 2: (240, 24), 3: (160, 32)}[ self.num_wg_mma ] self.use_block_sparsity = cutlass.const_expr(blocksparse_tensors is not None) self.has_piecewise_kv = cutlass.const_expr(piecewise_k is not None) if const_expr(self.use_block_sparsity): raise NotImplementedError("one-warpgroup SOL_ATTN path does not support block sparsity") if const_expr(self.has_piecewise_kv): raise NotImplementedError("one-warpgroup SOL_ATTN path does not support piecewise KV") self.use_scheduler_barrier = self.num_wg_mma == 2 self.use_tma_Q = self.arch >= Arch.sm_90 and not ( self.pack_gqa and self.tile_m % self.qhead_per_kvhead != 0 ) if const_expr(not self.use_tma_Q): raise NotImplementedError("one-warpgroup SOL_ATTN path requires TMA Q/O") # FP32 split partials require a direct register-to-global epilogue. # A BF16 split partial matches V/O dtype and can reuse the shared-memory # plus TMA-O epilogue. self.use_tma_O = ( self.sol_attn_num_splits == 1 or mO.element_type == self.dtype ) # Producer needs more registers when doing cp.async Q or KV loads if const_expr(self.num_wg_mma == 2 and (not self.use_tma_Q or not self.use_tma_KV)): self.num_mma_regs, self.num_producer_regs = 224, 40 if const_expr(self.sol_attn_mma_regs_override is not None): self.num_mma_regs = self.sol_attn_mma_regs_override self.rescale_O_before_gemm = False self._setup_attributes() # TODO: we prob don't need most of what's in _setup_attributes self.sQ_layout, self.sK_layout, self.sV_layout, self.sO_layout = [ sm90_utils.make_smem_layout(mX.element_type, LayoutEnum.ROW_MAJOR, shape, stage) for mX, shape, stage in [ (mQ, (self.tile_m, self.tile_hdim), None), (mK, (self.tile_n, self.tile_hdim), self.num_stages), (mV, (self.tile_n, self.tile_hdimv), self.num_stages), # sO always holds the BF16 PV epilogue tile. Split-KV's # global mO is an FP32 partial workspace, so derive this # shared-memory layout from V instead of global O. (mV, (self.tile_m, self.tile_hdimv), None), ] ] self.sP_layout = None if const_expr(not self.mma_pv_is_rs): self.sP_layout = sm90_utils.make_smem_layout( mV.element_type, LayoutEnum.ROW_MAJOR, (self.tile_m, self.tile_n) ) SharedStorage = self._get_shared_storage_cls() mQ_og, mO_og = mQ, mO if const_expr(self.pack_gqa): nheads_kv = mK.shape[2] mQ = pack_gqa_layout(mQ, self.qhead_per_kvhead, nheads_kv, head_idx=2) mO = pack_gqa_layout(mO, self.qhead_per_kvhead, nheads_kv, head_idx=2) if const_expr(mLSE is not None): mLSE = pack_gqa_layout(mLSE, self.qhead_per_kvhead, nheads_kv, head_idx=1) # TMA gmem_tiled_copy_Q = cpasync.CopyBulkTensorTileG2SOp() gmem_tiled_copy_KV = cpasync.CopyBulkTensorTileG2SOp() # Might multicast gmem_tiled_copy_O = cpasync.CopyBulkTensorTileS2GOp() self.tma_copy_bytes = { name: cute.size_in_bytes(mX.element_type, cute.select(layout, mode=[0, 1])) for name, mX, layout in [ ("Q", mQ, self.sQ_layout), ("K", mK, self.sK_layout), ("V", mV, self.sV_layout), ] } make_tiled_tma_atom_fn = ( partial(make_packgqa_tiled_tma_atom, qhead_per_kvhead=self.qhead_per_kvhead, head_idx=2) if const_expr(self.pack_gqa) else cpasync.make_tiled_tma_atom ) tma_atom_Q, tma_tensor_Q = None, None if const_expr(self.use_tma_Q): tma_atom_Q, tma_tensor_Q = make_tiled_tma_atom_fn( gmem_tiled_copy_Q, mQ_og if const_expr(self.pack_gqa) else mQ, self.sQ_layout, (self.tile_m, self.tile_hdim), # No mcast ) tma_atom_K, tma_tensor_K = None, None tma_atom_V, tma_tensor_V = None, None tma_atom_KC, tma_tensor_KC = None, None tma_atom_VC, tma_tensor_VC = None, None tma_atom_K2, tma_tensor_K2 = None, None tma_atom_V2, tma_tensor_V2 = None, None if const_expr(self.use_tma_KV): tma_atom_K, tma_tensor_K = cpasync.make_tiled_tma_atom( gmem_tiled_copy_KV, mK, cute.select(self.sK_layout, mode=[0, 1]), (self.tile_n, self.tile_hdim), 1, # No mcast for now ) tma_atom_V, tma_tensor_V = cpasync.make_tiled_tma_atom( gmem_tiled_copy_KV, mV, cute.select(self.sV_layout, mode=[0, 1]), (self.tile_n, self.tile_hdimv), 1, # No mcast for now ) tma_atom_KC, tma_tensor_KC = cpasync.make_tiled_tma_atom( gmem_tiled_copy_KV, mKC, cute.select(self.sK_layout, mode=[0, 1]), (self.tile_n, self.tile_hdim), 1, ) tma_atom_VC, tma_tensor_VC = cpasync.make_tiled_tma_atom( gmem_tiled_copy_KV, mVC, cute.select(self.sV_layout, mode=[0, 1]), (self.tile_n, self.tile_hdimv), 1, ) if const_expr(self.has_piecewise_kv): tma_atom_K2, tma_tensor_K2 = cpasync.make_tiled_tma_atom( gmem_tiled_copy_KV, piecewise_k, cute.select(self.sK_layout, mode=[0, 1]), (self.tile_n, self.tile_hdim), 1, ) tma_atom_V2, tma_tensor_V2 = cpasync.make_tiled_tma_atom( gmem_tiled_copy_KV, piecewise_v, cute.select(self.sV_layout, mode=[0, 1]), (self.tile_n, self.tile_hdimv), 1, ) tma_atom_O, tma_tensor_O = None, None if const_expr(self.use_tma_O): mO_tma = mO_og if const_expr(self.pack_gqa) else mO if const_expr(self.varlen_q): mO_tma = copy_utils.create_ragged_tensor_for_tma( mO_tma, ragged_dim=0, ptr_shift=True ) tma_atom_O, tma_tensor_O = make_tiled_tma_atom_fn( gmem_tiled_copy_O, mO_tma, self.sO_layout, (self.tile_m, self.tile_hdimv), # No mcast ) if const_expr(mCuSeqlensQ is not None or mSeqUsedQ is not None): TileScheduler = SingleTileVarlenScheduler else: TileScheduler = ( SingleTileScheduler if const_expr(not self.is_causal or self.is_local) else SingleTileLPTScheduler ) tile_sched_args = TileSchedulerArguments( cute.ceil_div(cute.size(mQ.shape[0]), self.tile_m), cute.size(mQ.shape[2]), cute.size(mQ.shape[3]) if const_expr(mCuSeqlensQ is None) else cute.size(mCuSeqlensQ.shape[0] - 1), self.sol_attn_num_splits, cute.size(mK.shape[0]) if const_expr(mPageTable is None) else mK.shape[0] * mPageTable.shape[1], mQ.shape[1], mV.shape[1], total_q=cute.size(mQ.shape[0]) if const_expr(mCuSeqlensQ is not None) else cute.size(mQ.shape[0]) * cute.size(mQ.shape[3]), tile_shape_mn=(self.tile_m, self.tile_n), mCuSeqlensQ=mCuSeqlensQ, mSeqUsedQ=mSeqUsedQ, qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, element_size=self.dtype.width // 8, is_persistent=False, lpt=self.is_causal or self.is_local, is_split_kv=self.sol_attn_num_splits > 1, ) tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args) grid_dim = TileScheduler.get_grid_shape(tile_sched_params) softmax_scale_log2, softmax_scale = utils.compute_softmax_scale_log2( softmax_scale, self.score_mod ) window_size_left = Int32(window_size_left) if window_size_left is not None else None window_size_right = Int32(window_size_right) if window_size_right is not None else None fastdiv_mods = utils.compute_fastdiv_mods( mQ, mK, self.qhead_per_kvhead, self.pack_gqa, aux_tensors, mPageTable ) self.kernel( tma_tensor_Q if const_expr(self.use_tma_Q) else mQ, tma_tensor_K if const_expr(self.use_tma_KV) else mK, tma_tensor_V if const_expr(self.use_tma_KV) else mV, tma_tensor_KC if const_expr(self.use_tma_KV) else mKC, tma_tensor_VC if const_expr(self.use_tma_KV) else mVC, tma_tensor_K2 if const_expr(self.has_piecewise_kv) else None, tma_tensor_V2 if const_expr(self.has_piecewise_kv) else None, tma_tensor_O if const_expr(self.use_tma_O) else mO, mGlobalThresh, mLSE, mCuSeqlensQ, mCuSeqlensK, mSeqUsedQ, mSeqUsedK, mPageTable, tma_atom_Q, tma_atom_K, tma_atom_V, tma_atom_KC, tma_atom_VC, tma_atom_K2, tma_atom_V2, tma_atom_O, softmax_scale_log2, softmax_scale, sink_range, window_size_left, window_size_right, learnable_sink, blocksparse_tensors, self.sQ_layout, self.sK_layout, self.sV_layout, self.sO_layout, self.sP_layout, self.gmem_tiled_copy_Q, self.gmem_tiled_copy_K, self.gmem_tiled_copy_V, self.gmem_tiled_copy_O, tiled_mma_qk, tiled_mma_pv, tile_sched_params, TileScheduler, SharedStorage, fastdiv_mods, ).launch( grid=grid_dim, block=[self.num_threads, 1, 1], stream=stream, min_blocks_per_mp=1, ) @cute.kernel def kernel( self, mQ: cute.Tensor, mK: cute.Tensor, mV: cute.Tensor, mKC: cute.Tensor, mVC: cute.Tensor, mK2: Optional[cute.Tensor], mV2: Optional[cute.Tensor], mO: cute.Tensor, mGlobalThresh: cute.Tensor, mLSE: Optional[cute.Tensor], mCuSeqlensQ: Optional[cute.Tensor], mCuSeqlensK: Optional[cute.Tensor], mSeqUsedQ: Optional[cute.Tensor], mSeqUsedK: Optional[cute.Tensor], mPageTable: Optional[cute.Tensor], tma_atom_Q: Optional[cute.CopyAtom], tma_atom_K: Optional[cute.CopyAtom], tma_atom_V: Optional[cute.CopyAtom], tma_atom_KC: Optional[cute.CopyAtom], tma_atom_VC: Optional[cute.CopyAtom], tma_atom_K2: Optional[cute.CopyAtom], tma_atom_V2: Optional[cute.CopyAtom], tma_atom_O: Optional[cute.CopyAtom], softmax_scale_log2: Float32, softmax_scale: Optional[Float32], sink_range: Int32, window_size_left: Optional[Int32], window_size_right: Optional[Int32], learnable_sink: Optional[cute.Tensor], blocksparse_tensors: Optional[BlockSparseTensors], sQ_layout: cute.ComposedLayout, sK_layout: cute.ComposedLayout, sV_layout: cute.ComposedLayout, sO_layout: cute.ComposedLayout, sP_layout: cute.ComposedLayout | None, gmem_tiled_copy_Q: cute.TiledCopy, gmem_tiled_copy_K: cute.TiledCopy, gmem_tiled_copy_V: cute.TiledCopy, gmem_tiled_copy_O: cute.TiledCopy, tiled_mma_qk: cute.TiledMma, tiled_mma_pv: cute.TiledMma, tile_sched_params: ParamsBase, TileScheduler: cutlass.Constexpr[Callable], SharedStorage: cutlass.Constexpr[Callable], fastdiv_mods=None, ): warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) # Prefetch tma descriptor if warp_idx == 0: for tma_atom in ( tma_atom_Q, tma_atom_K, tma_atom_V, tma_atom_KC, tma_atom_VC, tma_atom_K2, tma_atom_V2, tma_atom_O, ): if const_expr(tma_atom is not None): cpasync.prefetch_descriptor(tma_atom) smem = cutlass.utils.SmemAllocator() storage = smem.allocate(SharedStorage) # Mbarrier / pipeline init mbar_ptr_Q = storage.mbar_ptr_Q.data_ptr() ThreadCooperativeGroup = partial(pipeline.CooperativeGroup, pipeline.Agent.Thread) tma_warp = ThreadCooperativeGroup(1) load_threads = ThreadCooperativeGroup(self.num_threads_per_warp_group) mma_warps = ThreadCooperativeGroup(self.num_mma_threads // cute.arch.WARP_SIZE) if const_expr(self.use_tma_Q): pipeline_q = pipeline_custom.PipelineTmaAsync.create( barrier_storage=mbar_ptr_Q, num_stages=1, producer_group=tma_warp, consumer_group=mma_warps, tx_count=self.tma_copy_bytes["Q"], defer_sync=True, ) else: pipeline_q = pipeline_custom.PipelineCpAsync.create( barrier_storage=mbar_ptr_Q, num_stages=1, producer_group=load_threads, consumer_group=mma_warps, defer_sync=True, elect_one_release=True, syncwarp_before_release=False, ) if const_expr(self.use_tma_KV): pipeline_k = pipeline_custom.PipelineTmaAsync.create( barrier_storage=storage.mbar_ptr_K.data_ptr(), num_stages=self.num_stages, producer_group=tma_warp, consumer_group=mma_warps, tx_count=self.tma_copy_bytes["K"], defer_sync=True, ) pipeline_v = pipeline_custom.PipelineTmaAsync.create( barrier_storage=storage.mbar_ptr_V.data_ptr(), num_stages=self.num_stages, producer_group=tma_warp, consumer_group=mma_warps, tx_count=self.tma_copy_bytes["V"], defer_sync=True, ) else: pipeline_k = pipeline_custom.PipelineCpAsync.create( barrier_storage=storage.mbar_ptr_K.data_ptr(), num_stages=self.num_stages, producer_group=load_threads, consumer_group=mma_warps, defer_sync=True, elect_one_release=True, syncwarp_before_release=False, ) pipeline_v = pipeline_custom.PipelineCpAsync.create( barrier_storage=storage.mbar_ptr_V.data_ptr(), num_stages=self.num_stages, producer_group=load_threads, consumer_group=mma_warps, defer_sync=True, elect_one_release=True, syncwarp_before_release=False, ) # Cluster arrive after barrier init pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) # /////////////////////////////////////////////////////////////////////////////// # Get shared memory buffer # /////////////////////////////////////////////////////////////////////////////// sQ = storage.sQ.get_tensor(sQ_layout.outer, swizzle=sQ_layout.inner) sK = storage.sK.get_tensor(sK_layout.outer, swizzle=sK_layout.inner) if const_expr(not self.Q_in_regs): sV = storage.sV.get_tensor(sV_layout.outer, swizzle=sV_layout.inner) else: sV = storage.sQ.get_tensor( sV_layout.outer, swizzle=sV_layout.inner, dtype=mV.element_type ) # Transpose view of V to tensor with layout (head_dim_v, tile_n) for tiled mma sVt = layout_utils.transpose_view(sV) sP = None if const_expr(sP_layout is not None): sP = storage.sP.get_tensor(sP_layout.outer, swizzle=sP_layout.inner) # reuse sQ's data iterator sO = storage.sQ.get_tensor(sO_layout.outer, swizzle=sO_layout.inner, dtype=self.dtype) route_mask = storage.route_mask.get_tensor( cute.make_layout((4,)) ) route_sums = storage.route_sums.get_tensor(cute.make_layout((4, self.tile_n))) block_info = BlockInfo( self.tile_m, self.tile_n, self.is_causal, self.is_local, False, # is_split_kv window_size_left, window_size_right, qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, ) SeqlenInfoCls = partial( SeqlenInfoQK.create, seqlen_q_static=mQ.shape[0] if const_expr(not self.pack_gqa) else mQ.shape[0][1], seqlen_k_static=mK.shape[0] if const_expr(mPageTable is None) else mK.shape[0] * mPageTable.shape[1], mCuSeqlensQ=mCuSeqlensQ, mCuSeqlensK=mCuSeqlensK, mSeqUsedQ=mSeqUsedQ, mSeqUsedK=mSeqUsedK, # Don't need to pass in tile_mn because we won't access offset_padded ) AttentionMaskCls = partial( AttentionMask, self.tile_m, self.tile_n, window_size_left=window_size_left, window_size_right=window_size_right, qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, ) TileSchedulerCls = partial(TileScheduler.create, tile_sched_params) # Cluster wait before starting pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) cute.arch.setmaxregister_increase(self.num_mma_regs) self.mma_one_warpgroup_sol_attn_route_tma( tiled_mma_qk, tiled_mma_pv, mQ, mK, mV, mKC, mVC, mO, mLSE, sQ, sK, sV, sVt, sP, sO, tma_atom_Q, tma_atom_K, tma_atom_V, tma_atom_KC, tma_atom_VC, gmem_tiled_copy_O, tma_atom_O, pipeline_q, pipeline_k, pipeline_v, SeqlenInfoCls, AttentionMaskCls, TileSchedulerCls, mGlobalThresh, route_mask, route_sums, softmax_scale_log2, softmax_scale, sink_range, block_info, ) @cute.jit def epilogue_one_warpgroup_tma_o( self, acc_O: cute.Tensor, lse: cute.Tensor, mO: cute.Tensor, mLSE: Optional[cute.Tensor], sO: cute.Tensor, seqlen: SeqlenInfoQK, tma_atom_O: cute.CopyAtom, tiled_mma: cute.TiledMma, tidx: Int32, m_block: Int32, head_idx: Int32, batch_idx: Int32, ): """One-warpgroup TMA-O epilogue with an in-CTA store owner.""" rO = cute.make_fragment_like(acc_O, self.dtype) rO.store(acc_O.load().to(self.dtype)) cute.arch.barrier( barrier_id=int(NamedBarrierFwd.Epilogue), number_of_threads=self.num_epilogue_threads, ) smem_copy_atom_O = utils.get_smem_store_atom( self.arch.major * 10 + self.arch.minor, self.dtype ) smem_thr_copy_O = cute.make_tiled_copy_C(smem_copy_atom_O, tiled_mma).get_slice(tidx) taccOrO = smem_thr_copy_O.retile(rO) taccOsO = smem_thr_copy_O.partition_D(sO) cute.copy(smem_copy_atom_O, taccOrO, taccOsO) cO = cute.make_identity_tensor((self.tile_m, self.tile_hdimv)) if const_expr(mLSE is not None): mLSE_cur = mLSE[None, head_idx, batch_idx] gLSE = cute.local_tile(mLSE_cur, (self.tile_m,), (m_block,)) gLSE_expanded_layout = cute.append( gLSE.layout, cute.make_layout((self.tile_hdimv,), stride=(0,)) ) gLSE_expanded = cute.make_tensor(gLSE.iterator, gLSE_expanded_layout) thr_mma = tiled_mma.get_slice(tidx) taccOgLSE = layout_utils.reshape_acc_to_mn(thr_mma.partition_C(gLSE_expanded)) taccOcO = layout_utils.reshape_acc_to_mn(thr_mma.partition_C(cO)) t0accOcO = layout_utils.reshape_acc_to_mn(thr_mma.get_slice(0).partition_C(cO)) if taccOcO[0][1] == 0: for m in cutlass.range_constexpr(cute.size(taccOgLSE.shape[1])): if ( t0accOcO[m, 0][0] < seqlen.seqlen_q - m_block * self.tile_m - taccOcO[0][0] ): taccOgLSE[m, 0] = lse[m] mO_cur = mO[None, None, head_idx, batch_idx] cute.arch.fence_view_async_shared() cute.arch.barrier( barrier_id=int(NamedBarrierFwd.Epilogue), number_of_threads=self.num_epilogue_threads, ) gO = cute.local_tile(mO_cur, (self.tile_m, self.tile_hdimv), (m_block, 0)) store_O, _, _ = copy_utils.tma_get_copy_fn( tma_atom_O, 0, cute.make_layout(1), sO, gO, single_stage=True ) warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) if warp_idx == Int32(0): store_O() cute.arch.cp_async_bulk_commit_group() cute.arch.cp_async_bulk_wait_group(0, read=True) @cute.jit def epilogue_one_warpgroup_split_partial( self, acc_O: cute.Tensor, lse: cute.Tensor, mO: cute.Tensor, mLSE: cute.Tensor, seqlen: SeqlenInfoQK, tiled_mma: cute.TiledMma, tidx: Int32, m_block: Int32, partial_head_idx: Int32, batch_idx: Int32, ): """Write one normalized FP32 split partial and its natural-log LSE. ``mO`` and ``mLSE`` use a physical split-head dimension. The caller maps ``(split, head)`` to ``partial_head_idx``; a later combine kernel performs the log-sum-exp weighted reduction across that dimension. """ mO_cur = mO[None, None, partial_head_idx, batch_idx] gO = cute.local_tile( mO_cur, (self.tile_m, self.tile_hdimv), (m_block, 0) ) copy_atom = cute.make_copy_atom( cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=32, ) tiled_copy = cute.make_tiled_copy_C(copy_atom, tiled_mma) rO = cute.make_rmem_tensor_like(acc_O, Float32) rO.store(acc_O.load()) tOrO = tiled_copy.retile(rO) tOgO = tiled_copy.get_slice(tidx).partition_D(gO) cute.autovec_copy(tOrO, tOgO) mLSE_cur = mLSE[None, partial_head_idx, batch_idx] gLSE = cute.local_tile(mLSE_cur, (self.tile_m,), (m_block,)) gLSE_expanded_layout = cute.append( gLSE.layout, cute.make_layout((self.tile_hdimv,), stride=(0,)) ) gLSE_expanded = cute.make_tensor( gLSE.iterator, gLSE_expanded_layout ) thr_mma = tiled_mma.get_slice(tidx) taccOgLSE = layout_utils.reshape_acc_to_mn( thr_mma.partition_C(gLSE_expanded) ) cO = cute.make_identity_tensor((self.tile_m, self.tile_hdimv)) taccOcO = layout_utils.reshape_acc_to_mn(thr_mma.partition_C(cO)) t0accOcO = layout_utils.reshape_acc_to_mn( thr_mma.get_slice(0).partition_C(cO) ) if taccOcO[0][1] == 0: for m in cutlass.range_constexpr(cute.size(taccOgLSE.shape[1])): if ( t0accOcO[m, 0][0] < seqlen.seqlen_q - m_block * self.tile_m - taccOcO[0][0] ): taccOgLSE[m, 0] = lse[m] @cute.jit def mma_one_warpgroup_sol_attn_route_tma( self, tiled_mma_qk: cute.TiledMma, tiled_mma_pv: cute.TiledMma, mQ: cute.Tensor, mK: cute.Tensor, mV: cute.Tensor, mKC: cute.Tensor, mVC: cute.Tensor, mO: cute.Tensor, mLSE: Optional[cute.Tensor], sQ: cute.Tensor, sK: cute.Tensor, sV: cute.Tensor, sVt: cute.Tensor, sP: Optional[cute.Tensor], sO: cute.Tensor, tma_atom_Q: Optional[cute.CopyAtom], tma_atom_K: Optional[cute.CopyAtom], tma_atom_V: Optional[cute.CopyAtom], tma_atom_KC: Optional[cute.CopyAtom], tma_atom_VC: Optional[cute.CopyAtom], gmem_tiled_copy_O: cute.TiledCopy, tma_atom_O: Optional[cute.CopyAtom], pipeline_q: pipeline.PipelineAsync, pipeline_k: pipeline.PipelineAsync, pipeline_v: pipeline.PipelineAsync, SeqlenInfoCls: Callable, AttentionMaskCls: Callable, TileSchedulerCls: cutlass.Constexpr[Callable], mGlobalThresh: cute.Tensor, route_mask: cute.Tensor, route_sums: cute.Tensor, softmax_scale_log2: Float32, softmax_scale: Float32, sink_range: Int32, block_info: BlockInfo, ): """Run the fused route, approximate, and exact attention mainloop.""" tidx, _, _ = cute.arch.thread_idx() warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) if const_expr(not (self.use_tma_Q and self.use_tma_KV)): if tidx == Int32(0) and warp_idx == Int32(0): cute.printf("SOL_ATTN one-warpgroup path requires TMA Q/KV\n") else: q_producer_phase = Int32(1) q_consumer_phase = Int32(0) kv_producer_state = pipeline.make_pipeline_state( pipeline.PipelineUserType.Producer, self.num_stages ) kv_consumer_state = pipeline.make_pipeline_state( pipeline.PipelineUserType.Consumer, self.num_stages ) tile_scheduler = TileSchedulerCls() work_tile = tile_scheduler.initial_work_tile_info() if work_tile.is_valid_tile: m_block, head_idx, batch_idx, split_idx = work_tile.tile_idx partial_head_idx = ( head_idx + split_idx * mQ.shape[2] if const_expr(self.sol_attn_num_splits > 1) else head_idx ) seqlen = SeqlenInfoCls(batch_idx) head_idx_kv = ( head_idx // self.qhead_per_kvhead if const_expr(not self.pack_gqa) else head_idx ) mQ_cur = seqlen.offset_batch_Q(mQ, batch_idx, dim=3)[None, None, head_idx] mK_cur = seqlen.offset_batch_K(mK, batch_idx, dim=3)[ None, None, head_idx_kv ] mV_cur = seqlen.offset_batch_K(mV, batch_idx, dim=3)[ None, None, head_idx_kv ] mKC_cur = mKC[None, None, head_idx_kv, batch_idx] mVC_cur = mVC[None, None, head_idx_kv, batch_idx] gQ = cute.local_tile(mQ_cur, (self.tile_m, self.tile_hdim), (m_block, 0)) gK = cute.local_tile(mK_cur, (self.tile_n, self.tile_hdim), (None, 0)) gV = cute.local_tile(mV_cur, (self.tile_n, self.tile_hdimv), (None, 0)) gKC = cute.local_tile(mKC_cur, (self.tile_n, self.tile_hdim), (None, 0)) gVC = cute.local_tile(mVC_cur, (self.tile_n, self.tile_hdimv), (None, 0)) load_Q, _, _ = copy_utils.tma_get_copy_fn( tma_atom_Q, 0, cute.make_layout(1), gQ, sQ, single_stage=True ) tma_load_K_fn, _, _ = copy_utils.tma_get_copy_fn( tma_atom_K, 0, cute.make_layout(1), gK, sK ) tma_load_K_fn = copy_utils.tma_producer_copy_fn(tma_load_K_fn, pipeline_k) tma_load_V_fn, _, _ = copy_utils.tma_get_copy_fn( tma_atom_V, 0, cute.make_layout(1), gV, sV ) tma_load_V_fn = copy_utils.tma_producer_copy_fn(tma_load_V_fn, pipeline_v) tma_load_KC_fn, _, _ = copy_utils.tma_get_copy_fn( tma_atom_KC, 0, cute.make_layout(1), gKC, sK ) tma_load_KC_fn = copy_utils.tma_producer_copy_fn( tma_load_KC_fn, pipeline_k ) tma_load_VC_fn, _, _ = copy_utils.tma_get_copy_fn( tma_atom_VC, 0, cute.make_layout(1), gVC, sV ) tma_load_VC_fn = copy_utils.tma_producer_copy_fn( tma_load_VC_fn, pipeline_v ) if warp_idx == Int32(0): pipeline_q.producer_acquire_w_index_phase(0, q_producer_phase) load_Q(tma_bar_ptr=pipeline_q.sync_object_full.get_barrier(0)) pipeline_q.consumer_wait_w_index_phase(0, q_consumer_phase) warp_group_thread_layout = cute.make_layout( 1, stride=self.num_threads_per_warp_group ) thr_mma_qk = tiled_mma_qk.get_slice(tidx) wg_mma_qk = tiled_mma_qk.get_slice(warp_group_thread_layout(Int32(0))) wg_mma_pv = tiled_mma_pv.get_slice(warp_group_thread_layout(Int32(0))) _, tSrQ, tSrK = sm90_utils.partition_fragment_ABC( wg_mma_qk, (self.tile_m, self.tile_n, self.tile_hdim), sQ, sK ) mma_qk_fn = partial( self.sol_attn_qk_gemm_zero_init, tiled_mma_qk, (self.tile_m, self.tile_n), tSrQ, tSrK, ) acc_O, tOrP, tOrVt = sm90_utils.partition_fragment_ABC( wg_mma_pv, (self.tile_m, self.tile_hdimv, self.tile_n), sP, sVt ) mma_pv_fn = partial(sm90_utils.gemm_w_idx, tiled_mma_pv, acc_O, tOrP, tOrVt) smem_copy_atom_P = utils.get_smem_store_atom( self.arch.major * 10 + self.arch.minor, self.dtype ) smem_thr_copy_P = cute.make_tiled_copy_C( smem_copy_atom_P, tiled_mma_qk ).get_slice(tidx) tPsP = smem_thr_copy_P.partition_D(sP) if const_expr(sP is not None) else None smem_copy_params = SimpleNamespace( smem_thr_copy_P=smem_thr_copy_P, tPsP=tPsP, ) acc_O.fill(0.0) cS_route = cute.make_identity_tensor((self.tile_m, self.tile_n)) tScS_route_mn = layout_utils.reshape_acc_to_mn( thr_mma_qk.partition_C(cS_route) ) mask = AttentionMaskCls(seqlen) mask_fn = partial( mask.apply_mask, batch_idx=batch_idx, head_idx=head_idx, m_block=m_block, thr_mma=thr_mma_qk, mask_causal=self.is_causal, mask_local=self.is_local, aux_tensors=None, fastdiv_mods=None, ) score_mod_fn = None if const_expr(self.score_mod is not None): score_mod_fn = partial( self.apply_score_mod, thr_mma_qk, batch_idx, head_idx, m_block, softmax_scale=softmax_scale, aux_tensors=None, fastdiv_mods=None, ) softmax = Softmax.create( softmax_scale_log2, num_rows=acc_O.shape[0][0] * acc_O.shape[1], softmax_scale=softmax_scale, ) if const_expr(self.sol_attn_neutral_softmax_state): softmax.row_max.fill(-Float32.inf) softmax.row_sum.fill(0.0) exact_mma_one_n_block = partial( self.mma_one_n_block, mma_qk_fn=mma_qk_fn, pipeline_k=pipeline_k, pipeline_v=pipeline_v, acc_O=acc_O, tOrP=tOrP, smem_copy_params=smem_copy_params, softmax=softmax, score_mod_fn=score_mod_fn, score_scale_fn=None, check_inf=not self.sol_attn_assume_nonempty_rows, ) n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block) route_block_count = n_block_max - n_block_min if const_expr(self.sol_attn_static_num_full_route_groups >= 0): num_full_route_groups = Int32(self.sol_attn_static_num_full_route_groups) if const_expr(self.sol_attn_static_tail_valid_count > 0): tail_valid_count = Int32(self.sol_attn_static_tail_valid_count) else: tail_valid_count = Int32(0) elif const_expr(self.sol_attn_assume_full_route_groups): num_full_route_groups = cute.ceil_div( route_block_count, self.sol_attn_group_size ) tail_valid_count = Int32(0) else: num_full_route_groups = route_block_count // Int32(self.sol_attn_group_size) tail_valid_count = ( route_block_count - num_full_route_groups * Int32(self.sol_attn_group_size) ) num_route_groups = num_full_route_groups if tail_valid_count > Int32(0): num_route_groups += Int32(1) if const_expr(self.sol_attn_num_splits == 1): split_group_begin = Int32(0) split_num_route_groups = num_route_groups else: groups_per_split = ( num_route_groups + self.sol_attn_num_splits - 1 ) // self.sol_attn_num_splits split_group_begin = split_idx * groups_per_split split_group_end = cutlass.min( split_group_begin + groups_per_split, num_route_groups ) split_num_route_groups = cutlass.max( split_group_end - split_group_begin, Int32(0) ) O_should_accumulate = self.sol_attn_neutral_softmax_state for local_group_iter in cutlass.range( split_num_route_groups, unroll=1 ): group_iter = split_group_begin + local_group_iter group_start = n_block_min + group_iter * Int32(self.sol_attn_group_size) route_valid_count = Int32(self.sol_attn_group_size) if const_expr(not self.sol_attn_assume_full_route_groups): if group_iter == num_full_route_groups and tail_valid_count > Int32(0): route_valid_count = tail_valid_count route_col_offset = group_start - ( group_start // Int32(self.tile_n) ) * Int32(self.tile_n) route_n_block = group_start - route_col_offset route_tile = route_n_block // Int32(self.tile_n) has_next_route_group = ( local_group_iter + Int32(1) < split_num_route_groups ) next_route_tile = Int32(-1) if has_next_route_group: next_group_start = group_start + Int32(self.sol_attn_group_size) next_route_tile = next_group_start // Int32(self.tile_n) if warp_idx == Int32(0): if local_group_iter == Int32(0): pipeline_k.producer_acquire(kv_producer_state) tma_load_KC_fn( src_idx=route_tile, producer_state=kv_producer_state, ) else: previous_group_had_exact = ( (route_mask[0] != Int32(0)) or (route_mask[1] != Int32(0)) or (route_mask[2] != Int32(0)) or (route_mask[3] != Int32(0)) ) if not previous_group_had_exact: pipeline_k.producer_acquire(kv_producer_state) tma_load_KC_fn( src_idx=route_tile, producer_state=kv_producer_state, ) pipeline_v.producer_acquire(kv_producer_state) tma_load_VC_fn( src_idx=route_tile, producer_state=kv_producer_state, ) kv_producer_state.advance() pipeline_k.consumer_wait( kv_consumer_state, pipeline_k.consumer_try_wait(kv_consumer_state), ) acc_S = mma_qk_fn(B_idx=kv_consumer_state.index, wg_wait=-1) warpgroup.wait_group(0) pipeline_k.consumer_release(kv_consumer_state) mask0, mask1, mask2, mask3 = self.sol_attn_build_route_mask_from_acc( acc_S, route_sums, tScS_route_mn, m_block, group_start, route_valid_count, route_col_offset, seqlen, batch_idx, head_idx, mGlobalThresh, softmax_scale_log2, sink_range, False, route_mask_words_override=2, ) exact_mask0 = mask0 exact_mask1 = mask1 exact_mask2 = mask2 exact_mask3 = mask3 first_exact_n_block = group_start first_exact_exists = False if tidx == Int32(0): route_mask[0] = mask0 route_mask[1] = mask1 route_mask[2] = mask2 route_mask[3] = mask3 cute.arch.barrier( barrier_id=SOL_ATTN_ROUTE_MASK_BARRIER_ID, number_of_threads=self.num_mma_threads, ) mask0 = route_mask[0] mask1 = route_mask[1] mask2 = route_mask[2] mask3 = route_mask[3] exact_mask0 = mask0 exact_mask1 = mask1 exact_mask2 = mask2 exact_mask3 = mask3 first_exact_exists = ( (mask0 != Int32(0)) or (mask1 != Int32(0)) or (mask2 != Int32(0)) or (mask3 != Int32(0)) ) if mask0 != Int32(0): first_lowbit = mask0 & (Int32(0) - mask0) first_exact_n_block += sol_attn_selector.sol_attn_bfind_b32( first_lowbit ) exact_mask0 = mask0 & (mask0 - Int32(1)) elif mask1 != Int32(0): first_lowbit = mask1 & (Int32(0) - mask1) first_exact_n_block += Int32(32) + ( sol_attn_selector.sol_attn_bfind_b32(first_lowbit) ) exact_mask1 = mask1 & (mask1 - Int32(1)) elif mask2 != Int32(0): first_lowbit = mask2 & (Int32(0) - mask2) first_exact_n_block += Int32(64) + ( sol_attn_selector.sol_attn_bfind_b32(first_lowbit) ) exact_mask2 = mask2 & (mask2 - Int32(1)) elif mask3 != Int32(0): first_lowbit = mask3 & (Int32(0) - mask3) first_exact_n_block += Int32(96) + ( sol_attn_selector.sol_attn_bfind_b32(first_lowbit) ) exact_mask3 = mask3 & (mask3 - Int32(1)) if first_exact_exists and warp_idx == Int32(0): pipeline_k.producer_acquire(kv_producer_state) tma_load_K_fn( src_idx=first_exact_n_block, producer_state=kv_producer_state, ) if const_expr(self.score_mod is not None): score_mod_fn(acc_S, n_block=route_n_block, seqlen=seqlen) mask_fn( acc_S, n_block=route_n_block, mask_mod=self.mask_mod, mask_seqlen=not self.sol_attn_full_route_mask_seqlen_false, ) if const_expr(self.sol_attn_assume_full_route_groups): route_has_approx = (mask0 != Int32(-1)) or (mask1 != Int32(-1)) else: valid0 = route_valid_count if valid0 > Int32(32): valid0 = Int32(32) valid_bits0 = Int32(0) if valid0 > Int32(0): valid_bits0 = Int32(-1) if valid0 < Int32(32): valid_bits0 = (Int32(1) << valid0) - Int32(1) valid1 = route_valid_count - Int32(32) if valid1 < Int32(0): valid1 = Int32(0) if valid1 > Int32(32): valid1 = Int32(32) valid_bits1 = Int32(0) if valid1 > Int32(0): valid_bits1 = Int32(-1) if valid1 < Int32(32): valid_bits1 = (Int32(1) << valid1) - Int32(1) route_has_approx = ( ((mask0 & valid_bits0) != valid_bits0) or ((mask1 & valid_bits1) != valid_bits1) ) self.sol_attn_mask_route_approx_columns( acc_S, route_sums, tScS_route_mn, route_valid_count, route_col_offset, mask0, mask1, mask2, mask3, False, route_mask_words_override=self.sol_attn_group_words, ) pipeline_v.consumer_wait( kv_consumer_state, pipeline_v.consumer_try_wait(kv_consumer_state), ) if route_has_approx: row_sum_prev = None if const_expr( self.sol_attn_fast_route_lens and not self.sol_attn_full_block_row_sum_prescale ): row_sum_prev = cute.make_fragment_like(softmax.row_sum, Float32) row_sum_prev.store(softmax.row_sum.load()) if O_should_accumulate: if const_expr(self.sol_attn_full_block_row_sum_prescale): for r in cutlass.range( cute.size(softmax.row_sum), unroll_full=True ): softmax.row_sum[r] *= Float32(1.0 / self.tile_n) row_scale = softmax.online_softmax( acc_S, is_first=False, check_inf=not self.sol_attn_assume_nonempty_rows, ) softmax.rescale_O(acc_O, row_scale) if const_expr(self.sol_attn_full_block_row_sum_prescale): for r in cutlass.range( cute.size(softmax.row_sum), unroll_full=True ): softmax.row_sum[r] *= Float32(self.tile_n) elif const_expr(self.sol_attn_fast_route_lens): self.sol_attn_apply_route_current_lens_to_row_sum_fast( acc_S, tScS_route_mn, row_sum_prev, row_scale, group_start, route_valid_count, route_col_offset, seqlen, softmax, False, ) else: self.sol_attn_apply_route_current_lens_to_row_sum( acc_S, tScS_route_mn, group_start, route_valid_count, route_col_offset, seqlen, softmax, ) else: row_scale = softmax.online_softmax( acc_S, is_first=True, check_inf=not self.sol_attn_assume_nonempty_rows, ) if const_expr(self.sol_attn_full_block_row_sum_prescale): for r in cutlass.range( cute.size(softmax.row_sum), unroll_full=True ): softmax.row_sum[r] *= Float32(self.tile_n) elif const_expr(self.sol_attn_fast_route_lens): self.sol_attn_apply_route_current_lens_to_row_sum_fast( acc_S, tScS_route_mn, row_sum_prev, row_scale, group_start, route_valid_count, route_col_offset, seqlen, softmax, True, ) else: self.sol_attn_apply_route_current_lens_to_row_sum( acc_S, tScS_route_mn, group_start, route_valid_count, route_col_offset, seqlen, softmax, ) tOrP_acc = layout_utils.reshape_acc_to_frgA(acc_S) tOrP_cur = ( tOrP if const_expr(self.mma_pv_is_rs) else cute.make_rmem_tensor_like(tOrP_acc, self.dtype) ) utils.cvt_f16(tOrP_acc, tOrP_cur) if const_expr(not self.mma_pv_is_rs): tPrP = smem_copy_params.smem_thr_copy_P.retile(tOrP_cur) cute.copy( smem_copy_params.smem_thr_copy_P, tPrP, smem_copy_params.tPsP, ) cute.arch.fence_view_async_shared() cute.arch.sync_warp() if O_should_accumulate: sm90_utils.gemm_w_idx( tiled_mma_pv, acc_O, tOrP_cur, tOrVt, zero_init=False, B_idx=kv_consumer_state.index, wg_wait=-1, ) else: sm90_utils.gemm_w_idx( tiled_mma_pv, acc_O, tOrP_cur, tOrVt, zero_init=True, B_idx=kv_consumer_state.index, wg_wait=-1, ) warpgroup.wait_group(0) O_should_accumulate = True pipeline_v.consumer_release(kv_consumer_state) kv_consumer_state.advance() last_n_block = Int32(-1) if const_expr( (not self.sol_attn_assume_full_k_exact_blocks) or self.sol_attn_exact_mask_seqlen_last_only ): last_n_block = ( (seqlen.seqlen_k + Int32(self.tile_n - 1)) // Int32(self.tile_n) ) - Int32(1) if O_should_accumulate: ( kv_producer_state, kv_consumer_state, O_should_accumulate, _, ) = exact_stream.consume_exact_blocks( exact_mask0, exact_mask1, exact_mask2, exact_mask3, group_start, seqlen, kv_producer_state, kv_consumer_state, tma_load_K_fn, tma_load_V_fn, pipeline_k, pipeline_v, warp_idx == Int32(0), mma_pv_fn, exact_mma_one_n_block, mask_fn, score_mod_fn, O_should_accumulate, self.warp_scheduler_barrier_sync, self.warp_scheduler_barrier_arrive, not self.sol_attn_assume_full_k_exact_blocks, False, self.sol_attn_group_words, last_n_block, self.sol_attn_exact_mask_seqlen_last_only, first_exact_n_block, first_exact_exists, next_route_tile, tma_load_KC_fn, ) else: ( kv_producer_state, kv_consumer_state, O_should_accumulate, _, ) = exact_stream.consume_exact_blocks( exact_mask0, exact_mask1, exact_mask2, exact_mask3, group_start, seqlen, kv_producer_state, kv_consumer_state, tma_load_K_fn, tma_load_V_fn, pipeline_k, pipeline_v, warp_idx == Int32(0), mma_pv_fn, exact_mma_one_n_block, mask_fn, score_mod_fn, O_should_accumulate, self.warp_scheduler_barrier_sync, self.warp_scheduler_barrier_arrive, not self.sol_attn_assume_full_k_exact_blocks, True, self.sol_attn_group_words, last_n_block, self.sol_attn_exact_mask_seqlen_last_only, first_exact_n_block, first_exact_exists, next_route_tile, tma_load_KC_fn, ) pipeline_q.consumer_release_w_index(0) final_scale = softmax.finalize(sink_val=None) softmax.rescale_O(acc_O, final_scale) if const_expr(self.use_tma_O): self.epilogue_one_warpgroup_tma_o( acc_O, softmax.row_sum, mO, mLSE, sO, seqlen, tma_atom_O, tiled_mma_pv, tidx, m_block, partial_head_idx, batch_idx, ) else: self.epilogue_one_warpgroup_split_partial( acc_O, softmax.row_sum, mO, mLSE, seqlen, tiled_mma_pv, tidx, m_block, partial_head_idx, batch_idx, ) @cute.jit def mma_one_n_block( self, smem_pipe_read: pipeline.PipelineState | pipeline_custom.PipelineStateSimple, n_block: Int32, mma_qk_fn: Callable, mma_pv_fn: Callable, pipeline_k: pipeline.PipelineAsync, pipeline_v: pipeline.PipelineAsync, acc_O: cute.Tensor, tOrP: cute.Tensor, smem_copy_params: SimpleNamespace, softmax: Softmax, seqlen: SeqlenInfoQK, scores_scale: Optional[cute.Tensor] = None, score_mod_fn: Optional[Callable] = None, score_scale_fn: Optional[Callable] = None, mask_fn: Optional[Callable] = None, last_block_mask_fn: Optional[Callable] = None, last_n_block: Int32 = Int32(-1), is_first_n_block: cutlass.Constexpr = False, check_inf: cutlass.Constexpr = True, prefetch_next: cutlass.Constexpr = False, next_n_block: Int32 = Int32(-1), kv_producer_state=None, load_K: Optional[Callable] = None, issue_load=False, ): pipeline_k.consumer_wait(smem_pipe_read, pipeline_k.consumer_try_wait(smem_pipe_read)) acc_S = mma_qk_fn(B_idx=smem_pipe_read.index, wg_wait=-1) self.warp_scheduler_barrier_arrive() warpgroup.wait_group(0) pipeline_k.consumer_release(smem_pipe_read) # Reuse the released K stage while the current softmax and P@V run. if const_expr(prefetch_next): if issue_load and next_n_block >= Int32(0): pipeline_k.producer_acquire(kv_producer_state) load_K(src_idx=next_n_block, producer_state=kv_producer_state) if const_expr(score_scale_fn is not None): score_scale_fn(acc_S, n_block=n_block) if const_expr(score_mod_fn is not None): score_mod_fn(acc_S, n_block=n_block, seqlen=seqlen) if const_expr(mask_fn is not None): mask_fn(acc_S=acc_S, n_block=n_block) if const_expr(last_block_mask_fn is not None): if n_block == last_n_block: last_block_mask_fn(acc_S=acc_S, n_block=n_block) row_scale = softmax.online_softmax(acc_S, is_first=is_first_n_block, check_inf=check_inf) tOrP_acc = layout_utils.reshape_acc_to_frgA(acc_S) tOrP_cur = ( tOrP if const_expr(self.mma_pv_is_rs) else cute.make_rmem_tensor_like(tOrP_acc, self.dtype) ) utils.cvt_f16(tOrP_acc, tOrP_cur) if const_expr(not self.mma_pv_is_rs): tPrP = smem_copy_params.smem_thr_copy_P.retile(tOrP_cur) cute.copy(smem_copy_params.smem_thr_copy_P, tPrP, smem_copy_params.tPsP) softmax.rescale_O(acc_O, row_scale) if const_expr(not self.mma_pv_is_rs): cute.arch.fence_view_async_shared() cute.arch.sync_warp() pipeline_v.consumer_wait(smem_pipe_read, pipeline_v.consumer_try_wait(smem_pipe_read)) self.warp_scheduler_barrier_sync() mma_pv_fn(B_idx=smem_pipe_read.index, wg_wait=0) pipeline_v.consumer_release(smem_pipe_read) smem_pipe_read.advance() return smem_pipe_read @cute.jit def mma_init(self): warp_group_idx = utils.canonical_warp_group_idx(sync=False) if const_expr(self.use_scheduler_barrier): if warp_group_idx == 1: cute.arch.barrier_arrive( barrier_id=int(NamedBarrierFwd.WarpSchedulerWG1), number_of_threads=2 * self.num_threads_per_warp_group, ) @cute.jit def apply_score_mod( self, thr_mma_qk, batch_idx, head_idx, m_block, acc_S, n_block, softmax_scale, seqlen, aux_tensors: Optional[list] = None, fastdiv_mods=None, ): # Prepare index tensor cS = cute.make_identity_tensor((self.tile_m, self.tile_n)) cS = cute.domain_offset((m_block * self.tile_m, n_block * self.tile_n), cS) tScS = thr_mma_qk.partition_C(cS) apply_score_mod_inner( acc_S, tScS, self.score_mod, batch_idx, head_idx, softmax_scale, self.vec_size, self.qk_acc_dtype, aux_tensors, fastdiv_mods, seqlen_info=seqlen, constant_q_idx=None, qhead_per_kvhead=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, ) def warp_scheduler_barrier_sync(self): if const_expr(self.use_scheduler_barrier): cute.arch.barrier( barrier_id=int(NamedBarrierFwd.WarpSchedulerWG1) - 1 + utils.canonical_warp_group_idx(sync=False), number_of_threads=2 * self.num_threads_per_warp_group, ) def warp_scheduler_barrier_arrive(self): if const_expr(self.use_scheduler_barrier): assert self.num_wg_mma in [2, 3] cur_wg = utils.canonical_warp_group_idx(sync=False) - 1 if const_expr(self.num_wg_mma == 2): next_wg = 1 - cur_wg else: t = cur_wg + 1 next_wg = t % self.num_wg_mma cute.arch.barrier_arrive( barrier_id=int(NamedBarrierFwd.WarpSchedulerWG1) + next_wg, number_of_threads=2 * self.num_threads_per_warp_group, )