# 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.