Spaces:
Running on Zero
Running on Zero
| """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"] |