diff --git a/flash_attn/cute/flash_bwd_mla_dq_dqv_sm100.py b/flash_attn/cute/flash_bwd_mla_dq_dqv_sm100.py index e9275025c48..73280e00a6f 100644 --- a/flash_attn/cute/flash_bwd_mla_dq_dqv_sm100.py +++ b/flash_attn/cute/flash_bwd_mla_dq_dqv_sm100.py @@ -5,7 +5,7 @@ Performs both dQ = dS @ K and dQv = dS @ V, where K and V are gathered according to index tensor mIdxTopK. -This uses MQA with 128 heads. +This uses MQA with 1..128 heads and a fixed 128-row GEMM tile. Inputs: - dS: [batch, seqlen_q, nheads, top_k] or [total_q, nheads, top_k] @@ -59,17 +59,17 @@ def __init__( ): self.acc_dtype: Type[cutlass.Numeric] = acc_dtype self.nheads = nheads - assert self.nheads == 128, ( - "only 128 heads supported; will expand to include 64 heads in a future PR." - ) + # Head padding: see pack_gqa.padded_qheads_tma_source. + self.tile_m = 128 + assert 0 < self.nheads <= self.tile_m, f"at most {self.tile_m} heads, got {nheads}" self.head_dim_k = head_dim_k or 0 # when head_dim_k not provided, dQ is not computed self.head_dim_v = head_dim_v self.top_k = top_k self.tile_k = 128 self.cluster_shape_mn = (1, 2) - self.mma_tiler_dQ = (self.nheads, self.head_dim_k, self.tile_k) - self.mma_tiler_dQv = (self.nheads, self.head_dim_v // 2, self.tile_k) + self.mma_tiler_dQ = (self.tile_m, self.head_dim_k, self.tile_k) + self.mma_tiler_dQv = (self.tile_m, self.head_dim_v // 2, self.tile_k) self.num_mainloop_iters = self.top_k // self.tile_k self.arch = "sm_100" @@ -187,12 +187,14 @@ def static_reshape(t: cute.Tensor, *static_shapes) -> cute.Tensor: ), ) - mdS = static_reshape(mdS, self.nheads, self.top_k) - mdQv = static_reshape(mdQv, self.nheads, self.head_dim_v) + # Dynamic head extent for the fixed tile; see pack_gqa.padded_qheads_tma_source. + nheads = self.nheads if self.nheads == self.tile_m else Int32(self.nheads) + mdS = static_reshape(mdS, nheads, self.top_k) + mdQv = static_reshape(mdQv, nheads, self.head_dim_v) mV = static_reshape(mV, self.head_dim_v) mIdxTopK = static_reshape(mIdxTopK, self.top_k) if const_expr(self.compute_dQ): - mdQ = static_reshape(mdQ, self.nheads, self.head_dim_k) + mdQ = static_reshape(mdQ, nheads, self.head_dim_k) mK = static_reshape(mK, self.head_dim_k) # ---- layout info ---- diff --git a/flash_attn/cute/flash_bwd_mla_sm100.py b/flash_attn/cute/flash_bwd_mla_sm100.py index 00fddefacf9..3a0ade2ba4a 100644 --- a/flash_attn/cute/flash_bwd_mla_sm100.py +++ b/flash_attn/cute/flash_bwd_mla_sm100.py @@ -16,7 +16,11 @@ from quack import copy_utils, layout_utils -from flash_attn.cute.pack_gqa import pack_gqa_layout +from flash_attn.cute.pack_gqa import ( + pack_gqa_layout, + padded_qheads_tma_source, + regroup_padded_qheads, +) from flash_attn.cute.seqlen_info import SeqlenInfoQK from flash_attn.cute.block_info import BlockInfo import flash_attn.cute.blackwell_helpers as fa_sm100_utils @@ -52,12 +56,19 @@ def __init__( has_seqused_q: bool = False, disable_bitmask: bool = False, use_clc_scheduler: bool = True, + qhead_per_kvhead_valid: Optional[int] = None, ): use_cpasync_load_KV = True self.is_causal = is_causal self.is_local = False self.pack_gqa = True self.qhead_per_kvhead = qhead_per_kvhead + # Head padding and scaleP/dPsum contract: see pack_gqa.padded_qheads_tma_source. + if qhead_per_kvhead_valid is None: + qhead_per_kvhead_valid = qhead_per_kvhead + assert 0 < qhead_per_kvhead_valid <= qhead_per_kvhead + self.qhead_per_kvhead_valid = qhead_per_kvhead_valid + self.pad_qheads = qhead_per_kvhead_valid != qhead_per_kvhead self.nheads_kv = nheads_kv self.has_seqused_q = has_seqused_q self.use_tma_O = True @@ -65,7 +76,7 @@ def __init__( self.use_tma_KV = False self.topk_length = topk_length self.is_topk_gather = True - assert qhead_per_kvhead == 128 or qhead_per_kvhead == 64 + assert qhead_per_kvhead in (64, 128), f"sparse MLA bwd supports 64 or 128 heads, got {qhead_per_kvhead}" # user-provided option if topk indices guaranteed in bounds self.disable_bitmask = disable_bitmask @@ -211,7 +222,8 @@ def __init__( self.num_stages_dP = 1 self.num_stages_dPt = 1 self.num_stages_dV = 2 # == hdimv splits, for Umma <-> Async - self.num_epi_stages_dV = 8 # == 2 splits x 4 slots/split + # Per hdimv split: 2 warpgroup halves x 2 subtile parities. + self.num_epi_stages_dV = 4 self.num_stages_scaleP = 1 self.num_stages_dPsum = 1 @@ -251,7 +263,6 @@ def mbar_struct(num_stages): sdO_struct, sP_struct, sdS_struct, - sQvt_struct, sScaleP_struct, sdPsum_struct, ) = ( @@ -261,11 +272,28 @@ def mbar_struct(num_stages): (self.dtype, self.sdO_layout_staged), (self.dtype, self.sPt_layout_staged), (self.dtype, self.sdSt_layout_staged), - (self.dtype, self.sQvt_layout_staged), (self.dtype_scale, self.sScaleP_layout_staged), (self.dtype_scale, self.sdPsum_layout_staged), ] ) + # dV staging must not overwrite operands still in use by another split's MMA. + # If staging equals one dOt/Qvt stage, split s stages in place over stage s. + # Otherwise, both splits take turns using one staging area after the operands. + staging_elems = ( + math.prod(self.tile_dV) * self.num_epi_stages_dV * self.dtype_acc.width // self.dtype.width + ) + operand_elems = cute.cosize(self.sQvt_layout_staged) + assert self.num_stages_Qvt == self.num_hdimv_splits + stage_elems = operand_elems // self.num_hdimv_splits + if staging_elems == stage_elems: + self.sdV_split_offsets = [split * stage_elems for split in range(self.num_hdimv_splits)] + sQv_elems = operand_elems + else: + self.sdV_split_offsets = [operand_elems] * self.num_hdimv_splits + sQv_elems = operand_elems + staging_elems + sQvt_struct = cute.struct.Align[ + cute.struct.MemRange[self.dtype, sQv_elems], self.buffer_align_bytes + ] ( mbar_ptr_V_struct, # load V @@ -415,6 +443,12 @@ def __call__( ) topk_length_dynamic = mIndexTopk.shape[0] + # TMA source contract: see pack_gqa.padded_qheads_tma_source. + if const_expr(self.pad_qheads): + mQv_valid, mdO_valid, mP_valid, mdS_valid = [ + padded_qheads_tma_source(mX, self.qhead_per_kvhead_valid, head_idx=2) + for mX in (mQv, mdO, mP, mdS) + ] if const_expr(self.pack_gqa): mQv, mdO, mP, mdS, mScaleP = [ pack_gqa_layout(mX, self.qhead_per_kvhead, self.nheads_kv, head_idx=2) @@ -424,15 +458,18 @@ def __call__( ] if const_expr(mdPsum is not None): mdPsum = pack_gqa_layout(mdPsum, self.qhead_per_kvhead, self.nheads_kv, head_idx=1) + if const_expr(not self.pad_qheads): + mQv_valid, mdO_valid, mP_valid, mdS_valid = mQv, mdO, mP, mdS # ((h/h_k, s_q), dv, h_k, b) -> (dv, (h/h_k, s_q), h_k, b) # or ((h/h_k, total_q), dv, h_k) -> (dv, (h/h_k, total_q), h_k) + # *_valid: (h, dv, s_q, b) -> (dv, h, s_q, b). mma_operand_layout_transpose = ( [1, 0, 2, 3] if const_expr(mCuSeqlensQ is None) else [1, 0, 2] ) - mQvt, mdOt = [ + mQvt, mdOt, mQvt_valid, mdOt_valid = [ cute.make_tensor(mX.iterator, cute.select(mX.layout, mode=mma_operand_layout_transpose)) - for mX in (mQv, mdO) + for mX in (mQv, mdO, mQv_valid, mdO_valid) ] # fmt: off @@ -502,21 +539,37 @@ def __call__( ) cta_shape = cta_layout_vmnk.shape - def make_tma(make_fn, mX, smem_layout, mma_tiler, tiled_mma): - return make_fn(tma_load_op, mX, smem_layout, mma_tiler, tiled_mma, cta_shape) + def regroup(tensor, transposed=False): + if const_expr(self.pad_qheads): + # Transposed operands: undo the transpose, fold, and transpose back, mirroring + # how mdOt/mQvt derive from mdO/mQv. + if const_expr(transposed): + tensor = cute.make_tensor( + tensor.iterator, cute.select(tensor.layout, mode=mma_operand_layout_transpose) + ) + tensor = regroup_padded_qheads(tensor, self.qhead_per_kvhead, head_idx=2) + if const_expr(transposed): + tensor = cute.make_tensor( + tensor.iterator, cute.select(tensor.layout, mode=mma_operand_layout_transpose) + ) + return tensor + + def make_tma(make_fn, mX, smem_layout, mma_tiler, tiled_mma, transposed): + atom, tensor = make_fn(tma_load_op, mX, smem_layout, mma_tiler, tiled_mma, cta_shape) + return atom, regroup(tensor, transposed) A, B = cute.nvgpu.make_tiled_tma_atom_A, cute.nvgpu.make_tiled_tma_atom_B - # (atom_name, tensor_name, make_fn, m, smem_layout, mma_tiler, tiled_mma) + # (atom_name, tensor_name, make_fn, m, smem_layout, mma_tiler, tiled_mma, transposed) _tma_specs = [ - ("tma_atom_dO", "tma_tensor_dO", B, mdO, self.sdO_layout, self.mma_tiler_VdO, tiled_mma_VdO), - ("tma_atom_dOt", "tma_tensor_dOt", B, mdOt, self.sdOt_layout, self.mma_tiler_PtdOt, tiled_mma_PtdOt), - ("tma_atom_Qvt", "tma_tensor_Qvt", B, mQvt, self.sQvt_layout, self.mma_tiler_dStQvt, tiled_mma_dStQvt), + ("tma_atom_dO", "tma_tensor_dO", B, mdO_valid, self.sdO_layout, self.mma_tiler_VdO, tiled_mma_VdO, False), + ("tma_atom_dOt", "tma_tensor_dOt", B, mdOt_valid, self.sdOt_layout, self.mma_tiler_PtdOt, tiled_mma_PtdOt, True), + ("tma_atom_Qvt", "tma_tensor_Qvt", B, mQvt_valid, self.sQvt_layout, self.mma_tiler_dStQvt, tiled_mma_dStQvt, True), ] _tmas = {} - for atom_name, tensor_name, make_fn, m, smem_layout, mma_tiler, tiled_mma in _tma_specs: + for atom_name, tensor_name, make_fn, m, smem_layout, mma_tiler, tiled_mma, transposed in _tma_specs: _tmas[atom_name], _tmas[tensor_name] = ( - make_tma(make_fn, m, smem_layout, mma_tiler, tiled_mma) + make_tma(make_fn, m, smem_layout, mma_tiler, tiled_mma, transposed) ) (tma_atom_dO, tma_tensor_dO, @@ -526,10 +579,11 @@ def make_tma(make_fn, mX, smem_layout, mma_tiler, tiled_mma): # Make TMA load for P separately tma_atom_P, tma_tensor_P = cute.nvgpu.cpasync.make_tiled_tma_atom( cpasync.CopyBulkTensorTileG2SOp(), - mP, + mP_valid, self.sP_layout, self.tile_P, ) + tma_tensor_P = regroup(tma_tensor_P) # ==== TMA store ==== tma_store_op = cpasync.CopyBulkTensorTileS2GOp() @@ -545,8 +599,9 @@ def make_tma(make_fn, mX, smem_layout, mma_tiler, tiled_mma): self.dtype_dV, self.dV_layout_major, self.tile_dV, self.num_epi_stages_dV ) tma_atom_dS, tma_tensor_dS = cpasync.make_tiled_tma_atom( - tma_store_op, mdS, cute.select(sdS_layout_staged, mode=[0, 1]), self.tile_dS + tma_store_op, mdS_valid, cute.select(sdS_layout_staged, mode=[0, 1]), self.tile_dS ) + tma_tensor_dS = regroup(tma_tensor_dS) # fmt: on # ==== Allocate shared memory ==== @@ -805,10 +860,14 @@ def make_pipeline(cls, mbar_ptr, num_stages, producer, consumer, tx_count=None): (storage.sQv, sQvt_layout_staged), # {dOt, Qvt, dV} overlap ] ) - sdV = cute.make_tensor( - cute.recast_ptr(sdOt.iterator, sdV_layout_staged.inner, self.dtype_acc), sdV_layout_staged.outer - ) - assert cute.cosize(sdV) * self.dtype_acc.width // self.dtype.width == cute.cosize(sdOt) + # Staging placement: see _get_shared_storage_cls. + sdVs = [ + cute.make_tensor( + cute.recast_ptr(sdOt.iterator + offset, sdV_layout_staged.inner, self.dtype_acc), + sdV_layout_staged.outer, + ) + for offset in self.sdV_split_offsets + ] sScaleP = storage.sScaleP.get_tensor(sScaleP_layout_staged) sdPsum = storage.sdPsum.get_tensor(sdPsum_layout_staged) @@ -1053,7 +1112,7 @@ def make_pipeline(cls, mbar_ptr, num_stages, producer, consumer, tx_count=None): self.dVacc_store( mIndexTopk, mdV, - sdV, + sdVs, tdVtdV0, tdVtdV1, thr_mma_PtdOt, @@ -1972,7 +2031,7 @@ def dVacc_store( self, mIndexTopk: cute.Tensor, mdV: cute.Tensor, - sdV: cute.Tensor, + sdVs: list[cute.Tensor], tdVtdV0: cute.Tensor, tdVtdV1: cute.Tensor, thr_mma_PtdOt: cute.ThrMma, @@ -2020,15 +2079,14 @@ def dVacc_store( tiled_copy_r2s = tiled_copy_2d(self.dtype_acc, 4, 64) thr_copy_r2s = tiled_copy_r2s.get_slice(tidx % 64) - # ((4,1),1,8,(1,8)):((1,0),0,4,(0,2048)) - tRS_sdV = thr_copy_r2s.partition_D(sdV) - tiled_copy_s2r = copy_utils.tiled_copy_2d(self.dtype_acc, 8, self.num_epilogue_threads, 4) thr_copy_s2r = tiled_copy_s2r.get_slice(tidx) - # (V, M, N, STAGE) - tSR_sdV = thr_copy_s2r.partition_S(sdV) + # ((4,1),1,8,(1,8)):((1,0),0,4,(0,2048)) and (V, M, N, STAGE), per split + tRS_sdVs = [thr_copy_r2s.partition_D(sdV) for sdV in sdVs] + tSR_sdVs = [thr_copy_s2r.partition_S(sdV) for sdV in sdVs] + tRS_sdV, tSR_sdV = tRS_sdVs[0], tSR_sdVs[0] - cdV = cute.make_identity_tensor(cute.product_each(sdV.shape[:2])) + cdV = cute.make_identity_tensor(cute.product_each(sdVs[0].shape[:2])) # (V, M, N) tdVcdV = thr_copy_s2r.partition_S(cdV) @@ -2086,18 +2144,18 @@ def dVacc_store( tRS_rdV_cur = cute.make_tensor(tdVrdV_cur.iterator, tRS_rdV_cur_shape) - stage = 4 * split + 2 * wg_half + (i % 2) - cute.copy(tiled_copy_r2s, tRS_rdV_cur, tRS_sdV[None, None, None, stage]) + stage = 2 * wg_half + (i % 2) + cute.copy(tiled_copy_r2s, tRS_rdV_cur, tRS_sdVs[split][None, None, None, stage]) cute.arch.fence_view_async_shared() self.epi_barrier.arrive_and_wait() tSR_rdV = cute.make_rmem_tensor(tdVrdV_out_shape, dtype=self.dtype_acc) for w in cutlass.range_constexpr(2): - stage_out = 4 * split + 2 * w + (i % 2) + stage_out = 2 * w + (i % 2) cute.copy( tiled_copy_s2r, - tSR_sdV[None, None, None, stage_out], + tSR_sdVs[split][None, None, None, stage_out], tSR_rdV[None, None, None, w], ) diff --git a/flash_attn/cute/flash_fwd_mla_sm100.py b/flash_attn/cute/flash_fwd_mla_sm100.py index d3782302dea..23d61326382 100644 --- a/flash_attn/cute/flash_fwd_mla_sm100.py +++ b/flash_attn/cute/flash_fwd_mla_sm100.py @@ -18,7 +18,12 @@ from quack import copy_utils -from flash_attn.cute.pack_gqa import pack_gqa_layout, make_packgqa_tiled_tma_atom +from flash_attn.cute.pack_gqa import ( + pack_gqa_layout, + make_packgqa_tiled_tma_atom, + padded_qheads_tma_source, + regroup_padded_qheads, +) from flash_attn.cute.paged_kv import PagedKVManager from flash_attn.cute import utils as fa_utils from flash_attn.cute.seqlen_info import SeqlenInfoQK @@ -62,12 +67,20 @@ def __init__( disable_bitmask: bool = False, use_clc_scheduler: bool = True, has_qk: bool = True, + qhead_per_kvhead_valid: Optional[int] = None, ): self.is_causal = is_causal self.is_local = False self.pack_gqa = pack_gqa self.qhead_per_kvhead = qhead_per_kvhead assert qhead_per_kvhead <= 128 + # Head padding: see pack_gqa.padded_qheads_tma_source. + if qhead_per_kvhead_valid is None: + qhead_per_kvhead_valid = qhead_per_kvhead + assert 0 < qhead_per_kvhead_valid <= qhead_per_kvhead + assert qhead_per_kvhead_valid == qhead_per_kvhead or pack_gqa + self.qhead_per_kvhead_valid = qhead_per_kvhead_valid + self.pad_qheads = qhead_per_kvhead_valid != qhead_per_kvhead self.nheads_kv = nheads_kv self.use_tma_O = True self.use_cpasync_load_KV = use_cpasync_load_KV @@ -76,7 +89,7 @@ def __init__( self.is_topk_gather = is_topk_gather if is_topk_gather: assert pack_gqa - assert qhead_per_kvhead == 128, "require MQA 128 for DSA path" + assert qhead_per_kvhead == 128, "DSA path tiles one token x 128 packed Q heads" assert use_cpasync_load_KV # user-provided option if topk indices guaranteed in bounds self.disable_bitmask = disable_bitmask @@ -250,6 +263,13 @@ def __init__( f"Total TMEM columns allocated {self.total_tmem} exceeds capacity {self.tmem_alloc_cols}" ) + @cute.jit + def is_valid_qhead_row(self, packed_row) -> Boolean: + """Head guard for non-TMA accesses; see pack_gqa.padded_qheads_tma_source.""" + if const_expr(not self.pad_qheads): + return Boolean(True) + return packed_row % self.qhead_per_kvhead < self.qhead_per_kvhead_valid + def _get_shared_storage_cls(self): self.buffer_align_bytes = 1024 @@ -456,6 +476,14 @@ def __call__( mO_og = mO mP_og = mP + # TMA source contract: see pack_gqa.padded_qheads_tma_source. + if const_expr(self.pad_qheads): + mQ_valid, mQv_valid, mO_valid, mP_valid = [ + padded_qheads_tma_source(mX, self.qhead_per_kvhead_valid, head_idx=2) + if mX is not None + else None + for mX in (mQ, mQv, mO, mP) + ] if const_expr(self.pack_gqa): mQ, mQv, mO, mP, mRowMax = [ pack_gqa_layout(mX, self.qhead_per_kvhead, self.nheads_kv, head_idx=2) @@ -465,6 +493,8 @@ def __call__( ] if const_expr(mLSE is not None): mLSE = pack_gqa_layout(mLSE, self.qhead_per_kvhead, self.nheads_kv, head_idx=1) + if const_expr(not self.pad_qheads): + mQ_valid, mQv_valid, mO_valid, mP_valid = mQ, mQv, mO, mP # ==== Prepare MMAs ==== # (local_var, dtype_a, major_a, major_b, mma_tiler, operand_source_a) @@ -525,16 +555,19 @@ def __call__( ) cta_shape = cta_layout_vmnk.shape - def make_tma(make_fn, mX, smem_layout, mma_tiler, tiled_mma): - return make_fn(tma_load_op, mX, smem_layout, mma_tiler, tiled_mma, cta_shape) + def make_tma(make_fn, mX, smem_layout, mma_tiler, tiled_mma, packed_qheads): + atom, tensor = make_fn(tma_load_op, mX, smem_layout, mma_tiler, tiled_mma, cta_shape) + if const_expr(packed_qheads and self.pad_qheads): + tensor = regroup_padded_qheads(tensor, self.qhead_per_kvhead, head_idx=2) + return atom, tensor A, B = cute.nvgpu.make_tiled_tma_atom_A, cute.nvgpu.make_tiled_tma_atom_B # (atom_name, tensor_name, make_fn, m, smem_layout, mma_tiler, tiled_mma, kv_only) # fmt: off _tma_specs = [ - ("tma_atom_Q", "tma_tensor_Q", A, mQ, self.sQ_layout, self.mma_tiler_QK, tiled_mma_QK, False), - ("tma_atom_Qv", "tma_tensor_Qv", A, mQv, self.sQv_layout, self.mma_tiler_QvV, tiled_mma_QvV, False), + ("tma_atom_Q", "tma_tensor_Q", A, mQ_valid, self.sQ_layout, self.mma_tiler_QK, tiled_mma_QK, False), + ("tma_atom_Qv", "tma_tensor_Qv", A, mQv_valid, self.sQv_layout, self.mma_tiler_QvV, tiled_mma_QvV, False), ("tma_atom_K", "tma_tensor_K", B, mK, self.sK_layout, self.mma_tiler_QK, tiled_mma_QK, True), ("tma_atom_V", "tma_tensor_V", B, mV, self.sV_layout, self.mma_tiler_QvV, tiled_mma_QvV, True), ("tma_atom_Vt", "tma_tensor_Vt", B, mVt, self.sVt_layout, self.mma_tiler_PVt, tiled_mma_PVt, True), @@ -542,7 +575,7 @@ def make_tma(make_fn, mX, smem_layout, mma_tiler, tiled_mma): _tmas = {} for atom_name, tensor_name, make_fn, m, smem_layout, mma_tiler, tiled_mma, kv_only in _tma_specs: _tmas[atom_name], _tmas[tensor_name] = ( - make_tma(make_fn, m, smem_layout, mma_tiler, tiled_mma) + make_tma(make_fn, m, smem_layout, mma_tiler, tiled_mma, packed_qheads=not kv_only) if const_expr((not kv_only or self.use_tma_KV) and m is not None) else (None, None) ) @@ -561,11 +594,17 @@ def make_tma(make_fn, mX, smem_layout, mma_tiler, tiled_mma): and self.pack_gqa and self.cta_tile_m % self.qhead_per_kvhead == 0 ) - make_tiled_tma_atom_fn = ( - partial(make_packgqa_tiled_tma_atom, qhead_per_kvhead=self.qhead_per_kvhead, head_idx=2) - if const_expr(self.ragged_tma_O) - else cpasync.make_tiled_tma_atom - ) + assert not (self.ragged_tma_O and self.pad_qheads) + + def make_tiled_tma_atom_fn(op, mX, smem_layout, tiler): + if const_expr(self.ragged_tma_O): + return make_packgqa_tiled_tma_atom( + op, mX, smem_layout, tiler, qhead_per_kvhead=self.qhead_per_kvhead, head_idx=2 + ) + atom, tensor = cpasync.make_tiled_tma_atom(op, mX, smem_layout, tiler) + if const_expr(self.pad_qheads): + tensor = regroup_padded_qheads(tensor, self.qhead_per_kvhead, head_idx=2) + return atom, tensor # ==== Set up P smem -> gmem tma store ==== @@ -576,7 +615,7 @@ def make_tma(make_fn, mX, smem_layout, mma_tiler, tiled_mma): if const_expr(self.store_P): # TODO: add asserts - mP_tma = mP_og if const_expr(self.ragged_tma_O) else mP + mP_tma = mP_og if const_expr(self.ragged_tma_O) else mP_valid if const_expr(self.ragged_tma_O): mP_tma = copy_utils.create_ragged_tensor_for_tma( mP_tma, ragged_dim=0, ptr_shift=True @@ -600,7 +639,7 @@ def make_tma(make_fn, mX, smem_layout, mma_tiler, tiled_mma): ) if const_expr(self.use_tma_O): - mO_tma = mO_og if const_expr(self.ragged_tma_O) else mO + mO_tma = mO_og if const_expr(self.ragged_tma_O) else mO_valid if const_expr(self.ragged_tma_O): mO_tma = copy_utils.create_ragged_tensor_for_tma( mO_tma, ragged_dim=0, ptr_shift=True @@ -2658,6 +2697,7 @@ def softmax_loop( warp_idx, store_P=store_P, gRowMax=gRowMax, + packed_row0=cta_m_block * self.cta_tile_m, ) ### first iteration ### @@ -2794,6 +2834,7 @@ def softmax_step( is_first: Boolean = False, store_P: Optional[Callable] = None, gRowMax: Optional[cute.Tensor] = None, + packed_row0: Int32 = 0, ): leader_warp = warp_idx == 0 tSrP = cute.make_rmem_tensor(tSrS_t2r.shape, self.dtype_P) @@ -2833,7 +2874,7 @@ def softmax_step( row_max, acc_scale = softmax.update_row_max_from_local(row_max, is_first) if const_expr(gRowMax is not None): - if tidx < self.cta_tile_m: + if tidx < self.cta_tile_m and self.is_valid_qhead_row(packed_row0 + tidx): gRowMax[tidx, n_block] = row_max # note: acc_scales agree for paired threads @@ -3059,6 +3100,7 @@ def correction_loop( cta_m_block * self.cta_tile_m + tidx % self.cta_tile_m, self.qhead_per_kvhead, self.pack_gqa, + self.qhead_per_kvhead_valid if const_expr(self.pad_qheads) else None, ) row_max, row_sum = apply_learnable_sink(row_max, row_sum, sink_val, softmax_scale_log2) @@ -3091,7 +3133,9 @@ def correction_loop( if not acc_O_mn_row_is_zero_or_nan else -Float32.inf ) - if tidx < seqlen_q - cta_m_block * self.cta_tile_m: + if tidx < seqlen_q - cta_m_block * self.cta_tile_m and self.is_valid_qhead_row( + cta_m_block * self.cta_tile_m + tidx + ): gLSE[tidx] = lse row_idx = cta_m_block * self.cta_tile_m + tOicOi[0][0] diff --git a/flash_attn/cute/interface.py b/flash_attn/cute/interface.py index b6257bc43ce..577c676b98d 100644 --- a/flash_attn/cute/interface.py +++ b/flash_attn/cute/interface.py @@ -696,6 +696,12 @@ def _flash_attn_fwd( if softcap == 0.0: softcap = None qhead_per_kvhead = num_head // num_head_kv + # Sparse MLA head padding: see pack_gqa.padded_qheads_tma_source. + qhead_per_kvhead_valid = qhead_per_kvhead + if qv is not None and gather_kv_indices is not None and qhead_per_kvhead < 128: + assert num_head_kv == 1, "sparse MLA requires a single KV head" + qhead_per_kvhead = 128 + pack_gqa = True if pack_gqa is None: pack_gqa = qhead_per_kvhead > 1 @@ -1127,6 +1133,7 @@ def _flash_attn_fwd( head_dim, head_dim_v, qhead_per_kvhead, + qhead_per_kvhead_valid, causal, score_mod_hash, mask_mod_hash, @@ -1324,6 +1331,7 @@ def _flash_attn_fwd( has_cu_seqlens_q=cu_seqlens_q is not None, disable_bitmask=disable_sparse_kv_bitmask, has_qk=has_qk, + qhead_per_kvhead_valid=qhead_per_kvhead_valid, ) else: if use_dedicated_hd256_kernel: @@ -2743,7 +2751,13 @@ def _flash_attn_bwd_sparse_mla( nheads_kv, head_dim_v = v.shape[-2:] qhead_per_kvhead = nheads // nheads_kv gather_kv_length = gather_kv_indices.shape[-1] - assert nheads_kv == 1 and qhead_per_kvhead == 128, f"sparse MLA bwd: only MQA 128 supported for now" + assert nheads_kv == 1 and 0 < qhead_per_kvhead <= 128, ( + f"sparse MLA bwd: MQA with at most 128 heads, got {qhead_per_kvhead}" + ) + # Backward head padding: see pack_gqa.padded_qheads_tma_source. + qhead_per_kvhead_valid = qhead_per_kvhead + qhead_per_kvhead = 64 if qhead_per_kvhead <= 64 else 128 + pad_qheads = qhead_per_kvhead != qhead_per_kvhead_valid assert gather_kv_length % 128 == 0, f"sparse MLA bwd: {gather_kv_length=} must be divisible by 128" assert deterministic is False, "sparse MLA bwd: deterministic mode not yet supported" assert seqused_q is None and seqused_k is None, "sparse MLA bwd: seqused_q,k not yet supported" @@ -2819,26 +2833,28 @@ def _flash_attn_bwd_sparse_mla( _validate_tensor(dqv, "dqv", qv.shape, dtype, device) _validate_tensor(p, "p", p_shape, dtype, device) - if cu_seqlens_q is None: - dpsum = torch.empty(batch_size, seqlen_q, nheads, dtype=torch.float32, device=device) - else: - dpsum = torch.empty(total_q, nheads, dtype=torch.float32, device=device) - scale_p = torch.empty_like(row_max) + # Finite tile-width padding: see pack_gqa.padded_qheads_tma_source. + heads_shape = (batch_size, seqlen_q) if cu_seqlens_q is None else (total_q,) + alloc = torch.zeros if pad_qheads else torch.empty + dpsum = alloc(*heads_shape, qhead_per_kvhead, dtype=torch.float32, device=device) + scale_p = alloc(*row_max.shape[:-1], qhead_per_kvhead, dtype=torch.float32, device=device) dtype = torch2cute_dtype_map[dout.dtype] # Preprocess kernel: compute (o * dout).sum(dim=-1), scale_p. + # Padded counts use trivial packing: not all head counts divide the 128-row tile. + # Non-power-of-two tiles (e.g. 48 rows for 24 heads) produced incorrect dpsum. _bwd_preprocess( - out, dout, dpsum, lse, None, None, + out, dout, dpsum[..., :nheads], lse, None, None, cu_seqlens_q, seqused_q, None, dtype, head_dim, head_dim_v, m_block_size, row_max=row_max, - scale_p=scale_p, + scale_p=scale_p[..., :nheads], use_padded_offsets=False, nheads_major=True, pack_gqa=True, - qhead_per_kvhead=qhead_per_kvhead, - nheads_kv=nheads_kv, + qhead_per_kvhead=1 if pad_qheads else qhead_per_kvhead, + nheads_kv=nheads if pad_qheads else nheads_kv, softmax_scale=softmax_scale, fake_mode=fake_mode, ) @@ -2848,6 +2864,7 @@ def _flash_attn_bwd_sparse_mla( head_dim, head_dim_v, qhead_per_kvhead, + qhead_per_kvhead_valid, causal, cu_seqlens_q is None, cu_seqlens_k is None, @@ -2890,6 +2907,7 @@ def _flash_attn_bwd_sparse_mla( nheads_kv=nheads_kv, has_seqused_q=seqused_q is not None, disable_bitmask=disable_sparse_kv_bitmask, + qhead_per_kvhead_valid=qhead_per_kvhead_valid, ) fa_bwd_kernel = cute.compile( fa_bwd_obj, diff --git a/flash_attn/cute/pack_gqa.py b/flash_attn/cute/pack_gqa.py index 5b481b5e6fc..c6defd7cbc6 100644 --- a/flash_attn/cute/pack_gqa.py +++ b/flash_attn/cute/pack_gqa.py @@ -40,6 +40,50 @@ def pack_gqa_layout(T, qhead_per_kvhead, nheads_kv, head_idx): return cute.make_tensor(T.iterator, cute.make_layout(shape_packed, stride=stride_packed)) +def _heads_first_order(T, head_idx): + """Mode permutation swapping the seqlen mode (0) with the head mode; its own inverse.""" + return [head_idx, *range(1, head_idx), 0, *range(head_idx + 1, cute.rank(T))] + + +def padded_qheads_tma_source(T, qhead_per_kvhead_valid, head_idx): + """Return a heads-first TMA source with a dynamic extent equal to the real head count. + + .. note:: In-kernel Q-head padding (sparse MLA). + MQA only. Each tile covers one token and one top-k gather list: 128 heads in + forward and dQ/dQv, 64 or 128 in backward. The heads-first view + ``(nheads, ..., seqlen, ...)`` has a dynamic head extent so CuTe can tile it + without requiring divisibility. TMA zero-fills out-of-bounds loads and drops + out-of-bounds stores, avoiding padded operand copies in global memory. + ``regroup_padded_qheads`` folds the TMA coordinate tensor into the packed + ``(qhead, seqlen)`` layout. Non-TMA LSE, row_max, and learnable-sink accesses + need head guards: the hierarchical packed layout wraps padded heads into + the next token. The caller allocates dPsum/scaleP at tile width with finite + padding. TMA loads zero padded Q rows in forward and dO/P rows in backward. + These rows produce dS = 0 and contribute nothing to dK/dV. + """ + T = cute.make_tensor(T.iterator, cute.select(T.layout, mode=_heads_first_order(T, head_idx))) + shape = (cutlass.Int32(qhead_per_kvhead_valid), *T.shape[1:]) + return cute.make_tensor(T.iterator, cute.make_layout(shape, stride=T.stride)) + + +def regroup_padded_qheads(tma_tensor, qhead_per_kvhead, head_idx): + """Fold a ``padded_qheads_tma_source`` coordinate tensor into the pack-GQA layout + ``((qhead_per_kvhead, seqlen), ..., 1, ...)``: ``pack_gqa_layout`` with the tile's head + count, after undoing the heads-first permutation. + + Not ``pack_gqa_layout`` itself: its KV-head stride ``head_stride * qhead_per_kvhead`` + scales a TMA basis stride, which the DSL fails to lower (ICE). The size-1 KV-head mode + reuses the head basis instead. + """ + T = tma_tensor + order = _heads_first_order(T, head_idx) + shape = [T.shape[i] for i in order] + stride = [T.stride[i] for i in order] + shape[0], stride[0] = (qhead_per_kvhead, shape[0]), (stride[head_idx], stride[0]) + shape[head_idx] = 1 + return cute.make_tensor(T.iterator, cute.make_layout(tuple(shape), stride=tuple(stride))) + + def make_packgqa_tiled_tma_atom( op: cute.atom.CopyOp, gmem_tensor: cute.Tensor, diff --git a/flash_attn/cute/softmax.py b/flash_attn/cute/softmax.py index 4490d2b4da6..dd25de46cf0 100644 --- a/flash_attn/cute/softmax.py +++ b/flash_attn/cute/softmax.py @@ -2,7 +2,7 @@ import math import operator -from typing import Tuple +from typing import Optional, Tuple from dataclasses import dataclass import cutlass @@ -23,16 +23,24 @@ def load_learnable_sink( packed_row, qhead_per_kvhead: cutlass.Constexpr[int], pack_gqa: cutlass.Constexpr[bool], + qhead_per_kvhead_valid: cutlass.Constexpr[Optional[int]] = None, ) -> Float32: """Load the sink logit for one output row; see `apply_learnable_sink`. With pack_gqa the M tile interleaves the Q heads of one KV head, so `head_idx` is the KV head and the Q head is `packed_row % qhead_per_kvhead`; otherwise `head_idx` is already the Q head and `packed_row` is ignored. + Padded Q heads get -inf (no sink); see `pack_gqa.padded_qheads_tma_source`. """ if cutlass.const_expr(not pack_gqa): return Float32(learnable_sink[head_idx]) - return Float32(learnable_sink[head_idx * qhead_per_kvhead + packed_row % qhead_per_kvhead]) + if cutlass.const_expr(qhead_per_kvhead_valid is None): + qhead_per_kvhead_valid = qhead_per_kvhead + qhead = packed_row % qhead_per_kvhead + sink_val = -Float32.inf + if qhead < qhead_per_kvhead_valid: + sink_val = Float32(learnable_sink[head_idx * qhead_per_kvhead_valid + qhead]) + return sink_val @cute.jit diff --git a/tests/cute/test_flash_attn.py b/tests/cute/test_flash_attn.py index 74a8fe34f8a..f4c0ddfcf48 100644 --- a/tests/cute/test_flash_attn.py +++ b/tests/cute/test_flash_attn.py @@ -3376,12 +3376,16 @@ def check_canary(name, parent, pad_words): @pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize("causal", [False, True]) @pytest.mark.parametrize("shared_kv", [False, True]) +@pytest.mark.parametrize("nheads", [128, 64, 96, 24, 1]) @pytest.mark.parametrize("seqlen_q,seqlen_k", [(512, 512), (1024, 1024)]) @maybe_fake_tensor_mode(USE_FAKE_TENSOR) -def test_flash_attn_mla_sparse_bwd_sentinel(seqlen_q, seqlen_k, shared_kv, causal, dtype): +def test_flash_attn_mla_sparse_bwd_sentinel(seqlen_q, seqlen_k, nheads, shared_kv, causal, dtype): """Sparse-MLA backward with -1-padded gather_kv_indices, the padding any causal top-k selector produces for early queries. + nheads < 128 covers in-kernel head padding (pack_gqa.padded_qheads_tma_source): + 96 pads to 128 with a partial second CTA, 24 and 1 pad to the 64-head bwd tile. + Regression test for unguarded sentinel scatters: the dV/dK backward epilogues used to atomically accumulate at row -1 — out of bounds of the (batch-sliced) buffer — corrupting adjacent memory even though the addend @@ -3396,7 +3400,7 @@ def test_flash_attn_mla_sparse_bwd_sentinel(seqlen_q, seqlen_k, shared_kv, causa device = "cuda" torch.random.manual_seed(0) batch_size = 2 - nheads, nheads_kv, hdim, hdimv = 128, 1, 64, 512 + nheads_kv, hdim, hdimv = 1, 64, 512 topk_len = 256 q_ref = torch.randn(batch_size, seqlen_q, nheads, hdim, device=device, dtype=dtype).requires_grad_()