From 9fcab317525a9a61500b0cd1d583e399f21fc5b0 Mon Sep 17 00:00:00 2001 From: Caitlin Wang Date: Sat, 10 Oct 2026 20:26:31 -0700 Subject: [PATCH 1/3] [Perf][Attention] Run every sparse-prefill combine on the v1 combine Fold the nv_dev v1 combine (src/blackwell_prefill/combine.py) into src/sm100/fwd/combine.py and make it the public combine() for all callers: the Blackwell port route, the original FP8/BF16 path, the NVFP4 KV path and the Q8KV4 prefill. src/sm100/fwd/combine.py now has two sections, one per kernel: - Original combine: SparseAttentionForwardCombineF32F8 and its launcher combine_f32f8, unchanged apart from the names. It implements the FP32 and FP8 E4M3 STG.128 fake layouts as well as BF16/FP16. - v1 combine: SparseAttentionForwardCombine (moved verbatim) and its launcher combine(), the default. It takes BF16/FP16 partials only: the v1 kernel implements just that STG.128 fake layout. interface.py picks the launcher by partial dtype (_select_combine) on the two paths that accept a caller-chosen partial_dtype, the original FP8/BF16 path and the NVFP4 KV path. The Blackwell port route and Q8KV4 always write BF16 partials and call combine() directly. vLLM always writes BF16 partials, so production runs the v1 kernel. Other changes: - v1 launch: min_blocks_per_mp=0 (was 3 with output_scale + PDL), and a 3-stage O_partial ring on SM107 (2 elsewhere). - The Blackwell port route now accepts output_scale and passes it to the combine instead of falling back to the original path. - Delete src/blackwell_prefill/combine.py; drop the per-caller combine selection in interface.py and prefill_q8kv4/interface.py. Co-Authored-By: Claude Opus 5.5 --- python/fmha_sm100/cute/interface.py | 18 +- .../cute/src/blackwell_prefill/combine.py | 2354 ----------------- .../fmha_sm100/cute/src/sm100/fwd/combine.py | 2339 +++++++++++++++- python/fmha_sm100/prefill_q8kv4/interface.py | 4 - 4 files changed, 2342 insertions(+), 2373 deletions(-) delete mode 100644 python/fmha_sm100/cute/src/blackwell_prefill/combine.py diff --git a/python/fmha_sm100/cute/interface.py b/python/fmha_sm100/cute/interface.py index 35292211..eab24a6a 100644 --- a/python/fmha_sm100/cute/interface.py +++ b/python/fmha_sm100/cute/interface.py @@ -30,7 +30,7 @@ from cutlass import Float32, Int32 from cutlass.cute.runtime import from_dlpack -from src.sm100.fwd.combine import combine +from src.sm100.fwd.combine import combine, combine_f32f8 from src.sm100.fwd.atten_fwd import SparseAttentionForwardSm100 from src.sm100.fwd.atten_fwd_nvfp4_kv import SparseAttentionForwardNvfp4KvSm100 from src.sm100.prepare_scheduler import ( @@ -58,6 +58,13 @@ def _to_cute_int32_metadata(t: torch.Tensor): return to_cute_tensor_kvouter(t, assumed_align=4) +def _select_combine(partial_dtype: torch.dtype): + """The v1 combine takes BF16/FP16 partials; FP32 / FP8 E4M3 need combine_f32f8.""" + if partial_dtype in (torch.bfloat16, torch.float16): + return combine + return combine_f32f8 + + def _normalize_partial_dtype(partial_dtype: torch.dtype) -> torch.dtype: supported = {torch.float32, torch.bfloat16, torch.float16, torch.float8_e4m3fn} if partial_dtype not in supported: @@ -642,7 +649,6 @@ def _try_blackwell_sparse_prefill( if (not causal or schedule is None or page_table is None or partial_dtype != torch.bfloat16 or return_softmax_lse or return_temperature_lse - or output_scale is not None or enable_fp16_softmax or enable_2x_fp8 or qk_dtype != q.dtype or pv_dtype != q.dtype): return None @@ -674,7 +680,6 @@ def _try_blackwell_sparse_prefill( _validate_fwd_schedule(schedule, q=q, k2q_q_indices=k2q_q_indices, head_kv=k.shape[1]) from src.blackwell_prefill.atten_fwd_sm100 import run_pagekv - from src.blackwell_prefill.combine import combine as blackwell_combine partial = torch.empty((topK, *q.shape), dtype=torch.bfloat16, device=q.device) stats = torch.empty((topK, *q.shape[:2]), dtype=torch.float32, device=q.device) @@ -685,9 +690,10 @@ def _try_blackwell_sparse_prefill( k2q_row_ptr, k2q_q_indices, schedule, partial, stats, softmax_scale=softmax_scale, max_seqlen_q=max_seqlen_q, seqused_k=seqused_k, ) - blackwell_combine( + combine( partial, stats, out, None, cu_seqlens=cu_seqlens_q, split_counts=schedule.split_counts, use_pdl=True, raw_partial_stats=False, + output_scale=output_scale, ) return out @@ -1131,7 +1137,7 @@ def sparse_atten_nvfp4_kv_func( kv_layout=kv_layout, ) - combine( + _select_combine(O_partial.dtype)( O_partial, LSE_partial, O_out, @@ -1702,7 +1708,7 @@ def _sparse_atten_csr_varlen_forward( ) # Sparse Attention and Sparse Page Attention both use the varlen-Q # combine path; the kernel-written LSE_out is the final contract. - combine( + _select_combine(O_partial.dtype)( O_partial, LSE_partial, O_out, diff --git a/python/fmha_sm100/cute/src/blackwell_prefill/combine.py b/python/fmha_sm100/cute/src/blackwell_prefill/combine.py deleted file mode 100644 index 700e8886..00000000 --- a/python/fmha_sm100/cute/src/blackwell_prefill/combine.py +++ /dev/null @@ -1,2354 +0,0 @@ -"""Sparse forward combine kernel and public launcher. - -The epilogue converts the partial-output STG.128 fake layout to the real -output-column order. -""" - -import math -from functools import partial -from typing import Optional, Type - -import cutlass -from cutlass import Boolean, Float32, Int32, Int64, const_expr -import cutlass.cute as cute -from cutlass.cute import FastDivmodDivisor -from cutlass.cute.nvgpu import cpasync -import cuda.bindings.driver as cuda -import torch - -from src.common import utils -from src.common.aot_cache import save_aot, try_load_aot -from src.common.cute_dsl_utils import compile_with_timing -from src.common.cute_dsl_utils import assume_tensor_aligned, torch2cute_dtype_map -from src.common.pack_gqa import PackGQAComb -from src.common.seqlen_info import SeqlenInfo -from src.common.tma_utils import stg128_half_fake_col_to_real_col - - -class SparseAttentionForwardCombine: - def __init__( - self, - dtype: Type[cutlass.Numeric], - dtype_partial: Type[cutlass.Numeric], - head_dim: int, - tile_m: int = 8, - k_block_size: int = 64, - topk: int = 16, - num_threads: int = 256, - stages: int = 4, - use_pdl: bool = False, - min_blocks_per_mp: int = 0, - raw_partial_stats: bool = False, - ): - """ - Forward combine kernel for split attention computation. - - :param dtype: output data type - :param dtype_partial: partial accumulation data type - :param head_dim: head dimension - :param tile_m: m block size - :param k_block_size: k block size - :param topk: exact number of split partials - :param num_threads: number of threads - :param stages: number of pipeline stages - """ - self.dtype = dtype - self.dtype_partial = dtype_partial - self.head_dim = head_dim - self.tile_m = tile_m - self.k_block_size = k_block_size - self.topk = topk - self.num_threads = num_threads - self.is_even_k = head_dim % k_block_size == 0 - self.stages = stages - self.use_pdl = use_pdl - self.min_blocks_per_mp = min_blocks_per_mp - self.raw_partial_stats = raw_partial_stats - self.use_tma_partial = ( - raw_partial_stats - and dtype == cutlass.BFloat16 - and dtype_partial == cutlass.BFloat16 - and head_dim == 128 - and tile_m in (32, 64) - and k_block_size == 128 - and topk == 16 - and num_threads == tile_m * 4 - and stages == 2 - ) - - @staticmethod - def can_implement( - dtype, - dtype_partial, - head_dim, - tile_m, - k_block_size, - topk, - num_threads, - ) -> bool: - """Check if the kernel can be implemented with the given parameters.""" - if dtype not in [cutlass.Float16, cutlass.BFloat16, cutlass.Float32]: - return False - if dtype_partial not in [cutlass.Float16, cutlass.BFloat16]: - return False - if head_dim % 8 != 0: - return False - if num_threads % 32 != 0: - return False - if tile_m % 8 != 0: - return False - if topk > 256: - return False - if (tile_m * topk) % num_threads != 0: - return False - return True - - def _setup_attributes(self): - # GMEM copy setup for O partial - universal_copy_bits = 128 - async_copy_elems = universal_copy_bits // self.dtype_partial.width - assert self.k_block_size % async_copy_elems == 0 - - k_block_gmem = ( - 128 if self.k_block_size % 128 == 0 else (64 if self.k_block_size % 64 == 0 else 32) - ) - gmem_threads_per_row = k_block_gmem // async_copy_elems - assert self.num_threads % gmem_threads_per_row == 0 - - # Async copy atom for O partial load - atom_async_copy_partial = cute.make_copy_atom( - cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), - self.dtype_partial, - num_bits_per_copy=universal_copy_bits, - ) - tOpartial_layout = cute.make_ordered_layout( - (self.num_threads // gmem_threads_per_row, gmem_threads_per_row), - order=(1, 0), - ) - vOpartial_layout = cute.make_layout((1, async_copy_elems)) - self.gmem_tiled_copy_O_partial = cute.make_tiled_copy_tv( - atom_async_copy_partial, tOpartial_layout, vOpartial_layout - ) - - # GMEM copy setup for final O (use universal copy for store). - # Keep output stores independent from the partial-load copy atom. - output_copy_elems = universal_copy_bits // self.dtype.width - assert self.k_block_size % output_copy_elems == 0 - gmem_threads_per_row_o = k_block_gmem // output_copy_elems - assert self.num_threads % gmem_threads_per_row_o == 0 - atom_universal_copy = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), - self.dtype, - num_bits_per_copy=universal_copy_bits, - ) - tO_layout = cute.make_ordered_layout( - (self.num_threads // gmem_threads_per_row_o, gmem_threads_per_row_o), - order=(1, 0), - ) - vO_layout = cute.make_layout((1, output_copy_elems)) - self.gmem_tiled_copy_O = cute.make_tiled_copy_tv( - atom_universal_copy, - tO_layout, - vO_layout, - ) - # LSE copy setup with async copy (alignment = 1) - lse_copy_bits = Float32.width # 1 element per copy, width is in bits - m_block_smem = ( - 128 - if self.tile_m % 128 == 0 - else ( - 64 - if self.tile_m % 64 == 0 - else (32 if self.tile_m % 32 == 0 else (16 if self.tile_m % 16 == 0 else 8)) - ) - ) - gmem_threads_per_row_lse = m_block_smem - assert self.num_threads % gmem_threads_per_row_lse == 0 - - # Async copy atom for LSE load - atom_async_copy_lse = cute.make_copy_atom( - cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.ALWAYS), - Float32, - num_bits_per_copy=lse_copy_bits, - ) - tLSE_layout = cute.make_ordered_layout( - (self.num_threads // gmem_threads_per_row_lse, gmem_threads_per_row_lse), - order=(1, 0), - ) - vLSE_layout = cute.make_layout(1) - self.gmem_tiled_copy_LSE = cute.make_tiled_copy_tv( - atom_async_copy_lse, tLSE_layout, vLSE_layout - ) - - # /////////////////////////////////////////////////////////////////////////////// - # Shared memory - # /////////////////////////////////////////////////////////////////////////////// - - # Shared memory to register copy for LSE - self.smem_threads_per_col_lse = self.num_threads // m_block_smem - assert 32 % self.smem_threads_per_col_lse == 0 # Must divide warp size - - s2r_layout_atom_lse = cute.make_ordered_layout( - (self.smem_threads_per_col_lse, self.num_threads // self.smem_threads_per_col_lse), - order=(0, 1), - ) - self.s2r_tiled_copy_LSE = cute.make_tiled_copy_tv( - cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), Float32), - s2r_layout_atom_lse, - cute.make_layout(1), - ) - - # LSE shared memory layout with swizzling to avoid bank conflicts - # This works for kBlockMSmem = 8, 16, 32, 64, 128, no bank conflicts - if const_expr(m_block_smem == 8): - smem_lse_swizzle = cute.make_swizzle(5, 0, 5) - elif const_expr(m_block_smem == 16): - smem_lse_swizzle = cute.make_swizzle(4, 0, 4) - else: - smem_lse_swizzle = cute.make_swizzle(3, 2, 3) - lse_atom_splits = min(self.topk, 8) - smem_layout_atom_lse = cute.make_composed_layout( - smem_lse_swizzle, - 0, - cute.make_ordered_layout((lse_atom_splits, m_block_smem), order=(1, 0)), - ) - self.smem_layout_lse = cute.tile_to_shape( - smem_layout_atom_lse, (self.topk, self.tile_m), (0, 1) - ) - - # O_partial staging layout. The target raw-BF16 specialization uses a - # row-major tile that is shared by the TMA and consumer-thread views. - if const_expr(self.use_tma_partial): - self.smem_layout_o = cute.make_ordered_layout( - (self.tile_m, self.k_block_size, self.stages), - order=(1, 0, 2), - ) - self.smem_layout_o_tma = cute.make_ordered_layout( - (self.k_block_size, self.tile_m, self.stages), - order=(0, 1, 2), - ) - else: - smem_layout_atom_o = _get_cpasync_smem_layout_atom( - self.dtype_partial, self.k_block_size - ) - self.smem_layout_o = cute.tile_to_shape( - smem_layout_atom_o, - (self.tile_m, self.k_block_size, self.stages), - (0, 1, 2), - ) - self.smem_layout_o_tma = None - - @cute.jit - def __call__( - self, - mO_partial: cute.Tensor, - mLSE_partial: cute.Tensor, - mO: cute.Tensor, - mLSE: Optional[cute.Tensor] = None, - mLSE_temperature_partial: Optional[cute.Tensor] = None, - mLSE_temperature: Optional[cute.Tensor] = None, - cu_seqlens: Optional[cute.Tensor] = None, - seqused: Optional[cute.Tensor] = None, - num_splits_dynamic_ptr: Optional[cute.Tensor] = None, - varlen_batch_idx: Optional[cute.Tensor] = None, - semaphore_to_reset: Optional[cute.Tensor] = None, - mSplitCounts: Optional[cute.Tensor] = None, - mOutputScale: Optional[cute.Tensor] = None, - qhead_per_kvhead: Int32 = Int32(1), - # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI). - stream: cuda.CUstream = None, - ): - # Type checking - if const_expr(not (mO_partial.element_type == self.dtype_partial)): - raise TypeError("O partial tensor must match dtype_partial") - if const_expr(not (mO.element_type == self.dtype)): - raise TypeError("O tensor must match dtype") - if const_expr(mLSE_partial.element_type not in [Float32]): - raise TypeError("LSE partial tensor must be Float32") - if const_expr(mLSE is not None and mLSE.element_type not in [Float32]): - raise TypeError("LSE tensor must be Float32") - if const_expr( - mLSE_temperature_partial is not None - and mLSE_temperature_partial.element_type not in [Float32] - ): - raise TypeError("temperature LSE partial tensor must be Float32") - if const_expr(mLSE_temperature is not None and mLSE_temperature.element_type not in [Float32]): - raise TypeError("temperature LSE tensor must be Float32") - if const_expr((mLSE_temperature_partial is None) != (mLSE_temperature is None)): - raise ValueError( - "temperature LSE partial and output tensors must either both be provided or both be None" - ) - - # Shape validation - input tensors are in user format, need to be converted to kernel format - if const_expr(len(mO_partial.shape) not in [4, 5]): - raise ValueError( - "O partial tensor must have 4 or 5 dimensions: (num_splits, batch, seqlen, nheads, headdim) or (num_splits, total_q, nheads, headdim)" - ) - if const_expr(self.raw_partial_stats): - if const_expr(len(mLSE_partial.shape) != 4): - raise ValueError( - "Raw partial stats must have shape " - "(num_splits, total_q, nheads, 2)" - ) - if const_expr(cu_seqlens is None): - raise ValueError("Raw partial stats currently require the varlen combine path") - if const_expr(mLSE_temperature_partial is not None): - raise ValueError("Raw partial stats do not support independent temperature LSE") - elif const_expr(len(mLSE_partial.shape) not in [3, 4]): - raise ValueError( - "LSE partial tensor must have 3 or 4 dimensions: (num_splits, batch, seqlen, nheads) or (num_splits, total_q, nheads)" - ) - if const_expr(len(mO.shape) not in [3, 4]): - raise ValueError( - "O tensor must have 3 or 4 dimensions: (batch, seqlen, nheads, headdim) or (total_q, nheads, headdim)" - ) - if const_expr(mLSE is not None and len(mLSE.shape) not in [2, 3]): - raise ValueError( - "LSE tensor must have 2 or 3 dimensions: (batch, seqlen, nheads) or (total_q, nheads)" - ) - if const_expr(mLSE_temperature_partial is not None and len(mLSE_temperature_partial.shape) not in [3, 4]): - raise ValueError( - "temperature LSE partial tensor must have 3 or 4 dimensions: " - "(num_splits, batch, seqlen, nheads) or (num_splits, total_q, nheads)" - ) - if const_expr(mLSE_temperature is not None and len(mLSE_temperature.shape) not in [2, 3]): - raise ValueError( - "temperature LSE tensor must have 2 or 3 dimensions: " - "(batch, seqlen, nheads) or (total_q, nheads)" - ) - if const_expr(mSplitCounts is not None): - if const_expr(mSplitCounts.element_type not in [Int32]): - raise TypeError("split_counts tensor must be Int32") - if const_expr(cu_seqlens is not None): - if const_expr(len(mSplitCounts.shape) != 2): - raise ValueError("varlen split_counts tensor must have shape (total_q, nheads_kv)") - elif const_expr(len(mSplitCounts.shape) != 3): - raise ValueError("batched split_counts tensor must have shape (batch, seqlen, nheads_kv)") - if const_expr(mOutputScale is not None and mOutputScale.element_type not in [Float32]): - raise TypeError("output_scale tensor must be Float32") - - mO_partial, mO = [assume_tensor_aligned(t) for t in (mO_partial, mO)] - # (num_splits, b, seqlen, h, d) -> (seqlen, d, num_splits, h, b) - # or (num_splits, total_q, h, d) -> (total_q, d, num_splits, h) - O_partial_layout_transpose = ( - [2, 4, 0, 3, 1] if const_expr(cu_seqlens is None) else [1, 3, 0, 2] - ) - # (b, seqlen, h, d) -> (seqlen, d, h, b) or (total_q, h, d) -> (total_q, d, h) - mO_partial = cute.make_tensor( - mO_partial.iterator, cute.select(mO_partial.layout, mode=O_partial_layout_transpose) - ) - O_layout_transpose = [1, 3, 2, 0] if const_expr(cu_seqlens is None) else [0, 2, 1] - mO = cute.make_tensor(mO.iterator, cute.select(mO.layout, mode=O_layout_transpose)) - # (num_splits, b, h, seqlen) -> (seqlen, num_splits, h, b) - # Input is pre-transposed: [topK, B, Hq, Sq] with Sq innermost for K2-friendly reads. - # or (num_splits, total_q, h) -> (total_q, num_splits, h) - LSE_partial_layout_transpose = [3, 0, 2, 1] if const_expr(cu_seqlens is None) else [1, 0, 2] - mRawStats = None - mRowSum_partial = None - if const_expr(self.raw_partial_stats): - # [split, q, head, stat] -> [q, split, head, stat]. Keep the - # interleaved pair ABI while exposing scalar component views to K2. - mPartialStats = cute.make_tensor( - mLSE_partial.iterator, - cute.select(mLSE_partial.layout, mode=[1, 0, 2, 3]), - ) - mRawStats = mPartialStats - mLSE_partial = mPartialStats[None, None, None, 0] - mRowSum_partial = mPartialStats[None, None, None, 1] - else: - mLSE_partial = cute.make_tensor( - mLSE_partial.iterator, - cute.select(mLSE_partial.layout, mode=LSE_partial_layout_transpose), - ) - # (b, seqlen, h) -> (seqlen, h, b) or (total_q, h) -> (total_q, h) - LSE_layout_transpose = [1, 2, 0] if const_expr(cu_seqlens is None) else [0, 1] - mLSE = ( - cute.make_tensor(mLSE.iterator, cute.select(mLSE.layout, mode=LSE_layout_transpose)) - if mLSE is not None - else None - ) - mLSE_temperature_partial = ( - cute.make_tensor( - mLSE_temperature_partial.iterator, - cute.select(mLSE_temperature_partial.layout, mode=LSE_partial_layout_transpose), - ) - if mLSE_temperature_partial is not None - else None - ) - mLSE_temperature = ( - cute.make_tensor( - mLSE_temperature.iterator, - cute.select(mLSE_temperature.layout, mode=LSE_layout_transpose), - ) - if mLSE_temperature is not None - else None - ) - - # Determine if we have variable length sequences - varlen = const_expr(cu_seqlens is not None or seqused is not None) - - self._setup_attributes() - - # Flatten [split, q, head] into one TMA row coordinate, matching the - # CUDA combine descriptor. A separate split mode would require a - # D * total_q * heads stride; that stride reaches exactly 2**32 - # elements for the 512K target and wraps them in a CuTe DSL descriptor - # lowering. The flattened row extent remains representable and keeps - # every 32x128 transfer contiguous. - tma_atom_O_partial = None - mO_partial_tma = None - if const_expr(self.use_tma_partial): - total_flat = mO_partial.shape[0] * mO_partial.shape[3] - mO_partial_tma = cute.make_tensor( - mO_partial.iterator, - cute.make_layout( - (self.k_block_size, total_flat * self.topk), - stride=( - 1, - self.k_block_size, - ), - ), - ) - tma_atom_O_partial, mO_partial_tma = cpasync.make_tiled_tma_atom( - cpasync.CopyBulkTensorTileG2SOp(), - mO_partial_tma, - cute.select(self.smem_layout_o_tma, mode=[0, 1]), - (self.k_block_size, self.tile_m), - ) - - raw_stats_copy_atom = None - if const_expr(self.use_tma_partial): - raw_stats_copy_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), - Float32, - num_bits_per_copy=64, - ) - - smem_layout_lse_aux = ( - cute.make_layout((1,), stride=(1,)) - if const_expr(self.use_tma_partial) - else self.smem_layout_lse - ) - - # BF16/F16 output uses the register-shuffle epilogue and does not need - # the legacy shared-memory permutation buffer. Keep a one-element - # placeholder so one SharedStorage type can serve all specializations. - if const_expr(self.dtype in [cutlass.Float16, cutlass.BFloat16]): - smem_layout_perm = cute.make_layout((1,), stride=(1,)) - else: - smem_layout_perm = cute.make_ordered_layout( - (self.tile_m, self.k_block_size), order=(1, 0) - ) - - @cute.struct - class SharedStorage: - sLSE: cute.struct.Align[ - cute.struct.MemRange[Float32, cute.cosize(self.smem_layout_lse)], 128 - ] - sLSETemperature: cute.struct.Align[ - cute.struct.MemRange[Float32, cute.cosize(smem_layout_lse_aux)], 128 - ] - sMaxValidSplit: cute.struct.Align[cute.struct.MemRange[Int32, self.tile_m], 128] - sO: cute.struct.Align[ - cute.struct.MemRange[self.dtype_partial, cute.cosize(self.smem_layout_o)], 128 - ] - sO_perm: cute.struct.Align[ - cute.struct.MemRange[self.dtype, cute.cosize(smem_layout_perm)], 128 - ] - sOTmaMbar: cute.struct.Align[ - cute.struct.MemRange[Int64, self.stages], 128 - ] - smem_size = SharedStorage.size_in_bytes() - - @cute.struct - class TmaRawSharedStorage: - sOTmaMbar: cute.struct.Align[ - cute.struct.MemRange[Int64, self.stages], 128 - ] - sLSE: cute.struct.Align[ - cute.struct.MemRange[Float32, cute.cosize(self.smem_layout_lse)], 128 - ] - sMaxValidSplit: cute.struct.Align[ - cute.struct.MemRange[Int32, self.tile_m], 128 - ] - sO: cute.struct.Align[ - cute.struct.MemRange[ - self.dtype_partial, - cute.cosize(self.smem_layout_o), - ], - 128, - ] - tma_raw_smem_size = TmaRawSharedStorage.size_in_bytes() - - # Flatten concatenated varlen rows into one scheduling domain. Each - # combine row is independent, so this preserves the varlen layout - # while avoiding one total_q-sized grid per batch. - flatten_varlen = const_expr( - cu_seqlens is not None - and seqused is None - and varlen_batch_idx is None - ) - - # Grid: (ceil(rows * heads / tile_m), ceil(dim / k_block), domains) - # Head separated from rows enables future TMA on contiguous row tiles. - seqlen = mO_partial.shape[0] - num_head = mO_partial.shape[3] - batch_size = ( - Int32(1) - if const_expr(flatten_varlen) - else ( - mO_partial.shape[4] - if const_expr(cu_seqlens is None) - else Int32(cu_seqlens.shape[0] - 1) - ) - ) - - seqlen_divmod = FastDivmodDivisor(seqlen) - head_divmod = FastDivmodDivisor(num_head) - - grid_dim = ( - cute.ceil_div(seqlen * num_head, self.tile_m), - cute.ceil_div(self.head_dim, self.k_block_size), - batch_size, - ) - - if const_expr(self.use_tma_partial): - if const_expr(not flatten_varlen): - raise ValueError("TMA raw combine requires flattened varlen tensors") - self.kernel_tma_raw( - mO_partial_tma, - mRawStats, - mO, - mLSE, - semaphore_to_reset, - mSplitCounts, - mOutputScale, - qhead_per_kvhead, - tma_atom_O_partial, - raw_stats_copy_atom, - TmaRawSharedStorage, - self.smem_layout_lse, - self.smem_layout_o, - self.smem_layout_o_tma, - self.gmem_tiled_copy_O_partial, - self.gmem_tiled_copy_O, - head_divmod, - self.use_pdl, - ).launch( - grid=grid_dim, - block=[self.num_threads, 1, 1], - smem=tma_raw_smem_size, - stream=stream, - min_blocks_per_mp=self.min_blocks_per_mp, - use_pdl=self.use_pdl, - ) - else: - self.kernel( - mO_partial, - mO_partial_tma, - mLSE_partial, - mRowSum_partial, - mRawStats, - mO, - mLSE, - mLSE_temperature_partial, - mLSE_temperature, - cu_seqlens, - seqused, - num_splits_dynamic_ptr, - varlen_batch_idx, - semaphore_to_reset, - mSplitCounts, - mOutputScale, - qhead_per_kvhead, - tma_atom_O_partial, - raw_stats_copy_atom, - SharedStorage, - self.smem_layout_lse, - smem_layout_lse_aux, - self.smem_layout_o, - self.smem_layout_o_tma, - smem_layout_perm, - self.gmem_tiled_copy_O_partial, - self.gmem_tiled_copy_O, - self.gmem_tiled_copy_LSE, - self.s2r_tiled_copy_LSE, - seqlen_divmod, - head_divmod, - self.use_pdl, - varlen, - flatten_varlen, - ).launch( - grid=grid_dim, - block=[self.num_threads, 1, 1], - smem=smem_size, - stream=stream, - min_blocks_per_mp=self.min_blocks_per_mp, - use_pdl=self.use_pdl, - ) - - @cute.jit - def decode_flat_row_idx( - self, - idx: Int32, - head_divmod: FastDivmodDivisor, - ): - """Decode flattened tile rows under the H_q-innermost contract.""" - q_idx_local, head_idx = divmod(idx, head_divmod) - return q_idx_local, head_idx - - @cute.kernel - def kernel_tma_raw( - self, - mO_partial_tma: cute.Tensor, - mRawStats: cute.Tensor, - mO: cute.Tensor, - mLSE: Optional[cute.Tensor], - semaphore_to_reset: Optional[cute.Tensor], - mSplitCounts: Optional[cute.Tensor], - mOutputScale: Optional[cute.Tensor], - qhead_per_kvhead: Int32, - tma_atom_O_partial: cute.CopyAtom, - raw_stats_copy_atom: cute.CopyAtom, - SharedStorage: cutlass.Constexpr, - smem_layout_lse: cute.Layout | cute.ComposedLayout, - smem_layout_o: cute.Layout | cute.ComposedLayout, - smem_layout_o_tma: cute.Layout, - gmem_tiled_copy_O_partial: cute.TiledCopy, - gmem_tiled_copy_O: cute.TiledCopy, - head_divmod: FastDivmodDivisor, - use_pdl: cutlass.Constexpr[bool], - ): - """CUDA-shaped raw-stat combine specialization for the SM100 target.""" - tidx, _, _ = cute.arch.thread_idx() - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - m_block, k_block, maybe_virtual_batch = cute.arch.block_idx() - - if const_expr(use_pdl and semaphore_to_reset is not None): - if ( - tidx == 0 - and m_block == cute.arch.grid_dim()[0] - 1 - and k_block == cute.arch.grid_dim()[1] - 1 - and maybe_virtual_batch == cute.arch.grid_dim()[2] - 1 - ): - cute.arch.griddepcontrol_wait() - semaphore_to_reset[0] = 0 - if const_expr(use_pdl): - cute.arch.griddepcontrol_wait() - - smem = cutlass.utils.SmemAllocator() - storage = smem.allocate(SharedStorage) - sLSE = storage.sLSE.get_tensor(smem_layout_lse) - sMaxValidSplit = storage.sMaxValidSplit.get_tensor((self.tile_m,)) - sO = storage.sO.get_tensor(smem_layout_o) - sO_tma = storage.sO.get_tensor(smem_layout_o_tma) - sO_tma_mbar_ptr = storage.sOTmaMbar.data_ptr() - - total_flat = mRawStats.shape[0] * mRawStats.shape[2] - tile_begin = m_block * self.tile_m - output_scale = Float32(1.0) - if const_expr(mOutputScale is not None): - output_scale = mOutputScale[0] - - # Initialize the two transaction barriers and overlap the initial TMA - # loads with raw-stat reduction, matching the CUDA schedule. - gO_partial_tma = cute.local_tile( - mO_partial_tma, - (self.k_block_size, self.tile_m), - (k_block, None), - ) - tOsO_tma, tOgO_tma = cpasync.tma_partition( - tma_atom_O_partial, - 0, - cute.make_layout(1), - cute.group_modes(sO_tma, 0, 2), - cute.group_modes(gO_partial_tma, 0, 2), - ) - o_tma_copy_bytes = const_expr( - cute.size_in_bytes( - self.dtype_partial, - cute.select(smem_layout_o_tma, mode=[0, 1]), - ) - ) - row_tiles_per_split = total_flat // Int32(self.tile_m) - if tidx == 0: - for stage in cutlass.range(self.stages, unroll_full=True): - cute.arch.mbarrier_init(sO_tma_mbar_ptr + stage, 1) - cute.arch.mbarrier_init_fence() - cute.arch.sync_threads() - if warp_idx == 0: - for stage in cutlass.range(self.stages, unroll_full=True): - o_tma_barrier = sO_tma_mbar_ptr + stage - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - o_tma_barrier, - o_tma_copy_bytes, - ) - cute.copy( - tma_atom_O_partial, - tOgO_tma[ - None, - m_block + Int32(stage) * row_tiles_per_split, - ], - tOsO_tma[None, stage], - tma_bar_ptr=o_tma_barrier, - ) - - # Four adjacent lanes own one row and four interleaved split pairs. - reduce_lane = tidx % Int32(4) - reduce_row = tidx // Int32(4) - reduce_flat = tile_begin + reduce_row - row_valid = reduce_flat < total_flat - m_idx = Int32(0) - head_idx = Int32(0) - row_split_count = Int32(0) - if row_valid: - m_idx, head_idx = self.decode_flat_row_idx( - reduce_flat, - head_divmod, - ) - row_split_count = ( - mSplitCounts[m_idx, head_idx // qhead_per_kvhead] - if const_expr(mSplitCounts is not None) - else Int32(self.topk) - ) - - local_page_max = cute.make_rmem_tensor((4,), Float32) - local_page_sum = cute.make_rmem_tensor((4,), Float32) - raw_stats = cute.make_rmem_tensor((2,), Float32) - local_max = -Float32.inf - local_max_valid = Int32(-1) - for item in cutlass.range(4, unroll_full=True): - split = reduce_lane + item * Int32(4) - raw_stats[0] = -Float32.inf - raw_stats[1] = Float32(0.0) - if row_valid and split < row_split_count: - cute.copy( - raw_stats_copy_atom, - mRawStats[m_idx, split, head_idx, None], - raw_stats, - ) - stats_valid = ( - raw_stats[1] > Float32(0.0) - and raw_stats[1] < Float32.inf - and raw_stats[0] > -Float32.inf - and raw_stats[0] < Float32.inf - and raw_stats[0] == raw_stats[0] - and raw_stats[1] == raw_stats[1] - ) - page_max = raw_stats[0] if stats_valid else -Float32.inf - page_sum = raw_stats[1] if stats_valid else Float32(0.0) - local_page_max[item] = page_max - local_page_sum[item] = page_sum - local_max = max(local_max, page_max) - if stats_valid: - local_max_valid = split - - stats_max = cute.arch.warp_reduction_max( - local_max, - threads_in_group=4, - ) - max_valid_split = cute.arch.warp_reduction_max( - local_max_valid, - threads_in_group=4, - ) - stats_max_safe = ( - stats_max if max_valid_split >= Int32(0) else Float32(0.0) - ) - stats_sum = Float32(0.0) - for item in cutlass.range(4, unroll_full=True): - split_scale = cute.math.exp2( - local_page_max[item] - stats_max_safe, - fastmath=True, - ) - local_page_max[item] = split_scale - stats_sum += split_scale * local_page_sum[item] - stats_sum = cute.arch.warp_reduction_sum( - stats_sum, - threads_in_group=4, - ) - stats_sum_valid = ( - max_valid_split >= Int32(0) - and stats_sum > Float32(0.0) - and stats_sum < Float32.inf - and stats_sum == stats_sum - ) - inv_sum = Float32(0.0) - final_lse = -Float32.inf - if stats_sum_valid: - inv_sum = Float32(1.0) / stats_sum - final_lse = ( - stats_max + cute.math.log2(stats_sum, fastmath=True) - ) * Float32(math.log(2.0)) - for item in cutlass.range(4, unroll_full=True): - split = reduce_lane + item * Int32(4) - sLSE[split, reduce_row] = local_page_max[item] * inv_sum - if reduce_lane == Int32(0): - sMaxValidSplit[reduce_row] = max_valid_split - if const_expr(mLSE is not None): - if row_valid: - mLSE[m_idx, head_idx] = final_lse - - cute.arch.sync_threads() - block_max_valid_split = Int32(-1) - if const_expr(self.tile_m == 32): - if tidx < Int32(self.tile_m): - block_max_valid_split = cute.arch.warp_reduction_max( - sMaxValidSplit[tidx], - threads_in_group=self.tile_m, - ) - if tidx == 0: - sMaxValidSplit[0] = block_max_valid_split - else: - if tidx < Int32(32): - block_max_valid_split = cute.arch.warp_reduction_max( - max( - sMaxValidSplit[tidx], - sMaxValidSplit[tidx + Int32(32)], - ), - threads_in_group=32, - ) - if tidx == 0: - sMaxValidSplit[0] = block_max_valid_split - cute.arch.sync_threads() - block_max_valid_split = sMaxValidSplit[0] - - cO = cute.make_identity_tensor((self.tile_m, self.k_block_size)) - partial_copy = gmem_tiled_copy_O_partial.get_slice(tidx) - tOcO = partial_copy.partition_D(cO) - tOsO_partial = partial_copy.partition_D(sO) - num_rows = const_expr(cute.size(tOcO, mode=[1])) - tOrO_partial_row = cute.make_rmem_tensor_like( - tOsO_partial[None, 0, None, 0] - ) - tOrO = cute.make_rmem_tensor_like( - tOsO_partial[None, None, None, 0], - Float32, - ) - tOrO.fill(0.0) - - for split in cutlass.range(block_max_valid_split + 1, unroll=4): - stage = split % Int32(self.stages) - phase = (split // Int32(self.stages)) & Int32(1) - cute.arch.mbarrier_wait(sO_tma_mbar_ptr + stage, phase) - for row in cutlass.range(num_rows, unroll_full=True): - cute.autovec_copy( - tOsO_partial[None, row, None, stage], - tOrO_partial_row, - ) - weight = sLSE[split, tOcO[0, row, 0][0]] - if weight != Float32(0.0): - tOrO[None, row, None].store( - tOrO[None, row, None].load() - + weight - * tOrO_partial_row.load().to(Float32) - ) - - cute.arch.sync_threads() - split_to_load = split + Int32(self.stages) - if warp_idx == 0 and split_to_load <= block_max_valid_split: - o_tma_barrier = sO_tma_mbar_ptr + stage - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - o_tma_barrier, - o_tma_copy_bytes, - ) - cute.copy( - tma_atom_O_partial, - tOgO_tma[ - None, - m_block + split_to_load * row_tiles_per_split, - ], - tOsO_tma[None, stage], - tma_bar_ptr=o_tma_barrier, - ) - - # Transpose the four 32-bit BF16 pairs within each four-lane group, - # then retain one coalesced 128-bit output store per row. - output_copy = gmem_tiled_copy_O.get_slice(tidx) - tOcO_store = output_copy.partition_D(cO) - rO_row = cute.make_rmem_tensor( - tOcO_store[None, 0, 0].shape, - self.dtype, - ) - elems_per_store = const_expr( - cute.size(gmem_tiled_copy_O.layout_tv_tiled[1]) - ) - num_store_rows = const_expr(cute.size(tOcO_store, mode=[1])) - lane_in_pair_group = tidx % Int32(4) - for row in cutlass.range(num_store_rows, unroll_full=True): - for k in cutlass.range( - cute.size(tOcO_store, mode=[2]), - unroll_full=True, - ): - o0, o1 = tOrO[0, row, k], tOrO[1, row, k] - o2, o3 = tOrO[2, row, k], tOrO[3, row, k] - o4, o5 = tOrO[4, row, k], tOrO[5, row, k] - o6, o7 = tOrO[6, row, k], tOrO[7, row, k] - if const_expr(mOutputScale is not None): - o0, o1 = cute.arch.mul_packed_f32x2( - (o0, o1), - (output_scale, output_scale), - ) - o2, o3 = cute.arch.mul_packed_f32x2( - (o2, o3), - (output_scale, output_scale), - ) - o4, o5 = cute.arch.mul_packed_f32x2( - (o4, o5), - (output_scale, output_scale), - ) - o6, o7 = cute.arch.mul_packed_f32x2( - (o6, o7), - (output_scale, output_scale), - ) - word0 = utils.cvt_f16x2_f32(o0, o1, self.dtype) - word1 = utils.cvt_f16x2_f32(o2, o3, self.dtype) - word2 = utils.cvt_f16x2_f32(o4, o5, self.dtype) - word3 = utils.cvt_f16x2_f32(o6, o7, self.dtype) - - partner0 = cute.arch.shuffle_sync_bfly(word1, offset=1) - partner1 = cute.arch.shuffle_sync_bfly(word0, offset=1) - partner2 = cute.arch.shuffle_sync_bfly(word3, offset=1) - partner3 = cute.arch.shuffle_sync_bfly(word2, offset=1) - stage0, stage1 = word0, word1 - stage2, stage3 = word2, word3 - if lane_in_pair_group % Int32(2) != Int32(0): - stage0, stage2 = partner0, partner2 - if lane_in_pair_group % Int32(2) != Int32(1): - stage1, stage3 = partner1, partner3 - - partner0 = cute.arch.shuffle_sync_bfly(stage2, offset=2) - partner1 = cute.arch.shuffle_sync_bfly(stage3, offset=2) - partner2 = cute.arch.shuffle_sync_bfly(stage0, offset=2) - partner3 = cute.arch.shuffle_sync_bfly(stage1, offset=2) - out0, out1 = stage0, stage1 - out2, out3 = stage2, stage3 - if lane_in_pair_group // Int32(2) != Int32(0): - out0, out1 = partner0, partner1 - if lane_in_pair_group // Int32(2) != Int32(1): - out2, out3 = partner2, partner3 - - rO_pairs = cute.recast_tensor(rO_row, cutlass.Int32) - rO_pairs[0] = out0 - rO_pairs[1] = out1 - rO_pairs[2] = out2 - rO_pairs[3] = out3 - row_local = tOcO_store[0, row, 0][0] - flat_idx = tile_begin + row_local - if flat_idx < total_flat: - m_idx, head_idx = self.decode_flat_row_idx( - flat_idx, - head_divmod, - ) - mO_cur_copy = cute.tiled_divide( - mO[m_idx, None, head_idx], - (elems_per_store,), - ) - k_idx = tOcO_store[0, 0, k][1] // elems_per_store - cute.copy( - output_copy, - rO_row, - mO_cur_copy[None, k_idx], - ) - - @cute.kernel - def kernel( - self, - mO_partial: cute.Tensor, - mO_partial_tma: Optional[cute.Tensor], - mLSE_partial: cute.Tensor, - mRowSum_partial: Optional[cute.Tensor], - mRawStats: Optional[cute.Tensor], - mO: cute.Tensor, - mLSE: Optional[cute.Tensor], - mLSE_temperature_partial: Optional[cute.Tensor], - mLSE_temperature: Optional[cute.Tensor], - cu_seqlens: Optional[cute.Tensor], - seqused: Optional[cute.Tensor], - num_splits_dynamic_ptr: Optional[cute.Tensor], - varlen_batch_idx: Optional[cute.Tensor], - semaphore_to_reset: Optional[cute.Tensor], - mSplitCounts: Optional[cute.Tensor], - mOutputScale: Optional[cute.Tensor], - qhead_per_kvhead: Int32, - tma_atom_O_partial: Optional[cute.CopyAtom], - raw_stats_copy_atom: Optional[cute.CopyAtom], - SharedStorage: cutlass.Constexpr, - smem_layout_lse: cute.Layout | cute.ComposedLayout, - smem_layout_lse_aux: cute.Layout | cute.ComposedLayout, - smem_layout_o: cute.Layout | cute.ComposedLayout, - smem_layout_o_tma: Optional[cute.Layout], - smem_layout_perm: cute.Layout, - gmem_tiled_copy_O_partial: cute.TiledCopy, - gmem_tiled_copy_O: cute.TiledCopy, - gmem_tiled_copy_LSE: cute.TiledCopy, - s2r_tiled_copy_LSE: cute.TiledCopy, - seqlen_divmod: FastDivmodDivisor, - head_divmod: FastDivmodDivisor, - use_pdl: cutlass.Constexpr[bool], - varlen: cutlass.Constexpr[bool], - flatten_varlen: cutlass.Constexpr[bool], - ): - # Thread and block indices - tidx, _, _ = cute.arch.thread_idx() - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - m_block, k_block, maybe_virtual_batch = cute.arch.block_idx() - - if const_expr(flatten_varlen): - batch_idx = Int32(0) - else: - batch_idx = ( - varlen_batch_idx[maybe_virtual_batch] - if const_expr(varlen_batch_idx is not None) - else maybe_virtual_batch - ) - - # /////////////////////////////////////////////////////////////////////////////// - # Get shared memory buffer - # /////////////////////////////////////////////////////////////////////////////// - smem = cutlass.utils.SmemAllocator() - storage = smem.allocate(SharedStorage) - sLSE = storage.sLSE.get_tensor(smem_layout_lse) - sLSE_temperature = storage.sLSETemperature.get_tensor(smem_layout_lse_aux) - sMaxValidSplit = storage.sMaxValidSplit.get_tensor((self.tile_m,)) - sO = storage.sO.get_tensor(smem_layout_o) - sO_tma = None - sO_tma_mbar_ptr = None - if const_expr(self.use_tma_partial): - sO_tma = storage.sO.get_tensor(smem_layout_o_tma) - sO_tma_mbar_ptr = storage.sOTmaMbar.data_ptr() - sO_perm_buf = storage.sO_perm.get_tensor(smem_layout_perm) - - # Handle semaphore reset — wait for dependent grids first - if const_expr(use_pdl and semaphore_to_reset is not None): - if ( - tidx == 0 - and m_block == cute.arch.grid_dim()[0] - 1 - and k_block == cute.arch.grid_dim()[1] - 1 - and maybe_virtual_batch == cute.arch.grid_dim()[2] - 1 - ): - cute.arch.griddepcontrol_wait() - semaphore_to_reset[0] = 0 - - if const_expr(num_splits_dynamic_ptr is not None): - raise ValueError("K2 combine requires compile-time exact topK") - num_splits = Int32(self.topk) - # Concatenated varlen tensors already have total_q as their leading - # extent. Other modes retain SeqlenInfo-based per-batch slicing. - if const_expr(flatten_varlen): - seqlen = Int32(mO_partial.shape[0]) - offset = Int32(0) - else: - seqlen_info = SeqlenInfo.create( - batch_idx=batch_idx, - seqlen_static=mO_partial.shape[0], - cu_seqlens=cu_seqlens, - seqused=seqused, - # Don't need tile size because offset_padded is unused. - ) - seqlen, offset = seqlen_info.seqlen, seqlen_info.offset - - num_head = mO_partial.shape[3] - max_idx = seqlen * num_head - output_scale = Float32(1.0) - if const_expr(mOutputScale is not None): - output_scale = mOutputScale[0] - - if const_expr(not varlen) or m_block * self.tile_m < max_idx: - # Wait for dependent grids (e.g., the main attention kernel that produces O_partial/LSE_partial) - if const_expr(use_pdl): - cute.arch.griddepcontrol_wait() - - # =============================== - # Step 1: Load LSE_partial from gmem to shared memory - # =============================== - # `cLSE` (identity tensor for row/split coord tracking) is reused - # later in steps 4-5, so it must be defined on both branches. - cLSE = cute.make_identity_tensor((self.topk, self.tile_m)) - # Reshape mLSE_partial to PackGQA packed layout and delegate the - # tile load to PackGQAComb.load_LSE. The packed form folds (H_q, Sq) - # into one compound dim with H_q innermost (stride 1), so thread - # rows that vary along h_pos produce one-sector coalesced reads. - # Non-varlen path only — varlen keeps the original inline loop. - if const_expr(not varlen): - mLSE_partial_cur = seqlen_info.offset_batch(mLSE_partial, batch_idx, dim=3) - # mLSE_partial_cur: (H_q, topK, Sq) — after initial transpose - # [3,0,2,1] on [topK,B,Sq,H_q] and dropping B. - # Reorder to (H_q, Sq, topK) then group modes 0..1 for packed dim: - mLSE_partial_reord = cute.make_tensor( - mLSE_partial_cur.iterator, - cute.select(mLSE_partial_cur.layout, mode=[0, 2, 1]), - ) - mLSE_partial_packed = cute.group_modes(mLSE_partial_reord, 0, 2) - # shape ((H_q, Sq), topK) with H_q innermost. - packgqa = PackGQAComb( - m_block_size=self.tile_m, - head_dim_padded=0, # unused for LSE load - check_hdim_oob=False, # unused for LSE load - qhead_per_kvhead=1, # unused; num_heads_divmod is passed explicitly - ) - packgqa.load_LSE( - mLSE_partial_packed, - sLSE, - self.topk, - gmem_tiled_copy_LSE, - tidx, - m_block, - num_splits, - seqlen, - head_divmod, - mSplitCounts, - batch_idx, - qhead_per_kvhead, - ) - if const_expr(mLSE_temperature_partial is not None): - mLSE_temperature_partial_cur = seqlen_info.offset_batch( - mLSE_temperature_partial, batch_idx, dim=3) - mLSE_temperature_partial_reord = cute.make_tensor( - mLSE_temperature_partial_cur.iterator, - cute.select(mLSE_temperature_partial_cur.layout, mode=[0, 2, 1]), - ) - mLSE_temperature_partial_packed = cute.group_modes( - mLSE_temperature_partial_reord, 0, 2) - packgqa.load_LSE( - mLSE_temperature_partial_packed, - sLSE_temperature, - self.topk, - gmem_tiled_copy_LSE, - tidx, - m_block, - num_splits, - seqlen, - head_divmod, - mSplitCounts, - batch_idx, - qhead_per_kvhead, - ) - else: - # Varlen path keeps the same H_q-innermost flat-row contract: - # after transpose [1, 0, 2], mLSE_partial_cur is - # (q_local, split, head). - # mSplitCounts is the authoritative valid-split count per - # packed (q_abs, kv_head); masked splits stay at -inf and - # therefore drop out of the final kernel LSE_out reduction. - if const_expr(flatten_varlen): - mLSE_partial_cur = mLSE_partial - if const_expr(self.raw_partial_stats): - mRowSum_partial_cur = mRowSum_partial - else: - mLSE_partial_cur = seqlen_info.offset_batch( - mLSE_partial, batch_idx, dim=3 - ) - if const_expr(self.raw_partial_stats): - mRowSum_partial_cur = seqlen_info.offset_batch( - mRowSum_partial, batch_idx, dim=3 - ) - mLSE_partial_copy = cute.tiled_divide(mLSE_partial_cur, (1,)) - if const_expr(self.raw_partial_stats): - mRowSum_partial_copy = cute.tiled_divide(mRowSum_partial_cur, (1,)) - gmem_thr_copy_LSE = gmem_tiled_copy_LSE.get_slice(tidx) - tLSEsLSE = gmem_thr_copy_LSE.partition_D(sLSE) - tLSEsLSE_temperature = gmem_thr_copy_LSE.partition_D(sLSE_temperature) - tLSEcLSE = gmem_thr_copy_LSE.partition_S(cLSE) - if const_expr(mLSE_temperature_partial is not None): - if const_expr(flatten_varlen): - mLSE_temperature_partial_cur = mLSE_temperature_partial - else: - mLSE_temperature_partial_cur = seqlen_info.offset_batch( - mLSE_temperature_partial, batch_idx, dim=3 - ) - mLSE_temperature_partial_copy = cute.tiled_divide( - mLSE_temperature_partial_cur, (1,)) - - for m in cutlass.range(cute.size(tLSEcLSE, mode=[2]), unroll_full=True): - mi = tLSEcLSE[0, 0, m][1] - idx = m_block * self.tile_m + mi - if idx < max_idx: - m_idx, head_idx = self.decode_flat_row_idx(idx, head_divmod) - row_count = ( - mSplitCounts[offset + m_idx, head_idx // qhead_per_kvhead] - if const_expr(mSplitCounts is not None) - else num_splits - ) - mLSE_partial_cur_copy = mLSE_partial_copy[None, m_idx, None, head_idx] - if const_expr(self.raw_partial_stats): - mRowSum_partial_cur_copy = ( - mRowSum_partial_copy[None, m_idx, None, head_idx] - ) - if const_expr(mLSE_temperature_partial is not None): - mLSE_temperature_partial_cur_copy = ( - mLSE_temperature_partial_copy[None, m_idx, None, head_idx]) - for s in cutlass.range(cute.size(tLSEcLSE, mode=[1]), unroll_full=True): - si = tLSEcLSE[0, s, 0][0] - if si < num_splits and si < row_count: - cute.copy( - gmem_thr_copy_LSE, - mLSE_partial_cur_copy[None, si], - tLSEsLSE[None, s, m], - ) - if const_expr(self.raw_partial_stats): - cute.copy( - gmem_thr_copy_LSE, - mRowSum_partial_cur_copy[None, si], - tLSEsLSE_temperature[None, s, m], - ) - if const_expr(mLSE_temperature_partial is not None): - cute.copy( - gmem_thr_copy_LSE, - mLSE_temperature_partial_cur_copy[None, si], - tLSEsLSE_temperature[None, s, m], - ) - else: - tLSEsLSE[None, s, m].fill(-Float32.inf) - if const_expr(self.raw_partial_stats): - tLSEsLSE_temperature[None, s, m].fill(Float32(0.0)) - if const_expr(mLSE_temperature_partial is not None): - tLSEsLSE_temperature[None, s, m].fill(-Float32.inf) - else: - for s in cutlass.range(cute.size(tLSEcLSE, mode=[1]), unroll_full=True): - tLSEsLSE[None, s, m].fill(-Float32.inf) - if const_expr(self.raw_partial_stats): - tLSEsLSE_temperature[None, s, m].fill(Float32(0.0)) - if const_expr(mLSE_temperature_partial is not None): - tLSEsLSE_temperature[None, s, m].fill(-Float32.inf) - cute.arch.cp_async_commit_group() - - # =============================== - # Step 2: Load O_partial for pipeline stages - # =============================== - - gmem_thr_copy_O_partial = gmem_tiled_copy_O_partial.get_slice(tidx) - cO = cute.make_identity_tensor((self.tile_m, self.k_block_size)) - tOcO = gmem_thr_copy_O_partial.partition_D(cO) - tOsO_partial = gmem_thr_copy_O_partial.partition_D(sO) - - # Precompute per-row values for flattened (q_local, head) tiles. - num_rows = const_expr(cute.size(tOcO, mode=[1])) - tOmidx = cute.make_rmem_tensor(num_rows, cutlass.Int32) - tOhidx = cute.make_rmem_tensor(num_rows, cutlass.Int32) - if const_expr(not self.use_tma_partial): - tOSplitCount = cute.make_rmem_tensor(num_rows, cutlass.Int32) - tOrOptr = cute.make_rmem_tensor(num_rows, cutlass.Int64) - if const_expr(flatten_varlen): - mO_partial_cur = mO_partial - else: - mO_partial_cur = seqlen_info.offset_batch( - mO_partial, batch_idx, dim=4 - ) - for m in cutlass.range(num_rows, unroll_full=True): - mi = tOcO[0, m, 0][0] # m coordinate in tile - idx = m_block * self.tile_m + mi - if idx >= max_idx: - tOhidx[m] = -1 - tOmidx[m] = 0 - if const_expr(not self.use_tma_partial): - tOSplitCount[m] = 0 - tOrOptr[m] = cutlass.Int64(0) - else: - tOmidx[m], tOhidx[m] = self.decode_flat_row_idx(idx, head_divmod) - if const_expr(not self.use_tma_partial): - if const_expr(mSplitCounts is None): - tOSplitCount[m] = num_splits - elif const_expr(cu_seqlens is None): - tOSplitCount[m] = mSplitCounts[ - batch_idx, tOmidx[m], tOhidx[m] // qhead_per_kvhead - ] - else: - tOSplitCount[m] = mSplitCounts[ - offset + tOmidx[m], tOhidx[m] // qhead_per_kvhead - ] - tOrOptr[m] = utils.elem_pointer( - mO_partial_cur, - (tOmidx[m], k_block * self.k_block_size, 0, tOhidx[m]), - ).toint() - - tOpO = None - if const_expr(not self.is_even_k): - tOpO = cute.make_rmem_tensor(cute.size(tOcO, mode=[2]), Boolean) - for k in cutlass.range(cute.size(tOpO), unroll_full=True): - tOpO[k] = tOcO[0, 0, k][1] < mO_partial.shape[1] - k_block * self.k_block_size - if const_expr(self.use_tma_partial): - gO_partial_tma = cute.local_tile( - mO_partial_tma, - (self.k_block_size, self.tile_m), - (k_block, m_block, None), - ) - tOsO_tma, tOgO_tma = cpasync.tma_partition( - tma_atom_O_partial, - 0, - cute.make_layout(1), - cute.group_modes(sO_tma, 0, 2), - cute.group_modes(gO_partial_tma, 0, 2), - ) - o_tma_copy_bytes = const_expr( - cute.size_in_bytes( - self.dtype_partial, - cute.select(smem_layout_o_tma, mode=[0, 1]), - ) - ) - if tidx == 0: - for stage in cutlass.range(self.stages, unroll_full=True): - cute.arch.mbarrier_init(sO_tma_mbar_ptr + stage, 1) - cute.arch.mbarrier_init_fence() - cute.arch.sync_threads() - if warp_idx == 0: - for stage in cutlass.range(self.stages, unroll_full=True): - o_tma_barrier = sO_tma_mbar_ptr + stage - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - o_tma_barrier, - o_tma_copy_bytes, - ) - cute.copy( - tma_atom_O_partial, - tOgO_tma[None, stage], - tOsO_tma[None, stage], - tma_bar_ptr=o_tma_barrier, - ) - else: - load_O_partial = partial( - self.load_O_partial, - gmem_tiled_copy_O_partial, - tOrOptr, - tOsO_partial, - tOhidx, - tOSplitCount, - tOpO, - tOcO, - mO_partial_cur.layout, - ) - - # Load the initial cp.async stages of O_partial. - for stage in cutlass.range(self.stages - 1, unroll_full=True): - if stage < num_splits: - load_O_partial(stage, stage) - cute.arch.cp_async_commit_group() - - # =============================== - # Step 3: Load and transpose LSE from smem to registers - # =============================== - - # Wait for LSE and initial O partial stages to complete - if const_expr(self.use_tma_partial): - cute.arch.cp_async_wait_group(0) - else: - cute.arch.cp_async_wait_group(self.stages - 1) - cute.arch.sync_threads() - - s2r_thr_copy_LSE = s2r_tiled_copy_LSE.get_slice(tidx) - ts2rsLSE = s2r_thr_copy_LSE.partition_S(sLSE) - ts2rrLSE = cute.make_rmem_tensor_like(ts2rsLSE) - cute.copy(s2r_tiled_copy_LSE, ts2rsLSE, ts2rrLSE) - if const_expr( - self.raw_partial_stats or mLSE_temperature_partial is not None - ): - ts2rsLSE_temperature = s2r_thr_copy_LSE.partition_S(sLSE_temperature) - ts2rrLSE_temperature = cute.make_rmem_tensor_like(ts2rsLSE_temperature) - cute.copy( - s2r_tiled_copy_LSE, - ts2rsLSE_temperature, - ts2rrLSE_temperature, - ) - - # =============================== - # Step 4: Compute final LSE along split dimension - # =============================== - - final_lse = cute.make_rmem_tensor(cute.size(ts2rrLSE, mode=[2]), Float32) - ts2rcLSE = s2r_thr_copy_LSE.partition_D(cLSE) - # We compute the max valid split for each row to short-circuit the computation later - max_valid_split = cute.make_rmem_tensor(cute.size(ts2rrLSE, mode=[2]), Int32) - assert cute.size(ts2rrLSE, mode=[0]) == 1 - # Compute max, scales, and final LSE for each row. Invalid splits - # have already been filled with -inf, so Step 5 can write the - # kernel-native LSE_out directly. - for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True): - threads_per_col = const_expr(self.smem_threads_per_col_lse) - if const_expr(self.raw_partial_stats): - # Sanitize the raw pair before reductions so NaN/Inf cannot - # contaminate a neighboring split in the thread group. - max_valid_idx = Int32(-1) - for s in cutlass.range( - cute.size(ts2rrLSE, mode=[1]), unroll_full=True - ): - scaled_max_log2 = ts2rrLSE[0, s, m] - row_sum = ts2rrLSE_temperature[0, s, m] - stats_valid = ( - row_sum > Float32(0.0) - and row_sum < Float32.inf - and scaled_max_log2 > -Float32.inf - and scaled_max_log2 < Float32.inf - and row_sum == row_sum - and scaled_max_log2 == scaled_max_log2 - ) - ts2rrLSE[0, s, m] = ( - scaled_max_log2 if stats_valid else -Float32.inf - ) - ts2rrLSE_temperature[0, s, m] = ( - row_sum if stats_valid else Float32(0.0) - ) - if stats_valid: - max_valid_idx = ts2rcLSE[0, s, 0][0] - - stats_max = cute.arch.warp_reduction_max( - ts2rrLSE[None, None, m] - .load() - .reduce( - cute.ReductionOp.MAX, - init_val=-Float32.inf, - reduction_profile=0, - ), - threads_in_group=threads_per_col, - ) - max_valid_split[m] = cute.arch.warp_reduction_max( - max_valid_idx, threads_in_group=threads_per_col - ) - stats_max_safe = ( - stats_max - if max_valid_split[m] >= Int32(0) - else Float32(0.0) - ) - stats_sum = Float32(0.0) - for s in cutlass.range( - cute.size(ts2rrLSE, mode=[1]), unroll_full=True - ): - row_sum = ts2rrLSE_temperature[0, s, m] - split_scale = Float32(0.0) - if row_sum > Float32(0.0): - split_scale = cute.math.exp2( - ts2rrLSE[0, s, m] - stats_max_safe, - fastmath=True, - ) - stats_sum += split_scale * row_sum - # Raw partial-O already includes the local row sum, so - # its combine weight excludes row_sum from the numerator. - ts2rrLSE[0, s, m] = split_scale - stats_sum = cute.arch.warp_reduction_sum( - stats_sum, threads_in_group=threads_per_col - ) - stats_sum_valid = ( - max_valid_split[m] >= Int32(0) - and stats_sum > Float32(0.0) - and stats_sum < Float32.inf - and stats_sum == stats_sum - ) - inv_sum = Float32(0.0) - if stats_sum_valid: - inv_sum = Float32(1.0) / stats_sum - final_lse[m] = ( - stats_max - + cute.math.log2(stats_sum, fastmath=True) - ) * Float32(math.log(2.0)) - else: - final_lse[m] = -Float32.inf - ts2rrLSE[None, None, m].store( - ts2rrLSE[None, None, m].load() * inv_sum - ) - else: - # Find max normalized-partial LSE value across splits. - lse_max = cute.arch.warp_reduction_max( - ts2rrLSE[None, None, m] - .load() - .reduce( - cute.ReductionOp.MAX, - init_val=-Float32.inf, - reduction_profile=0, - ), - threads_in_group=threads_per_col, - ) - max_valid_idx = -1 - for s in cutlass.range( - cute.size(ts2rrLSE, mode=[1]), unroll_full=True - ): - if ts2rrLSE[0, s, m] != -Float32.inf: - max_valid_idx = ts2rcLSE[0, s, 0][0] - max_valid_split[m] = cute.arch.warp_reduction_max( - max_valid_idx, threads_in_group=threads_per_col - ) - lse_max_cur = Float32(0.0) if lse_max == -Float32.inf else lse_max - LOG2_E = math.log2(math.e) - lse_sum_cur = Float32(0.0) - for s in cutlass.range( - cute.size(ts2rrLSE, mode=[1]), unroll_full=True - ): - scale = cute.math.exp2( - ts2rrLSE[0, s, m] * LOG2_E - - lse_max_cur * LOG2_E, - fastmath=True, - ) - lse_sum_cur += scale - ts2rrLSE[0, s, m] = scale - lse_sum_cur = cute.arch.warp_reduction_sum( - lse_sum_cur, threads_in_group=threads_per_col - ) - inv_sum = Float32(0.0) - if ( - max_valid_split[m] < 0 - or lse_sum_cur == Float32(0.0) - or lse_sum_cur != lse_sum_cur - ): - final_lse[m] = -Float32.inf - else: - final_lse[m] = ( - cute.math.log(lse_sum_cur, fastmath=True) + lse_max - ) - inv_sum = Float32(1.0) / lse_sum_cur - ts2rrLSE[None, None, m].store( - ts2rrLSE[None, None, m].load() * inv_sum - ) - # Store the scales exp(lse - lse_logsum) back to smem - cute.copy(s2r_tiled_copy_LSE, ts2rrLSE, ts2rsLSE) - - if const_expr(mLSE_temperature_partial is not None): - final_lse_temperature = cute.make_rmem_tensor( - cute.size(ts2rrLSE_temperature, mode=[2]), Float32) - for m in cutlass.range(cute.size(ts2rrLSE_temperature, mode=[2]), unroll_full=True): - threads_per_col = const_expr(self.smem_threads_per_col_lse) - lse_temperature_max = cute.arch.warp_reduction_max( - ts2rrLSE_temperature[None, None, m] - .load() - .reduce( - cute.ReductionOp.MAX, - init_val=-Float32.inf, - reduction_profile=0, - ), - threads_in_group=threads_per_col, - ) - lse_temperature_max_cur = ( - 0.0 if lse_temperature_max == -Float32.inf else lse_temperature_max - ) - LOG2_E = math.log2(math.e) - lse_temperature_sum_cur = 0.0 - for s in cutlass.range( - cute.size(ts2rrLSE_temperature, mode=[1]), unroll_full=True): - scale = cute.math.exp2( - ts2rrLSE_temperature[0, s, m] * LOG2_E - - (lse_temperature_max_cur * LOG2_E), - fastmath=True, - ) - lse_temperature_sum_cur += scale - lse_temperature_sum_cur = cute.arch.warp_reduction_sum( - lse_temperature_sum_cur, threads_in_group=threads_per_col - ) - if ( - max_valid_split[m] < 0 - or lse_temperature_sum_cur == 0.0 - or lse_temperature_sum_cur != lse_temperature_sum_cur - ): - final_lse_temperature[m] = -Float32.inf - else: - final_lse_temperature[m] = ( - cute.math.log(lse_temperature_sum_cur, fastmath=True) - + lse_temperature_max - ) - - # Store max valid split to smem - for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True): - if ts2rcLSE[0, 0, m][0] == 0: # Only thread responsible for s=0 writes - mi = ts2rcLSE[0, 0, m][1] - if mi < self.tile_m: - sMaxValidSplit[mi] = max_valid_split[m] - - # =============================== - # Step 5: Store final LSE to gmem - # This writeback is the authoritative LSE_out returned by the - # public Sparse Attention / Sparse Page Attention interface. - # =============================== - - if const_expr(mLSE is not None): - if const_expr(cu_seqlens is None): - mLSE_cur = mLSE[None, None, batch_idx] - else: - mLSE_cur = cute.domain_offset((offset, 0), mLSE) - if const_expr(mLSE_temperature is not None): - if const_expr(cu_seqlens is None): - mLSE_temperature_cur = mLSE_temperature[None, None, batch_idx] - else: - mLSE_temperature_cur = cute.domain_offset( - (offset, 0), mLSE_temperature) - if k_block == 0: # Only first k_block writes LSE when mLSE is provided - for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True): - if ts2rcLSE[0, 0, m][0] == 0: # Only thread responsible for s=0 writes - mi = ts2rcLSE[0, 0, m][1] - idx = m_block * self.tile_m + mi - if idx < max_idx: - m_idx, head_idx = self.decode_flat_row_idx(idx, head_divmod) - mLSE_cur[m_idx, head_idx] = final_lse[m] - if const_expr(mLSE_temperature is not None): - mLSE_temperature_cur[m_idx, head_idx] = ( - final_lse_temperature[m]) - - # =============================== - # Step 6: Read O_partial and accumulate final O - # =============================== - - cute.arch.sync_threads() - - if const_expr(self.use_tma_partial): - # TMA stages are CTA-wide resources, so every consumer thread - # must execute the same split count and barrier sequence. - block_max_valid_split = Int32(-1) - if tidx < Int32(self.tile_m): - block_max_valid_split = cute.arch.warp_reduction_max( - sMaxValidSplit[tidx], threads_in_group=self.tile_m - ) - if tidx == 0: - sMaxValidSplit[0] = block_max_valid_split - cute.arch.sync_threads() - thr_max_valid_split = sMaxValidSplit[0] - else: - # cp.async loads are thread-private, so each thread can stop at - # the largest valid split among its own rows. - thr_max_valid_split = sMaxValidSplit[tOcO[0, 0, 0][0]] - for m in cutlass.range( - 1, cute.size(tOcO, mode=[1]), unroll_full=True - ): - thr_max_valid_split = max( - thr_max_valid_split, - sMaxValidSplit[tOcO[0, m, 0][0]], - ) - - tOrO_partial = cute.make_rmem_tensor_like(tOsO_partial[None, None, None, 0]) - tOrO = cute.make_rmem_tensor_like(tOrO_partial, Float32) - tOrO.fill(0.0) - - stage_load = Int32(self.stages - 1) - stage_compute = Int32(0) - - # Main accumulation loop - for s in cutlass.range(thr_max_valid_split + 1, unroll=4): - if const_expr(self.use_tma_partial): - stage_compute = s % Int32(self.stages) - stage_phase = (s // Int32(self.stages)) & Int32(1) - cute.arch.mbarrier_wait( - sO_tma_mbar_ptr + stage_compute, - stage_phase, - ) - - # Get scales for this split - scale = cute.make_rmem_tensor(num_rows, Float32) - for m in cutlass.range(num_rows, unroll_full=True): - scale[m] = sLSE[s, tOcO[0, m, 0][0]] - - if const_expr(self.use_tma_partial): - cute.autovec_copy( - tOsO_partial[None, None, None, stage_compute], - tOrO_partial, - ) - else: - # Load the next thread-private cp.async stage. - split_to_load = s + self.stages - 1 - if split_to_load <= thr_max_valid_split: - load_O_partial(split_to_load, stage_load) - cute.arch.cp_async_commit_group() - stage_load = ( - 0 if stage_load == self.stages - 1 else stage_load + 1 - ) - cute.arch.cp_async_wait_group(self.stages - 1) - cute.autovec_copy( - tOsO_partial[None, None, None, stage_compute], - tOrO_partial, - ) - stage_compute = ( - 0 - if stage_compute == self.stages - 1 - else stage_compute + 1 - ) - - # Accumulate scaled partial results - for m in cutlass.range(num_rows, unroll_full=True): - if tOhidx[m] >= 0 and scale[m] > 0.0: - tOrO[None, m, None].store( - tOrO[None, m, None].load() - + scale[m] * tOrO_partial[None, m, None].load().to(Float32) - ) - - if const_expr(self.use_tma_partial): - # Match CUDA's CTA-wide stage handoff: no consumer may - # release a stage while another warp is still reading it. - cute.arch.sync_threads() - split_to_load = s + self.stages - if warp_idx == 0 and split_to_load <= thr_max_valid_split: - o_tma_barrier = sO_tma_mbar_ptr + stage_compute - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - o_tma_barrier, - o_tma_copy_bytes, - ) - cute.copy( - tma_atom_O_partial, - tOgO_tma[None, split_to_load], - tOsO_tma[None, stage_compute], - tma_bar_ptr=o_tma_barrier, - ) - - # Flush any outstanding async-copy groups before the local Step-7 - # permutation buffer is read on the tail of the kernel. - if const_expr(not self.use_tma_partial): - cute.arch.cp_async_wait_group(0) - cute.arch.sync_threads() - - # =============================== - # Step 7: Write final O to gmem (fake→real via SMEM) - # =============================== - - if const_expr(cu_seqlens is None): - mO_cur = mO[None, None, None, batch_idx] - elif const_expr(flatten_varlen): - mO_cur = mO - else: - mO_cur = cute.domain_offset((offset, 0, 0), mO) - mO_cur = utils.domain_offset_aligned((0, k_block * self.k_block_size, 0), mO_cur) - num_vals = const_expr(cute.size(tOcO, mode=[0])) - if const_expr(not use_pdl): - # Direct / standalone calls don't participate in the K1->K2 - # dependency chain. Use a simple per-element real-column store - # path here to keep mixed-shape launches stable. - for m in cutlass.range(num_rows, unroll_full=True): - if tOhidx[m] >= 0: - for k in cutlass.range(cute.size(tOcO, mode=[2]), unroll_full=True): - if const_expr(self.is_even_k) or tOpO[k]: - for v in cutlass.range(num_vals, unroll_full=True): - fake_col = tOcO[v, 0, k][1] - real_col = stg128_half_fake_col_to_real_col(fake_col) - o_val = tOrO[v, m, k] - if const_expr(mOutputScale is not None): - o_val = o_val * output_scale - mO_cur[tOmidx[m], real_col, tOhidx[m]] = o_val.to(self.dtype) - elif const_expr(self.dtype in [cutlass.BFloat16, cutlass.Float16]): - # The BF16/F16 fake layout is a 4x4 transpose of 32-bit - # pairs inside each 32-column group. Transpose the pair - # registers with two XOR-shuffle stages, then retain the - # original coalesced 128-bit GMEM store. - assert num_vals == 8 - gmem_thr_copy_O = gmem_tiled_copy_O.get_slice(tidx) - tOcO_store = gmem_thr_copy_O.partition_D(cO) - rO = cute.make_rmem_tensor(tOcO_store.shape, self.dtype) - elems_per_store = const_expr( - cute.size(gmem_tiled_copy_O.layout_tv_tiled[1]) - ) - num_store_rows = const_expr(cute.size(tOcO_store, mode=[1])) - assert num_store_rows == num_rows - lane_in_pair_group = tidx % Int32(4) - - for m in cutlass.range(num_store_rows, unroll_full=True): - for k in cutlass.range( - cute.size(tOcO_store, mode=[2]), unroll_full=True - ): - o0, o1 = tOrO[0, m, k], tOrO[1, m, k] - o2, o3 = tOrO[2, m, k], tOrO[3, m, k] - o4, o5 = tOrO[4, m, k], tOrO[5, m, k] - o6, o7 = tOrO[6, m, k], tOrO[7, m, k] - if const_expr(mOutputScale is not None): - o0, o1 = cute.arch.mul_packed_f32x2( - (o0, o1), (output_scale, output_scale) - ) - o2, o3 = cute.arch.mul_packed_f32x2( - (o2, o3), (output_scale, output_scale) - ) - o4, o5 = cute.arch.mul_packed_f32x2( - (o4, o5), (output_scale, output_scale) - ) - o6, o7 = cute.arch.mul_packed_f32x2( - (o6, o7), (output_scale, output_scale) - ) - word0 = utils.cvt_f16x2_f32(o0, o1, self.dtype) - word1 = utils.cvt_f16x2_f32(o2, o3, self.dtype) - word2 = utils.cvt_f16x2_f32(o4, o5, self.dtype) - word3 = utils.cvt_f16x2_f32(o6, o7, self.dtype) - - partner0 = cute.arch.shuffle_sync_bfly(word1, offset=1) - partner1 = cute.arch.shuffle_sync_bfly(word0, offset=1) - partner2 = cute.arch.shuffle_sync_bfly(word3, offset=1) - partner3 = cute.arch.shuffle_sync_bfly(word2, offset=1) - stage0, stage1 = word0, word1 - stage2, stage3 = word2, word3 - if lane_in_pair_group % Int32(2) != Int32(0): - stage0, stage2 = partner0, partner2 - if lane_in_pair_group % Int32(2) != Int32(1): - stage1, stage3 = partner1, partner3 - - partner0 = cute.arch.shuffle_sync_bfly(stage2, offset=2) - partner1 = cute.arch.shuffle_sync_bfly(stage3, offset=2) - partner2 = cute.arch.shuffle_sync_bfly(stage0, offset=2) - partner3 = cute.arch.shuffle_sync_bfly(stage1, offset=2) - out0, out1 = stage0, stage1 - out2, out3 = stage2, stage3 - if lane_in_pair_group // Int32(2) != Int32(0): - out0, out1 = partner0, partner1 - if lane_in_pair_group // Int32(2) != Int32(1): - out2, out3 = partner2, partner3 - - rO_pairs = cute.recast_tensor( - rO[None, m, k], cutlass.Int32 - ) - rO_pairs[0] = out0 - rO_pairs[1] = out1 - rO_pairs[2] = out2 - rO_pairs[3] = out3 - - for m in cutlass.range(num_store_rows, unroll_full=True): - row_local = tOcO_store[0, m, 0][0] - idx = m_block * self.tile_m + row_local - if idx < max_idx: - m_idx, head_idx = self.decode_flat_row_idx( - idx, head_divmod - ) - mO_cur_copy = cute.tiled_divide( - mO_cur[m_idx, None, head_idx], (elems_per_store,) - ) - for k in cutlass.range( - cute.size(tOcO_store, mode=[2]), unroll_full=True - ): - k_idx = ( - tOcO_store[0, 0, k][1] // elems_per_store - ) - cute.copy( - gmem_thr_copy_O, - rO[None, m, k], - mO_cur_copy[None, k_idx], - ) - else: - # 7a: fp32 accumulator -> output dtype SMEM with fake→real - # permutation. The dedicated permutation buffer stays separate - # from the O_partial pipeline staging buffer. - sO_perm = sO_perm_buf - - if const_expr(self.dtype in [cutlass.BFloat16, cutlass.Float16]): - # O_partial uses a dtype-specific STG.128 fake layout, but - # sO_perm is in the final O dtype. For all supported fake - # layouts, adjacent fake pairs map to adjacent real columns, - # so write the final BF16/F16 O pair as one 32-bit SMEM store. - assert num_vals % 2 == 0 - r2s_o_pair_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), - cutlass.Int32, - num_bits_per_copy=32, - ) - rO_pair_word = cute.make_rmem_tensor((1,), cutlass.Int32) - sO_perm_i32_base = cute.make_ptr( - dtype=cutlass.Int32, - value=sO_perm.iterator.toint(), - mem_space=sO_perm.iterator.memspace, - assumed_align=4, - ) - sO_perm_i32_row_stride = Int32((self.k_block_size + 16) // 2) - for m in cutlass.range(num_rows, unroll_full=True): - row_local = tOcO[0, m, 0][0] - if tOhidx[m] >= 0: - for k in cutlass.range(cute.size(tOcO, mode=[2]), unroll_full=True): - for v_pair in cutlass.range(num_vals // 2, unroll_full=True): - v = v_pair * 2 - fake_col = tOcO[v, 0, k][1] - real_col = stg128_half_fake_col_to_real_col(fake_col) - o0 = tOrO[v, m, k] - o1 = tOrO[v + 1, m, k] - if const_expr(mOutputScale is not None): - o0, o1 = cute.arch.mul_packed_f32x2( - (o0, o1), - (output_scale, output_scale), - ) - rO_pair_word[0] = utils.cvt_f16x2_f32(o0, o1, self.dtype) - smem_pair_ptr = cute.make_ptr( - dtype=cutlass.Int32, - value=( - sO_perm_i32_base.toint() - + Int64( - row_local * sO_perm_i32_row_stride - + real_col // Int32(2) - ) - * Int64(4) - ), - mem_space=sO_perm.iterator.memspace, - assumed_align=4, - ) - sO_pair = cute.make_tensor( - smem_pair_ptr, - cute.make_layout((1,), stride=(1,)), - ) - cute.copy(r2s_o_pair_atom, rO_pair_word, sO_pair) - else: - # 7a: iterate over ALL val elements in mode[0]. - # tOcO[v, m, k][1] gives different fake_col for each v. - r2s_o_scalar_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), - self.dtype, - num_bits_per_copy=self.dtype.width, - ) - rO_scalar = cute.make_rmem_tensor((1,), self.dtype) - for m in cutlass.range(num_rows, unroll_full=True): - row_local = tOcO[0, m, 0][0] - if tOhidx[m] >= 0: - for k in cutlass.range(cute.size(tOcO, mode=[2]), unroll_full=True): - for v in cutlass.range(num_vals, unroll_full=True): - fake_col = tOcO[v, 0, k][1] - real_col = stg128_half_fake_col_to_real_col(fake_col) - o_val = tOrO[v, m, k] - if const_expr(mOutputScale is not None): - o_val = o_val * output_scale - rO_scalar[0] = o_val.to(self.dtype) - smem_ptr = utils.elem_pointer(sO_perm, (row_local, real_col)) - smem_scalar_ptr = cute.make_ptr( - dtype=self.dtype, - value=smem_ptr.toint(), - mem_space=sO_perm.iterator.memspace, - assumed_align=self.dtype.width // 8, - ) - sO_scalar = cute.make_tensor( - smem_scalar_ptr, - cute.make_layout((1,), stride=(1,)), - ) - cute.copy(r2s_o_scalar_atom, rO_scalar, sO_scalar) - - cute.arch.sync_threads() - - # 7b: SMEM (real order, output dtype) → GMEM - gmem_thr_copy_O = gmem_tiled_copy_O.get_slice(tidx) - tOcO_store = gmem_thr_copy_O.partition_D(cO) - tOsO_store = gmem_thr_copy_O.partition_D(sO_perm) - rO = cute.make_rmem_tensor(tOcO_store.shape, self.dtype) - elems_per_store = const_expr(cute.size(gmem_tiled_copy_O.layout_tv_tiled[1])) - num_store_rows = const_expr(cute.size(tOcO_store, mode=[1])) - tOpO_store = None - if const_expr(not self.is_even_k): - tOpO_store = cute.make_rmem_tensor(cute.size(tOcO_store, mode=[2]), Boolean) - for k in cutlass.range(cute.size(tOpO_store), unroll_full=True): - tOpO_store[k] = ( - tOcO_store[0, 0, k][1] - < mO_partial.shape[1] - k_block * self.k_block_size - ) - - # Read output dtype from SMEM (now in real column order). - for m in cutlass.range(num_store_rows, unroll_full=True): - for k in cutlass.range(cute.size(tOcO_store, mode=[2]), unroll_full=True): - if const_expr(self.is_even_k) or tOpO_store[k]: - cute.autovec_copy(tOsO_store[None, m, k], rO[None, m, k]) - - # Write bf16 to GMEM using gmem_tiled_copy_O (same as original FA Step 7) - for m in cutlass.range(num_store_rows, unroll_full=True): - row_local = tOcO_store[0, m, 0][0] - idx = m_block * self.tile_m + row_local - if idx < max_idx: - m_idx, head_idx = self.decode_flat_row_idx(idx, head_divmod) - mO_cur_copy = cute.tiled_divide( - mO_cur[m_idx, None, head_idx], (elems_per_store,) - ) - for k in cutlass.range(cute.size(tOcO_store, mode=[2]), unroll_full=True): - k_idx = tOcO_store[0, 0, k][1] // elems_per_store - if const_expr(self.is_even_k) or tOpO_store[k]: - cute.copy(gmem_thr_copy_O, rO[None, m, k], mO_cur_copy[None, k_idx]) - - @cute.jit - def load_O_partial( - self, - gmem_tiled_copy_O_partial: cute.TiledCopy, - tOrOptr: cute.Tensor, - tOsO_partial: cute.Tensor, - tOhidx: cute.Tensor, - tOSplitCount: cute.Tensor, - tOpO: Optional[cute.Tensor], - tOcO: cute.Tensor, - mO_cur_partial_layout: cute.Layout, - split: Int32, - stage: Int32, - ) -> None: - elems_per_load = const_expr(cute.size(gmem_tiled_copy_O_partial.layout_tv_tiled[1])) - tOsO_partial_cur = tOsO_partial[None, None, None, stage] - for m in cutlass.range(cute.size(tOcO, [1]), unroll_full=True): - if tOhidx[m] >= 0: - o_gmem_ptr = cute.make_ptr( - tOsO_partial.element_type, - tOrOptr[m], - cute.AddressSpace.gmem, - assumed_align=16, - ) - mO_partial_cur = cute.make_tensor( - o_gmem_ptr, cute.slice_(mO_cur_partial_layout, (0, None, None, 0)) - ) - mO_partial_cur_copy = cute.tiled_divide(mO_partial_cur, (elems_per_load,)) - for k in cutlass.range(cute.size(tOcO, mode=[2]), unroll_full=True): - k_idx = tOcO[0, 0, k][1] // elems_per_load - if split < tOSplitCount[m] and (const_expr(tOpO is None) or tOpO[k]): - cute.copy( - gmem_tiled_copy_O_partial, - mO_partial_cur_copy[None, k_idx, split], - tOsO_partial_cur[None, m, k], - ) - else: - tOsO_partial_cur[None, m, k].fill(0) - - -def _get_cutlass_dtype(torch_dtype: torch.dtype): - if torch_dtype not in torch2cute_dtype_map: - raise TypeError(f"Unsupported dtype: {torch_dtype}") - return torch2cute_dtype_map[torch_dtype] - - -_combine_compile_cache = {} - - -def _get_cpasync_smem_layout_atom(dtype: Type[cutlass.Numeric], k_dim: int) -> cute.ComposedLayout: - dtype_byte = const_expr(dtype.width // 8) - bytes_per_row = const_expr(k_dim * dtype_byte) - smem_k_block_size = ( - const_expr( - 128 - if bytes_per_row % 128 == 0 - else (64 if bytes_per_row % 64 == 0 else (32 if bytes_per_row % 32 == 0 else 16)) - ) - // dtype_byte - ) - swizzle_bits = ( - 4 - if smem_k_block_size == 128 - else (3 if smem_k_block_size == 64 else (2 if smem_k_block_size == 32 else 1)) - ) - swizzle_base = 2 if dtype_byte == 4 else (3 if dtype_byte == 2 else 4) - return cute.make_composed_layout( - cute.make_swizzle(swizzle_bits, swizzle_base, swizzle_base), - 0, - cute.make_ordered_layout( - (8 if const_expr(k_dim % 32 == 0) else 16, smem_k_block_size), - order=(1, 0), - ), - ) - - -def combine( - o_partial_fake, - lse_partial, - o_out, - lse_out, - *, - lse_temperature_partial=None, - lse_temperature_out=None, - cu_seqlens=None, - seqused=None, - split_counts=None, - output_scale=None, - use_pdl=False, - raw_partial_stats=False, -): - """K2: merge sparse forward split partials into the final output. - - STG.128 fake-layout handling remains an internal implementation detail. - When lse_out is provided, the kernel writes the final authoritative - log-sum-exp for each query row/head directly into that tensor. - - Args: - o_partial_fake: - Batched: [num_splits, batch, Sq, head_q, dim] - Varlen: [num_splits, total_q, head_q, dim] - lse_partial: - Batched: [num_splits, batch, Sq, head_q] - Varlen: [num_splits, total_q, head_q] - o_out: - Batched: [batch, Sq, head_q, dim] - Varlen: [total_q, head_q, dim] - lse_out: - Batched: [batch, Sq, head_q] - Varlen: [total_q, head_q] - lse_temperature_partial: - Optional temperature-scaled LSE partial with the same shape as - lse_partial. - lse_temperature_out: - Optional temperature-scaled final LSE with the same shape as - lse_out. - cu_seqlens: Optional [batch + 1] int32 for varlen-Q combine. - seqused: Optional [batch] int32 effective lengths for combine. - split_counts: Optional int32 rowwise valid split counts prepared from - q2k metadata. Batched: [batch, seqlen, head_kv]. Varlen: - [total_q, head_kv]. - output_scale: Optional fp32 tensor with at least one element. When - provided, the final O accumulator is multiplied once before store. - use_pdl: When True, wait on PDL dependencies from the producer K1 - kernel. When False, launch without PDL waits. - raw_partial_stats: Interpret lse_partial as interleaved FP32 - (scaled_max_log2, row_sum) pairs and combine raw partial-O. - """ - D = o_partial_fake.shape[-1] - num_splits = o_partial_fake.shape[0] - raw_partial_stats = bool(raw_partial_stats) - return_temperature_lse = ( - lse_temperature_partial is not None or lse_temperature_out is not None - ) - if (lse_temperature_partial is None) != (lse_temperature_out is None): - raise ValueError( - "lse_temperature_partial and lse_temperature_out must either both be provided or both be None" - ) - if raw_partial_stats: - expected_stats_shape = (*o_partial_fake.shape[:-1], 2) - if tuple(lse_partial.shape) != tuple(expected_stats_shape): - raise ValueError( - "raw lse_partial must have shape " - f"{expected_stats_shape}, got {tuple(lse_partial.shape)}" - ) - if cu_seqlens is None or o_partial_fake.ndim != 4: - raise ValueError("raw_partial_stats currently requires varlen partial tensors") - if return_temperature_lse: - raise ValueError("raw_partial_stats does not support independent temperature LSE") - else: - expected_lse_shape = o_partial_fake.shape[:-1] - if tuple(lse_partial.shape) != tuple(expected_lse_shape): - raise ValueError( - "lse_partial shape must match o_partial_fake without head_dim, " - f"got {tuple(lse_partial.shape)} vs {tuple(expected_lse_shape)}" - ) - if lse_partial.dtype != torch.float32: - raise TypeError(f"lse_partial must be torch.float32, got {lse_partial.dtype}") - if lse_temperature_partial is not None and lse_temperature_partial.shape != lse_partial.shape: - raise ValueError( - "lse_temperature_partial must have the same shape as lse_partial, " - f"got {lse_temperature_partial.shape} vs {lse_partial.shape}" - ) - if lse_temperature_out is not None: - if lse_out is None: - raise ValueError("lse_temperature_out requires lse_out") - if lse_temperature_out.shape != lse_out.shape: - raise ValueError( - "lse_temperature_out must have the same shape as lse_out, " - f"got {lse_temperature_out.shape} vs {lse_out.shape}" - ) - if lse_temperature_out.dtype != torch.float32 or lse_temperature_partial.dtype != torch.float32: - raise TypeError("temperature LSE tensors must be torch.float32") - - if o_partial_fake.dtype not in (torch.bfloat16, torch.float16): - raise TypeError( - "o_partial_fake dtype must be torch.bfloat16 or torch.float16, " - f"got {o_partial_fake.dtype}" - ) - partial_dtype = _get_cutlass_dtype(o_partial_fake.dtype) - out_dtype = _get_cutlass_dtype(o_out.dtype) - if output_scale is not None: - if output_scale.dtype != torch.float32: - raise TypeError(f"output_scale must be torch.float32, got {output_scale.dtype}") - if output_scale.numel() < 1: - raise ValueError("output_scale must contain at least one element") - if output_scale.device != o_out.device: - raise ValueError("output_scale must be on the same device as o_out") - output_scale = output_scale.contiguous() - if split_counts is not None: - if split_counts.dtype != torch.int32: - raise TypeError(f"split_counts must be torch.int32, got {split_counts.dtype}") - if o_out.ndim == 4: - if split_counts.ndim != 3: - raise ValueError( - f"batched split_counts must have shape [batch, seqlen, head_kv], got {split_counts.shape}" - ) - if split_counts.shape[:2] != o_out.shape[:2]: - raise ValueError( - f"split_counts shape {split_counts.shape} must match batch/seqlen of o_out {o_out.shape}" - ) - else: - if cu_seqlens is None: - raise ValueError("split_counts with varlen output requires cu_seqlens") - if split_counts.ndim != 2: - raise ValueError( - f"varlen split_counts must have shape [total_q, head_kv], got {split_counts.shape}" - ) - if split_counts.shape[0] != o_out.shape[0]: - raise ValueError( - f"split_counts total_q ({split_counts.shape[0]}) must match o_out total_q " - f"({o_out.shape[0]})" - ) - if o_out.shape[-2] % split_counts.shape[-1] != 0: - raise ValueError( - f"o_out heads ({o_out.shape[-2]}) must be divisible by split_counts heads ({split_counts.shape[-1]})" - ) - qheadperkv = o_out.shape[-2] // split_counts.shape[-1] - else: - qheadperkv = 1 - if cu_seqlens is not None: - if cu_seqlens.dtype != torch.int32: - raise TypeError(f"cu_seqlens must be torch.int32, got {cu_seqlens.dtype}") - if cu_seqlens.ndim != 1: - raise ValueError(f"cu_seqlens must be rank-1, got {cu_seqlens.shape}") - if not cu_seqlens.is_contiguous(): - raise ValueError("cu_seqlens must be contiguous") - if seqused is not None: - if seqused.dtype != torch.int32: - raise TypeError(f"seqused must be torch.int32, got {seqused.dtype}") - if seqused.ndim != 1: - raise ValueError(f"seqused must be rank-1, got {seqused.shape}") - if not seqused.is_contiguous(): - raise ValueError("seqused must be contiguous") - - k_block_size = 128 if D > 64 else 64 - tile_m = 64 if raw_partial_stats else 32 - num_threads = tile_m * 4 - has_cu_seqlens = cu_seqlens is not None - has_seqused = seqused is not None - has_lse = lse_out is not None - has_split_counts = split_counts is not None - has_output_scale = output_scale is not None - min_blocks_per_mp = 3 if has_output_scale and use_pdl else 0 - - capability = torch.cuda.get_device_capability(o_out.device) - stages = 2 - key = ( - "blackwell_prefill_combine", - capability, - stages, - D, - k_block_size, - tile_m, - num_threads, - num_splits, - partial_dtype, - out_dtype, - has_cu_seqlens, - has_seqused, - has_lse, - bool(return_temperature_lse), - has_split_counts, - has_output_scale, - use_pdl, - min_blocks_per_mp, - raw_partial_stats, - ) - if key not in _combine_compile_cache: - loaded = try_load_aot(key) - if loaded is not None: - _combine_compile_cache[key] = loaded - else: - from quack.compile_utils import make_fake_tensor - - kernel = SparseAttentionForwardCombine( - dtype=out_dtype, - dtype_partial=partial_dtype, - head_dim=D, - tile_m=tile_m, - k_block_size=k_block_size, - topk=num_splits, - num_threads=num_threads, - use_pdl=use_pdl, - min_blocks_per_mp=min_blocks_per_mp, - stages=stages, - raw_partial_stats=raw_partial_stats, - ) - div = 128 // partial_dtype.width - if has_cu_seqlens: - total_q, nheads = (cute.sym_int64() for _ in range(2)) - mO_partial = make_fake_tensor( - partial_dtype, (num_splits, total_q, nheads, D), divisibility=div - ) - mLSE_partial = make_fake_tensor( - Float32, - ( - (num_splits, total_q, nheads, 2) - if raw_partial_stats - else (num_splits, total_q, nheads) - ), - divisibility=2 if raw_partial_stats else 1, - leading_dim=3 if raw_partial_stats else 2, - ) - mO = make_fake_tensor( - out_dtype, (total_q, nheads, D), divisibility=128 // out_dtype.width - ) - mLSE = ( - make_fake_tensor(Float32, (total_q, nheads), divisibility=1, leading_dim=1) - if has_lse - else None - ) - mLSE_temperature_partial = ( - make_fake_tensor( - Float32, (num_splits, total_q, nheads), divisibility=1, leading_dim=2 - ) - if return_temperature_lse - else None - ) - mLSE_temperature = ( - make_fake_tensor(Float32, (total_q, nheads), divisibility=1, leading_dim=1) - if return_temperature_lse - else None - ) - else: - batch, sq, nheads = (cute.sym_int64() for _ in range(3)) - mO_partial = make_fake_tensor( - partial_dtype, (num_splits, batch, sq, nheads, D), divisibility=div - ) - mLSE_partial = make_fake_tensor( - Float32, (num_splits, batch, sq, nheads), divisibility=1, leading_dim=3 - ) - mO = make_fake_tensor( - out_dtype, (batch, sq, nheads, D), divisibility=128 // out_dtype.width - ) - mLSE = ( - make_fake_tensor(Float32, (batch, sq, nheads), divisibility=1, leading_dim=2) - if has_lse - else None - ) - mLSE_temperature_partial = ( - make_fake_tensor( - Float32, (num_splits, batch, sq, nheads), divisibility=1, leading_dim=3 - ) - if return_temperature_lse - else None - ) - mLSE_temperature = ( - make_fake_tensor(Float32, (batch, sq, nheads), divisibility=1, leading_dim=2) - if return_temperature_lse - else None - ) - if not has_split_counts: - mSplitCounts = None - elif has_cu_seqlens: - total_q_ctr, nheads_kv = (cute.sym_int64() for _ in range(2)) - mSplitCounts = make_fake_tensor( - Int32, (total_q_ctr, nheads_kv), divisibility=1, leading_dim=1 - ) - else: - nheads_kv = cute.sym_int64() - mSplitCounts = make_fake_tensor( - Int32, (batch, sq, nheads_kv), divisibility=1, leading_dim=2 - ) - mOutputScale = ( - make_fake_tensor(Float32, (cute.sym_int64(),), divisibility=1, leading_dim=0) - if has_output_scale - else None - ) - stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - - _combine_compile_cache[key] = compile_with_timing( - kernel, - mO_partial, - mLSE_partial, - mO, - mLSE, - mLSE_temperature_partial, - mLSE_temperature, - None - if cu_seqlens is None - else make_fake_tensor(Int32, (cute.sym_int64(),), divisibility=1, leading_dim=0), - None - if seqused is None - else make_fake_tensor(Int32, (cute.sym_int64(),), divisibility=1, leading_dim=0), - None, - None, - None, - mSplitCounts, - mOutputScale, - Int32(qheadperkv), - stream, - options="--enable-tvm-ffi", - ) - save_aot(key, _combine_compile_cache[key]) - - with torch.cuda.nvtx.range("K2_Combine"): - _combine_compile_cache[key]( - o_partial_fake, - lse_partial, - o_out, - lse_out, - lse_temperature_partial, - lse_temperature_out, - cu_seqlens, - seqused, - None, - None, - None, - split_counts, - output_scale, - qheadperkv, - ) diff --git a/python/fmha_sm100/cute/src/sm100/fwd/combine.py b/python/fmha_sm100/cute/src/sm100/fwd/combine.py index c60ed95d..45f5c0d1 100644 --- a/python/fmha_sm100/cute/src/sm100/fwd/combine.py +++ b/python/fmha_sm100/cute/src/sm100/fwd/combine.py @@ -1,10 +1,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 MiniMax # SPDX-License-Identifier: MIT -"""Sparse forward combine kernel and public launcher. +"""Sparse forward combine (K2) kernels and launchers. -This keeps the local fake-layout -> real-layout epilogue needed by the lean -sparse forward path. +- ``combine``: the default. Runs ``SparseAttentionForwardCombine`` (the nv_dev + v1 combine); BF16/FP16 partials only. +- ``combine_f32f8``: runs ``SparseAttentionForwardCombineF32F8`` (the original + combine), which also implements the FP32 and FP8 E4M3 STG.128 fake layouts. + Callers that may write FP32 / FP8 partials pick it by partial dtype. """ # Modified Step 7: O_out write with SMEM fake->real column permutation. @@ -22,7 +25,8 @@ from cutlass import Float32, Int32, Int64, Boolean, const_expr from src.common import utils -from src.common.cute_dsl_utils import assume_tensor_aligned, torch2cute_dtype_map +from src.common.aot_cache import save_aot, try_load_aot +from src.common.cute_dsl_utils import assume_tensor_aligned, compile_with_timing, torch2cute_dtype_map from src.common.seqlen_info import SeqlenInfo from cutlass.cute import FastDivmodDivisor @@ -34,7 +38,12 @@ ) -class SparseAttentionForwardCombine: +# ----------------------------------------------------------------------------- +# Original combine (combine_f32f8): FP32 / FP8 E4M3 partials +# ----------------------------------------------------------------------------- + + +class SparseAttentionForwardCombineF32F8: def __init__( self, dtype: Type[cutlass.Numeric], @@ -1161,6 +1170,11 @@ def load_O_partial( tOsO_partial_cur[None, m, k].fill(0) +# ----------------------------------------------------------------------------- +# Shared helpers +# ----------------------------------------------------------------------------- + + def _get_cutlass_dtype(torch_dtype: torch.dtype): if torch_dtype not in torch2cute_dtype_map: raise TypeError(f"Unsupported dtype: {torch_dtype}") @@ -1197,7 +1211,7 @@ def _get_cpasync_smem_layout_atom(dtype: Type[cutlass.Numeric], k_dim: int) -> c ) -def combine( +def combine_f32f8( o_partial_fake, lse_partial, o_out, @@ -1211,7 +1225,9 @@ def combine( output_scale=None, use_pdl=False, ): - """K2: merge sparse forward split partials into the final output. + """K2 on the original combine kernel, for FP32 / FP8 E4M3 partials. + + It also accepts BF16/FP16, but ``combine`` is the faster choice there. STG.128 fake-layout handling remains an internal implementation detail. When lse_out is provided, the kernel writes the final authoritative @@ -1348,7 +1364,7 @@ def combine( stages = 3 if is_sm107 else 2 key = ( - "combine", + "combine_f32f8", capability, stages, D, @@ -1375,7 +1391,7 @@ def combine( else: from quack.compile_utils import make_fake_tensor - kernel = SparseAttentionForwardCombine( + kernel = SparseAttentionForwardCombineF32F8( dtype=out_dtype, dtype_partial=partial_dtype, head_dim=D, @@ -1510,3 +1526,2308 @@ def combine( output_scale, qheadperkv, ) + + +# ----------------------------------------------------------------------------- +# v1 combine (combine, the default): BF16/FP16 partials +# ----------------------------------------------------------------------------- + + +class SparseAttentionForwardCombine: + def __init__( + self, + dtype: Type[cutlass.Numeric], + dtype_partial: Type[cutlass.Numeric], + head_dim: int, + tile_m: int = 8, + k_block_size: int = 64, + topk: int = 16, + num_threads: int = 256, + stages: int = 4, + use_pdl: bool = False, + min_blocks_per_mp: int = 0, + raw_partial_stats: bool = False, + ): + """ + Forward combine kernel for split attention computation. + + :param dtype: output data type + :param dtype_partial: partial accumulation data type + :param head_dim: head dimension + :param tile_m: m block size + :param k_block_size: k block size + :param topk: exact number of split partials + :param num_threads: number of threads + :param stages: number of pipeline stages + """ + self.dtype = dtype + self.dtype_partial = dtype_partial + self.head_dim = head_dim + self.tile_m = tile_m + self.k_block_size = k_block_size + self.topk = topk + self.num_threads = num_threads + self.is_even_k = head_dim % k_block_size == 0 + self.stages = stages + self.use_pdl = use_pdl + self.min_blocks_per_mp = min_blocks_per_mp + self.raw_partial_stats = raw_partial_stats + self.use_tma_partial = ( + raw_partial_stats + and dtype == cutlass.BFloat16 + and dtype_partial == cutlass.BFloat16 + and head_dim == 128 + and tile_m in (32, 64) + and k_block_size == 128 + and topk == 16 + and num_threads == tile_m * 4 + and stages == 2 + ) + + @staticmethod + def can_implement( + dtype, + dtype_partial, + head_dim, + tile_m, + k_block_size, + topk, + num_threads, + ) -> bool: + """Check if the kernel can be implemented with the given parameters.""" + if dtype not in [cutlass.Float16, cutlass.BFloat16, cutlass.Float32]: + return False + if dtype_partial not in [cutlass.Float16, cutlass.BFloat16]: + return False + if head_dim % 8 != 0: + return False + if num_threads % 32 != 0: + return False + if tile_m % 8 != 0: + return False + if topk > 256: + return False + if (tile_m * topk) % num_threads != 0: + return False + return True + + def _setup_attributes(self): + # GMEM copy setup for O partial + universal_copy_bits = 128 + async_copy_elems = universal_copy_bits // self.dtype_partial.width + assert self.k_block_size % async_copy_elems == 0 + + k_block_gmem = ( + 128 if self.k_block_size % 128 == 0 else (64 if self.k_block_size % 64 == 0 else 32) + ) + gmem_threads_per_row = k_block_gmem // async_copy_elems + assert self.num_threads % gmem_threads_per_row == 0 + + # Async copy atom for O partial load + atom_async_copy_partial = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), + self.dtype_partial, + num_bits_per_copy=universal_copy_bits, + ) + tOpartial_layout = cute.make_ordered_layout( + (self.num_threads // gmem_threads_per_row, gmem_threads_per_row), + order=(1, 0), + ) + vOpartial_layout = cute.make_layout((1, async_copy_elems)) + self.gmem_tiled_copy_O_partial = cute.make_tiled_copy_tv( + atom_async_copy_partial, tOpartial_layout, vOpartial_layout + ) + + # GMEM copy setup for final O (use universal copy for store). + # Keep output stores independent from the partial-load copy atom. + output_copy_elems = universal_copy_bits // self.dtype.width + assert self.k_block_size % output_copy_elems == 0 + gmem_threads_per_row_o = k_block_gmem // output_copy_elems + assert self.num_threads % gmem_threads_per_row_o == 0 + atom_universal_copy = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.dtype, + num_bits_per_copy=universal_copy_bits, + ) + tO_layout = cute.make_ordered_layout( + (self.num_threads // gmem_threads_per_row_o, gmem_threads_per_row_o), + order=(1, 0), + ) + vO_layout = cute.make_layout((1, output_copy_elems)) + self.gmem_tiled_copy_O = cute.make_tiled_copy_tv( + atom_universal_copy, + tO_layout, + vO_layout, + ) + # LSE copy setup with async copy (alignment = 1) + lse_copy_bits = Float32.width # 1 element per copy, width is in bits + m_block_smem = ( + 128 + if self.tile_m % 128 == 0 + else ( + 64 + if self.tile_m % 64 == 0 + else (32 if self.tile_m % 32 == 0 else (16 if self.tile_m % 16 == 0 else 8)) + ) + ) + gmem_threads_per_row_lse = m_block_smem + assert self.num_threads % gmem_threads_per_row_lse == 0 + + # Async copy atom for LSE load + atom_async_copy_lse = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.ALWAYS), + Float32, + num_bits_per_copy=lse_copy_bits, + ) + tLSE_layout = cute.make_ordered_layout( + (self.num_threads // gmem_threads_per_row_lse, gmem_threads_per_row_lse), + order=(1, 0), + ) + vLSE_layout = cute.make_layout(1) + self.gmem_tiled_copy_LSE = cute.make_tiled_copy_tv( + atom_async_copy_lse, tLSE_layout, vLSE_layout + ) + + # /////////////////////////////////////////////////////////////////////////////// + # Shared memory + # /////////////////////////////////////////////////////////////////////////////// + + # Shared memory to register copy for LSE + self.smem_threads_per_col_lse = self.num_threads // m_block_smem + assert 32 % self.smem_threads_per_col_lse == 0 # Must divide warp size + + s2r_layout_atom_lse = cute.make_ordered_layout( + (self.smem_threads_per_col_lse, self.num_threads // self.smem_threads_per_col_lse), + order=(0, 1), + ) + self.s2r_tiled_copy_LSE = cute.make_tiled_copy_tv( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), Float32), + s2r_layout_atom_lse, + cute.make_layout(1), + ) + + # LSE shared memory layout with swizzling to avoid bank conflicts + # This works for kBlockMSmem = 8, 16, 32, 64, 128, no bank conflicts + if const_expr(m_block_smem == 8): + smem_lse_swizzle = cute.make_swizzle(5, 0, 5) + elif const_expr(m_block_smem == 16): + smem_lse_swizzle = cute.make_swizzle(4, 0, 4) + else: + smem_lse_swizzle = cute.make_swizzle(3, 2, 3) + lse_atom_splits = min(self.topk, 8) + smem_layout_atom_lse = cute.make_composed_layout( + smem_lse_swizzle, + 0, + cute.make_ordered_layout((lse_atom_splits, m_block_smem), order=(1, 0)), + ) + self.smem_layout_lse = cute.tile_to_shape( + smem_layout_atom_lse, (self.topk, self.tile_m), (0, 1) + ) + + # O_partial staging layout. The target raw-BF16 specialization uses a + # row-major tile that is shared by the TMA and consumer-thread views. + if const_expr(self.use_tma_partial): + self.smem_layout_o = cute.make_ordered_layout( + (self.tile_m, self.k_block_size, self.stages), + order=(1, 0, 2), + ) + self.smem_layout_o_tma = cute.make_ordered_layout( + (self.k_block_size, self.tile_m, self.stages), + order=(0, 1, 2), + ) + else: + smem_layout_atom_o = _get_cpasync_smem_layout_atom( + self.dtype_partial, self.k_block_size + ) + self.smem_layout_o = cute.tile_to_shape( + smem_layout_atom_o, + (self.tile_m, self.k_block_size, self.stages), + (0, 1, 2), + ) + self.smem_layout_o_tma = None + + @cute.jit + def __call__( + self, + mO_partial: cute.Tensor, + mLSE_partial: cute.Tensor, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor] = None, + mLSE_temperature_partial: Optional[cute.Tensor] = None, + mLSE_temperature: Optional[cute.Tensor] = None, + cu_seqlens: Optional[cute.Tensor] = None, + seqused: Optional[cute.Tensor] = None, + num_splits_dynamic_ptr: Optional[cute.Tensor] = None, + varlen_batch_idx: Optional[cute.Tensor] = None, + semaphore_to_reset: Optional[cute.Tensor] = None, + mSplitCounts: Optional[cute.Tensor] = None, + mOutputScale: Optional[cute.Tensor] = None, + qhead_per_kvhead: Int32 = Int32(1), + # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI). + stream: cuda.CUstream = None, + ): + # Type checking + if const_expr(not (mO_partial.element_type == self.dtype_partial)): + raise TypeError("O partial tensor must match dtype_partial") + if const_expr(not (mO.element_type == self.dtype)): + raise TypeError("O tensor must match dtype") + if const_expr(mLSE_partial.element_type not in [Float32]): + raise TypeError("LSE partial tensor must be Float32") + if const_expr(mLSE is not None and mLSE.element_type not in [Float32]): + raise TypeError("LSE tensor must be Float32") + if const_expr( + mLSE_temperature_partial is not None + and mLSE_temperature_partial.element_type not in [Float32] + ): + raise TypeError("temperature LSE partial tensor must be Float32") + if const_expr(mLSE_temperature is not None and mLSE_temperature.element_type not in [Float32]): + raise TypeError("temperature LSE tensor must be Float32") + if const_expr((mLSE_temperature_partial is None) != (mLSE_temperature is None)): + raise ValueError( + "temperature LSE partial and output tensors must either both be provided or both be None" + ) + + # Shape validation - input tensors are in user format, need to be converted to kernel format + if const_expr(len(mO_partial.shape) not in [4, 5]): + raise ValueError( + "O partial tensor must have 4 or 5 dimensions: (num_splits, batch, seqlen, nheads, headdim) or (num_splits, total_q, nheads, headdim)" + ) + if const_expr(self.raw_partial_stats): + if const_expr(len(mLSE_partial.shape) != 4): + raise ValueError( + "Raw partial stats must have shape " + "(num_splits, total_q, nheads, 2)" + ) + if const_expr(cu_seqlens is None): + raise ValueError("Raw partial stats currently require the varlen combine path") + if const_expr(mLSE_temperature_partial is not None): + raise ValueError("Raw partial stats do not support independent temperature LSE") + elif const_expr(len(mLSE_partial.shape) not in [3, 4]): + raise ValueError( + "LSE partial tensor must have 3 or 4 dimensions: (num_splits, batch, seqlen, nheads) or (num_splits, total_q, nheads)" + ) + if const_expr(len(mO.shape) not in [3, 4]): + raise ValueError( + "O tensor must have 3 or 4 dimensions: (batch, seqlen, nheads, headdim) or (total_q, nheads, headdim)" + ) + if const_expr(mLSE is not None and len(mLSE.shape) not in [2, 3]): + raise ValueError( + "LSE tensor must have 2 or 3 dimensions: (batch, seqlen, nheads) or (total_q, nheads)" + ) + if const_expr(mLSE_temperature_partial is not None and len(mLSE_temperature_partial.shape) not in [3, 4]): + raise ValueError( + "temperature LSE partial tensor must have 3 or 4 dimensions: " + "(num_splits, batch, seqlen, nheads) or (num_splits, total_q, nheads)" + ) + if const_expr(mLSE_temperature is not None and len(mLSE_temperature.shape) not in [2, 3]): + raise ValueError( + "temperature LSE tensor must have 2 or 3 dimensions: " + "(batch, seqlen, nheads) or (total_q, nheads)" + ) + if const_expr(mSplitCounts is not None): + if const_expr(mSplitCounts.element_type not in [Int32]): + raise TypeError("split_counts tensor must be Int32") + if const_expr(cu_seqlens is not None): + if const_expr(len(mSplitCounts.shape) != 2): + raise ValueError("varlen split_counts tensor must have shape (total_q, nheads_kv)") + elif const_expr(len(mSplitCounts.shape) != 3): + raise ValueError("batched split_counts tensor must have shape (batch, seqlen, nheads_kv)") + if const_expr(mOutputScale is not None and mOutputScale.element_type not in [Float32]): + raise TypeError("output_scale tensor must be Float32") + + mO_partial, mO = [assume_tensor_aligned(t) for t in (mO_partial, mO)] + # (num_splits, b, seqlen, h, d) -> (seqlen, d, num_splits, h, b) + # or (num_splits, total_q, h, d) -> (total_q, d, num_splits, h) + O_partial_layout_transpose = ( + [2, 4, 0, 3, 1] if const_expr(cu_seqlens is None) else [1, 3, 0, 2] + ) + # (b, seqlen, h, d) -> (seqlen, d, h, b) or (total_q, h, d) -> (total_q, d, h) + mO_partial = cute.make_tensor( + mO_partial.iterator, cute.select(mO_partial.layout, mode=O_partial_layout_transpose) + ) + O_layout_transpose = [1, 3, 2, 0] if const_expr(cu_seqlens is None) else [0, 2, 1] + mO = cute.make_tensor(mO.iterator, cute.select(mO.layout, mode=O_layout_transpose)) + # (num_splits, b, h, seqlen) -> (seqlen, num_splits, h, b) + # Input is pre-transposed: [topK, B, Hq, Sq] with Sq innermost for K2-friendly reads. + # or (num_splits, total_q, h) -> (total_q, num_splits, h) + LSE_partial_layout_transpose = [3, 0, 2, 1] if const_expr(cu_seqlens is None) else [1, 0, 2] + mRawStats = None + mRowSum_partial = None + if const_expr(self.raw_partial_stats): + # [split, q, head, stat] -> [q, split, head, stat]. Keep the + # interleaved pair ABI while exposing scalar component views to K2. + mPartialStats = cute.make_tensor( + mLSE_partial.iterator, + cute.select(mLSE_partial.layout, mode=[1, 0, 2, 3]), + ) + mRawStats = mPartialStats + mLSE_partial = mPartialStats[None, None, None, 0] + mRowSum_partial = mPartialStats[None, None, None, 1] + else: + mLSE_partial = cute.make_tensor( + mLSE_partial.iterator, + cute.select(mLSE_partial.layout, mode=LSE_partial_layout_transpose), + ) + # (b, seqlen, h) -> (seqlen, h, b) or (total_q, h) -> (total_q, h) + LSE_layout_transpose = [1, 2, 0] if const_expr(cu_seqlens is None) else [0, 1] + mLSE = ( + cute.make_tensor(mLSE.iterator, cute.select(mLSE.layout, mode=LSE_layout_transpose)) + if mLSE is not None + else None + ) + mLSE_temperature_partial = ( + cute.make_tensor( + mLSE_temperature_partial.iterator, + cute.select(mLSE_temperature_partial.layout, mode=LSE_partial_layout_transpose), + ) + if mLSE_temperature_partial is not None + else None + ) + mLSE_temperature = ( + cute.make_tensor( + mLSE_temperature.iterator, + cute.select(mLSE_temperature.layout, mode=LSE_layout_transpose), + ) + if mLSE_temperature is not None + else None + ) + + # Determine if we have variable length sequences + varlen = const_expr(cu_seqlens is not None or seqused is not None) + + self._setup_attributes() + + # Flatten [split, q, head] into one TMA row coordinate, matching the + # CUDA combine descriptor. A separate split mode would require a + # D * total_q * heads stride; that stride reaches exactly 2**32 + # elements for the 512K target and wraps them in a CuTe DSL descriptor + # lowering. The flattened row extent remains representable and keeps + # every 32x128 transfer contiguous. + tma_atom_O_partial = None + mO_partial_tma = None + if const_expr(self.use_tma_partial): + total_flat = mO_partial.shape[0] * mO_partial.shape[3] + mO_partial_tma = cute.make_tensor( + mO_partial.iterator, + cute.make_layout( + (self.k_block_size, total_flat * self.topk), + stride=( + 1, + self.k_block_size, + ), + ), + ) + tma_atom_O_partial, mO_partial_tma = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + mO_partial_tma, + cute.select(self.smem_layout_o_tma, mode=[0, 1]), + (self.k_block_size, self.tile_m), + ) + + raw_stats_copy_atom = None + if const_expr(self.use_tma_partial): + raw_stats_copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + Float32, + num_bits_per_copy=64, + ) + + smem_layout_lse_aux = ( + cute.make_layout((1,), stride=(1,)) + if const_expr(self.use_tma_partial) + else self.smem_layout_lse + ) + + # BF16/F16 output uses the register-shuffle epilogue and does not need + # the legacy shared-memory permutation buffer. Keep a one-element + # placeholder so one SharedStorage type can serve all specializations. + if const_expr(self.dtype in [cutlass.Float16, cutlass.BFloat16]): + smem_layout_perm = cute.make_layout((1,), stride=(1,)) + else: + smem_layout_perm = cute.make_ordered_layout( + (self.tile_m, self.k_block_size), order=(1, 0) + ) + + @cute.struct + class SharedStorage: + sLSE: cute.struct.Align[ + cute.struct.MemRange[Float32, cute.cosize(self.smem_layout_lse)], 128 + ] + sLSETemperature: cute.struct.Align[ + cute.struct.MemRange[Float32, cute.cosize(smem_layout_lse_aux)], 128 + ] + sMaxValidSplit: cute.struct.Align[cute.struct.MemRange[Int32, self.tile_m], 128] + sO: cute.struct.Align[ + cute.struct.MemRange[self.dtype_partial, cute.cosize(self.smem_layout_o)], 128 + ] + sO_perm: cute.struct.Align[ + cute.struct.MemRange[self.dtype, cute.cosize(smem_layout_perm)], 128 + ] + sOTmaMbar: cute.struct.Align[ + cute.struct.MemRange[Int64, self.stages], 128 + ] + smem_size = SharedStorage.size_in_bytes() + + @cute.struct + class TmaRawSharedStorage: + sOTmaMbar: cute.struct.Align[ + cute.struct.MemRange[Int64, self.stages], 128 + ] + sLSE: cute.struct.Align[ + cute.struct.MemRange[Float32, cute.cosize(self.smem_layout_lse)], 128 + ] + sMaxValidSplit: cute.struct.Align[ + cute.struct.MemRange[Int32, self.tile_m], 128 + ] + sO: cute.struct.Align[ + cute.struct.MemRange[ + self.dtype_partial, + cute.cosize(self.smem_layout_o), + ], + 128, + ] + tma_raw_smem_size = TmaRawSharedStorage.size_in_bytes() + + # Flatten concatenated varlen rows into one scheduling domain. Each + # combine row is independent, so this preserves the varlen layout + # while avoiding one total_q-sized grid per batch. + flatten_varlen = const_expr( + cu_seqlens is not None + and seqused is None + and varlen_batch_idx is None + ) + + # Grid: (ceil(rows * heads / tile_m), ceil(dim / k_block), domains) + # Head separated from rows enables future TMA on contiguous row tiles. + seqlen = mO_partial.shape[0] + num_head = mO_partial.shape[3] + batch_size = ( + Int32(1) + if const_expr(flatten_varlen) + else ( + mO_partial.shape[4] + if const_expr(cu_seqlens is None) + else Int32(cu_seqlens.shape[0] - 1) + ) + ) + + seqlen_divmod = FastDivmodDivisor(seqlen) + head_divmod = FastDivmodDivisor(num_head) + + grid_dim = ( + cute.ceil_div(seqlen * num_head, self.tile_m), + cute.ceil_div(self.head_dim, self.k_block_size), + batch_size, + ) + + if const_expr(self.use_tma_partial): + if const_expr(not flatten_varlen): + raise ValueError("TMA raw combine requires flattened varlen tensors") + self.kernel_tma_raw( + mO_partial_tma, + mRawStats, + mO, + mLSE, + semaphore_to_reset, + mSplitCounts, + mOutputScale, + qhead_per_kvhead, + tma_atom_O_partial, + raw_stats_copy_atom, + TmaRawSharedStorage, + self.smem_layout_lse, + self.smem_layout_o, + self.smem_layout_o_tma, + self.gmem_tiled_copy_O_partial, + self.gmem_tiled_copy_O, + head_divmod, + self.use_pdl, + ).launch( + grid=grid_dim, + block=[self.num_threads, 1, 1], + smem=tma_raw_smem_size, + stream=stream, + min_blocks_per_mp=self.min_blocks_per_mp, + use_pdl=self.use_pdl, + ) + else: + self.kernel( + mO_partial, + mO_partial_tma, + mLSE_partial, + mRowSum_partial, + mRawStats, + mO, + mLSE, + mLSE_temperature_partial, + mLSE_temperature, + cu_seqlens, + seqused, + num_splits_dynamic_ptr, + varlen_batch_idx, + semaphore_to_reset, + mSplitCounts, + mOutputScale, + qhead_per_kvhead, + tma_atom_O_partial, + raw_stats_copy_atom, + SharedStorage, + self.smem_layout_lse, + smem_layout_lse_aux, + self.smem_layout_o, + self.smem_layout_o_tma, + smem_layout_perm, + self.gmem_tiled_copy_O_partial, + self.gmem_tiled_copy_O, + self.gmem_tiled_copy_LSE, + self.s2r_tiled_copy_LSE, + seqlen_divmod, + head_divmod, + self.use_pdl, + varlen, + flatten_varlen, + ).launch( + grid=grid_dim, + block=[self.num_threads, 1, 1], + smem=smem_size, + stream=stream, + min_blocks_per_mp=self.min_blocks_per_mp, + use_pdl=self.use_pdl, + ) + + @cute.jit + def decode_flat_row_idx( + self, + idx: Int32, + head_divmod: FastDivmodDivisor, + ): + """Decode flattened tile rows under the H_q-innermost contract.""" + q_idx_local, head_idx = divmod(idx, head_divmod) + return q_idx_local, head_idx + + @cute.kernel + def kernel_tma_raw( + self, + mO_partial_tma: cute.Tensor, + mRawStats: cute.Tensor, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + semaphore_to_reset: Optional[cute.Tensor], + mSplitCounts: Optional[cute.Tensor], + mOutputScale: Optional[cute.Tensor], + qhead_per_kvhead: Int32, + tma_atom_O_partial: cute.CopyAtom, + raw_stats_copy_atom: cute.CopyAtom, + SharedStorage: cutlass.Constexpr, + smem_layout_lse: cute.Layout | cute.ComposedLayout, + smem_layout_o: cute.Layout | cute.ComposedLayout, + smem_layout_o_tma: cute.Layout, + gmem_tiled_copy_O_partial: cute.TiledCopy, + gmem_tiled_copy_O: cute.TiledCopy, + head_divmod: FastDivmodDivisor, + use_pdl: cutlass.Constexpr[bool], + ): + """CUDA-shaped raw-stat combine specialization for the SM100 target.""" + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + m_block, k_block, maybe_virtual_batch = cute.arch.block_idx() + + if const_expr(use_pdl and semaphore_to_reset is not None): + if ( + tidx == 0 + and m_block == cute.arch.grid_dim()[0] - 1 + and k_block == cute.arch.grid_dim()[1] - 1 + and maybe_virtual_batch == cute.arch.grid_dim()[2] - 1 + ): + cute.arch.griddepcontrol_wait() + semaphore_to_reset[0] = 0 + if const_expr(use_pdl): + cute.arch.griddepcontrol_wait() + + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + sLSE = storage.sLSE.get_tensor(smem_layout_lse) + sMaxValidSplit = storage.sMaxValidSplit.get_tensor((self.tile_m,)) + sO = storage.sO.get_tensor(smem_layout_o) + sO_tma = storage.sO.get_tensor(smem_layout_o_tma) + sO_tma_mbar_ptr = storage.sOTmaMbar.data_ptr() + + total_flat = mRawStats.shape[0] * mRawStats.shape[2] + tile_begin = m_block * self.tile_m + output_scale = Float32(1.0) + if const_expr(mOutputScale is not None): + output_scale = mOutputScale[0] + + # Initialize the two transaction barriers and overlap the initial TMA + # loads with raw-stat reduction, matching the CUDA schedule. + gO_partial_tma = cute.local_tile( + mO_partial_tma, + (self.k_block_size, self.tile_m), + (k_block, None), + ) + tOsO_tma, tOgO_tma = cpasync.tma_partition( + tma_atom_O_partial, + 0, + cute.make_layout(1), + cute.group_modes(sO_tma, 0, 2), + cute.group_modes(gO_partial_tma, 0, 2), + ) + o_tma_copy_bytes = const_expr( + cute.size_in_bytes( + self.dtype_partial, + cute.select(smem_layout_o_tma, mode=[0, 1]), + ) + ) + row_tiles_per_split = total_flat // Int32(self.tile_m) + if tidx == 0: + for stage in cutlass.range(self.stages, unroll_full=True): + cute.arch.mbarrier_init(sO_tma_mbar_ptr + stage, 1) + cute.arch.mbarrier_init_fence() + cute.arch.sync_threads() + if warp_idx == 0: + for stage in cutlass.range(self.stages, unroll_full=True): + o_tma_barrier = sO_tma_mbar_ptr + stage + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + o_tma_barrier, + o_tma_copy_bytes, + ) + cute.copy( + tma_atom_O_partial, + tOgO_tma[ + None, + m_block + Int32(stage) * row_tiles_per_split, + ], + tOsO_tma[None, stage], + tma_bar_ptr=o_tma_barrier, + ) + + # Four adjacent lanes own one row and four interleaved split pairs. + reduce_lane = tidx % Int32(4) + reduce_row = tidx // Int32(4) + reduce_flat = tile_begin + reduce_row + row_valid = reduce_flat < total_flat + m_idx = Int32(0) + head_idx = Int32(0) + row_split_count = Int32(0) + if row_valid: + m_idx, head_idx = self.decode_flat_row_idx( + reduce_flat, + head_divmod, + ) + row_split_count = ( + mSplitCounts[m_idx, head_idx // qhead_per_kvhead] + if const_expr(mSplitCounts is not None) + else Int32(self.topk) + ) + + local_page_max = cute.make_rmem_tensor((4,), Float32) + local_page_sum = cute.make_rmem_tensor((4,), Float32) + raw_stats = cute.make_rmem_tensor((2,), Float32) + local_max = -Float32.inf + local_max_valid = Int32(-1) + for item in cutlass.range(4, unroll_full=True): + split = reduce_lane + item * Int32(4) + raw_stats[0] = -Float32.inf + raw_stats[1] = Float32(0.0) + if row_valid and split < row_split_count: + cute.copy( + raw_stats_copy_atom, + mRawStats[m_idx, split, head_idx, None], + raw_stats, + ) + stats_valid = ( + raw_stats[1] > Float32(0.0) + and raw_stats[1] < Float32.inf + and raw_stats[0] > -Float32.inf + and raw_stats[0] < Float32.inf + and raw_stats[0] == raw_stats[0] + and raw_stats[1] == raw_stats[1] + ) + page_max = raw_stats[0] if stats_valid else -Float32.inf + page_sum = raw_stats[1] if stats_valid else Float32(0.0) + local_page_max[item] = page_max + local_page_sum[item] = page_sum + local_max = max(local_max, page_max) + if stats_valid: + local_max_valid = split + + stats_max = cute.arch.warp_reduction_max( + local_max, + threads_in_group=4, + ) + max_valid_split = cute.arch.warp_reduction_max( + local_max_valid, + threads_in_group=4, + ) + stats_max_safe = ( + stats_max if max_valid_split >= Int32(0) else Float32(0.0) + ) + stats_sum = Float32(0.0) + for item in cutlass.range(4, unroll_full=True): + split_scale = cute.math.exp2( + local_page_max[item] - stats_max_safe, + fastmath=True, + ) + local_page_max[item] = split_scale + stats_sum += split_scale * local_page_sum[item] + stats_sum = cute.arch.warp_reduction_sum( + stats_sum, + threads_in_group=4, + ) + stats_sum_valid = ( + max_valid_split >= Int32(0) + and stats_sum > Float32(0.0) + and stats_sum < Float32.inf + and stats_sum == stats_sum + ) + inv_sum = Float32(0.0) + final_lse = -Float32.inf + if stats_sum_valid: + inv_sum = Float32(1.0) / stats_sum + final_lse = ( + stats_max + cute.math.log2(stats_sum, fastmath=True) + ) * Float32(math.log(2.0)) + for item in cutlass.range(4, unroll_full=True): + split = reduce_lane + item * Int32(4) + sLSE[split, reduce_row] = local_page_max[item] * inv_sum + if reduce_lane == Int32(0): + sMaxValidSplit[reduce_row] = max_valid_split + if const_expr(mLSE is not None): + if row_valid: + mLSE[m_idx, head_idx] = final_lse + + cute.arch.sync_threads() + block_max_valid_split = Int32(-1) + if const_expr(self.tile_m == 32): + if tidx < Int32(self.tile_m): + block_max_valid_split = cute.arch.warp_reduction_max( + sMaxValidSplit[tidx], + threads_in_group=self.tile_m, + ) + if tidx == 0: + sMaxValidSplit[0] = block_max_valid_split + else: + if tidx < Int32(32): + block_max_valid_split = cute.arch.warp_reduction_max( + max( + sMaxValidSplit[tidx], + sMaxValidSplit[tidx + Int32(32)], + ), + threads_in_group=32, + ) + if tidx == 0: + sMaxValidSplit[0] = block_max_valid_split + cute.arch.sync_threads() + block_max_valid_split = sMaxValidSplit[0] + + cO = cute.make_identity_tensor((self.tile_m, self.k_block_size)) + partial_copy = gmem_tiled_copy_O_partial.get_slice(tidx) + tOcO = partial_copy.partition_D(cO) + tOsO_partial = partial_copy.partition_D(sO) + num_rows = const_expr(cute.size(tOcO, mode=[1])) + tOrO_partial_row = cute.make_rmem_tensor_like( + tOsO_partial[None, 0, None, 0] + ) + tOrO = cute.make_rmem_tensor_like( + tOsO_partial[None, None, None, 0], + Float32, + ) + tOrO.fill(0.0) + + for split in cutlass.range(block_max_valid_split + 1, unroll=4): + stage = split % Int32(self.stages) + phase = (split // Int32(self.stages)) & Int32(1) + cute.arch.mbarrier_wait(sO_tma_mbar_ptr + stage, phase) + for row in cutlass.range(num_rows, unroll_full=True): + cute.autovec_copy( + tOsO_partial[None, row, None, stage], + tOrO_partial_row, + ) + weight = sLSE[split, tOcO[0, row, 0][0]] + if weight != Float32(0.0): + tOrO[None, row, None].store( + tOrO[None, row, None].load() + + weight + * tOrO_partial_row.load().to(Float32) + ) + + cute.arch.sync_threads() + split_to_load = split + Int32(self.stages) + if warp_idx == 0 and split_to_load <= block_max_valid_split: + o_tma_barrier = sO_tma_mbar_ptr + stage + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + o_tma_barrier, + o_tma_copy_bytes, + ) + cute.copy( + tma_atom_O_partial, + tOgO_tma[ + None, + m_block + split_to_load * row_tiles_per_split, + ], + tOsO_tma[None, stage], + tma_bar_ptr=o_tma_barrier, + ) + + # Transpose the four 32-bit BF16 pairs within each four-lane group, + # then retain one coalesced 128-bit output store per row. + output_copy = gmem_tiled_copy_O.get_slice(tidx) + tOcO_store = output_copy.partition_D(cO) + rO_row = cute.make_rmem_tensor( + tOcO_store[None, 0, 0].shape, + self.dtype, + ) + elems_per_store = const_expr( + cute.size(gmem_tiled_copy_O.layout_tv_tiled[1]) + ) + num_store_rows = const_expr(cute.size(tOcO_store, mode=[1])) + lane_in_pair_group = tidx % Int32(4) + for row in cutlass.range(num_store_rows, unroll_full=True): + for k in cutlass.range( + cute.size(tOcO_store, mode=[2]), + unroll_full=True, + ): + o0, o1 = tOrO[0, row, k], tOrO[1, row, k] + o2, o3 = tOrO[2, row, k], tOrO[3, row, k] + o4, o5 = tOrO[4, row, k], tOrO[5, row, k] + o6, o7 = tOrO[6, row, k], tOrO[7, row, k] + if const_expr(mOutputScale is not None): + o0, o1 = cute.arch.mul_packed_f32x2( + (o0, o1), + (output_scale, output_scale), + ) + o2, o3 = cute.arch.mul_packed_f32x2( + (o2, o3), + (output_scale, output_scale), + ) + o4, o5 = cute.arch.mul_packed_f32x2( + (o4, o5), + (output_scale, output_scale), + ) + o6, o7 = cute.arch.mul_packed_f32x2( + (o6, o7), + (output_scale, output_scale), + ) + word0 = utils.cvt_f16x2_f32(o0, o1, self.dtype) + word1 = utils.cvt_f16x2_f32(o2, o3, self.dtype) + word2 = utils.cvt_f16x2_f32(o4, o5, self.dtype) + word3 = utils.cvt_f16x2_f32(o6, o7, self.dtype) + + partner0 = cute.arch.shuffle_sync_bfly(word1, offset=1) + partner1 = cute.arch.shuffle_sync_bfly(word0, offset=1) + partner2 = cute.arch.shuffle_sync_bfly(word3, offset=1) + partner3 = cute.arch.shuffle_sync_bfly(word2, offset=1) + stage0, stage1 = word0, word1 + stage2, stage3 = word2, word3 + if lane_in_pair_group % Int32(2) != Int32(0): + stage0, stage2 = partner0, partner2 + if lane_in_pair_group % Int32(2) != Int32(1): + stage1, stage3 = partner1, partner3 + + partner0 = cute.arch.shuffle_sync_bfly(stage2, offset=2) + partner1 = cute.arch.shuffle_sync_bfly(stage3, offset=2) + partner2 = cute.arch.shuffle_sync_bfly(stage0, offset=2) + partner3 = cute.arch.shuffle_sync_bfly(stage1, offset=2) + out0, out1 = stage0, stage1 + out2, out3 = stage2, stage3 + if lane_in_pair_group // Int32(2) != Int32(0): + out0, out1 = partner0, partner1 + if lane_in_pair_group // Int32(2) != Int32(1): + out2, out3 = partner2, partner3 + + rO_pairs = cute.recast_tensor(rO_row, cutlass.Int32) + rO_pairs[0] = out0 + rO_pairs[1] = out1 + rO_pairs[2] = out2 + rO_pairs[3] = out3 + row_local = tOcO_store[0, row, 0][0] + flat_idx = tile_begin + row_local + if flat_idx < total_flat: + m_idx, head_idx = self.decode_flat_row_idx( + flat_idx, + head_divmod, + ) + mO_cur_copy = cute.tiled_divide( + mO[m_idx, None, head_idx], + (elems_per_store,), + ) + k_idx = tOcO_store[0, 0, k][1] // elems_per_store + cute.copy( + output_copy, + rO_row, + mO_cur_copy[None, k_idx], + ) + + @cute.kernel + def kernel( + self, + mO_partial: cute.Tensor, + mO_partial_tma: Optional[cute.Tensor], + mLSE_partial: cute.Tensor, + mRowSum_partial: Optional[cute.Tensor], + mRawStats: Optional[cute.Tensor], + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + mLSE_temperature_partial: Optional[cute.Tensor], + mLSE_temperature: Optional[cute.Tensor], + cu_seqlens: Optional[cute.Tensor], + seqused: Optional[cute.Tensor], + num_splits_dynamic_ptr: Optional[cute.Tensor], + varlen_batch_idx: Optional[cute.Tensor], + semaphore_to_reset: Optional[cute.Tensor], + mSplitCounts: Optional[cute.Tensor], + mOutputScale: Optional[cute.Tensor], + qhead_per_kvhead: Int32, + tma_atom_O_partial: Optional[cute.CopyAtom], + raw_stats_copy_atom: Optional[cute.CopyAtom], + SharedStorage: cutlass.Constexpr, + smem_layout_lse: cute.Layout | cute.ComposedLayout, + smem_layout_lse_aux: cute.Layout | cute.ComposedLayout, + smem_layout_o: cute.Layout | cute.ComposedLayout, + smem_layout_o_tma: Optional[cute.Layout], + smem_layout_perm: cute.Layout, + gmem_tiled_copy_O_partial: cute.TiledCopy, + gmem_tiled_copy_O: cute.TiledCopy, + gmem_tiled_copy_LSE: cute.TiledCopy, + s2r_tiled_copy_LSE: cute.TiledCopy, + seqlen_divmod: FastDivmodDivisor, + head_divmod: FastDivmodDivisor, + use_pdl: cutlass.Constexpr[bool], + varlen: cutlass.Constexpr[bool], + flatten_varlen: cutlass.Constexpr[bool], + ): + # Thread and block indices + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + m_block, k_block, maybe_virtual_batch = cute.arch.block_idx() + + if const_expr(flatten_varlen): + batch_idx = Int32(0) + else: + batch_idx = ( + varlen_batch_idx[maybe_virtual_batch] + if const_expr(varlen_batch_idx is not None) + else maybe_virtual_batch + ) + + # /////////////////////////////////////////////////////////////////////////////// + # Get shared memory buffer + # /////////////////////////////////////////////////////////////////////////////// + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + sLSE = storage.sLSE.get_tensor(smem_layout_lse) + sLSE_temperature = storage.sLSETemperature.get_tensor(smem_layout_lse_aux) + sMaxValidSplit = storage.sMaxValidSplit.get_tensor((self.tile_m,)) + sO = storage.sO.get_tensor(smem_layout_o) + sO_tma = None + sO_tma_mbar_ptr = None + if const_expr(self.use_tma_partial): + sO_tma = storage.sO.get_tensor(smem_layout_o_tma) + sO_tma_mbar_ptr = storage.sOTmaMbar.data_ptr() + sO_perm_buf = storage.sO_perm.get_tensor(smem_layout_perm) + + # Handle semaphore reset — wait for dependent grids first + if const_expr(use_pdl and semaphore_to_reset is not None): + if ( + tidx == 0 + and m_block == cute.arch.grid_dim()[0] - 1 + and k_block == cute.arch.grid_dim()[1] - 1 + and maybe_virtual_batch == cute.arch.grid_dim()[2] - 1 + ): + cute.arch.griddepcontrol_wait() + semaphore_to_reset[0] = 0 + + if const_expr(num_splits_dynamic_ptr is not None): + raise ValueError("K2 combine requires compile-time exact topK") + num_splits = Int32(self.topk) + # Concatenated varlen tensors already have total_q as their leading + # extent. Other modes retain SeqlenInfo-based per-batch slicing. + if const_expr(flatten_varlen): + seqlen = Int32(mO_partial.shape[0]) + offset = Int32(0) + else: + seqlen_info = SeqlenInfo.create( + batch_idx=batch_idx, + seqlen_static=mO_partial.shape[0], + cu_seqlens=cu_seqlens, + seqused=seqused, + # Don't need tile size because offset_padded is unused. + ) + seqlen, offset = seqlen_info.seqlen, seqlen_info.offset + + num_head = mO_partial.shape[3] + max_idx = seqlen * num_head + output_scale = Float32(1.0) + if const_expr(mOutputScale is not None): + output_scale = mOutputScale[0] + + if const_expr(not varlen) or m_block * self.tile_m < max_idx: + # Wait for dependent grids (e.g., the main attention kernel that produces O_partial/LSE_partial) + if const_expr(use_pdl): + cute.arch.griddepcontrol_wait() + + # =============================== + # Step 1: Load LSE_partial from gmem to shared memory + # =============================== + # `cLSE` (identity tensor for row/split coord tracking) is reused + # later in steps 4-5, so it must be defined on both branches. + cLSE = cute.make_identity_tensor((self.topk, self.tile_m)) + # Reshape mLSE_partial to PackGQA packed layout and delegate the + # tile load to PackGQAComb.load_LSE. The packed form folds (H_q, Sq) + # into one compound dim with H_q innermost (stride 1), so thread + # rows that vary along h_pos produce one-sector coalesced reads. + # Non-varlen path only — varlen keeps the original inline loop. + if const_expr(not varlen): + mLSE_partial_cur = seqlen_info.offset_batch(mLSE_partial, batch_idx, dim=3) + # mLSE_partial_cur: (H_q, topK, Sq) — after initial transpose + # [3,0,2,1] on [topK,B,Sq,H_q] and dropping B. + # Reorder to (H_q, Sq, topK) then group modes 0..1 for packed dim: + mLSE_partial_reord = cute.make_tensor( + mLSE_partial_cur.iterator, + cute.select(mLSE_partial_cur.layout, mode=[0, 2, 1]), + ) + mLSE_partial_packed = cute.group_modes(mLSE_partial_reord, 0, 2) + # shape ((H_q, Sq), topK) with H_q innermost. + packgqa = PackGQAComb( + m_block_size=self.tile_m, + head_dim_padded=0, # unused for LSE load + check_hdim_oob=False, # unused for LSE load + qhead_per_kvhead=1, # unused; num_heads_divmod is passed explicitly + ) + packgqa.load_LSE( + mLSE_partial_packed, + sLSE, + self.topk, + gmem_tiled_copy_LSE, + tidx, + m_block, + num_splits, + seqlen, + head_divmod, + mSplitCounts, + batch_idx, + qhead_per_kvhead, + ) + if const_expr(mLSE_temperature_partial is not None): + mLSE_temperature_partial_cur = seqlen_info.offset_batch( + mLSE_temperature_partial, batch_idx, dim=3) + mLSE_temperature_partial_reord = cute.make_tensor( + mLSE_temperature_partial_cur.iterator, + cute.select(mLSE_temperature_partial_cur.layout, mode=[0, 2, 1]), + ) + mLSE_temperature_partial_packed = cute.group_modes( + mLSE_temperature_partial_reord, 0, 2) + packgqa.load_LSE( + mLSE_temperature_partial_packed, + sLSE_temperature, + self.topk, + gmem_tiled_copy_LSE, + tidx, + m_block, + num_splits, + seqlen, + head_divmod, + mSplitCounts, + batch_idx, + qhead_per_kvhead, + ) + else: + # Varlen path keeps the same H_q-innermost flat-row contract: + # after transpose [1, 0, 2], mLSE_partial_cur is + # (q_local, split, head). + # mSplitCounts is the authoritative valid-split count per + # packed (q_abs, kv_head); masked splits stay at -inf and + # therefore drop out of the final kernel LSE_out reduction. + if const_expr(flatten_varlen): + mLSE_partial_cur = mLSE_partial + if const_expr(self.raw_partial_stats): + mRowSum_partial_cur = mRowSum_partial + else: + mLSE_partial_cur = seqlen_info.offset_batch( + mLSE_partial, batch_idx, dim=3 + ) + if const_expr(self.raw_partial_stats): + mRowSum_partial_cur = seqlen_info.offset_batch( + mRowSum_partial, batch_idx, dim=3 + ) + mLSE_partial_copy = cute.tiled_divide(mLSE_partial_cur, (1,)) + if const_expr(self.raw_partial_stats): + mRowSum_partial_copy = cute.tiled_divide(mRowSum_partial_cur, (1,)) + gmem_thr_copy_LSE = gmem_tiled_copy_LSE.get_slice(tidx) + tLSEsLSE = gmem_thr_copy_LSE.partition_D(sLSE) + tLSEsLSE_temperature = gmem_thr_copy_LSE.partition_D(sLSE_temperature) + tLSEcLSE = gmem_thr_copy_LSE.partition_S(cLSE) + if const_expr(mLSE_temperature_partial is not None): + if const_expr(flatten_varlen): + mLSE_temperature_partial_cur = mLSE_temperature_partial + else: + mLSE_temperature_partial_cur = seqlen_info.offset_batch( + mLSE_temperature_partial, batch_idx, dim=3 + ) + mLSE_temperature_partial_copy = cute.tiled_divide( + mLSE_temperature_partial_cur, (1,)) + + for m in cutlass.range(cute.size(tLSEcLSE, mode=[2]), unroll_full=True): + mi = tLSEcLSE[0, 0, m][1] + idx = m_block * self.tile_m + mi + if idx < max_idx: + m_idx, head_idx = self.decode_flat_row_idx(idx, head_divmod) + row_count = ( + mSplitCounts[offset + m_idx, head_idx // qhead_per_kvhead] + if const_expr(mSplitCounts is not None) + else num_splits + ) + mLSE_partial_cur_copy = mLSE_partial_copy[None, m_idx, None, head_idx] + if const_expr(self.raw_partial_stats): + mRowSum_partial_cur_copy = ( + mRowSum_partial_copy[None, m_idx, None, head_idx] + ) + if const_expr(mLSE_temperature_partial is not None): + mLSE_temperature_partial_cur_copy = ( + mLSE_temperature_partial_copy[None, m_idx, None, head_idx]) + for s in cutlass.range(cute.size(tLSEcLSE, mode=[1]), unroll_full=True): + si = tLSEcLSE[0, s, 0][0] + if si < num_splits and si < row_count: + cute.copy( + gmem_thr_copy_LSE, + mLSE_partial_cur_copy[None, si], + tLSEsLSE[None, s, m], + ) + if const_expr(self.raw_partial_stats): + cute.copy( + gmem_thr_copy_LSE, + mRowSum_partial_cur_copy[None, si], + tLSEsLSE_temperature[None, s, m], + ) + if const_expr(mLSE_temperature_partial is not None): + cute.copy( + gmem_thr_copy_LSE, + mLSE_temperature_partial_cur_copy[None, si], + tLSEsLSE_temperature[None, s, m], + ) + else: + tLSEsLSE[None, s, m].fill(-Float32.inf) + if const_expr(self.raw_partial_stats): + tLSEsLSE_temperature[None, s, m].fill(Float32(0.0)) + if const_expr(mLSE_temperature_partial is not None): + tLSEsLSE_temperature[None, s, m].fill(-Float32.inf) + else: + for s in cutlass.range(cute.size(tLSEcLSE, mode=[1]), unroll_full=True): + tLSEsLSE[None, s, m].fill(-Float32.inf) + if const_expr(self.raw_partial_stats): + tLSEsLSE_temperature[None, s, m].fill(Float32(0.0)) + if const_expr(mLSE_temperature_partial is not None): + tLSEsLSE_temperature[None, s, m].fill(-Float32.inf) + cute.arch.cp_async_commit_group() + + # =============================== + # Step 2: Load O_partial for pipeline stages + # =============================== + + gmem_thr_copy_O_partial = gmem_tiled_copy_O_partial.get_slice(tidx) + cO = cute.make_identity_tensor((self.tile_m, self.k_block_size)) + tOcO = gmem_thr_copy_O_partial.partition_D(cO) + tOsO_partial = gmem_thr_copy_O_partial.partition_D(sO) + + # Precompute per-row values for flattened (q_local, head) tiles. + num_rows = const_expr(cute.size(tOcO, mode=[1])) + tOmidx = cute.make_rmem_tensor(num_rows, cutlass.Int32) + tOhidx = cute.make_rmem_tensor(num_rows, cutlass.Int32) + if const_expr(not self.use_tma_partial): + tOSplitCount = cute.make_rmem_tensor(num_rows, cutlass.Int32) + tOrOptr = cute.make_rmem_tensor(num_rows, cutlass.Int64) + if const_expr(flatten_varlen): + mO_partial_cur = mO_partial + else: + mO_partial_cur = seqlen_info.offset_batch( + mO_partial, batch_idx, dim=4 + ) + for m in cutlass.range(num_rows, unroll_full=True): + mi = tOcO[0, m, 0][0] # m coordinate in tile + idx = m_block * self.tile_m + mi + if idx >= max_idx: + tOhidx[m] = -1 + tOmidx[m] = 0 + if const_expr(not self.use_tma_partial): + tOSplitCount[m] = 0 + tOrOptr[m] = cutlass.Int64(0) + else: + tOmidx[m], tOhidx[m] = self.decode_flat_row_idx(idx, head_divmod) + if const_expr(not self.use_tma_partial): + if const_expr(mSplitCounts is None): + tOSplitCount[m] = num_splits + elif const_expr(cu_seqlens is None): + tOSplitCount[m] = mSplitCounts[ + batch_idx, tOmidx[m], tOhidx[m] // qhead_per_kvhead + ] + else: + tOSplitCount[m] = mSplitCounts[ + offset + tOmidx[m], tOhidx[m] // qhead_per_kvhead + ] + tOrOptr[m] = utils.elem_pointer( + mO_partial_cur, + (tOmidx[m], k_block * self.k_block_size, 0, tOhidx[m]), + ).toint() + + tOpO = None + if const_expr(not self.is_even_k): + tOpO = cute.make_rmem_tensor(cute.size(tOcO, mode=[2]), Boolean) + for k in cutlass.range(cute.size(tOpO), unroll_full=True): + tOpO[k] = tOcO[0, 0, k][1] < mO_partial.shape[1] - k_block * self.k_block_size + if const_expr(self.use_tma_partial): + gO_partial_tma = cute.local_tile( + mO_partial_tma, + (self.k_block_size, self.tile_m), + (k_block, m_block, None), + ) + tOsO_tma, tOgO_tma = cpasync.tma_partition( + tma_atom_O_partial, + 0, + cute.make_layout(1), + cute.group_modes(sO_tma, 0, 2), + cute.group_modes(gO_partial_tma, 0, 2), + ) + o_tma_copy_bytes = const_expr( + cute.size_in_bytes( + self.dtype_partial, + cute.select(smem_layout_o_tma, mode=[0, 1]), + ) + ) + if tidx == 0: + for stage in cutlass.range(self.stages, unroll_full=True): + cute.arch.mbarrier_init(sO_tma_mbar_ptr + stage, 1) + cute.arch.mbarrier_init_fence() + cute.arch.sync_threads() + if warp_idx == 0: + for stage in cutlass.range(self.stages, unroll_full=True): + o_tma_barrier = sO_tma_mbar_ptr + stage + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + o_tma_barrier, + o_tma_copy_bytes, + ) + cute.copy( + tma_atom_O_partial, + tOgO_tma[None, stage], + tOsO_tma[None, stage], + tma_bar_ptr=o_tma_barrier, + ) + else: + load_O_partial = partial( + self.load_O_partial, + gmem_tiled_copy_O_partial, + tOrOptr, + tOsO_partial, + tOhidx, + tOSplitCount, + tOpO, + tOcO, + mO_partial_cur.layout, + ) + + # Load the initial cp.async stages of O_partial. + for stage in cutlass.range(self.stages - 1, unroll_full=True): + if stage < num_splits: + load_O_partial(stage, stage) + cute.arch.cp_async_commit_group() + + # =============================== + # Step 3: Load and transpose LSE from smem to registers + # =============================== + + # Wait for LSE and initial O partial stages to complete + if const_expr(self.use_tma_partial): + cute.arch.cp_async_wait_group(0) + else: + cute.arch.cp_async_wait_group(self.stages - 1) + cute.arch.sync_threads() + + s2r_thr_copy_LSE = s2r_tiled_copy_LSE.get_slice(tidx) + ts2rsLSE = s2r_thr_copy_LSE.partition_S(sLSE) + ts2rrLSE = cute.make_rmem_tensor_like(ts2rsLSE) + cute.copy(s2r_tiled_copy_LSE, ts2rsLSE, ts2rrLSE) + if const_expr( + self.raw_partial_stats or mLSE_temperature_partial is not None + ): + ts2rsLSE_temperature = s2r_thr_copy_LSE.partition_S(sLSE_temperature) + ts2rrLSE_temperature = cute.make_rmem_tensor_like(ts2rsLSE_temperature) + cute.copy( + s2r_tiled_copy_LSE, + ts2rsLSE_temperature, + ts2rrLSE_temperature, + ) + + # =============================== + # Step 4: Compute final LSE along split dimension + # =============================== + + final_lse = cute.make_rmem_tensor(cute.size(ts2rrLSE, mode=[2]), Float32) + ts2rcLSE = s2r_thr_copy_LSE.partition_D(cLSE) + # We compute the max valid split for each row to short-circuit the computation later + max_valid_split = cute.make_rmem_tensor(cute.size(ts2rrLSE, mode=[2]), Int32) + assert cute.size(ts2rrLSE, mode=[0]) == 1 + # Compute max, scales, and final LSE for each row. Invalid splits + # have already been filled with -inf, so Step 5 can write the + # kernel-native LSE_out directly. + for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True): + threads_per_col = const_expr(self.smem_threads_per_col_lse) + if const_expr(self.raw_partial_stats): + # Sanitize the raw pair before reductions so NaN/Inf cannot + # contaminate a neighboring split in the thread group. + max_valid_idx = Int32(-1) + for s in cutlass.range( + cute.size(ts2rrLSE, mode=[1]), unroll_full=True + ): + scaled_max_log2 = ts2rrLSE[0, s, m] + row_sum = ts2rrLSE_temperature[0, s, m] + stats_valid = ( + row_sum > Float32(0.0) + and row_sum < Float32.inf + and scaled_max_log2 > -Float32.inf + and scaled_max_log2 < Float32.inf + and row_sum == row_sum + and scaled_max_log2 == scaled_max_log2 + ) + ts2rrLSE[0, s, m] = ( + scaled_max_log2 if stats_valid else -Float32.inf + ) + ts2rrLSE_temperature[0, s, m] = ( + row_sum if stats_valid else Float32(0.0) + ) + if stats_valid: + max_valid_idx = ts2rcLSE[0, s, 0][0] + + stats_max = cute.arch.warp_reduction_max( + ts2rrLSE[None, None, m] + .load() + .reduce( + cute.ReductionOp.MAX, + init_val=-Float32.inf, + reduction_profile=0, + ), + threads_in_group=threads_per_col, + ) + max_valid_split[m] = cute.arch.warp_reduction_max( + max_valid_idx, threads_in_group=threads_per_col + ) + stats_max_safe = ( + stats_max + if max_valid_split[m] >= Int32(0) + else Float32(0.0) + ) + stats_sum = Float32(0.0) + for s in cutlass.range( + cute.size(ts2rrLSE, mode=[1]), unroll_full=True + ): + row_sum = ts2rrLSE_temperature[0, s, m] + split_scale = Float32(0.0) + if row_sum > Float32(0.0): + split_scale = cute.math.exp2( + ts2rrLSE[0, s, m] - stats_max_safe, + fastmath=True, + ) + stats_sum += split_scale * row_sum + # Raw partial-O already includes the local row sum, so + # its combine weight excludes row_sum from the numerator. + ts2rrLSE[0, s, m] = split_scale + stats_sum = cute.arch.warp_reduction_sum( + stats_sum, threads_in_group=threads_per_col + ) + stats_sum_valid = ( + max_valid_split[m] >= Int32(0) + and stats_sum > Float32(0.0) + and stats_sum < Float32.inf + and stats_sum == stats_sum + ) + inv_sum = Float32(0.0) + if stats_sum_valid: + inv_sum = Float32(1.0) / stats_sum + final_lse[m] = ( + stats_max + + cute.math.log2(stats_sum, fastmath=True) + ) * Float32(math.log(2.0)) + else: + final_lse[m] = -Float32.inf + ts2rrLSE[None, None, m].store( + ts2rrLSE[None, None, m].load() * inv_sum + ) + else: + # Find max normalized-partial LSE value across splits. + lse_max = cute.arch.warp_reduction_max( + ts2rrLSE[None, None, m] + .load() + .reduce( + cute.ReductionOp.MAX, + init_val=-Float32.inf, + reduction_profile=0, + ), + threads_in_group=threads_per_col, + ) + max_valid_idx = -1 + for s in cutlass.range( + cute.size(ts2rrLSE, mode=[1]), unroll_full=True + ): + if ts2rrLSE[0, s, m] != -Float32.inf: + max_valid_idx = ts2rcLSE[0, s, 0][0] + max_valid_split[m] = cute.arch.warp_reduction_max( + max_valid_idx, threads_in_group=threads_per_col + ) + lse_max_cur = Float32(0.0) if lse_max == -Float32.inf else lse_max + LOG2_E = math.log2(math.e) + lse_sum_cur = Float32(0.0) + for s in cutlass.range( + cute.size(ts2rrLSE, mode=[1]), unroll_full=True + ): + scale = cute.math.exp2( + ts2rrLSE[0, s, m] * LOG2_E + - lse_max_cur * LOG2_E, + fastmath=True, + ) + lse_sum_cur += scale + ts2rrLSE[0, s, m] = scale + lse_sum_cur = cute.arch.warp_reduction_sum( + lse_sum_cur, threads_in_group=threads_per_col + ) + inv_sum = Float32(0.0) + if ( + max_valid_split[m] < 0 + or lse_sum_cur == Float32(0.0) + or lse_sum_cur != lse_sum_cur + ): + final_lse[m] = -Float32.inf + else: + final_lse[m] = ( + cute.math.log(lse_sum_cur, fastmath=True) + lse_max + ) + inv_sum = Float32(1.0) / lse_sum_cur + ts2rrLSE[None, None, m].store( + ts2rrLSE[None, None, m].load() * inv_sum + ) + # Store the scales exp(lse - lse_logsum) back to smem + cute.copy(s2r_tiled_copy_LSE, ts2rrLSE, ts2rsLSE) + + if const_expr(mLSE_temperature_partial is not None): + final_lse_temperature = cute.make_rmem_tensor( + cute.size(ts2rrLSE_temperature, mode=[2]), Float32) + for m in cutlass.range(cute.size(ts2rrLSE_temperature, mode=[2]), unroll_full=True): + threads_per_col = const_expr(self.smem_threads_per_col_lse) + lse_temperature_max = cute.arch.warp_reduction_max( + ts2rrLSE_temperature[None, None, m] + .load() + .reduce( + cute.ReductionOp.MAX, + init_val=-Float32.inf, + reduction_profile=0, + ), + threads_in_group=threads_per_col, + ) + lse_temperature_max_cur = ( + 0.0 if lse_temperature_max == -Float32.inf else lse_temperature_max + ) + LOG2_E = math.log2(math.e) + lse_temperature_sum_cur = 0.0 + for s in cutlass.range( + cute.size(ts2rrLSE_temperature, mode=[1]), unroll_full=True): + scale = cute.math.exp2( + ts2rrLSE_temperature[0, s, m] * LOG2_E + - (lse_temperature_max_cur * LOG2_E), + fastmath=True, + ) + lse_temperature_sum_cur += scale + lse_temperature_sum_cur = cute.arch.warp_reduction_sum( + lse_temperature_sum_cur, threads_in_group=threads_per_col + ) + if ( + max_valid_split[m] < 0 + or lse_temperature_sum_cur == 0.0 + or lse_temperature_sum_cur != lse_temperature_sum_cur + ): + final_lse_temperature[m] = -Float32.inf + else: + final_lse_temperature[m] = ( + cute.math.log(lse_temperature_sum_cur, fastmath=True) + + lse_temperature_max + ) + + # Store max valid split to smem + for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True): + if ts2rcLSE[0, 0, m][0] == 0: # Only thread responsible for s=0 writes + mi = ts2rcLSE[0, 0, m][1] + if mi < self.tile_m: + sMaxValidSplit[mi] = max_valid_split[m] + + # =============================== + # Step 5: Store final LSE to gmem + # This writeback is the authoritative LSE_out returned by the + # public Sparse Attention / Sparse Page Attention interface. + # =============================== + + if const_expr(mLSE is not None): + if const_expr(cu_seqlens is None): + mLSE_cur = mLSE[None, None, batch_idx] + else: + mLSE_cur = cute.domain_offset((offset, 0), mLSE) + if const_expr(mLSE_temperature is not None): + if const_expr(cu_seqlens is None): + mLSE_temperature_cur = mLSE_temperature[None, None, batch_idx] + else: + mLSE_temperature_cur = cute.domain_offset( + (offset, 0), mLSE_temperature) + if k_block == 0: # Only first k_block writes LSE when mLSE is provided + for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True): + if ts2rcLSE[0, 0, m][0] == 0: # Only thread responsible for s=0 writes + mi = ts2rcLSE[0, 0, m][1] + idx = m_block * self.tile_m + mi + if idx < max_idx: + m_idx, head_idx = self.decode_flat_row_idx(idx, head_divmod) + mLSE_cur[m_idx, head_idx] = final_lse[m] + if const_expr(mLSE_temperature is not None): + mLSE_temperature_cur[m_idx, head_idx] = ( + final_lse_temperature[m]) + + # =============================== + # Step 6: Read O_partial and accumulate final O + # =============================== + + cute.arch.sync_threads() + + if const_expr(self.use_tma_partial): + # TMA stages are CTA-wide resources, so every consumer thread + # must execute the same split count and barrier sequence. + block_max_valid_split = Int32(-1) + if tidx < Int32(self.tile_m): + block_max_valid_split = cute.arch.warp_reduction_max( + sMaxValidSplit[tidx], threads_in_group=self.tile_m + ) + if tidx == 0: + sMaxValidSplit[0] = block_max_valid_split + cute.arch.sync_threads() + thr_max_valid_split = sMaxValidSplit[0] + else: + # cp.async loads are thread-private, so each thread can stop at + # the largest valid split among its own rows. + thr_max_valid_split = sMaxValidSplit[tOcO[0, 0, 0][0]] + for m in cutlass.range( + 1, cute.size(tOcO, mode=[1]), unroll_full=True + ): + thr_max_valid_split = max( + thr_max_valid_split, + sMaxValidSplit[tOcO[0, m, 0][0]], + ) + + tOrO_partial = cute.make_rmem_tensor_like(tOsO_partial[None, None, None, 0]) + tOrO = cute.make_rmem_tensor_like(tOrO_partial, Float32) + tOrO.fill(0.0) + + stage_load = Int32(self.stages - 1) + stage_compute = Int32(0) + + # Main accumulation loop + for s in cutlass.range(thr_max_valid_split + 1, unroll=4): + if const_expr(self.use_tma_partial): + stage_compute = s % Int32(self.stages) + stage_phase = (s // Int32(self.stages)) & Int32(1) + cute.arch.mbarrier_wait( + sO_tma_mbar_ptr + stage_compute, + stage_phase, + ) + + # Get scales for this split + scale = cute.make_rmem_tensor(num_rows, Float32) + for m in cutlass.range(num_rows, unroll_full=True): + scale[m] = sLSE[s, tOcO[0, m, 0][0]] + + if const_expr(self.use_tma_partial): + cute.autovec_copy( + tOsO_partial[None, None, None, stage_compute], + tOrO_partial, + ) + else: + # Load the next thread-private cp.async stage. + split_to_load = s + self.stages - 1 + if split_to_load <= thr_max_valid_split: + load_O_partial(split_to_load, stage_load) + cute.arch.cp_async_commit_group() + stage_load = ( + 0 if stage_load == self.stages - 1 else stage_load + 1 + ) + cute.arch.cp_async_wait_group(self.stages - 1) + cute.autovec_copy( + tOsO_partial[None, None, None, stage_compute], + tOrO_partial, + ) + stage_compute = ( + 0 + if stage_compute == self.stages - 1 + else stage_compute + 1 + ) + + # Accumulate scaled partial results + for m in cutlass.range(num_rows, unroll_full=True): + if tOhidx[m] >= 0 and scale[m] > 0.0: + tOrO[None, m, None].store( + tOrO[None, m, None].load() + + scale[m] * tOrO_partial[None, m, None].load().to(Float32) + ) + + if const_expr(self.use_tma_partial): + # Match CUDA's CTA-wide stage handoff: no consumer may + # release a stage while another warp is still reading it. + cute.arch.sync_threads() + split_to_load = s + self.stages + if warp_idx == 0 and split_to_load <= thr_max_valid_split: + o_tma_barrier = sO_tma_mbar_ptr + stage_compute + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + o_tma_barrier, + o_tma_copy_bytes, + ) + cute.copy( + tma_atom_O_partial, + tOgO_tma[None, split_to_load], + tOsO_tma[None, stage_compute], + tma_bar_ptr=o_tma_barrier, + ) + + # Flush any outstanding async-copy groups before the local Step-7 + # permutation buffer is read on the tail of the kernel. + if const_expr(not self.use_tma_partial): + cute.arch.cp_async_wait_group(0) + cute.arch.sync_threads() + + # =============================== + # Step 7: Write final O to gmem (fake→real via SMEM) + # =============================== + + if const_expr(cu_seqlens is None): + mO_cur = mO[None, None, None, batch_idx] + elif const_expr(flatten_varlen): + mO_cur = mO + else: + mO_cur = cute.domain_offset((offset, 0, 0), mO) + mO_cur = utils.domain_offset_aligned((0, k_block * self.k_block_size, 0), mO_cur) + num_vals = const_expr(cute.size(tOcO, mode=[0])) + if const_expr(not use_pdl): + # Direct / standalone calls don't participate in the K1->K2 + # dependency chain. Use a simple per-element real-column store + # path here to keep mixed-shape launches stable. + for m in cutlass.range(num_rows, unroll_full=True): + if tOhidx[m] >= 0: + for k in cutlass.range(cute.size(tOcO, mode=[2]), unroll_full=True): + if const_expr(self.is_even_k) or tOpO[k]: + for v in cutlass.range(num_vals, unroll_full=True): + fake_col = tOcO[v, 0, k][1] + real_col = stg128_half_fake_col_to_real_col(fake_col) + o_val = tOrO[v, m, k] + if const_expr(mOutputScale is not None): + o_val = o_val * output_scale + mO_cur[tOmidx[m], real_col, tOhidx[m]] = o_val.to(self.dtype) + elif const_expr(self.dtype in [cutlass.BFloat16, cutlass.Float16]): + # The BF16/F16 fake layout is a 4x4 transpose of 32-bit + # pairs inside each 32-column group. Transpose the pair + # registers with two XOR-shuffle stages, then retain the + # original coalesced 128-bit GMEM store. + assert num_vals == 8 + gmem_thr_copy_O = gmem_tiled_copy_O.get_slice(tidx) + tOcO_store = gmem_thr_copy_O.partition_D(cO) + rO = cute.make_rmem_tensor(tOcO_store.shape, self.dtype) + elems_per_store = const_expr( + cute.size(gmem_tiled_copy_O.layout_tv_tiled[1]) + ) + num_store_rows = const_expr(cute.size(tOcO_store, mode=[1])) + assert num_store_rows == num_rows + lane_in_pair_group = tidx % Int32(4) + + for m in cutlass.range(num_store_rows, unroll_full=True): + for k in cutlass.range( + cute.size(tOcO_store, mode=[2]), unroll_full=True + ): + o0, o1 = tOrO[0, m, k], tOrO[1, m, k] + o2, o3 = tOrO[2, m, k], tOrO[3, m, k] + o4, o5 = tOrO[4, m, k], tOrO[5, m, k] + o6, o7 = tOrO[6, m, k], tOrO[7, m, k] + if const_expr(mOutputScale is not None): + o0, o1 = cute.arch.mul_packed_f32x2( + (o0, o1), (output_scale, output_scale) + ) + o2, o3 = cute.arch.mul_packed_f32x2( + (o2, o3), (output_scale, output_scale) + ) + o4, o5 = cute.arch.mul_packed_f32x2( + (o4, o5), (output_scale, output_scale) + ) + o6, o7 = cute.arch.mul_packed_f32x2( + (o6, o7), (output_scale, output_scale) + ) + word0 = utils.cvt_f16x2_f32(o0, o1, self.dtype) + word1 = utils.cvt_f16x2_f32(o2, o3, self.dtype) + word2 = utils.cvt_f16x2_f32(o4, o5, self.dtype) + word3 = utils.cvt_f16x2_f32(o6, o7, self.dtype) + + partner0 = cute.arch.shuffle_sync_bfly(word1, offset=1) + partner1 = cute.arch.shuffle_sync_bfly(word0, offset=1) + partner2 = cute.arch.shuffle_sync_bfly(word3, offset=1) + partner3 = cute.arch.shuffle_sync_bfly(word2, offset=1) + stage0, stage1 = word0, word1 + stage2, stage3 = word2, word3 + if lane_in_pair_group % Int32(2) != Int32(0): + stage0, stage2 = partner0, partner2 + if lane_in_pair_group % Int32(2) != Int32(1): + stage1, stage3 = partner1, partner3 + + partner0 = cute.arch.shuffle_sync_bfly(stage2, offset=2) + partner1 = cute.arch.shuffle_sync_bfly(stage3, offset=2) + partner2 = cute.arch.shuffle_sync_bfly(stage0, offset=2) + partner3 = cute.arch.shuffle_sync_bfly(stage1, offset=2) + out0, out1 = stage0, stage1 + out2, out3 = stage2, stage3 + if lane_in_pair_group // Int32(2) != Int32(0): + out0, out1 = partner0, partner1 + if lane_in_pair_group // Int32(2) != Int32(1): + out2, out3 = partner2, partner3 + + rO_pairs = cute.recast_tensor( + rO[None, m, k], cutlass.Int32 + ) + rO_pairs[0] = out0 + rO_pairs[1] = out1 + rO_pairs[2] = out2 + rO_pairs[3] = out3 + + for m in cutlass.range(num_store_rows, unroll_full=True): + row_local = tOcO_store[0, m, 0][0] + idx = m_block * self.tile_m + row_local + if idx < max_idx: + m_idx, head_idx = self.decode_flat_row_idx( + idx, head_divmod + ) + mO_cur_copy = cute.tiled_divide( + mO_cur[m_idx, None, head_idx], (elems_per_store,) + ) + for k in cutlass.range( + cute.size(tOcO_store, mode=[2]), unroll_full=True + ): + k_idx = ( + tOcO_store[0, 0, k][1] // elems_per_store + ) + cute.copy( + gmem_thr_copy_O, + rO[None, m, k], + mO_cur_copy[None, k_idx], + ) + else: + # 7a: fp32 accumulator -> output dtype SMEM with fake→real + # permutation. The dedicated permutation buffer stays separate + # from the O_partial pipeline staging buffer. + sO_perm = sO_perm_buf + + if const_expr(self.dtype in [cutlass.BFloat16, cutlass.Float16]): + # O_partial uses a dtype-specific STG.128 fake layout, but + # sO_perm is in the final O dtype. For all supported fake + # layouts, adjacent fake pairs map to adjacent real columns, + # so write the final BF16/F16 O pair as one 32-bit SMEM store. + assert num_vals % 2 == 0 + r2s_o_pair_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + cutlass.Int32, + num_bits_per_copy=32, + ) + rO_pair_word = cute.make_rmem_tensor((1,), cutlass.Int32) + sO_perm_i32_base = cute.make_ptr( + dtype=cutlass.Int32, + value=sO_perm.iterator.toint(), + mem_space=sO_perm.iterator.memspace, + assumed_align=4, + ) + sO_perm_i32_row_stride = Int32((self.k_block_size + 16) // 2) + for m in cutlass.range(num_rows, unroll_full=True): + row_local = tOcO[0, m, 0][0] + if tOhidx[m] >= 0: + for k in cutlass.range(cute.size(tOcO, mode=[2]), unroll_full=True): + for v_pair in cutlass.range(num_vals // 2, unroll_full=True): + v = v_pair * 2 + fake_col = tOcO[v, 0, k][1] + real_col = stg128_half_fake_col_to_real_col(fake_col) + o0 = tOrO[v, m, k] + o1 = tOrO[v + 1, m, k] + if const_expr(mOutputScale is not None): + o0, o1 = cute.arch.mul_packed_f32x2( + (o0, o1), + (output_scale, output_scale), + ) + rO_pair_word[0] = utils.cvt_f16x2_f32(o0, o1, self.dtype) + smem_pair_ptr = cute.make_ptr( + dtype=cutlass.Int32, + value=( + sO_perm_i32_base.toint() + + Int64( + row_local * sO_perm_i32_row_stride + + real_col // Int32(2) + ) + * Int64(4) + ), + mem_space=sO_perm.iterator.memspace, + assumed_align=4, + ) + sO_pair = cute.make_tensor( + smem_pair_ptr, + cute.make_layout((1,), stride=(1,)), + ) + cute.copy(r2s_o_pair_atom, rO_pair_word, sO_pair) + else: + # 7a: iterate over ALL val elements in mode[0]. + # tOcO[v, m, k][1] gives different fake_col for each v. + r2s_o_scalar_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.dtype, + num_bits_per_copy=self.dtype.width, + ) + rO_scalar = cute.make_rmem_tensor((1,), self.dtype) + for m in cutlass.range(num_rows, unroll_full=True): + row_local = tOcO[0, m, 0][0] + if tOhidx[m] >= 0: + for k in cutlass.range(cute.size(tOcO, mode=[2]), unroll_full=True): + for v in cutlass.range(num_vals, unroll_full=True): + fake_col = tOcO[v, 0, k][1] + real_col = stg128_half_fake_col_to_real_col(fake_col) + o_val = tOrO[v, m, k] + if const_expr(mOutputScale is not None): + o_val = o_val * output_scale + rO_scalar[0] = o_val.to(self.dtype) + smem_ptr = utils.elem_pointer(sO_perm, (row_local, real_col)) + smem_scalar_ptr = cute.make_ptr( + dtype=self.dtype, + value=smem_ptr.toint(), + mem_space=sO_perm.iterator.memspace, + assumed_align=self.dtype.width // 8, + ) + sO_scalar = cute.make_tensor( + smem_scalar_ptr, + cute.make_layout((1,), stride=(1,)), + ) + cute.copy(r2s_o_scalar_atom, rO_scalar, sO_scalar) + + cute.arch.sync_threads() + + # 7b: SMEM (real order, output dtype) → GMEM + gmem_thr_copy_O = gmem_tiled_copy_O.get_slice(tidx) + tOcO_store = gmem_thr_copy_O.partition_D(cO) + tOsO_store = gmem_thr_copy_O.partition_D(sO_perm) + rO = cute.make_rmem_tensor(tOcO_store.shape, self.dtype) + elems_per_store = const_expr(cute.size(gmem_tiled_copy_O.layout_tv_tiled[1])) + num_store_rows = const_expr(cute.size(tOcO_store, mode=[1])) + tOpO_store = None + if const_expr(not self.is_even_k): + tOpO_store = cute.make_rmem_tensor(cute.size(tOcO_store, mode=[2]), Boolean) + for k in cutlass.range(cute.size(tOpO_store), unroll_full=True): + tOpO_store[k] = ( + tOcO_store[0, 0, k][1] + < mO_partial.shape[1] - k_block * self.k_block_size + ) + + # Read output dtype from SMEM (now in real column order). + for m in cutlass.range(num_store_rows, unroll_full=True): + for k in cutlass.range(cute.size(tOcO_store, mode=[2]), unroll_full=True): + if const_expr(self.is_even_k) or tOpO_store[k]: + cute.autovec_copy(tOsO_store[None, m, k], rO[None, m, k]) + + # Write bf16 to GMEM using gmem_tiled_copy_O (same as original FA Step 7) + for m in cutlass.range(num_store_rows, unroll_full=True): + row_local = tOcO_store[0, m, 0][0] + idx = m_block * self.tile_m + row_local + if idx < max_idx: + m_idx, head_idx = self.decode_flat_row_idx(idx, head_divmod) + mO_cur_copy = cute.tiled_divide( + mO_cur[m_idx, None, head_idx], (elems_per_store,) + ) + for k in cutlass.range(cute.size(tOcO_store, mode=[2]), unroll_full=True): + k_idx = tOcO_store[0, 0, k][1] // elems_per_store + if const_expr(self.is_even_k) or tOpO_store[k]: + cute.copy(gmem_thr_copy_O, rO[None, m, k], mO_cur_copy[None, k_idx]) + + @cute.jit + def load_O_partial( + self, + gmem_tiled_copy_O_partial: cute.TiledCopy, + tOrOptr: cute.Tensor, + tOsO_partial: cute.Tensor, + tOhidx: cute.Tensor, + tOSplitCount: cute.Tensor, + tOpO: Optional[cute.Tensor], + tOcO: cute.Tensor, + mO_cur_partial_layout: cute.Layout, + split: Int32, + stage: Int32, + ) -> None: + elems_per_load = const_expr(cute.size(gmem_tiled_copy_O_partial.layout_tv_tiled[1])) + tOsO_partial_cur = tOsO_partial[None, None, None, stage] + for m in cutlass.range(cute.size(tOcO, [1]), unroll_full=True): + if tOhidx[m] >= 0: + o_gmem_ptr = cute.make_ptr( + tOsO_partial.element_type, + tOrOptr[m], + cute.AddressSpace.gmem, + assumed_align=16, + ) + mO_partial_cur = cute.make_tensor( + o_gmem_ptr, cute.slice_(mO_cur_partial_layout, (0, None, None, 0)) + ) + mO_partial_cur_copy = cute.tiled_divide(mO_partial_cur, (elems_per_load,)) + for k in cutlass.range(cute.size(tOcO, mode=[2]), unroll_full=True): + k_idx = tOcO[0, 0, k][1] // elems_per_load + if split < tOSplitCount[m] and (const_expr(tOpO is None) or tOpO[k]): + cute.copy( + gmem_tiled_copy_O_partial, + mO_partial_cur_copy[None, k_idx, split], + tOsO_partial_cur[None, m, k], + ) + else: + tOsO_partial_cur[None, m, k].fill(0) + + +def combine( + o_partial_fake, + lse_partial, + o_out, + lse_out, + *, + lse_temperature_partial=None, + lse_temperature_out=None, + cu_seqlens=None, + seqused=None, + split_counts=None, + output_scale=None, + use_pdl=False, + raw_partial_stats=False, +): + """K2: merge sparse forward split partials into the final output. + + Runs the v1 kernel, which takes BF16/FP16 partials only; use + ``combine_f32f8`` for FP32 / FP8 E4M3 partials. + + STG.128 fake-layout handling remains an internal implementation detail. + When lse_out is provided, the kernel writes the final authoritative + log-sum-exp for each query row/head directly into that tensor. + + Args: + o_partial_fake: + Batched: [num_splits, batch, Sq, head_q, dim] + Varlen: [num_splits, total_q, head_q, dim] + lse_partial: + Batched: [num_splits, batch, Sq, head_q] + Varlen: [num_splits, total_q, head_q] + o_out: + Batched: [batch, Sq, head_q, dim] + Varlen: [total_q, head_q, dim] + lse_out: + Batched: [batch, Sq, head_q] + Varlen: [total_q, head_q] + lse_temperature_partial: + Optional temperature-scaled LSE partial with the same shape as + lse_partial. + lse_temperature_out: + Optional temperature-scaled final LSE with the same shape as + lse_out. + cu_seqlens: Optional [batch + 1] int32 for varlen-Q combine. + seqused: Optional [batch] int32 effective lengths for combine. + split_counts: Optional int32 rowwise valid split counts prepared from + q2k metadata. Batched: [batch, seqlen, head_kv]. Varlen: + [total_q, head_kv]. + output_scale: Optional fp32 tensor with at least one element. When + provided, the final O accumulator is multiplied once before store. + use_pdl: When True, wait on PDL dependencies from the producer K1 + kernel. When False, launch without PDL waits. + raw_partial_stats: Interpret lse_partial as interleaved FP32 + (scaled_max_log2, row_sum) pairs and combine raw partial-O. + """ + D = o_partial_fake.shape[-1] + num_splits = o_partial_fake.shape[0] + raw_partial_stats = bool(raw_partial_stats) + return_temperature_lse = ( + lse_temperature_partial is not None or lse_temperature_out is not None + ) + if (lse_temperature_partial is None) != (lse_temperature_out is None): + raise ValueError( + "lse_temperature_partial and lse_temperature_out must either both be provided or both be None" + ) + if raw_partial_stats: + expected_stats_shape = (*o_partial_fake.shape[:-1], 2) + if tuple(lse_partial.shape) != tuple(expected_stats_shape): + raise ValueError( + "raw lse_partial must have shape " + f"{expected_stats_shape}, got {tuple(lse_partial.shape)}" + ) + if cu_seqlens is None or o_partial_fake.ndim != 4: + raise ValueError("raw_partial_stats currently requires varlen partial tensors") + if return_temperature_lse: + raise ValueError("raw_partial_stats does not support independent temperature LSE") + else: + expected_lse_shape = o_partial_fake.shape[:-1] + if tuple(lse_partial.shape) != tuple(expected_lse_shape): + raise ValueError( + "lse_partial shape must match o_partial_fake without head_dim, " + f"got {tuple(lse_partial.shape)} vs {tuple(expected_lse_shape)}" + ) + if lse_partial.dtype != torch.float32: + raise TypeError(f"lse_partial must be torch.float32, got {lse_partial.dtype}") + if lse_temperature_partial is not None and lse_temperature_partial.shape != lse_partial.shape: + raise ValueError( + "lse_temperature_partial must have the same shape as lse_partial, " + f"got {lse_temperature_partial.shape} vs {lse_partial.shape}" + ) + if lse_temperature_out is not None: + if lse_out is None: + raise ValueError("lse_temperature_out requires lse_out") + if lse_temperature_out.shape != lse_out.shape: + raise ValueError( + "lse_temperature_out must have the same shape as lse_out, " + f"got {lse_temperature_out.shape} vs {lse_out.shape}" + ) + if lse_temperature_out.dtype != torch.float32 or lse_temperature_partial.dtype != torch.float32: + raise TypeError("temperature LSE tensors must be torch.float32") + + if o_partial_fake.dtype not in (torch.bfloat16, torch.float16): + raise TypeError( + "o_partial_fake dtype must be torch.bfloat16 or torch.float16, " + f"got {o_partial_fake.dtype}; use combine_f32f8 for FP32 / FP8 partials" + ) + partial_dtype = _get_cutlass_dtype(o_partial_fake.dtype) + out_dtype = _get_cutlass_dtype(o_out.dtype) + if output_scale is not None: + if output_scale.dtype != torch.float32: + raise TypeError(f"output_scale must be torch.float32, got {output_scale.dtype}") + if output_scale.numel() < 1: + raise ValueError("output_scale must contain at least one element") + if output_scale.device != o_out.device: + raise ValueError("output_scale must be on the same device as o_out") + output_scale = output_scale.contiguous() + if split_counts is not None: + if split_counts.dtype != torch.int32: + raise TypeError(f"split_counts must be torch.int32, got {split_counts.dtype}") + if o_out.ndim == 4: + if split_counts.ndim != 3: + raise ValueError( + f"batched split_counts must have shape [batch, seqlen, head_kv], got {split_counts.shape}" + ) + if split_counts.shape[:2] != o_out.shape[:2]: + raise ValueError( + f"split_counts shape {split_counts.shape} must match batch/seqlen of o_out {o_out.shape}" + ) + else: + if cu_seqlens is None: + raise ValueError("split_counts with varlen output requires cu_seqlens") + if split_counts.ndim != 2: + raise ValueError( + f"varlen split_counts must have shape [total_q, head_kv], got {split_counts.shape}" + ) + if split_counts.shape[0] != o_out.shape[0]: + raise ValueError( + f"split_counts total_q ({split_counts.shape[0]}) must match o_out total_q " + f"({o_out.shape[0]})" + ) + if o_out.shape[-2] % split_counts.shape[-1] != 0: + raise ValueError( + f"o_out heads ({o_out.shape[-2]}) must be divisible by split_counts heads ({split_counts.shape[-1]})" + ) + qheadperkv = o_out.shape[-2] // split_counts.shape[-1] + else: + qheadperkv = 1 + if cu_seqlens is not None: + if cu_seqlens.dtype != torch.int32: + raise TypeError(f"cu_seqlens must be torch.int32, got {cu_seqlens.dtype}") + if cu_seqlens.ndim != 1: + raise ValueError(f"cu_seqlens must be rank-1, got {cu_seqlens.shape}") + if not cu_seqlens.is_contiguous(): + raise ValueError("cu_seqlens must be contiguous") + if seqused is not None: + if seqused.dtype != torch.int32: + raise TypeError(f"seqused must be torch.int32, got {seqused.dtype}") + if seqused.ndim != 1: + raise ValueError(f"seqused must be rank-1, got {seqused.shape}") + if not seqused.is_contiguous(): + raise ValueError("seqused must be contiguous") + + k_block_size = 128 if D > 64 else 64 + tile_m = 64 if raw_partial_stats else 32 + num_threads = tile_m * 4 + has_cu_seqlens = cu_seqlens is not None + has_seqused = seqused is not None + has_lse = lse_out is not None + has_split_counts = split_counts is not None + has_output_scale = output_scale is not None + # min_blocks_per_mp > 1 caps registers and sets a shared-memory carveout; + # with ptxas choosing the register budget this kernel is as fast or faster + # on B200/B300, and on SM107 a min-blocks carveout fails the launch (Xid 13). + min_blocks_per_mp = 0 + + capability = torch.cuda.get_device_capability(o_out.device) + # Rubin (SM107) takes a 3-stage O_partial ring; Blackwell keeps 2. + stages = 3 if capability == (10, 7) else 2 + key = ( + "combine", + capability, + stages, + D, + k_block_size, + tile_m, + num_threads, + num_splits, + partial_dtype, + out_dtype, + has_cu_seqlens, + has_seqused, + has_lse, + bool(return_temperature_lse), + has_split_counts, + has_output_scale, + use_pdl, + min_blocks_per_mp, + raw_partial_stats, + ) + if key not in _combine_compile_cache: + loaded = try_load_aot(key) + if loaded is not None: + _combine_compile_cache[key] = loaded + else: + from quack.compile_utils import make_fake_tensor + + kernel = SparseAttentionForwardCombine( + dtype=out_dtype, + dtype_partial=partial_dtype, + head_dim=D, + tile_m=tile_m, + k_block_size=k_block_size, + topk=num_splits, + num_threads=num_threads, + use_pdl=use_pdl, + min_blocks_per_mp=min_blocks_per_mp, + stages=stages, + raw_partial_stats=raw_partial_stats, + ) + div = 128 // partial_dtype.width + if has_cu_seqlens: + total_q, nheads = (cute.sym_int64() for _ in range(2)) + mO_partial = make_fake_tensor( + partial_dtype, (num_splits, total_q, nheads, D), divisibility=div + ) + mLSE_partial = make_fake_tensor( + Float32, + ( + (num_splits, total_q, nheads, 2) + if raw_partial_stats + else (num_splits, total_q, nheads) + ), + divisibility=2 if raw_partial_stats else 1, + leading_dim=3 if raw_partial_stats else 2, + ) + mO = make_fake_tensor( + out_dtype, (total_q, nheads, D), divisibility=128 // out_dtype.width + ) + mLSE = ( + make_fake_tensor(Float32, (total_q, nheads), divisibility=1, leading_dim=1) + if has_lse + else None + ) + mLSE_temperature_partial = ( + make_fake_tensor( + Float32, (num_splits, total_q, nheads), divisibility=1, leading_dim=2 + ) + if return_temperature_lse + else None + ) + mLSE_temperature = ( + make_fake_tensor(Float32, (total_q, nheads), divisibility=1, leading_dim=1) + if return_temperature_lse + else None + ) + else: + batch, sq, nheads = (cute.sym_int64() for _ in range(3)) + mO_partial = make_fake_tensor( + partial_dtype, (num_splits, batch, sq, nheads, D), divisibility=div + ) + mLSE_partial = make_fake_tensor( + Float32, (num_splits, batch, sq, nheads), divisibility=1, leading_dim=3 + ) + mO = make_fake_tensor( + out_dtype, (batch, sq, nheads, D), divisibility=128 // out_dtype.width + ) + mLSE = ( + make_fake_tensor(Float32, (batch, sq, nheads), divisibility=1, leading_dim=2) + if has_lse + else None + ) + mLSE_temperature_partial = ( + make_fake_tensor( + Float32, (num_splits, batch, sq, nheads), divisibility=1, leading_dim=3 + ) + if return_temperature_lse + else None + ) + mLSE_temperature = ( + make_fake_tensor(Float32, (batch, sq, nheads), divisibility=1, leading_dim=2) + if return_temperature_lse + else None + ) + if not has_split_counts: + mSplitCounts = None + elif has_cu_seqlens: + total_q_ctr, nheads_kv = (cute.sym_int64() for _ in range(2)) + mSplitCounts = make_fake_tensor( + Int32, (total_q_ctr, nheads_kv), divisibility=1, leading_dim=1 + ) + else: + nheads_kv = cute.sym_int64() + mSplitCounts = make_fake_tensor( + Int32, (batch, sq, nheads_kv), divisibility=1, leading_dim=2 + ) + mOutputScale = ( + make_fake_tensor(Float32, (cute.sym_int64(),), divisibility=1, leading_dim=0) + if has_output_scale + else None + ) + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + + _combine_compile_cache[key] = compile_with_timing( + kernel, + mO_partial, + mLSE_partial, + mO, + mLSE, + mLSE_temperature_partial, + mLSE_temperature, + None + if cu_seqlens is None + else make_fake_tensor(Int32, (cute.sym_int64(),), divisibility=1, leading_dim=0), + None + if seqused is None + else make_fake_tensor(Int32, (cute.sym_int64(),), divisibility=1, leading_dim=0), + None, + None, + None, + mSplitCounts, + mOutputScale, + Int32(qheadperkv), + stream, + options="--enable-tvm-ffi", + ) + save_aot(key, _combine_compile_cache[key], sources=[__name__]) + + with torch.cuda.nvtx.range("K2_Combine"): + _combine_compile_cache[key]( + o_partial_fake, + lse_partial, + o_out, + lse_out, + lse_temperature_partial, + lse_temperature_out, + cu_seqlens, + seqused, + None, + None, + None, + split_counts, + output_scale, + qheadperkv, + ) diff --git a/python/fmha_sm100/prefill_q8kv4/interface.py b/python/fmha_sm100/prefill_q8kv4/interface.py index 06a7bbdf..21417a69 100644 --- a/python/fmha_sm100/prefill_q8kv4/interface.py +++ b/python/fmha_sm100/prefill_q8kv4/interface.py @@ -200,10 +200,6 @@ def run_prefill( out = torch.empty((total_q, num_q_heads, _HEAD_DIM), dtype=torch.bfloat16, **options) _, combine = _sparse_stack() - from ..sparse_fmha_adapter import _supports_blackwell_prefill - - if _supports_blackwell_prefill(q.device, topk=topk): - from src.blackwell_prefill.combine import combine load_extension(q.device, block_scale_shift).run( q, k_cache, From a28a79871e8d8b59914efb8f886a2e4f5a6fd4b0 Mon Sep 17 00:00:00 2001 From: Caitlin Wang Date: Sun, 11 Oct 2026 16:14:18 -0700 Subject: [PATCH 2/3] perf(indexer): port Rubin Q8KV8 prefill indexer optimizations (MiniMax-AI/MSA#17) SM107-specific changes: - SM107MmaFP8Op (K=64) for the FP8 UMMA - TMEM load-and-reduce epilogue on SM107 - 168 registers for every warp role - K loader prefetches the next page-table entry - two Q stages - planner caps the page chunk (8..128) to fill the clusters Changes on all architectures: - define the physical page index before the dynamic K-load loop on the non-SM107 path (port of MiniMax-AI/MSA bc5f83a) - keep the four persistent task-bucket loops rolled (cutlass.range(..., unroll=1)) instead of unrolling them at trace time, which mainly shrinks the generated code Rubin (SM107), full 128-case prefill-indexer suite, score + top-k per call, weighted mean versus dev e441976: - H=1: 54.71 -> 49.34 us (about 1.11x); short queries (<512) about -22% - H=2: 88.04 -> 82.15 us (about 1.07x) Top-k output is bitwise identical in all 128 cases. The prefill and top-k tests in test_q8_indexer.py pass on SM107; on SM100 (B200) they passed before the loop-rolling change. Suggested-by: Fan Shi Co-Authored-By: Claude Opus 5.5 Signed-off-by: Caitlin Wang --- .../fmha_sm100/cute/q8_indexer_interface.py | 27 +++- .../cute/src/sm100/q8kv8_indexer_prefill.py | 153 ++++++++++++++---- .../src/sm100/q8kv8_indexer_prefill_plan.py | 8 + 3 files changed, 152 insertions(+), 36 deletions(-) diff --git a/python/fmha_sm100/cute/q8_indexer_interface.py b/python/fmha_sm100/cute/q8_indexer_interface.py index 00c51162..c1e5ec72 100644 --- a/python/fmha_sm100/cute/q8_indexer_interface.py +++ b/python/fmha_sm100/cute/q8_indexer_interface.py @@ -602,6 +602,7 @@ class _PrefillPlanState: plan_error: torch.Tensor task_capacity: int num_candidate_q_tiles: int + max_page_chunk: int total_q: int num_heads: int scores: torch.Tensor @@ -641,12 +642,19 @@ def _run_prefill_plan(state: _PrefillPlanState) -> None: Int32(state.num_candidate_q_tiles), Int32(state.task_capacity), Int32(state.num_heads), + Int32(state.max_page_chunk), cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), options="--enable-tvm-ffi", ), ) reset_compiled(state.task_counts, state.plan_error) - build_compiled(*build_args, state.num_candidate_q_tiles, state.task_capacity, state.num_heads) + build_compiled( + *build_args, + state.num_candidate_q_tiles, + state.task_capacity, + state.num_heads, + state.max_page_chunk, + ) class BatchPrefillIndexerQ8KV8Wrapper: @@ -712,6 +720,22 @@ def plan( total_rows = total_q * self._num_heads q_tile_capacity = -(-total_rows // q_tile) + batch - 1 task_capacity = q_tile_capacity * -(-max_pages // _PREFILL_TASK_CAPACITY_PAGE_CHUNK) + # On SM107, cap the planner's page chunk so every Q tile splits into at + # most clusters // q_tiles tasks: short-query batches then fill the + # machine in one wave instead of leaving most clusters idle. Host-known + # upper bounds only (no D2H); long-query batches keep the default chunk. + if _require_supported_device(device) == (10, 7): + num_clusters = _sm_count(device) // Q8KV8PrefillIndexerSm100.cta_group_size + chunks_per_q_tile = max(1, num_clusters // q_tile_capacity) + max_page_chunk = max( + Q8KV8PrefillIndexerPlanBuild.min_page_chunk, + min( + Q8KV8PrefillIndexerPlanBuild.large_page_chunk, + -(-max_pages // chunks_per_q_tile), + ), + ) + else: + max_page_chunk = Q8KV8PrefillIndexerPlanBuild.large_page_chunk options = {"device": device} self._state = _PrefillPlanState( cu_seqlens_q=cu_seqlens_q, @@ -732,6 +756,7 @@ def plan( plan_error=torch.empty((1,), dtype=torch.int32, **options), task_capacity=task_capacity, num_candidate_q_tiles=batch * -(-max_seqlen_q * self._num_heads // q_tile), + max_page_chunk=max_page_chunk, total_q=total_q, num_heads=self._num_heads, scores=torch.empty((total_rows, max_pages), dtype=torch.float32, **options), diff --git a/python/fmha_sm100/cute/src/sm100/q8kv8_indexer_prefill.py b/python/fmha_sm100/cute/src/sm100/q8kv8_indexer_prefill.py index 273980eb..b91a4aeb 100644 --- a/python/fmha_sm100/cute/src/sm100/q8kv8_indexer_prefill.py +++ b/python/fmha_sm100/cute/src/sm100/q8kv8_indexer_prefill.py @@ -42,6 +42,26 @@ class Q8KV8PrefillIndexerSm100: # batch index, and (q_rows - 1) | (page_count - 1) << task_page_count_shift. task_descriptor_words = 5 task_page_count_shift = 8 + # Operand formats with a dedicated SM107 MMA instruction (K=64 per + # instruction). Any other format, and every other architecture, goes + # through the generic F8F6F4 instruction at K=32. + sm107_fp8_dtypes = (cutlass.Float8E4M3FN, cutlass.Float8E5M2) + + @classmethod + def uses_sm107_fp8_mma( + cls, + compute_capability: tuple[int, int], + io_dtype=cutlass.Float8E4M3FN, + ) -> bool: + return compute_capability == (10, 7) and io_dtype in cls.sm107_fp8_dtypes + + @classmethod + def mma_inst_k( + cls, + compute_capability: tuple[int, int], + io_dtype=cutlass.Float8E4M3FN, + ) -> int: + return 64 if cls.uses_sm107_fp8_mma(compute_capability, io_dtype) else 32 def __init__( self, @@ -59,11 +79,22 @@ def __init__( self.io_dtype = cutlass.Float8E4M3FN self.acc_dtype = cutlass.Float32 self.compute_capability = compute_capability + # Two Q stages on SM107 so the next task's Q tile can land while the + # current task's pages are still being consumed; the MMA warp reads + # its Q operand from the current consumer stage. + self.q_stages = 2 if compute_capability == (10, 7) else 1 self.num_persistent_clusters = num_persistent_clusters - self.use_tmem_load_reduce = compute_capability == (10, 3) + self.use_tmem_load_reduce = compute_capability in ((10, 3), (10, 7)) self.cluster_shape_mnk = (2, 1, 1) self.mma_tiler_mnk = (self.q_tile, self.k_tile, self.head_dim) - self.mma_inst_shape_mnk = (self.q_tile, self.k_tile, 32) + self.use_sm107_fp8_mma = self.uses_sm107_fp8_mma( + compute_capability, self.io_dtype + ) + self.mma_inst_shape_mnk = ( + self.q_tile, + self.k_tile, + self.mma_inst_k(compute_capability, self.io_dtype), + ) self.score_warp_ids = tuple(range(8)) self.q_load_warp_id = 8 @@ -73,8 +104,18 @@ def __init__( self.score_consumer_warps = 4 self.score_worker_groups = 2 self.threads_per_cta = cute.arch.WARP_SIZE * 12 - self.num_regs_score = 192 - self.num_regs_other = 32 + # Register budgets per warp role. The launch allocates 168 registers + # per thread (64K / 384 threads). On SM107 every role keeps that: the + # score warps fit in 168, and growing them to 192 starved the loader + # and MMA warps. SM100 and SM103 keep the original split. If a split is + # used, 8 * 32 * (score - 168) <= 4 * 32 * (168 - other) must hold or + # the score warps spin forever in setmaxregister_increase. + if compute_capability == (10, 7): + self.num_regs_score = 168 + self.num_regs_other = 168 + else: + self.num_regs_score = 192 + self.num_regs_other = 32 self.tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols("sm_100") @cute.jit @@ -90,16 +131,28 @@ def __call__( stream: cuda.CUstream = None, ): """Build descriptors and launch the 2-CTA kernel.""" - op = tcgen05.MmaF8F6F4Op( - self.io_dtype, - self.io_dtype, - self.acc_dtype, - self.mma_inst_shape_mnk, - tcgen05.CtaGroup.TWO, - tcgen05.OperandSource.SMEM, - OperandMajorMode.K, - OperandMajorMode.K, - ) + if cutlass.const_expr(self.use_sm107_fp8_mma): + op = tcgen05.SM107MmaFP8Op( + self.io_dtype, + self.io_dtype, + self.acc_dtype, + self.mma_inst_shape_mnk, + tcgen05.CtaGroup.TWO, + tcgen05.OperandSource.SMEM, + OperandMajorMode.K, + OperandMajorMode.K, + ) + else: + op = tcgen05.MmaF8F6F4Op( + self.io_dtype, + self.io_dtype, + self.acc_dtype, + self.mma_inst_shape_mnk, + tcgen05.CtaGroup.TWO, + tcgen05.OperandSource.SMEM, + OperandMajorMode.K, + OperandMajorMode.K, + ) tiled_mma = cute.make_tiled_mma(op) q_smem_layout = sm100_utils.make_smem_layout_a( tiled_mma, @@ -422,7 +475,7 @@ def kernel( acc_consumer_state.advance() tile_base = cutlass.Int32(0) worker_base = cutlass.Int32(0) - for rev_bucket in cutlass.range_constexpr(self.num_task_buckets): + for rev_bucket in cutlass.range(self.num_task_buckets, unroll=1): bucket = cutlass.Int32(self.num_task_buckets - 1 - rev_bucket) task_count = mTaskCounts[bucket] task_idx = cluster_idx - worker_base @@ -503,7 +556,7 @@ def kernel( self.q_stages, ) worker_base = cutlass.Int32(0) - for rev_bucket in cutlass.range_constexpr(self.num_task_buckets): + for rev_bucket in cutlass.range(self.num_task_buckets, unroll=1): bucket = cutlass.Int32(self.num_task_buckets - 1 - rev_bucket) task_count = mTaskCounts[bucket] task_idx = cluster_idx - worker_base @@ -565,7 +618,7 @@ def kernel( self.k_stages, ) worker_base = cutlass.Int32(0) - for rev_bucket in cutlass.range_constexpr(self.num_task_buckets): + for rev_bucket in cutlass.range(self.num_task_buckets, unroll=1): bucket = cutlass.Int32(self.num_task_buckets - 1 - rev_bucket) task_count = mTaskCounts[bucket] task_idx = cluster_idx - worker_base @@ -584,20 +637,47 @@ def kernel( task_idx, mTaskDescriptors, ) - for tile_idx in cutlass.range(page_count, unroll=1): - logical_page = page_begin + tile_idx - physical_page = mPageTable[batch_idx, logical_page] - k_pipe.producer_acquire(k_producer_state) - cute.copy( - k_tma_atom, - tKgK[(None, 0, physical_page)], - tKsK[(None, k_producer_state.index)], - tma_bar_ptr=k_pipe.producer_get_barrier( - k_producer_state - ), - mcast_mask=k_mcast_mask, - ) - k_producer_state.advance() + if cutlass.const_expr(self.compute_capability == (10, 7)): + # The physical page index is a dependent global load on + # the TMA issue path; fetch the next page's index while + # this page's TMA is issued so the two latencies overlap. + # The last iteration re-reads the final entry (in bounds). + last_page = page_begin + page_count - cutlass.Int32(1) + physical_page = mPageTable[batch_idx, page_begin] + for tile_idx in cutlass.range(page_count, unroll=1): + logical_page = page_begin + tile_idx + next_page = cutlass.min( + logical_page + cutlass.Int32(1), last_page + ) + next_physical_page = mPageTable[batch_idx, next_page] + k_pipe.producer_acquire(k_producer_state) + cute.copy( + k_tma_atom, + tKgK[(None, 0, physical_page)], + tKsK[(None, k_producer_state.index)], + tma_bar_ptr=k_pipe.producer_get_barrier( + k_producer_state + ), + mcast_mask=k_mcast_mask, + ) + k_producer_state.advance() + physical_page = next_physical_page + else: + physical_page = cutlass.Int32(0) + for tile_idx in cutlass.range(page_count, unroll=1): + logical_page = page_begin + tile_idx + physical_page = mPageTable[batch_idx, logical_page] + k_pipe.producer_acquire(k_producer_state) + cute.copy( + k_tma_atom, + tKgK[(None, 0, physical_page)], + tKsK[(None, k_producer_state.index)], + tma_bar_ptr=k_pipe.producer_get_barrier( + k_producer_state + ), + mcast_mask=k_mcast_mask, + ) + k_producer_state.advance() task_idx += cutlass.Int32(self.num_persistent_clusters) worker_base = ( worker_base + task_count @@ -620,9 +700,7 @@ def kernel( num_k_blocks = cute.size(tCrQ, mode=[2]) if is_leader_cta: worker_base = cutlass.Int32(0) - for rev_bucket in cutlass.range_constexpr( - self.num_task_buckets - ): + for rev_bucket in cutlass.range(self.num_task_buckets, unroll=1): bucket = cutlass.Int32( self.num_task_buckets - 1 - rev_bucket ) @@ -665,7 +743,12 @@ def kernel( tiled_mma, tCtAcc, tCrQ[ - (None, None, k_block_idx, 0) + ( + None, + None, + k_block_idx, + q_consumer_state.index, + ) ], tCrK[ ( diff --git a/python/fmha_sm100/cute/src/sm100/q8kv8_indexer_prefill_plan.py b/python/fmha_sm100/cute/src/sm100/q8kv8_indexer_prefill_plan.py index caca7b26..100ffc0d 100644 --- a/python/fmha_sm100/cute/src/sm100/q8kv8_indexer_prefill_plan.py +++ b/python/fmha_sm100/cute/src/sm100/q8kv8_indexer_prefill_plan.py @@ -57,6 +57,7 @@ class Q8KV8PrefillIndexerPlanBuild: num_buckets = Q8KV8PrefillIndexerSm100.num_task_buckets split_page_chunk = 64 large_page_chunk = 128 + min_page_chunk = 8 split_q_tile_threshold = 60 descriptor_words = Q8KV8PrefillIndexerSm100.task_descriptor_words page_count_shift = Q8KV8PrefillIndexerSm100.task_page_count_shift @@ -74,6 +75,7 @@ def __call__( num_candidate_q_tiles: cutlass.Int32, task_capacity: cutlass.Int32, num_heads: cutlass.Int32, + max_page_chunk: cutlass.Int32, stream: cuda.CUstream = None, ) -> None: self.kernel( @@ -85,6 +87,7 @@ def __call__( mPlanError, task_capacity, num_heads, + max_page_chunk, ).launch( grid=(num_candidate_q_tiles, 1, 1), block=(self.threads_per_cta, 1, 1), @@ -137,6 +140,7 @@ def kernel( mPlanError: cute.Tensor, task_capacity: cutlass.Int32, num_heads: cutlass.Int32, + max_page_chunk: cutlass.Int32, ) -> None: lane_idx = cute.arch.lane_idx() candidate_idx, _, _ = cute.arch.block_idx() @@ -196,6 +200,10 @@ def kernel( page_chunk = cutlass.Int32(self.split_page_chunk) if num_q_tiles >= cutlass.Int32(self.split_q_tile_threshold): page_chunk = cutlass.Int32(self.large_page_chunk) + # Host-computed cap keeps the task count near the resident cluster + # count when few Q tiles would otherwise leave most clusters idle. + if page_chunk > max_page_chunk: + page_chunk = max_page_chunk if page_chunk > num_pages: page_chunk = num_pages page_begin = cutlass.Int32(0) From ac00774dd300b945d60164643a390f9d75a26fba Mon Sep 17 00:00:00 2001 From: Caitlin Wang Date: Sun, 11 Oct 2026 16:55:26 -0700 Subject: [PATCH 3/3] feat(indexer): let Q8KV8 prefill callers own the score buffer BatchPrefillIndexerQ8KV8Wrapper always allocated its own score buffer and its only public call ran the TopK as well. A caller that selects pages itself (vLLM's shared TopK over decode and prefill rows) needs the scores written into its own buffer, without the wrapper's TopK. - plan(..., score_out=None): optional caller-owned fp32 buffer of shape [total_q * num_heads, >= max_pages] with contiguous pages per row; a wider row pitch is accepted, so a slice of a larger buffer works. The kernel layout is already dynamic, so no recompile and no kernel change. - run_scores(q, k_cache): public scores-only call; run() is unchanged. Tests: score_out matches the wrapper-owned scores bitwise at an 8192-page pitch, leaves other entries untouched, runs no TopK, and run() TopK is unchanged; invalid score_out buffers are rejected. Prefill and TopK tests in test_q8_indexer.py pass on SM107 (41 passed). Co-Authored-By: Claude Opus 5.5 Signed-off-by: Caitlin Wang --- python/fmha_sm100/cute/README.md | 9 ++++ .../fmha_sm100/cute/q8_indexer_interface.py | 47 ++++++++++++++++++- python/fmha_sm100/cute/test_q8_indexer.py | 45 +++++++++++++++++- 3 files changed, 99 insertions(+), 2 deletions(-) diff --git a/python/fmha_sm100/cute/README.md b/python/fmha_sm100/cute/README.md index e360ac6b..49610e14 100644 --- a/python/fmha_sm100/cute/README.md +++ b/python/fmha_sm100/cute/README.md @@ -442,6 +442,15 @@ prefill.plan( topk_indices = prefill.run(index_q, index_k_cache) # [total_q, 4, 16] int32 ``` +A caller that runs its own page selection (for example vLLM's shared TopK over +decode and prefill rows) passes its score buffer to `plan` and scores without +the TopK: + +```python +prefill.plan(..., score_out=scores[rows_lo:rows_hi]) # fp32 [total_q * num_heads, >= max_pages] +prefill.run_scores(index_q, index_k_cache) # writes pages before each row's local page +``` + ### Q8KV Input And Output Contract - `index_q`: `[num_tokens, num_heads, 128]`, `torch.float8_e4m3fn`, diff --git a/python/fmha_sm100/cute/q8_indexer_interface.py b/python/fmha_sm100/cute/q8_indexer_interface.py index c1e5ec72..1ad06a52 100644 --- a/python/fmha_sm100/cute/q8_indexer_interface.py +++ b/python/fmha_sm100/cute/q8_indexer_interface.py @@ -249,6 +249,24 @@ def bind_indexer_module_loader(loader) -> None: _indexer_module_loader = loader +def _check_score_out( + score_out: torch.Tensor, rows: int, max_pages: int, device: torch.device +) -> None: + """Validate a caller-owned prefill score buffer: fp32 ``[rows, >= max_pages]``.""" + + _check_cuda_tensor( + score_out, name="score_out", dtype=torch.float32, device=device, contiguous=False + ) + if score_out.ndim != 2 or score_out.shape[0] != rows or score_out.shape[1] < max_pages: + raise ValueError( + f"score_out must have shape [{rows}, >= {max_pages}], got {tuple(score_out.shape)}" + ) + if score_out.stride(1) != 1 or score_out.stride(0) < score_out.shape[1]: + raise ValueError("score_out must keep each row's pages contiguous and rows disjoint") + if score_out.data_ptr() % _SCALAR_ALIGNMENT: + raise ValueError(f"score_out must be {_SCALAR_ALIGNMENT}-byte aligned") + + def _topk_select( scores: torch.Tensor, lengths: torch.Tensor, out: torch.Tensor, *, use_pdl: bool = True ) -> torch.Tensor: @@ -672,6 +690,9 @@ class BatchPrefillIndexerQ8KV8Wrapper: total_q=total_q, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, ) topk_indices = wrapper.run(q, k_cache) + + A caller that selects pages itself passes its own score buffer to ``plan`` + (``score_out``) and calls ``run_scores``, which skips the TopK. """ def __init__(self, *, num_heads: int = 1) -> None: @@ -687,12 +708,19 @@ def plan( total_q: int, max_seqlen_q: int, max_seqlen_k: int, + score_out: torch.Tensor | None = None, ) -> None: """Size reusable buffers from host bounds and build device tasks. ``cu_seqlens_q`` is ``[batch + 1]`` int32 starting at zero. The host bounds must cover the device metadata; they size buffers and the plan grid but never enter a compile key. Call outside CUDA Graph capture. + + ``score_out`` optionally replaces the wrapper's score buffer: fp32 + ``[total_q * num_heads, >= ceil(max_seqlen_k / 128)]`` with contiguous + pages per row; the row pitch may be wider, so a slice of a larger + buffer works. Scores are written in place; entries the kernel does not + write keep their contents. """ _check_cuda_tensor(cu_seqlens_q, name="cu_seqlens_q", dtype=torch.int32) @@ -736,6 +764,8 @@ def plan( ) else: max_page_chunk = Q8KV8PrefillIndexerPlanBuild.large_page_chunk + if score_out is not None: + _check_score_out(score_out, total_rows, max_pages, device) options = {"device": device} self._state = _PrefillPlanState( cu_seqlens_q=cu_seqlens_q, @@ -759,7 +789,11 @@ def plan( max_page_chunk=max_page_chunk, total_q=total_q, num_heads=self._num_heads, - scores=torch.empty((total_rows, max_pages), dtype=torch.float32, **options), + scores=( + score_out + if score_out is not None + else torch.empty((total_rows, max_pages), dtype=torch.float32, **options) + ), num_valid_pages=torch.empty((total_rows,), dtype=torch.int32, **options), topk_indices=torch.empty( (total_q, self._num_heads, _TOP_K), dtype=torch.int32, **options @@ -774,6 +808,17 @@ def replan(self) -> None: raise RuntimeError("plan() must be called before replan()") _run_prefill_plan(self._state) + def run_scores(self, q: torch.Tensor, k_cache: torch.Tensor) -> torch.Tensor: + """Score every row's historical pages without the TopK; return the score buffer. + + Each row's pages before its local page are written. The local page and + later pages keep their previous contents, so a caller selecting its own + pages pre-fills the buffer (for example with ``-inf``) and forces the + local page. + """ + + return self._run_scores(q, k_cache) + def _run_scores(self, q: torch.Tensor, k_cache: torch.Tensor) -> torch.Tensor: """Write historical page scores; the local page and later pages stay untouched.""" diff --git a/python/fmha_sm100/cute/test_q8_indexer.py b/python/fmha_sm100/cute/test_q8_indexer.py index 439939de..8fde5fe9 100644 --- a/python/fmha_sm100/cute/test_q8_indexer.py +++ b/python/fmha_sm100/cute/test_q8_indexer.py @@ -836,7 +836,7 @@ def _make_prefill( } -def _plan_prefill(wrapper, inputs): +def _plan_prefill(wrapper, inputs, **kwargs): wrapper.plan( inputs["cu_seqlens_q"], inputs["seq_lens"], @@ -844,6 +844,7 @@ def _plan_prefill(wrapper, inputs): total_q=inputs["total_q"], max_seqlen_q=inputs["max_seqlen_q"], max_seqlen_k=inputs["max_seqlen_k"], + **kwargs, ) @@ -964,6 +965,48 @@ def test_prefill_plan_reuse_replan_and_cuda_graph(num_heads): assert torch.equal(wrapper._state.scores.nan_to_num(POISON), expected_scores) +@pytest.mark.parametrize("num_heads", PREFILL_HEADS) +def test_prefill_score_out_writes_caller_buffer(num_heads): + """``score_out`` takes the scores at the caller's pitch; ``run_scores`` skips the TopK.""" + + inputs = _make_prefill([300, 45, 900], [100, 2000, 0], seed=11, num_heads=num_heads) + q, k_cache = inputs["q"], inputs["k"].cache + owned = BatchPrefillIndexerQ8KV8Wrapper(num_heads=num_heads) + _plan_prefill(owned, inputs) + expected = _run_prefill_scores(owned, q, k_cache).clone() + expected_topk = owned.run(q, k_cache).clone() + + # A row-offset slice of a wider buffer, the way vLLM passes its shared scores. + rows, width = expected.shape + shared = torch.full((rows + 3, 8192), float("nan"), device="cuda") + score_out = shared[3:] + wrapper = BatchPrefillIndexerQ8KV8Wrapper(num_heads=num_heads) + _plan_prefill(wrapper, inputs, score_out=score_out) + wrapper._state.topk_indices.fill_(-777) + assert wrapper.run_scores(q, k_cache).data_ptr() == score_out.data_ptr() + torch.cuda.synchronize() + assert (wrapper._state.topk_indices == -777).all() + assert torch.equal(score_out[:, :width].nan_to_num(POISON), expected.nan_to_num(POISON)) + assert torch.isnan(score_out[:, width:]).all() + assert torch.isnan(shared[:3]).all() + assert torch.equal(wrapper.run(q, k_cache), expected_topk) + + +def test_prefill_rejects_invalid_score_out(): + inputs = _make_prefill([10, 20], [0, 100], seed=9) + rows = inputs["total_q"] + max_pages = -(-inputs["max_seqlen_k"] // PAGE_SIZE) + for bad in ( + torch.empty((rows, max_pages), dtype=torch.bfloat16, device="cuda"), + torch.empty((rows + 1, max_pages), device="cuda"), + torch.empty((rows, max_pages - 1), device="cuda"), + torch.empty((max_pages + 1, rows), device="cuda").t(), + torch.empty((rows, max_pages)), + ): + with pytest.raises((TypeError, ValueError), match="score_out"): + _plan_prefill(BatchPrefillIndexerQ8KV8Wrapper(), inputs, score_out=bad) + + def test_prefill_compile_is_shape_independent(): """Batch, total_q, lengths, and varlen distribution reuse one compiled kernel set."""