LTX-2.5 attention forms and the gated dual-rotary Diffusers adapter

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)

flashrt_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",

flashrt_structures/adapters/diffusers_gated_rotary_attention.py

@@ -0,0 +1,244 @@
+"""Attention adapter for gated dual-rotary Diffusers attention hosts.
+
+The audio+video joint-transformer form: Q and K are RMS-normalised after
+projection, rotary embeddings arrive as *separate* query/key boundaries
+(cross-modal sites rotate each side with its own table), the attention
+output may pass a per-head sigmoid gate computed from the pre-attention
+hidden states, and the processor owns the whole half — there is no
+residual or rescale state on the attention module itself. The stock
+Diffusers adapter refuses this family (its processor-state contract does
+not exist here), so the family gets its own adapter with the same shape:
+reproduce the projection half faithfully, capture real Q/K/V at every
+called site, bind the dense attention family per site, and replace only
+the attention dispatch. Projections, norms, rope, gating, and the output
+projection remain the host's own modules.
+"""
+
+from __future__ import annotations
+
+import inspect
+
+import torch
+
+from ..impls.attention_core import bind_dense_attention_best
+
+
+def _compatible_site(module, processor) -> tuple[bool, str]:
+    """Whether ``module`` exposes the gated dual-rotary processor contract."""
+    if not callable(processor):
+        return False, "processor is not callable"
+    try:
+        parameters = inspect.signature(processor.__call__).parameters
+    except (TypeError, ValueError, AttributeError):
+        return False, "processor call signature is not inspectable"
+    for name in ("query_rotary_emb", "key_rotary_emb"):
+        if name not in parameters:
+            return False, f"processor has no {name!r} boundary"
+    for attr in ("to_q", "to_k", "to_v", "norm_q", "norm_k"):
+        if not isinstance(getattr(module, attr, None), torch.nn.Module):
+            return False, f"attention lacks callable slot {attr!r}"
+    try:
+        out_proj, out_drop = module.to_out[0], module.to_out[1]
+    except (AttributeError, IndexError, KeyError, TypeError):
+        return False, "attention lacks the to_out[projection, dropout] slots"
+    if not all(isinstance(part, torch.nn.Module)
+               for part in (out_proj, out_drop)):
+        return False, "attention output slots are not modules"
+    heads = getattr(module, "heads", None)
+    if not isinstance(heads, int) or heads <= 0:
+        return False, "attention lacks a positive integer head count"
+    if not hasattr(module, "to_gate_logits"):
+        return False, "attention lacks the gate-logits slot"
+    if getattr(module, "rope_type", None) not in ("interleaved", "split"):
+        return False, "attention rope type is not a recognised form"
+    return True, ""
+
+
+def _apply_rope(attn, query, key, query_rotary_emb, key_rotary_emb):
+    if query_rotary_emb is None:
+        return query, key
+    from diffusers.models.transformers.transformer_ltx2 import (
+        apply_interleaved_rotary_emb, apply_split_rotary_emb)
+    k_rope = key_rotary_emb if key_rotary_emb is not None else query_rotary_emb
+    apply = (apply_interleaved_rotary_emb if attn.rope_type == "interleaved"
+             else apply_split_rotary_emb)
+    return apply(query, query_rotary_emb), apply(key, k_rope)
+
+
+def _qkv(attn, hidden_states, encoder_hidden_states,
+         query_rotary_emb, key_rotary_emb):
+    """Reproduce the host projection half; return SDPA-layout Q/K/V."""
+    context = (hidden_states if encoder_hidden_states is None
+               else encoder_hidden_states)
+    query = attn.norm_q(attn.to_q(hidden_states))
+    key = attn.norm_k(attn.to_k(context))
+    value = attn.to_v(context)
+    query, key = _apply_rope(attn, query, key, query_rotary_emb,
+                             key_rotary_emb)
+    head_dim = query.shape[-1] // attn.heads
+    query = query.unflatten(2, (attn.heads, head_dim)).transpose(1, 2)
+    key = key.unflatten(2, (attn.heads, head_dim)).transpose(1, 2)
+    value = value.unflatten(2, (attn.heads, head_dim)).transpose(1, 2)
+    return query, key, value
+
+
+class _Recorder:
+    def __init__(self, original, rows):
+        self.original = original
+        self.rows = rows
+
+    def __call__(self, attn, hidden_states, encoder_hidden_states=None,
+                 attention_mask=None, query_rotary_emb=None,
+                 key_rotary_emb=None, *args, **kwargs):
+        query, key, value = _qkv(
+            attn, hidden_states, encoder_hidden_states,
+            query_rotary_emb, key_rotary_emb)
+        self.rows.append({
+            "q": query.detach(),
+            "key": key.detach(),
+            "value": value.detach(),
+            "mask": (attention_mask.detach()
+                     if attention_mask is not None else None),
+        })
+        return self.original(
+            attn, hidden_states, encoder_hidden_states, attention_mask,
+            query_rotary_emb, key_rotary_emb, *args, **kwargs)
+
+
+class _FlashRTGatedRotaryAttnProcessor:
+    """Host processor with only the attention dispatch replaced."""
+
+    def __init__(self, core, original):
+        self.core = core
+        self.original = original
+
+    def __call__(self, attn, hidden_states, encoder_hidden_states=None,
+                 attention_mask=None, query_rotary_emb=None,
+                 key_rotary_emb=None, *args, **kwargs):
+        if attention_mask is not None and not getattr(
+                self.core, "allowed_ranges", ()):
+            return self.original(
+                attn, hidden_states, encoder_hidden_states, attention_mask,
+                query_rotary_emb, key_rotary_emb, *args, **kwargs)
+        gate_logits = None
+        if attn.to_gate_logits is not None:
+            gate_logits = attn.to_gate_logits(hidden_states)
+        query, key, value = _qkv(
+            attn, hidden_states, encoder_hidden_states,
+            query_rotary_emb, key_rotary_emb)
+        projection_dtype = query.dtype
+        guard = getattr(self.core, "_frt_guard", None)
+        accepted_dtypes = tuple(getattr(guard, "dtypes", ()) or ())
+        if accepted_dtypes and projection_dtype not in accepted_dtypes:
+            return self.original(
+                attn, hidden_states, encoder_hidden_states, attention_mask,
+                query_rotary_emb, key_rotary_emb, *args, **kwargs)
+        out = self.core(query, key, value)
+        out = out.transpose(1, 2).flatten(2, 3).to(projection_dtype)
+        if gate_logits is not None:
+            out = out.unflatten(2, (attn.heads, -1))
+            out = out * (2.0 * torch.sigmoid(gate_logits)).unsqueeze(-1)
+            out = out.flatten(2, 3)
+        out = attn.to_out[0](out)
+        out = attn.to_out[1](out)
+        return out
+
+
+class DiffusersGatedRotaryAttentionAdapter:
+    """Route gated dual-rotary Diffusers processors through the family.
+
+    Which executable form serves the seam is the family's decision, and
+    preferring a quantized one is a precision decision -- so it arrives
+    from the active scheme's ``attention_forms``, the same way the
+    gated-delta adapter reads its projection format. ``prefer`` is the
+    direct form of the same choice for a caller assembling this adapter
+    by hand; the scheme wins when both are given, because the scheme is
+    what the deployment selected.
+    """
+
+    __name__ = "diffusers_gated_rotary_attention"
+    scheme_aware = True
+
+    def __init__(self, prefer=()):
+        self.prefer = tuple(prefer)
+
+    def __call__(self, model, forward, *, prefix_cadence: bool = False,
+                 scheme=None):
+        del prefix_cadence
+        prefer = tuple(getattr(scheme, "attention_forms", ()) or self.prefer)
+        sites = []
+        for path, module in model.named_modules():
+            processor = getattr(module, "processor", None)
+            compatible, _ = _compatible_site(module, processor)
+            if compatible:
+                sites.append((path, module, processor))
+        if not sites:
+            return None
+
+        refused = []
+        captures = [[] for _ in sites]
+        for (_, module, original), rows in zip(sites, captures):
+            module.processor = _Recorder(original, rows)
+        try:
+            with torch.no_grad():
+                forward()
+        finally:
+            for _, module, original in sites:
+                module.processor = original
+
+        routes = []
+        observed = {}
+        variants = {}
+        for (path, module, original), rows in zip(sites, captures):
+            if not rows:
+                refused.append((
+                    f"{path}.processor",
+                    "attention_core gated-rotary: compatible processor was "
+                    "not called during calibration",
+                ))
+                continue
+            try:
+                core = bind_dense_attention_best(rows, prefer=prefer)
+            except ValueError as exc:
+                refused.append((f"{path}.processor", str(exc)[:160]))
+                continue
+            if core is None:
+                refused.append((
+                    f"{path}.processor",
+                    "attention_core gated-rotary: no family variant serves "
+                    "the captured head dimension or mask form",
+                ))
+                continue
+            routed = _FlashRTGatedRotaryAttnProcessor(core, original)
+            routes.append((module, original, routed))
+            observed[f"{path}.processor::attention_core"] = core
+            variants[f"{path}.processor"] = {
+                "bound": getattr(core, "_frt_variant", "fa2"),
+                "superseded": list(getattr(core, "_frt_variant_trail", ())),
+            }
+        if not routes:
+            return {}, None, {"refused": refused}
+
+        def enable() -> None:
+            for module, _, routed in routes:
+                module.processor = routed
+
+        def disable() -> None:
+            for module, original, _ in routes:
+                module.processor = original
+
+        def release() -> None:
+            # Reverting put the host processors back; this gives back the
+            # memory. The closures above keep working afterwards because
+            # they close over this list rather than over its contents.
+            routes.clear()
+
+        enable()
+        return {}, None, {
+            "revert": [disable],
+            "release": [release],
+            "observed": observed,
+            "toggle": (enable, disable),
+            "refused": refused,
+            "attention_variants": variants,
+        }