| |
| |
|
|
| 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 |
| ] |
| |
| 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: |
| |
| |
| 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 |
| 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") |
| |
| |
| |
| self.use_tma_O = ( |
| self.sol_attn_num_splits == 1 or mO.element_type == self.dtype |
| ) |
| |
| 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() |
| |
| 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), |
| |
| |
| |
| (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) |
|
|
| |
| gmem_tiled_copy_Q = cpasync.CopyBulkTensorTileG2SOp() |
| gmem_tiled_copy_KV = cpasync.CopyBulkTensorTileG2SOp() |
| 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), |
| ) |
| 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, |
| ) |
| 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, |
| ) |
| 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), |
| ) |
| 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()) |
| |
| 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) |
|
|
| |
| 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, |
| ) |
|
|
| |
| pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) |
|
|
| |
| |
| |
| 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 |
| ) |
| |
| 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) |
| |
| 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, |
| 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, |
| |
| ) |
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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, |
| ): |
| |
| 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, |
| ) |
|
|
