Add package source and torch-universal build variant

CARD.md

@@ -0,0 +1,21 @@
+# 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.

README.md

@@ -1,9 +1,32 @@
-# flashrt/flashrt-flex-attention-train
+# flashrt-flex-attention-train
 
-This repository is a compatibility mirror for older `kernels` clients
-that resolve repositories through the default Hugging Face model repo API.
+FlexAttention replacement training package for PI-style dual-expert
+transformers.
 
-Canonical Kernel Hub repo: https://huggingface.co/kernels/flashrt/flashrt-flex-attention-train
+Hub repo: `flashrt/flashrt-flex-attention-train`
 
-Do not edit this mirror by hand. It is generated from the Kernel Hub
-`vN` branches and contains the same `build/**` artifacts.
+## 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.

VALIDATION.md

@@ -0,0 +1,48 @@
+# Validation
+
+Correctness smoke:
+
+```bash
+python flashrt-flex-attention-train/tests/test_flashrt_flex_attention_train.py --backend source --mode smoke
+python flashrt-flex-attention-train/tests/test_flashrt_flex_attention_train.py --backend source --mode full
+```
+
+Installed artifact smoke:
+
+```bash
+python flashrt-flex-attention-train/tests/test_flashrt_flex_attention_train.py --backend installed --mode full
+python scripts/prebuild_check.py --package flashrt-flex-attention-train --check-config
+```
+
+Microbenchmark and acceptance gates:
+
+```bash
+python flashrt-flex-attention-train/benchmarks/benchmark.py --device cuda --dtype bf16 --mode all --output /tmp/flex_attention_a100.json
+python flashrt-flex-attention-train/benchmarks/benchmark.py --device cuda --dtype bf16 --mode all --require-gates
+```
+
+Shape/tile matrix:
+
+```bash
+python flashrt-flex-attention-train/benchmarks/shape_matrix.py --presets a100 --output /tmp/a100_flex_matrix.jsonl
+python flashrt-flex-attention-train/benchmarks/shape_matrix.py --presets consumer --output /tmp/consumer_flex_matrix.jsonl
+```
+
+Minimum gate for connecting native kernels:
+
+- forward time <= `0.95 * SDPA`
+- forward+backward time <= `0.95 * SDPA`
+- peak memory <= `1.03 * SDPA`
+- forward max abs diff <= `2e-3`
+- gradient norm relative diff <= `1e-2`
+
+Higher gate for publishing native kernels:
+
+- full PI052 step speedup >= `1.08x` on A100 and RTX 5090
+- isolated attention forward+backward speedup >= `1.25x`
+- no text/flow step regression
+- unsupported shapes automatically fall back to SDPA
+
+The current package provides the reference/fallback path and therefore is not
+expected to pass the performance gates until native CUDA fwd/bwd kernels are
+added behind the same API.

benchmarks/README.md

@@ -0,0 +1,22 @@
+# Benchmarks
+
+`benchmark.py` compares the package API against the current two-call SDPA
+baseline at PI052 FlexAttention replacement shapes. It records forward time, forward+backward time,
+peak memory, forward max absolute diff, and gradient norm relative diff.
+
+The default shape is the measured PI052 flow-only shape:
+
+- `B=4`
+- `heads=8`
+- `head_dim=256`
+- `prefix_len=700`
+- `action_blocks=5`
+- `action_block_size=50`
+
+For hardware/tile sweeps, run one process per shape and preset to avoid
+`torch.compile` guard reuse noise:
+
+```bash
+python flashrt-flex-attention-train/benchmarks/shape_matrix.py --presets a100 --output /tmp/a100_flex_matrix.jsonl
+python flashrt-flex-attention-train/benchmarks/shape_matrix.py --presets consumer --output /tmp/consumer_flex_matrix.jsonl
+```

benchmarks/RESULTS.md

@@ -0,0 +1,91 @@
+# Results
+
+## Headline (RTX 5090, real PI052 shapes, GQA kv_heads=1)
+
+At the REAL attention shapes (q 8 heads / kv 1 head, D=256, bf16) the
+materialized-logits `manual` backend (cuBLAS GEMMs + one compiled masked
+softmax, native autograd) beats every fused path on fwd+bwd:
+
+| shape | SDPA fwd/fwd+bwd | best flex fwd | best flex fwd+bwd | manual fwd/fwd+bwd | manual vs flex (fwd+bwd) |
+| --- | ---: | ---: | ---: | ---: | ---: |
+| b4_p700_k5_c50 | 1.581 / 4.006 | 0.262 | 2.597 | 0.360 / **1.285** | **2.0x** |
+| b2_p700_k5_c50 | 0.775 / 2.011 | 0.166 | 2.486 | 0.221 / **0.864** | **2.9x** |
+| b1_p700_k5_c50 | 0.453 / 1.292 | 0.139 | 2.180 | 0.127 / **0.977** | 2.2x |
+| b4_p512_k5_c50 | 0.967 / 2.601 | 0.179 | 1.934 | 0.245 / **1.057** | 1.8x |
+| b4_p896_k5_c50 | 2.332 / 5.667 | 0.323 | 2.634 | 0.515 / **1.854** | 1.4x |
+| b4_p700_k1_c50 | 1.009 / 2.660 | ERROR | ERROR | 0.279 / **0.975** | flex NaN/err at K=1 |
+| b4_p700_k8_c50 | 2.116 / 5.212 | 0.301 | 2.673 | 0.497 / **1.719** | 1.6x |
+
+(5090, torch 2.11.0+cu128, median of 10 after 5 warmup, harness includes
+per-iter leaf clones + loss; flex rows are the best of
+{default, bwd_shrunk_only} x {64x64, 128x128}.)
+
+- vs the SDPA dense-mask baseline the manual backend is **2.3-3.1x** on
+  fwd+bwd, on every shape.
+- manual peak (fwd) memory is **0.35-0.43x** of the SDPA path — the dense
+  additive mask the SDPA path materializes costs more than the transient
+  logits.
+- `TORCHINDUCTOR_MAX_AUTOTUNE=1` trims manual fwd another ~14%
+  (0.360 -> 0.311 at b4_p700_k5); fwd+bwd unchanged.
+
+## Why (mechanism, not vibes)
+
+- **GQA is the fidelity key.** With 8 kv heads (earlier sweeps) flex
+  fwd+bwd looked best. At the real kv_heads=1: the flex BACKWARD
+  autotune has no valid config on the 5090 (needs 112 KB shared memory,
+  hardware 101 KB) and even the shrunken-tile fallback runs 2.2-2.7 ms —
+  the backward dominates everything. The manual backward is four cuBLAS
+  GEMMs + a fused softmax-grad chain: 0.5-1.4 ms.
+- flex **forward** stays the fastest single direction (autotune,
+  128x128 or 64x64 masks) — but you cannot combine flex-fwd with
+  manual-bwd, and fwd+bwd is what training pays.
+- Precision class: manual materializes logits in bf16, so outputs differ
+  from the fp32-softmax fused paths by up to ~1.2e-2 max-abs (flex class:
+  ~2e-3). Model-level parity gates (loss rel <= 1e-3, grads <= 1%) are
+  the ship test; the fwd-diff gate for manual is tracked separately.
+
+## Per-architecture verdict (final, 2026-07-09)
+
+| arch | local microbench | real-model E2E | verdict |
+| --- | --- | --- | --- |
+| RTX 5090 (sm120 consumer) | manual 2.3-3.1x vs SDPA, 1.4-2.9x vs flex | flow -7.8% / text -11.6% | **manual ON** |
+| A100 (sm80) | manual wins, incl. vs the repeat-interleave production baseline (B1/P1024: 1.98 vs 3.60 ms) | LOSES (text 436 vs 386 ms; action-only scope and eager-vs-compiled both exonerated) | manual OFF — integration-level interaction, parked |
+| H200 (sm90) | manual LOSES the microbench outright (1.80 vs sdpa_repeat 1.18 at B1/P1024) — Hopper FMHA is too strong | not needed | manual OFF |
+| RTX PRO 6000 (sm120 workstation) | pending | pending | expected ON (5090 arch family) |
+
+`flex_attention(impl="auto")` encodes this: manual only on sm120-class
+CUDA with no dropout, SDPA elsewhere.
+
+## Dispatch recommendation (as of 2026-07-09)
+
+- 5090-class training (fwd+bwd): `manual` everywhere;
+  `TORCHINDUCTOR_MAX_AUTOTUNE` optional (+14% fwd).
+- fwd-only (inference prefill): flex `default` autotune, mask 64x64
+  (128x128 for P>=896-class shapes).
+- A100/H100/H200: SDPA path (see the verdict table).
+
+## Next levers (not yet implemented)
+
+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
+   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
+   remaining gap is confirmed on both archs.
+
+## History
+
+Earlier 8-kv-head sweeps (superseded — wrong KV shape for PI052):
+5090 flex fwd+bwd best 1.48-2.55 ms vs SDPA 3.97; A100 matrix at 8 kv
+heads showed flex positive at real shapes with 64x64 masks. Those runs
+also established: torch autotune beats every hand preset at 8 heads;
+`torch_default_explicit` NaNs at K=1; ROWS_GUARANTEED_SAFE /
+BLOCKS_ARE_CONTIGUOUS / PRESCALE_QK no help on 5090.
+
+No native CUDA performance results are claimed yet.

benchmarks/benchmark.py

