diff --git a/scripts/benchmarks/README.md b/scripts/benchmarks/README.md index dfd60ddac6..0d6ac62e4f 100644 --- a/scripts/benchmarks/README.md +++ b/scripts/benchmarks/README.md @@ -65,11 +65,47 @@ the CB to FA integration that are unrelated to which FA version you use: as aliases for `max_length_q` / `max_length_k`. Without this rename CB's model kwargs never bind and FA is called with `max_seqlen_q=None`. +## Install + +Prereqs: `torch >= 2.5` for `torch.nn.attention.flex_attention`. The +Triton backend that flex_attention uses by default runs on Ampere, +Hopper, and Blackwell -- no separate install. + +FA4 (CuTeDSL) targets Hopper and Blackwell only. `qwen3_flex_inference.py` +auto-enables FA4 on supported GPUs and falls back to the Triton +`flex_attention` backend elsewhere. `--fa4_prefill` forces on (warns + +falls back if the GPU does not support it); `--no-fa4_prefill` forces off. + +CUDA 13 (recommended, used on B200 / RTX 50xx): + + pip install --index-url https://download.pytorch.org/whl/cu130 torch + pip install "flash-attn-4[cu13]" + +CUDA 12 (H100 boxes still on cu12): + + pip install torch # default index is cu12 + pip install flash-attn-4 + +Pin `flash-attn-4==4.0.0b9` to match this benchmark. The `[cu13]` +extra pulls in `nvidia-cutlass-dsl` built for CUDA 13. + +| GPU | arch | sm | Auto FA4 | Triton flex_attention | +|--------------|-----------|-------|----------|------------------------| +| A100 | Ampere | sm_80 | off (uses Triton) | Works | +| H100 / H200 | Hopper | sm_90 | on | Works | +| RTX 50xx | Blackwell | sm_120 | on | Works | +| B200 / GB200 | Blackwell | sm_100 | on | Works | + +The transformers continuous-batching path's FA4 wiring (the +`flash_attn_fa4_shim.py` monkey-patches and the +`site-packages/flash_attn/__init__.py` namespace shim that makes FA4 +visible under the FA2 import name) is covered below under "Known +integration notes". + ## Reproduce ```bash pip install unsloth "transformers>=4.57" "trl>=0.25" peft vllm -uv pip install --no-deps flash-attn-4==4.0.0b9 # Generation microbenchmark (32 prompts, 512 new tokens each) CUDA_VISIBLE_DEVICES=2 python scripts/benchmarks/cb_vs_vllm_generation.py \ diff --git a/scripts/benchmarks/qwen3_flex_inference.py b/scripts/benchmarks/qwen3_flex_inference.py index d94f695d4b..e9ca603350 100644 --- a/scripts/benchmarks/qwen3_flex_inference.py +++ b/scripts/benchmarks/qwen3_flex_inference.py @@ -457,7 +457,7 @@ class FlexInference: max_new_tokens = 512, decode_kernel_options = None, prefill_kernel_options = None, - fa4_prefill = False, + fa4_prefill = None, base_model = None, peft_model = None, ): @@ -477,6 +477,24 @@ class FlexInference: self.max_seq_length = max_seq_length self.page_size = page_size self.max_new_tokens = max_new_tokens + # FA4 CuTeDSL kernels ship for Hopper (sm_90) and Blackwell (sm_100, + # sm_120) only. `fa4_prefill=None` means auto-detect: enable where + # supported, silently fall back to the Triton flex_attention backend + # elsewhere. Explicit `fa4_prefill=True` on sub-Hopper still falls + # back, but warns -- the user asked for a kernel that isn't there. + if fa4_prefill is None or fa4_prefill: + major, _ = torch.cuda.get_device_capability(self.device) + supported = major >= 9 + if fa4_prefill and not supported: + import warnings + warnings.warn( + f"--fa4_prefill needs Hopper (sm_90) or Blackwell " + f"(sm_100 / sm_120); found sm_{major}0. Falling back to " + f"the Triton flex_attention backend.", + RuntimeWarning, + stacklevel = 2, + ) + fa4_prefill = supported self.fa4_prefill = fa4_prefill # On SM100 (Blackwell), FA4 via flex_attention requires Q block = 256, # KV block = 128. See attention-gym `get_flash_block_size`. @@ -866,10 +884,12 @@ def main(): ) p.add_argument( "--fa4_prefill", - action = "store_true", + default = None, + action = argparse.BooleanOptionalAction, help = ( "Use BLOCK_SIZE=(256,128) + BACKEND=FLASH on prefill to unlock the " - "CuTeDSL FA4 kernel on Blackwell (SM100)." + "CuTeDSL FA4 kernel. Default auto-enables on Hopper (sm_90) and " + "Blackwell (sm_100, sm_120); use --no-fa4_prefill to force off." ), ) p.add_argument( diff --git a/tests/test_fa4_capability_guard.py b/tests/test_fa4_capability_guard.py new file mode 100644 index 0000000000..c6ac8da7a9 --- /dev/null +++ b/tests/test_fa4_capability_guard.py @@ -0,0 +1,158 @@ +"""Unit test for the FA4 capability guard in FlexInference.__init__. + +Runs on any GPU (and on CPU) because we monkey-patch +`torch.cuda.get_device_capability` and stub out the page-table / model +patching that the constructor does after the guard. +""" +import os +import sys +import types +import warnings +import unittest +from unittest import mock + +import torch + + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +BENCH_DIR = os.path.join(REPO_ROOT, "scripts", "benchmarks") +if BENCH_DIR not in sys.path: + sys.path.insert(0, BENCH_DIR) + +# qwen3_flex_inference imports heavy siblings (flex_paged_attention). +# Stub the PageTable and patch_qwen3_model the constructor calls after the +# guard so we don't need a real model / CUDA device. +import qwen3_flex_inference as qfi # noqa: E402 + + +class _FakePageTable: + def __init__(self, *a, **kw): + pass + + def create_causal_blockmask(self, *a, **kw): + return None + + +class _FakeTokenizer: + eos_token_id = 0 + + +def _make_fake_model(device_str="cpu"): + m = types.SimpleNamespace() + m.device = torch.device(device_str) + return m + + +def _build(fa4_prefill, cc_major, cc_minor=0): + """Construct a FlexInference with the guard exercised. + + Returns the instance. Patches torch.cuda.get_device_capability, + torch.zeros (to avoid CUDA allocation), PageTable, and + patch_qwen3_model so __init__ can run to completion without a real + model. + """ + fake_model = _make_fake_model("cpu") + fake_tok = _FakeTokenizer() + + _real_zeros = torch.zeros + + def _fake_zeros(*a, **kw): + kw.pop("device", None) + return _real_zeros(*a, **kw) + + with mock.patch.object( + torch.cuda, "get_device_capability", return_value=(cc_major, cc_minor) + ), mock.patch.object(qfi, "PageTable", _FakePageTable), mock.patch.object( + qfi, "patch_qwen3_model", lambda *a, **kw: None + ), mock.patch.object(torch, "zeros", _fake_zeros): + return qfi.FlexInference( + model=fake_model, + tokenizer=fake_tok, + max_batch_size=2, + max_seq_length=128, + n_pages=4, + page_size=128, + max_new_tokens=16, + fa4_prefill=fa4_prefill, + ) + + +def _fa4_warnings(caught): + return [ + w for w in caught + if issubclass(w.category, RuntimeWarning) + and "fa4_prefill" in str(w.message) + ] + + +class TestFA4CapabilityGuard(unittest.TestCase): + # --- explicit opt-in: --fa4_prefill=True --- + def test_explicit_on_sub_hopper_disables_and_warns(self): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + fi = _build(fa4_prefill=True, cc_major=8) + self.assertTrue( + _fa4_warnings(caught), + f"expected RuntimeWarning about fa4_prefill, got {caught!r}", + ) + self.assertIs(fi.fa4_prefill, False) + self.assertEqual(fi.prefill_q_block, 128) + self.assertNotIn("BACKEND", fi.prefill_kernel_options) + + def _assert_fa4_enabled(self, cc_major, fa4_prefill): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + fi = _build(fa4_prefill=fa4_prefill, cc_major=cc_major) + self.assertEqual( + _fa4_warnings(caught), [], + f"unexpected fa4 RuntimeWarning on sm_{cc_major}0 " + f"with fa4_prefill={fa4_prefill}: {caught!r}", + ) + self.assertIs(fi.fa4_prefill, True) + self.assertEqual(fi.prefill_q_block, 256) + self.assertEqual(fi.prefill_kernel_options.get("BACKEND"), "FLASH") + + def test_explicit_on_hopper_enables(self): + self._assert_fa4_enabled(cc_major=9, fa4_prefill=True) + + def test_explicit_on_blackwell_sm100_enables(self): + self._assert_fa4_enabled(cc_major=10, fa4_prefill=True) + + def test_explicit_on_blackwell_sm120_enables(self): + self._assert_fa4_enabled(cc_major=12, fa4_prefill=True) + + # --- auto-detect: fa4_prefill is None --- + def test_auto_on_sub_hopper_disables_silently(self): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + fi = _build(fa4_prefill=None, cc_major=8) + self.assertEqual( + _fa4_warnings(caught), [], + f"auto-detect must not warn on unsupported GPU: {caught!r}", + ) + self.assertIs(fi.fa4_prefill, False) + self.assertEqual(fi.prefill_q_block, 128) + self.assertNotIn("BACKEND", fi.prefill_kernel_options) + + def test_auto_on_hopper_enables(self): + self._assert_fa4_enabled(cc_major=9, fa4_prefill=None) + + def test_auto_on_blackwell_sm100_enables(self): + self._assert_fa4_enabled(cc_major=10, fa4_prefill=None) + + def test_auto_on_blackwell_sm120_enables(self): + self._assert_fa4_enabled(cc_major=12, fa4_prefill=None) + + # --- explicit opt-out: --no-fa4_prefill / fa4_prefill=False --- + def test_explicit_off_on_blackwell_stays_off(self): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + fi = _build(fa4_prefill=False, cc_major=10) + self.assertEqual(_fa4_warnings(caught), []) + self.assertIs(fi.fa4_prefill, False) + self.assertEqual(fi.prefill_q_block, 128) + self.assertNotIn("BACKEND", fi.prefill_kernel_options) + + +if __name__ == "__main__": + unittest.main()