| # flashrt-flex-attention-train | |
| FlexAttention replacement training package for PI-style dual-expert | |
| transformers. | |
| Hub repo: `flashrt/flashrt-flex-attention-train` | |
| ## Public API | |
| - `flex_attention` | |
| - `flex_attention_forward` | |
| - `reference_flex_attention` | |
| - `build_block_sparse_bool_masks` | |
| - `backend_marker` | |
| ## Scope | |
| This package locks the public Tensor API and correctness harness for a native | |
| replacement of the PI052 FlexAttention/SDPA attention path: | |
| - prefix self-attention rows | |
| - action-to-prefix rows plus block-diagonal action rows | |
| - `head_dim=256` | |
| - BF16 forward/backward through PyTorch autograd fallback | |
| - detached-prefix semantics for action rows reading prefix K/V | |
| - prefix mask, prefix padding mask, action block mask, and action padding mask | |
| - automatic SDPA fallback for unsupported shapes | |
| The current implementation is the SDPA-backed training reference. It is meant | |
| to be the stable integration target for native CUDA fwd/bwd kernels; no native | |
| performance claim is made until the benchmark gates in `VALIDATION.md` pass on | |
| both A100 and RTX 5090. | |