Restore package README; record v2 verdict and structural downgrade

README.md

@@ -1,33 +1,32 @@
-# pi052 Training Optimization Workspace
+# flashrt-flex-attention-train
 
-This workspace records the pi052 training optimization plan, baseline tests,
-profiles, and environment notes. It is intentionally separate from the
-`lerobot/` source tree so optimization experiments do not mix with upstream
-code changes.
+FlexAttention replacement training package for PI-style dual-expert
+transformers.
 
-Current policy:
+Hub repo: `flashrt/flashrt-flex-attention-train`
 
-- Main test machine: RTX 5080 server at `internal.mindon.tech:20001`.
-- Thor server is parked for now; it is not the training test target.
-- Hugging Face Jobs are a final validation path only, and only A100 is planned.
-- HF Jobs budget target: keep early validation under a few dozen dollars.
-- Prefer China-friendly mirrors and ModelScope for downloads.
+## Public API
 
-Directory layout:
+- `flex_attention`
+- `flex_attention_forward`
+- `reference_flex_attention`
+- `build_block_sparse_bool_masks`
+- `backend_marker`
 
-- `docs/`: plans, environment records, and test protocols.
-- `scripts/`: helper scripts for remote setup, probing, and later test runs.
-- `results/`: JSON/CSV/text outputs from tests.
-- `profiles/`: torch profiler traces and summaries.
-- `env/`: requirements and environment lock notes.
+## Scope
 
-No pi052 model code should be edited from this workspace directly. Source
-changes should happen in `../lerobot/` after the baseline tests are in place.
+This package locks the public Tensor API and correctness harness for a native
+replacement of the PI052 FlexAttention/SDPA attention path:
 
-Primary planning documents:
+- 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
 
-- `docs/plan.md`: phase-level strategy.
-- `docs/proof_tests.md`: exact tests and proof thresholds before coding.
-- `docs/baseline_tests.md`: shorter baseline checklist.
-- `docs/status.md`: what has already been prepared.
-- `docs/proof_results_initial.md`: first synthetic proof-test results.
+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.

benchmarks/RESULTS.md

@@ -64,16 +64,27 @@ CUDA with no dropout, SDPA elsewhere.
   (128x128 for P>=896-class shapes).
 - A100/H100/H200: SDPA path (see the verdict table).
 
-## Next levers (not yet implemented)
+## Measured and closed follow-up levers (2026-07-09)
+
+1. **bf16-saved-p custom autograd (`manual_attention_part_v2`): REJECTED
+   end-to-end.** Against the EAGER composed part it wins 1.32-1.48x
+   fwd+bwd (fused softmax-gradient chain) — but the production form is
+   the composed part under `torch.compile`, where AOTAutograd's
+   partitioner already generates a fused backward and manages what gets
+   saved; the hand-written Function blocks that whole-graph treatment
+   and measured slower on the real model (text step 341.6 vs 333.3 ms
+   on a 5090). The symbol stays exported for API stability and as the
+   documented negative. Lesson (third time this project): compare
+   against the PRODUCTION form of the baseline, not its eager form.
+2. **Structural 3-GEMM split: DOWNGRADED.** The synthetic harness's
+   half/half att pattern overstates it — on the real model, flow-step
+   prefixes are fully bidirectional (no maskable quadrant) and text-step
+   causal spans vary per sample (no uniform split). The remaining
+   uniform win is the action cross-chunk block (~7% of attention
+   FLOPs); not scheduled.
 
-1. Structural 3-GEMM split: the prefix part computes a dense P x P even
-   though the (att=0 rows x att=1 cols) quadrant is fully masked
-   (~20% wasted FLOPs at p700) and the action part's cross-chunk block
-   is near-empty. Splitting Q into (group-0, group-1, action) rows cuts
-   both.
-2. Custom autograd saving bf16 probabilities only (halves saved-activation
-   bytes and backward traffic).
-3. Native CUDA fused kernel: manual sits at ~20% of bf16 peak on the
+## Next levers (not yet implemented)
+1. Native CUDA fused kernel: manual sits at ~20% of bf16 peak on the
    5090 (harness-inclusive); an FA2-style specialized kernel
    (D=256, GQA, prefix-dense + action-block) targeting 40-50% would be
    another ~2x. Entry per house rules: only after 1-2 are in and the