[misc] handle flash attention 3.0 with recent transformers version (>= 4.56.0)

[ArdalanM]

Summary

verl/utils/attention_utils.py falls back to transformers.modeling_flash_attention_utils when flash_attn isn't installed, which is how we support running with FlashAttention-3 (flash_attn_interface) instead of FA2. That fallback needs transformers>=4.56.0 (_pad_input/_unpad_input weren't importable before that). On older transformers it failed with a confusing bare ImportError instead of a clear message.

This PR catches that and raises a clear error telling the user to install flash_attn or upgrade transformers.

Test plan

Tested with FlashAttention-3 and transformers 5.12 via examples/cispo_trainer/run_qwen3_8b_fsdp.sh:

#!/usr/bin/env bash
# Smoke test for examples/cispo_trainer/run_qwen3_8b_fsdp.sh on the FlashAttention-3 fallback
# added to verl/utils/attention_utils.py.
#
# Scaled down for a 2-GPU box: tiny GSM8K/MATH slices, short sequences, a single training
# step, and FSDP param/optimizer offload to fit the Adam optimizer state.
#
# Usage:
#   ./test_cispo_qwen3_8b_fa3_smoketest.sh

# ---- run the real training recipe for a single step, forcing FlashAttention-3 ----
MODEL_PATH=Qwen/Qwen3.5-0.8B \
NNODES=1 \
NGPUS_PER_NODE=2 \
TRAIN_BATCH_SIZE=8 \
PPO_MINI_BATCH_SIZE=8 \
MAX_PROMPT_LENGTH=256 \
MAX_RESPONSE_LENGTH=256 \
PPO_MAX_TOKEN_LEN_PER_GPU=1024 \
ROLLOUT_TP=1 \
ROLLOUT_GPU_MEM_UTIL=0.15 \
ROLLOUT_N=2 \
TOTAL_EPOCHS=1 \
PROJECT_NAME=verl_cispo_smoketest \
EXPERIMENT_NAME=fa3_smoketest \
    bash ./run_qwen3_8b_fsdp.sh \
    trainer.logger='["console"]' \
    trainer.val_before_train=False \
    trainer.total_training_steps=1 \
    +actor_rollout_ref.model.override_config.attn_implementation=flash_attention_3 \
    actor_rollout_ref.actor.fsdp_config.param_offload=True \
    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \
    actor_rollout_ref.rollout.max_num_seqs=64 \
    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=1024 \
    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=1024

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.

---

ardalan.mehrani seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[gemini-code-assist]

Code Review

This pull request introduces a fallback mechanism in verl/utils/attention_utils.py to support environments where flash_attn is not installed (such as FlashAttention-3 only environments) by importing equivalent functions from transformers (>=4.56.0) and einops. The reviewer noted that repeatedly raising and catching ImportError exceptions on every call to _get_attention_functions() in the training hot path introduces significant performance overhead, and suggested caching the resolved imports to avoid this.

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.

[gemini-code-assist]

In FA3-only environments (where flash_attn is not installed), every call to pad_input, unpad_input, etc., will invoke _get_attention_functions(), which attempts to import flash_attn.bert_padding, raises an ImportError, catches it, and then imports from transformers. Raising and catching exceptions on every single step in the training hot path introduces significant performance overhead. Caching the imported functions after the first successful resolution avoids this overhead entirely.

    if _index_first_axis is not None:
        return _index_first_axis, _pad_input, _rearrange, _unpad_input

    if is_torch_npu_available(check_device=False):
        from verl.utils.npu_flash_attn_utils import index_first_axis, pad_input, rearrange, unpad_input
    else:
        try:
            from flash_attn.bert_padding import index_first_axis, pad_input, rearrange, unpad_input
        except ImportError as e:
            # FlashAttention-2 (flash_attn) is not installed, e.g. FA3-only environments.
            # transformers ships equivalent implementations with matching signatures/returns, but
            # _pad_input/_unpad_input are only importable at module level since transformers==4.56.0
            # (https://github.com/huggingface/transformers/pull/40002); _index_first_axis has been
            # available since transformers==4.53.0 (https://github.com/huggingface/transformers/pull/38972).
            # rearrange has no transformers equivalent - flash_attn.bert_padding.rearrange is itself just
            # a re-export of einops.rearrange, so we import it directly from einops (a transformers dep).
            from einops import rearrange

            try:
                from transformers.modeling_flash_attention_utils import (
                    _index_first_axis as index_first_axis,
                    _pad_input as pad_input,
                    _unpad_input as unpad_input,
                )
            except ImportError:
                raise ImportError(
                    "Neither flash_attn nor a compatible transformers (>=4.56.0) providing "
                    "_index_first_axis/_pad_input/_unpad_input was found. Install flash_attn "
                    "or upgrade transformers to >=4.56.0."
                ) from e

    _index_first_axis, _pad_input, _rearrange, _unpad_input = index_first_axis, pad_input, rearrange, unpad_input

[wuxibin89]

@ArdalanM Please fix pre-commit fail: https://github.com/verl-project/verl/blob/main/CONTRIBUTING.md#code-linting-and-formatting