xai-org/x-algorithm · error · RuntimeError

FlashAttention2 has been removed. Please use fa_version='3'

Error message

FlashAttention2 has been removed. Please use fa_version='3' instead.

What it means

_get_attn_impl guards the 'flash_attn' arm: FlashAttention-2 has been removed, so requesting attn_impl='flash_attn' together with fa_version=2 raises RuntimeError and directs you to fa_version=3 (which routes to import_attention_fa3 and therefore requires the fa3 kernels to be present).

Source

Thrown at phoenix/xrex/models/layers.py:399

        return query_heads, key_heads

    @hk.transparent
    def _get_attn_impl(
        self,
        extra_attn_kwargs: dict[str, jax.Array | None],
        *,
        mask: jax.Array | None,
        seqpack_layout: SequencePackedLayout | None,
    ) -> tuple[Type[Attention], dict[str, jax.Array | None]]:
        match self.config.attn_impl:
            case "jax_attn" | "jax_attn_interleaved":
                extra_attn_kwargs["masks"] = mask
                attn_class = JaxAttention
            case "pallas_attn":
                attn_class = PallasAttention
            case "flash_attn":
                if int(self.config.fa_version) == 2:
                    raise RuntimeError(
                        "FlashAttention2 has been removed. Please use fa_version='3' instead."
                    )
                attn_class = import_attention_fa3()
            case "pallas_ranker_attn":
                attn_class = PallasRankerAttention
            case "pallas_ranker_attn_infer":
                attn_class = PallasRankerAttentionInference
            case "pallas_ranker_varlen_attn":
                assert seqpack_layout is not None
                extra_attn_kwargs["cu_seqlens"] = jnp.asarray(
                    seqpack_layout.cu_seqlens, dtype=jnp.int32
                )
                attn_class = PallasRankerVarlenAttention
            case "cutedsl_ranker_varlen_attn":
                assert seqpack_layout is not None, (
                    "cutedsl_ranker_varlen_attn requires seqpack_layout"
                )
                assert seqpack_layout.block_sparse is not None, (

View on GitHub (pinned to 24c60942c5)

Solutions

1. Set config.fa_version=3 alongside attn_impl='flash_attn'.
2. If FA3 kernels are unavailable in your tree, use 'jax_attn' or 'pallas_attn' instead.
3. Sanitize loaded configs to rewrite fa_version 2 to 3 (or switch impl) at load time.

Example fix

# before
attn_impl: flash_attn
fa_version: 2

# after
attn_impl: flash_attn
fa_version: 3

if config.attn_impl == "flash_attn" and int(config.fa_version) == 2:
    config.fa_version = 3  # migrate FA2 -> FA3

When it happens

Trigger: config.attn_impl='flash_attn' with config.fa_version=2 (or a config defaulting fa_version to 2); older configs that used FA2 flash attention.

Common situations: Legacy configs after the FA2 removal; forgetting to bump fa_version when migrating from FA2 to FA3.

Related errors

• FA2 has been removed but {self.config.attn_impl=} requires i
• attn_impl='flash_attn' selects the fa3 training-kernel arm (
• attn_logit_cap_method {method!r} is not supported by JaxAtte
• Invalid attention implementation: {self.config.attn_impl}
• Invalid argument {arg!r}, not a key=value replacement and no

AI-assisted analysis of xai-org/x-algorithm@24c60942c5 (2026-08-28). Data as JSON: /api/errors/bd7a82b95d6a08c5. Report an issue: GitHub.