[model] support FlashAttention-4 (flash_attn: fa4)#10639
Conversation
Add FlashAttention-4 as a selectable attention backend, mirroring the existing fa2 dispatch path so it applies uniformly to all supported models. Resolves the pre-existing `# FIXME compatibility fa3/fa4` in data/collator.py. - AttentionFunction.FA4 + dispatch to transformers' "flash_attention_4" (guarded by is_flash_attn_4_available(); graceful message on old transformers). - Generalized vision-tower fallback: fa4 (v4.0.0.beta21) crashes in its backward-preprocess kernel for any attention head_dim not a multiple of 32. Root cause (flash_attn/cute/flash_bwd_preprocess.py): head_dim is padded to a multiple of 32 (line 68); when the real head_dim differs, the OOB copy-predicate (check_hdim_v_oob, line 71) is passed to cute.copy WITHOUT slicing it per row-tile (lines 367-369, 381-382), unlike the forward (flash_fwd.py:478) and main-backward (flash_bwd.py:1161) kernels which slice it -- producing a CuTe shape mismatch. Verified on B30Z: head_dim 64/128/256 pass, 72/80 fail, forward-only always works. Vision towers with such head_dim (Qwen3-VL / Qwen3.5 ViT: 1152/16=72) route the vision tower to fa2 and the LLM to fa4 via the transformers dict-form _attn_implementation; text / LLM-only models run fully on fa4. Detection is generic (_fa4_vision_needs_fa2, %32), not a model-type allow-list. - Varlen speed fix: transformers passes max_seqlen_q/k as a CUDA tensor on packed paths, but fa4's cute API is typed Optional[int] and builds its backward JIT compile-key from it -- a tensor hashes by identity and recompiles every step (~12x slowdown). Coerce to int at the fa4 entry point (idempotent; no-op if fa4 absent). - collator FLASH_ATTN_IMPLS, Qwen3.5 packing-patch gate extended to fa4, two example configs, GPU-free unit tests. Validated with seed-fixed fa2-vs-fa4 A/B runs on B30Z (Blackwell/B300-class, sm_103a), flash-attn-4 4.0.0b21, transformers 5.6.0. Numerics aligned (Pearson 0.974, mean |dloss|<0.01, no drift) on Qwen3-VL-8B (ViT->fa2 routing) and Qwen3-VL-4B (whole model fa4) over 1000 steps. Speed on Qwen3-4B text with packing: fa4 +12.9% at 8192 tokens/step and +18.0% at 16384 tokens/step (attention speedup grows with sequence length).
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for FlashAttention-4 (fa4), including a fallback mechanism to route vision towers with unsupported head dimensions to FlashAttention-2 and a patch to prevent per-step backward recompiles by coercing sequence length tensors to integers. The review feedback suggests two robust improvements: adding a safety check to prevent a potential division-by-zero error when calculating the attention head dimension, and using inspect.signature to dynamically bind arguments in the FlashAttention-4 patch instead of relying on hardcoded positional indices.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| hidden_size = getattr(sub_config, "hidden_size", None) | ||
| num_heads = getattr(sub_config, "num_attention_heads", None) or getattr(sub_config, "num_heads", None) | ||
| if hidden_size and num_heads: | ||
| return int(hidden_size) // int(num_heads) |
There was a problem hiding this comment.
If num_heads is configured as '0' (as a string) or if it is somehow parsed as 0, bool(num_heads) might evaluate to True (for '0'), leading to a ZeroDivisionError when calculating hidden_size // num_heads. Adding a robust check to ensure int(num_heads) > 0 prevents potential runtime crashes.
| hidden_size = getattr(sub_config, "hidden_size", None) | |
| num_heads = getattr(sub_config, "num_attention_heads", None) or getattr(sub_config, "num_heads", None) | |
| if hidden_size and num_heads: | |
| return int(hidden_size) // int(num_heads) | |
| hidden_size = getattr(sub_config, "hidden_size", None) | |
| num_heads = getattr(sub_config, "num_attention_heads", None) or getattr(sub_config, "num_heads", None) | |
| if hidden_size and num_heads and int(num_heads) > 0: | |
| return int(hidden_size) // int(num_heads) |
| # Public signature: flash_attn_varlen_func(q, k, v, qv=None, cu_seqlens_q=None, | ||
| # cu_seqlens_k=None, max_seqlen_q=None, max_seqlen_k=None, ...) -> positional indices 6, 7. | ||
| max_seqlen_q_pos, max_seqlen_k_pos = 6, 7 | ||
|
|
||
| @functools.wraps(orig) | ||
| def flash_attn_varlen_func_int_seqlen(*args, **kwargs): | ||
| for name in ("max_seqlen_q", "max_seqlen_k"): | ||
| value = kwargs.get(name) | ||
| if isinstance(value, torch.Tensor): | ||
| kwargs[name] = int(value.item()) | ||
|
|
||
| if args: | ||
| args = list(args) | ||
| for pos in (max_seqlen_q_pos, max_seqlen_k_pos): | ||
| if len(args) > pos and isinstance(args[pos], torch.Tensor): | ||
| args[pos] = int(args[pos].item()) | ||
|
|
||
| args = tuple(args) | ||
|
|
||
| return orig(*args, **kwargs) |
There was a problem hiding this comment.
Hardcoding the positional indices 6 and 7 for max_seqlen_q and max_seqlen_k is fragile and highly sensitive to any future signature changes in flash-attn-4 (especially since it is currently in beta). Using inspect.signature to dynamically bind the arguments is a much more robust and maintainable approach.
| # Public signature: flash_attn_varlen_func(q, k, v, qv=None, cu_seqlens_q=None, | |
| # cu_seqlens_k=None, max_seqlen_q=None, max_seqlen_k=None, ...) -> positional indices 6, 7. | |
| max_seqlen_q_pos, max_seqlen_k_pos = 6, 7 | |
| @functools.wraps(orig) | |
| def flash_attn_varlen_func_int_seqlen(*args, **kwargs): | |
| for name in ("max_seqlen_q", "max_seqlen_k"): | |
| value = kwargs.get(name) | |
| if isinstance(value, torch.Tensor): | |
| kwargs[name] = int(value.item()) | |
| if args: | |
| args = list(args) | |
| for pos in (max_seqlen_q_pos, max_seqlen_k_pos): | |
| if len(args) > pos and isinstance(args[pos], torch.Tensor): | |
| args[pos] = int(args[pos].item()) | |
| args = tuple(args) | |
| return orig(*args, **kwargs) | |
| # Public signature: flash_attn_varlen_func(q, k, v, qv=None, cu_seqlens_q=None, | |
| # cu_seqlens_k=None, max_seqlen_q=None, max_seqlen_k=None, ...) | |
| import inspect | |
| sig = inspect.signature(orig) | |
| @functools.wraps(orig) | |
| def flash_attn_varlen_func_int_seqlen(*args, **kwargs): | |
| try: | |
| bound_args = sig.bind(*args, **kwargs) | |
| for name in ("max_seqlen_q", "max_seqlen_k"): | |
| if name in bound_args.arguments: | |
| value = bound_args.arguments[name] | |
| if isinstance(value, torch.Tensor): | |
| bound_args.arguments[name] = int(value.item()) | |
| return orig(*bound_args.args, **bound_args.kwargs) | |
| except TypeError: | |
| return orig(*args, **kwargs) |
1 participant
What does this PR do?
Adds FlashAttention-4 as a selectable attention backend via
flash_attn: fa4, mirroring theexisting
fa2dispatch path so it applies uniformly to every supported model. It also resolves thepre-existing
# FIXME compatibility fa3/fa4indata/collator.py.What's included
AttentionFunction.FA4enum value + dispatch to transformers'"flash_attention_4"implementation, guarded by
is_flash_attn_4_available()(graceful warning + fallback when fa4 or anew-enough transformers is unavailable — never a hard crash).
backward-preprocess kernel for any attention
head_dimthat is not a multiple of 32 (the kernel'shead-dim padding granularity). Empirically verified on B30Z: d=72/80 fail at the first backward; d=64/128/256
pass; the forward is fine for all — the bug is backward-only. Root cause traced in the flash-attention source
(see "Root-cause analysis" below): the preprocess kernel builds an out-of-bounds head-dim copy predicate but
passes it to
cute.copywithout slicing it per row-tile, unlike every other (forward and main-backward)kernel, producing a CuTe shape-mismatch (
'cute.copy' op expects pred to have compatible shape with (1,(3)) but got (1,(2,3))). Models whose vision tower has such ahead_dim(e.g. Qwen3-VL, Qwen3.5: ViThead_dim=72) route only the vision tower to fa2 and the language modelto fa4 via the transformers dict-form
_attn_implementation; text / LLM-only models run entirely on fa4.Detection is generic (
_fa4_vision_needs_fa2), not a model-type allow-list, so it covers current and futureVLMs.
max_seqlen_q/kasa CUDA tensor, but fa4's cute API is typed
Optional[int]and builds its backward JIT compile-keyfrom it. A tensor hashes by object identity, so the backward kernel recompiles every step (~12×
slowdown vs fa2). We coerce it to
intat the fa4 entry point (_patch_fa4_varlen_int_seqlen,idempotent, no-op when fa4 is absent).
FLASH_ATTN_IMPLSrefactor (fa2/fa4 share the varlen unpadded-packing contract), Qwen3.5GPU packing-patch gate extended to fa4, two example configs, and GPU-free unit tests.
Experiments (seed-fixed fa2-vs-fa4 A/B)
Hardware: NVIDIA B30Z (Blackwell / B300-class,
sm_103a); flash-attn-44.0.0b21; transformers 5.6.0;bf16;
seed=42; batch size 1; identical data ordering per pair (0 token-mismatches, verified).Speed (Qwen3-4B, text,
neat_packing, uniform sequences, 1000 steps each). fa4 acceleratesattention, whose cost scales O(L²), so its benefit grows with sequence length:
At short sequence lengths the gap is small (attention is a minor fraction of the step): a separate
Qwen3.5-9B multimodal SFT run at ~580 tokens/step measured only ~+3% steady-state — expected, and the
reason the benchmark above targets long packed sequences where fa4 is designed to help.
Numerics (Qwen3-VL, 1000 steps, real multimodal SFT). fa4 matches fa2 with no drift:
The 8B run confirms the generalized ViT→fa2 routing fires (and the model trains without the fa4-backward
crash); the 4B run is a negative control — same
qwen3_vlarchitecture,head_dim=64, so the wholemodel runs on fa4 with no fallback line. Per-step loss differences are bf16 attention-kernel-reorder
noise (concentrated on near-zero-loss steps); the aggregate trajectories are essentially identical.
Notes / limitations
flash-attn-4wheel. For VLMs whose vision tower has ahead_dimnot divisible by 32, fa4 currently accelerates only the language model (the ViT stays onfa2 — a flash-attn-4 backward-kernel bug, analyzed below, not a LlamaFactory one).
src/llamafactory/v1/) has a separateAttentionFunctionand is intentionally leftout of scope; fa4 there can be a follow-up.
Before submitting
tests/model/model_utils/test_attention.py).Files changed (8)
src/llamafactory/extras/constants.pyFA4 = "fa4"enum membersrc/llamafactory/hparams/model_args.pyflash_attnhelp text lists fa4src/llamafactory/model/model_utils/attention.py_fa4_vision_needs_fa2generalized routing,_patch_fa4_varlen_int_seqlen, print armsrc/llamafactory/data/collator.pyFLASH_ATTN_IMPLSconstant; fa2/fa4 share packing path; removed stale FIXMEsrc/llamafactory/model/patcher.pytests/model/model_utils/test_attention.pyexamples/train_full/qwen3vl_full_sft_fa4.yaml,qwen3_full_sft_fa4.yamlVerification commands (reviewer can reproduce)
Appendix — Root-cause analysis of the fa4 vision-tower crash (source-backed)
All line numbers are from the official flash-attention repo at tag
fa4-v4.0.0.beta21(
flash_attn/cute/), the wheel used in these experiments.Symptom. With
flash_attn: fa4, a Qwen3-VL / Qwen3.5 run aborts at the first backward with a CuTeDSL error:
Where the head_dim gate is set.
FlashAttentionBackwardPreprocess.__init__pads head_dim to amultiple of 32 and flags the out-of-bounds (OOB) case:
So
check_hdim_v_oobis True exactly whenhead_dim % 32 != 0— e.g. 72 → padded 96 (OOB), 80 → 96(OOB); 64/128/256 are already multiples of 32 (not OOB).
can_implementonly checkshead_dim % 8 == 0(line 97), so 72/80 are accepted at build time and only fail at runtime.
The bug. In the OOB path the kernel builds a copy predicate once and binds it to a
partialwithoutslicing it per row-tile (
m), then applies that partial to per-mtensor slices:copy_utils.copyforwardspredstraight tocute.copyunchanged(
quack/copy_utils.py:224-236). The predicate frompredicate_kcarries an extravmode(shape
(2,3)in the k-modes), but the single-msource slice expects a predicate of shape(3)→the
(1,(3))vs(1,(2,3))mismatch.Why the forward and the main backward are unaffected. Every other kernel slices the predicate by
mbefore the copy, so the shapes match:
The forward also pads to a multiple of 16 (
flash_fwd.py:82), a different granularity. Hence: forwardworks for all head_dims; the main backward works; only
flash_bwd_preprocess.pyhas the unsliced-predicatebug, and only on the
check_hdim_v_oob(head_dim % 32 != 0) path. This matches the empirical resultexactly (64/128/256 pass, 72/80 fail, forward-only always compiles).
How the PR handles it.
_fa4_vision_needs_fa2()routes a vision tower to fa2 whenever its head_dim isnot a multiple of
_FA4_HEAD_DIM_MULTIPLE = 32(the true kernel trigger), leaving the LLM on fa4. This isminimal and precise: it does not needlessly downgrade vision towers whose head_dim is a multiple of 32
(e.g. 96), and it catches the real offenders (Qwen3-VL / Qwen3.5 ViT head_dim 72). A fix could instead be
sent upstream to flash-attention (slice the predicate per-
minflash_bwd_preprocess.pyas the otherkernels do); until a fixed wheel ships, routing the vision tower to fa2 is the safe path.