| # flashrt-flex-attention-train | |
| FlexAttention replacement API for PI-style prefix/action training masks. | |
| ## Available functions | |
| - `flex_attention(q, k, v, prefix_len, action_block_size, ...)` | |
| - `flex_attention_forward(...)` | |
| - `reference_flex_attention(...)` | |
| - `build_block_sparse_bool_masks(...)` | |
| ## Acceptance status | |
| - Reference/SDPA autograd path: available. | |
| - CUDA optimized implementation: pending A100/5090 acceptance. | |
| - Precision mode: bf16/fp32 training reference, no FP8/FP4. | |
| - Fallback: unsupported shapes route to SDPA. | |
| Use this package to lock Lerobot/PI052 integration and run correctness and | |
| benchmark gates before replacing the internal reference path with optimized | |
| CUDA kernels. | |