feat: enable FP8 large-head attention on SM80

flashinfer/aot.py

@@ -257,9 +257,9 @@ def gen_attention(
     head_dim_ckv = 512
     head_dim_kpe = 64
 
-    # For 16-bit KV, head_dim > 256 FA2 modules use the Ampere+ large-head path.
-    # NVFP4 KV large-head is validated for FA2 batch prefill on SM8+; other
-    # one-byte large-head modules stay SM100+-only until validated separately.
+    # For 16-bit and FP8 KV, head_dim > 256 FA2 modules use the Ampere+
+    # large-head path. NVFP4 KV large-head is validated for FA2 batch prefill
+    # on SM8+, while NVFP4 decode remains SM100+-only.
     from .jit.core import current_compilation_context
 
     has_sm8_or_newer = any(
@@ -268,7 +268,6 @@ def gen_attention(
     has_sm10_or_newer = any(
         major >= 10 for major, _ in current_compilation_context.TARGET_CUDA_ARCHS
     )
-
     # FA2 MHA / MQA / GQA
     for (
         (head_dim_qk, head_dim_vo),
@@ -285,12 +284,8 @@ def gen_attention(
     ):
         large_head = head_dim_qk > 256 or head_dim_vo > 256
         nvfp4_large_head = large_head and _is_nvfp4_kv_dtype(dtype_kv)
-        if large_head:
-            if dtype_kv.itemsize == 1 and not nvfp4_large_head:
-                if not has_sm10_or_newer:
-                    continue
-            elif not has_sm8_or_newer:
-                continue
+        if large_head and not has_sm8_or_newer:
+            continue
         yield from gen_fa2(
             dtype_qo=dtype_qo,
             dtype_kv=dtype_kv,
@@ -1123,7 +1118,7 @@ def parse_head_dim(head_dim: str) -> Tuple[int, int]:
 def get_default_config():
     """Get default AOT configuration"""
     return {
-        # Note: head_dim=512 (FA2 prefill/decode, SM100+) excluded to reduce
+        # Note: head_dim=512 (FA2 prefill/decode, SM80+) excluded to reduce
         # space in the jit-cache wheel.
         "fa2_head_dim": [(64, 64), (128, 128), (256, 256)],
         "fa3_head_dim": [(192, 128), (128, 128), (64, 64), (256, 256)],

flashinfer/jit/attention/modules.py

@@ -1258,14 +1258,13 @@ def _fa2_head_dim_nvcc_flags(
 ) -> Optional[List[str]]:
     """Return arch flags for FA2 large-head modules.
 
-    For 16-bit KV, head_dim > 256 uses the Ampere+ large-head path. NVFP4 KV
-    can opt into the same arch set only for validated FA2 prefill read paths.
-    Other one-byte large-head modules remain restricted to SM100+ until those
-    variants are validated separately.
+    For 16-bit and FP8 KV, head_dim > 256 uses the Ampere+ large-head path.
+    NVFP4 KV can opt into the same arch set only for validated FA2 prefill
+    read paths; NVFP4 large-head decode remains restricted to SM100+.
     """
     if head_dim_qk > 256 or head_dim_vo > 256:
-        if dtype_kv.itemsize == 1:
-            if not (allow_nvfp4_sm8_large_head and _is_nvfp4_kv_dtype(dtype_kv)):
+        if _is_nvfp4_kv_dtype(dtype_kv):
+            if not allow_nvfp4_sm8_large_head:
                 return current_compilation_context.get_nvcc_flags_list(
                     supported_major_versions=[10, 11, 12]
                 )

tests/attention/test_batch_decode_kernels.py

@@ -39,16 +39,6 @@ def skip_if_head_dim_unsupported(head_dim: int):
         pytest.skip("16-bit FA2 head_dim > 256 is only supported on SM80 or newer")
 
 
-def skip_if_head_dim_dtype_unsupported(head_dim: int, kv_dtype: torch.dtype):
-    skip_if_head_dim_unsupported(head_dim)
-    if (
-        head_dim > 256
-        and kv_dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
-        and get_compute_capability(torch.device("cuda:0"))[0] < 10
-    ):
-        pytest.skip("head_dim > 256 with FP8 KV is only validated on SM100 or newer")
-
-
 def skip_if_nvfp4_large_head_decode_unsupported(head_dim: int):
     if head_dim > 256 and get_compute_capability(torch.device("cuda:0"))[0] < 10:
         pytest.skip(
@@ -130,7 +120,7 @@ def _run_batch_decode_with_paged_kv_cache_case(
             pytest.skip("cuTile decode fp8 KV not covered yet.")
         if head_dim > 256:
             pytest.skip("cuTile decode head_dim>256 not covered yet.")
-    skip_if_head_dim_dtype_unsupported(head_dim, kv_dtype)
+    skip_if_head_dim_unsupported(head_dim)
     q = torch.randn(batch_size, num_qo_heads, head_dim, device="cuda:0", dtype=q_dtype)
     num_pages_per_seq = (kv_len + page_size - 1) // page_size
     total_num_pages = num_pages_per_seq * batch_size

tests/attention/test_fp8_prefill.py

@@ -22,13 +22,13 @@
 
 
 def head_dim_512_supported() -> bool:
-    # head_dim > 256 is only supported on SM100+.
-    return get_compute_capability(torch.device("cuda:0"))[0] >= 10
+    # FP8 FA2 head_dim > 256 uses the Ampere+ large-head path.
+    return get_compute_capability(torch.device("cuda:0"))[0] >= 8
 
 
 def skip_if_head_dim_unsupported(head_dim: int):
     if head_dim > 256 and not head_dim_512_supported():
-        pytest.skip("head_dim > 256 is only supported on SM100 or newer")
+        pytest.skip("FP8 FA2 head_dim > 256 is only supported on SM80 or newer")
 
 
 @pytest.mark.parametrize("batch_size", [12, 17])

tests/jit/test_jit_cpp_ext.py

@@ -272,6 +272,16 @@ def test_customize_batch_prefill_nvfp4_large_head_uses_prefill_flags(
         attention_modules._fa2_head_dim_nvcc_flags(512, 512, torch.uint8)
 
 
+def test_fa2_fp8_large_head_uses_sm80_flags(monkeypatch):
+    monkeypatch.setattr(
+        attention_modules.current_compilation_context, "TARGET_CUDA_ARCHS", {(8, 0)}
+    )
+
+    flags = attention_modules._fa2_head_dim_nvcc_flags(512, 512, torch.float8_e4m3fn)
+    assert flags is not None
+    assert any("sm_80" in flag for flag in flags)
+
+
 @pytest.mark.parametrize(
     ("head_dim_qk", "head_dim_vo", "supported"),
     [