File size: 1,070 Bytes
5e8586f 42a06d3 5e8586f 42a06d3 5e8586f 42a06d3 5e8586f b22e03e 5e8586f b22e03e 5e8586f b22e03e 5e8586f b22e03e 5e8586f b22e03e 5e8586f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 | # 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.
|