{"record":{"id":"d3441a439936b87f","repo":"xai-org/x-algorithm","slug":"attn-impl-flash-attn-selects-the-fa3-training-ke","errorCode":null,"errorMessage":"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').","messagePattern":"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'\\)\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"phoenix/xrex/models/attention.py","lineNumber":178,"sourceCode":"                segment_ids,\n                causal=self.config.causal,\n                sm_scale=self.scale_config.attn_output_scale(self.config.key_size),\n                cap=self.config.attn_logit_cap,\n                backward_pass_impl=\"triton_split\",\n                interpret=False,\n            ), None\n\n        return sharded_mha, ()\n\n\ndef import_attention_fa3():\n    import importlib.util\n\n    if importlib.util.find_spec(\"xrex.models.attention_fa3\") is not None:\n        from xrex.models import attention_fa3\n\n        return attention_fa3.FLASH_ATTN_IMPL\n    raise NotImplementedError(\n        \"attn_impl='flash_attn' selects the fa3 training-kernel arm \"\n        \"(xrex.models.attention_fa3, backed by the compiled xrex/cuda/fa3 \"\n        \"kernels), which is not available in this tree. Choose another \"\n        \"attn_impl (e.g. 'jax_attn' or 'pallas_ranker_attn').\"\n    )\n\n\ndef make_segment_ids_monotonic(segment_ids):\n    segment_ids -= jnp.expand_dims(segment_ids[:, 0], axis=-1)\n    segment_ids += jnp.expand_dims(jnp.pad(jnp.cumsum(segment_ids[:-1, -1] + 1), (1, 0)), axis=-1)\n    return segment_ids\n\n\ndef identity_segment_ids(segment_ids):\n    return segment_ids\n\n\ndef qk_tie_segment_ids(segment_ids, segment_ids_k):","sourceCodeStart":160,"sourceCodeEnd":196,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/models/attention.py#L160-L196","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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."],"exampleFix":"# before\nconfig.attn_impl = \"flash_attn\"\n\n# after\nconfig.attn_impl = \"pallas_ranker_attn\"","handlingStrategy":"fallback","validationCode":"import importlib.util\nif importlib.util.find_spec(\"xrex.models.attention_fa3\") is None:\n    config.attn_impl = \"jax_attn\"  # or \"pallas_ranker_attn\"","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Probe optional kernels with importlib.util.find_spec before selecting attn_impl.","Document which attn_impl values are available in each tree/environment."],"tags":["attention","fa3","optional-dependency","cuda","feature-unavailable"],"backgroundTag":"optional-feature-not-installed","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}