diff --git a/flash_attn/cute/blackwell_helpers.py b/flash_attn/cute/blackwell_helpers.py index 07d6054..95f14e3 100644 --- a/flash_attn/cute/blackwell_helpers.py +++ b/flash_attn/cute/blackwell_helpers.py @@ -10,6 +10,30 @@ import flash_attn.cute.mma_sm100_desc as sm100_desc +def tcgen05_mma_kind(op) -> str: + """PTX tcgen05.mma kind string for the given MMA op.""" + if isinstance(op, tcgen05.mma.MmaF16BF16Op): + return "f16" + if isinstance(op, tcgen05.mma.MmaTF32Op): + return "tf32" + if isinstance(op, tcgen05.mma.MmaI8Op): + return "i8" + fp8_ops = tuple( + getattr(tcgen05.mma, name) + for name in ("MmaFP8Op", "MmaF8F6F4Op") + if hasattr(tcgen05.mma, name) + ) + if fp8_ops and isinstance(op, fp8_ops): + return "f8f6f4" + if hasattr(tcgen05.mma, "MmaMXF8Op") and isinstance(op, tcgen05.mma.MmaMXF8Op): + return "mxf8f6f4" + if hasattr(tcgen05.mma, "MmaMXF4Op") and isinstance(op, tcgen05.mma.MmaMXF4Op): + return "mxf4" + if hasattr(tcgen05.mma, "MmaMXF4NVF4Op") and isinstance(op, tcgen05.mma.MmaMXF4NVF4Op): + return "mxf4nvf4" + raise TypeError(f"Unsupported tcgen05 MMA op kind: {type(op).__name__}") + + @cute.jit def gemm_w_idx( tiled_mma: cute.TiledMma, @@ -394,6 +418,7 @@ def gemm_ptx_partial( sA_layout = sA.layout if sA is not None else tCrA.layout sB_layout = sB.layout idesc: int = const_expr(sm100_desc.mma_op_to_idesc(op)) + kind = tcgen05_mma_kind(op) if const_expr(not is_ts): sA_swizzle = sA.iterator.type.swizzle_type smem_desc_base_a: int = const_expr( @@ -477,7 +502,7 @@ def gemm_ptx_partial( f"mov.b64 smem_desc_a, {{smem_desc_a_lo_start, smem_desc_a_hi}};\n\t" f"mov.b64 smem_desc_b, {{smem_desc_b_lo_start, smem_desc_b_hi}};\n\t" "setp.ne.b32 p, $2, 0;\n\t" - f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::f16 [tmem_acc], smem_desc_a, smem_desc_b, idesc, {pred_str};\n\t" + f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], smem_desc_a, smem_desc_b, idesc, {pred_str};\n\t" + "".join( ( # f"add.u32 smem_desc_a_lo, smem_desc_a_lo, {hex(offset_a_diff[k - 1])};\n\t" @@ -486,7 +511,7 @@ def gemm_ptx_partial( f"add.u32 smem_desc_b_lo, smem_desc_b_lo_start, {hex(offset_b[k])};\n\t" f"mov.b64 smem_desc_a, {{smem_desc_a_lo, smem_desc_a_hi}};\n\t" f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t" - f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::f16 [tmem_acc], smem_desc_a, smem_desc_b, idesc, 1;\n\t" + f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], smem_desc_a, smem_desc_b, idesc, 1;\n\t" ) for k in range(1, cute.size(tCrA.shape[2])) ) @@ -554,7 +579,7 @@ def gemm_ptx_partial( f"mov.b32 smem_desc_b_hi, {hex(smem_desc_b_hi)};\n\t" f"mov.b64 smem_desc_b, {{smem_desc_b_lo_start, smem_desc_b_hi}};\n\t" "setp.ne.b32 p, $2, 0;\n\t" - f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::f16 [tmem_acc], [tmem_a], smem_desc_b, idesc, {pred_str};\n\t" + f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], [tmem_a], smem_desc_b, idesc, {pred_str};\n\t" + "".join( ( # f"add.u32 tmem_a, tmem_a, {hex(offset_a_diff[k - 1])};\n\t" @@ -562,7 +587,7 @@ def gemm_ptx_partial( f"add.u32 smem_desc_b_lo, smem_desc_b_lo_start, {hex(offset_b[k])};\n\t" f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t" # f"@leader_thread tcgen05.mma.cta_group::1.kind::f16 [tmem_acc], [tmem_a], smem_desc_b, idesc, 1;\n\t" - f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::f16 [tmem_acc], [tmem_a + {hex(offset_a[k])}], smem_desc_b, idesc, 1;\n\t" + f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], [tmem_a + {hex(offset_a[k])}], smem_desc_b, idesc, 1;\n\t" ) for k in range( 1, @@ -575,7 +600,7 @@ def gemm_ptx_partial( ( f"add.u32 smem_desc_b_lo, smem_desc_b_lo, {hex(offset_b_diff[k - 1])};\n\t" f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t" - f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::f16 [tmem_acc], [tmem_a + {hex(offset_a[k])}], smem_desc_b, idesc, 1;\n\t" + f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], [tmem_a + {hex(offset_a[k])}], smem_desc_b, idesc, 1;\n\t" ) for k in range(split_arrive_idx, cute.size(tCrA.shape[2])) ) @@ -783,6 +808,7 @@ def gemm_ptx_precomputed( mbar_phase: Optional[Int32] = None, zero_init: bool | Boolean = False, cta_group: int = 1, + kind: str = "f16", ) -> None: # acc_tmem_addr += acc_offset is_ts = const_expr(smem_desc_base_a is None) @@ -842,7 +868,7 @@ def gemm_ptx_precomputed( f"mov.b64 smem_desc_a, {{smem_desc_a_lo_start, smem_desc_a_hi}};\n\t" f"mov.b64 smem_desc_b, {{smem_desc_b_lo_start, smem_desc_b_hi}};\n\t" "setp.ne.b32 p, $2, 0;\n\t" - f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::f16 [tmem_acc], smem_desc_a, smem_desc_b, idesc, {pred_str};\n\t" + f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], smem_desc_a, smem_desc_b, idesc, {pred_str};\n\t" + "".join( ( # f"add.u32 smem_desc_a_lo, smem_desc_a_lo, {hex(offset_a_diff[k - 1])};\n\t" @@ -851,7 +877,7 @@ def gemm_ptx_precomputed( f"add.s32 smem_desc_b_lo, smem_desc_b_lo_start, {hex(offset_b[k])};\n\t" f"mov.b64 smem_desc_a, {{smem_desc_a_lo, smem_desc_a_hi}};\n\t" f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t" - f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::f16 [tmem_acc], smem_desc_a, smem_desc_b, idesc, 1;\n\t" + f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], smem_desc_a, smem_desc_b, idesc, 1;\n\t" ) for k in range(1, num_k_tile) ) @@ -911,7 +937,7 @@ def gemm_ptx_precomputed( f"mov.b32 smem_desc_b_hi, {hex(smem_desc_b_hi)};\n\t" f"mov.b64 smem_desc_b, {{smem_desc_b_lo_start, smem_desc_b_hi}};\n\t" "setp.ne.b32 p, $2, 0;\n\t" - f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::f16 [tmem_acc], [tmem_a], smem_desc_b, idesc, {pred_str};\n\t" + f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], [tmem_a], smem_desc_b, idesc, {pred_str};\n\t" + "".join( ( # f"add.u32 tmem_a, tmem_a, {hex(offset_a_diff[k - 1])};\n\t" @@ -919,7 +945,7 @@ def gemm_ptx_precomputed( f"add.u32 smem_desc_b_lo, smem_desc_b_lo_start, {hex(offset_b[k])};\n\t" f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t" # f"@leader_thread tcgen05.mma.cta_group::1.kind::f16 [tmem_acc], [tmem_a], smem_desc_b, idesc, 1;\n\t" - f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::f16 [tmem_acc], [tmem_a + {hex(offset_a[k])}], smem_desc_b, idesc, 1;\n\t" + f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], [tmem_a + {hex(offset_a[k])}], smem_desc_b, idesc, 1;\n\t" ) for k in range( 1, @@ -933,7 +959,7 @@ def gemm_ptx_precomputed( # f"add.u32 smem_desc_b_lo, smem_desc_b_lo, {hex(offset_b_diff[k - 1])};\n\t" f"add.u32 smem_desc_b_lo, smem_desc_b_lo_start, {hex(offset_b[k])};\n\t" f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t" - f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::f16 [tmem_acc], [tmem_a + {hex(offset_a[k])}], smem_desc_b, idesc, 1;\n\t" + f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], [tmem_a + {hex(offset_a[k])}], smem_desc_b, idesc, 1;\n\t" ) for k in range(num_k_tile // 4 * 3, num_k_tile) ) @@ -1019,6 +1045,7 @@ def gemm_ptx_precomputed_varname( smem_offset: int, zero_init: bool | Boolean = False, cta_group: int = 1, + kind: str = "f16", ) -> None: is_ts = False num_k_tile = cute.size(tCrB_layout.shape[2]) @@ -1066,8 +1093,8 @@ def gemm_ptx_precomputed_varname( for k in range(1, num_k_tile) ) + "setp.ne.b32 p, $1, 0;\n\t" - # f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::f16 [tmem_acc], {smem_var_name_prefix}_0, smem_desc_b, idesc, {pred_str};\n\t" - f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::f16 [tmem_acc], {smem_var_name_prefix}_0, smem_desc_b_0, {idesc_var_name}, {pred_str};\n\t" + # f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], {smem_var_name_prefix}_0, smem_desc_b, idesc, {pred_str};\n\t" + f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], {smem_var_name_prefix}_0, smem_desc_b_0, {idesc_var_name}, {pred_str};\n\t" + "".join( ( # f"mov.b64 {{smem_desc_a_lo, smem_desc_a_hi}}, {smem_var_name_prefix}_{k};\n\t" @@ -1075,9 +1102,9 @@ def gemm_ptx_precomputed_varname( # f"add.s32 smem_desc_b_lo, smem_desc_b_lo_start, {hex(offset_b[k])};\n\t" # f"mov.b64 {smem_var_name_prefix}_{k}, {{smem_desc_a_lo, smem_desc_a_hi}};\n\t" # f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t" - # f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::f16 [tmem_acc], {smem_var_name_prefix}_{k}, smem_desc_b, idesc, 1;\n\t" - # f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::f16 [tmem_acc], {smem_var_name_prefix}_{k}, smem_desc_b, {idesc_var_name}, 1;\n\t" - f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::f16 [tmem_acc], {smem_var_name_prefix}_{k}, smem_desc_b_{k}, {idesc_var_name}, 1;\n\t" + # f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], {smem_var_name_prefix}_{k}, smem_desc_b, idesc, 1;\n\t" + # f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], {smem_var_name_prefix}_{k}, smem_desc_b, {idesc_var_name}, 1;\n\t" + f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], {smem_var_name_prefix}_{k}, smem_desc_b_{k}, {idesc_var_name}, 1;\n\t" ) for k in range(1, num_k_tile) ) diff --git a/flash_attn/cute/flash_fwd_sm100.py b/flash_attn/cute/flash_fwd_sm100.py index 4f931a5..995116e 100644 --- a/flash_attn/cute/flash_fwd_sm100.py +++ b/flash_attn/cute/flash_fwd_sm100.py @@ -112,6 +112,15 @@ class SfS2TCopies: (True, False, 256, False): {"ex2_emu_freq": 14, "ex2_emu_res": 6, "ex2_emu_start_frg": 0, "num_regs_softmax": 256, "num_regs_correction": 160}, (True, True, 256, False): {"ex2_emu_freq": 14, "ex2_emu_res": 6, "ex2_emu_start_frg": 0, "num_regs_softmax": 256, "num_regs_correction": 160}, } +_FP8_TUNING_CONFIG = { + # Currently equal to the base entry for this key; kept as an + # explicit pin so the fp8-validated values survive future + # retuning of the bf16 table. + (False, True, 128, False): { + "ex2_emu_freq": 16, + "ex2_emu_start_frg": 1, + }, +} # === END TUNING KNOBS === @@ -481,7 +490,7 @@ def _setup_attributes(self): kv_stage = 3 v_mma_stage = kv_stage # TODO: revisit hard-coded stage counts for mxfp8 - if self.v_dtype.width == 8: + if self.v_dequant: kv_stage = 2 v_mma_stage = ( 1 if self.has_bias and self.q_stage == 2 else kv_stage @@ -555,6 +564,14 @@ def __call__( """ # setup static attributes before smem/grid/tma computation self.q_dtype = mQ.element_type + # Plain-fp8 path only: qk-blockscaled Q is also e4m3 but keeps a + # bf16 P store and stock softmax scaling, so the fp8 P pre-scale + # machinery must not key on dtype width alone. + self.fp8_plain = ( + self.q_dtype.width == 8 + and not self.qk_blockscaled + and not self.v_dequant + ) self.k_dtype = mK.element_type self.v_dtype = mV.element_type self.o_dtype = mO.element_type @@ -648,8 +665,13 @@ def __call__( raise TypeError(f"Type mismatch: {self.q_dtype} != {self.k_dtype}") if const_expr(not self.qk_blockscaled and not self.v_dequant and self.q_dtype != self.v_dtype): raise TypeError(f"Type mismatch: {self.q_dtype} != {self.v_dtype}") - if const_expr(mBias is not None and self.bias_dtype != self.v_mma_dtype): - raise TypeError(f"Type mismatch: {self.v_mma_dtype} != {self.bias_dtype}") + if const_expr( + mBias is not None + and self.bias_dtype not in (cutlass.BFloat16, cutlass.Float16) + ): + raise TypeError( + f"Attention bias must be BF16 or FP16, got {self.bias_dtype}" + ) if const_expr( self.qk_blockscaled and (self.sfk_dtype is None or self.qk_sf_vec_size is None) ): @@ -657,6 +679,20 @@ def __call__( if const_expr(self.qk_blockscaled and self.sfq_dtype != self.sfk_dtype): raise TypeError(f"Type mismatch: {self.sfq_dtype} != {self.sfk_dtype}") + if const_expr(self.fp8_plain): + fp8_tune = _FP8_TUNING_CONFIG.get( + ( + self.use_2cta_instrs, + self.is_causal, + self.head_dim_padded, + self.is_sm103, + ), + {}, + ) + self._tune = {**self._tune, **fp8_tune} + if const_expr("ex2_emu_freq" in fp8_tune): + self.enable_ex2_emu = fp8_tune["ex2_emu_freq"] > 0 + self._setup_attributes() # Compute SF TMEM offsets for blockscaled @@ -2653,6 +2689,7 @@ def mma( qk_mma_op, pv_mma_op = tiled_mma_qk.op, tiled_mma_pv.op qk_mma_idesc, pv_mma_idesc = sm100_desc.mma_op_to_idesc(qk_mma_op), sm100_desc.mma_op_to_idesc(pv_mma_op) + qk_mma_kind = sm100_utils.tcgen05_mma_kind(qk_mma_op) q_smem_base = sm100_desc.smem_desc_base_from_tensor(sQ, sm100_desc.Major.K) k_smem_base = sm100_desc.smem_desc_base_from_tensor(sK, sm100_desc.Major.K) v_smem_base = sm100_desc.smem_desc_base_from_tensor(sV, sm100_desc.Major.MN) @@ -2693,6 +2730,7 @@ def mma( smem_offset=-sQ_stage_stride if stage == 0 else sQ_stage_stride, zero_init=True, cta_group=self.cta_group_size, + kind=qk_mma_kind, ) for stage in range(self.q_stage) ] @@ -3130,7 +3168,12 @@ def softmax_loop( ) tStScale_r2t = thr_tmem_store_scale.partition_D(tStScale) tmem_store_atom = cute.make_copy_atom( - tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(16)), Float32 + tcgen05.copy.St32x32bOp( + tcgen05.copy.Repetition( + 8 if const_expr(self.fp8_plain) else 16 + ) + ), + Float32, ) thr_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP).get_slice(tidx) tStP_r2t = thr_tmem_store.partition_D(tStP) # (((16,32),1),1,4) @@ -3226,9 +3269,20 @@ def softmax_loop( softmax = SoftmaxSm100.create( softmax_scale_log2, - rescale_threshold=8.0 if const_expr(self.q_dtype.width == 16) else 0.0, + # fp8 P values are pre-scaled by 2**max_offset, so + # lazy-rescale slack is bounded by the e4m3 finite + # range: threshold <= log2(448) - max_offset = 0.807. + # Past that bound, a row-max jump that stays under the + # threshold skips the rescale and clips P; the relative + # bias makes such jumps routine. + rescale_threshold=( + 0.75 + if const_expr(self.fp8_plain) + else (8.0 if const_expr(self.q_dtype.width == 16) else 0.0) + ), softmax_scale=softmax_scale, store_row_max=self.store_row_max, + max_offset=8 if const_expr(self.fp8_plain) else 0, ) softmax.reset() @@ -3568,8 +3622,8 @@ def softmax_step( bias_si_phase = bias_si_consumer_state.phase bias_si_stage = bias_si_consumer_state.index + stage pipeline_bias.consumer_wait_w_index_phase(bias_si_stage, bias_si_phase) - tBrS = cute.make_tensor(tSrS_t2r.iterator, cute.make_fragment_like(tS2RsBias[None, None, None, 0].layout)) if const_expr(self.bias_block_size == 128) or tidx < self.bias_block_size: + tBrS = cute.make_tensor(tSrS_t2r.iterator, cute.make_fragment_like(tS2RsBias[None, None, None, 0].layout)) for i in cutlass.range_constexpr(cute.size(tS2RsBias.shape[2])): tBrS_cur = tBrS[None, 0, i] tS2RsBias_cur = tS2RsBias[None, 0, i, bias_si_stage] @@ -3737,6 +3791,17 @@ def correction_loop( cutlass.pipeline.PipelineUserType.Producer, 1 ) + max_offset = ( + Float32(8.0) + if const_expr(self.fp8_plain) + else Float32(0.0) + ) + max_offset_scale = ( + Float32(256.0) + if const_expr(self.fp8_plain) + else Float32(1.0) + ) + work_tile = tile_scheduler.initial_work_tile_info() while work_tile.is_valid_tile: m_block, head_idx, batch_idx, split_idx = work_tile.tile_idx @@ -3870,10 +3935,13 @@ def correction_loop( if row_max == -Float32.inf: # It's possible to have an empty row with splitKV. row_max = sink_val * (LOG2_E / softmax_scale_log2) - row_sum = Float32(1.0) + row_sum = max_offset_scale else: row_sum += cute.math.exp2( - sink_val * LOG2_E - row_max * softmax_scale_log2, fastmath=True + sink_val * LOG2_E + - row_max * softmax_scale_log2 + + max_offset, + fastmath=True, ) acc_O_mn_row_is_zero_or_nan = row_sum == 0.0 or row_sum != row_sum stats[stage] = (row_sum, row_max, acc_O_mn_row_is_zero_or_nan) @@ -3978,7 +4046,12 @@ def correction_loop( # cute.printf("row_sum = {}, row_max = {}, acc_O_mn_row_is_zero_or_nan = {}\n", row_sum, row_max, acc_O_mn_row_is_zero_or_nan) LN2 = math.log(2.0) lse = ( - (row_max * softmax_scale_log2 + cute.math.log2(row_sum, fastmath=True)) * LN2 + ( + row_max * softmax_scale_log2 + + cute.math.log2(row_sum, fastmath=True) + - max_offset + ) + * LN2 if not acc_O_mn_row_is_zero_or_nan else -Float32.inf ) diff --git a/flash_attn/cute/interface.py b/flash_attn/cute/interface.py index d8d84e4..786377b 100644 --- a/flash_attn/cute/interface.py +++ b/flash_attn/cute/interface.py @@ -385,7 +385,11 @@ def _flash_attn_fwd( assert q.dtype in [torch.float16, torch.bfloat16], "Q/K must be float16 or bfloat16" assert q.dtype == k.dtype, "Q and K must have the same dtype" else: - assert q.dtype in [torch.float16, torch.bfloat16], "inputs must be float16 or bfloat16" + assert q.dtype in [ + torch.float16, + torch.bfloat16, + torch.float8_e4m3fn, + ], "inputs must be float16, bfloat16, or float8_e4m3fn" assert q.dtype == k.dtype == v.dtype, "inputs must have the same dtype" for t in [cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k]: if t is not None: @@ -415,6 +419,14 @@ def _flash_attn_fwd( ) ), "inputs must be on CUDA device" arch = _get_device_arch() if _arch is None else _arch + if q.dtype == torch.float8_e4m3fn and not blockscaled: + assert arch in [100, 103], ( + "unscaled FP8 attention is only supported on SM100/SM103" + ) + assert not (head_dim == 256 and head_dim_v == 256), ( + "unscaled FP8 attention does not support head_dim 256 " + "(the dedicated hd256 kernel has no fp8 path)" + ) assert arch // 10 in [8, 9, 10, 11, 12], "Unsupported compute capability. Supported: 8.x, 9.x, 10.x, 11.x, 12.x" assert num_head % num_head_kv == 0, "num_head must be divisible by num_head_kv" alignment = 16 // q.element_size() @@ -425,7 +437,12 @@ def _flash_attn_fwd( if softcap == 0.0: softcap = None - out_torch_dtype = torch.bfloat16 if v_blockscaled else (v.dtype if blockscaled else q.dtype) + plain_fp8 = q.dtype == torch.float8_e4m3fn and not blockscaled + out_torch_dtype = ( + torch.bfloat16 + if v_blockscaled or plain_fp8 + else (v.dtype if blockscaled else q.dtype) + ) device = q.device q_batch_seqlen_shape = (batch_size, seqlen_q) if cu_seqlens_q is None else (total_q,) lse_shape = (batch_size, num_head, seqlen_q) if cu_seqlens_q is None else (num_head, total_q) diff --git a/flash_attn/cute/softmax.py b/flash_attn/cute/softmax.py index 1f0306d..3e04d41 100644 --- a/flash_attn/cute/softmax.py +++ b/flash_attn/cute/softmax.py @@ -170,6 +170,7 @@ def rescale_O(self, acc_O: cute.Tensor, row_scale: cute.Tensor) -> None: class SoftmaxSm100(Softmax): rescale_threshold: cutlass.Constexpr[float] = 0.0 row_max_true: Optional[cute.Tensor] = None + max_offset: cutlass.Constexpr[int] = 0 @staticmethod def create( @@ -177,6 +178,7 @@ def create( rescale_threshold: cutlass.Constexpr[float] = 0.0, softmax_scale: Float32 | None = None, store_row_max: cutlass.Constexpr[bool] = False, + max_offset: cutlass.Constexpr[int] = 0, ): num_rows = 1 arch = 100 @@ -194,6 +196,7 @@ def create( softmax_scale, rescale_threshold=rescale_threshold, row_max_true=row_max_true, + max_offset=max_offset, ) @cute.jit @@ -274,11 +277,12 @@ def scale_subtract_rowmax( ): assert cute.size(acc_S_row.shape) % 2 == 0, "acc_S_row must have an even number of elements" row_max_scaled = row_max * self.scale_log2 + bias = Float32(self.max_offset) - row_max_scaled for i in cutlass.range(0, cute.size(acc_S_row.shape), 2, unroll_full=True): acc_S_row[i], acc_S_row[i + 1] = cute.arch.fma_packed_f32x2( (acc_S_row[i], acc_S_row[i + 1]), (self.scale_log2, self.scale_log2), - (-row_max_scaled, -row_max_scaled), + (bias, bias), ) @cute.jit