structures: LTX-2.5 attention forms and host bindings

docs/ltx25_structures_design.md

@@ -0,0 +1,125 @@
+# LTX-2.5 on the structures layer — design
+
+Scope: promote the LTX-2.5 integration (attention swap, W4A4 FFN chain,
+quantize-on-adopt) from model-private modules into catalog structures, so the
+same regions attach to any host that binds them — the native frontend here,
+and diffusers-hosted checkpoints through a binding, with no model-specific
+code in any impl.
+
+Everything below follows `docs/structures.md`: specs name positions, bindings
+place them on a host, impls decide what runs there. Nothing in this design
+adds a second vocabulary for calibration or qualification.
+
+## 1. What generalizes, and to which structure
+
+| Region (measured on LTX-2.5) | Catalog structure | Status |
+|---|---|---|
+| unmasked self/cross attention, head_dim 128 | `attention_core` | new backend impl |
+| GELU FFN pair (proj → tanh-GELU → down) | `vision_ffn` | new backend impl |
+| per-linear NVFP4 weight adoption | `quantize_on_adopt` | new scheme + binding attribute |
+| adaLN (rms · (1+scale) + shift, per-token tables) | host stage | future `adaln_producer` backend |
+| q/k RMSNorm + RoPE | host stage | future `qk_norm_rope` backend |
+| two-stage denoise pipeline | `video_generation_pipeline` | new pipeline binding |
+
+The audio branch (head_dim 64, short sequences) stays on the host attention
+path by measurement: quantized attention loses to SDPA at those shapes, so
+the binding simply does not claim those sites.
+
+## 2. `attention_core` backend: `sage2_qk_int8_pv_fp8`
+
+Executable form: per-warp INT8 quantization of Q, per-block INT8 of K,
+per-channel FP8 of V, one fused attention kernel, bf16 out. The kernel and
+its quantizers ship as one Hub artifact; the impl reads the supported head
+dims and layouts from the artifact instead of duplicating capability
+knowledge, exactly as the FA2 backend does.
+
+Qualification (all decided from real captures, refusal is legible):
+
+- head_dim must be advertised by the artifact (128 today); other dims return
+  no binding so the host keeps its own attention,
+- masked sites are not claimed — a mask that packs to a dense run can ride
+  the existing packed-KV plan later; today the masked path stays host,
+- scratch (int8/fp8 staging + output) is allocated per shape and shared
+  across all same-shaped sites; call sequences are pointer-stable, so the
+  region is CUDA-graph capturable.
+
+Parity gate: the spec's `real_distribution` rule. Measured on the target
+model this backend holds ~0.9992 cosine per call against an fp32 reference,
+and matched-input single-forward parity sits inside the noise floor of any
+same-precision kernel substitution; the latency rule is satisfied with
+2.0-2.4x over the strongest SDPA backend at the model's sequence lengths.
+
+## 3. `vision_ffn` backend: `w4a4_nvfp4_cutlass`
+
+Executable form, three launches replacing six:
+
+    activation quantize (bf16 -> NVFP4 + block scales)
+    up GEMM with bias + tanh-GELU + NVFP4 output epilogue
+    down GEMM (bf16 out; bias added when the slot carries one)
+
+Weight slots come from the spec; the impl accepts either origin:
+
+- **prequantized hosts** (checkpoint ships NVFP4): dequantize with the
+  reference kernel, requantize into the executable layout at adopt,
+- **bf16 hosts**: direct quantize at adopt (~seconds for a 22B model).
+
+Qualification:
+
+- both dims divisible by 16; rows padded to 128 through a staging buffer when
+  the host batches oddly — the GEMM rejects unaligned M *without writing
+  output*, so the impl owns the pad rather than trusting a return code,
+- adopt is layer-by-layer so peak memory stays near the fp4 footprint,
+- parity is gated against the plain-torch reference on real captures; on the
+  target model the chain holds the same distance from a bf16 golden as the
+  host's own W4A4 path while being 1.25-1.3x faster.
+
+## 4. `quantize_on_adopt`: the site list is a binding attribute
+
+The measured result that shapes this design: blanket adoption of every large
+linear visibly damages output, while adopting exactly the checkpoint
+author's calibrated selection (per-block attention/FFN linears, minus the
+final blocks; never adaLN producers, connectors, or patch/readout
+projections) matches bf16 quality. That selection is knowledge about the
+*host*, not about any impl — so it lives in the host binding as an explicit
+site list, and the scheme refuses to adopt outside it unless the caller
+overrides deliberately. A prequantized checkpoint is itself the receipt for
+that list.
+
+## 5. Pipeline binding
+
+`video_generation_pipeline`, same family as the existing video hosts:
+condition encoding (text tower + connector stack, slower cadence, embeddings
+cacheable per prompt), latent preparation, the fixed-step denoise loop, an
+optional latent upsample stage, and VAE decode. Hot-path segments classify
+per the coverage contract; attention and FFN regions point at the structures
+above, adaLN/RoPE stay declared host stages until their structures land, and
+the denoise loop is the graph-capture boundary.
+
+Two facts from bring-up that the binding must carry as attributes rather
+than rediscover:
+
+- the distilled checkpoint generation wants single-pass denoising — guidance
+  and modality-isolation scales at 1.0 — and defaults that re-enable extra
+  passes triple the step cost silently,
+- with the transformer resident, decode tiling must be budgeted against the
+  memory decode will actually see, not a pre-build snapshot.
+
+## 6. Measured context (RTX 5090, 1536x1024x121f unless noted)
+
+- native frontend: denoise 23.9s -> 11.7s (2.04x) with attention + FFN +
+  compile + whole-loop capture; per-step 1068.6 -> 491.6ms (stage 1),
+  5111.7 -> 2596ms (stage 2)
+- diffusers-hosted bf16 checkpoint, single-pass distilled schedule:
+  254 -> 54s end-to-end (3.4x; per-step 3.85x) with adopt + attention swap +
+  per-block compile, quality matched to the bf16 baseline by frame
+  inspection; at 768x512x49f the gap to the offload baseline is >10x
+- adopt cost: ~6s for 1176 linears of a 22B transformer
+
+## 7. Sequencing
+
+1. attention backend impl + gate records
+2. vision_ffn backend impl + gate records
+3. quantize_on_adopt site-list attribute + host binding
+4. pipeline binding with coverage classification
+5. adaLN / qk-norm-RoPE structures (removes the two biggest remaining host
+   stages; profiled at ~35% of a denoise step on the diffusers host)

