xai-org/x-algorithm · error · NotImplementedError

FA2 has been removed but {self.config.attn_impl=} requires i

Error message

FA2 has been removed but {self.config.attn_impl=} requires it

What it means

The attention layer's __call__ explicitly blocks attn_impl='recsys_flash_attn': FlashAttention-2 support has been removed from the tree, and this impl required it, so the config is rejected at forward time with a NotImplementedError naming the offending attn_impl.

Source

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

            wd_mask = self.scale_config.ln_weight_decay_mask
            query_heads = norm_fn(
                query_heads, "q_norm", qk_lr, create_scale=True, weight_decay_mask=wd_mask
            )
            key_heads = norm_fn(
                key_heads, "k_norm", qk_lr, create_scale=True, weight_decay_mask=wd_mask
            )

        if add_rope:
            query_heads, key_heads = self._apply_rope(
                query_heads,
                key_heads,
                positions=positions,
            )
            query_heads = checkpoint_name(query_heads, "query_heads_rope")
            key_heads = checkpoint_name(key_heads, "key_heads_rope")

        if self.config.attn_impl == "recsys_flash_attn":
            raise NotImplementedError(
                f"FA2 has been removed but {self.config.attn_impl=} requires it"
            )

        extra_attn_kwargs: dict[str, jax.Array | None] = {}

        temp = jnp.ones(segment_ids.shape, dtype=jnp.int32)

        attn_class, extra_attn_kwargs = self._get_attn_impl(
            extra_attn_kwargs,
            mask=mask,
            seqpack_layout=seqpack_layout,
        )
        tmp = attn_class(
            self.config,
            self.scale_config,
            sharding_context=self.sharding_context,
        )(
            query_heads,

View on GitHub (pinned to 24c60942c5)

Solutions

1. Update the model config to a supported attn_impl such as 'jax_attn', 'pallas_attn', or 'flash_attn' with fa_version=3.
2. Regenerate or edit the saved config rather than reusing the one embedded in the old checkpoint.
3. Add a config-migration check that rewrites legacy impl names on load.

Example fix

# before
attn_impl: recsys_flash_attn

# after
attn_impl: jax_attn

if config.attn_impl == "recsys_flash_attn":
    raise SystemExit("FA2-based impl removed; update config to 'jax_attn' or 'flash_attn' with fa_version=3")

When it happens

Trigger: Loading or running a model whose config still has attn_impl='recsys_flash_attn'; older checkpoints/configs created before the FA2 removal.

Common situations: Resuming from a legacy checkpoint or YAML config after upgrading to a version that dropped FA2; configs inherited from an older experiment.

Related errors

• FlashAttention2 has been removed. Please use fa_version='3'
• Please override this method for specific attention impl.
• attn_impl='flash_attn' selects the fa3 training-kernel arm (
• attn_logit_cap_method {method!r} is not supported by JaxAtte
• Mask dimensionality {mask.ndim} must match logits dimensiona

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