xai-org/x-algorithm · error · NotImplementedError

attn_impl='flash_attn' selects the fa3 training-kernel arm (

Error message

attn_impl='flash_attn' selects the fa3 training-kernel arm (xrex.models.attention_fa3, backed by the compiled xrex/cuda/fa3 kernels), which is not available in this tree. Choose another attn_impl (e.g. 'jax_attn' or 'pallas_ranker_attn').

What it means

import_attention_fa3 looks for the module xrex.models.attention_fa3 via importlib.util.find_spec; when it is absent it raises NotImplementedError. This module backs the fa3 training kernels (compiled xrex/cuda/fa3), which are not shipped in this tree, so selecting attn_impl='flash_attn' cannot be honored here.

Source

Thrown at phoenix/xrex/models/attention.py:178

                segment_ids,
                causal=self.config.causal,
                sm_scale=self.scale_config.attn_output_scale(self.config.key_size),
                cap=self.config.attn_logit_cap,
                backward_pass_impl="triton_split",
                interpret=False,
            ), None

        return sharded_mha, ()


def import_attention_fa3():
    import importlib.util

    if importlib.util.find_spec("xrex.models.attention_fa3") is not None:
        from xrex.models import attention_fa3

        return attention_fa3.FLASH_ATTN_IMPL
    raise NotImplementedError(
        "attn_impl='flash_attn' selects the fa3 training-kernel arm "
        "(xrex.models.attention_fa3, backed by the compiled xrex/cuda/fa3 "
        "kernels), which is not available in this tree. Choose another "
        "attn_impl (e.g. 'jax_attn' or 'pallas_ranker_attn')."
    )


def make_segment_ids_monotonic(segment_ids):
    segment_ids -= jnp.expand_dims(segment_ids[:, 0], axis=-1)
    segment_ids += jnp.expand_dims(jnp.pad(jnp.cumsum(segment_ids[:-1, -1] + 1), (1, 0)), axis=-1)
    return segment_ids


def identity_segment_ids(segment_ids):
    return segment_ids


def qk_tie_segment_ids(segment_ids, segment_ids_k):

View on GitHub (pinned to 24c60942c5)

Solutions

Switch attn_impl to 'jax_attn' or 'pallas_ranker_attn' as the message suggests.
Build/install the fa3 CUDA kernels and ensure xrex.models.attention_fa3 is importable.
Gate configs that request flash_attn behind a capability check (find_spec) with a fallback impl.

Example fix

# before
config.attn_impl = "flash_attn"

# after
config.attn_impl = "pallas_ranker_attn"

Defensive patterns

Strategy: fallback

Validate before calling

import importlib.util
if importlib.util.find_spec("xrex.models.attention_fa3") is None:
    config.attn_impl = "jax_attn"  # or "pallas_ranker_attn"

Prevention

Probe optional kernels with importlib.util.find_spec before selecting attn_impl.
Document which attn_impl values are available in each tree/environment.

When it happens

Trigger: Setting config.attn_impl='flash_attn' (with fa_version 3) in a checkout that lacks xrex/models/attention_fa3.py or the compiled xrex/cuda/fa3 extension; _get_attn_impl then calls import_attention_fa3.

Common situations: Porting a config from an internal build that ships FA3 kernels to an open/lean tree; missing CUDA extension build step; using a machine without the compiled kernels.

Related errors

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