support sm100 fwd hdim256

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

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
 

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,