docs/ltx25_usage.md

@@ -121,3 +121,133 @@ Measured on 5090 at 1536×1024×121f (median, video self-attention site,
 S=24576): SDPA-cudnn 42.4ms, sage2 17.4ms, sage3 13.0ms. End-to-end stage-2
 denoise per step: 5.11s (SDPA) → 3.84s (sage2), with output quality equivalent
 under matched-input single-forward cosine and frame inspection.
+
+## The same model through the structures layer
+
+The runtime above drives the official pipeline. The transformer is also
+reachable as an ordinary Diffusers host, where the structures layer attaches
+to it without a model-specific path:
+
+```python
+from flash_rt import structures
+
+plan = structures.attach(model, forward, scheme="nvfp4_balance")
+print(plan.report())          # bound seams, gate results, ledger
+plan.detach()                 # restores the host exactly
+```
+
+`attach` discovers the seams, calibrates on one real forward, gates accuracy
+and latency per family, and keeps the host path wherever a gate declines.
+Nothing here is LTX-specific: the attention seam is recognised by the
+processor contract (separate query/key rotary boundaries, per-head gating),
+not by a model or class name.
+
+`scheme=` selects the precision profile, and two are relevant here:
+
+| scheme | what it quantizes |
+|---|---|
+| `"nvfp4_balance"` | the projection GEMMs (W4A4); attention keeps the family's precision-first order |
+| `"nvfp4_balance_sage"` | the same, and allows the quantized attention forms to be weighed first |
+
+The split is deliberate. A quantized attention form trades a bounded error
+for speed, so it is a precision decision like any other and arrives through
+the profile a deployment selected — never through a device check or a host
+binding. Naming a form does not force it: the family still qualifies the
+shape, speed-gates the result, and falls through to the published order when
+the installed package does not serve the site.
+
+To tune one seam without registering a profile, name the forms directly:
+
+```python
+plan = structures.attach(model, forward, scheme="nvfp4_balance",
+                         attention_forms=("sage2",))
+```
+
+Either way the answer comes back measured. `plan.report()` prints each
+family's accuracy band and the paired latency it was judged on, and the
+attachment can be reverted exactly, so the way to decide between these is to
+run both and read the two reports rather than to take this table's word for
+it on a different card.
+
+### End to end
+
+What a request costs, wall clock, same prompt and seed, distilled
+single-pass recipe. The baseline is the unmodified host: a 44GB bf16
+checkpoint that does not fit on a 32GB part, so it runs with weight
+offloading, which is what a user of this model on this class of card
+actually starts from.
+
+| Request | Host (offload) | `"nvfp4_balance"` | `"nvfp4_balance_sage"` |
+|---|---|---|---|
+| 768×512×49f | 99.8 s | 6.0 s (16.6×), peak 29.9 GB | **5.7 s (17.5×)**, peak 26.8 GB |
+| 1536×1024×121f | 181.6 s | **87.9 s (2.07×)**, peak 28.0 GB | does not fit, see below |
+
+Medians of three warm runs, eager. Frames are inspection-equivalent to the
+host's own output.
+
+The full-size row is the honest one to read closely. Attaching all 48 blocks
+succeeds and each block's gate measures 1.48×, but the assembled pipeline
+sits close to the limit of a 32GB part: 23.9 GB resident before the request
+begins, 28.0 GB at peak, and video-VAE tiling is needed for decode to have
+room at all. Three things account for the distance between 2.07× here and
+what the same hardware reaches with a hand-assembled configuration:
+
+- this measurement is eager, with no compilation of the block stack;
+- the quantized attention profile does not fit at this size yet — the
+  per-shape staging pool costs about 3 GB on top, which the assembly step
+  runs out of;
+- the audio feed-forward stays at host precision throughout, 3.0 GB across
+  the model, because its 126-row calls sit outside the fused chain's
+  128-row alignment and the seam declines them rather than produce an
+  unwritten output.
+
+### Where the time goes, one transformer block
+
+The table below is a diagnostic, not the result: it says which family earned
+which part of the request time above. Real checkpoint weights, real captured
+deployment inputs, paired alternating timing inside the gate.
+
+| Site shape | `scheme=` | Block latency | Attention family | Peak memory |
+|---|---|---|---|---|
+| S=24576 (1536×1024×121f) | host, unattached | 134.3 ms | — | 12.2 GB |
+| | `"nvfp4_balance"` | 117.1 ms (1.15×) | BF16 form bound, declined at 1.006× | 8.2 GB |
+| | `"nvfp4_balance_sage"` | **89.8 ms (1.49×)** | activated, 1.259× | at the 32GB ceiling |
+| S=2688 (768×512×49f) | host, unattached | 10.2 ms | — | 2.3 GB |
+| | `"nvfp4_balance"` | 8.2 ms (1.25×) | declined | 1.7 GB |
+| | `"nvfp4_balance_sage"` | 8.0 ms (1.28×) | activated, 1.022× | 4.7 GB |
+
+Matched-forward cosine against the host's own output is 0.99999 in every row,
+and `detach` restores it bit-exactly (max-abs 0.0). Three things in that table
+are easy to misread, so they are worth stating:
+
+- **A block ratio is not a kernel ratio.** The same attention that measures
+  2.34× on its own (45.9 → 19.6 ms at S=24576) shows up as 1.259× for the
+  attention unit, because the unit is judged against the whole block. The
+  27 ms it saves is the same 27 ms in both numbers.
+- **The projections-only profile leaves the BF16 attention form bound and
+  declined.** That form measures 46.1 ms against the host's 45.9 ms here, so
+  the gate is right to keep the host path; nothing about the quantized forms
+  is being judged in that row.
+- **Peak memory moves in both directions.** Quantizing the projections takes
+  it from 12.2 to 8.2 GB. Preferring quantized attention gives some back,
+  because each attention site owns its staging and quantization workspace:
+  four sites at S=24576 reach the ceiling of a 32GB part. Pooling those
+  workspaces is the open item before that profile is usable at full size.
+
+### How the whole-model figures were produced
+
+Blocks are attached one at a time, because a 44GB bf16 checkpoint is not
+resident on a 32GB part: each block is materialized alone, attached on its
+own real inputs, and its host weights released before the next. The
+feed-forward seams are bound explicitly rather than by discovery, because
+`vision_ffn` does not claim this host's shape — its projections carry no
+bias and its norm sits outside the seam, both of which the structure's
+boundary requires. Whether to widen that boundary is a catalog decision.
+
+### Kernel availability is the package's own statement
+
+The forms read their envelope from the installed artifact. The sage3 package
+publishes head_dim 128 only in its CUDA 13 builds; on a CUDA 12.8 host it
+advertises head_dim 64, so a 128-wide site is refused there and the ladder
+falls through — visible on the refusal trail rather than as a silent
+slowdown. Nothing in this repository keeps a second table of that.

