File size: 700 Bytes
b22e03e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 | # 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.
|