"""FlashAttention varlen shim for PPMem on ZeroGPU.

The PPMem attention code (ppmem/model/attention.py) imports
``flash_attn_varlen_func`` from the ``flash_attn_interface`` module — the
FlashAttention-3 package name. FlashAttention-3 does not ship prebuilt wheels
for the RTX PRO 6000 (sm_120) yet: FA3 targets Hopper's TMEM/tcgen05
instructions and is known not to work on Blackwell.

FlashAttention-2's ``flash_attn_varlen_func`` has the same parameter names
used at both call sites (``query, key, value, cu_seqlens_q=, cu_seqlens_k=,
max_seqlen_q=, max_seqlen_k=, causal=``) and — unlike FA3, which returns a
tuple — returns the plain output tensor both call sites expect. The cu_seqlens
tensors built by ``FlexFlashAttnVarlen.init_attn_params`` are int32, matching
FA2's requirement.

So re-export the FA2 function under the FA3 module name; ppmem's guarded
``magi_attention`` import keeps its fallback branch (this one) active in the
absence of the MagiAttention package.
"""

from flash_attn import flash_attn_varlen_func

__all__ = ["flash_attn_varlen_func"]