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.