VeOmni Flash Attention Custom Name Adapter (Transformers 5.x)

VeOmni Flash Attention Custom Name Adapter (Transformers 5.x)#

Problem Background#

VeOmni uses custom attention implementation names:

  • veomni_flash_attention_2_with_sp

  • veomni_flash_attention_3_with_sp

  • veomni_flash_attention_4_with_sp

These names are registered into ALL_ATTENTION_FUNCTIONS and routed to VeOmni’s SP-aware attention wrapper.

With Transformers 5.x, model init and flash-attention preload logic may still call transformers.modeling_flash_attention_utils._lazy_imports(...) for the configured implementation string. For non-native names, _lazy_imports falls back to hub-kernel loading and can fail with:

ValueError: Could not find the currently requested flash attention implementation at veomni_flash_attention_2_with_sp

even though VeOmni already registered the custom attention function.

Why This Happens#

The failure path is:

  1. Model config keeps VeOmni custom name in _attn_implementation.

  2. Transformers flash preload code tries to resolve low-level flash kernels from the implementation string.

  3. Custom VeOmni names are not hub kernel identifiers.

  4. Hub fallback returns no valid kernel entry for this name.

  5. _lazy_imports raises before normal ALL_ATTENTION_FUNCTIONS dispatch takes effect.

Adapter Strategy Implemented#

Instead of patching _lazy_imports directly, VeOmni patches:

transformers.integrations.hub_kernels.load_and_register_attn_kernel

and intercepts VeOmni custom names only. This compatibility adapter is applied only when transformers>=5.0.0.

For VeOmni names, the adapter returns a local kernel-like object exposing:

  • flash_attn_func

  • flash_attn_varlen_func

mapped to local FA2/FA3/FA4 backends:

  • veomni_flash_attention_2_with_sp -> flash_attn.flash_attn_func / flash_attn.flash_attn_varlen_func

  • veomni_flash_attention_3_with_sp -> flash_attn_interface.flash_attn_func / flash_attn_interface.flash_attn_varlen_func

  • veomni_flash_attention_4_with_sp -> flash_attn.cute.flash_attn_func / flash_attn.cute.flash_attn_varlen_func

For simplicity, paged VeOmni aliases (for example paged|veomni_flash_attention_2_with_sp) are not handled by this adapter.

All non-VeOmni implementations are delegated to the original Transformers loader unchanged.

Design Goals#

  • Keep VeOmni custom implementation names unchanged.

  • Keep existing VeOmni ALL_ATTENTION_FUNCTIONS.register(...) behavior unchanged.

  • Avoid hub-kernel lookup for VeOmni private names.

  • Minimize patch surface by touching a single integration point.

  • Fail fast with clear ImportError when required FA backend is missing.

Expected Runtime Behavior#

After import veomni:

  • VeOmni custom names remain registered in ALL_ATTENTION_FUNCTIONS.

  • _lazy_imports("veomni_flash_attention_2_with_sp") and _lazy_imports("veomni_flash_attention_4_with_sp") can resolve through the adapter.

  • No spurious “kernel hub name not found” error for VeOmni custom names.

  • Paged VeOmni aliases are outside the adapter scope.

Notes#

  • This adapter is a compatibility bridge for Transformers 5.x behavior around flash preload.

  • It does not change VeOmni SP attention semantics.

  • It does not require the kernels Python package for VeOmni custom names.

  • FA2 and FA3 have dedicated branches in _lazy_imports and are resolved directly without reaching the hub-kernel path. The adapter is therefore a no-op for those two in practice, but is kept for safety.

  • FA4 (veomni_flash_attention_4_with_sp) has no such branch in _lazy_imports and always falls through to the hub-kernel path. The adapter is the critical component that makes FA4 usable.

  • FA4 requires the flash-attn-cute package (flash_attn.cute). To install FA4: uv sync --extra gpu --extra fa4.