Fix flash-attn2 hub-kernel attribute names in attention dispatch (#14012)

[Alex-Wengg]

What does this PR do?

Fixes #14012.

The hub-kernel registry (_HUB_KERNELS_REGISTRY in attention_dispatch.py) points the FLASH_HUB and FLASH_VARLEN_HUB backends (both kernels-community/flash-attn2) at attribute paths that don't exist in that kernel:

• flash_attn_interface._wrapped_flash_attn_forward / _wrapped_flash_attn_backward
• flash_attn_interface._wrapped_flash_attn_varlen_forward / _wrapped_flash_attn_varlen_backward

kernels-community/flash-attn2 exposes these without the _wrapped_ prefix (_flash_attn_forward, _flash_attn_varlen_forward, …), so _resolve_kernel_attr raises AttributeError: ... does not define attribute path 'flash_attn_interface._wrapped_flash_attn_varlen_forward' when the flash-attn2 hub backend is used.

This is corroborated by the FLASH_3_HUB config directly above, which already uses the non-_wrapped_ names (flash_attn_interface._flash_attn_forward / _flash_attn_backward).

Fix

Use the actual (non-_wrapped_) attribute names for the flash-attn2 hub configs. Four strings; no behavior change beyond resolving the correct symbols.

Verification

Introspecting the kernel (no flash-attn compile, no model download):

from kernels import get_kernel
fi = get_kernel("kernels-community/flash-attn2", version=1).flash_attn_interface

# names the registry currently looks for — all MISSING:
[hasattr(fi, n) for n in ("_wrapped_flash_attn_forward", "_wrapped_flash_attn_backward",
                          "_wrapped_flash_attn_varlen_forward", "_wrapped_flash_attn_varlen_backward")]
# -> [False, False, False, False]

# names this PR uses — all PRESENT:
[hasattr(fi, n) for n in ("_flash_attn_forward", "_flash_attn_backward",
                          "_flash_attn_varlen_forward", "_flash_attn_varlen_backward")]
# -> [True, True, True, True]

After the fix, _maybe_download_kernel_for_backend(...) resolves both forward and backward for FLASH_HUB and FLASH_VARLEN_HUB (no AttributeError). Verified on kernels==0.12.3.

AI assistance (Claude) was used to investigate and draft this change; the diff and verification were reviewed and run locally.

@sayakpaul

[sayakpaul]

Will be fixed in huggingface/kernels-community#976. Cc: @DN6 @drbh