From 2407955bd1c3385da852321746f3ba79a1db7da4 Mon Sep 17 00:00:00 2001 From: Kate Cheng Date: Wed, 23 Sep 2026 22:12:27 -0700 Subject: [PATCH 01/13] Add SM90 FP8 MMA forward for head dimension 512 Signed-off-by: Kate Cheng --- flash_attn/cute/flash_fwd_sm90.py | 488 +++++++++++++++++++++++------- flash_attn/cute/interface.py | 87 ++++-- tests/cute/test_flash_attn.py | 433 ++++++++++++++++++++++++-- 3 files changed, 863 insertions(+), 145 deletions(-) diff --git a/flash_attn/cute/flash_fwd_sm90.py b/flash_attn/cute/flash_fwd_sm90.py index df449df02e8..798e78921f6 100644 --- a/flash_attn/cute/flash_fwd_sm90.py +++ b/flash_attn/cute/flash_fwd_sm90.py @@ -66,13 +66,21 @@ def __init__( f"Fused quant output not implemented for {type(self).__name__}" ) self.intra_wg_overlap = intra_wg_overlap + # Match the FA3 FP8 path: keep softmax probabilities in registers for PV. + # SM90 FP8 needs a different accumulator-to-A-fragment mapping than FP16. self.mma_pv_is_rs = mma_pv_is_rs self.buffer_align_bytes = 1024 self.use_tma_KV = not paged_kv_non_tma self.is_split_kv = is_split_kv - # FP8-KV scale-fold path (SM90: fp16 Q + (fp8-e4m3 -> fp16) paged K/V -> fp16 O). + # SM90 compensated FP8-MMA path: native fp16/bf16 Q, E4M3 paged K/V, + # FP8 WGMMA QK/PV, FP32 accumulation, and native fp16/bf16 O. self.fp8_kv_dequant = fp8_kv_dequant self.kv_dtype = kv_dtype if kv_dtype is not None else self.dtype + # Hopper FP8 WGMMA needs FP8 on both inputs. Q is quantized once per output + # tile. P is scaled before its FP8 cast to preserve long-context softmax + # values; the reciprocal is folded into final output normalization. + self.fp8_q_scale = 1.0 + self.fp8_p_scale = 256.0 assert self.use_tma_KV or not (self.check_hdim_oob or self.check_hdim_v_oob), ( "Paged KV does not support irregular head dim" ) @@ -105,9 +113,10 @@ def _get_smem_layout_atom(self): def _get_tiled_mma(self): self.atom_layout_n = 2 if self.tile_hdim > 256 or self.tile_hdimv > 256 else 1 + mma_dtype = self.kv_dtype if self.fp8_kv_dequant else self.dtype tiled_mma_qk = sm90_utils_basic.make_trivial_tiled_mma( - self.dtype, - self.dtype, + mma_dtype, + mma_dtype, warpgroup.OperandMajorMode.K, warpgroup.OperandMajorMode.K, Float32, @@ -115,10 +124,12 @@ def _get_tiled_mma(self): tiler_mn=(64, self.tile_n), ) tiled_mma_pv = sm90_utils_basic.make_trivial_tiled_mma( - self.dtype, - self.dtype, + mma_dtype, + mma_dtype, warpgroup.OperandMajorMode.K, - warpgroup.OperandMajorMode.MN, + warpgroup.OperandMajorMode.K + if self.fp8_kv_dequant + else warpgroup.OperandMajorMode.MN, Float32, atom_layout_mnk=( self.tile_m // 64, @@ -149,7 +160,22 @@ def _get_shared_storage_cls(self): mbar_ptr_V_struct = cute.struct.MemRange[cutlass.Int64, self.num_stages * 2] if const_expr(self.fp8_kv_dequant): - # sStage: the fp8 staging buffer; its num_stages=1 TMA pipeline gates GMEM->sStage. + sQ_input_struct = cute.struct.Align[ + cute.struct.MemRange[self.q_dtype, cute.cosize(self.sQ_layout)], + self.buffer_align_bytes, + ] + sQ_mma_struct = cute.struct.Align[ + cute.struct.MemRange[self.kv_dtype, cute.cosize(self.sQ_mma_layout)], + self.buffer_align_bytes, + ] + sK_mma_struct = cute.struct.Align[ + cute.struct.MemRange[self.kv_dtype, cute.cosize(self.sK_layout)], + self.buffer_align_bytes, + ] + sV_mma_struct = cute.struct.Align[ + cute.struct.MemRange[self.kv_dtype, cute.cosize(self.sV_mma_layout)], + self.buffer_align_bytes, + ] sStage_struct = cute.struct.Align[ cute.struct.MemRange[self.kv_dtype, cute.cosize(self.sStage_layout)], self.buffer_align_bytes, @@ -157,17 +183,16 @@ def _get_shared_storage_cls(self): mbar_ptr_Stage_struct = cute.struct.MemRange[cutlass.Int64, 1 * 2] @cute.struct - class SharedStorageQKVDequant: + class SharedStorageQKVFp8Mma: mbar_ptr_Q: mbar_ptr_Q_struct - mbar_ptr_K: mbar_ptr_K_struct - mbar_ptr_V: mbar_ptr_V_struct mbar_ptr_Stage: mbar_ptr_Stage_struct - sV: sV_struct - sQ: sQ_struct - sK: sK_struct + sQ: sQ_input_struct + sQ_mma: sQ_mma_struct + sK: sK_mma_struct + sV_mma: sV_mma_struct sStage: sStage_struct - return SharedStorageQKVDequant + return SharedStorageQKVFp8Mma @cute.struct class SharedStorageQKV: @@ -221,7 +246,7 @@ def __call__( """Configures and launches the flash attention kernel. mQ/mK/mV/mO has same data types(supports fp16 and bf16) and same layout - (except the fp8-KV scale-fold path: fp16 Q + fp8 e4m3 K/V -> fp16 O): + (except FP8 MMA: native fp16/bf16 Q, E4M3 K/V, native fp16/bf16 O): (batch_size, seqlen_q, num_head, head_dim):(_, _, _, 1) """ if const_expr(not self.fp8_kv_dequant): @@ -234,10 +259,11 @@ def __call__( ) else: # FP8-KV: K/V are fp8 (rejected by the symmetric _check_type). Q/O may each be - # fp16 OR bf16 -- compute is always fp16 (the single-instruction fp8->fp16 fast - # path). Q is narrowed bf16->fp16 in-kernel and the fp32 accumulator is cast to - # mO's dtype in the epilogue, so the Q/O tensor dtypes are independent of compute. - assert self.dtype == Float16, "fp8_kv_dequant computes in fp16" + # fp16 or bf16. Q and softmax P use compensated pairs of E4M3 values; + # QK and PV both execute as FP8 WGMMA with fp32 accumulation. + assert self.dtype == Float16, ( + "fp8_kv_dequant requires fp16 auxiliary conversion layouts" + ) assert mQ.element_type in (Float16, BFloat16), "fp8_kv_dequant expects fp16 or bf16 Q" if const_expr(self.is_split_kv): assert mO.element_type == Float32, ( @@ -248,8 +274,7 @@ def __call__( "fp8_kv_dequant expects fp16 or bf16 O" ) - # Q input / O output dtypes (may differ from the fp16 compute dtype on the fp8-KV - # path). q_dtype gates the in-kernel bf16->fp16 Q narrow; o_dtype drives the + # Q input / O output dtypes. q_dtype drives Q quantization and o_dtype drives the # epilogue cast. SplitKV writes fp32 partials via acc_O directly, so o_dtype stays # the compute dtype there (the combine kernel produces the final bf16/fp16 O). self.q_dtype = mQ.element_type @@ -284,6 +309,10 @@ def __call__( self.num_mma_threads = tiled_mma_qk.size self.num_threads_per_warp_group = 128 self.num_wg_mma = self.num_mma_threads // self.num_threads_per_warp_group + if const_expr(self.fp8_kv_dequant): + assert self.num_wg_mma == 2, ( + "SM90 d512 FP8 MMA requires exactly two MMA warp groups" + ) assert self.num_wg_mma in [1, 2, 3] self.num_threads = self.num_threads_per_warp_group * (self.num_wg_mma + 1) self.num_producer_threads = 32 @@ -311,10 +340,8 @@ def __call__( self.p_shared_across_wg = self.atom_layout_n == 2 self._setup_attributes() # TODO: we prob don't need most of what's in _setup_attributes - # Per-tensor smem dtypes. FP8-KV is the one special case: sK/sV hold the fp16 - # dequant target (self.dtype), not mK/mV's fp8 source. - sK_dtype = self.dtype if const_expr(self.fp8_kv_dequant) else mK.element_type - sV_dtype = self.dtype if const_expr(self.fp8_kv_dequant) else mV.element_type + sK_dtype = self.kv_dtype if const_expr(self.fp8_kv_dequant) else mK.element_type + sV_dtype = self.kv_dtype if const_expr(self.fp8_kv_dequant) else mV.element_type self.sQ_layout, self.sK_layout, self.sV_layout, self.sO_layout = [ sm90_utils.make_smem_layout(dtype, LayoutEnum.ROW_MAJOR, shape, stage) for dtype, shape, stage in [ @@ -328,9 +355,21 @@ def __call__( (self.o_dtype, (self.tile_m, self.tile_hdimv), None), ] ] - # fp8 staging tile: ONE shared buffer (tile_hdim == tile_hdimv so it serves K and V). + self.sQ_mma_layout = None + self.sV_mma_layout = None self.sStage_layout = None if const_expr(self.fp8_kv_dequant): + self.sQ_mma_layout = sm90_utils.make_smem_layout( + self.kv_dtype, LayoutEnum.ROW_MAJOR, (self.tile_m, self.tile_hdim), None + ) + # Construct WGMMA-B directly as (output_dim, reduction_dim), K-major. + self.sV_mma_layout = sm90_utils.make_smem_layout( + self.kv_dtype, + LayoutEnum.ROW_MAJOR, + (self.tile_hdimv, self.tile_n), + self.num_stages, + ) + # One row-major staging tile is shared by K and V TMA loads. self.sStage_layout = sm90_utils.make_smem_layout( self.kv_dtype, LayoutEnum.ROW_MAJOR, (self.tile_n, self.tile_hdim), 1 ) @@ -378,7 +417,8 @@ def __call__( tma_atom_K, tma_tensor_K = None, None tma_atom_V, tma_tensor_V = None, None if const_expr(self.use_tma_KV): - # FP8-KV: K and V from GMEM land in sStage for the consumer to dequant(instead of sK/sV) + # FP8 K/V land in sStage for the consumer's WGMMA-layout + # rearrangement instead of being widened into sK/sV. sK_tma_box = cute.select( self.sStage_layout if const_expr(self.fp8_kv_dequant) else self.sK_layout, mode=[0, 1], @@ -487,6 +527,8 @@ def __call__( self.sV_layout, self.sO_layout, self.sP_layout, + self.sQ_mma_layout, + self.sV_mma_layout, self.sStage_layout, self.gmem_tiled_copy_Q, self.gmem_tiled_copy_K, @@ -538,6 +580,8 @@ def kernel( sV_layout: cute.ComposedLayout, sO_layout: cute.ComposedLayout, sP_layout: cute.ComposedLayout | None, + sQ_mma_layout: cute.ComposedLayout | None, + sV_mma_layout: cute.ComposedLayout | None, sStage_layout: cute.ComposedLayout | None, gmem_tiled_copy_Q: cute.TiledCopy, gmem_tiled_copy_K: cute.TiledCopy, @@ -651,28 +695,37 @@ def kernel( # typed with the compute dtype, so the TMA'd bf16 Q must be read back as q_dtype # (passing dtype here, like sO below) -- otherwise bf16 bytes get read as fp16. sQ = storage.sQ.get_tensor(sQ_layout.outer, swizzle=sQ_layout.inner, dtype=self.q_dtype) - # FP8-KV bf16 Q: fp16 compute view aliasing sQ's smem (both 2 bytes, so the - # element layout/swizzle is identical -- reinterpret the same storage as fp16, the - # same way sO reuses sQ's storage below). _narrow_q_to_compute() casts sQ bf16->fp16 - # in place per Q tile; the QK WGMMA reads sQ_mma. When Q is already fp16, sQ_mma is sQ. sQ_mma = sQ - if const_expr(self.fp8_kv_dequant and self.q_dtype != self.dtype): - sQ_mma = storage.sQ.get_tensor( - sQ_layout.outer, swizzle=sQ_layout.inner, dtype=self.dtype + sQ_res = sQ + if const_expr(self.fp8_kv_dequant): + sQ_mma = storage.sQ_mma.get_tensor( + sQ_mma_layout.outer, swizzle=sQ_mma_layout.inner + ) + # Reuse the now-dead FP16 Q input bytes for the FP8 residual term. + sQ_res = storage.sQ.get_tensor( + sQ_mma_layout.outer, swizzle=sQ_mma_layout.inner, dtype=self.kv_dtype ) sK = storage.sK.get_tensor(sK_layout.outer, swizzle=sK_layout.inner) - if const_expr(not self.Q_in_regs): + if const_expr(self.fp8_kv_dequant): + sVt = storage.sV_mma.get_tensor( + sV_mma_layout.outer, swizzle=sV_mma_layout.inner + ) + sV = layout_utils.transpose_view(sVt) + elif const_expr(not self.Q_in_regs): sV = storage.sV.get_tensor(sV_layout.outer, swizzle=sV_layout.inner) + sVt = layout_utils.transpose_view(sV) else: sV = storage.sQ.get_tensor( sV_layout.outer, swizzle=sV_layout.inner, dtype=mV.element_type ) - # Transpose view of V to tensor with layout (head_dim_v, tile_n) for tiled mma - sVt = layout_utils.transpose_view(sV) + sVt = layout_utils.transpose_view(sV) sP = None + sP_res = None if const_expr(sP_layout is not None): sP = storage.sP.get_tensor(sP_layout.outer, swizzle=sP_layout.inner) - # fp8 staging buffer (TMA dst, dequant src). + if const_expr(self.fp8_kv_dequant): + sP_res = storage.sP_res.get_tensor(sP_layout.outer, swizzle=sP_layout.inner) + # FP8 staging buffer: TMA destination and WGMMA-layout rearrangement source. sStage = None if const_expr(self.fp8_kv_dequant): sStage = storage.sStage.get_tensor(sStage_layout.outer, swizzle=sStage_layout.inner) @@ -766,8 +819,10 @@ def kernel( sK, sVt, sP, + sP_res, sO, sQ_mma, + sQ_res, learnable_sink, pipeline_k, pipeline_v, @@ -830,7 +885,8 @@ def load_fp8( head_idx // self.qhead_per_kvhead if const_expr(not self.pack_gqa) else head_idx ) - # Q stays fp16 (no dequant) + # Q remains in its native fp16/bf16 source dtype until compensated + # E4M3 preparation in the MMA warp group. load_Q = None pack_gqa = None if const_expr(self.use_tma_Q): @@ -1258,10 +1314,11 @@ def mma( sK: cute.Tensor, sVt: cute.Tensor, sP: Optional[cute.Tensor], + sP_res: Optional[cute.Tensor], sO: cute.Tensor, - # FP8-KV bf16 Q: fp16 compute view aliasing sQ (built in kernel() from the - # region-local sQ_layout). Equals sQ when Q is already fp16. + # FP8 high and residual Q terms used by the compensated QK WGMMA. sQ_mma: cute.Tensor, + sQ_res: cute.Tensor, learnable_sink: Optional[cute.Tensor], pipeline_k: pipeline.PipelineAsync, pipeline_v: pipeline.PipelineAsync, @@ -1279,8 +1336,7 @@ def mma( aux_data: AuxData = AuxData(), fastdiv_mods=None, num_splits: Int32 = Int32(1), - # FP8-KV cast inputs (None on the non-fp8 path): sV (fp16 V write target), - # sStage, pipeline_stage. + # FP8-KV staging inputs (None on the non-fp8 path). sV: Optional[cute.Tensor] = None, sStage: Optional[cute.Tensor] = None, pipeline_stage: Optional[pipeline.PipelineAsync] = None, @@ -1294,20 +1350,46 @@ def mma( thr_mma_qk = tiled_mma_qk.get_slice(tidx) wg_mma_qk = tiled_mma_qk.get_slice(warp_group_thread_layout(warp_group_idx)) wg_mma_pv = tiled_mma_pv.get_slice(warp_group_thread_layout(warp_group_idx)) - # FP8-KV bf16 Q: the QK WGMMA needs fp16 operands. sQ holds the bf16 Q as TMA'd; - # sQ_mma is the fp16 compute view (same smem) built in kernel(). _narrow_q_to_compute() - # casts sQ bf16->fp16 in place per tile before the QK reads sQ_mma. - narrow_q = const_expr(self.fp8_kv_dequant and self.q_dtype != self.dtype) + quantize_q = const_expr(self.fp8_kv_dequant) _, tSrQ, tSrK = sm90_utils.partition_fragment_ABC( wg_mma_qk, (self.tile_m, self.tile_n, self.tile_hdim), sQ_mma, sK ) - mma_qk_fn = partial( - sm90_utils.gemm_zero_init, tiled_mma_qk, (self.tile_m, self.tile_n), tSrQ, tSrK - ) + if const_expr(self.fp8_kv_dequant): + _, tSrQ_res, _ = sm90_utils.partition_fragment_ABC( + wg_mma_qk, (self.tile_m, self.tile_n, self.tile_hdim), sQ_res, sK + ) + mma_qk_fn = partial( + self._gemm_qk_fp8_compensated, + tiled_mma_qk, + (self.tile_m, self.tile_n), + tSrQ, + tSrQ_res, + tSrK, + ) + else: + mma_qk_fn = partial( + sm90_utils.gemm_zero_init, + tiled_mma_qk, + (self.tile_m, self.tile_n), + tSrQ, + tSrK, + ) acc_O, tOrP, tOrVt = sm90_utils.partition_fragment_ABC( wg_mma_pv, (self.tile_m, self.tile_hdimv, self.tile_n), sP, sVt ) - mma_pv_fn = partial(sm90_utils.gemm_w_idx, tiled_mma_pv, acc_O, tOrP, tOrVt) + tOrP_res = None + if const_expr(self.fp8_kv_dequant): + tOrP_res = cute.make_rmem_tensor_like(tOrP) + mma_pv_fn = partial( + self._gemm_pv_fp8_compensated, + tiled_mma_pv, + acc_O, + tOrP, + tOrP_res, + tOrVt, + ) + else: + mma_pv_fn = partial(sm90_utils.gemm_w_idx, tiled_mma_pv, acc_O, tOrP, tOrVt) # /////////////////////////////////////////////////////////////////////////////// # Smem copy atom tiling @@ -1317,7 +1399,9 @@ def mma( ) smem_thr_copy_P = cute.make_tiled_copy_C(smem_copy_atom_P, tiled_mma_qk).get_slice(tidx) tPsP = smem_thr_copy_P.partition_D(sP) if const_expr(sP is not None) else None - smem_copy_params = SimpleNamespace(smem_thr_copy_P=smem_thr_copy_P, tPsP=tPsP) + smem_copy_params = SimpleNamespace( + smem_thr_copy_P=smem_thr_copy_P, tPsP=tPsP, tOrP_res=tOrP_res + ) self.mma_init() @@ -1326,48 +1410,91 @@ def mma( pipeline.PipelineUserType.Consumer, self.num_stages ) - # FP8-KV MMA WGs side dequant setup: - # K: cooperative 256-thread(2 WGs) copy of the full tile; - # V: each WG copies only its own hdimv-half. + # FP8 staging -> FP8 WGMMA buffers. K keeps its orientation; V is explicitly + # transposed so the PV B operand is K-major as required by Hopper FP8 WGMMA. dequant_params = None if const_expr(self.fp8_kv_dequant): sStage2 = sStage[None, None, 0] sK2 = sK[None, None, 0] - sV2 = sV[None, None, 0] + sStageVt2 = layout_utils.transpose_view(sStage2) + sVt2 = sVt[None, None, 0] half = const_expr(min(256, self.tile_hdimv)) - VEC = const_expr(8) # 8 fp16 == 128-bit store; 8 fp8 == 64-bit load + VEC = const_expr(8) tiled_copy_K = copy_utils.tiled_copy_2d( - self.dtype, self.tile_hdim // VEC, self.num_mma_threads, num_copy_elems=VEC - ) - tiled_copy_V = copy_utils.tiled_copy_2d( - self.dtype, half // VEC, self.num_threads_per_warp_group, num_copy_elems=VEC + self.kv_dtype, + self.tile_hdim // VEC, + self.num_mma_threads, + num_copy_elems=VEC, ) thr_copy_K = tiled_copy_K.get_slice(tidx) - thr_copy_V = tiled_copy_V.get_slice(tidx % self.num_threads_per_warp_group) - # WG-local hdimv-half sub-views (tile coord (0, warp_group_idx)). - gStage_V = cute.local_tile(sStage2, (self.tile_n, half), (0, warp_group_idx)) - gV = cute.local_tile(sV2, (self.tile_n, half), (0, warp_group_idx)) + # Build both WG-local output-dimension halves with static coordinates. + # A dynamic local_tile coordinate loses the 128-bit alignment proof + # required by STSM even though both 256-column offsets are aligned. + gStage_V0 = cute.local_tile(sStageVt2, (half, self.tile_n), (0, 0)) + gStage_V1 = cute.local_tile(sStageVt2, (half, self.tile_n), (1, 0)) + gV0 = cute.local_tile(sVt2, (half, self.tile_n), (0, 0)) + gV1 = cute.local_tile(sVt2, (half, self.tile_n), (1, 0)) + + # Port of FA3's FP8 V transpose. LDSM.T reads 64x8 byte tiles as + # 16-bit matrices, PRMT interleaves their bytes, and STSM writes the + # K-major WGMMA layout without scalar shared-memory traffic. + ldsm_atom = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp(True, 4), self.kv_dtype + ) + stsm_atom = cute.make_copy_atom( + cute.nvgpu.warp.StMatrix8x8x16bOp(False, 4), self.kv_dtype + ) + ldsm_copy = cute.make_tiled_copy_tv( + ldsm_atom, + cute.make_layout((32, 4, 1, 1), stride=(4, 1, 0, 0)), + cute.make_layout((2, 2, 1, 4), stride=(1, 2, 16, 4)), + ) + stsm_copy = cute.make_tiled_copy_tv( + stsm_atom, + cute.make_layout((8, 4, 4, 1), stride=(4, 1, 32, 0)), + # FA3's faster STSM layout avoids shared-memory bank conflicts. + # It permutes V's output columns, so each completed PV must apply + # the inverse permutation to its accumulator registers. + cute.make_layout((1, 4, 2, 2), stride=(0, 1, 4, 8)), + ) + transpose_tid = tidx % self.num_threads_per_warp_group + ldsm_thr = ldsm_copy.get_slice(transpose_tid) + stsm_thr = stsm_copy.get_slice(transpose_tid) + tStage_V0 = ldsm_thr.partition_S(cute.flat_divide(gStage_V0, (64, 8))) + tStage_V1 = ldsm_thr.partition_S(cute.flat_divide(gStage_V1, (64, 8))) + tsV0 = stsm_thr.partition_D(cute.flat_divide(gV0, (8, 16))) + tsV1 = stsm_thr.partition_D(cute.flat_divide(gV1, (8, 16))) + tStage_V0 = cute.group_modes(tStage_V0, 1, cute.rank(tStage_V0)) + tStage_V1 = cute.group_modes(tStage_V1, 1, cute.rank(tStage_V1)) + tsV0 = cute.group_modes(tsV0, 1, cute.rank(tsV0)) + tsV1 = cute.group_modes(tsV1, 1, cute.rank(tsV1)) dequant_params = SimpleNamespace( pipeline_stage=pipeline_stage, warp_group_idx=warp_group_idx, tStage_K=thr_copy_K.partition_S(sStage2), tsK=thr_copy_K.partition_D(sK2), - tStage_V=thr_copy_V.partition_S(gStage_V), - tsV=thr_copy_V.partition_D(gV), + tStage_V0=tStage_V0, + tStage_V1=tStage_V1, + tsV0=tsV0, + tsV1=tsV1, + ldsm_copy=ldsm_copy, + stsm_copy=stsm_copy, ) - # FP8-KV bf16 Q in-place narrow partitions (256-thread cooperative copy of the - # full Q tile; src is the bf16 sQ, dst the fp16 compute view over the same smem). q_narrow_params = None - if const_expr(narrow_q): + if const_expr(quantize_q): VEC_Q = const_expr(8) tiled_copy_Q = copy_utils.tiled_copy_2d( - self.dtype, self.tile_hdim // VEC_Q, self.num_mma_threads, num_copy_elems=VEC_Q + self.kv_dtype, + self.tile_hdim // VEC_Q, + self.num_mma_threads, + num_copy_elems=VEC_Q, ) thr_copy_Q = tiled_copy_Q.get_slice(tidx) q_narrow_params = SimpleNamespace( tQsrc=thr_copy_Q.partition_S(sQ), tQdst=thr_copy_Q.partition_D(sQ_mma), + tQres=thr_copy_Q.partition_D(sQ_res), ) tile_scheduler = TileSchedulerCls() @@ -1446,6 +1573,8 @@ def mma( qk_descale, v_descale_tile = self._load_effective_descales( descale_tensors, batch_idx, head_idx_kv ) + qk_descale = qk_descale / Float32(self.fp8_q_scale) + v_descale_tile = v_descale_tile / Float32(self.fp8_p_scale) if const_expr(self.score_mod is None): softmax_scale_log2_eff = softmax_scale_log2 * qk_descale softmax_scale_eff = None @@ -1542,10 +1671,8 @@ def mma( n_block_max = n_block_max_full n_block_max_orig = n_block_max pipeline_q.consumer_wait_w_index_phase(0, q_consumer_phase) - # FP8-KV bf16 Q: cast the just-arrived bf16 Q tile -> fp16 in place (with a - # 256-thread barrier) before any QK WGMMA reads it. No-op when Q is already fp16. - if const_expr(narrow_q): - self._narrow_q_to_compute(q_narrow_params) + if const_expr(quantize_q): + self._quantize_q_to_fp8(q_narrow_params) # For performance reason, we separate out two kinds of iterations: # those that need masking on S, and those that don't. # We need masking on S for the very last block when K and V has length not multiple of tile_n. @@ -1804,40 +1931,198 @@ def _load_effective_descales(self, descale_tensors, batch_idx: Int32, kv_head_id return qk_descale, v_descale @cute.jit - def _narrow_q_to_compute(self, q_narrow_params): - """In-place narrow the bf16 Q tile -> fp16 compute dtype. src/dst alias the same - smem (both 2 bytes), so each thread reads bf16 into registers, converts, and writes - fp16 back. A 256-thread NarrowQ barrier publishes the fp16 Q to both MMA warpgroups - before the QK WGMMA. Called once per Q tile (Q is reused across all KV blocks).""" - # The in-place aliasing (sQ_mma reinterprets sQ's storage) is only valid when the - # Q-input and compute dtypes have the same width -- bf16<->fp16 are both 16-bit. A - # different width would make sQ_mma's elements overlap sQ's incorrectly. - assert self.q_dtype.width == self.dtype.width, ( - "in-place Q narrow requires q_dtype and compute dtype of equal width" - ) + def _gemm_two_fp8_a( + self, + tiled_mma, + acc, + tA_hi, + tA_res, + tB, + zero_init: cutlass.Constexpr[bool], + wg_wait: cutlass.Constexpr[int], + ): + """Issue high and residual FP8 A terms in one WGMMA commit group.""" + warpgroup.fence() + mma_atom = cute.make_mma_atom(tiled_mma.op) + mma_atom.set(warpgroup.Field.ACCUMULATE, not zero_init) + for k in cutlass.range_constexpr(cute.size(tA_hi.shape[2])): + cute.gemm(mma_atom, acc, tA_hi[None, None, k], tB[None, None, k], acc) + mma_atom.set(warpgroup.Field.ACCUMULATE, True) + for k in cutlass.range_constexpr(cute.size(tA_res.shape[2])): + cute.gemm(mma_atom, acc, tA_res[None, None, k], tB[None, None, k], acc) + warpgroup.commit_group() + if const_expr(wg_wait >= 0): + warpgroup.wait_group(wg_wait) + + @cute.jit + def _gemm_qk_fp8_compensated( + self, tiled_mma, shape, tQ_hi, tQ_res, tK, B_idx=None, wg_wait: int = -1 + ): + acc = cute.make_rmem_tensor(tiled_mma.partition_shape_C(shape), Float32) + rK = tK if const_expr(B_idx is None) else tK[None, None, None, B_idx] + self._gemm_two_fp8_a(tiled_mma, acc, tQ_hi, tQ_res, rK, True, wg_wait) + return acc + + @cute.jit + def _gemm_pv_fp8_compensated( + self, + tiled_mma, + acc, + tP_hi, + tP_res, + tV, + zero_init, + B_idx=None, + wg_wait: int = -1, + ): + rV = tV if const_expr(B_idx is None) else tV[None, None, None, B_idx] + self._gemm_two_fp8_a(tiled_mma, acc, tP_hi, tP_res, rV, zero_init, wg_wait) + + @cute.jit + def _quantize_q_to_fp8(self, q_narrow_params): + """Quantize the TMA-loaded FP16/BF16 Q tile into FP8 shared memory once.""" qp = q_narrow_params rS = cute.make_rmem_tensor_like(qp.tQsrc) rD = cute.make_rmem_tensor_like(qp.tQdst) + rR = cute.make_rmem_tensor_like(qp.tQres) cute.autovec_copy(qp.tQsrc, rS) - rD.store(rS.load().to(self.dtype)) + q_scaled = rS.load().to(Float32) * Float32(self.fp8_q_scale) + q_hi = q_scaled.to(self.kv_dtype) + rD.store(q_hi) + rR.store((q_scaled - q_hi.to(Float32)).to(self.kv_dtype)) cute.autovec_copy(rD, qp.tQdst) cute.arch.fence_view_async_shared() cute.arch.barrier( barrier_id=int(NamedBarrierFwd.NarrowQ), number_of_threads=self.num_mma_threads ) + # All input FP16/BF16 values have been captured in registers, so the Q input + # allocation can now safely hold the compact residual FP8 tile. + cute.autovec_copy(rR, qp.tQres) + cute.arch.fence_view_async_shared() + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.NarrowQ), number_of_threads=self.num_mma_threads + ) @cute.jit - def _cast_tile_fp8_to_f16(self, pipeline_stage, tStage, tsDst, phase): - """Cast one staged fp8 tile -> fp16 smem without applying a descale.""" + def _copy_staged_fp8_tile(self, pipeline_stage, tStage, tsDst, phase): + """Copy/reorder one staged FP8 tile into its FP8 WGMMA shared-memory layout.""" rS = cute.make_rmem_tensor_like(tStage) rD = cute.make_rmem_tensor_like(tsDst) pipeline_stage.consumer_wait_w_index_phase(0, phase) cute.autovec_copy(tStage, rS) - rD.store(utils.cvt_fp8_to_f16_packed(rS.load(), self.dtype)) + rD.store(rS.load()) cute.autovec_copy(rD, tsDst) pipeline_stage.consumer_release_w_index(0) cute.arch.fence_view_async_shared() + @cute.jit + def _transpose_fp8_v_tile(self, dp, tStage_V, tsV): + for i in cutlass.range_constexpr(cute.size(tStage_V.shape[1])): + dst = tsV[None, i] + # The STSM TV layout guarantees every lane's address is 16-byte + # aligned; retain that fact after partition/slicing in the DSL IR. + dst = cute.make_tensor(dst.iterator.align(16), dst.layout) + rV = cute.make_rmem_tensor_like(dst) + cute.copy(dp.ldsm_copy, tStage_V[None, i], rV) + rV_u32 = cute.recast_tensor(rV, Int32) + for j in cutlass.range_constexpr(cute.size(rV_u32) // 2): + upper = rV_u32[2 * j] + lower = rV_u32[2 * j + 1] + rV_u32[2 * j] = cute.arch.prmt(upper, lower, Int32(0x6420)) + rV_u32[2 * j + 1] = cute.arch.prmt(upper, lower, Int32(0x7531)) + cute.copy(dp.stsm_copy, rV, dst) + + @cute.jit + def _transpose_staged_fp8_v(self, dp, phase): + """FA3-style LDSM.T -> PRMT -> STSM transpose for one WG's V half.""" + dp.pipeline_stage.consumer_wait_w_index_phase(0, phase) + if dp.warp_group_idx == 0: + self._transpose_fp8_v_tile(dp, dp.tStage_V0, dp.tsV0) + else: + self._transpose_fp8_v_tile(dp, dp.tStage_V1, dp.tsV1) + dp.pipeline_stage.consumer_release_w_index(0) + cute.arch.fence_view_async_shared() + + @cute.jit + def _permute_Cregs_fp8(self, acc): + """Match FA3's QK-C to FP8 PV-A register permutation.""" + acc_u64 = cute.recast_tensor(acc, cutlass.Int64) + flat = cute.make_tensor(acc_u64.iterator, cute.make_layout(cute.size(acc_u64))) + for i in cutlass.range_constexpr(cute.size(flat) // 4): + tmp = flat[4 * i + 1] + flat[4 * i + 1] = flat[4 * i + 2] + flat[4 * i + 2] = tmp + + @cute.jit + def _permute_output_fp8(self, acc): + """Undo the output-column permutation from FA3's conflict-free V STSM.""" + # The hdim-512 Python kernel's PV accumulator carries the local FP8 + # output permutation as well. Undo that first, then apply FA3's + # cross-lane V-column permutation below. The order is significant. + frag32 = cute.group_modes(acc, 1, cute.rank(acc)) + for mi in cutlass.range_constexpr(cute.size(frag32.shape[1])): + for j in cutlass.range_constexpr(cute.size(frag32.shape[0][1])): + for i in cutlass.range_constexpr(frag32.shape[0][2] // 2): + tmp = frag32[((1, j, 2 * i), mi)] + frag32[((1, j, 2 * i), mi)] = frag32[((0, j, 2 * i + 1), mi)] + frag32[((0, j, 2 * i + 1), mi)] = tmp + + frag = cute.group_modes(cute.recast_tensor(acc, cutlass.Int64), 1, cute.rank(acc)) + n8 = frag.shape[0][2] + assert n8 % 2 == 0 + quad_idx = cute.arch.lane_idx() % 4 + lane_03 = quad_idx == 0 or quad_idx == 3 + upper_src = Int32(0) + if quad_idx == 1: + upper_src = Int32(2) + elif quad_idx == 2: + upper_src = Int32(3) + elif quad_idx == 3: + upper_src = Int32(1) + lower_src = upper_src ^ Int32(2) + for mi in cutlass.range_constexpr(cute.size(frag.shape[1])): + for j in cutlass.range_constexpr(cute.size(frag.shape[0][1])): + for i in cutlass.range_constexpr(n8 // 2): + upper = frag[((0, j, 2 * i), mi)] + lower = frag[((0, j, 2 * i + 1), mi)] + upper0, lower0 = upper, lower + if not lane_03: + upper0, lower0 = lower, upper + upper0 = utils.shuffle_sync(upper0, upper_src, width=4) + lower0 = utils.shuffle_sync(lower0, lower_src, width=4) + if lane_03: + frag[((0, j, 2 * i), mi)] = upper0 + frag[((0, j, 2 * i + 1), mi)] = lower0 + else: + frag[((0, j, 2 * i), mi)] = lower0 + frag[((0, j, 2 * i + 1), mi)] = upper0 + + @cute.jit + def _reshape_acc_to_fp8_frgA(self, acc): + """SM90 FP8 form of FA3 convert_layout_acc_Aregs.""" + layout = acc.layout + n8 = layout.shape[0][2] + assert n8 % 4 == 0 + frg_layout = cute.make_layout( + ((4, 2, 2), layout.shape[1], (n8 // 4, layout.shape[2])), + stride=( + (1, 4, 8), + layout.stride[1], + (16, layout.stride[2]), + ), + ) + return cute.make_tensor(acc.iterator, frg_layout) + + @cute.jit + def _prepare_fp8_probabilities(self, acc_S, tOrP, tOrP_res): + """Convert softmax probabilities into compensated FP8 PV A fragments.""" + self._permute_Cregs_fp8(acc_S) + tP_acc = self._reshape_acc_to_fp8_frgA(acc_S) + p_scaled = tP_acc.load() * Float32(self.fp8_p_scale) + p_hi = p_scaled.to(self.kv_dtype) + tOrP.store(p_hi) + tOrP_res.store((p_scaled - p_hi.to(Float32)).to(self.kv_dtype)) + @cute.jit def first_half_block_overlap_fp8( self, @@ -1856,11 +2141,10 @@ def first_half_block_overlap_fp8( is_first_block: bool = False, dequant_params: Optional[SimpleNamespace] = None, ): - """FP8-KV first-half block: cooperative 256-thread cast of the first K tile - + DequantK barrier, then the synchronous QK.""" + """FP8-KV first half: rearrange K, synchronize, then issue QK.""" dp = dequant_params nthr = const_expr(self.num_mma_threads) - self._cast_tile_fp8_to_f16(dp.pipeline_stage, dp.tStage_K, dp.tsK, Int32(0)) + self._copy_staged_fp8_tile(dp.pipeline_stage, dp.tStage_K, dp.tsK, Int32(0)) cute.arch.barrier(barrier_id=int(NamedBarrierFwd.DequantK), number_of_threads=nthr) acc_S = mma_qk_fn(B_idx=kv_consumer_state.index, wg_wait=0) @@ -1868,9 +2152,7 @@ def first_half_block_overlap_fp8( score_mod_fn(acc_S, n_block=n_block, seqlen=seqlen) mask_fn(acc_S, n_block=n_block, mask_seqlen=True) row_scale = softmax.online_softmax(acc_S, is_first=is_first_block) - tOrP_acc = layout_utils.reshape_acc_to_frgA(acc_S) - tOrP_cur = tOrP - tOrP_cur.store(tOrP_acc.load().to(self.dtype)) + self._prepare_fp8_probabilities(acc_S, tOrP, smem_copy_params.tOrP_res) if const_expr(self.rescale_O_before_gemm): acc_O.fill(0.0) scores_scale.store(row_scale.load()) @@ -1911,18 +2193,18 @@ def last_half_block_overlap_fp8( acc_O: Optional[cute.Tensor] = None, dequant_params: Optional[SimpleNamespace] = None, ): - """FP8-KV final PV: each WG casts its own hdimv-half of the last V tile - + WG-local DequantV barrier, then the synchronous PV.""" + """FP8-KV final PV: each WG rearranges its V half before the synchronous PV.""" if const_expr(self.rescale_O_before_gemm): softmax.rescale_O(acc_O, scores_scale) dp = dequant_params nthr_wg = const_expr(self.num_threads_per_warp_group) - self._cast_tile_fp8_to_f16(dp.pipeline_stage, dp.tStage_V, dp.tsV, Int32(1)) + self._transpose_staged_fp8_v(dp, Int32(1)) if dp.warp_group_idx == 0: cute.arch.barrier(barrier_id=int(NamedBarrierFwd.DequantV0), number_of_threads=nthr_wg) else: cute.arch.barrier(barrier_id=int(NamedBarrierFwd.DequantV1), number_of_threads=nthr_wg) mma_pv_fn(B_idx=kv_consumer_state.index, zero_init=zero_init, wg_wait=0) + self._permute_output_fp8(acc_O) kv_consumer_state.advance() return kv_consumer_state @@ -2090,7 +2372,7 @@ def mma_one_n_block_intrawg_overlap_fp8( nthr = const_expr(self.num_mma_threads) nthr_wg = const_expr(self.num_threads_per_warp_group) # ---- 1. cast K[n] -> sK (cooperative 256-thread, staging phase 1) ---- - self._cast_tile_fp8_to_f16(dp.pipeline_stage, dp.tStage_K, dp.tsK, Int32(1)) + self._copy_staged_fp8_tile(dp.pipeline_stage, dp.tStage_K, dp.tsK, Int32(1)) cute.arch.barrier(barrier_id=int(NamedBarrierFwd.DequantK), number_of_threads=nthr) # ---- 2. S = Q @ K.T (async) ---- self.warp_scheduler_barrier_sync() @@ -2098,7 +2380,7 @@ def mma_one_n_block_intrawg_overlap_fp8( if const_expr(self.rescale_O_before_gemm): softmax.rescale_O(acc_O, scores_scale) # ---- 3. cast V[n_prev] -> sV (WG-split, phase 0) WHILE QK[n] runs ---- - self._cast_tile_fp8_to_f16(dp.pipeline_stage, dp.tStage_V, dp.tsV, Int32(0)) + self._transpose_staged_fp8_v(dp, Int32(0)) if dp.warp_group_idx == 0: cute.arch.barrier(barrier_id=int(NamedBarrierFwd.DequantV0), number_of_threads=nthr_wg) else: @@ -2113,9 +2395,7 @@ def mma_one_n_block_intrawg_overlap_fp8( mask_fn(acc_S=acc_S, n_block=n_block) row_scale = softmax.online_softmax(acc_S, check_inf=check_inf) warpgroup.wait_group(0) - tOrP_acc = layout_utils.reshape_acc_to_frgA(acc_S) - tOrP_cur = tOrP - utils.cvt_f16(tOrP_acc, tOrP_cur) + self._prepare_fp8_probabilities(acc_S, tOrP, smem_copy_params.tOrP_res) if const_expr(not self.rescale_O_before_gemm): softmax.rescale_O(acc_O, row_scale) if const_expr(self.rescale_O_before_gemm): diff --git a/flash_attn/cute/interface.py b/flash_attn/cute/interface.py index 12dca296b1d..4fccaf0a25d 100644 --- a/flash_attn/cute/interface.py +++ b/flash_attn/cute/interface.py @@ -777,10 +777,12 @@ def _flash_attn_fwd( "inputs must be float16, bfloat16, fp8 e4m3fn, or fp8 e5m2" ) if fp8_kv_dequant: - # FP8-KV in-kernel dequant: fp16/bf16 Q + fp8 e4m3 K/V (dequantized in-kernel to - # fp16), fp16/bf16 O. Compute is always fp16 (the single-instruction fp8->fp16 fast - # path); a bf16 Q is narrowed to fp16 in-kernel and the fp32 accumulator is cast to - # O's dtype in the epilogue, so Q/O dtypes are independent of compute. SM90 only. + # SM90 FP8 MMA: fp16/bf16 Q is represented as compensated E4M3 pairs; + # paged E4M3 K/V remain FP8 WGMMA operands, and accumulation is FP32. + # The epilogue casts FP32 accumulators to the native fp16/bf16 output dtype. + assert q is not None and qv is None, ( + "fp8_kv_dequant requires Q and does not support QV-only input" + ) assert q.dtype in (torch.float16, torch.bfloat16), ( "fp8_kv_dequant requires fp16 or bf16 Q" ) @@ -788,6 +790,18 @@ def _flash_attn_fwd( "fp8_kv_dequant requires fp8 e4m3 K/V" ) assert page_table is not None, "fp8_kv_dequant requires paged KV (page_table)" + if dynamic_causal is not None: + raise NotImplementedError( + "FA4 SM90 FP8 MMA does not support dynamic_causal" + ) + if block_sparse_tensors is not None: + raise NotImplementedError( + "FA4 SM90 FP8 MMA does not support block sparsity" + ) + assert head_dim == head_dim_v == 512, ( + "FA4 SM90 FP8 MMA currently requires head_dim == head_dim_v == 512; " + f"got head_dim={head_dim}, head_dim_v={head_dim_v}." + ) else: input_tensors = {"q": q, "k": k, "v": v, "qv": qv} present = {name: t for name, t in input_tensors.items() if t is not None} @@ -852,7 +866,7 @@ def _flash_attn_fwd( pack_gqa = qhead_per_kvhead > 1 is_fp8 = v.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) and not fp8_kv_dequant - if is_fp8 and requires_grad: + if (is_fp8 or fp8_kv_dequant) and requires_grad: raise NotImplementedError("FA4 CuTe FP8 backward is not supported yet (forward-only).") if output_scale is not None: assert output_scale.dtype == torch.float32, "output_scale must be float32" @@ -931,23 +945,44 @@ def _flash_attn_fwd( "q_descale/k_descale/v_descale are only supported for FP8 inputs" ) - if fp8_kv_dequant and q_descale is None: - # fp16 Q has no q-scale; materialize an identity (1.0) q_descale so the - # DescaleTensors struct has all three fields present (identity q is a no-op). - _ref_descale = k_descale if k_descale is not None else v_descale - q_descale = ( - torch.ones_like(_ref_descale) - if _ref_descale is not None - else torch.ones((batch_size, num_head_kv), dtype=torch.float32, device=q.device) + if fp8_kv_dequant and any( + t is None for t in (q_descale, k_descale, v_descale) + ): + # DescaleTensors' MLIR reconstruction packs present values positionally. + # Materialize every missing slot so configurations such as V-only scaling + # cannot shift V into the K field and silently corrupt QK/PV scaling. + _ref_descale = next( + (t for t in (q_descale, k_descale, v_descale) if t is not None), + None, ) + if _ref_descale is None: + identity_descale = torch.ones( + (batch_size, num_head_kv), dtype=torch.float32, device=q.device + ) + else: + identity_descale = torch.ones_like( + _ref_descale, memory_format=torch.contiguous_format + ) + q_descale, k_descale, v_descale = [ + t if t is not None else identity_descale + for t in (q_descale, k_descale, v_descale) + ] + # Expanded scalar scales can have a zero stride even when every size is + # one; Tensor.contiguous() may preserve that layout. Copy only those + # tensors whose CuTe dynamic leading dimension is not actually unit-stride. + q_descale, k_descale, v_descale = [ + t.clone(memory_format=torch.contiguous_format) + if t is not None and t.stride(-1) != 1 + else t + for t in (q_descale, k_descale, v_descale) + ] dtype = torch2cute_dtype_map[q_dtype] kv_dtype = torch2cute_dtype_map[k.dtype] if fp8_kv_dequant else dtype if fp8_kv_dequant: assert arch // 10 == 9, "fp8_kv_dequant is an SM90-only forward (compute capability 9.x)" - # Compute is fp16 regardless of the (fp16/bf16) Q dtype: the fp8->fp16 cvt is the - # only single-instruction widening on SM90. The kernel derives the Q/O tensor - # dtypes from mQ/mO (which keep q.dtype / out_torch_dtype), narrowing bf16 Q to - # fp16 in-kernel and casting the fp32 accumulator to O's dtype in the epilogue. + # This dtype selects fp16 auxiliary layouts/conversions. QK and PV use + # compensated E4M3 operands with FP8 WGMMA and FP32 accumulation; mQ/mO + # retain the caller's fp16/bf16 input and output dtypes. dtype = torch2cute_dtype_map[torch.float16] if is_fp8: assert arch // 10 == 10, "FP8 is only supported on SM100 (compute capability 10.x) for FA4 CuTe." @@ -1000,6 +1035,10 @@ def _flash_attn_fwd( intra_wg_overlap = fwd_cfg.intra_wg_overlap if fp8_kv_dequant: + assert intra_wg_overlap, ( + "FA4 SM90 FP8 MMA requires intra_wg_overlap=True until a compatible " + "non-overlap control flow is implemented." + ) # Force RS-mode PV: frees sP for the fp8 staging buffer, keeping smem within # budget at d=512 while preserving intra-WG overlap. mma_pv_is_rs = True @@ -1483,10 +1522,10 @@ def _flash_attn_fwd( fa_logging.get_fa_log_level(), output_quant_key, fp8_kv_dequant, - # fp8_kv_dequant forces compute dtype = fp16, so the Q/O tensor dtypes (which the - # kernel derives from mQ/mO and which select the in-kernel narrow/widen) are no - # longer captured by `dtype` above -- key on them explicitly. Redundant elsewhere. - q.dtype, + # FP8 MMA uses a fixed fp16 auxiliary dtype, so `dtype` no longer captures + # native Q/O tensor dtypes. Key on them explicitly for quantization and + # epilogue selection. This is redundant for the other paths. + q_dtype, out_torch_dtype, # the decode kernel is a different class with a different grid and smem plan decode_splits, @@ -1532,7 +1571,7 @@ def _flash_attn_fwd( lse_tensor = to_cute_tensor(lse, assumed_align=4) q_descale_tensor, k_descale_tensor, v_descale_tensor = ( - to_cute_tensor(t, assumed_align=4, leading_dim=1) + to_cute_tensor(t, assumed_align=4, leading_dim=1) if t is not None else None for t in (q_descale, k_descale, v_descale) ) descale_tensors_tensor = ( @@ -1879,8 +1918,8 @@ def _flash_attn_fwd( for t in (q_call, k_call, v_call, qv_call) ] elif fp8_kv_dequant: - # FP8-KV dequant: Q is fp16/bf16 but K/V are fp8 e4m3 -- apply the same uint8 - # FFI workaround to K/V only (the compiled kernel takes them as uint8). + # FP8 MMA has native fp16/bf16 Q and E4M3 K/V. Apply the uint8 FFI + # workaround to K/V only (the compiled kernel takes them as uint8). k_call = k_call.view(torch.uint8) v_call = v_call.view(torch.uint8) out_call = out.detach() diff --git a/tests/cute/test_flash_attn.py b/tests/cute/test_flash_attn.py index b022daea0a1..9a0814c6f20 100644 --- a/tests/cute/test_flash_attn.py +++ b/tests/cute/test_flash_attn.py @@ -2801,26 +2801,28 @@ def test_flash_attn_sm100_hdim256_seqused_paged_kv_fwd( @pytest.mark.parametrize("mha_type", ["mha", "gqa"]) @pytest.mark.parametrize("k_scale,v_scale", [(1.0, 1.0), (0.5, 0.25)]) @pytest.mark.parametrize("num_splits", [1, 2]) +@pytest.mark.parametrize("q_dtype", [torch.float16, torch.bfloat16]) @maybe_fake_tensor_mode(USE_FAKE_TENSOR) -def test_flash_attn_kvcache_fp8_dequant_sm90(num_splits, k_scale, v_scale, mha_type): - """SM90 fp16-Q + fp8-KV-cache dequant paged forward (d=512). +def test_flash_attn_kvcache_fp8_dequant_sm90(num_splits, k_scale, v_scale, mha_type, q_dtype): + """SM90 fp16/bf16-Q + fp8-KV-cache Tensor Core paged forward (d=512). - fp16 Q with per-tensor FP8 (e4m3) paged K/V cast in-kernel to fp16 (the - compute/output dtype). Per-(batch, kv_head) descales are folded by the kernel + fp16/bf16 Q and per-tensor FP8 (e4m3) paged K/V execute QK and PV with FP8 WGMMA. + Compensated FP8 pairs preserve Q and softmax-P accuracy without widening the + K/V operands to fp16. Per-(batch, kv_head) descales are folded by the kernel (q*k into the score scale, v into final output normalization/final_scale). The reference reads the exact paged FP8 bytes the kernel indexes (gathered through page_table, cast to fp16) and applies the same descales via attention_ref (dequantize-then-attend). num_splits=2 exercises the SplitKV path (fp32 - partials combined back to fp16). Q is true fp16 -> q_descale is identity - (None on both sides). page_size must equal tile_n (64 for d=512 on SM90). + partials combined back to the Q dtype). Q is not pre-quantized, so q_descale is + identity (None on both sides). page_size must equal tile_n (64 for d=512 on SM90). """ device, d, page_size, causal = "cuda", 512, 64, True batch_size, nheads, seqlen_q, seqlen_k = 2, 4, 5, 256 nheads_k = nheads if mha_type == "mha" else 2 torch.random.manual_seed(0) - # fp16 Q (true precision -> no q_descale); bf16 paged K/V from the shared builder. - q = torch.randn(batch_size, seqlen_q, nheads, d, device=device, dtype=torch.float16) + # True-precision Q (no q_descale); bf16 paged K/V from the shared builder. + q = torch.randn(batch_size, seqlen_q, nheads, d, device=device, dtype=q_dtype) _, _, page_table, k_cache_paged, v_cache_paged, _ = _generate_block_kvcache( seqlen_k, page_size, batch_size, nheads_k, d, d, device, torch.bfloat16, torch.bfloat16 ) @@ -2840,21 +2842,22 @@ def test_flash_attn_kvcache_fp8_dequant_sm90(num_splits, k_scale, v_scale, mha_t arange = rearrange(torch.arange(seqlen_k, device=device), "s -> 1 s") key_padding_mask = arange < rearrange(cache_seqlens, "b -> b 1") - # Reference reads the EXACT fp8 bytes the kernel reads: gather the paged fp8 through the - # page table into dense (b, seqlen_k, nheads_k, d), then to fp16 (the kernel casts - # fp8 -> fp16). attention_ref applies k/v_descale = dequant; out is fp16 (== kernel O dtype). + # Reference reads the exact FP8 bytes the kernel reads: gather the paged FP8 through + # the page table, then widen only the reference copy. attention_ref applies the K/V + # descales while the kernel folds the same scales around its FP8 MMA operations. gather = lambda paged: rearrange( paged[page_table.flatten()], "(b n) p ... -> b (n p) ...", b=batch_size )[:, :seqlen_k].to(torch.float16) k_ref, v_ref = gather(k_paged_fp8), gather(v_paged_fp8) common = dict(causal=causal, k_descale=k_descale, v_descale=v_descale) out_ref, _ = attention_ref(q, k_ref, v_ref, None, key_padding_mask, **common) - # out_pt models fp16 arithmetic (NO fp8 intermediate_dtype) -> sets the error bar. + # out_pt models input-type arithmetic (without FP8 intermediates) and sets the error bar. + k_pt, v_pt = k_ref.to(q_dtype), v_ref.to(q_dtype) out_pt, _ = attention_ref( - q, k_ref, v_ref, None, key_padding_mask, upcast=False, reorder_ops=True, **common + q, k_pt, v_pt, None, key_padding_mask, upcast=False, reorder_ops=True, **common ) - # ---- kernel under test: fp16 Q + fp8 paged K/V, cast in-kernel ---- + # ---- kernel under test: fp16/bf16 Q + fp8 paged K/V, FP8 Tensor Core MMA ---- out, *_ = _flash_attn_fwd( q=q, k=k_paged_fp8, @@ -2875,12 +2878,408 @@ def test_flash_attn_kvcache_fp8_dequant_sm90(num_splits, k_scale, v_scale, mha_t print(f"Output max diff: {(out - out_ref).abs().max().item()}") print(f"Output mean diff: {(out - out_ref).abs().mean().item()}") + print( + "Output cosine: " + f"{torch.nn.functional.cosine_similarity(out.float().flatten(), out_ref.float().flatten(), dim=0).item()}" + ) print(f"Pytorch max diff: {(out_pt - out_ref).abs().max().item()}") print(f"Pytorch mean diff: {(out_pt - out_ref).abs().mean().item()}") - # Same max/mean multiplier style as test_flash_attn_kvcache (fp8 bar). - assert (out - out_ref).abs().max().item() <= 4 * (out_pt - out_ref).abs().max().item() + 1e-5 - assert (out - out_ref).abs().mean().item() <= 3 * (out_pt - out_ref).abs().mean().item() + # FP8 products have a wider numerical envelope than fp16 MMA even with the + # compensated representation. Bound both worst-case and aggregate error. + assert (out - out_ref).abs().max().item() <= 8 * (out_pt - out_ref).abs().max().item() + 1e-5 + assert (out - out_ref).abs().mean().item() <= 6 * (out_pt - out_ref).abs().mean().item() + + +@pytest.mark.skipif(not IS_SM90, reason="fp8-KV Tensor Core forward is SM90-only") +@maybe_fake_tensor_mode(USE_FAKE_TENSOR) +def test_flash_attn_kvcache_fp8_mma_v_layout_sm90(): + """FP8 PV must preserve token-to-output mapping after the K-major V transpose.""" + d, page_size = 512, 64 + q = torch.zeros(1, 1, 1, d, device="cuda", dtype=torch.float16) + k = torch.zeros(1, page_size, 1, d, device="cuda", dtype=torch.float8_e4m3fn) + v_fp16 = torch.zeros(1, page_size, 1, d, device="cuda", dtype=torch.float16) + token_idx = torch.arange(page_size, device="cuda") + v_fp16[0, token_idx, 0, token_idx] = 1 + v = v_fp16.to(torch.float8_e4m3fn) + page_table = torch.zeros(1, 1, device="cuda", dtype=torch.int32) + seqused_k = torch.full((1,), page_size, device="cuda", dtype=torch.int32) + descale = torch.ones(1, 1, device="cuda", dtype=torch.float32) + + out, *_ = _flash_attn_fwd( + q=q, + k=k, + v=v, + causal=False, + page_table=page_table, + seqused_k=seqused_k, + max_seqlen_q=1, + max_seqlen_k=page_size, + k_descale=descale, + v_descale=descale, + fp8_kv_dequant=True, + ) + if is_fake_mode(): + return + + expected = torch.zeros_like(out) + expected[..., :page_size] = 1 / page_size + torch.testing.assert_close(out, expected, atol=0, rtol=0) + + +def _make_sm90_fp8_mma_inputs(head_dim=512, seqlen_k=64, requires_grad=False): + """Build the smallest paged FP8-KV input accepted by the SM90 interface.""" + page_size = 64 + num_pages = math.ceil(seqlen_k / page_size) + q = torch.zeros( + 1, + 1, + 1, + head_dim, + device="cuda", + dtype=torch.float16, + requires_grad=requires_grad, + ) + k = torch.zeros( + num_pages, + page_size, + 1, + head_dim, + device="cuda", + dtype=torch.float8_e4m3fn, + ) + v = torch.zeros_like(k) + page_table = torch.arange(num_pages, device="cuda", dtype=torch.int32).unsqueeze(0) + seqused_k = torch.tensor([seqlen_k], device="cuda", dtype=torch.int32) + descale = torch.ones((), device="cuda", dtype=torch.float32).expand(1, 1) + return q, k, v, page_table, seqused_k, descale + + +@pytest.mark.skipif(not IS_SM90, reason="FP8-KV Tensor Core forward is SM90-only") +@pytest.mark.parametrize("head_dim", [128, 256, 384]) +def test_flash_attn_kvcache_fp8_mma_rejects_unsupported_head_dim_sm90(head_dim): + q, k, v, page_table, seqused_k, descale = _make_sm90_fp8_mma_inputs(head_dim) + with pytest.raises(AssertionError, match="head_dim == head_dim_v == 512"): + _flash_attn_fwd( + q=q, + k=k, + v=v, + page_table=page_table, + seqused_k=seqused_k, + max_seqlen_q=1, + max_seqlen_k=64, + k_descale=descale, + v_descale=descale, + fp8_kv_dequant=True, + ) + + +@pytest.mark.skipif(not IS_SM90, reason="FP8-KV Tensor Core forward is SM90-only") +def test_flash_attn_kvcache_fp8_mma_rejects_nonoverlap_sm90(): + q, k, v, page_table, seqused_k, descale = _make_sm90_fp8_mma_inputs() + with pytest.raises(AssertionError, match="intra_wg_overlap=True"): + _flash_attn_fwd( + q=q, + k=k, + v=v, + page_table=page_table, + seqused_k=seqused_k, + max_seqlen_q=1, + max_seqlen_k=64, + k_descale=descale, + v_descale=descale, + fp8_kv_dequant=True, + intra_wg_overlap=False, + ) + + +@pytest.mark.skipif(not IS_SM90, reason="FP8-KV Tensor Core forward is SM90-only") +def test_flash_attn_kvcache_fp8_mma_rejects_requires_grad_sm90(): + q, k, v, page_table, seqused_k, descale = _make_sm90_fp8_mma_inputs( + requires_grad=True + ) + with pytest.raises(NotImplementedError, match="forward-only"): + _flash_attn_fwd( + q=q, + k=k, + v=v, + page_table=page_table, + seqused_k=seqused_k, + max_seqlen_q=1, + max_seqlen_k=64, + k_descale=descale, + v_descale=descale, + fp8_kv_dequant=True, + ) + + +@pytest.mark.skipif(not IS_SM90, reason="FP8-KV Tensor Core forward is SM90-only") +def test_flash_attn_kvcache_fp8_mma_rejects_qv_only_sm90(): + qv, k, v, page_table, seqused_k, descale = _make_sm90_fp8_mma_inputs() + with pytest.raises(AssertionError, match="requires Q and does not support QV-only"): + _flash_attn_fwd( + q=None, + qv=qv, + k=k, + v=v, + page_table=page_table, + seqused_k=seqused_k, + max_seqlen_q=1, + max_seqlen_k=64, + k_descale=descale, + v_descale=descale, + fp8_kv_dequant=True, + ) + + +@pytest.mark.skipif(not IS_SM90, reason="FP8-KV Tensor Core forward is SM90-only") +@pytest.mark.parametrize( + "unsupported_kwarg,error", + [ + ("dynamic_causal", "does not support dynamic_causal"), + ("block_sparse_tensors", "does not support block sparsity"), + ], +) +def test_flash_attn_kvcache_fp8_mma_rejects_unsupported_scheduling_sm90( + unsupported_kwarg, error +): + q, k, v, page_table, seqused_k, descale = _make_sm90_fp8_mma_inputs() + unsupported_value = ( + torch.zeros(1, device="cuda", dtype=torch.bool) + if unsupported_kwarg == "dynamic_causal" + else object() + ) + with pytest.raises(NotImplementedError, match=error): + _flash_attn_fwd( + q=q, + k=k, + v=v, + page_table=page_table, + seqused_k=seqused_k, + max_seqlen_q=1, + max_seqlen_k=64, + k_descale=descale, + v_descale=descale, + fp8_kv_dequant=True, + **{unsupported_kwarg: unsupported_value}, + ) + + +@pytest.mark.skipif(not IS_SM90, reason="FP8-KV Tensor Core forward is SM90-only") +@pytest.mark.parametrize("seqlen_k", [1, 63, 65, 193]) +@maybe_fake_tensor_mode(USE_FAKE_TENSOR) +def test_flash_attn_kvcache_fp8_mma_partial_kv_blocks_sm90(seqlen_k): + q, k, v, page_table, seqused_k, descale = _make_sm90_fp8_mma_inputs( + seqlen_k=seqlen_k + ) + # Make padded tokens conspicuously different so an incorrect tail mask fails. + values = ((torch.arange(seqlen_k, device="cuda") % 5) - 2).to(torch.float16) + v.view(-1, 1, 512)[:seqlen_k].copy_( + values[:, None, None].expand(seqlen_k, 1, 512).to(torch.float8_e4m3fn) + ) + v.view(-1, 1, 512)[seqlen_k:].fill_(16) + + out, *_ = _flash_attn_fwd( + q=q, + k=k, + v=v, + causal=False, + page_table=page_table, + seqused_k=seqused_k, + max_seqlen_q=1, + max_seqlen_k=seqlen_k, + q_descale=None, + k_descale=descale, + v_descale=descale, + fp8_kv_dequant=True, + ) + if is_fake_mode(): + return + + assert descale.stride(-1) == 0 # Exercise broadcast scale canonicalization. + expected = values.float().mean().expand_as(out.float()) + torch.testing.assert_close(out.float(), expected, atol=2e-2, rtol=2e-2) + + +@pytest.mark.skipif(not IS_SM90, reason="FP8-KV Tensor Core forward is SM90-only") +@maybe_fake_tensor_mode(USE_FAKE_TENSOR) +def test_flash_attn_kvcache_fp8_mma_v_only_descale_sm90(): + """A missing middle K scale must not shift V scale into the QK scale slot.""" + q, k, v, page_table, seqused_k, _ = _make_sm90_fp8_mma_inputs() + v.fill_(1) + v_descale = torch.full((1, 1), 0.25, device="cuda", dtype=torch.float32) + + out, *_ = _flash_attn_fwd( + q=q, + k=k, + v=v, + causal=False, + page_table=page_table, + seqused_k=seqused_k, + max_seqlen_q=1, + max_seqlen_k=64, + q_descale=None, + k_descale=None, + v_descale=v_descale, + fp8_kv_dequant=True, + ) + if is_fake_mode(): + return + + expected = torch.full_like(out.float(), 0.25) + torch.testing.assert_close(out.float(), expected, atol=2e-2, rtol=2e-2) + + +@pytest.mark.skipif(not IS_SM90, reason="FP8-KV Tensor Core forward is SM90-only") +@maybe_fake_tensor_mode(USE_FAKE_TENSOR) +def test_flash_attn_kvcache_fp8_mma_long_context_splitkv_sm90(): + """Validate compensated FP8 QK/PV accuracy at a 4K target context.""" + torch.manual_seed(0) + batch_size, seqlen_q, seqlen_k = 1, 1, 4096 + nheads, nheads_kv, d, page_size = 2, 1, 512, 64 + q = torch.randn( + batch_size, seqlen_q, nheads, d, device="cuda", dtype=torch.float16 + ) + k_source = torch.randn( + batch_size * seqlen_k // page_size, + page_size, + nheads_kv, + d, + device="cuda", + dtype=torch.bfloat16, + ) + v_source = torch.randn_like(k_source) + k_scale, v_scale = 0.5, 0.25 + k = (k_source / k_scale).to(torch.float8_e4m3fn) + v = (v_source / v_scale).to(torch.float8_e4m3fn) + page_table = torch.arange( + seqlen_k // page_size, device="cuda", dtype=torch.int32 + ).unsqueeze(0) + seqused_k = torch.tensor([seqlen_k], device="cuda", dtype=torch.int32) + k_descale = torch.full((1, 1), k_scale, device="cuda", dtype=torch.float32) + v_descale = torch.full((1, 1), v_scale, device="cuda", dtype=torch.float32) + + k_ref = k.view(batch_size, seqlen_k, nheads_kv, d).to(torch.float16) + v_ref = v.view(batch_size, seqlen_k, nheads_kv, d).to(torch.float16) + out_ref, _ = attention_ref( + q, + k_ref, + v_ref, + None, + None, + causal=False, + k_descale=k_descale, + v_descale=v_descale, + ) + out_pt, _ = attention_ref( + q, + k_ref, + v_ref, + None, + None, + causal=False, + upcast=False, + reorder_ops=True, + k_descale=k_descale, + v_descale=v_descale, + ) + + out, *_ = _flash_attn_fwd( + q=q, + k=k, + v=v, + causal=False, + page_table=page_table, + seqused_k=seqused_k, + max_seqlen_q=seqlen_q, + max_seqlen_k=seqlen_k, + num_splits=4, + k_descale=k_descale, + v_descale=v_descale, + fp8_kv_dequant=True, + ) + if is_fake_mode(): + return + + assert (out - out_ref).abs().max().item() <= ( + 8 * (out_pt - out_ref).abs().max().item() + 1e-5 + ) + assert (out - out_ref).abs().mean().item() <= ( + 8 * (out_pt - out_ref).abs().mean().item() + 1e-5 + ) + + +@pytest.mark.skipif(not IS_SM90, reason="FP8-KV Tensor Core forward is SM90-only") +@pytest.mark.parametrize("q_magnitude", [448.0, 512.0]) +@maybe_fake_tensor_mode(USE_FAKE_TENSOR) +def test_flash_attn_kvcache_fp8_mma_large_q_sm90(q_magnitude): + """Stress the fixed Q scale at and beyond one E4M3 component's range. + + Q is represented as two compensated E4M3 components. The 512 case therefore + also verifies that the residual component recovers values above 448 instead + of silently relying on a single E4M3 conversion. K is inversely scaled so + the attention-score envelope remains well conditioned. + """ + torch.manual_seed(0) + batch_size, seqlen_q, seqlen_k = 1, 1, 256 + nheads, nheads_kv, d, page_size = 2, 1, 512, 64 + q = torch.empty( + batch_size, seqlen_q, nheads, d, device="cuda", dtype=torch.float16 + ).uniform_(-q_magnitude, q_magnitude) + k_source = torch.empty( + seqlen_k // page_size, + page_size, + nheads_kv, + d, + device="cuda", + dtype=torch.float16, + ).uniform_(-1, 1) + v_source = torch.empty_like(k_source).uniform_(-1, 1) + k_scale, v_scale = 1.0 / q_magnitude, 0.25 + k = k_source.to(torch.float8_e4m3fn) + v = (v_source / v_scale).to(torch.float8_e4m3fn) + page_table = torch.arange( + seqlen_k // page_size, device="cuda", dtype=torch.int32 + ).unsqueeze(0) + seqused_k = torch.tensor([seqlen_k], device="cuda", dtype=torch.int32) + k_descale = torch.full((1, 1), k_scale, device="cuda", dtype=torch.float32) + v_descale = torch.full((1, 1), v_scale, device="cuda", dtype=torch.float32) + + k_ref = k.view(batch_size, seqlen_k, nheads_kv, d).to(torch.float16) + v_ref = v.view(batch_size, seqlen_k, nheads_kv, d).to(torch.float16) + out_ref, _ = attention_ref( + q, + k_ref, + v_ref, + None, + None, + causal=False, + k_descale=k_descale, + v_descale=v_descale, + ) + out, *_ = _flash_attn_fwd( + q=q, + k=k, + v=v, + causal=False, + page_table=page_table, + seqused_k=seqused_k, + max_seqlen_q=seqlen_q, + max_seqlen_k=seqlen_k, + k_descale=k_descale, + v_descale=v_descale, + fp8_kv_dequant=True, + ) + if is_fake_mode(): + return + + if q_magnitude > 448: + assert q.abs().max().item() > 448 + assert torch.isfinite(out).all() + assert (out.float() - out_ref.float()).abs().max().item() < 1e-3 + assert torch.nn.functional.cosine_similarity( + out.float().flatten(), out_ref.float().flatten(), dim=0 + ).item() > 0.9999 @pytest.mark.skipif(not IS_SM90, reason="SM90 hdim 512 forward") From 991185c33d6ee361c8b46a40a7b16db0c91eea00 Mon Sep 17 00:00:00 2001 From: Kate Cheng Date: Wed, 30 Sep 2026 00:20:36 -0700 Subject: [PATCH 02/13] Harden SM90 FP8 MMA specialization contracts Signed-off-by: Kate Cheng --- flash_attn/cute/flash_fwd_sm90.py | 68 ++++++++-- flash_attn/cute/interface.py | 45 ++++--- tests/cute/test_flash_attn.py | 202 ++++++++++++++++++++++++++++-- 3 files changed, 267 insertions(+), 48 deletions(-) diff --git a/flash_attn/cute/flash_fwd_sm90.py b/flash_attn/cute/flash_fwd_sm90.py index 798e78921f6..d8982527daf 100644 --- a/flash_attn/cute/flash_fwd_sm90.py +++ b/flash_attn/cute/flash_fwd_sm90.py @@ -75,6 +75,10 @@ def __init__( # SM90 compensated FP8-MMA path: native fp16/bf16 Q, E4M3 paged K/V, # FP8 WGMMA QK/PV, FP32 accumulation, and native fp16/bf16 O. self.fp8_kv_dequant = fp8_kv_dequant + assert not self.fp8_kv_dequant or self.mma_pv_is_rs, ( + "SM90 FP8 KV dequant requires RS-mode PV because its shared storage " + "does not allocate sP" + ) self.kv_dtype = kv_dtype if kv_dtype is not None else self.dtype # Hopper FP8 WGMMA needs FP8 on both inputs. Q is quantized once per output # tile. P is scaled before its FP8 cast to preserve long-context softmax @@ -720,11 +724,8 @@ def kernel( ) sVt = layout_utils.transpose_view(sV) sP = None - sP_res = None if const_expr(sP_layout is not None): sP = storage.sP.get_tensor(sP_layout.outer, swizzle=sP_layout.inner) - if const_expr(self.fp8_kv_dequant): - sP_res = storage.sP_res.get_tensor(sP_layout.outer, swizzle=sP_layout.inner) # FP8 staging buffer: TMA destination and WGMMA-layout rearrangement source. sStage = None if const_expr(self.fp8_kv_dequant): @@ -819,7 +820,6 @@ def kernel( sK, sVt, sP, - sP_res, sO, sQ_mma, sQ_res, @@ -915,10 +915,42 @@ def load_fp8( n_block_min, n_block_max = block_info.get_n_block_min_max( seqlen, m_block, split_idx, num_splits ) + # A bidirectional entry in a dynamic-causal batch owns a slice of + # the full KV range, not the causal range. This must mirror the + # consumer's range calculation exactly: each emitted K/V stage has + # a matching consumer wait. + if const_expr(self._mDynamicCausal is not None): + psc = self._mDynamicCausal[batch_idx] + if not psc: + n_block_max_full = cute.ceil_div(seqlen.seqlen_k, self.tile_n) + if const_expr(self.is_split_kv): + num_n_blocks_per_split = cute.ceil_div( + n_block_max_full, num_splits + ) + n_block_min = split_idx * num_n_blocks_per_split + n_block_max = cutlass.min( + n_block_min + num_n_blocks_per_split, n_block_max_full + ) + else: + n_block_min = Int32(0) + n_block_max = n_block_max_full + + # An empty SplitKV partition still needs the same one K and one V + # pipeline transaction as the consumer's first/last-half sequence. + # Use physical page 0 as a known-valid dummy source; its result is + # discarded by the matching consumer-side empty-partition path. + is_empty_split = n_block_min >= n_block_max + if is_empty_split: + n_block_min = Int32(0) + n_block_max = Int32(1) # ---- Prologue emit: stage K[n_block_max - 1] (warp 0 only) ---- + page_idx = Int32(0) if warp_idx_in_wg == 0: - page_idx = mPageTable[batch_idx, n_block_max - 1] + if is_empty_split: + page_idx = Int32(0) + else: + page_idx = mPageTable[batch_idx, n_block_max - 1] pipeline_stage.producer_acquire_w_index_phase(0, stage_producer_phase) stage_copy_K(src_idx=page_idx, dst_idx=0, tma_bar_ptr=stage_full_bar) stage_producer_phase ^= 1 @@ -955,7 +987,11 @@ def load_fp8( stage_producer_phase ^= 1 # ---- Epilogue emit: stage V[n_block_min] ---- - page_idx = mPageTable[batch_idx, n_block_min] + page_idx = Int32(0) + if is_empty_split: + page_idx = Int32(0) + else: + page_idx = mPageTable[batch_idx, n_block_min] pipeline_stage.producer_acquire_w_index_phase(0, stage_producer_phase) stage_copy_V(src_idx=page_idx, dst_idx=0, tma_bar_ptr=stage_full_bar) stage_producer_phase ^= 1 @@ -1314,7 +1350,6 @@ def mma( sK: cute.Tensor, sVt: cute.Tensor, sP: Optional[cute.Tensor], - sP_res: Optional[cute.Tensor], sO: cute.Tensor, # FP8 high and residual Q terms used by the compensated QK WGMMA. sQ_mma: cute.Tensor, @@ -1669,7 +1704,14 @@ def mma( else: n_block_min = Int32(0) n_block_max = n_block_max_full - n_block_max_orig = n_block_max + # Mirror load_fp8: empty partitions consume a known-valid dummy + # K/V pair so the pipeline transaction count stays matched. Their + # partial result is suppressed before SplitKV combine below. + is_empty_split = n_block_min >= n_block_max + if const_expr(self.fp8_kv_dequant): + if is_empty_split: + n_block_min = Int32(0) + n_block_max = Int32(1) pipeline_q.consumer_wait_w_index_phase(0, q_consumer_phase) if const_expr(quantize_q): self._quantize_q_to_fp8(q_narrow_params) @@ -1829,11 +1871,11 @@ def mma( row_scale = softmax_tile.finalize(final_scale=v_descale_tile, sink_val=sink_val) softmax_tile.rescale_O(acc_O, row_scale) - # Override empty splits so combine kernel gives zero weight - if const_expr(self.is_split_kv): - if n_block_min >= n_block_max_orig: - acc_O.fill(Float32(0.0)) - softmax_tile.row_sum.fill(-Float32.inf) + # Make an empty partition contribute zero weight to SplitKV combine. + # This also defines an all-empty sequence as a zero output. + if is_empty_split: + acc_O.fill(Float32(0.0)) + softmax_tile.row_sum.fill(-Float32.inf) # /////////////////////////////////////////////////////////////////////////////// # Epilogue diff --git a/flash_attn/cute/interface.py b/flash_attn/cute/interface.py index 4fccaf0a25d..102d9703e08 100644 --- a/flash_attn/cute/interface.py +++ b/flash_attn/cute/interface.py @@ -790,18 +790,19 @@ def _flash_attn_fwd( "fp8_kv_dequant requires fp8 e4m3 K/V" ) assert page_table is not None, "fp8_kv_dequant requires paged KV (page_table)" - if dynamic_causal is not None: - raise NotImplementedError( - "FA4 SM90 FP8 MMA does not support dynamic_causal" - ) if block_sparse_tensors is not None: raise NotImplementedError( "FA4 SM90 FP8 MMA does not support block sparsity" ) - assert head_dim == head_dim_v == 512, ( - "FA4 SM90 FP8 MMA currently requires head_dim == head_dim_v == 512; " - f"got head_dim={head_dim}, head_dim_v={head_dim_v}." - ) + if head_dim != 512 or head_dim_v != 512: + raise NotImplementedError( + "FA4 SM90 FP8 MMA requires head_dim == head_dim_v == 512; " + f"got head_dim={head_dim}, head_dim_v={head_dim_v}." + ) + if intra_wg_overlap is False: + raise NotImplementedError( + "FA4 SM90 FP8 MMA requires intra_wg_overlap=True" + ) else: input_tensors = {"q": q, "k": k, "v": v, "qv": qv} present = {name: t for name, t in input_tensors.items() if t is not None} @@ -1035,25 +1036,22 @@ def _flash_attn_fwd( intra_wg_overlap = fwd_cfg.intra_wg_overlap if fp8_kv_dequant: - assert intra_wg_overlap, ( - "FA4 SM90 FP8 MMA requires intra_wg_overlap=True until a compatible " - "non-overlap control flow is implemented." - ) + if not intra_wg_overlap: + raise NotImplementedError( + "FA4 SM90 FP8 MMA requires intra_wg_overlap=True until a compatible " + "non-overlap control flow is implemented." + ) # Force RS-mode PV: frees sP for the fp8 staging buffer, keeping smem within # budget at d=512 while preserving intra-WG overlap. mma_pv_is_rs = True # The SM90 fp8-KV-dequant producer is TMA-only (no cp.async fallback): a paged - # page_size != tile_n gives use_tma_KV=False -> a None TMA atom. Assert here so - # it fails with a clear message instead of crashing later. - assert page_size == tile_n, ( - f"FA4 SM90 fp8-KV-dequant requires the paged-KV page_size == tile_n ({tile_n}); " - f"got page_size={page_size}. " - ) - # This path shares one staging buffer between K and V. - assert head_dim == head_dim_v, ( - "FA4 SM90 fp8-KV-dequant requires head_dim == head_dim_v (one staging " - f"buffer serves K and V); got head_dim={head_dim}, head_dim_v={head_dim_v}." - ) + # page_size != tile_n gives use_tma_KV=False -> a None TMA atom. Reject it here + # with a clear message instead of crashing later. + if page_size != tile_n: + raise NotImplementedError( + f"FA4 SM90 FP8 MMA requires the paged-KV page_size == tile_n ({tile_n}); " + f"got page_size={page_size}." + ) seqlen_q_packgqa = max_seqlen_q * (qhead_per_kvhead if pack_gqa else 1) max_m_blocks_leq_one = seqlen_q_packgqa <= q_stage * tile_m @@ -1475,6 +1473,7 @@ def _flash_attn_fwd( cu_seqlens_k is None, seqused_q is None, seqused_k is None, + dynamic_causal is not None, page_table is not None, window_size_left is not None, window_size_right is not None, diff --git a/tests/cute/test_flash_attn.py b/tests/cute/test_flash_attn.py index 9a0814c6f20..ef43083432a 100644 --- a/tests/cute/test_flash_attn.py +++ b/tests/cute/test_flash_attn.py @@ -2959,7 +2959,7 @@ def _make_sm90_fp8_mma_inputs(head_dim=512, seqlen_k=64, requires_grad=False): @pytest.mark.parametrize("head_dim", [128, 256, 384]) def test_flash_attn_kvcache_fp8_mma_rejects_unsupported_head_dim_sm90(head_dim): q, k, v, page_table, seqused_k, descale = _make_sm90_fp8_mma_inputs(head_dim) - with pytest.raises(AssertionError, match="head_dim == head_dim_v == 512"): + with pytest.raises(NotImplementedError, match="head_dim == head_dim_v == 512"): _flash_attn_fwd( q=q, k=k, @@ -2977,7 +2977,7 @@ def test_flash_attn_kvcache_fp8_mma_rejects_unsupported_head_dim_sm90(head_dim): @pytest.mark.skipif(not IS_SM90, reason="FP8-KV Tensor Core forward is SM90-only") def test_flash_attn_kvcache_fp8_mma_rejects_nonoverlap_sm90(): q, k, v, page_table, seqused_k, descale = _make_sm90_fp8_mma_inputs() - with pytest.raises(AssertionError, match="intra_wg_overlap=True"): + with pytest.raises(NotImplementedError, match="intra_wg_overlap=True"): _flash_attn_fwd( q=q, k=k, @@ -3034,21 +3034,13 @@ def test_flash_attn_kvcache_fp8_mma_rejects_qv_only_sm90(): @pytest.mark.skipif(not IS_SM90, reason="FP8-KV Tensor Core forward is SM90-only") @pytest.mark.parametrize( - "unsupported_kwarg,error", - [ - ("dynamic_causal", "does not support dynamic_causal"), - ("block_sparse_tensors", "does not support block sparsity"), - ], + "unsupported_kwarg,error", [("block_sparse_tensors", "does not support block sparsity")] ) def test_flash_attn_kvcache_fp8_mma_rejects_unsupported_scheduling_sm90( unsupported_kwarg, error ): q, k, v, page_table, seqused_k, descale = _make_sm90_fp8_mma_inputs() - unsupported_value = ( - torch.zeros(1, device="cuda", dtype=torch.bool) - if unsupported_kwarg == "dynamic_causal" - else object() - ) + unsupported_value = object() with pytest.raises(NotImplementedError, match=error): _flash_attn_fwd( q=q, @@ -3065,6 +3057,66 @@ def test_flash_attn_kvcache_fp8_mma_rejects_unsupported_scheduling_sm90( ) +@pytest.mark.skipif(not IS_SM90, reason="FP8-KV Tensor Core forward is SM90-only") +@pytest.mark.parametrize("num_splits", [1, 2]) +@maybe_fake_tensor_mode(USE_FAKE_TENSOR) +def test_flash_attn_kvcache_fp8_mma_learnable_sink_gqa_sm90(num_splits): + """FP8-MMA applies one nonzero sink logit per Q head, including SplitKV.""" + torch.manual_seed(0) + batch_size, seqlen_q, seqlen_k = 1, 7, 128 + nheads, nheads_kv, d, page_size = 4, 2, 512, 64 + q = torch.randn(batch_size, seqlen_q, nheads, d, device="cuda", dtype=torch.float16) + k_source = torch.randn( + batch_size * seqlen_k // page_size, + page_size, + nheads_kv, + d, + device="cuda", + dtype=torch.float16, + ) + v_source = torch.randn_like(k_source) + k_scale, v_scale = 0.5, 0.25 + k = (k_source / k_scale).to(torch.float8_e4m3fn) + v = (v_source / v_scale).to(torch.float8_e4m3fn) + page_table = torch.tensor([[0, 1]], device="cuda", dtype=torch.int32) + seqused_k = torch.tensor([seqlen_k], device="cuda", dtype=torch.int32) + k_descale = torch.full((batch_size, nheads_kv), k_scale, device="cuda", dtype=torch.float32) + v_descale = torch.full((batch_size, nheads_kv), v_scale, device="cuda", dtype=torch.float32) + sink = torch.tensor([1.25, -0.5, 0.75, -1.0], device="cuda", dtype=torch.float32) + + k_ref = k.view(batch_size, seqlen_k, nheads_kv, d).to(torch.float16) + v_ref = v.view(batch_size, seqlen_k, nheads_kv, d).to(torch.float16) + out_ref, _ = attention_ref( + q, + k_ref, + v_ref, + None, + None, + causal=True, + k_descale=k_descale, + v_descale=v_descale, + learnable_sink=sink, + ) + + out, *_ = _flash_attn_fwd( + q=q, + k=k, + v=v, + causal=True, + page_table=page_table, + seqused_k=seqused_k, + max_seqlen_q=seqlen_q, + max_seqlen_k=seqlen_k, + num_splits=num_splits, + k_descale=k_descale, + v_descale=v_descale, + learnable_sink=sink, + fp8_kv_dequant=True, + ) + if not is_fake_mode(): + torch.testing.assert_close(out, out_ref, atol=8e-2, rtol=8e-2) + + @pytest.mark.skipif(not IS_SM90, reason="FP8-KV Tensor Core forward is SM90-only") @pytest.mark.parametrize("seqlen_k", [1, 63, 65, 193]) @maybe_fake_tensor_mode(USE_FAKE_TENSOR) @@ -3130,6 +3182,66 @@ def test_flash_attn_kvcache_fp8_mma_v_only_descale_sm90(): torch.testing.assert_close(out.float(), expected, atol=2e-2, rtol=2e-2) +@pytest.mark.skipif(not IS_SM90, reason="FP8-KV Tensor Core forward is SM90-only") +@maybe_fake_tensor_mode(USE_FAKE_TENSOR) +def test_flash_attn_kvcache_fp8_mma_empty_splitkv_partition_sm90(): + """Empty SplitKV partitions stage a dummy tile without indexing page_table OOB.""" + torch.manual_seed(0) + batch_size, seqlen_q, max_seqlen_k = 2, 2, 64 + nheads, nheads_kv, d, page_size = 2, 1, 512, 64 + q = torch.randn(batch_size, seqlen_q, nheads, d, device="cuda", dtype=torch.float16) + k_source = torch.randn( + batch_size, + page_size, + nheads_kv, + d, + device="cuda", + dtype=torch.float16, + ) + v_source = torch.randn_like(k_source) + k_scale, v_scale = 0.5, 0.25 + k = (k_source / k_scale).to(torch.float8_e4m3fn) + v = (v_source / v_scale).to(torch.float8_e4m3fn) + page_table = torch.tensor([[0], [1]], device="cuda", dtype=torch.int32) + # Batch item 0 has no valid keys; item 1 has one tile. Four splits make + # every partition of item 0 and three partitions of item 1 empty. + seqused_k = torch.tensor([0, max_seqlen_k], device="cuda", dtype=torch.int32) + k_descale = torch.full((batch_size, nheads_kv), k_scale, device="cuda", dtype=torch.float32) + v_descale = torch.full((batch_size, nheads_kv), v_scale, device="cuda", dtype=torch.float32) + + out, *_ = _flash_attn_fwd( + q=q, + k=k, + v=v, + causal=False, + page_table=page_table, + seqused_k=seqused_k, + max_seqlen_q=seqlen_q, + max_seqlen_k=max_seqlen_k, + num_splits=4, + k_descale=k_descale, + v_descale=v_descale, + fp8_kv_dequant=True, + ) + if is_fake_mode(): + return + + k_ref = k[1:].to(torch.float16) + v_ref = v[1:].to(torch.float16) + out_ref, _ = attention_ref( + q[1:], + k_ref, + v_ref, + None, + None, + causal=False, + k_descale=k_descale[1:], + v_descale=v_descale[1:], + ) + torch.testing.assert_close(out[0], torch.zeros_like(out[0]), atol=0, rtol=0) + torch.testing.assert_close(out[1], out_ref[0], atol=8e-2, rtol=8e-2) + + @pytest.mark.skipif(not IS_SM90, reason="FP8-KV Tensor Core forward is SM90-only") @maybe_fake_tensor_mode(USE_FAKE_TENSOR) def test_flash_attn_kvcache_fp8_mma_long_context_splitkv_sm90(): @@ -3209,6 +3321,72 @@ def test_flash_attn_kvcache_fp8_mma_long_context_splitkv_sm90(): ) +@pytest.mark.skipif(not IS_SM90, reason="FP8-KV Tensor Core forward is SM90-only") +@pytest.mark.parametrize("num_splits", [1, 4]) +@maybe_fake_tensor_mode(USE_FAKE_TENSOR) +def test_flash_attn_kvcache_fp8_mma_dynamic_causal_mixed_batch_sm90(num_splits): + """A mixed causal/bidirectional batch must stage the same tile range it consumes.""" + torch.manual_seed(0) + batch_size, seqlen_q, seqlen_k = 2, 64, 128 + nheads, nheads_kv, d, page_size = 2, 1, 512, 64 + q = torch.randn(batch_size, seqlen_q, nheads, d, device="cuda", dtype=torch.float16) + k_source = torch.randn( + batch_size * seqlen_k // page_size, + page_size, + nheads_kv, + d, + device="cuda", + dtype=torch.float16, + ) + v_source = torch.randn_like(k_source) + k_scale, v_scale = 0.5, 0.25 + k = (k_source / k_scale).to(torch.float8_e4m3fn) + v = (v_source / v_scale).to(torch.float8_e4m3fn) + page_table = torch.tensor([[0, 1], [2, 3]], device="cuda", dtype=torch.int32) + seqused_k = torch.full((batch_size,), seqlen_k, device="cuda", dtype=torch.int32) + dynamic_causal = torch.tensor([1, 0], device="cuda", dtype=torch.int32) + k_descale = torch.full((batch_size, nheads_kv), k_scale, device="cuda", dtype=torch.float32) + v_descale = torch.full((batch_size, nheads_kv), v_scale, device="cuda", dtype=torch.float32) + + k_ref = k.view(batch_size, seqlen_k, nheads_kv, d).to(torch.float16) + v_ref = v.view(batch_size, seqlen_k, nheads_kv, d).to(torch.float16) + out_causal, _ = attention_ref( + q[:1], k_ref[:1], v_ref[:1], None, None, causal=True, + k_descale=k_descale[:1], v_descale=v_descale[:1], + ) + out_bidirectional, _ = attention_ref( + q[1:], k_ref[1:], v_ref[1:], None, None, causal=False, + k_descale=k_descale[1:], v_descale=v_descale[1:], + ) + out_ref = torch.cat((out_causal, out_bidirectional), dim=0) + + out_static_ref, _ = attention_ref( + q, k_ref, v_ref, None, None, causal=True, + k_descale=k_descale, v_descale=v_descale, + ) + # Alternate optional argument presence with otherwise identical shapes so + # the JIT cache cannot reuse a static-causal signature for dynamic causal. + for causal_flags in (None, dynamic_causal, None, dynamic_causal): + out, *_ = _flash_attn_fwd( + q=q, + k=k, + v=v, + causal=True, + dynamic_causal=causal_flags, + page_table=page_table, + seqused_k=seqused_k, + max_seqlen_q=seqlen_q, + max_seqlen_k=seqlen_k, + num_splits=num_splits, + k_descale=k_descale, + v_descale=v_descale, + fp8_kv_dequant=True, + ) + if not is_fake_mode(): + expected = out_static_ref if causal_flags is None else out_ref + torch.testing.assert_close(out, expected, atol=8e-2, rtol=8e-2) + + @pytest.mark.skipif(not IS_SM90, reason="FP8-KV Tensor Core forward is SM90-only") @pytest.mark.parametrize("q_magnitude", [448.0, 512.0]) @maybe_fake_tensor_mode(USE_FAKE_TENSOR) From 386976191ed66fec1841bb6e04031ec105334c05 Mon Sep 17 00:00:00 2001 From: Simon Veitner Date: Tue, 13 Oct 2026 13:09:56 +0000 Subject: [PATCH 03/13] [Cute,Fwd,Sm90] Load FP8 K with TMA into a two-slot sK FP8 K went through the single sStage buffer: TMA into sStage, then a 256-thread register copy into sK and a named barrier. sStage and sK have the same layout, so the copy only moved bytes. It also left a cross- warpgroup WAR on sK: the next block's copy can overwrite sK while the other warpgroup's QK WGMMA still reads it. Only an injected delay triggers it today. Load K with TMA straight into sK, which now has two slots with their own TMA pipeline. Each warpgroup waits for its slot, issues QK and releases the slot after its own QK wait. sStage holds only V, so the V transposes take their phase from the consumer pipeline state. The slot parity restarts per tile, so a trace-time assert rejects persistent or CLC schedulers, which would run several tiles per CTA. Output is bitwise identical to 8dbd390 on decode, split, MTP, causal and non-causal prefill, empty-split and sliding-window cases. H200, Gemma 4 TP1 heads: decode b32 16K 842 -> 747 us, prefill 8K causal 6483 -> 6017 us (do_bench); geomean over 14 vLLM-style shapes 0.909x. Shared memory grows to about 225 of 227 KiB. Signed-off-by: Simon Veitner Co-Authored-By: Claude Opus 5.5 (1M context) (cherry picked from commit 8dbc8b84254b514c7a211d5afa57eb4e68028739) Signed-off-by: yunhsuanc --- flash_attn/cute/flash_fwd_sm90.py | 129 +++++++++++++++++------------- flash_attn/cute/named_barrier.py | 4 +- 2 files changed, 75 insertions(+), 58 deletions(-) diff --git a/flash_attn/cute/flash_fwd_sm90.py b/flash_attn/cute/flash_fwd_sm90.py index d8982527daf..ecd8e3a8986 100644 --- a/flash_attn/cute/flash_fwd_sm90.py +++ b/flash_attn/cute/flash_fwd_sm90.py @@ -38,6 +38,7 @@ from flash_attn.cute.named_barrier import NamedBarrierFwd from quack.cute_dsl_utils import ParamsBase from flash_attn.cute.tile_scheduler import ( + SchedulingMode, TileSchedulerArguments, SingleTileScheduler, SingleTileLPTScheduler, @@ -85,6 +86,8 @@ def __init__( # values; the reciprocal is folded into final output normalization. self.fp8_q_scale = 1.0 self.fp8_p_scale = 256.0 + # FP8 K is TMA'd straight into two sK slots; sStage holds only V. + self.num_k_slots_fp8 = 2 assert self.use_tma_KV or not (self.check_hdim_oob or self.check_hdim_v_oob), ( "Paged KV does not support irregular head dim" ) @@ -185,11 +188,13 @@ def _get_shared_storage_cls(self): self.buffer_align_bytes, ] mbar_ptr_Stage_struct = cute.struct.MemRange[cutlass.Int64, 1 * 2] + mbar_ptr_K_fp8_struct = cute.struct.MemRange[cutlass.Int64, self.num_k_slots_fp8 * 2] @cute.struct class SharedStorageQKVFp8Mma: mbar_ptr_Q: mbar_ptr_Q_struct mbar_ptr_Stage: mbar_ptr_Stage_struct + mbar_ptr_K: mbar_ptr_K_fp8_struct sQ: sQ_input_struct sQ_mma: sQ_mma_struct sK: sK_mma_struct @@ -352,7 +357,11 @@ def __call__( # sQ holds the (possibly bf16) Q tensor as TMA'd; the fp16 compute view # is an in-place recast (both are 2 bytes) -- see mma(). (self.q_dtype, (self.tile_m, self.tile_hdim), None), - (sK_dtype, (self.tile_n, self.tile_hdim), self.num_stages), + ( + sK_dtype, + (self.tile_n, self.tile_hdim), + self.num_k_slots_fp8 if const_expr(self.fp8_kv_dequant) else self.num_stages, + ), (sV_dtype, (self.tile_n, self.tile_hdimv), self.num_stages), # sO uses the output dtype (o_dtype); for splitkv o_dtype == compute dtype # (the fp32 partials are written from acc_O directly, not via sO). @@ -373,7 +382,7 @@ def __call__( (self.tile_hdimv, self.tile_n), self.num_stages, ) - # One row-major staging tile is shared by K and V TMA loads. + # One row-major staging tile for V TMA loads; K is TMA'd into sK. self.sStage_layout = sm90_utils.make_smem_layout( self.kv_dtype, LayoutEnum.ROW_MAJOR, (self.tile_n, self.tile_hdim), 1 ) @@ -421,8 +430,7 @@ def __call__( tma_atom_K, tma_tensor_K = None, None tma_atom_V, tma_tensor_V = None, None if const_expr(self.use_tma_KV): - # FP8 K/V land in sStage for the consumer's WGMMA-layout - # rearrangement instead of being widened into sK/sV. + # FP8 K lands in sK and V in sStage (same layout); neither is widened. sK_tma_box = cute.select( self.sStage_layout if const_expr(self.fp8_kv_dequant) else self.sK_layout, mode=[0, 1], @@ -494,6 +502,11 @@ def __call__( cu_total_splits_m_blocks_ptr=mCuTotalSplitsMBlocks, ) tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args) + if const_expr(self.fp8_kv_dequant): + assert not tile_sched_args.is_persistent and ( + getattr(tile_sched_params, "scheduling_mode", SchedulingMode.STATIC) + == SchedulingMode.STATIC + ), "FP8 sK slot parity (_kslot) restarts per tile: one tile per CTA only" grid_dim = TileScheduler.get_grid_shape(tile_sched_params) softmax_scale_log2, softmax_scale = utils.compute_softmax_scale_log2( softmax_scale, self.score_mod @@ -642,7 +655,14 @@ def kernel( pipeline_stage = None if const_expr(self.fp8_kv_dequant): - pipeline_k = None + pipeline_k = pipeline_custom.PipelineTmaAsync.create( + barrier_storage=storage.mbar_ptr_K.data_ptr(), + num_stages=self.num_k_slots_fp8, + producer_group=tma_warp, + consumer_group=mma_warps, + tx_count=self.tma_copy_bytes["K"], + defer_sync=True, + ) pipeline_v = None pipeline_stage = pipeline_custom.PipelineTmaAsync.create( barrier_storage=storage.mbar_ptr_Stage.data_ptr(), @@ -854,11 +874,13 @@ def load_fp8( mK: cute.Tensor, mV: cute.Tensor, sQ: cute.Tensor, + sK: cute.Tensor, sStage: cute.Tensor, tma_atom_Q: Optional[cute.CopyAtom], tma_atom_K: Optional[cute.CopyAtom], tma_atom_V: Optional[cute.CopyAtom], pipeline_q: pipeline.PipelineAsync, + pipeline_k: pipeline.PipelineAsync, pipeline_stage: pipeline.PipelineAsync, gmem_tiled_copy_Q: cute.TiledCopy, mPageTable: Optional[cute.Tensor], @@ -867,8 +889,7 @@ def load_fp8( TileSchedulerCls: Callable, num_splits: Int32 = Int32(1), ): - """FP8-KV producer: warp 0 TMAs fp8 K/V tiles from GMEM -> the single staging buffer - sStage, staggered to match the consumer's intra-WG overlap.""" + """FP8-KV producer: warp 0 TMAs fp8 K tiles into sK slots and V tiles into sStage.""" warp_idx_in_wg = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4 tidx, _, _ = cute.arch.thread_idx() stage_full_bar = pipeline_stage.sync_object_full.get_barrier(0) @@ -905,8 +926,8 @@ def load_fp8( gK = cute.local_tile(mK_cur, (self.tile_n, self.tile_hdim), (0, 0, None)) gV = cute.local_tile(mV_cur, (self.tile_n, self.tile_hdimv), (0, 0, None)) - stage_copy_K, _, _ = copy_utils.tma_get_copy_fn( - tma_atom_K, 0, cute.make_layout(1), gK, sStage + load_K, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_K, 0, cute.make_layout(1), gK, sK ) stage_copy_V, _, _ = copy_utils.tma_get_copy_fn( tma_atom_V, 0, cute.make_layout(1), gV, sStage @@ -944,16 +965,20 @@ def load_fp8( n_block_min = Int32(0) n_block_max = Int32(1) - # ---- Prologue emit: stage K[n_block_max - 1] (warp 0 only) ---- + # ---- Prologue emit: load K[n_block_max - 1] (warp 0 only) ---- page_idx = Int32(0) if warp_idx_in_wg == 0: if is_empty_split: page_idx = Int32(0) else: page_idx = mPageTable[batch_idx, n_block_max - 1] - pipeline_stage.producer_acquire_w_index_phase(0, stage_producer_phase) - stage_copy_K(src_idx=page_idx, dst_idx=0, tma_bar_ptr=stage_full_bar) - stage_producer_phase ^= 1 + ks, kp = self._kslot(n_block_max - 1, n_block_max) + pipeline_k.producer_acquire_w_index_phase(ks, kp ^ 1) + load_K( + src_idx=page_idx, + dst_idx=ks, + tma_bar_ptr=pipeline_k.sync_object_full.get_barrier(ks), + ) # Q load (warp 0 for TMA, all 128 for pack_gqa cp.async). if const_expr(self.use_tma_Q): @@ -968,7 +993,7 @@ def load_fp8( pipeline_q.producer_commit_w_index(0) q_producer_phase ^= 1 - # ---- Mainloop + epilogue: stage K[n] then V[n_prev], then V[min]. All TMA + # ---- Mainloop + epilogue: load K[n] then V[n_prev], then V[min]. All TMA # staging is warp-0 only; the other producer warps just advance the tile # scheduler in lockstep below. ---- if warp_idx_in_wg == 0: @@ -977,10 +1002,14 @@ def load_fp8( n_block = n_block_prev - 1 page_idx = mPageTable[batch_idx, n_block] page_idx_prev = mPageTable[batch_idx, n_block_prev] - # stage K[n] - pipeline_stage.producer_acquire_w_index_phase(0, stage_producer_phase) - stage_copy_K(src_idx=page_idx, dst_idx=0, tma_bar_ptr=stage_full_bar) - stage_producer_phase ^= 1 + # load K[n] into its sK slot + ks, kp = self._kslot(n_block, n_block_max) + pipeline_k.producer_acquire_w_index_phase(ks, kp ^ 1) + load_K( + src_idx=page_idx, + dst_idx=ks, + tma_bar_ptr=pipeline_k.sync_object_full.get_barrier(ks), + ) # stage V[n_prev] pipeline_stage.producer_acquire_w_index_phase(0, stage_producer_phase) stage_copy_V(src_idx=page_idx_prev, dst_idx=0, tma_bar_ptr=stage_full_bar) @@ -1035,11 +1064,13 @@ def load( mK, mV, sQ, + sK, sStage, tma_atom_Q, tma_atom_K, tma_atom_V, pipeline_q, + pipeline_k, pipeline_stage, gmem_tiled_copy_Q, mPageTable, @@ -1450,18 +1481,9 @@ def mma( dequant_params = None if const_expr(self.fp8_kv_dequant): sStage2 = sStage[None, None, 0] - sK2 = sK[None, None, 0] sStageVt2 = layout_utils.transpose_view(sStage2) sVt2 = sVt[None, None, 0] half = const_expr(min(256, self.tile_hdimv)) - VEC = const_expr(8) - tiled_copy_K = copy_utils.tiled_copy_2d( - self.kv_dtype, - self.tile_hdim // VEC, - self.num_mma_threads, - num_copy_elems=VEC, - ) - thr_copy_K = tiled_copy_K.get_slice(tidx) # Build both WG-local output-dimension halves with static coordinates. # A dynamic local_tile coordinate loses the 128-bit alignment proof # required by STSM even though both 256-column offsets are aligned. @@ -1506,8 +1528,7 @@ def mma( dequant_params = SimpleNamespace( pipeline_stage=pipeline_stage, warp_group_idx=warp_group_idx, - tStage_K=thr_copy_K.partition_S(sStage2), - tsK=thr_copy_K.partition_D(sK2), + nmax=Int32(0), tStage_V0=tStage_V0, tStage_V1=tStage_V1, tsV0=tsV0, @@ -1712,6 +1733,7 @@ def mma( if is_empty_split: n_block_min = Int32(0) n_block_max = Int32(1) + dequant_params.nmax = n_block_max pipeline_q.consumer_wait_w_index_phase(0, q_consumer_phase) if const_expr(quantize_q): self._quantize_q_to_fp8(q_narrow_params) @@ -1972,6 +1994,15 @@ def _load_effective_descales(self, descale_tensors, batch_idx: Int32, kv_head_id v_descale = Float32(descale_tensors.v_descale[batch_idx, kv_head_idx]) return qk_descale, v_descale + @cute.jit + def _kslot(self, n_block, n_block_max): + """sK slot and its fill parity for K[n_block]; fills run n_block_max - 1 down.""" + slot = n_block % 2 + first = n_block_max - 1 + if first % 2 != slot: + first = first - 1 + return slot, ((first - n_block) // 2) % 2 + @cute.jit def _gemm_two_fp8_a( self, @@ -2045,18 +2076,6 @@ def _quantize_q_to_fp8(self, q_narrow_params): barrier_id=int(NamedBarrierFwd.NarrowQ), number_of_threads=self.num_mma_threads ) - @cute.jit - def _copy_staged_fp8_tile(self, pipeline_stage, tStage, tsDst, phase): - """Copy/reorder one staged FP8 tile into its FP8 WGMMA shared-memory layout.""" - rS = cute.make_rmem_tensor_like(tStage) - rD = cute.make_rmem_tensor_like(tsDst) - pipeline_stage.consumer_wait_w_index_phase(0, phase) - cute.autovec_copy(tStage, rS) - rD.store(rS.load()) - cute.autovec_copy(rD, tsDst) - pipeline_stage.consumer_release_w_index(0) - cute.arch.fence_view_async_shared() - @cute.jit def _transpose_fp8_v_tile(self, dp, tStage_V, tsV): for i in cutlass.range_constexpr(cute.size(tStage_V.shape[1])): @@ -2183,12 +2202,12 @@ def first_half_block_overlap_fp8( is_first_block: bool = False, dequant_params: Optional[SimpleNamespace] = None, ): - """FP8-KV first half: rearrange K, synchronize, then issue QK.""" + """FP8-KV first half: wait for K in its sK slot, issue QK, release the slot.""" dp = dequant_params - nthr = const_expr(self.num_mma_threads) - self._copy_staged_fp8_tile(dp.pipeline_stage, dp.tStage_K, dp.tsK, Int32(0)) - cute.arch.barrier(barrier_id=int(NamedBarrierFwd.DequantK), number_of_threads=nthr) - acc_S = mma_qk_fn(B_idx=kv_consumer_state.index, wg_wait=0) + ks, kp = self._kslot(n_block, dp.nmax) + pipeline_k.consumer_wait_w_index_phase(ks, kp) + acc_S = mma_qk_fn(B_idx=ks, wg_wait=0) + pipeline_k.consumer_release_w_index(ks) if const_expr(score_mod_fn is not None): score_mod_fn(acc_S, n_block=n_block, seqlen=seqlen) @@ -2240,7 +2259,7 @@ def last_half_block_overlap_fp8( softmax.rescale_O(acc_O, scores_scale) dp = dequant_params nthr_wg = const_expr(self.num_threads_per_warp_group) - self._transpose_staged_fp8_v(dp, Int32(1)) + self._transpose_staged_fp8_v(dp, kv_consumer_state.phase) if dp.warp_group_idx == 0: cute.arch.barrier(barrier_id=int(NamedBarrierFwd.DequantV0), number_of_threads=nthr_wg) else: @@ -2405,24 +2424,23 @@ def mma_one_n_block_intrawg_overlap_fp8( check_inf: cutlass.Constexpr = True, dequant_params: Optional[SimpleNamespace] = None, ): - # FP8-KV intra-WG overlap: cast K[n] for QK[n], then cast V[n_prev] + # FP8-KV intra-WG overlap: wait for K[n] for QK[n], then cast V[n_prev] # under the QK WGMMA for PV[n_prev]. Separate from the generic overlap # to keep that path unchanged. smem_pipe_read_v = smem_pipe_read.clone() smem_pipe_read.advance() dp = dequant_params - nthr = const_expr(self.num_mma_threads) nthr_wg = const_expr(self.num_threads_per_warp_group) - # ---- 1. cast K[n] -> sK (cooperative 256-thread, staging phase 1) ---- - self._copy_staged_fp8_tile(dp.pipeline_stage, dp.tStage_K, dp.tsK, Int32(1)) - cute.arch.barrier(barrier_id=int(NamedBarrierFwd.DequantK), number_of_threads=nthr) + # ---- 1. wait for K[n] in its sK slot ---- + ks, kp = self._kslot(n_block, dp.nmax) + pipeline_k.consumer_wait_w_index_phase(ks, kp) # ---- 2. S = Q @ K.T (async) ---- self.warp_scheduler_barrier_sync() - acc_S = mma_qk_fn(B_idx=smem_pipe_read.index, wg_wait=-1) + acc_S = mma_qk_fn(B_idx=ks, wg_wait=-1) if const_expr(self.rescale_O_before_gemm): softmax.rescale_O(acc_O, scores_scale) - # ---- 3. cast V[n_prev] -> sV (WG-split, phase 0) WHILE QK[n] runs ---- - self._transpose_staged_fp8_v(dp, Int32(0)) + # ---- 3. cast V[n_prev] -> sV (WG-split) WHILE QK[n] runs ---- + self._transpose_staged_fp8_v(dp, smem_pipe_read_v.phase) if dp.warp_group_idx == 0: cute.arch.barrier(barrier_id=int(NamedBarrierFwd.DequantV0), number_of_threads=nthr_wg) else: @@ -2431,6 +2449,7 @@ def mma_one_n_block_intrawg_overlap_fp8( mma_pv_fn(B_idx=smem_pipe_read_v.index, wg_wait=-1) self.warp_scheduler_barrier_arrive() warpgroup.wait_group(1) + pipeline_k.consumer_release_w_index(ks) if const_expr(score_mod_fn is not None): score_mod_fn(acc_S, n_block=n_block, seqlen=seqlen) if const_expr(mask_fn is not None): diff --git a/flash_attn/cute/named_barrier.py b/flash_attn/cute/named_barrier.py index c903f66d3aa..75c94cbbf88 100644 --- a/flash_attn/cute/named_barrier.py +++ b/flash_attn/cute/named_barrier.py @@ -10,9 +10,7 @@ class NamedBarrierFwd(enum.IntEnum): WarpSchedulerWG3 = enum.auto() PFull = enum.auto() PEmpty = enum.auto() - # FP8-KV consumer-side dequant (flash_fwd_sm90.py). DequantK: 256-thread sync so the - # full fp16 K is visible to both MMA warpgroups before QK. DequantV0/DequantV1: - # per-warpgroup 128-thread sync (each WG dequants only its own V half). + # FP8-KV (flash_fwd_sm90.py): DequantK is unused; DequantV0/1 sync each WG's V half. DequantK = enum.auto() DequantV0 = enum.auto() DequantV1 = enum.auto() From 74e90c6f092f6c66a331d44c26e6aeaf5189e6d1 Mon Sep 17 00:00:00 2001 From: Simon Veitner Date: Tue, 13 Oct 2026 13:09:56 +0000 Subject: [PATCH 04/13] [Cute,Fwd,Sm90] Transpose FP8 V between the two QK halves The V[n-1] transpose was issued after all 32 QK QGMMAs, and a cycle trace showed it never overlapped them. Issue it between the two Q terms of the QK sum instead, inside the same commit group, so it runs while the first 16 QGMMAs execute. _gemm_two_fp8_a and _gemm_qk_fp8_compensated now name their operands first and second, because the residual term can also go first. Output is bitwise identical to the previous commit. H200: decode b32 16K 747 -> 736 us, prefill 8K 6017 -> 5869 us; geomean over 14 shapes 0.896x of 8dbd390. Signed-off-by: Simon Veitner Co-Authored-By: Claude Opus 5.5 (1M context) (cherry picked from commit 1ebb85f55df23b022133619156e96b615cfdfbc1) Signed-off-by: yunhsuanc --- flash_attn/cute/flash_fwd_sm90.py | 47 +++++++++++++++++++------------ 1 file changed, 29 insertions(+), 18 deletions(-) diff --git a/flash_attn/cute/flash_fwd_sm90.py b/flash_attn/cute/flash_fwd_sm90.py index ecd8e3a8986..ce47f0410db 100644 --- a/flash_attn/cute/flash_fwd_sm90.py +++ b/flash_attn/cute/flash_fwd_sm90.py @@ -2008,32 +2008,36 @@ def _gemm_two_fp8_a( self, tiled_mma, acc, - tA_hi, - tA_res, + tA_first, + tA_second, tB, zero_init: cutlass.Constexpr[bool], wg_wait: cutlass.Constexpr[int], + mid=None, ): - """Issue high and residual FP8 A terms in one WGMMA commit group.""" + """Issue two FP8 A terms in one WGMMA commit group; mid() runs between them.""" warpgroup.fence() mma_atom = cute.make_mma_atom(tiled_mma.op) mma_atom.set(warpgroup.Field.ACCUMULATE, not zero_init) - for k in cutlass.range_constexpr(cute.size(tA_hi.shape[2])): - cute.gemm(mma_atom, acc, tA_hi[None, None, k], tB[None, None, k], acc) + for k in cutlass.range_constexpr(cute.size(tA_first.shape[2])): + cute.gemm(mma_atom, acc, tA_first[None, None, k], tB[None, None, k], acc) mma_atom.set(warpgroup.Field.ACCUMULATE, True) - for k in cutlass.range_constexpr(cute.size(tA_res.shape[2])): - cute.gemm(mma_atom, acc, tA_res[None, None, k], tB[None, None, k], acc) + if const_expr(mid is not None): + mid() + warpgroup.fence() + for k in cutlass.range_constexpr(cute.size(tA_second.shape[2])): + cute.gemm(mma_atom, acc, tA_second[None, None, k], tB[None, None, k], acc) warpgroup.commit_group() if const_expr(wg_wait >= 0): warpgroup.wait_group(wg_wait) @cute.jit def _gemm_qk_fp8_compensated( - self, tiled_mma, shape, tQ_hi, tQ_res, tK, B_idx=None, wg_wait: int = -1 + self, tiled_mma, shape, tQ_first, tQ_second, tK, B_idx=None, wg_wait: int = -1, mid=None ): acc = cute.make_rmem_tensor(tiled_mma.partition_shape_C(shape), Float32) rK = tK if const_expr(B_idx is None) else tK[None, None, None, B_idx] - self._gemm_two_fp8_a(tiled_mma, acc, tQ_hi, tQ_res, rK, True, wg_wait) + self._gemm_two_fp8_a(tiled_mma, acc, tQ_first, tQ_second, rK, True, wg_wait, mid=mid) return acc @cute.jit @@ -2425,7 +2429,7 @@ def mma_one_n_block_intrawg_overlap_fp8( dequant_params: Optional[SimpleNamespace] = None, ): # FP8-KV intra-WG overlap: wait for K[n] for QK[n], then cast V[n_prev] - # under the QK WGMMA for PV[n_prev]. Separate from the generic overlap + # between the two Q terms of QK[n] for PV[n_prev]. Separate from the generic overlap # to keep that path unchanged. smem_pipe_read_v = smem_pipe_read.clone() smem_pipe_read.advance() @@ -2434,17 +2438,24 @@ def mma_one_n_block_intrawg_overlap_fp8( # ---- 1. wait for K[n] in its sK slot ---- ks, kp = self._kslot(n_block, dp.nmax) pipeline_k.consumer_wait_w_index_phase(ks, kp) - # ---- 2. S = Q @ K.T (async) ---- + + # ---- 3. cast V[n_prev] -> sV (WG-split) while the first Q term's QGMMAs run ---- + def _vstep(): + self._transpose_staged_fp8_v(dp, smem_pipe_read_v.phase) + if dp.warp_group_idx == 0: + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.DequantV0), number_of_threads=nthr_wg + ) + else: + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.DequantV1), number_of_threads=nthr_wg + ) + + # ---- 2. S = Q @ K.T (async), with step 3 between its two Q terms ---- self.warp_scheduler_barrier_sync() - acc_S = mma_qk_fn(B_idx=ks, wg_wait=-1) + acc_S = mma_qk_fn(B_idx=ks, wg_wait=-1, mid=_vstep) if const_expr(self.rescale_O_before_gemm): softmax.rescale_O(acc_O, scores_scale) - # ---- 3. cast V[n_prev] -> sV (WG-split) WHILE QK[n] runs ---- - self._transpose_staged_fp8_v(dp, smem_pipe_read_v.phase) - if dp.warp_group_idx == 0: - cute.arch.barrier(barrier_id=int(NamedBarrierFwd.DequantV0), number_of_threads=nthr_wg) - else: - cute.arch.barrier(barrier_id=int(NamedBarrierFwd.DequantV1), number_of_threads=nthr_wg) # ---- 4. O += P @ V (async), then wait QK then PV ---- mma_pv_fn(B_idx=smem_pipe_read_v.index, wg_wait=-1) self.warp_scheduler_barrier_arrive() From 587074190209e24286d0728d4be35fdddaa28cd3 Mon Sep 17 00:00:00 2001 From: Simon Veitner Date: Tue, 13 Oct 2026 13:09:56 +0000 Subject: [PATCH 05/13] [Cute,Fwd,Sm90] Promote FP8 PV sums to FP32 every KV block FP8 WGMMA accumulates below FP32 (PTX ISA 9.7.16.5.2). The PV sum of a whole split stayed in one WGMMA accumulator, so long splits lost mass: q = 0, V = 1 gave 0.500 at 32K keys and 0.125 at 128K with num_splits=1. On real Gemma 4 activations at 16K the relative L2 error was 5.5e-2 (layer 5) and 2.3e-1 (layer 59). Run each block's PV as four 64-column passes into a zero-initialized 32-register scratch, and add the scratch to acc_O with FP32 adds. Afterwards q = 0, V = 1 gives 1.000 at every length, and layers 5 and 59 reach 3.5e-3 and 2.1e-3. The new test checks the 32K and 128K cases. After commits 1 and 2 the branch runs at 0.958x of 8dbd390; this commit alone costs 1.101x because the K copy addresses then spill. On this commit one fp16-Q case of test_flash_attn_kvcache_fp8_dequant_sm90 misses its max-error bound by one fp16 ulp (3.17e-3 vs 2.94e-3; 8dbd390 already reaches 2.93e-3 with num_splits=2). Its error comes from QK, and the next commit halves it. Signed-off-by: Simon Veitner Co-Authored-By: Claude Opus 5.5 (1M context) (cherry picked from commit 3cf312fb12836b9c6b61286df9cf880106e031a7) Signed-off-by: yunhsuanc --- flash_attn/cute/flash_fwd_sm90.py | 73 +++++++++++++++++++++++-------- tests/cute/test_flash_attn.py | 31 +++++++++++++ 2 files changed, 86 insertions(+), 18 deletions(-) diff --git a/flash_attn/cute/flash_fwd_sm90.py b/flash_attn/cute/flash_fwd_sm90.py index ce47f0410db..f3f8b5828b2 100644 --- a/flash_attn/cute/flash_fwd_sm90.py +++ b/flash_attn/cute/flash_fwd_sm90.py @@ -1444,15 +1444,16 @@ def mma( wg_mma_pv, (self.tile_m, self.tile_hdimv, self.tile_n), sP, sVt ) tOrP_res = None + pv_scratch = None if const_expr(self.fp8_kv_dequant): + # FP8 WGMMA sums below FP32 (PTX 9.7.16.5.2): promote PV every block. + assert self.rescale_O_before_gemm tOrP_res = cute.make_rmem_tensor_like(tOrP) + pv_scratch = self._make_pv_scratch( + sVt, warp_group_thread_layout(warp_group_idx), pv_cols=64 + ) mma_pv_fn = partial( - self._gemm_pv_fp8_compensated, - tiled_mma_pv, - acc_O, - tOrP, - tOrP_res, - tOrVt, + self._gemm_pv_fp8_promoted, pv_scratch, acc_O, tOrP, tOrP_res ) else: mma_pv_fn = partial(sm90_utils.gemm_w_idx, tiled_mma_pv, acc_O, tOrP, tOrVt) @@ -2040,20 +2041,56 @@ def _gemm_qk_fp8_compensated( self._gemm_two_fp8_a(tiled_mma, acc, tQ_first, tQ_second, rK, True, wg_wait, mid=mid) return acc + def _make_pv_scratch(self, sVt, wg_thr_layout, pv_cols: int): + """PV column chunks: one 64 x pv_cols WGMMA per WG into a shared FP32 scratch.""" + half = min(256, self.tile_hdimv) + assert half == 256 and self.tile_hdimv == 2 * half and half % pv_cols == 0 + tiled_mma = sm90_utils_basic.make_trivial_tiled_mma( + self.kv_dtype, + self.kv_dtype, + warpgroup.OperandMajorMode.K, + warpgroup.OperandMajorMode.K, + Float32, + atom_layout_mnk=(self.tile_m // 64, 2, 1), + tiler_mn=(64, pv_cols), + a_source=warpgroup.OperandSource.RMEM, + ) + wg_mma = tiled_mma.get_slice(wg_thr_layout) + # Rows [c * pv_cols, (c + 1) * pv_cols) of each WG's half of sVt, K-major. + chunk_layout = cute.make_layout( + ((8, pv_cols // 8, 2), (self.tile_n, 1), (1, 1)), + stride=((self.tile_n, 8 * self.tile_n, half * self.tile_n), (1, 0), (0, 0)), + ) + tVs = [] + acc = None + for c in range(half // pv_cols): + sVt_c = cute.make_tensor(sVt.iterator + c * pv_cols * self.tile_n, chunk_layout) + acc, _, tV_c = sm90_utils.partition_fragment_ABC( + wg_mma, (self.tile_m, 2 * pv_cols, self.tile_n), None, sVt_c + ) + tVs.append(tV_c) + return SimpleNamespace(tiled_mma=tiled_mma, acc=acc, tVs=tVs, n8=pv_cols // 8) + @cute.jit - def _gemm_pv_fp8_compensated( - self, - tiled_mma, - acc, - tP_hi, - tP_res, - tV, - zero_init, - B_idx=None, - wg_wait: int = -1, + def _gemm_pv_fp8_promoted( + self, pv, acc_O, tP_hi, tP_res, zero_init=False, B_idx=None, wg_wait: int = -1 ): - rV = tV if const_expr(B_idx is None) else tV[None, None, None, B_idx] - self._gemm_two_fp8_a(tiled_mma, acc, tP_hi, tP_res, rV, zero_init, wg_wait) + """O += P @ V per 64-key block, added in FP32; zero_init and wg_wait are unused.""" + for c in cutlass.range_constexpr(len(pv.tVs)): + rV = pv.tVs[c] if const_expr(B_idx is None) else pv.tVs[c][None, None, None, B_idx] + self._gemm_two_fp8_a(pv.tiled_mma, pv.acc, tP_hi, tP_res, rV, True, 0) + self._pv_fold(pv, acc_O, c) + + @cute.jit + def _pv_fold(self, pv, acc_O, c: cutlass.Constexpr[int]): + fo = cute.group_modes(acc_O, 1, cute.rank(acc_O)) + fc = cute.group_modes(pv.acc, 1, cute.rank(pv.acc)) + for mi in cutlass.range_constexpr(cute.size(fc.shape[1])): + for e in cutlass.range_constexpr(2): + for r in cutlass.range_constexpr(2): + for j in cutlass.range_constexpr(pv.n8): + idx = ((e, r, c * pv.n8 + j), mi) + fo[idx] = fo[idx] + fc[((e, r, j), mi)] @cute.jit def _quantize_q_to_fp8(self, q_narrow_params): diff --git a/tests/cute/test_flash_attn.py b/tests/cute/test_flash_attn.py index ef43083432a..3dd54ad47d6 100644 --- a/tests/cute/test_flash_attn.py +++ b/tests/cute/test_flash_attn.py @@ -3321,6 +3321,37 @@ def test_flash_attn_kvcache_fp8_mma_long_context_splitkv_sm90(): ) +@pytest.mark.skipif(not IS_SM90, reason="FP8-KV Tensor Core forward is SM90-only") +@pytest.mark.parametrize("seqlen_k", [32768, 131072]) +@maybe_fake_tensor_mode(USE_FAKE_TENSOR) +def test_flash_attn_kvcache_fp8_mma_long_split_keeps_mass_sm90(seqlen_k): + """q = 0 and V = 1 give exactly 1 at any length in one split.""" + nheads, nheads_kv, d, page_size = 2, 1, 512, 64 + n_pages = seqlen_k // page_size + q = torch.zeros(1, 1, nheads, d, device="cuda", dtype=torch.bfloat16) + shape = (n_pages, page_size, nheads_kv, d) + k = torch.zeros(shape, device="cuda").to(torch.float8_e4m3fn) + v = torch.ones(shape, device="cuda").to(torch.float8_e4m3fn) + page_table = torch.arange(n_pages, device="cuda", dtype=torch.int32).unsqueeze(0) + seqused_k = torch.tensor([seqlen_k], device="cuda", dtype=torch.int32) + out, *_ = _flash_attn_fwd( + q=q, + k=k, + v=v, + page_table=page_table, + seqused_k=seqused_k, + max_seqlen_q=1, + max_seqlen_k=seqlen_k, + softmax_scale=1.0, + num_splits=1, + fp8_kv_dequant=True, + ) + if is_fake_mode(): + return + ones = torch.ones_like(out, dtype=torch.float32) + torch.testing.assert_close(out.float(), ones, atol=1e-3, rtol=0) + + @pytest.mark.skipif(not IS_SM90, reason="FP8-KV Tensor Core forward is SM90-only") @pytest.mark.parametrize("num_splits", [1, 4]) @maybe_fake_tensor_mode(USE_FAKE_TENSOR) From 062518f75cfa74bf8cf60a693ed89ba91267e7f9 Mon Sep 17 00:00:00 2001 From: Simon Veitner Date: Tue, 13 Oct 2026 13:09:56 +0000 Subject: [PATCH 06/13] [Cute,Fwd,Sm90] Issue the residual Q term first in FP8 QK Each QGMMA truncates its addends relative to the running sum, so the residual Q terms lose most of their bits when they are added after the high terms. Issue the 16 residual k-steps first, while the sum is still small, by passing the residual Q fragment as the first operand. Same instructions and registers, no measurable time change (0.958x of 8dbd390 after commits 1 to 4, as after commits 1 to 3). On top of commits 1 to 3: test_flash_attn_fp8_kv_cache.py goes from 10 to 4 failures of 14, real Gemma 4 16K error from 3.5e-3 to 3.0e-3 (layer 5) and 2.1e-3 to 1.7e-3 (layer 59), and the fp16-Q case of test_flash_attn_kvcache_fp8_dequant_sm90 from 3.174e-3 to 1.465e-3 max error. Signed-off-by: Simon Veitner Co-Authored-By: Claude Opus 5.5 (1M context) (cherry picked from commit 73ff1378309c958106536dfe601ccffe64982f76) Signed-off-by: yunhsuanc --- flash_attn/cute/flash_fwd_sm90.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/flash_attn/cute/flash_fwd_sm90.py b/flash_attn/cute/flash_fwd_sm90.py index f3f8b5828b2..92ab02dadee 100644 --- a/flash_attn/cute/flash_fwd_sm90.py +++ b/flash_attn/cute/flash_fwd_sm90.py @@ -1428,8 +1428,9 @@ def mma( self._gemm_qk_fp8_compensated, tiled_mma_qk, (self.tile_m, self.tile_n), - tSrQ, + # Residual terms first, while the running sum is still small. tSrQ_res, + tSrQ, tSrK, ) else: From af38305c52b16d11f1bad49f7bd1f84ab5285120 Mon Sep 17 00:00:00 2001 From: Simon Veitner Date: Tue, 13 Oct 2026 13:09:56 +0000 Subject: [PATCH 07/13] [Cute,Fwd,Sm90] Split the FP8 QK sum in causal and local builds The 4 remaining test_flash_attn_fp8_kv_cache.py failures (logits with std about 22) come from the QK sum: 32 QGMMAs into one accumulator. Sum k-steps 0 to 7 into acc_S and 8 to 15 into the PV scratch, which is idle while QK runs, then wait and add the scratch to acc_S before the first PV pass. The first block uses its own 32-register temporary. Both keep the residual-first order of the previous commit. Only causal and local builds get the split: in non-causal builds ptxas serializes the WGMMAs (C7511), which ran 1.33 to 2.03x slower. With the gate, ptxas -v reports no C7511 for causal, local and non-causal builds (num_splits 1 and 4); without it, the non-causal builds report C7511. With commits 1 to 5 all 14 tests pass, real Gemma 4 16K error is 1.97e-3 (layer 5) and 1.69e-3 (layer 59), the fp16-Q case of test_flash_attn_kvcache_fp8_dequant_sm90 drops to 9.8e-4 max error, and time is 0.982x of 8dbd390 (causal). Non-causal builds keep the accuracy of the previous commit: the 14 tests made non-causal still fail 4. Signed-off-by: Simon Veitner Co-Authored-By: Claude Opus 5.5 (1M context) (cherry picked from commit 3d4be5f852305eac8a36d7b7126f7d85051f3279) Signed-off-by: yunhsuanc --- flash_attn/cute/flash_fwd_sm90.py | 75 +++++++++++++++++++++++++++++-- 1 file changed, 71 insertions(+), 4 deletions(-) diff --git a/flash_attn/cute/flash_fwd_sm90.py b/flash_attn/cute/flash_fwd_sm90.py index 92ab02dadee..aaedfd7d643 100644 --- a/flash_attn/cute/flash_fwd_sm90.py +++ b/flash_attn/cute/flash_fwd_sm90.py @@ -76,6 +76,8 @@ def __init__( # SM90 compensated FP8-MMA path: native fp16/bf16 Q, E4M3 paged K/V, # FP8 WGMMA QK/PV, FP32 accumulation, and native fp16/bf16 O. self.fp8_kv_dequant = fp8_kv_dequant + # Split the FP8 QK sum in two; non-causal builds would serialize WGMMAs (ptxas C7511). + self.qk_split = fp8_kv_dequant and (self.is_causal or self.is_local) assert not self.fp8_kv_dequant or self.mma_pv_is_rs, ( "SM90 FP8 KV dequant requires RS-mode PV because its shared storage " "does not allocate sP" @@ -1588,6 +1590,7 @@ def mma( check_inf=True, scores_scale=scores_scale, **(dict(dequant_params=dequant_params) if const_expr(self.fp8_kv_dequant) else {}), + **(dict(qk_split_acc=pv_scratch.acc) if const_expr(self.qk_split) else {}), ) process_first_half_block = partial( @@ -2035,13 +2038,69 @@ def _gemm_two_fp8_a( @cute.jit def _gemm_qk_fp8_compensated( - self, tiled_mma, shape, tQ_first, tQ_second, tK, B_idx=None, wg_wait: int = -1, mid=None + self, + tiled_mma, + shape, + tQ_first, + tQ_second, + tK, + B_idx=None, + wg_wait: int = -1, + mid=None, + split_acc=None, ): acc = cute.make_rmem_tensor(tiled_mma.partition_shape_C(shape), Float32) rK = tK if const_expr(B_idx is None) else tK[None, None, None, B_idx] - self._gemm_two_fp8_a(tiled_mma, acc, tQ_first, tQ_second, rK, True, wg_wait, mid=mid) + if const_expr(not self.qk_split): + self._gemm_two_fp8_a(tiled_mma, acc, tQ_first, tQ_second, rK, True, wg_wait, mid=mid) + elif const_expr(split_acc is None): + # First block: own second accumulator, always waited for and folded here. + acc2 = cute.make_rmem_tensor(tiled_mma.partition_shape_C(shape), Float32) + self._gemm_qk_fp8_split(tiled_mma, acc, acc2, tQ_first, tQ_second, rK, None) + warpgroup.wait_group(0) + self._fold_qk_split(acc, acc2) + else: + # wg_wait is unused: the caller waits and folds split_acc before reusing it. + self._gemm_qk_fp8_split(tiled_mma, acc, split_acc, tQ_first, tQ_second, rK, mid) return acc + @cute.jit + def _gemm_qk_fp8_split( + self, + tiled_mma, + acc, + acc2, + tA_first, + tA_second, + rK, + mid, + ): + """One commit group: k-steps in the first half sum into acc, the rest into acc2.""" + nk = const_expr(cute.size(tA_first.shape[2])) + warpgroup.fence() + a1 = cute.make_mma_atom(tiled_mma.op) + a2 = cute.make_mma_atom(tiled_mma.op) + a1.set(warpgroup.Field.ACCUMULATE, False) + a2.set(warpgroup.Field.ACCUMULATE, False) + for t in cutlass.range_constexpr(2): + tA = tA_first if const_expr(t == 0) else tA_second + if const_expr(t == 1 and mid is not None): + mid() + warpgroup.fence() + for k in cutlass.range_constexpr(nk): + if const_expr(k < nk // 2): + cute.gemm(a1, acc, tA[None, None, k], rK[None, None, k], acc) + a1.set(warpgroup.Field.ACCUMULATE, True) + else: + cute.gemm(a2, acc2, tA[None, None, k], rK[None, None, k], acc2) + a2.set(warpgroup.Field.ACCUMULATE, True) + warpgroup.commit_group() + + @cute.jit + def _fold_qk_split(self, acc, acc2): + for i in cutlass.range_constexpr(cute.size(acc)): + acc[i] = acc[i] + acc2[i] + def _make_pv_scratch(self, sVt, wg_thr_layout, pv_cols: int): """PV column chunks: one 64 x pv_cols WGMMA per WG into a shared FP32 scratch.""" half = min(256, self.tile_hdimv) @@ -2465,6 +2524,7 @@ def mma_one_n_block_intrawg_overlap_fp8( mask_fn: Optional[Callable] = None, check_inf: cutlass.Constexpr = True, dequant_params: Optional[SimpleNamespace] = None, + qk_split_acc: Optional[cute.Tensor] = None, ): # FP8-KV intra-WG overlap: wait for K[n] for QK[n], then cast V[n_prev] # between the two Q terms of QK[n] for PV[n_prev]. Separate from the generic overlap @@ -2491,10 +2551,17 @@ def _vstep(): # ---- 2. S = Q @ K.T (async), with step 3 between its two Q terms ---- self.warp_scheduler_barrier_sync() - acc_S = mma_qk_fn(B_idx=ks, wg_wait=-1, mid=_vstep) + if const_expr(self.qk_split): + acc_S = mma_qk_fn(B_idx=ks, wg_wait=-1, mid=_vstep, split_acc=qk_split_acc) + else: + acc_S = mma_qk_fn(B_idx=ks, wg_wait=-1, mid=_vstep) if const_expr(self.rescale_O_before_gemm): softmax.rescale_O(acc_O, scores_scale) - # ---- 4. O += P @ V (async), then wait QK then PV ---- + # ---- 4. O += P @ V; the promoted passes wait for QK and PV ---- + if const_expr(self.qk_split): + # The PV passes reuse the scratch: QK must land and fold first. + warpgroup.wait_group(0) + self._fold_qk_split(acc_S, qk_split_acc) mma_pv_fn(B_idx=smem_pipe_read_v.index, wg_wait=-1) self.warp_scheduler_barrier_arrive() warpgroup.wait_group(1) From 30d694edde0ca92e6c3097d3cc9e55d7483ba798 Mon Sep 17 00:00:00 2001 From: Simon Veitner Date: Tue, 13 Oct 2026 13:09:56 +0000 Subject: [PATCH 08/13] [Cute,Fwd,Sm90] Document FP8 MMA accumulation and fp16 Q limits Comments said the FP8 path accumulates in FP32. FP8 WGMMA sums below FP32 within an instruction chain (PTX ISA 9.7.16.5.2); only the per-block PV promotion adds in FP32. State that, state that fp16 Q keeps 8 of its 11 significant bits in the E4M3 pair, and give the real range of the large-Q test (values up to 896, at most 2.2% error). Signed-off-by: Simon Veitner Co-Authored-By: Claude Opus 5.5 (1M context) (cherry picked from commit 322a9af1689fa42b181a729a54dfd8159832939b) Signed-off-by: yunhsuanc --- flash_attn/cute/flash_fwd_sm90.py | 5 ++--- flash_attn/cute/interface.py | 9 +++------ tests/cute/test_flash_attn.py | 3 +-- 3 files changed, 6 insertions(+), 11 deletions(-) diff --git a/flash_attn/cute/flash_fwd_sm90.py b/flash_attn/cute/flash_fwd_sm90.py index aaedfd7d643..1a4ac103478 100644 --- a/flash_attn/cute/flash_fwd_sm90.py +++ b/flash_attn/cute/flash_fwd_sm90.py @@ -73,8 +73,7 @@ def __init__( self.buffer_align_bytes = 1024 self.use_tma_KV = not paged_kv_non_tma self.is_split_kv = is_split_kv - # SM90 compensated FP8-MMA path: native fp16/bf16 Q, E4M3 paged K/V, - # FP8 WGMMA QK/PV, FP32 accumulation, and native fp16/bf16 O. + # SM90 FP8-MMA path: FP8 WGMMA sums below FP32, so PV is promoted per KV block. self.fp8_kv_dequant = fp8_kv_dequant # Split the FP8 QK sum in two; non-causal builds would serialize WGMMAs (ptxas C7511). self.qk_split = fp8_kv_dequant and (self.is_causal or self.is_local) @@ -271,7 +270,7 @@ def __call__( else: # FP8-KV: K/V are fp8 (rejected by the symmetric _check_type). Q/O may each be # fp16 or bf16. Q and softmax P use compensated pairs of E4M3 values; - # QK and PV both execute as FP8 WGMMA with fp32 accumulation. + # QK and PV run as FP8 WGMMA, which sums below FP32; PV is promoted per block. assert self.dtype == Float16, ( "fp8_kv_dequant requires fp16 auxiliary conversion layouts" ) diff --git a/flash_attn/cute/interface.py b/flash_attn/cute/interface.py index 102d9703e08..373c11801cb 100644 --- a/flash_attn/cute/interface.py +++ b/flash_attn/cute/interface.py @@ -777,15 +777,14 @@ def _flash_attn_fwd( "inputs must be float16, bfloat16, fp8 e4m3fn, or fp8 e5m2" ) if fp8_kv_dequant: - # SM90 FP8 MMA: fp16/bf16 Q is represented as compensated E4M3 pairs; - # paged E4M3 K/V remain FP8 WGMMA operands, and accumulation is FP32. - # The epilogue casts FP32 accumulators to the native fp16/bf16 output dtype. + # SM90 FP8 MMA: E4M3 Q pairs and K/V; WGMMA sums below FP32 (PTX 9.7.16.5.2). assert q is not None and qv is None, ( "fp8_kv_dequant requires Q and does not support QV-only input" ) assert q.dtype in (torch.float16, torch.bfloat16), ( "fp8_kv_dequant requires fp16 or bf16 Q" ) + # fp16 Q keeps 8 of its 11 significant bits in the E4M3 pair; bf16 Q is exact. assert k.dtype == v.dtype == torch.float8_e4m3fn, ( "fp8_kv_dequant requires fp8 e4m3 K/V" ) @@ -981,9 +980,7 @@ def _flash_attn_fwd( kv_dtype = torch2cute_dtype_map[k.dtype] if fp8_kv_dequant else dtype if fp8_kv_dequant: assert arch // 10 == 9, "fp8_kv_dequant is an SM90-only forward (compute capability 9.x)" - # This dtype selects fp16 auxiliary layouts/conversions. QK and PV use - # compensated E4M3 operands with FP8 WGMMA and FP32 accumulation; mQ/mO - # retain the caller's fp16/bf16 input and output dtypes. + # fp16 selects the auxiliary layouts; mQ/mO keep the caller's fp16/bf16 dtypes. dtype = torch2cute_dtype_map[torch.float16] if is_fp8: assert arch // 10 == 10, "FP8 is only supported on SM100 (compute capability 10.x) for FA4 CuTe." diff --git a/tests/cute/test_flash_attn.py b/tests/cute/test_flash_attn.py index 3dd54ad47d6..276ecfad106 100644 --- a/tests/cute/test_flash_attn.py +++ b/tests/cute/test_flash_attn.py @@ -3425,8 +3425,7 @@ def test_flash_attn_kvcache_fp8_mma_large_q_sm90(q_magnitude): """Stress the fixed Q scale at and beyond one E4M3 component's range. Q is represented as two compensated E4M3 components. The 512 case therefore - also verifies that the residual component recovers values above 448 instead - of silently relying on a single E4M3 conversion. K is inversely scaled so + also covers values above 448 (up to 896, at most 2.2% error). K is inversely scaled so the attention-score envelope remains well conditioned. """ torch.manual_seed(0) From 11e6121b1bc117752effbbcce1b76ae72287ac1e Mon Sep 17 00:00:00 2001 From: Simon Veitner Date: Tue, 13 Oct 2026 13:09:56 +0000 Subject: [PATCH 09/13] [Cute,Fwd,Sm90] Fence sStage reads before releasing it to TMA The V transpose reads sStage with LDSM (generic proxy), and the next V tile is written into sStage by TMA (async proxy). The proxy fence ran after the release, and only lane 0 of each warp arrives on the empty barrier, so the reads were not ordered before the TMA write (PTX ISA 8.6, 9.7.9.26.2 and 9.7.14.4). Issue the fence and a __syncwarp before the release. The fence still precedes the per-warpgroup barrier that makes the transposed sV visible to the PV WGMMA. No hang or corruption was observed without this; it closes the gap in the memory model. The release-then-fence order comes from the PR 164 dequant path and 8dbd390 kept it. On its own this commit fixes the V path only: on 8dbd390 the K copy releases sStage in the same order until the first commit of this series removes that copy. Output is bitwise identical to the previous commit and time is unchanged (0.981x of 8dbd390 for the whole branch). Signed-off-by: Simon Veitner Co-Authored-By: Claude Opus 5.5 (1M context) (cherry picked from commit 34337517bfd7a77bc539c0c625a2ae611fe230ca) Signed-off-by: yunhsuanc --- flash_attn/cute/flash_fwd_sm90.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/flash_attn/cute/flash_fwd_sm90.py b/flash_attn/cute/flash_fwd_sm90.py index 1a4ac103478..c51ba0152b9 100644 --- a/flash_attn/cute/flash_fwd_sm90.py +++ b/flash_attn/cute/flash_fwd_sm90.py @@ -2201,8 +2201,10 @@ def _transpose_staged_fp8_v(self, dp, phase): self._transpose_fp8_v_tile(dp, dp.tStage_V0, dp.tsV0) else: self._transpose_fp8_v_tile(dp, dp.tStage_V1, dp.tsV1) - dp.pipeline_stage.consumer_release_w_index(0) + # Order the LDSM reads before the next V TMA; lane 0 of each warp releases. cute.arch.fence_view_async_shared() + cute.arch.sync_warp() + dp.pipeline_stage.consumer_release_w_index(0) @cute.jit def _permute_Cregs_fp8(self, acc): From 6d71c3f914eb7943f49647ea4828c6b2fe9176c6 Mon Sep 17 00:00:00 2001 From: Kate Cheng Date: Wed, 14 Oct 2026 15:41:03 -0700 Subject: [PATCH 10/13] Format SM90 FP8 forward kernel Signed-off-by: Kate Cheng --- flash_attn/cute/flash_fwd_sm90.py | 28 +++++++--------------------- 1 file changed, 7 insertions(+), 21 deletions(-) diff --git a/flash_attn/cute/flash_fwd_sm90.py b/flash_attn/cute/flash_fwd_sm90.py index c51ba0152b9..e298e3134a9 100644 --- a/flash_attn/cute/flash_fwd_sm90.py +++ b/flash_attn/cute/flash_fwd_sm90.py @@ -135,9 +135,7 @@ def _get_tiled_mma(self): mma_dtype, mma_dtype, warpgroup.OperandMajorMode.K, - warpgroup.OperandMajorMode.K - if self.fp8_kv_dequant - else warpgroup.OperandMajorMode.MN, + warpgroup.OperandMajorMode.K if self.fp8_kv_dequant else warpgroup.OperandMajorMode.MN, Float32, atom_layout_mnk=( self.tile_m // 64, @@ -320,9 +318,7 @@ def __call__( self.num_threads_per_warp_group = 128 self.num_wg_mma = self.num_mma_threads // self.num_threads_per_warp_group if const_expr(self.fp8_kv_dequant): - assert self.num_wg_mma == 2, ( - "SM90 d512 FP8 MMA requires exactly two MMA warp groups" - ) + assert self.num_wg_mma == 2, "SM90 d512 FP8 MMA requires exactly two MMA warp groups" assert self.num_wg_mma in [1, 2, 3] self.num_threads = self.num_threads_per_warp_group * (self.num_wg_mma + 1) self.num_producer_threads = 32 @@ -723,18 +719,14 @@ def kernel( sQ_mma = sQ sQ_res = sQ if const_expr(self.fp8_kv_dequant): - sQ_mma = storage.sQ_mma.get_tensor( - sQ_mma_layout.outer, swizzle=sQ_mma_layout.inner - ) + sQ_mma = storage.sQ_mma.get_tensor(sQ_mma_layout.outer, swizzle=sQ_mma_layout.inner) # Reuse the now-dead FP16 Q input bytes for the FP8 residual term. sQ_res = storage.sQ.get_tensor( sQ_mma_layout.outer, swizzle=sQ_mma_layout.inner, dtype=self.kv_dtype ) sK = storage.sK.get_tensor(sK_layout.outer, swizzle=sK_layout.inner) if const_expr(self.fp8_kv_dequant): - sVt = storage.sV_mma.get_tensor( - sV_mma_layout.outer, swizzle=sV_mma_layout.inner - ) + sVt = storage.sV_mma.get_tensor(sV_mma_layout.outer, swizzle=sV_mma_layout.inner) sV = layout_utils.transpose_view(sVt) elif const_expr(not self.Q_in_regs): sV = storage.sV.get_tensor(sV_layout.outer, swizzle=sV_layout.inner) @@ -927,9 +919,7 @@ def load_fp8( gK = cute.local_tile(mK_cur, (self.tile_n, self.tile_hdim), (0, 0, None)) gV = cute.local_tile(mV_cur, (self.tile_n, self.tile_hdimv), (0, 0, None)) - load_K, _, _ = copy_utils.tma_get_copy_fn( - tma_atom_K, 0, cute.make_layout(1), gK, sK - ) + load_K, _, _ = copy_utils.tma_get_copy_fn(tma_atom_K, 0, cute.make_layout(1), gK, sK) stage_copy_V, _, _ = copy_utils.tma_get_copy_fn( tma_atom_V, 0, cute.make_layout(1), gV, sStage ) @@ -946,9 +936,7 @@ def load_fp8( if not psc: n_block_max_full = cute.ceil_div(seqlen.seqlen_k, self.tile_n) if const_expr(self.is_split_kv): - num_n_blocks_per_split = cute.ceil_div( - n_block_max_full, num_splits - ) + num_n_blocks_per_split = cute.ceil_div(n_block_max_full, num_splits) n_block_min = split_idx * num_n_blocks_per_split n_block_max = cutlass.min( n_block_min + num_n_blocks_per_split, n_block_max_full @@ -1454,9 +1442,7 @@ def mma( pv_scratch = self._make_pv_scratch( sVt, warp_group_thread_layout(warp_group_idx), pv_cols=64 ) - mma_pv_fn = partial( - self._gemm_pv_fp8_promoted, pv_scratch, acc_O, tOrP, tOrP_res - ) + mma_pv_fn = partial(self._gemm_pv_fp8_promoted, pv_scratch, acc_O, tOrP, tOrP_res) else: mma_pv_fn = partial(sm90_utils.gemm_w_idx, tiled_mma_pv, acc_O, tOrP, tOrVt) From fe4f9ce1ff1a394d13ec961c914eca562dc60a76 Mon Sep 17 00:00:00 2001 From: Simon Veitner Date: Mon, 26 Oct 2026 16:02:37 +0100 Subject: [PATCH 11/13] [Cute,Fwd,Sm90] Keep the FP8 P tail: scale the P residual in its own PV pass P = e4m3(256p) + e4m3(256p - hi) rounds both terms to 0 for p below 2^-18 of the running row max, so PV drops those keys while row_sum still counts them. Blocks run last to first, so a dominant recent key cuts the earlier context out of the output. Scale the residual by 32 (|residual| <= 8 for p <= 1, so 256 < 448) and run it as a second PV pass through the same FP32 scratch, unscaled at the fold. Keys now reach PV down to 2^-23 of the row max. Measured on H200 against 6d71c3f: - Gemma-4 31B FP8, 121K prefill, 10 global layers: per-row error p99 0.29-0.42% (was 0.86-2.97%; merge base 0.26-0.31%), max 0.48-0.98% (was 2.5-7.9%). - New test_flash_attn_kvcache_fp8_mma_keeps_softmax_tail_sm90 fails before (+24%, -100%) and passes after; FP8 suite 41 passed, 2 skipped. - Kernel time +6% to +13% (prefill 8K x 123K +7%, 2K x 33K +13%, decode b64 x 33K +10%, b1 x 123K +6%). Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Simon Veitner Signed-off-by: Kate Cheng --- flash_attn/cute/flash_fwd_sm90.py | 40 +++++++++++++++++++++++----- tests/cute/test_flash_attn.py | 43 +++++++++++++++++++++++++++++++ 2 files changed, 77 insertions(+), 6 deletions(-) diff --git a/flash_attn/cute/flash_fwd_sm90.py b/flash_attn/cute/flash_fwd_sm90.py index e298e3134a9..ff6cc1d4bdf 100644 --- a/flash_attn/cute/flash_fwd_sm90.py +++ b/flash_attn/cute/flash_fwd_sm90.py @@ -87,6 +87,10 @@ def __init__( # values; the reciprocal is folded into final output normalization. self.fp8_q_scale = 1.0 self.fp8_p_scale = 256.0 + # The P residual gets its own scale and PV pass, so a probability down to + # 2^-23 of the row max survives instead of rounding to 0 below 2^-18. + # |residual| <= 8 for p <= 1, so 8 * 32 = 256 stays below E4M3's 448. + self.fp8_p_res_scale = 32.0 # FP8 K is TMA'd straight into two sK slots; sStage holds only V. self.num_k_slots_fp8 = 2 assert self.use_tma_KV or not (self.check_hdim_oob or self.check_hdim_v_oob), ( @@ -2120,14 +2124,34 @@ def _make_pv_scratch(self, sVt, wg_thr_layout, pv_cols: int): def _gemm_pv_fp8_promoted( self, pv, acc_O, tP_hi, tP_res, zero_init=False, B_idx=None, wg_wait: int = -1 ): - """O += P @ V per 64-key block, added in FP32; zero_init and wg_wait are unused.""" + """O += P @ V per 64-key block, added in FP32; zero_init and wg_wait are unused. + + P_hi and the scaled P residual run as separate passes through the same + scratch, so the residual is unscaled in FP32 when it is folded into O. + """ for c in cutlass.range_constexpr(len(pv.tVs)): rV = pv.tVs[c] if const_expr(B_idx is None) else pv.tVs[c][None, None, None, B_idx] - self._gemm_two_fp8_a(pv.tiled_mma, pv.acc, tP_hi, tP_res, rV, True, 0) - self._pv_fold(pv, acc_O, c) + self._gemm_one_fp8_a(pv.tiled_mma, pv.acc, tP_hi, rV) + self._pv_fold(pv, acc_O, c, 1.0) + self._gemm_one_fp8_a(pv.tiled_mma, pv.acc, tP_res, rV) + self._pv_fold(pv, acc_O, c, 1.0 / self.fp8_p_res_scale) @cute.jit - def _pv_fold(self, pv, acc_O, c: cutlass.Constexpr[int]): + def _gemm_one_fp8_a(self, tiled_mma, acc, tA, tB): + """acc = tA @ tB as one synchronous FP8 WGMMA commit group.""" + warpgroup.fence() + mma_atom = cute.make_mma_atom(tiled_mma.op) + mma_atom.set(warpgroup.Field.ACCUMULATE, False) + for k in cutlass.range_constexpr(cute.size(tA.shape[2])): + cute.gemm(mma_atom, acc, tA[None, None, k], tB[None, None, k], acc) + mma_atom.set(warpgroup.Field.ACCUMULATE, True) + warpgroup.commit_group() + warpgroup.wait_group(0) + + @cute.jit + def _pv_fold( + self, pv, acc_O, c: cutlass.Constexpr[int], scale: cutlass.Constexpr[float] + ): fo = cute.group_modes(acc_O, 1, cute.rank(acc_O)) fc = cute.group_modes(pv.acc, 1, cute.rank(pv.acc)) for mi in cutlass.range_constexpr(cute.size(fc.shape[1])): @@ -2135,7 +2159,10 @@ def _pv_fold(self, pv, acc_O, c: cutlass.Constexpr[int]): for r in cutlass.range_constexpr(2): for j in cutlass.range_constexpr(pv.n8): idx = ((e, r, c * pv.n8 + j), mi) - fo[idx] = fo[idx] + fc[((e, r, j), mi)] + if const_expr(scale == 1.0): + fo[idx] = fo[idx] + fc[((e, r, j), mi)] + else: + fo[idx] = fo[idx] + fc[((e, r, j), mi)] * Float32(scale) @cute.jit def _quantize_q_to_fp8(self, q_narrow_params): @@ -2270,7 +2297,8 @@ def _prepare_fp8_probabilities(self, acc_S, tOrP, tOrP_res): p_scaled = tP_acc.load() * Float32(self.fp8_p_scale) p_hi = p_scaled.to(self.kv_dtype) tOrP.store(p_hi) - tOrP_res.store((p_scaled - p_hi.to(Float32)).to(self.kv_dtype)) + p_res = (p_scaled - p_hi.to(Float32)) * Float32(self.fp8_p_res_scale) + tOrP_res.store(p_res.to(self.kv_dtype)) @cute.jit def first_half_block_overlap_fp8( diff --git a/tests/cute/test_flash_attn.py b/tests/cute/test_flash_attn.py index 276ecfad106..37c1d16afbd 100644 --- a/tests/cute/test_flash_attn.py +++ b/tests/cute/test_flash_attn.py @@ -3352,6 +3352,49 @@ def test_flash_attn_kvcache_fp8_mma_long_split_keeps_mass_sm90(seqlen_k): torch.testing.assert_close(out.float(), ones, atol=1e-3, rtol=0) +@pytest.mark.skipif(not IS_SM90, reason="FP8-KV Tensor Core forward is SM90-only") +@pytest.mark.parametrize("delta", [12.0, 13.0]) +@maybe_fake_tensor_mode(USE_FAKE_TENSOR) +def test_flash_attn_kvcache_fp8_mma_keeps_softmax_tail_sm90(delta): + """The last key leads every earlier key by `delta`; the earlier keys hold the tail. + + KV blocks run last to first, so each earlier p = exp(-delta) is quantized against + the last key. It must reach PV as well as row_sum, or the output drops toward 0. + Every tail key has the same p, so one E4M3 rounding (up to 6.25%) hits them all. + """ + seqlen_k, nheads, nheads_kv, d, page_size = 32768, 2, 1, 512, 64 + n_pages = seqlen_k // page_size + q = torch.zeros(1, 1, nheads, d, device="cuda", dtype=torch.bfloat16) + q[..., 0] = 1.0 + k = torch.zeros(seqlen_k, nheads_kv, d, device="cuda") + v = torch.ones(seqlen_k, nheads_kv, d, device="cuda") + k[-1, :, 0] = 1.0 + v[-1] = 0.0 + shape = (n_pages, page_size, nheads_kv, d) + k = k.to(torch.float8_e4m3fn).view(shape) + v = v.to(torch.float8_e4m3fn).view(shape) + page_table = torch.arange(n_pages, device="cuda", dtype=torch.int32).unsqueeze(0) + seqused_k = torch.tensor([seqlen_k], device="cuda", dtype=torch.int32) + out, *_ = _flash_attn_fwd( + q=q, + k=k, + v=v, + page_table=page_table, + seqused_k=seqused_k, + max_seqlen_q=1, + max_seqlen_k=seqlen_k, + softmax_scale=delta, + causal=True, + num_splits=1, + fp8_kv_dequant=True, + ) + if is_fake_mode(): + return + tail = (seqlen_k - 1) * math.exp(-delta) + expected = torch.full_like(out, tail / (tail + 1.0), dtype=torch.float32) + torch.testing.assert_close(out.float(), expected, atol=0, rtol=7e-2) + + @pytest.mark.skipif(not IS_SM90, reason="FP8-KV Tensor Core forward is SM90-only") @pytest.mark.parametrize("num_splits", [1, 4]) @maybe_fake_tensor_mode(USE_FAKE_TENSOR) From 8cd24f774500ba27da77891af881b7cc7d564dcb Mon Sep 17 00:00:00 2001 From: Simon Veitner Date: Mon, 26 Oct 2026 15:35:28 +0100 Subject: [PATCH 12/13] [Cute,Fwd,Sm90] Document the |q| <= 448 range of the FP8 MMA path Q is quantized as an E4M3 pair at fixed scale 1.0 with no range check. Measured on H200: rel error 0.2% up to |q| = 448, 2% at 768-896, 15% at 1024, 65% at 2048. Say so where callers look. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Simon Veitner Signed-off-by: Kate Cheng --- flash_attn/cute/interface.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/flash_attn/cute/interface.py b/flash_attn/cute/interface.py index 373c11801cb..2997d4f065e 100644 --- a/flash_attn/cute/interface.py +++ b/flash_attn/cute/interface.py @@ -712,6 +712,10 @@ def _flash_attn_fwd( dequant = out_fp8 * output_scale. SM100/SM110 only. compile_only: If True, compile the selected kernel and return without launching it. + fp8_kv_dequant: SM90 only. fp16/bf16 Q with paged E4M3 K/V, head_dim 512, + page_size 64. Q is split into two E4M3 terms at a fixed scale of 1, so + the caller must keep |q| <= 448: above 448 the error grows about 10x, + and above 896 Q saturates silently. Fold larger Q into k_descale. """ aux_scalars = tuple(aux_scalars) if aux_scalars else None requires_grad = any( From 2648f1b7b0f17d229d1fe7f9981a1d0bae714c6e Mon Sep 17 00:00:00 2001 From: Simon Veitner Date: Mon, 26 Oct 2026 15:35:28 +0100 Subject: [PATCH 13/13] [Cute,Fwd,Sm90] Keep the sink LSE for fully masked rows without SplitKV The empty-split override now runs on every SM90 launch, so a non-split row with no visible key and a learnable sink returns LSE = -inf instead of the sink logit. Restrict the LSE override to SplitKV and FP8, as the merge base did for the non-FP8 path. O is unchanged. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Simon Veitner Signed-off-by: Kate Cheng --- flash_attn/cute/flash_fwd_sm90.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/flash_attn/cute/flash_fwd_sm90.py b/flash_attn/cute/flash_fwd_sm90.py index ff6cc1d4bdf..b5a659f9c0e 100644 --- a/flash_attn/cute/flash_fwd_sm90.py +++ b/flash_attn/cute/flash_fwd_sm90.py @@ -1891,7 +1891,10 @@ def mma( # This also defines an all-empty sequence as a zero output. if is_empty_split: acc_O.fill(Float32(0.0)) - softmax_tile.row_sum.fill(-Float32.inf) + # A non-split launch keeps the learnable sink's LSE, as before; + # FP8 consumed a dummy K/V block, so its LSE is not meaningful. + if const_expr(self.is_split_kv or self.fp8_kv_dequant): + softmax_tile.row_sum.fill(-Float32.inf) # /////////////////////////////////////////////////////////////////////////////// # Epilogue