@@ -0,0 +1,779 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+import time
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT / "torch-ext"))
+import flashrt_flex_attention_train as flex_ops  # noqa: E402
+
+
+FLEX_TILE_PRESETS = {
+    "default": None,
+    # fwd fully autotuned; backward constrained to the consumer-GPU tiles
+    # (autotuned backward at GQA/D=256 exceeds 5090 shared memory).
+    "bwd_shrunk_only": {
+        "bwd_BLOCK_M1": 32,
+        "bwd_BLOCK_N1": 64,
+        "bwd_BLOCK_M2": 64,
+        "bwd_BLOCK_N2": 32,
+    },
+    "torch_default_explicit": {
+        "fwd_BLOCK_M": 32,
+        "fwd_BLOCK_N": 64,
+        "fwd_num_stages": 2,
+        "fwd_num_warps": 4,
+        "bwd_BLOCK_M1": 32,
+        "bwd_BLOCK_N1": 64,
+        "bwd_BLOCK_M2": 64,
+        "bwd_BLOCK_N2": 32,
+        "bwd_num_stages": 1,
+        "bwd_num_warps": 4,
+    },
+    "a100_d256_bwd_32x64": {
+        "fwd_BLOCK_M": 32,
+        "fwd_BLOCK_N": 64,
+        "fwd_num_stages": 2,
+        "fwd_num_warps": 4,
+        "bwd_BLOCK_M1": 32,
+        "bwd_BLOCK_N1": 64,
+        "bwd_BLOCK_M2": 64,
+        "bwd_BLOCK_N2": 32,
+        "bwd_num_stages": 3,
+        "bwd_num_warps": 4,
+    },
+    "a100_d256_bwd_32x128": {
+        "fwd_BLOCK_M": 32,
+        "fwd_BLOCK_N": 64,
+        "fwd_num_stages": 2,
+        "fwd_num_warps": 4,
+        "bwd_BLOCK_M1": 32,
+        "bwd_BLOCK_N1": 128,
+        "bwd_BLOCK_M2": 128,
+        "bwd_BLOCK_N2": 32,
+        "bwd_num_stages": 3,
+        "bwd_num_warps": 8,
+    },
+    "a100_d256_bwd_64x64": {
+        "fwd_BLOCK_M": 64,
+        "fwd_BLOCK_N": 64,
+        "fwd_num_stages": 3,
+        "fwd_num_warps": 4,
+        "bwd_BLOCK_M1": 64,
+        "bwd_BLOCK_N1": 64,
+        "bwd_BLOCK_M2": 64,
+        "bwd_BLOCK_N2": 64,
+        "bwd_num_stages": 3,
+        "bwd_num_warps": 4,
+    },
+    "a100_d256_bwd_64x128": {
+        "fwd_BLOCK_M": 64,
+        "fwd_BLOCK_N": 128,
+        "fwd_num_stages": 3,
+        "fwd_num_warps": 4,
+        "bwd_BLOCK_M1": 64,
+        "bwd_BLOCK_N1": 128,
+        "bwd_BLOCK_M2": 128,
+        "bwd_BLOCK_N2": 64,
+        "bwd_num_stages": 3,
+        "bwd_num_warps": 8,
+    },
+    "a100_d256_bwd_write_dq_false": {
+        "fwd_BLOCK_M": 32,
+        "fwd_BLOCK_N": 64,
+        "fwd_num_stages": 2,
+        "fwd_num_warps": 4,
+        "bwd_BLOCK_M1": 32,
+        "bwd_BLOCK_N1": 64,
+        "bwd_BLOCK_M2": 64,
+        "bwd_BLOCK_N2": 32,
+        "bwd_num_stages": 3,
+        "bwd_num_warps": 4,
+        "WRITE_DQ": False,
+    },
+    "a100_d256_prescale_safe": {
+        "fwd_BLOCK_M": 32,
+        "fwd_BLOCK_N": 64,
+        "fwd_num_stages": 2,
+        "fwd_num_warps": 4,
+        "bwd_BLOCK_M1": 32,
+        "bwd_BLOCK_N1": 64,
+        "bwd_BLOCK_M2": 64,
+        "bwd_BLOCK_N2": 32,
+        "bwd_num_stages": 3,
+        "bwd_num_warps": 4,
+        "PRESCALE_QK": True,
+        "ROWS_GUARANTEED_SAFE": True,
+    },
+    "a100_d256_contig_safe": {
+        "fwd_BLOCK_M": 32,
+        "fwd_BLOCK_N": 64,
+        "fwd_num_stages": 2,
+        "fwd_num_warps": 4,
+        "bwd_BLOCK_M1": 32,
+        "bwd_BLOCK_N1": 64,
+        "bwd_BLOCK_M2": 64,
+        "bwd_BLOCK_N2": 32,
+        "bwd_num_stages": 3,
+        "bwd_num_warps": 4,
+        "ROWS_GUARANTEED_SAFE": True,
+        "BLOCKS_ARE_CONTIGUOUS": True,
+    },
+    "a100_d256_contig_prescale": {
+        "fwd_BLOCK_M": 32,
+        "fwd_BLOCK_N": 64,
+        "fwd_num_stages": 2,
+        "fwd_num_warps": 4,
+        "bwd_BLOCK_M1": 32,
+        "bwd_BLOCK_N1": 64,
+        "bwd_BLOCK_M2": 64,
+        "bwd_BLOCK_N2": 32,
+        "bwd_num_stages": 3,
+        "bwd_num_warps": 4,
+        "PRESCALE_QK": True,
+        "ROWS_GUARANTEED_SAFE": True,
+        "BLOCKS_ARE_CONTIGUOUS": True,
+    },
+    "a100_d256_contig_write_dq_false": {
+        "fwd_BLOCK_M": 32,
+        "fwd_BLOCK_N": 64,
+        "fwd_num_stages": 2,
+        "fwd_num_warps": 4,
+        "bwd_BLOCK_M1": 32,
+        "bwd_BLOCK_N1": 64,
+        "bwd_BLOCK_M2": 64,
+        "bwd_BLOCK_N2": 32,
+        "bwd_num_stages": 3,
+        "bwd_num_warps": 4,
+        "ROWS_GUARANTEED_SAFE": True,
+        "BLOCKS_ARE_CONTIGUOUS": True,
+        "WRITE_DQ": False,
+    },
+}
+
+
+def bench(fn, warmup: int, iters: int) -> float:
+    for _ in range(warmup):
+        fn()
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+    times = []
+    for _ in range(iters):
+        if torch.cuda.is_available():
+            t0 = torch.cuda.Event(enable_timing=True)
+            t1 = torch.cuda.Event(enable_timing=True)
+            t0.record()
+            fn()
+            t1.record()
+            torch.cuda.synchronize()
+            times.append(t0.elapsed_time(t1))
+        else:
+            start = time.perf_counter()
+            fn()
+            times.append((time.perf_counter() - start) * 1000.0)
+    times.sort()
+    return times[len(times) // 2]
+
+
+def peak_bytes(fn, device: str) -> int:
+    if not device.startswith("cuda"):
+        fn()
+        return 0
+    torch.cuda.reset_peak_memory_stats()
+    fn()
+    torch.cuda.synchronize()
+    return int(torch.cuda.max_memory_allocated())
+
+
+def best_nested_speedup(nested: dict) -> dict | None:
+    best = None
+    for block_key, by_preset in nested.items():
+        for preset, speedup in by_preset.items():
+            if not isinstance(speedup, float):
+                continue
+            if best is None or speedup > best["speedup"]:
+                best = {"block_mask": block_key, "preset": preset, "speedup": speedup}
+    return best
+
+
+def make_inputs(args):
+    dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float32
+    action_len = args.action_blocks * args.action_block_size
+    total = args.prefix_len + action_len
+    torch.manual_seed(args.seed)
+    q = torch.randn(args.batch, args.heads, total, args.head_dim, device=args.device, dtype=dtype)
+    kv_heads = args.kv_heads or args.heads
+    k = torch.randn(args.batch, kv_heads, total, args.head_dim, device=args.device, dtype=dtype)
+    v = torch.randn_like(k)
+    prefix_valid = torch.ones(args.batch, args.prefix_len, device=args.device, dtype=torch.bool)
+    prefix_att = torch.zeros_like(prefix_valid)
+    prefix_att[:, args.prefix_len // 2 :] = True
+    if args.no_prefix_mask:
+        prefix_valid = None
+        prefix_att = None
+    return q, k, v, prefix_valid, prefix_att
+
+
+def make_flex_bundle(args, prefix_valid, prefix_att, block_size: tuple[int, int]):
+    try:
+        from torch.nn.attention.flex_attention import create_block_mask, flex_attention
+    except Exception:
+        return None
+
+    compiled_mask = torch.compile(create_block_mask, dynamic=False)
+    batch = args.batch
+    prefix_len = args.prefix_len
+    action_len = args.action_blocks * args.action_block_size
+    total_len = prefix_len + action_len
+    chunk = args.action_block_size
+    pad = prefix_valid
+    cum = torch.cumsum(prefix_att.to(torch.long), dim=1) if prefix_att is not None else None
+
+    def prefix_rows(b, h, q_idx, kv_idx):
+        if cum is None:
+            return kv_idx < prefix_len
+        kv_p = kv_idx.clamp(max=prefix_len - 1)
+        ok = (cum[b, kv_p] <= cum[b, q_idx]) & pad[b, kv_p] & pad[b, q_idx]
+        return (kv_idx < prefix_len) & ok
+
+    def action_rows(b, h, q_idx, kv_idx):
+        if pad is None:
+            to_prefix = kv_idx < prefix_len
+        else:
+            kv_p = kv_idx.clamp(max=prefix_len - 1)
+            to_prefix = (kv_idx < prefix_len) & pad[b, kv_p]
+        same_block = (q_idx // chunk) == ((kv_idx - prefix_len) // chunk)
+        return to_prefix | ((kv_idx >= prefix_len) & same_block)
+
+    block_prefix = compiled_mask(
+        prefix_rows,
+        B=batch,
+        H=None,
+        Q_LEN=prefix_len,
+        KV_LEN=total_len,
+        device=torch.device(args.device),
+        BLOCK_SIZE=block_size,
+    )
+    block_action = compiled_mask(
+        action_rows,
+        B=batch,
+        H=None,
+        Q_LEN=action_len,
+        KV_LEN=total_len,
+        device=torch.device(args.device),
+        BLOCK_SIZE=block_size,
+    )
+    compiled_calls = {}
+    scale = args.head_dim**-0.5
+    gqa = (args.kv_heads or args.heads) != args.heads
+    for name, options in FLEX_TILE_PRESETS.items():
+        def prefix_call(q, k, v, block_mask, options=options):
+            return flex_attention(
+                q,
+                k,
+                v,
+                block_mask=block_mask,
+                scale=scale,
+                enable_gqa=gqa,
+                kernel_options=options,
+            )
+
+        def action_call(q, k, v, block_mask, options=options):
+            return flex_attention(
+                q,
+                k,
+                v,
+                block_mask=block_mask,
+                scale=scale,
+                enable_gqa=gqa,
+                kernel_options=options,
+            )
+
+        compiled_calls[name] = (
+            torch.compile(prefix_call, dynamic=False),
+            torch.compile(action_call, dynamic=False),
+        )
+    return compiled_calls, block_prefix, block_action
+
+
+def _manual_part(qs, ks, vs, m, scale):
+    """Materialized-logits attention: cuBLAS GEMMs + fused masked softmax.
+
+    Exact SDPA semantics (fp32 softmax); grouped-query handled as a strided
+    batched GEMM over (kv_head, group*Sq) without materializing repeated K/V.
+    """
+    B, H, Sq, D = qs.shape
+    Hk = ks.shape[1]
+    if Hk != H:
+        g = H // Hk
+        q2 = qs.reshape(B, Hk, g * Sq, D)
+        logits = (q2 @ ks.transpose(-1, -2)).reshape(B, H, Sq, -1)
+    else:
+        logits = qs @ ks.transpose(-1, -2)
+    logits = logits * scale + m
+    p = logits.float().softmax(dim=-1).to(qs.dtype)
+    if Hk != H:
+        out = (p.reshape(B, Hk, g * Sq, -1) @ vs).reshape(B, H, Sq, D)
+    else:
+        out = p @ vs
+    return out
+
+
+_manual_part_compiled = None
+
+
+def get_manual_part():
+    global _manual_part_compiled
+    if _manual_part_compiled is None:
+        _manual_part_compiled = torch.compile(_manual_part, dynamic=False)
+    return _manual_part_compiled
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--dtype", choices=["bf16", "fp32"], default="bf16")
+    parser.add_argument("--batch", type=int, default=4)
+    parser.add_argument("--heads", type=int, default=8)
+    parser.add_argument("--kv-heads", type=int, default=None, help="KV heads for GQA (default: same as --heads)")
+    parser.add_argument("--head-dim", type=int, default=256)
+    parser.add_argument("--prefix-len", type=int, default=700)
+    parser.add_argument("--action-blocks", type=int, default=5)
+    parser.add_argument("--action-block-size", type=int, default=50)
+    parser.add_argument("--warmup", type=int, default=10)
+    parser.add_argument("--iters", type=int, default=30)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--mode", choices=["fwd", "fwdbwd", "all"], default="all")
+    parser.add_argument(
+        "--backend",
+        default="all",
+        help="comma-separated subset of {package, torch-flex, manual} or 'all'",
+    )
+    parser.add_argument("--flex-preset", choices=sorted(FLEX_TILE_PRESETS), default="a100_d256_bwd_32x64")
+    parser.add_argument("--sweep-flex-presets", action="store_true")
+    parser.add_argument("--block-mask-q", type=int, default=128)
+    parser.add_argument("--block-mask-kv", type=int, default=128)
+    parser.add_argument("--sweep-block-mask-sizes", action="store_true")
+    parser.add_argument("--output")
+    parser.add_argument("--require-gates", action="store_true")
+    parser.add_argument("--no-prefix-mask", action="store_true")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    if args.device.startswith("cuda") and not torch.cuda.is_available():
+        raise SystemExit("CUDA requested but not available")
+    backends = (
+        {"package", "torch-flex", "manual"}
+        if args.backend == "all"
+        else {b.strip() for b in args.backend.split(",") if b.strip()}
+    )
+
+    q, k, v, prefix_valid, prefix_att = make_inputs(args)
+    action_len = args.action_blocks * args.action_block_size
+    scale = args.head_dim**-0.5
+    pm, am = flex_ops.build_block_sparse_bool_masks(
+        prefix_valid,
+        prefix_att,
+        batch=args.batch,
+        prefix_len=args.prefix_len,
+        action_len=action_len,
+        action_block_size=args.action_block_size,
+        device=q.device,
+    )
+    full = torch.cat([pm, am], dim=1)
+    add_mask = torch.where(
+        full[:, None],
+        torch.zeros((), device=q.device, dtype=q.dtype),
+        torch.full((), flex_ops.MASK_VALUE_F32, device=q.device, dtype=q.dtype),
+    )
+    block_sizes = [(args.block_mask_q, args.block_mask_kv)]
+    if args.sweep_block_mask_sizes:
+        block_sizes = [(q, kv) for q in (16, 32, 64, 128) for kv in (32, 64, 128)]
+    flex_bundles = {
+        f"{q}x{kv}": make_flex_bundle(args, prefix_valid, prefix_att, (q, kv))
+        for q, kv in block_sizes
+    } if "torch-flex" in backends else {}
+    flex_presets = sorted(FLEX_TILE_PRESETS) if args.sweep_flex_presets else [args.flex_preset]
+
+    gqa = (args.kv_heads or args.heads) != args.heads
+
+    def sdpa_fwd():
+        out_p = F.scaled_dot_product_attention(
+            q[:, :, : args.prefix_len],
+            k,
+            v,
+            attn_mask=add_mask[:, :, : args.prefix_len],
+            scale=scale,
+            enable_gqa=gqa,
+        )
+        kd = torch.cat([k[:, :, : args.prefix_len].detach(), k[:, :, args.prefix_len :]], dim=2)
+        vd = torch.cat([v[:, :, : args.prefix_len].detach(), v[:, :, args.prefix_len :]], dim=2)
+        out_a = F.scaled_dot_product_attention(
+            q[:, :, args.prefix_len :],
+            kd,
+            vd,
+            attn_mask=add_mask[:, :, args.prefix_len :],
+            scale=scale,
+            enable_gqa=gqa,
+        )
+        return torch.cat([out_p, out_a], dim=2)
+
+    manual_part = get_manual_part() if "manual" in backends else None
+
+    def manual_eager_fwd():
+        out_p = _manual_part(q[:, :, : args.prefix_len], k, v, add_mask[:, :, : args.prefix_len], scale)
+        kd = torch.cat([k[:, :, : args.prefix_len].detach(), k[:, :, args.prefix_len :]], dim=2)
+        vd = torch.cat([v[:, :, : args.prefix_len].detach(), v[:, :, args.prefix_len :]], dim=2)
+        out_a = _manual_part(
+            q[:, :, args.prefix_len :], kd, vd, add_mask[:, :, args.prefix_len :], scale
+        )
+        return torch.cat([out_p, out_a], dim=2)
+
+    def manual_eager_fwdbwd():
+        qq = q.detach().clone().requires_grad_(True)
+        kk = k.detach().clone().requires_grad_(True)
+        vv = v.detach().clone().requires_grad_(True)
+        out_p = _manual_part(
+            qq[:, :, : args.prefix_len], kk, vv, add_mask[:, :, : args.prefix_len], scale
+        )
+        kd = torch.cat([kk[:, :, : args.prefix_len].detach(), kk[:, :, args.prefix_len :]], dim=2)
+        vd = torch.cat([vv[:, :, : args.prefix_len].detach(), vv[:, :, args.prefix_len :]], dim=2)
+        out_a = _manual_part(
+            qq[:, :, args.prefix_len :], kd, vd, add_mask[:, :, args.prefix_len :], scale
+        )
+        torch.cat([out_p, out_a], dim=2).float().square().mean().backward()
+
+    def _repeat_kv(t):
+        # the model's current path: materialize K/V to the q-head count
+        return t.repeat_interleave(args.heads // t.shape[1], dim=1)
+
+    def sdpa_repeat_fwd():
+        kr, vr = _repeat_kv(k), _repeat_kv(v)
+        out_p = F.scaled_dot_product_attention(
+            q[:, :, : args.prefix_len], kr, vr, attn_mask=add_mask[:, :, : args.prefix_len], scale=scale
+        )
+        kd = torch.cat([kr[:, :, : args.prefix_len].detach(), kr[:, :, args.prefix_len :]], dim=2)
+        vd = torch.cat([vr[:, :, : args.prefix_len].detach(), vr[:, :, args.prefix_len :]], dim=2)
+        out_a = F.scaled_dot_product_attention(
+            q[:, :, args.prefix_len :], kd, vd, attn_mask=add_mask[:, :, args.prefix_len :], scale=scale
+        )
+        return torch.cat([out_p, out_a], dim=2)
+
+    def manual_fwd():
+        out_p = manual_part(q[:, :, : args.prefix_len], k, v, add_mask[:, :, : args.prefix_len], scale)
+        kd = torch.cat([k[:, :, : args.prefix_len].detach(), k[:, :, args.prefix_len :]], dim=2)
+        vd = torch.cat([v[:, :, : args.prefix_len].detach(), v[:, :, args.prefix_len :]], dim=2)
+        out_a = manual_part(
+            q[:, :, args.prefix_len :], kd, vd, add_mask[:, :, args.prefix_len :], scale
+        )
+        return torch.cat([out_p, out_a], dim=2)
+
+    def package_fwd():
+        return flex_ops.flex_attention(
+            q,
+            k,
+            v,
+            prefix_len=args.prefix_len,
+            action_block_size=args.action_block_size,
+            prefix_valid=prefix_valid,
+            prefix_att=prefix_att,
+            scale=scale,
+        )
+
+    def torch_flex_fwd(preset: str, block_key: str):
+        if not flex_bundles:
+            raise RuntimeError("PyTorch FlexAttention is unavailable")
+        compiled_calls, block_prefix, block_action = flex_bundles[block_key]
+        prefix_call, action_call = compiled_calls[preset]
+        out_p = prefix_call(
+            q[:, :, : args.prefix_len],
+            k,
+            v,
+            block_prefix,
+        )
+        kd = torch.cat([k[:, :, : args.prefix_len].detach(), k[:, :, args.prefix_len :]], dim=2)
+        vd = torch.cat([v[:, :, : args.prefix_len].detach(), v[:, :, args.prefix_len :]], dim=2)
+        out_a = action_call(
+            q[:, :, args.prefix_len :],
+            kd,
+            vd,
+            block_action,
+        )
+        return torch.cat([out_p, out_a], dim=2)
+
+    report = {
+        "gpu": torch.cuda.get_device_name() if args.device.startswith("cuda") else "cpu",
+        "torch": torch.__version__,
+        "shape": {
+            "B": args.batch,
+            "heads": args.heads,
+            "kv_heads": args.kv_heads or args.heads,
+            "head_dim": args.head_dim,
+            "prefix_len": args.prefix_len,
+            "action_len": action_len,
+            "action_block_size": args.action_block_size,
+        },
+        "backend": args.backend,
+        "flex_presets": flex_presets,
+        "block_mask_sizes": list(flex_bundles) if flex_bundles else [],
+        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
+    }
+    report["package_has_native_ops"] = False
+    report["package_native_supported"] = False
+
+    with torch.no_grad():
+        sdpa_out = sdpa_fwd().float()
+        if "package" in backends:
+            report["package_fwd_max_abs_diff"] = float((sdpa_out - package_fwd().float()).abs().max())
+        if "manual" in backends:
+            report["manual_fwd_max_abs_diff"] = float((sdpa_out - manual_fwd().float()).abs().max())
+        if flex_bundles:
+            report["torch_flex_fwd_max_abs_diff"] = {}
+            for block_key in flex_bundles:
+                report["torch_flex_fwd_max_abs_diff"][block_key] = {}
+            for block_key in flex_bundles:
+              for preset in flex_presets:
+                try:
+                    report["torch_flex_fwd_max_abs_diff"][block_key][preset] = float(
+                        (sdpa_out - torch_flex_fwd(preset, block_key).float()).abs().max()
+                    )
+                except Exception as exc:
+                    report["torch_flex_fwd_max_abs_diff"][block_key][preset] = f"ERROR: {type(exc).__name__}: {exc}"
+        if args.mode in {"fwd", "all"}:
+            report["sdpa_fwd_ms"] = bench(sdpa_fwd, args.warmup, args.iters)
+            report["sdpa_peak_bytes"] = peak_bytes(sdpa_fwd, args.device)
+            if "package" in backends:
+                report["package_fwd_ms"] = bench(package_fwd, args.warmup, args.iters)
+                report["package_fwd_speedup"] = report["sdpa_fwd_ms"] / report["package_fwd_ms"]
+                report["package_peak_bytes"] = peak_bytes(package_fwd, args.device)
+            if "manual" in backends:
+                report["manual_fwd_ms"] = bench(manual_fwd, args.warmup, args.iters)
+                report["manual_fwd_speedup"] = report["sdpa_fwd_ms"] / report["manual_fwd_ms"]
+                report["manual_peak_bytes"] = peak_bytes(manual_fwd, args.device)
+            if gqa:
+                report["sdpa_repeat_fwd_ms"] = bench(sdpa_repeat_fwd, args.warmup, args.iters)
+            if "manual" in backends:
+                report["manual_eager_fwd_ms"] = bench(manual_eager_fwd, args.warmup, args.iters)
+            if flex_bundles:
+                report["torch_flex_fwd_ms"] = {}
+                report["torch_flex_fwd_speedup"] = {}
+                for block_key in flex_bundles:
+                    report["torch_flex_fwd_ms"][block_key] = {}
+                    report["torch_flex_fwd_speedup"][block_key] = {}
+                for block_key in flex_bundles:
+                  for preset in flex_presets:
+                    try:
+                        ms = bench(lambda preset=preset, block_key=block_key: torch_flex_fwd(preset, block_key), args.warmup, args.iters)
+                    except Exception as exc:
+                        report["torch_flex_fwd_ms"][block_key][preset] = f"ERROR: {type(exc).__name__}: {exc}"
+                    else:
+                        report["torch_flex_fwd_ms"][block_key][preset] = ms
+                        report["torch_flex_fwd_speedup"][block_key][preset] = report["sdpa_fwd_ms"] / ms
+                report["best_torch_flex_fwd"] = best_nested_speedup(report["torch_flex_fwd_speedup"])
+
+    if args.mode in {"fwdbwd", "all"}:
+        def sdpa_fwdbwd():
+            qq = q.detach().clone().requires_grad_(True)
+            kk = k.detach().clone().requires_grad_(True)
+            vv = v.detach().clone().requires_grad_(True)
+            out_p = F.scaled_dot_product_attention(
+                qq[:, :, : args.prefix_len],
+                kk,
+                vv,
+                attn_mask=add_mask[:, :, : args.prefix_len],
+                scale=scale,
+                enable_gqa=gqa,
+            )
+            kd = torch.cat([kk[:, :, : args.prefix_len].detach(), kk[:, :, args.prefix_len :]], dim=2)
+            vd = torch.cat([vv[:, :, : args.prefix_len].detach(), vv[:, :, args.prefix_len :]], dim=2)
+            out_a = F.scaled_dot_product_attention(
+                qq[:, :, args.prefix_len :],
+                kd,
+                vd,
+                attn_mask=add_mask[:, :, args.prefix_len :],
+                scale=scale,
+                enable_gqa=gqa,
+            )
+            torch.cat([out_p, out_a], dim=2).float().square().mean().backward()
+
+        def manual_fwdbwd():
+            qq = q.detach().clone().requires_grad_(True)
+            kk = k.detach().clone().requires_grad_(True)
+            vv = v.detach().clone().requires_grad_(True)
+            out_p = manual_part(
+                qq[:, :, : args.prefix_len], kk, vv, add_mask[:, :, : args.prefix_len], scale
+            )
+            kd = torch.cat([kk[:, :, : args.prefix_len].detach(), kk[:, :, args.prefix_len :]], dim=2)
+            vd = torch.cat([vv[:, :, : args.prefix_len].detach(), vv[:, :, args.prefix_len :]], dim=2)
+            out_a = manual_part(
+                qq[:, :, args.prefix_len :], kd, vd, add_mask[:, :, args.prefix_len :], scale
+            )
+            torch.cat([out_p, out_a], dim=2).float().square().mean().backward()
+
+        def sdpa_repeat_fwdbwd():
+            qq = q.detach().clone().requires_grad_(True)
+            kk = k.detach().clone().requires_grad_(True)
+            vv = v.detach().clone().requires_grad_(True)
+            kr, vr = _repeat_kv(kk), _repeat_kv(vv)
+            out_p = F.scaled_dot_product_attention(
+                qq[:, :, : args.prefix_len], kr, vr, attn_mask=add_mask[:, :, : args.prefix_len], scale=scale
+            )
+            kd = torch.cat([kr[:, :, : args.prefix_len].detach(), kr[:, :, args.prefix_len :]], dim=2)
+            vd = torch.cat([vr[:, :, : args.prefix_len].detach(), vr[:, :, args.prefix_len :]], dim=2)
+            out_a = F.scaled_dot_product_attention(
+                qq[:, :, args.prefix_len :], kd, vd, attn_mask=add_mask[:, :, args.prefix_len :], scale=scale
+            )
+            torch.cat([out_p, out_a], dim=2).float().square().mean().backward()
+
+        def package_fwdbwd():
+            qq = q.detach().clone().requires_grad_(True)
+            kk = k.detach().clone().requires_grad_(True)
+            vv = v.detach().clone().requires_grad_(True)
+            flex_ops.flex_attention(
+                qq,
+                kk,
+                vv,
+                prefix_len=args.prefix_len,
+                action_block_size=args.action_block_size,
+                prefix_valid=prefix_valid,
+                prefix_att=prefix_att,
+                scale=scale,
+            ).float().square().mean().backward()
+
+        def torch_flex_fwdbwd(preset: str, block_key: str):
+            if not flex_bundles:
+                raise RuntimeError("PyTorch FlexAttention is unavailable")
+            compiled_calls, block_prefix, block_action = flex_bundles[block_key]
+            prefix_call, action_call = compiled_calls[preset]
+            qq = q.detach().clone().requires_grad_(True)
+            kk = k.detach().clone().requires_grad_(True)
+            vv = v.detach().clone().requires_grad_(True)
+            out_p = prefix_call(
+                qq[:, :, : args.prefix_len],
+                kk,
+                vv,
+                block_prefix,
+            )
+            kd = torch.cat([kk[:, :, : args.prefix_len].detach(), kk[:, :, args.prefix_len :]], dim=2)
+            vd = torch.cat([vv[:, :, : args.prefix_len].detach(), vv[:, :, args.prefix_len :]], dim=2)
+            out_a = action_call(
+                qq[:, :, args.prefix_len :],
+                kd,
+                vd,
+                block_action,
+            )
+            torch.cat([out_p, out_a], dim=2).float().square().mean().backward()
+
+        report["sdpa_fwdbwd_ms"] = bench(sdpa_fwdbwd, args.warmup, args.iters)
+        if "package" in backends:
+            report["package_fwdbwd_ms"] = bench(package_fwdbwd, args.warmup, args.iters)
+            report["package_fwdbwd_speedup"] = report["sdpa_fwdbwd_ms"] / report["package_fwdbwd_ms"]
+        if "manual" in backends:
+            report["manual_fwdbwd_ms"] = bench(manual_fwdbwd, args.warmup, args.iters)
+            report["manual_fwdbwd_speedup"] = report["sdpa_fwdbwd_ms"] / report["manual_fwdbwd_ms"]
+        if gqa:
+            report["sdpa_repeat_fwdbwd_ms"] = bench(sdpa_repeat_fwdbwd, args.warmup, args.iters)
+        if "manual" in backends:
+            report["manual_eager_fwdbwd_ms"] = bench(manual_eager_fwdbwd, args.warmup, args.iters)
+        if flex_bundles:
+            report["torch_flex_fwdbwd_ms"] = {}
+            report["torch_flex_fwdbwd_speedup"] = {}
+            for block_key in flex_bundles:
+                report["torch_flex_fwdbwd_ms"][block_key] = {}
+                report["torch_flex_fwdbwd_speedup"][block_key] = {}
+            for block_key in flex_bundles:
+              for preset in flex_presets:
+                try:
+                    ms = bench(lambda preset=preset, block_key=block_key: torch_flex_fwdbwd(preset, block_key), args.warmup, args.iters)
+                except Exception as exc:
+                    report["torch_flex_fwdbwd_ms"][block_key][preset] = f"ERROR: {type(exc).__name__}: {exc}"
+                else:
+                    report["torch_flex_fwdbwd_ms"][block_key][preset] = ms
+                    report["torch_flex_fwdbwd_speedup"][block_key][preset] = report["sdpa_fwdbwd_ms"] / ms
+            report["best_torch_flex_fwdbwd"] = best_nested_speedup(report["torch_flex_fwdbwd_speedup"])
+
+        if "package" not in backends:
+            _finish(report, args)
+            return
+        q1 = q.detach().clone().requires_grad_(True)
+        k1 = k.detach().clone().requires_grad_(True)
+        v1 = v.detach().clone().requires_grad_(True)
+        q2 = q.detach().clone().requires_grad_(True)
+        k2 = k.detach().clone().requires_grad_(True)
+        v2 = v.detach().clone().requires_grad_(True)
+        out1_p = F.scaled_dot_product_attention(
+            q1[:, :, : args.prefix_len], k1, v1, attn_mask=add_mask[:, :, : args.prefix_len], scale=scale
+        )
+        k1d = torch.cat([k1[:, :, : args.prefix_len].detach(), k1[:, :, args.prefix_len :]], dim=2)
+        v1d = torch.cat([v1[:, :, : args.prefix_len].detach(), v1[:, :, args.prefix_len :]], dim=2)
+        out1_a = F.scaled_dot_product_attention(
+            q1[:, :, args.prefix_len :], k1d, v1d, attn_mask=add_mask[:, :, args.prefix_len :], scale=scale
+        )
+        torch.cat([out1_p, out1_a], dim=2).float().square().mean().backward()
+        flex_ops.flex_attention(
+            q2,
+            k2,
+            v2,
+            prefix_len=args.prefix_len,
+            action_block_size=args.action_block_size,
+            prefix_valid=prefix_valid,
+            prefix_att=prefix_att,
+            scale=scale,
+        ).float().square().mean().backward()
+        denom = torch.linalg.vector_norm(torch.cat([q1.grad.flatten(), k1.grad.flatten(), v1.grad.flatten()])).clamp_min(1e-12)
+        numer = torch.linalg.vector_norm(
+            torch.cat([(q1.grad - q2.grad).flatten(), (k1.grad - k2.grad).flatten(), (v1.grad - v2.grad).flatten()])
+        )
+        report["package_grad_norm_rel_diff"] = float(numer / denom)
+
+    _finish(report, args)
+
+
+def _finish(report, args):
+    gates = {
+        "package_fwd_max_abs_diff": report.get("package_fwd_max_abs_diff", 0.0) <= 2e-3,
+        "package_grad_norm_rel_diff": report.get("package_grad_norm_rel_diff", 0.0) <= 1e-2,
+    }
+    if "torch_flex_fwd_max_abs_diff" in report:
+        vals = [v for by_block in report["torch_flex_fwd_max_abs_diff"].values() for v in by_block.values()]
+        gates["torch_flex_fwd_max_abs_diff"] = any(isinstance(v, float) and v <= 2e-3 for v in vals)
+    if "manual_fwd_max_abs_diff" in report:
+        gates["manual_fwd_max_abs_diff"] = report["manual_fwd_max_abs_diff"] <= 2e-3
+    if "manual_fwd_speedup" in report:
+        gates["manual_fwd_speedup_ge_1p0526"] = report["manual_fwd_speedup"] >= (1.0 / 0.95)
+    if "manual_fwdbwd_speedup" in report:
+        gates["manual_fwdbwd_speedup_ge_1p0526"] = report["manual_fwdbwd_speedup"] >= (1.0 / 0.95)
+    if "package_fwd_speedup" in report:
+        gates["package_fwd_speedup_ge_1p0526"] = report["package_fwd_speedup"] >= (1.0 / 0.95)
+        if report.get("sdpa_peak_bytes", 0) > 0:
+            gates["package_peak_memory_le_sdpa_plus_3pct"] = report["package_peak_bytes"] <= int(report["sdpa_peak_bytes"] * 1.03)
+    if "package_fwdbwd_speedup" in report:
+        gates["package_fwdbwd_speedup_ge_1p0526"] = report["package_fwdbwd_speedup"] >= (1.0 / 0.95)
+    if "torch_flex_fwd_speedup" in report:
+        vals = [v for by_block in report["torch_flex_fwd_speedup"].values() for v in by_block.values() if isinstance(v, float)]
+        gates["torch_flex_fwd_speedup_ge_1p0526"] = bool(vals) and max(vals) >= (1.0 / 0.95)
+    if "torch_flex_fwdbwd_speedup" in report:
+        vals = [v for by_block in report["torch_flex_fwdbwd_speedup"].values() for v in by_block.values() if isinstance(v, float)]
+        gates["torch_flex_fwdbwd_speedup_ge_1p0526"] = bool(vals) and max(vals) >= (1.0 / 0.95)
+    report["gates"] = gates
+
+    text = json.dumps(report, indent=2)
+    if args.output:
+        Path(args.output).write_text(text + "\n", encoding="utf-8")
+    print(text)
+    if args.require_gates and not all(gates.values()):
+        raise SystemExit("one or more Flex attention acceptance gates failed")
+
+
+if __name__ == "__main__":
+    main()

benchmarks/shape_matrix.py

@@ -0,0 +1,169 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import json
+import subprocess
+import sys
+from pathlib import Path
+
+
+ROOT = Path(__file__).resolve().parents[1]
+BENCH = ROOT / "benchmarks" / "benchmark.py"
+
+
+DEFAULT_SHAPES = [
+    {"name": "pi052_b4_p700_k5_c50", "batch": 4, "prefix_len": 700, "action_blocks": 5, "action_block_size": 50},
+    {"name": "pi052_b2_p700_k5_c50", "batch": 2, "prefix_len": 700, "action_blocks": 5, "action_block_size": 50},
+    {"name": "pi052_b1_p700_k5_c50", "batch": 1, "prefix_len": 700, "action_blocks": 5, "action_block_size": 50},
+    {"name": "pi052_b4_p512_k5_c50", "batch": 4, "prefix_len": 512, "action_blocks": 5, "action_block_size": 50},
+    {"name": "pi052_b4_p896_k5_c50", "batch": 4, "prefix_len": 896, "action_blocks": 5, "action_block_size": 50},
+    {"name": "pi052_b4_p700_k1_c50", "batch": 4, "prefix_len": 700, "action_blocks": 1, "action_block_size": 50},
+    {"name": "pi052_b4_p700_k8_c50", "batch": 4, "prefix_len": 700, "action_blocks": 8, "action_block_size": 50},
+]
+
+
+def parse_presets(text: str) -> list[str]:
+    if text == "a100":
+        return [
+            "default",
+            "torch_default_explicit",
+            "a100_d256_bwd_32x64",
+            "a100_d256_bwd_32x128",
+            "a100_d256_bwd_64x64",
+            "a100_d256_contig_safe",
+            "a100_d256_contig_prescale",
+            "a100_d256_contig_write_dq_false",
+        ]
+    if text == "consumer":
+        return [
+            "default",
+            "torch_default_explicit",
+            "a100_d256_bwd_32x64",
+            "a100_d256_bwd_64x128",
+            "a100_d256_contig_safe",
+            "a100_d256_contig_prescale",
+        ]
+    return [x.strip() for x in text.split(",") if x.strip()]
+
+
+def parse_block_sizes(text: str) -> list[tuple[int, int]]:
+    if text == "default":
+        return [(128, 128)]
+    if text == "a100":
+        return [(64, 64), (64, 128), (128, 64), (128, 128)]
+    if text == "full":
+        return [(q, kv) for q in (16, 32, 64, 128) for kv in (32, 64, 128)]
+    out = []
+    for item in text.split(","):
+        item = item.strip().lower()
+        if not item:
+            continue
+        q, kv = item.split("x", 1)
+        out.append((int(q), int(kv)))
+    return out
+
+
+def load_shapes(path: str | None) -> list[dict]:
+    if path is None:
+        return DEFAULT_SHAPES
+    data = json.loads(Path(path).read_text(encoding="utf-8"))
+    if not isinstance(data, list):
+        raise SystemExit("shape file must be a JSON list")
+    return data
+
+
+def extract_json(stdout: str) -> dict:
+    start = stdout.find("{")
+    end = stdout.rfind("}")
+    if start < 0 or end < start:
+        return {"raw_stdout": stdout}
+    return json.loads(stdout[start : end + 1])
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--dtype", choices=["bf16", "fp32"], default="bf16")
+    parser.add_argument("--heads", type=int, default=8)
+    parser.add_argument("--kv-heads", type=int, default=None)
+    parser.add_argument("--head-dim", type=int, default=256)
+    parser.add_argument("--warmup", type=int, default=5)
+    parser.add_argument("--iters", type=int, default=10)
+    parser.add_argument("--mode", choices=["fwd", "fwdbwd", "all"], default="all")
+    parser.add_argument("--backend", default="torch-flex", help="comma-separated subset of {package, torch-flex, manual} or 'all'")
+    parser.add_argument("--presets", default="consumer", help="'consumer', 'a100', or comma-separated preset names")
+    parser.add_argument("--block-mask-sizes", default="default", help="'default', 'a100', 'full', or comma-separated QxKV sizes")
+    parser.add_argument("--shapes-json")
+    parser.add_argument("--output", default=str(ROOT / "benchmarks" / "matrix_results.jsonl"))
+    parser.add_argument("--fail-fast", action="store_true")
+    args = parser.parse_args()
+
+    shapes = load_shapes(args.shapes_json)
+    presets = parse_presets(args.presets)
+    block_sizes = parse_block_sizes(args.block_mask_sizes)
+    output = Path(args.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    output.write_text("", encoding="utf-8")
+
+    for shape in shapes:
+        for preset in presets:
+            for block_q, block_kv in block_sizes:
+                cmd = [
+                    sys.executable,
+                    str(BENCH),
+                    "--device",
+                    args.device,
+                    "--dtype",
+                    args.dtype,
+                    "--batch",
+                    str(shape["batch"]),
+                    "--heads",
+                    str(args.heads),
+                    "--kv-heads",
+                    str(args.kv_heads if args.kv_heads is not None else args.heads),
+                    "--head-dim",
+                    str(args.head_dim),
+                    "--prefix-len",
+                    str(shape["prefix_len"]),
+                    "--action-blocks",
+                    str(shape["action_blocks"]),
+                    "--action-block-size",
+                    str(shape["action_block_size"]),
+                    "--warmup",
+                    str(args.warmup),
+                    "--iters",
+                    str(args.iters),
+                    "--mode",
+                    args.mode,
+                    "--backend",
+                    args.backend,
+                    "--flex-preset",
+                    preset,
+                    "--block-mask-q",
+                    str(block_q),
+                    "--block-mask-kv",
+                    str(block_kv),
+                ]
+                proc = subprocess.run(cmd, cwd=ROOT.parent, text=True, capture_output=True)
+                row = {
+                    "shape_name": shape.get("name", ""),
+                    "preset": preset,
+                    "block_mask": f"{block_q}x{block_kv}",
+                    "returncode": proc.returncode,
+                }
+                if proc.returncode == 0:
+                    row.update(extract_json(proc.stdout))
+                else:
+                    row["stdout"] = proc.stdout[-4000:]
+                    row["stderr"] = proc.stderr[-4000:]
+                with output.open("a", encoding="utf-8") as f:
+                    f.write(json.dumps(row, sort_keys=True) + "\n")
+                print(json.dumps(row, sort_keys=True))
+                if proc.returncode != 0 and args.fail_fast:
+                    return proc.returncode
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

build.toml

@@ -0,0 +1,28 @@
+[general]
+name = "flashrt-flex-attention-train"
+license = "Apache-2.0"
+version = 1
+backends = ["cuda"]
+
+[general.cuda]
+minver = "12.8"
+
+[general.hub]
+repo-id = "flashrt/flashrt-flex-attention-train"
+
+[torch]
+include = ["csrc"]
+src = [
+  "torch-ext/torch_binding.cpp",
+  "torch-ext/torch_binding.h",
+]
+
+[kernel.flashrt_flex_attention_train_stub]
+backend = "cuda"
+depends = ["torch"]
+include = ["csrc"]
+cuda-minver = "12.8"
+src = [
+  "csrc/stub.cu",
+  "csrc/stub.cuh",
+]

build/torch-universal/flashrt_flex_attention_train/__init__.py

@@ -0,0 +1,492 @@
+"""FlashRT Flex-style block-sparse attention training API.
+
+The public function implements the PI052 prefix/action mask pattern:
+
+* prefix query rows use the original K/V tensors, so prefix losses keep normal
+  gradients into prefix K/V;
+* action query rows read detached prefix K/V plus normal action K/V by default,
+  matching the current training semantics.
+
+Unsupported shapes route to the SDPA reference path. Native CUDA kernels are
+not exposed until a shape-specialized implementation beats SDPA on the target
+A100/5090 validation matrix.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+
+try:
+    from ._ops import ops
+
+    _HAS_OPS = hasattr(ops, "_flashrt_training_package_marker")
+except Exception:  # source-tree tests before kernel-builder creates _ops.py
+    ops = None
+    _HAS_OPS = False
+
+
+MASK_VALUE_F32 = -2.3819763e38
+
+
+def _use_ops(namespace_ops) -> None:
+    """Install a manually built extension (dev/testing path)."""
+    global ops, _HAS_OPS
+    ops = namespace_ops
+    _HAS_OPS = hasattr(ops, "_flashrt_training_package_marker")
+
+
+def _check_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None:
+    if q.dim() != 4 or k.dim() != 4 or v.dim() != 4:
+        raise ValueError("q, k, and v must be shaped (B, H, S, D)")
+    if q.shape[0] != k.shape[0] or q.shape[0] != v.shape[0]:
+        raise ValueError("q, k, and v batch dimensions must match")
+    if k.shape != v.shape:
+        raise ValueError("k and v shapes must match")
+    if q.shape[2] != k.shape[2] or q.shape[3] != k.shape[3]:
+        raise ValueError("q, k, and v sequence/head_dim dimensions must match")
+    if q.device != k.device or q.device != v.device:
+        raise ValueError("q, k, and v must be on the same device")
+
+
+def _as_valid(mask: Optional[torch.Tensor], batch: int, length: int, device: torch.device) -> torch.Tensor:
+    if mask is None:
+        return torch.ones((batch, length), dtype=torch.bool, device=device)
+    if mask.shape != (batch, length):
+        raise ValueError(f"mask must be shaped {(batch, length)}, got {tuple(mask.shape)}")
+    return mask.to(device=device, dtype=torch.bool)
+
+
+def build_block_sparse_bool_masks(
+    prefix_valid: Optional[torch.Tensor],
+    prefix_att: Optional[torch.Tensor],
+    *,
+    batch: int,
+    prefix_len: int,
+    action_len: int,
+    action_block_size: int,
+    non_fast_prefix_len: Optional[int] = None,
+    action_valid: Optional[torch.Tensor] = None,
+    device: Optional[torch.device] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Build boolean masks for the split FlexAttention SDPA calls.
+
+    Returns ``(prefix_rows, action_rows)`` with shapes ``(B, P, S)`` and
+    ``(B, A, S)``. Boolean True means the key/value position is visible.
+
+    ``prefix_att`` follows Lerobot's cumulative-block convention: prefix key
+    ``j`` is visible to prefix query ``i`` when ``cumsum(prefix_att)[j] <=
+    cumsum(prefix_att)[i]`` and both rows are valid. When omitted, prefix rows
+    attend to all valid prefix tokens.
+    """
+    if action_block_size <= 0:
+        raise ValueError("action_block_size must be positive")
+    if prefix_len < 0 or action_len < 0:
+        raise ValueError("prefix_len and action_len must be non-negative")
+    total_len = prefix_len + action_len
+    dev = device
+    if dev is None:
+        for t in (prefix_valid, prefix_att, action_valid):
+            if t is not None:
+                dev = t.device
+                break
+    if dev is None:
+        dev = torch.device("cpu")
+
+    p_valid = _as_valid(prefix_valid, batch, prefix_len, dev)
+    a_valid = _as_valid(action_valid, batch, action_len, dev)
+
+    if prefix_att is None:
+        prefix_rows = p_valid[:, :, None] & p_valid[:, None, :]
+    else:
+        if prefix_att.shape != (batch, prefix_len):
+            raise ValueError(
+                f"prefix_att must be shaped {(batch, prefix_len)}, got {tuple(prefix_att.shape)}"
+            )
+        cum = torch.cumsum(prefix_att.to(device=dev, dtype=torch.long), dim=1)
+        prefix_rows = (cum[:, None, :] <= cum[:, :, None]) & p_valid[:, :, None] & p_valid[:, None, :]
+
+    prefix_pad = torch.zeros((batch, prefix_len, action_len), dtype=torch.bool, device=dev)
+    prefix_rows = torch.cat([prefix_rows, prefix_pad], dim=2)
+
+    nf = prefix_len if non_fast_prefix_len is None else int(non_fast_prefix_len)
+    nf = max(0, min(nf, prefix_len))
+    action_to_prefix = torch.zeros((batch, action_len, prefix_len), dtype=torch.bool, device=dev)
+    if nf > 0:
+        action_to_prefix[:, :, :nf] = p_valid[:, None, :nf]
+    action_to_prefix &= a_valid[:, :, None]
+
+    q_block = torch.arange(action_len, device=dev) // int(action_block_size)
+    kv_block = q_block
+    action_block = q_block[:, None] == kv_block[None, :]
+    action_block = action_block[None, :, :].expand(batch, -1, -1)
+    action_block = action_block & a_valid[:, :, None] & a_valid[:, None, :]
+    action_rows = torch.cat([action_to_prefix, action_block], dim=2)
+
+    if prefix_rows.shape != (batch, prefix_len, total_len):
+        raise AssertionError("internal prefix mask shape error")
+    if action_rows.shape != (batch, action_len, total_len):
+        raise AssertionError("internal action mask shape error")
+    return prefix_rows, action_rows
+
+
+def _bool_to_sdpa_mask(mask: torch.Tensor, q: torch.Tensor) -> torch.Tensor:
+    value = MASK_VALUE_F32
+    if q.dtype.is_floating_point:
+        finfo = torch.finfo(q.dtype)
+        value = max(MASK_VALUE_F32, finfo.min)
+    return torch.where(
+        mask[:, None, :, :],
+        torch.zeros((), dtype=q.dtype, device=q.device),
+        torch.full((), value, dtype=q.dtype, device=q.device),
+    )
+
+
+def _slice_attention_mask(
+    attention_mask: torch.Tensor,
+    start: int,
+    end: int,
+    q: torch.Tensor,
+) -> torch.Tensor:
+    if attention_mask.dim() == 3:
+        mask = attention_mask[:, start:end, :]
+        if mask.dtype == torch.bool:
+            return mask[:, None, :, :]
+        return mask[:, None, :, :].to(dtype=q.dtype)
+    if attention_mask.dim() == 4:
+        mask = attention_mask[:, :, start:end, :]
+        return mask if mask.dtype == torch.bool else mask.to(dtype=q.dtype)
+    raise ValueError("attention_mask must be (B, S, S) or (B, 1|H, S, S)")
+
+
+def _sdpa(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    mask: Optional[torch.Tensor],
+    *,
+    scale: Optional[float],
+    dropout_p: float,
+    enable_gqa: bool,
+) -> torch.Tensor:
+    kwargs = {"attn_mask": mask, "dropout_p": float(dropout_p), "scale": scale}
+    if enable_gqa:
+        kwargs["enable_gqa"] = True
+    try:
+        return F.scaled_dot_product_attention(q, k, v, **kwargs)
+    except TypeError:
+        kwargs.pop("enable_gqa", None)
+        return F.scaled_dot_product_attention(q, k, v, **kwargs)
+
+
+def reference_flex_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    prefix_len: int,
+    action_block_size: int,
+    attention_mask: Optional[torch.Tensor] = None,
+    prefix_valid: Optional[torch.Tensor] = None,
+    prefix_att: Optional[torch.Tensor] = None,
+    non_fast_prefix_len: Optional[int] = None,
+    action_valid: Optional[torch.Tensor] = None,
+    detach_prefix_kv_for_action: bool = True,
+    scale: Optional[float] = None,
+    dropout_p: float = 0.0,
+    enable_gqa: Optional[bool] = None,
+) -> torch.Tensor:
+    """SDPA reference for the PI052 FlexAttention replacement shape.
+
+    Args:
+        q, k, v: ``(B, Hq/Hkv, S, D)`` tensors.
+        prefix_len: number of prefix rows/columns at the start of sequence.
+        action_block_size: size of each block-diagonal action segment.
+        attention_mask: optional prebuilt additive or boolean full mask.
+        prefix_valid: optional ``(B, P)`` valid prefix positions.
+        prefix_att: optional ``(B, P)`` cumulative-block markers.
+        non_fast_prefix_len: prefix columns visible to action rows.
+        action_valid: optional ``(B, A)`` valid action positions.
+        detach_prefix_kv_for_action: detach prefix K/V on the action-row path.
+        scale: SDPA scale. Defaults to ``D ** -0.5``.
+        dropout_p: SDPA dropout probability.
+        enable_gqa: pass SDPA GQA mode when q heads and kv heads differ.
+    """
+    _check_qkv(q, k, v)
+    batch, _, total_len, head_dim = q.shape
+    if not (0 <= int(prefix_len) <= total_len):
+        raise ValueError("prefix_len must be in [0, S]")
+    prefix_len = int(prefix_len)
+    action_len = total_len - prefix_len
+    if scale is None:
+        scale = head_dim**-0.5
+    if enable_gqa is None:
+        enable_gqa = q.shape[1] != k.shape[1]
+
+    q_prefix = q[:, :, :prefix_len, :]
+    q_action = q[:, :, prefix_len:, :]
+    k_prefix = k[:, :, :prefix_len, :]
+    k_action = k[:, :, prefix_len:, :]
+    v_prefix = v[:, :, :prefix_len, :]
+    v_action = v[:, :, prefix_len:, :]
+
+    if attention_mask is None:
+        prefix_bool, action_bool = build_block_sparse_bool_masks(
+            prefix_valid,
+            prefix_att,
+            batch=batch,
+            prefix_len=prefix_len,
+            action_len=action_len,
+            action_block_size=action_block_size,
+            non_fast_prefix_len=non_fast_prefix_len,
+            action_valid=action_valid,
+            device=q.device,
+        )
+        prefix_mask = _bool_to_sdpa_mask(prefix_bool, q)
+        action_mask = _bool_to_sdpa_mask(action_bool, q)
+    else:
+        prefix_mask = _slice_attention_mask(attention_mask, 0, prefix_len, q)
+        action_mask = _slice_attention_mask(attention_mask, prefix_len, total_len, q)
+
+    out_parts = []
+    if prefix_len:
+        out_parts.append(
+            _sdpa(
+                q_prefix,
+                k,
+                v,
+                prefix_mask,
+                scale=scale,
+                dropout_p=dropout_p,
+                enable_gqa=bool(enable_gqa),
+            )
+        )
+    if action_len:
+        prefix_k = k_prefix.detach() if detach_prefix_kv_for_action else k_prefix
+        prefix_v = v_prefix.detach() if detach_prefix_kv_for_action else v_prefix
+        k_for_action = torch.cat([prefix_k, k_action], dim=2)
+        v_for_action = torch.cat([prefix_v, v_action], dim=2)
+        out_parts.append(
+            _sdpa(
+                q_action,
+                k_for_action,
+                v_for_action,
+                action_mask,
+                scale=scale,
+                dropout_p=dropout_p,
+                enable_gqa=bool(enable_gqa),
+            )
+        )
+    if not out_parts:
+        return q.new_empty(q.shape)
+    return torch.cat(out_parts, dim=2) if len(out_parts) == 2 else out_parts[0]
+
+
+def _manual_attention_part(qs, ks, vs, mask, scale):
+    """Materialized-logits attention part: cuBLAS GEMMs + fused masked softmax.
+
+    Same math as SDPA with an additive mask (fp32 softmax; logits stored in
+    the io dtype between the GEMM and the softmax). Grouped queries run as a
+    strided batched GEMM over the KV heads, so a 1-head K/V is never
+    repeated. At PI052 training shapes (GQA 8:1, D=256, bf16) this beats
+    both SDPA-with-dense-mask (2.3-3.1x) and the best FlexAttention
+    configuration (1.4-2.9x) on fwd+bwd — see benchmarks/RESULTS.md.
+    """
+    B, H, Sq, D = qs.shape
+    Hk = ks.shape[1]
+    if Hk != H:
+        g = H // Hk
+        q2 = qs.reshape(B, Hk, g * Sq, D)
+        logits = (q2 @ ks.transpose(-1, -2)).reshape(B, H, Sq, -1)
+    else:
+        logits = qs @ ks.transpose(-1, -2)
+    logits = logits * scale
+    if mask is not None:
+        logits = logits + mask
+    p = logits.float().softmax(dim=-1).to(qs.dtype)
+    if Hk != H:
+        out = (p.reshape(B, Hk, g * Sq, -1) @ vs).reshape(B, H, Sq, D)
+    else:
+        out = p @ vs
+    return out
+
+
+# Public alias: integrations (e.g. the LeRobot pi052 flag) consume the raw
+# per-part op and assemble masks/splits themselves.
+manual_attention_part = _manual_attention_part
+
+_manual_part_compiled = None
+
+
+def _get_manual_part():
+    global _manual_part_compiled
+    if _manual_part_compiled is None:
+        _manual_part_compiled = torch.compile(_manual_attention_part, dynamic=False)
+    return _manual_part_compiled
+
+
+def manual_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    prefix_len: int,
+    action_block_size: int,
+    attention_mask: Optional[torch.Tensor] = None,
+    prefix_valid: Optional[torch.Tensor] = None,
+    prefix_att: Optional[torch.Tensor] = None,
+    non_fast_prefix_len: Optional[int] = None,
+    action_valid: Optional[torch.Tensor] = None,
+    detach_prefix_kv_for_action: bool = True,
+    scale: Optional[float] = None,
+    dropout_p: float = 0.0,
+    compile_part: bool = True,
+) -> torch.Tensor:
+    """Materialized-logits implementation of :func:`reference_flex_attention`.
+
+    Same mask semantics and prefix/action split; each part runs through
+    :func:`_manual_attention_part` instead of SDPA. ``dropout_p`` must be 0
+    (training attention dropout is unused in PI052); other values raise so
+    callers fall back explicitly.
+    """
+    if dropout_p:
+        raise ValueError("manual_attention does not support dropout; use the reference path")
+    _check_qkv(q, k, v)
+    batch, _, total_len, head_dim = q.shape
+    if not (0 <= int(prefix_len) <= total_len):
+        raise ValueError("prefix_len must be in [0, S]")
+    prefix_len = int(prefix_len)
+    action_len = total_len - prefix_len
+    if scale is None:
+        scale = head_dim**-0.5
+
+    if attention_mask is None:
+        prefix_bool, action_bool = build_block_sparse_bool_masks(
+            prefix_valid,
+            prefix_att,
+            batch=batch,
+            prefix_len=prefix_len,
+            action_len=action_len,
+            action_block_size=action_block_size,
+            non_fast_prefix_len=non_fast_prefix_len,
+            action_valid=action_valid,
+            device=q.device,
+        )
+        prefix_mask = _bool_to_sdpa_mask(prefix_bool, q)
+        action_mask = _bool_to_sdpa_mask(action_bool, q)
+    else:
+        prefix_mask = _slice_attention_mask(attention_mask, 0, prefix_len, q)
+        action_mask = _slice_attention_mask(attention_mask, prefix_len, total_len, q)
+        if prefix_mask.dtype == torch.bool:
+            prefix_mask = _bool_to_sdpa_mask(prefix_mask[:, 0], q)
+        if action_mask.dtype == torch.bool:
+            action_mask = _bool_to_sdpa_mask(action_mask[:, 0], q)
+
+    part = _get_manual_part() if compile_part else _manual_attention_part
+    out_parts = []
+    if prefix_len:
+        out_parts.append(part(q[:, :, :prefix_len, :], k, v, prefix_mask, scale))
+    if action_len:
+        k_prefix = k[:, :, :prefix_len, :]
+        v_prefix = v[:, :, :prefix_len, :]
+        if detach_prefix_kv_for_action:
+            k_prefix = k_prefix.detach()
+            v_prefix = v_prefix.detach()
+        k_for_action = torch.cat([k_prefix, k[:, :, prefix_len:, :]], dim=2)
+        v_for_action = torch.cat([v_prefix, v[:, :, prefix_len:, :]], dim=2)
+        out_parts.append(part(q[:, :, prefix_len:, :], k_for_action, v_for_action, action_mask, scale))
+    if not out_parts:
+        return q.new_empty(q.shape)
+    return torch.cat(out_parts, dim=2) if len(out_parts) == 2 else out_parts[0]
+
+
+def flex_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    prefix_len: int,
+    action_block_size: int,
+    attention_mask: Optional[torch.Tensor] = None,
+    prefix_valid: Optional[torch.Tensor] = None,
+    prefix_att: Optional[torch.Tensor] = None,
+    non_fast_prefix_len: Optional[int] = None,
+    action_valid: Optional[torch.Tensor] = None,
+    detach_prefix_kv_for_action: bool = True,
+    scale: Optional[float] = None,
+    dropout_p: float = 0.0,
+    enable_gqa: Optional[bool] = None,
+    force_fallback: bool = False,
+    impl: str = "sdpa",
+) -> torch.Tensor:
+    """Flex-style block-sparse attention.
+
+    ``impl="sdpa"`` (default) keeps the SDPA reference path;
+    ``impl="manual"`` routes through the materialized-logits
+    implementation; ``impl="auto"`` picks manual only where it has been
+    measured to win end-to-end — consumer Blackwell (sm120-class) with
+    no dropout. On A100 (sm80) the manual math wins microbenches but
+    loses training-step integration, and on H100/H200 (sm90) the fused
+    FMHA kernels win outright, so auto keeps SDPA there.
+    """
+    _ = force_fallback
+    if impl == "auto":
+        sm120 = q.is_cuda and torch.cuda.get_device_capability(q.device)[0] == 12
+        impl = "manual" if (sm120 and not dropout_p) else "sdpa"
+    if impl == "manual":
+        return manual_attention(
+            q,
+            k,
+            v,
+            prefix_len=prefix_len,
+            action_block_size=action_block_size,
+            attention_mask=attention_mask,
+            prefix_valid=prefix_valid,
+            prefix_att=prefix_att,
+            non_fast_prefix_len=non_fast_prefix_len,
+            action_valid=action_valid,
+            detach_prefix_kv_for_action=detach_prefix_kv_for_action,
+            scale=scale,
+            dropout_p=dropout_p,
+        )
+    return reference_flex_attention(
+        q,
+        k,
+        v,
+        prefix_len=prefix_len,
+        action_block_size=action_block_size,
+        attention_mask=attention_mask,
+        prefix_valid=prefix_valid,
+        prefix_att=prefix_att,
+        non_fast_prefix_len=non_fast_prefix_len,
+        action_valid=action_valid,
+        detach_prefix_kv_for_action=detach_prefix_kv_for_action,
+        scale=scale,
+        dropout_p=dropout_p,
+        enable_gqa=enable_gqa,
+    )
+
+
+def flex_attention_forward(*args, **kwargs) -> torch.Tensor:
+    """Forward-only compatibility wrapper."""
+    return flex_attention(*args, **kwargs)
+
+
+def backend_marker(x: torch.Tensor) -> torch.Tensor:
+    if ops is None:
+        return x
+    return ops._flashrt_training_package_marker(x)
+
+
+__all__ = [
+    "MASK_VALUE_F32",
+    "backend_marker",
+    "build_block_sparse_bool_masks",
+    "flex_attention",
+    "flex_attention_forward",
+    "manual_attention",
+    "manual_attention_part",
+    "reference_flex_attention",
+]

csrc/README.md

@@ -0,0 +1,7 @@
+# csrc
+
+CUDA extension placeholder for the FlexAttention replacement package.
+
+The v1 package currently exposes the stable Python Tensor API and SDPA-backed
+autograd fallback. Native CUDA forward/backward kernels should be added here
+without changing the public functions in `torch-ext/flashrt_flex_attention_train`.

csrc/stub.cu

@@ -0,0 +1,3 @@
+// SPDX-License-Identifier: Apache-2.0
+
+#include "stub.cuh"

csrc/stub.cuh

@@ -0,0 +1,3 @@
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once

flake.nix

@@ -0,0 +1,17 @@
+{
+  description = "Flake for FlashRT Flex attention training kernels";
+
+  inputs = {
+    kernel-builder.url = "github:huggingface/kernels/432702bfbfbb17d3a1bd2c2743d004e21e769ab7";
+  };
+
+  outputs =
+    {
+      self,
+      kernel-builder,
+    }:
+    kernel-builder.lib.genKernelFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}

tests/test_flashrt_flex_attention_train.py

@@ -0,0 +1,182 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import importlib
+import sys
+
+import torch
+import torch.nn.functional as F
+
+
+def load_ops(backend: str, artifact: str | None):
+    if backend == "source":
+        sys.path.insert(0, "flashrt-flex-attention-train/torch-ext")
+        try:
+            return importlib.import_module("flashrt_flex_attention_train")
+        finally:
+            sys.path.remove("flashrt-flex-attention-train/torch-ext")
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return importlib.import_module("flashrt_flex_attention_train")
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
+
+
+def _shape(mode: str):
+    if mode == "full" and torch.cuda.is_available():
+        return "cuda", torch.bfloat16, 2, 8, 13, 6, 256
+    return "cpu", torch.float32, 2, 2, 5, 4, 16
+
+
+def test_matches_explicit_sdpa(flex_ops, mode: str) -> None:
+    device, dtype, bsz, heads, prefix, action, dim = _shape(mode)
+    torch.manual_seed(11)
+    q = torch.randn(bsz, heads, prefix + action, dim, device=device, dtype=dtype, requires_grad=True)
+    k = torch.randn_like(q, requires_grad=True)
+    v = torch.randn_like(q, requires_grad=True)
+    prefix_valid = torch.ones(bsz, prefix, device=device, dtype=torch.bool)
+    prefix_valid[0, -1] = False
+    prefix_att = torch.zeros(bsz, prefix, device=device, dtype=torch.bool)
+    prefix_att[:, prefix // 2 :] = True
+    action_valid = torch.ones(bsz, action, device=device, dtype=torch.bool)
+    action_valid[1, -1] = False
+
+    out = flex_ops.flex_attention(
+        q,
+        k,
+        v,
+        prefix_len=prefix,
+        action_block_size=2,
+        prefix_valid=prefix_valid,
+        prefix_att=prefix_att,
+        action_valid=action_valid,
+        non_fast_prefix_len=prefix - 1,
+    )
+
+    pm, am = flex_ops.build_block_sparse_bool_masks(
+        prefix_valid,
+        prefix_att,
+        batch=bsz,
+        prefix_len=prefix,
+        action_len=action,
+        action_block_size=2,
+        non_fast_prefix_len=prefix - 1,
+        action_valid=action_valid,
+        device=q.device,
+    )
+    pm = torch.where(
+        pm[:, None],
+        torch.zeros((), device=q.device, dtype=q.dtype),
+        torch.full((), flex_ops.MASK_VALUE_F32, device=q.device, dtype=q.dtype),
+    )
+    am = torch.where(
+        am[:, None],
+        torch.zeros((), device=q.device, dtype=q.dtype),
+        torch.full((), flex_ops.MASK_VALUE_F32, device=q.device, dtype=q.dtype),
+    )
+    expected_p = F.scaled_dot_product_attention(q[:, :, :prefix], k, v, attn_mask=pm, scale=dim**-0.5)
+    kd = torch.cat([k[:, :, :prefix].detach(), k[:, :, prefix:]], dim=2)
+    vd = torch.cat([v[:, :, :prefix].detach(), v[:, :, prefix:]], dim=2)
+    expected_a = F.scaled_dot_product_attention(q[:, :, prefix:], kd, vd, attn_mask=am, scale=dim**-0.5)
+    expected = torch.cat([expected_p, expected_a], dim=2)
+    tol = 2e-3 if dtype == torch.bfloat16 else 1e-5
+    torch.testing.assert_close(out, expected, atol=tol, rtol=tol)
+
+
+def test_detached_prefix_semantics(flex_ops, mode: str) -> None:
+    device, dtype, bsz, heads, prefix, action, dim = _shape(mode)
+    torch.manual_seed(17)
+    q = torch.randn(bsz, heads, prefix + action, dim, device=device, dtype=dtype, requires_grad=True)
+    k = torch.randn_like(q, requires_grad=True)
+    v = torch.randn_like(q, requires_grad=True)
+    out = flex_ops.flex_attention(q, k, v, prefix_len=prefix, action_block_size=2)
+    out[:, :, prefix:].float().square().mean().backward()
+    assert torch.count_nonzero(k.grad[:, :, :prefix]) == 0
+    assert torch.count_nonzero(v.grad[:, :, :prefix]) == 0
+    assert torch.count_nonzero(k.grad[:, :, prefix:]) > 0
+    assert torch.count_nonzero(v.grad[:, :, prefix:]) > 0
+
+
+def test_dense_attention_mask_path(flex_ops, mode: str) -> None:
+    device, dtype, bsz, heads, prefix, action, dim = _shape(mode)
+    torch.manual_seed(23)
+    q = torch.randn(bsz, heads, prefix + action, dim, device=device, dtype=dtype)
+    k = torch.randn_like(q)
+    v = torch.randn_like(q)
+    pm, am = flex_ops.build_block_sparse_bool_masks(
+        None,
+        None,
+        batch=bsz,
+        prefix_len=prefix,
+        action_len=action,
+        action_block_size=2,
+        device=q.device,
+    )
+    full = torch.cat([pm, am], dim=1)
+    mask = torch.where(
+        full[:, None],
+        torch.zeros((), device=q.device, dtype=q.dtype),
+        torch.full((), flex_ops.MASK_VALUE_F32, device=q.device, dtype=q.dtype),
+    )
+    out_from_dense = flex_ops.flex_attention(q, k, v, prefix_len=prefix, action_block_size=2, attention_mask=mask)
+    out_from_parts = flex_ops.flex_attention(q, k, v, prefix_len=prefix, action_block_size=2)
+    tol = 2e-3 if dtype == torch.bfloat16 else 1e-5
+    torch.testing.assert_close(out_from_dense, out_from_parts, atol=tol, rtol=tol)
+
+
+def test_manual_matches_reference(flex_ops, mode: str) -> None:
+    device, dtype, bsz, heads, prefix, action, dim = _shape(mode)
+    kv_heads = 1 if mode == "full" else heads  # GQA on the real-shape run
+    torch.manual_seed(29)
+    total = prefix + action
+    q1 = torch.randn(bsz, heads, total, dim, device=device, dtype=dtype, requires_grad=True)
+    k1 = torch.randn(bsz, kv_heads, total, dim, device=device, dtype=dtype, requires_grad=True)
+    v1 = torch.randn_like(k1, requires_grad=True)
+    q2 = q1.detach().clone().requires_grad_(True)
+    k2 = k1.detach().clone().requires_grad_(True)
+    v2 = v1.detach().clone().requires_grad_(True)
+    prefix_att = torch.zeros(bsz, prefix, device=device, dtype=torch.bool)
+    prefix_att[:, prefix // 2 :] = True
+    kwargs = dict(prefix_len=prefix, action_block_size=2, prefix_att=prefix_att)
+
+    ref = flex_ops.flex_attention(q1, k1, v1, **kwargs)
+    got = flex_ops.manual_attention(q2, k2, v2, compile_part=device != "cpu", **kwargs)
+    # bf16-logits class: the manual path stores logits in the io dtype
+    # between the GEMM and the fp32 softmax.
+    tol = 2e-2 if dtype == torch.bfloat16 else 1e-5
+    torch.testing.assert_close(got, ref, atol=tol, rtol=tol)
+
+    ref.float().square().mean().backward()
+    got.float().square().mean().backward()
+    for a, b in ((q1, q2), (k1, k2), (v1, v2)):
+        denom = torch.linalg.vector_norm(a.grad.float()).clamp_min(1e-12)
+        rel = torch.linalg.vector_norm((a.grad - b.grad).float()) / denom
+        assert float(rel) <= 2e-2, f"grad rel diff {float(rel)}"
+
+    # impl dispatch reaches the same path
+    via_impl = flex_ops.flex_attention(
+        q2.detach(), k2.detach(), v2.detach(), impl="manual", **kwargs
+    )
+    torch.testing.assert_close(via_impl, got.detach(), atol=tol, rtol=tol)
+
+
+def run(flex_ops, mode: str) -> None:
+    test_matches_explicit_sdpa(flex_ops, mode)
+    test_detached_prefix_semantics(flex_ops, mode)
+    test_dense_attention_mask_path(flex_ops, mode)
+    test_manual_matches_reference(flex_ops, mode)
+    x = torch.ones(1)
+    torch.testing.assert_close(flex_ops.backend_marker(x), x)
+    print(f"flashrt-flex-attention-train {mode}: passed")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="source")
+    parser.add_argument("--artifact")
+    parser.add_argument("--mode", choices=["smoke", "full"], default="smoke")
+    args = parser.parse_args()
+    run(load_ops(args.backend, args.artifact), args.mode)

torch-ext/README.md

@@ -0,0 +1,7 @@
+# torch-ext
+
+PyTorch bindings and Python package for `flashrt-flex-attention-train`.
+
+`torch_binding.cpp` registers a package marker so HF `kernel-builder` can load
+the extension. The Python module provides the SDPA fallback and the stable
+training API used by PI052 integration.

torch-ext/flashrt_flex_attention_train/__init__.py

@@ -0,0 +1,492 @@
+"""FlashRT Flex-style block-sparse attention training API.
+
+The public function implements the PI052 prefix/action mask pattern:
+
+* prefix query rows use the original K/V tensors, so prefix losses keep normal
+  gradients into prefix K/V;
+* action query rows read detached prefix K/V plus normal action K/V by default,
+  matching the current training semantics.
+
+Unsupported shapes route to the SDPA reference path. Native CUDA kernels are
+not exposed until a shape-specialized implementation beats SDPA on the target
+A100/5090 validation matrix.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+
+try:
+    from ._ops import ops
+
+    _HAS_OPS = hasattr(ops, "_flashrt_training_package_marker")
+except Exception:  # source-tree tests before kernel-builder creates _ops.py
+    ops = None
+    _HAS_OPS = False
+
+
+MASK_VALUE_F32 = -2.3819763e38
+
+
+def _use_ops(namespace_ops) -> None:
+    """Install a manually built extension (dev/testing path)."""
+    global ops, _HAS_OPS
+    ops = namespace_ops
+    _HAS_OPS = hasattr(ops, "_flashrt_training_package_marker")
+
+
+def _check_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None:
+    if q.dim() != 4 or k.dim() != 4 or v.dim() != 4:
+        raise ValueError("q, k, and v must be shaped (B, H, S, D)")
+    if q.shape[0] != k.shape[0] or q.shape[0] != v.shape[0]:
+        raise ValueError("q, k, and v batch dimensions must match")
+    if k.shape != v.shape:
+        raise ValueError("k and v shapes must match")
+    if q.shape[2] != k.shape[2] or q.shape[3] != k.shape[3]:
+        raise ValueError("q, k, and v sequence/head_dim dimensions must match")
+    if q.device != k.device or q.device != v.device:
+        raise ValueError("q, k, and v must be on the same device")
+
+
+def _as_valid(mask: Optional[torch.Tensor], batch: int, length: int, device: torch.device) -> torch.Tensor:
+    if mask is None:
+        return torch.ones((batch, length), dtype=torch.bool, device=device)
+    if mask.shape != (batch, length):
+        raise ValueError(f"mask must be shaped {(batch, length)}, got {tuple(mask.shape)}")
+    return mask.to(device=device, dtype=torch.bool)
+
+
+def build_block_sparse_bool_masks(
+    prefix_valid: Optional[torch.Tensor],
+    prefix_att: Optional[torch.Tensor],
+    *,
+    batch: int,
+    prefix_len: int,
+    action_len: int,
+    action_block_size: int,
+    non_fast_prefix_len: Optional[int] = None,
+    action_valid: Optional[torch.Tensor] = None,
+    device: Optional[torch.device] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Build boolean masks for the split FlexAttention SDPA calls.
+
+    Returns ``(prefix_rows, action_rows)`` with shapes ``(B, P, S)`` and
+    ``(B, A, S)``. Boolean True means the key/value position is visible.
+
+    ``prefix_att`` follows Lerobot's cumulative-block convention: prefix key
+    ``j`` is visible to prefix query ``i`` when ``cumsum(prefix_att)[j] <=
+    cumsum(prefix_att)[i]`` and both rows are valid. When omitted, prefix rows
+    attend to all valid prefix tokens.
+    """
+    if action_block_size <= 0:
+        raise ValueError("action_block_size must be positive")
+    if prefix_len < 0 or action_len < 0:
+        raise ValueError("prefix_len and action_len must be non-negative")
+    total_len = prefix_len + action_len
+    dev = device
+    if dev is None:
+        for t in (prefix_valid, prefix_att, action_valid):
+            if t is not None:
+                dev = t.device
+                break
+    if dev is None:
+        dev = torch.device("cpu")
+
+    p_valid = _as_valid(prefix_valid, batch, prefix_len, dev)
+    a_valid = _as_valid(action_valid, batch, action_len, dev)
+
+    if prefix_att is None:
+        prefix_rows = p_valid[:, :, None] & p_valid[:, None, :]
+    else:
+        if prefix_att.shape != (batch, prefix_len):
+            raise ValueError(
+                f"prefix_att must be shaped {(batch, prefix_len)}, got {tuple(prefix_att.shape)}"
+            )
+        cum = torch.cumsum(prefix_att.to(device=dev, dtype=torch.long), dim=1)
+        prefix_rows = (cum[:, None, :] <= cum[:, :, None]) & p_valid[:, :, None] & p_valid[:, None, :]
+
+    prefix_pad = torch.zeros((batch, prefix_len, action_len), dtype=torch.bool, device=dev)
+    prefix_rows = torch.cat([prefix_rows, prefix_pad], dim=2)
+
+    nf = prefix_len if non_fast_prefix_len is None else int(non_fast_prefix_len)
+    nf = max(0, min(nf, prefix_len))
+    action_to_prefix = torch.zeros((batch, action_len, prefix_len), dtype=torch.bool, device=dev)
+    if nf > 0:
+        action_to_prefix[:, :, :nf] = p_valid[:, None, :nf]
+    action_to_prefix &= a_valid[:, :, None]
+
+    q_block = torch.arange(action_len, device=dev) // int(action_block_size)
+    kv_block = q_block
+    action_block = q_block[:, None] == kv_block[None, :]
+    action_block = action_block[None, :, :].expand(batch, -1, -1)
+    action_block = action_block & a_valid[:, :, None] & a_valid[:, None, :]
+    action_rows = torch.cat([action_to_prefix, action_block], dim=2)
+
+    if prefix_rows.shape != (batch, prefix_len, total_len):
+        raise AssertionError("internal prefix mask shape error")
+    if action_rows.shape != (batch, action_len, total_len):
+        raise AssertionError("internal action mask shape error")
+    return prefix_rows, action_rows
+
+
+def _bool_to_sdpa_mask(mask: torch.Tensor, q: torch.Tensor) -> torch.Tensor:
+    value = MASK_VALUE_F32
+    if q.dtype.is_floating_point:
+        finfo = torch.finfo(q.dtype)
+        value = max(MASK_VALUE_F32, finfo.min)
+    return torch.where(
+        mask[:, None, :, :],
+        torch.zeros((), dtype=q.dtype, device=q.device),
+        torch.full((), value, dtype=q.dtype, device=q.device),
+    )
+
+
+def _slice_attention_mask(
+    attention_mask: torch.Tensor,
+    start: int,
+    end: int,
+    q: torch.Tensor,
+) -> torch.Tensor:
+    if attention_mask.dim() == 3:
+        mask = attention_mask[:, start:end, :]
+        if mask.dtype == torch.bool:
+            return mask[:, None, :, :]
+        return mask[:, None, :, :].to(dtype=q.dtype)
+    if attention_mask.dim() == 4:
+        mask = attention_mask[:, :, start:end, :]
+        return mask if mask.dtype == torch.bool else mask.to(dtype=q.dtype)
+    raise ValueError("attention_mask must be (B, S, S) or (B, 1|H, S, S)")
+
+
+def _sdpa(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    mask: Optional[torch.Tensor],
+    *,
+    scale: Optional[float],
+    dropout_p: float,
+    enable_gqa: bool,
+) -> torch.Tensor:
+    kwargs = {"attn_mask": mask, "dropout_p": float(dropout_p), "scale": scale}
+    if enable_gqa:
+        kwargs["enable_gqa"] = True
+    try:
+        return F.scaled_dot_product_attention(q, k, v, **kwargs)
+    except TypeError:
+        kwargs.pop("enable_gqa", None)
+        return F.scaled_dot_product_attention(q, k, v, **kwargs)
+
+
+def reference_flex_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    prefix_len: int,
+    action_block_size: int,
+    attention_mask: Optional[torch.Tensor] = None,
+    prefix_valid: Optional[torch.Tensor] = None,
+    prefix_att: Optional[torch.Tensor] = None,
+    non_fast_prefix_len: Optional[int] = None,
+    action_valid: Optional[torch.Tensor] = None,
+    detach_prefix_kv_for_action: bool = True,
+    scale: Optional[float] = None,
+    dropout_p: float = 0.0,
+    enable_gqa: Optional[bool] = None,
+) -> torch.Tensor:
+    """SDPA reference for the PI052 FlexAttention replacement shape.
+
+    Args:
+        q, k, v: ``(B, Hq/Hkv, S, D)`` tensors.
+        prefix_len: number of prefix rows/columns at the start of sequence.
+        action_block_size: size of each block-diagonal action segment.
+        attention_mask: optional prebuilt additive or boolean full mask.
+        prefix_valid: optional ``(B, P)`` valid prefix positions.
+        prefix_att: optional ``(B, P)`` cumulative-block markers.
+        non_fast_prefix_len: prefix columns visible to action rows.
+        action_valid: optional ``(B, A)`` valid action positions.
+        detach_prefix_kv_for_action: detach prefix K/V on the action-row path.
+        scale: SDPA scale. Defaults to ``D ** -0.5``.
+        dropout_p: SDPA dropout probability.
+        enable_gqa: pass SDPA GQA mode when q heads and kv heads differ.
+    """
+    _check_qkv(q, k, v)
+    batch, _, total_len, head_dim = q.shape
+    if not (0 <= int(prefix_len) <= total_len):
+        raise ValueError("prefix_len must be in [0, S]")
+    prefix_len = int(prefix_len)
+    action_len = total_len - prefix_len
+    if scale is None:
+        scale = head_dim**-0.5
+    if enable_gqa is None:
+        enable_gqa = q.shape[1] != k.shape[1]
+
+    q_prefix = q[:, :, :prefix_len, :]
+    q_action = q[:, :, prefix_len:, :]
+    k_prefix = k[:, :, :prefix_len, :]
+    k_action = k[:, :, prefix_len:, :]
+    v_prefix = v[:, :, :prefix_len, :]
+    v_action = v[:, :, prefix_len:, :]
+
+    if attention_mask is None:
+        prefix_bool, action_bool = build_block_sparse_bool_masks(
+            prefix_valid,
+            prefix_att,
+            batch=batch,
+            prefix_len=prefix_len,
+            action_len=action_len,
+            action_block_size=action_block_size,
+            non_fast_prefix_len=non_fast_prefix_len,
+            action_valid=action_valid,
+            device=q.device,
+        )
+        prefix_mask = _bool_to_sdpa_mask(prefix_bool, q)
+        action_mask = _bool_to_sdpa_mask(action_bool, q)
+    else:
+        prefix_mask = _slice_attention_mask(attention_mask, 0, prefix_len, q)
+        action_mask = _slice_attention_mask(attention_mask, prefix_len, total_len, q)
+
+    out_parts = []
+    if prefix_len:
+        out_parts.append(
+            _sdpa(
+                q_prefix,
+                k,
+                v,
+                prefix_mask,
+                scale=scale,
+                dropout_p=dropout_p,
+                enable_gqa=bool(enable_gqa),
+            )
+        )
+    if action_len:
+        prefix_k = k_prefix.detach() if detach_prefix_kv_for_action else k_prefix
+        prefix_v = v_prefix.detach() if detach_prefix_kv_for_action else v_prefix
+        k_for_action = torch.cat([prefix_k, k_action], dim=2)
+        v_for_action = torch.cat([prefix_v, v_action], dim=2)
+        out_parts.append(
+            _sdpa(
+                q_action,
+                k_for_action,
+                v_for_action,
+                action_mask,
+                scale=scale,
+                dropout_p=dropout_p,
+                enable_gqa=bool(enable_gqa),
+            )
+        )
+    if not out_parts:
+        return q.new_empty(q.shape)
+    return torch.cat(out_parts, dim=2) if len(out_parts) == 2 else out_parts[0]
+
+
+def _manual_attention_part(qs, ks, vs, mask, scale):
+    """Materialized-logits attention part: cuBLAS GEMMs + fused masked softmax.
+
+    Same math as SDPA with an additive mask (fp32 softmax; logits stored in
+    the io dtype between the GEMM and the softmax). Grouped queries run as a
+    strided batched GEMM over the KV heads, so a 1-head K/V is never
+    repeated. At PI052 training shapes (GQA 8:1, D=256, bf16) this beats
+    both SDPA-with-dense-mask (2.3-3.1x) and the best FlexAttention
+    configuration (1.4-2.9x) on fwd+bwd — see benchmarks/RESULTS.md.
+    """
+    B, H, Sq, D = qs.shape
+    Hk = ks.shape[1]
+    if Hk != H:
+        g = H // Hk
+        q2 = qs.reshape(B, Hk, g * Sq, D)
+        logits = (q2 @ ks.transpose(-1, -2)).reshape(B, H, Sq, -1)
+    else:
+        logits = qs @ ks.transpose(-1, -2)
+    logits = logits * scale
+    if mask is not None:
+        logits = logits + mask
+    p = logits.float().softmax(dim=-1).to(qs.dtype)
+    if Hk != H:
+        out = (p.reshape(B, Hk, g * Sq, -1) @ vs).reshape(B, H, Sq, D)
+    else:
+        out = p @ vs
+    return out
+
+
+# Public alias: integrations (e.g. the LeRobot pi052 flag) consume the raw
+# per-part op and assemble masks/splits themselves.
+manual_attention_part = _manual_attention_part
+
+_manual_part_compiled = None
+
+
+def _get_manual_part():
+    global _manual_part_compiled
+    if _manual_part_compiled is None:
+        _manual_part_compiled = torch.compile(_manual_attention_part, dynamic=False)
+    return _manual_part_compiled
+
+
+def manual_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    prefix_len: int,
+    action_block_size: int,
+    attention_mask: Optional[torch.Tensor] = None,
+    prefix_valid: Optional[torch.Tensor] = None,
+    prefix_att: Optional[torch.Tensor] = None,
+    non_fast_prefix_len: Optional[int] = None,
+    action_valid: Optional[torch.Tensor] = None,
+    detach_prefix_kv_for_action: bool = True,
+    scale: Optional[float] = None,
+    dropout_p: float = 0.0,
+    compile_part: bool = True,
+) -> torch.Tensor:
+    """Materialized-logits implementation of :func:`reference_flex_attention`.
+
+    Same mask semantics and prefix/action split; each part runs through
+    :func:`_manual_attention_part` instead of SDPA. ``dropout_p`` must be 0
+    (training attention dropout is unused in PI052); other values raise so
+    callers fall back explicitly.
+    """
+    if dropout_p:
+        raise ValueError("manual_attention does not support dropout; use the reference path")
+    _check_qkv(q, k, v)
+    batch, _, total_len, head_dim = q.shape
+    if not (0 <= int(prefix_len) <= total_len):
+        raise ValueError("prefix_len must be in [0, S]")
+    prefix_len = int(prefix_len)
+    action_len = total_len - prefix_len
+    if scale is None:
+        scale = head_dim**-0.5
+
+    if attention_mask is None:
+        prefix_bool, action_bool = build_block_sparse_bool_masks(
+            prefix_valid,
+            prefix_att,
+            batch=batch,
+            prefix_len=prefix_len,
+            action_len=action_len,
+            action_block_size=action_block_size,
+            non_fast_prefix_len=non_fast_prefix_len,
+            action_valid=action_valid,
+            device=q.device,
+        )
+        prefix_mask = _bool_to_sdpa_mask(prefix_bool, q)
+        action_mask = _bool_to_sdpa_mask(action_bool, q)
+    else:
+        prefix_mask = _slice_attention_mask(attention_mask, 0, prefix_len, q)
+        action_mask = _slice_attention_mask(attention_mask, prefix_len, total_len, q)
+        if prefix_mask.dtype == torch.bool:
+            prefix_mask = _bool_to_sdpa_mask(prefix_mask[:, 0], q)
+        if action_mask.dtype == torch.bool:
+            action_mask = _bool_to_sdpa_mask(action_mask[:, 0], q)
+
+    part = _get_manual_part() if compile_part else _manual_attention_part
+    out_parts = []
+    if prefix_len:
+        out_parts.append(part(q[:, :, :prefix_len, :], k, v, prefix_mask, scale))
+    if action_len:
+        k_prefix = k[:, :, :prefix_len, :]
+        v_prefix = v[:, :, :prefix_len, :]
+        if detach_prefix_kv_for_action:
+            k_prefix = k_prefix.detach()
+            v_prefix = v_prefix.detach()
+        k_for_action = torch.cat([k_prefix, k[:, :, prefix_len:, :]], dim=2)
+        v_for_action = torch.cat([v_prefix, v[:, :, prefix_len:, :]], dim=2)
+        out_parts.append(part(q[:, :, prefix_len:, :], k_for_action, v_for_action, action_mask, scale))
+    if not out_parts:
+        return q.new_empty(q.shape)
+    return torch.cat(out_parts, dim=2) if len(out_parts) == 2 else out_parts[0]
+
+
+def flex_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    prefix_len: int,
+    action_block_size: int,
+    attention_mask: Optional[torch.Tensor] = None,
+    prefix_valid: Optional[torch.Tensor] = None,
+    prefix_att: Optional[torch.Tensor] = None,
+    non_fast_prefix_len: Optional[int] = None,
+    action_valid: Optional[torch.Tensor] = None,
+    detach_prefix_kv_for_action: bool = True,
+    scale: Optional[float] = None,
+    dropout_p: float = 0.0,
+    enable_gqa: Optional[bool] = None,
+    force_fallback: bool = False,
+    impl: str = "sdpa",
+) -> torch.Tensor:
+    """Flex-style block-sparse attention.
+
+    ``impl="sdpa"`` (default) keeps the SDPA reference path;
+    ``impl="manual"`` routes through the materialized-logits
+    implementation; ``impl="auto"`` picks manual only where it has been
+    measured to win end-to-end — consumer Blackwell (sm120-class) with
+    no dropout. On A100 (sm80) the manual math wins microbenches but
+    loses training-step integration, and on H100/H200 (sm90) the fused
+    FMHA kernels win outright, so auto keeps SDPA there.
+    """
+    _ = force_fallback
+    if impl == "auto":
+        sm120 = q.is_cuda and torch.cuda.get_device_capability(q.device)[0] == 12
+        impl = "manual" if (sm120 and not dropout_p) else "sdpa"
+    if impl == "manual":
+        return manual_attention(
+            q,
+            k,
+            v,
+            prefix_len=prefix_len,
+            action_block_size=action_block_size,
+            attention_mask=attention_mask,
+            prefix_valid=prefix_valid,
+            prefix_att=prefix_att,
+            non_fast_prefix_len=non_fast_prefix_len,
+            action_valid=action_valid,
+            detach_prefix_kv_for_action=detach_prefix_kv_for_action,
+            scale=scale,
+            dropout_p=dropout_p,
+        )
+    return reference_flex_attention(
+        q,
+        k,
+        v,
+        prefix_len=prefix_len,
+        action_block_size=action_block_size,
+        attention_mask=attention_mask,
+        prefix_valid=prefix_valid,
+        prefix_att=prefix_att,
+        non_fast_prefix_len=non_fast_prefix_len,
+        action_valid=action_valid,
+        detach_prefix_kv_for_action=detach_prefix_kv_for_action,
+        scale=scale,
+        dropout_p=dropout_p,
+        enable_gqa=enable_gqa,
+    )
+
+
+def flex_attention_forward(*args, **kwargs) -> torch.Tensor:
+    """Forward-only compatibility wrapper."""
+    return flex_attention(*args, **kwargs)
+
+
+def backend_marker(x: torch.Tensor) -> torch.Tensor:
+    if ops is None:
+        return x
+    return ops._flashrt_training_package_marker(x)
+
+
+__all__ = [
+    "MASK_VALUE_F32",
+    "backend_marker",
+    "build_block_sparse_bool_masks",
+    "flex_attention",
+    "flex_attention_forward",
+    "manual_attention",
+    "manual_attention_part",
+    "reference_flex_attention",
+]

torch-ext/torch_binding.cpp

@@ -0,0 +1,18 @@
+// SPDX-License-Identifier: Apache-2.0
+
+#include <torch/all.h>
+#include <torch/library.h>
+
+#include "registration.h"
+#include "torch_binding.h"
+
+torch::Tensor flashrt_training_package_marker(torch::Tensor x) { return x; }
+
+TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
+  ops.def("_flashrt_training_package_marker(Tensor x) -> Tensor");
+  ops.impl("_flashrt_training_package_marker",
+           c10::DispatchKey::CompositeExplicitAutograd,
+           &flashrt_training_package_marker);
+}
+
+REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

torch-ext/torch_binding.h

@@ -0,0 +1,7 @@
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include <torch/all.h>
+
+torch::Tensor flashrt_training_package_marker(torch::Tensor x);