8dfa076ee4
New scripts under scripts/benchmarks/:
- flash_attn_fa4_shim.py: monkey-patches that let transformers CB dispatch to
Flash Attention 4 on Blackwell (sm_100). CB's ContinuousBatchProcessor
otherwise emits a 4D paged attention mask for flash_attention_2 (which then
breaks _flash_attention_forward's _upad_input branch), and passes
max_seqlen_q instead of max_length_q. The shim skips the mask for FA and
accepts both names.
- persistent_cb.py: replaces model.generate_batch with a version that reuses
one ContinuousBatchingManager across calls, avoiding the per-step
PagedAttentionCache realloc. Wired up behind --persistent_cb on the tpaged
and standalone scripts.
- qwen3_grpo_naive.py: vanilla HF model.generate + TRL GRPOTrainer, no vLLM
and no CB. Mirrors the TRL docs example. Useful as a third column in the
comparison and also as a "will this at least converge" sanity check.
Adds --attn_impl and --persistent_cb flags to the existing generation and
training scripts. No changes to Unsloth internals.
Updated README.md with the FA install recipe (flash-attn-4==4.0.0b9, plus
a small site-packages shim that re-exports FA4's cute.* symbols under the
FA2 flash_attn namespace so transformers' is_flash_attn_2_available() and
_lazy_imports("flash_attention_2") succeed on B200).
Benchmark numbers on a single B200, Qwen3-4B-Base LoRA rank 32, bf16:
Generation microbenchmark (32 prompts, 512 new tokens):
vLLM 7224 decode tok/s (100%)
CB paged|sdpa 527 decode tok/s ( 7.3%)
CB paged|flash_attention_2 (FA4) 709 decode tok/s ( 9.8%)
CB paged|flash_attention_2 persistent 529 decode tok/s ( 7.3%)
GRPO training (max_steps=20, num_generations=2, per_device_batch=2):
vLLM colocated 136.6 s peak 157 GB
naive TRL (HF generate) 910.0 s peak 15 GB
CB SDPA 1521.5 s peak 98 GB (prior run)
CB FA4 1470.7 s peak 82 GB
CB FA4 + persistent 1562.1 s peak 87 GB
CB FA4 + ng=4 persistent 1597.0 s peak 94 GB
FA4 is a real ~1.4x improvement over SDPA for CB decode throughput but the
50% of vLLM target is still not reached. The remaining gap is driven by
CUDA graph capture (which ContinuousBatchingManager still NotImplementedErrors
on) and vLLM's scheduler being more efficient for decode-heavy GRPO rollouts.
Naive TRL generate is the honest small-rig baseline: 6.7x slower than vLLM
at 10% of the VRAM footprint, and ~1.7x faster than CB here.
96 lines
3.5 KiB
Python
96 lines
3.5 KiB
Python
"""Make transformers' continuous batching dispatch to Flash Attention 4 on B200.
|
|
|
|
Two integration gaps between transformers' continuous batching (CB) and the
|
|
FA2 varlen path make `attn_implementation="flash_attention_2"` fail out of the
|
|
box in transformers 4.57 even with a working FA varlen kernel:
|
|
|
|
1. CB's `ContinuousBatchProcessor` creates a 4D paged attention mask of shape
|
|
`[1, 1, q_len, k_len]` for every attention implementation except
|
|
`"paged_attention"`. `_flash_attention_forward` then enters the
|
|
`if attention_mask is not None:` branch and calls `_upad_input`, which
|
|
expects a 2D mask and fails (the FA4 kernel ultimately asserts on
|
|
`cu_seqlens_k.shape`).
|
|
|
|
2. CB passes `max_seqlen_q`/`max_seqlen_k` as model kwargs while
|
|
`_flash_attention_forward` names the parameters `max_length_q`/
|
|
`max_length_k`. The former therefore never bind and the varlen branch
|
|
inside `_flash_attention_forward` invokes the FA kernel with
|
|
`max_seqlen_q=None`.
|
|
|
|
This module patches around both, and (via the sibling
|
|
`site-packages/flash_attn/__init__.py` shim) points the FA2 varlen dispatch
|
|
at FA4's Blackwell-capable kernel. Call `apply()` once, before
|
|
`model.generate_batch` is invoked.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import functools
|
|
|
|
|
|
_APPLIED = False
|
|
|
|
|
|
def apply() -> None:
|
|
global _APPLIED
|
|
if _APPLIED:
|
|
return
|
|
|
|
import transformers # noqa: F401 - force load
|
|
from transformers.generation.continuous_batching import continuous_api as _cb
|
|
from transformers import modeling_flash_attention_utils as _fa_utils
|
|
|
|
_patch_return_attention_mask(_cb)
|
|
_patch_flash_attention_forward(_fa_utils)
|
|
|
|
_APPLIED = True
|
|
|
|
|
|
def _patch_return_attention_mask(cb_module) -> None:
|
|
"""Don't materialise a 4D attention mask when the kernel is FA varlen.
|
|
|
|
The existing `return_attention_mask` only skips the mask for
|
|
`"paged_attention"`. We extend the skip set to `"flash_attention_2"`
|
|
(and `"flash_attention_3"` for future-proofing) because the varlen path
|
|
relies on `cu_seq_lens_*` and reads no mask.
|
|
"""
|
|
_SKIP_MASK_IMPLS = {
|
|
"paged_attention",
|
|
"flash_attention_2",
|
|
"flash_attention_3",
|
|
}
|
|
|
|
def return_attention_mask(self) -> bool:
|
|
return self.config._attn_implementation not in _SKIP_MASK_IMPLS
|
|
|
|
cb_module.ContinuousBatchProcessor.return_attention_mask = return_attention_mask
|
|
|
|
|
|
def _patch_flash_attention_forward(fa_utils_module) -> None:
|
|
"""Accept CB's `max_seqlen_q`/`max_seqlen_k` kwargs as aliases.
|
|
|
|
transformers names the parameters `max_length_q`/`max_length_k`, but CB
|
|
(and most downstream call sites) name them `max_seqlen_q`/
|
|
`max_seqlen_k`. We rename at the boundary so callers on either side work
|
|
unchanged.
|
|
"""
|
|
original = fa_utils_module._flash_attention_forward
|
|
|
|
@functools.wraps(original)
|
|
def wrapper(*args, **kwargs):
|
|
if kwargs.get("max_length_q") is None and "max_seqlen_q" in kwargs:
|
|
kwargs["max_length_q"] = kwargs.pop("max_seqlen_q")
|
|
if kwargs.get("max_length_k") is None and "max_seqlen_k" in kwargs:
|
|
kwargs["max_length_k"] = kwargs.pop("max_seqlen_k")
|
|
return original(*args, **kwargs)
|
|
|
|
fa_utils_module._flash_attention_forward = wrapper
|
|
|
|
# Some integration modules imported the function by name before we
|
|
# patched. Re-bind the most common consumers so they pick up the wrapper.
|
|
try:
|
|
from transformers.integrations import flash_attention as _flash_integration
|
|
|
|
_flash_integration._flash_attention_forward = wrapper
|
|
except Exception:
|
|
pass
|