flash_rt/structures/adapters/__init__.py

@@ -12,6 +12,9 @@
     register_qkv_rope_adapter,
 )
 from .diffusers_attention import DiffusersAttentionAdapter
+from .diffusers_gated_rotary_attention import (
+    DiffusersGatedRotaryAttentionAdapter,
+)
 from .diffusers_rotary_attention import DiffusersRotaryAttentionAdapter
 from .factored_two_way_attention import FactoredTwoWayAttentionAdapter
 from .factored_qk_norm_rope import FactoredQkNormRopeAdapter
@@ -33,6 +36,7 @@
 register_qkv_rope_adapter(PackedQkvRopeAdapter())
 register_attention_adapter(GemmaAttentionAdapter())
 register_attention_adapter(FactoredTwoWayAttentionAdapter())
+register_attention_adapter(DiffusersGatedRotaryAttentionAdapter())
 register_attention_adapter(DiffusersRotaryAttentionAdapter())
 register_attention_adapter(DiffusersAttentionAdapter())
 # the fused-layer form is tried first; it refuses cleanly (missing
@@ -43,6 +47,7 @@
 
 __all__ = [
     "DiffusersAttentionAdapter",
+    "DiffusersGatedRotaryAttentionAdapter",
     "DiffusersRotaryAttentionAdapter",
     "GemmaAttentionAdapter",
     "TransformersGatedDeltaAdapter",