diff --git a/flash_attn/cute/flash_fwd_sm100.py b/flash_attn/cute/flash_fwd_sm100.py index 6c9c20d0b76..d31ed7097b0 100644 --- a/flash_attn/cute/flash_fwd_sm100.py +++ b/flash_attn/cute/flash_fwd_sm100.py @@ -2,13 +2,13 @@ # - BF16 & FP16 dtype # - noncausal & causal attention # - MHA, GQA, MQA -# - hdim 64, 96, 128, (192, 128). +# - hdim 64, 96, 128, (192, 128), 256. # - varlen # - sliding window # - split-kv # Unsupported features that will be added later: # - page size != 128 -# - more hdim (192, 256) +# - more hdim (192) # Based on the cutlass example and cute-dsl example: # https://github.com/NVIDIA/cutlass/tree/main/examples/77_blackwell_fmha # https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/blackwell/fmha.py @@ -86,6 +86,10 @@ (False, True, 128, True): {"ex2_emu_freq": 0, "ex2_emu_start_frg": 0, "num_regs_softmax": 176, "num_regs_correction": 64}, (True, False, 192, True): {"ex2_emu_freq": 0, "ex2_emu_start_frg": 0, "num_regs_softmax": 176, "num_regs_correction": 64}, (False, True, 192, True): {"ex2_emu_freq": 0, "ex2_emu_start_frg": 0, "num_regs_softmax": 176, "num_regs_correction": 72}, + (False, False, 256, False): {"ex2_emu_freq": 16, "ex2_emu_start_frg": 0, "num_regs_softmax": 176, "num_regs_correction": 88}, + (False, True, 256, False): {"ex2_emu_freq": 16, "ex2_emu_start_frg": 0, "num_regs_softmax": 176, "num_regs_correction": 88}, + (False, False, 256, True): {"ex2_emu_freq": 0, "ex2_emu_start_frg": 0, "num_regs_softmax": 176, "num_regs_correction": 80}, + (False, True, 256, True): {"ex2_emu_freq": 0, "ex2_emu_start_frg": 0, "num_regs_softmax": 176, "num_regs_correction": 80}, } # === END TUNING KNOBS === @@ -173,12 +177,13 @@ def __init__( is_sm103 = self.arch >= Arch.sm_103 and self.arch <= Arch.sm_103f self.is_sm103 = is_sm103 # enable_ex2_emu is derived: True if tuning config has freq > 0, else fallback to default logic - _default_enable_ex2_emu = (self.head_dim_padded <= 128 or (self.head_dim_padded == 192 and self.use_2cta_instrs and not self.is_causal and not self.is_local)) and not is_sm103 + _default_enable_ex2_emu = (self.head_dim_padded <= 128 or (self.head_dim_padded == 192 and self.use_2cta_instrs and not self.is_causal and not self.is_local) or self.head_dim_padded == 256) and not is_sm103 self.enable_ex2_emu = _default_enable_ex2_emu self.s0_s1_barrier = False self.overlap_sO_sQ = ( (self.head_dim_padded == 192 and self.head_dim_v_padded >= 64) or - (self.head_dim_v_padded >= 128 and self.is_split_kv) + (self.head_dim_v_padded >= 128 and self.is_split_kv) or + self.head_dim_padded >= 256 or self.head_dim_v_padded >= 256 ) if self.overlap_sO_sQ: self.is_persistent = False diff --git a/flash_attn/cute/interface.py b/flash_attn/cute/interface.py index ef624677f01..9ccbb4caae8 100644 --- a/flash_attn/cute/interface.py +++ b/flash_attn/cute/interface.py @@ -111,9 +111,10 @@ def _validate_head_dims(head_dim: int, head_dim_v: int, compute_capability: int, f"head_dim and head_dim_v must be between 8 and 256 and divisible by {alignment}." ) elif compute_capability in [10, 11]: - assert (is_standard_range or is_deepseek_shape) and head_dim % alignment == 0 and head_dim_v % alignment == 0, ( + is_sm100_range = 8 <= head_dim <= 256 and 8 <= head_dim_v <= 256 + assert (is_sm100_range or is_deepseek_shape) and head_dim % alignment == 0 and head_dim_v % alignment == 0, ( f"(head_dim, head_dim_v)=({head_dim}, {head_dim_v}) is not supported on SM100/SM110. " - f"head_dim and head_dim_v must be between 8 and 128 and divisible by {alignment}, or (192, 128) for DeepSeek." + f"head_dim and head_dim_v must be between 8 and 256 and divisible by {alignment}, or (192, 128) for DeepSeek." ) @@ -159,6 +160,13 @@ def _tile_size_fwd_sm90(head_dim, head_dim_v, is_causal, is_local, sparse_block_ tile_n = 64 if is_local else 80 return FwdConfig(128, tile_n, True, True) +def _tile_size_fwd_sm100(head_dim, head_dim_v, is_causal, is_local): + """Return FwdConfig for SM100 forward.""" + if head_dim <= 192: + return FwdConfig(128, 128, True, True) + else: # hdim 256 + return FwdConfig(128, 64, True, True) + @dataclass(frozen=True) class BwdConfig: m_block_size: int @@ -476,6 +484,8 @@ def _flash_attn_fwd( fwd_cfg = FwdConfig(128, 128, True, True) else: fwd_cfg = FwdConfig(128, 64, True, True) + elif arch // 10 in [10, 11]: + fwd_cfg = _tile_size_fwd_sm100(head_dim, head_dim_v, causal, local) elif arch // 10 == 8: fwd_cfg = FwdConfig(128, 64, True, True) # SM80, should tune elif arch // 10 == 9: @@ -499,7 +509,10 @@ def _flash_attn_fwd( max_seqlen_k = seqlen_k seqlen_q_packgqa = max_seqlen_q * qhead_per_kvhead if arch // 10 == 10: - q_stage = 2 if seqlen_q_packgqa > tile_m else 1 + if head_dim > 192 or head_dim_v > 192: # hdim 256 + q_stage = 1 + else: + q_stage = 2 if seqlen_q_packgqa > tile_m else 1 else: q_stage = 1 diff --git a/tests/cute/test_flash_attn.py b/tests/cute/test_flash_attn.py index 69e6308fb60..d4d6449041d 100644 --- a/tests/cute/test_flash_attn.py +++ b/tests/cute/test_flash_attn.py @@ -253,7 +253,7 @@ def test_flash_attn_output( # SplitKV not supported on SM90 - skip this iteration if IS_SM90 and num_splits > 1: continue - if IS_SM100 and (d >= 192 and dv >= 192): # hdim 192 and 256 not support on SM100 + if IS_SM100 and (d == 192 and dv == 192): # hdim 192 not support on SM100 continue out, lse = flash_attn_func( q,