diff --git a/docs/ltx25_structures_design.md b/docs/ltx25_structures_design.md new file mode 100644 index 00000000..3ae0b453 --- /dev/null +++ b/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) diff --git a/docs/ltx25_usage.md b/docs/ltx25_usage.md index d4f80633..2c77ff7d 100644 --- a/docs/ltx25_usage.md +++ b/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. diff --git a/flash_rt/structures/adapters/__init__.py b/flash_rt/structures/adapters/__init__.py index d6369f68..d73fba99 100644 --- a/flash_rt/structures/adapters/__init__.py +++ b/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", diff --git a/flash_rt/structures/adapters/diffusers_gated_rotary_attention.py b/flash_rt/structures/adapters/diffusers_gated_rotary_attention.py new file mode 100644 index 00000000..72e34afc --- /dev/null +++ b/flash_rt/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, + } diff --git a/flash_rt/structures/autobuild.py b/flash_rt/structures/autobuild.py index 10b07e3a..7a9d6870 100644 --- a/flash_rt/structures/autobuild.py +++ b/flash_rt/structures/autobuild.py @@ -50,6 +50,20 @@ _GATED_DELTA_ADAPTERS: list = [] +class _AttentionOverride: + """The caller's own answer to the question a scheme usually answers. + + Adapters read ``attention_forms`` off the scheme; an explicit argument + is the same statement made directly, so it arrives the same way rather + than through a second parameter every adapter would have to learn. + """ + + __slots__ = ("attention_forms",) + + def __init__(self, forms): + self.attention_forms = tuple(forms) + + def register_attention_adapter(adapter) -> None: """Register a host-family attention adapter (callable).""" _ATTENTION_ADAPTERS.append(adapter) @@ -115,6 +129,15 @@ class AutoPlan: #: anything. A seam that cannot be turned off cannot be measured. toggles: list[tuple[Callable[[], None], Callable[[], None]]] = field( default_factory=list) + #: callables that drop an adapter's own hold on the forms it bound. + #: Reverting a routed seam puts the host processor back but does not + #: free anything: the bound form stays reachable through ``observed`` + #: and through the very closures that reverted it. That is right while + #: the form might still be activated and wrong once it cannot be, and + #: the difference is a form's whole working set — at long sequence + #: lengths, hundreds of megabytes per site, held for something that + #: will never run. + releases: list[Callable[[], None]] = field(default_factory=list) #: ``flash_rt.core.precision_spec.ModelPrecisionSpec`` for the scales #: this plan baked in — the repo's introspection format, not a private #: one, so ``plan.precision_spec`` reads like ``rt.precision_spec`` @@ -128,6 +151,22 @@ def disable_routed(self) -> None: for _, off in self.toggles: off() + def release_routed(self) -> None: + """Give back routed forms that will not be activated. Idempotent. + + Called when the gate has settled and no routed unit won: the host + processors are already back, and what remains is memory held for a + form the gate declined. The revert callables stay in place and stay + correct — an adapter's release empties its own route list, so those + closures survive with nothing left to undo. + """ + self.disable_routed() + for release in self.releases: + release() + self.releases.clear() + self.toggles.clear() + self.observed.clear() + def abort(self) -> None: """Roll back everything this plan touched without an attach. @@ -544,6 +583,7 @@ def auto_swaps( percentile: float = 99.9, max_samples: int | None = None, scheme: str | Any = "auto", + attention_forms: Sequence[str] | None = None, verbose: bool = False, stream_store: Any = None, ) -> AutoPlan: @@ -1263,8 +1303,25 @@ def _stream_window_undo(): # sample entry, where that callable takes a sample — # the whole point of normalising the three ways in was # that nothing downstream should see the difference - result = adapter(model, thunks[0], - prefix_cadence=prefix_cadence) + # same convention as the gated-delta adapters: an + # adapter that declares scheme awareness receives the + # active scheme, because which executable form may + # serve its seam is a precision decision and the + # scheme is where those are stated + if getattr(adapter, "scheme_aware", False): + # an explicit argument outranks the profile: it is + # the caller answering the same question directly, + # which is how a deployment tunes one seam without + # having to register a scheme for it + result = adapter( + model, thunks[0], + prefix_cadence=prefix_cadence, + scheme=(_AttentionOverride(attention_forms) + if attention_forms is not None + else scheme_obj)) + else: + result = adapter(model, thunks[0], + prefix_cadence=prefix_cadence) except (ValueError, RuntimeError) as refusal: plan.notes.setdefault("refused", []).append( ("attention_core", str(refusal)[:200])) @@ -1294,6 +1351,7 @@ def _stream_window_undo(): "attention_core_variants", {}).update( extras["attention_variants"]) plan.revert.extend(extras.get("revert", ())) + plan.releases.extend(extras.get("release", ())) if extras.get("toggle") is not None: plan.toggles.append(extras["toggle"]) if update is not None: diff --git a/flash_rt/structures/bindings/ltx25_dit.yaml b/flash_rt/structures/bindings/ltx25_dit.yaml new file mode 100644 index 00000000..1d310ac1 --- /dev/null +++ b/flash_rt/structures/bindings/ltx25_dit.yaml @@ -0,0 +1,27 @@ +binding: ltx25_dit +structure: vision_ffn +dims: {D: 4096, F: 16384, layers: 48} +variant: {activation: gelu} # host applies the tanh approximation +boundary_dtype: bf16 + +m_profile: + denoise_stage1: {m_class: large, phase: denoise} # M = video tokens, first distilled phase + denoise_stage2: {m_class: large, phase: denoise} # M = 4x tokens after latent upsample + +weights_map: + diffusers: + layout: out_in + w_fc1: "transformer_blocks.{i}.ff.net.0.proj.weight" + b_fc1: {const: zeros} # ff_bias=false on this host family + w_fc2: "transformer_blocks.{i}.ff.net.2.weight" + b_fc2: {const: zeros} + w_norm: {const: ones} # norm2 is non-affine; scale/shift come from the + b_norm: {const: zeros} # per-token adaptive tables, outside this seam + +quantized_layers: "0-41" # the checkpoint's own quantization whitelist: + # the last six blocks stay at host precision + +hosts: + diffusers: + module_path: "transformer_blocks.{i}.ff" + versions: ">=0.40" diff --git a/flash_rt/structures/bindings/ltx25_video_pipeline.yaml b/flash_rt/structures/bindings/ltx25_video_pipeline.yaml new file mode 100644 index 00000000..195c14e3 --- /dev/null +++ b/flash_rt/structures/bindings/ltx25_video_pipeline.yaml @@ -0,0 +1,155 @@ +binding: ltx25_video_pipeline +structure: video_generation_pipeline + +# Joint audio+video DiT host. The distilled checkpoint's official recipe is +# a single pass per step: every guidance/modality scale is 1.0, so no +# classifier-free or modality-isolation passes exist on the hot path. A +# binding that re-enables them changes the program being measured, not a +# knob on the same program. +recipe: + guidance_scale: 1.0 + audio_guidance_scale: 1.0 + modality_scale: 1.0 + audio_modality_scale: 1.0 + phases: + - {name: stage1, steps: 8} + - {name: stage2, steps: 3, after: latent_upsample} + +# Sites the checkpoint family itself quantizes; a scheme refuses to adopt +# outside this list unless explicitly overridden. The last six transformer +# blocks, every adaptive-norm single, the modality connectors, and the +# patchify/proj_out boundary stay at host precision. +quantization: + adopt_blocks: "0-41" + adopt_modules: + - "attn1" + - "attn2" + - "audio_attn1" + - "audio_attn2" + - "audio_to_video_attn" + - "video_to_audio_attn" + - "ff" + - "audio_ff" + adopt_leaves: + attention: ["to_q", "to_k", "to_v", "to_out.0"] + feed_forward: ["net.0.proj", "net.2"] + +stages: + condition_encode: + seam: "text encoder and per-modality connectors" + capture: eager_or_offloaded + residency: staged # encoder weights leave the device before denoise + latent_prepare: + seam: "video/audio latent initialization from the request seed" + capture: eager + denoise: + seam: "joint transformer forward under the ancestral solver loop" + loop_steps: "@recipe.phases" + capture: host_dependent + decode: + seam: "video VAE decode and audio vocoder" + capture: eager + tiling_budget: resident_aware # decode tiling must be sized against + # memory the resident transformer holds + +coverage: + contract: complete_hot_path + hot_path: + - prompt_encode + - modality_connectors + - latent_initialization + - denoise_control + - adaptive_norm_tables + - per_head_gating + - qk_norm_rope + - video_self_attention + - video_cross_attention + - audio_attention + - cross_modal_attention + - video_ffn + - audio_ffn + - attention_projections + - latent_upsample + - scheduler_update + - vae_decode + segments: + - name: prompt_encode + stage: condition_encode + classification: host_stage + seam: "text encoder forward" + - name: modality_connectors + stage: condition_encode + classification: host_stage + seam: "per-modality connector projections over encoder output" + - name: latent_initialization + stage: latent_prepare + classification: state_region + seam: "seeded video/audio latents and their patchified views" + - name: denoise_control + stage: denoise + classification: control + seam: "two-phase ancestral Euler schedule, single pass per step" + - name: adaptive_norm_tables + stage: denoise + classification: host_stage + seam: "per-token scale/shift/gate table lookups around every block" + - name: per_head_gating + stage: denoise + classification: host_stage + seam: "sigmoid per-head gate on attention outputs" + - name: qk_norm_rope + stage: denoise + classification: host_stage + seam: "per-head Q/K RMSNorm and rotary application" + - name: video_self_attention + stage: denoise + classification: structure + seam: "video branch attn1, dense unmasked, head_dim 128" + structures: [attention_core] + - name: video_cross_attention + stage: denoise + classification: structure + seam: "video branch attn2 over connector tokens, head_dim 128" + structures: [attention_core] + - name: audio_attention + stage: denoise + classification: host_stage + seam: "audio branch attention, head_dim 64 — outside the measured + net-win band of the claimed attention forms" + - name: cross_modal_attention + stage: denoise + classification: host_stage + seam: "audio-to-video and video-to-audio attention" + - name: video_ffn + stage: denoise + classification: structure + seam: "video branch feed-forward, tanh-GELU" + structures: [vision_ffn] + - name: audio_ffn + stage: denoise + classification: host_stage + seam: "audio branch feed-forward — small-M band measured net-negative + for the W4A4 form, retained on the host" + - name: attention_projections + stage: denoise + classification: structure + seam: "Q/K/V/O and feed-forward projections on the checkpoint's own + quantization whitelist" + structures: [linear_proj] + - name: latent_upsample + stage: denoise + classification: host_stage + seam: "latent upsampler between the two distilled phases" + - name: scheduler_update + stage: denoise + classification: state_region + seam: "ancestral solver latent update" + - name: vae_decode + stage: decode + classification: host_stage + seam: "tiled video VAE decode and audio vocoder" + +hosts: + diffusers: + module_path: "diffusers.LTX2VideoTransformer3DModel" + versions: ">=0.40" diff --git a/flash_rt/structures/frontdoor.py b/flash_rt/structures/frontdoor.py index aa64a1e8..585579d4 100644 --- a/flash_rt/structures/frontdoor.py +++ b/flash_rt/structures/frontdoor.py @@ -83,7 +83,8 @@ #: hands the consumer a dtype it refuses — which the ledger would report #: as a family that fell back, from a split this gate created itself. _CHAIN = "negotiated_fp8_chain" -_ROUTED = "attention_core_routed" +_ROUTED_SUFFIX = "_routed" +_ROUTED = "attention_core" + _ROUTED_SUFFIX def _cuda_time_ms(fn: Callable[[], Any], warmup: int = 3, @@ -294,6 +295,7 @@ def attach( iters: int = 10, on_guard_fail: str = "fallback", scheme: str | Any = "auto", + attention_forms: Sequence[str] | None = None, negotiate_fp8: bool = True, verbose: bool = True, ) -> Plan: @@ -348,8 +350,8 @@ def say(msg: str) -> None: plan = auto_swaps(model, forward, structures=structures, observations=observations, percentile=percentile, max_samples=max_samples, prefix_cadence=prefix_cadence, - scheme=scheme, negotiate_fp8=negotiate_fp8, - verbose=verbose) + scheme=scheme, attention_forms=attention_forms, + negotiate_fp8=negotiate_fp8, verbose=verbose) if not plan.swaps and not plan.toggles: plan.revert_all() return Plan({}, {}, {"digest": "none", "seams": 0, @@ -411,6 +413,9 @@ def scored() -> torch.Tensor | None: f"the host module, so this unit did not run: " f"{led['seams_fell_back'][:3]}") continue + stat["shape"] = _measured_shape(paths, plan) + stat["per_shape_rule"] = _per_shape_rule(name) + stat["memory"] = _memory_delta(paths) timing = _paired_ab(eval_thunk, arm.on, arm.off, rounds=rounds, iters=iters) stat["e2e"] = timing @@ -418,7 +423,11 @@ def scored() -> torch.Tensor | None: stat["outcome"] = "refused" stat["reason"] = ( f"no net win ({timing['speedup']:.3f}x, spread " - f"{timing['spread']:.3f}) at {_shape_note(plan)}") + f"{timing['spread']:.3f}){_memory_note(stat['memory'])} " + f"at {stat['shape']}" + + (" — this unit's latency qualification is per shape, " + "so the verdict holds for that shape and no other" + if stat["per_shape_rule"] else "")) continue stat["outcome"] = "activated" if routed: @@ -484,6 +493,14 @@ def scored() -> torch.Tensor | None: if not winners and not routed_winner: plan.revert_all() say("outcome: whole-host refusal — model left untouched") + elif not routed_winner: + # Some units won and the routed ones did not. Reverting them put + # the host processors back, but the forms themselves were still + # reachable — through the plan, for as long as the caller holds it + # to detach with. A declined form's working set is not small, so + # holding it made an attachment cost more memory than the host it + # replaced at exactly the shapes where that matters most. + plan.release_routed() activated = dict(winners) if routed_winner: @@ -562,6 +579,143 @@ def _round(metrics: Mapping[str, Any]) -> dict[str, Any]: for k, v in metrics.items()} +def _per_shape_rule(unit: str) -> bool: + """Whether this unit's spec says its latency verdict is per shape. + + Thirteen catalog specs declare ``latency.per_shape: true`` and nothing + has ever read it, so every verdict has been recorded as though it were + universal. Reading it does not change any decision; it changes what a + refusal claims, which is the difference between "this form does not + help" and "this form did not help at the one size it was measured at". + """ + from .registry import load + + try: + spec = load(unit[:-len(_ROUTED_SUFFIX)] + if unit.endswith(_ROUTED_SUFFIX) else unit) + except (KeyError, ValueError): + return False + latency = (spec.gates or {}).get("latency") or {} + return bool(latency.get("per_shape")) + + +def _measured_shape(paths, plan: AutoPlan) -> str: + """The form the bound seams were armed for, read off their guards. + + ``m_profile`` carries this for hosts whose bindings declare it and is + empty for the rest, where the refusal then reads "rows unrecorded" — + a shape-scoped verdict with no shape on it. The guards know: every + bound seam declares the row count or capacity it was armed for. + """ + rows = set() + for module in (paths.values() if hasattr(paths, "values") else paths): + guard = getattr(module, "_frt_guard", None) + if guard is None: + continue + if getattr(guard, "rows", None): + rows.add(f"rows={guard.rows}") + elif getattr(guard, "row_capacity", None): + rows.add(f"rows<={guard.row_capacity}") + if rows: + return ", ".join(sorted(rows)) + declared = sorted({m for s in plan.seams for m in (s.m_profile or ())}) + return f"rows={declared}" if declared else "rows unrecorded" + + +def _tensors_of(module, depth: int = 2): + """Every device tensor a bound form holds, however it holds it. + + Parameters and buffers are the easy half. The rest — staging, packed + weights, a workspace object an artifact handed back — hangs off plain + attributes and off containers inside them, which is exactly where a + form's working set lives, so a walk that stops at ``parameters()`` + reports the largest holdings as zero. + """ + if module is None: + return + if torch.is_tensor(module): + yield module + return + if depth <= 0: + return + if isinstance(module, torch.nn.Module): + yield from module.parameters() + yield from module.buffers() + for sub in module.modules(): + for value in vars(sub).values(): + yield from _tensors_of(value, depth - 1) + return + if isinstance(module, (tuple, list, set)): + for item in module: + yield from _tensors_of(item, depth - 1) + return + if isinstance(module, Mapping): + for item in module.values(): + yield from _tensors_of(item, depth - 1) + return + for value in getattr(module, "__dict__", {}).values(): + yield from _tensors_of(value, depth - 1) + + +def _device_bytes(modules, exclude: set[int] | None = None) -> tuple[int, set]: + """Device bytes these hold, counting each storage once. + + Deduplication by pointer is the whole point: a family that pools one + staging set across its sites would otherwise be counted once per site + and read as many times its real cost. ``exclude`` keeps a retained host + module out of its replacement's total — the bound form holds it for + fallback, and counting it on both sides would report a form that halves + the weights as costing nothing. + """ + seen, total = set(), 0 + exclude = exclude or set() + for module in modules: + for tensor in _tensors_of(module): + ptr = tensor.data_ptr() + if not tensor.is_cuda or ptr in seen or ptr in exclude: + continue + seen.add(ptr) + total += tensor.numel() * tensor.element_size() + return total, seen + + +def _memory_delta(paths) -> dict[str, int]: + """What this unit changes about resident memory. Recorded, not judged. + + The latency gate has always decided alone, which is right until a unit + that loses on time turns out to hold — or to have saved — gigabytes. + Today that consequence does not appear in the receipt at all, so it + cannot even be discussed; a refusal reads as free when it may be the + most expensive line in the run. + + ``bound`` is what the unit's forms hold, ``host`` what the modules they + replace hold, and the difference is what activating this unit does to + resident memory: negative when a quantized form replaces host weights, + positive when a form brings a working set the host did not need. + """ + if not torch.cuda.is_available(): + return {} + modules = list(paths.values() if hasattr(paths, "values") else paths) + hosts = [] + for module in modules: + host = getattr(module, "_frt_host", None) + hosts.append(host() if callable(host) else None) + host_bytes, host_ptrs = _device_bytes(hosts) + bound, _ = _device_bytes(modules, exclude=host_ptrs) + return {"bound_bytes": bound, "host_bytes": host_bytes, + "delta_bytes": bound - host_bytes} + + +def _memory_note(memory) -> str: + """The memory consequence, in the refusal itself.""" + if not memory: + return "" + delta = memory.get("delta_bytes", 0) + if abs(delta) < 64 << 20: + return "" + return f", {delta / 2 ** 30:+.2f} GiB resident" + + def _shape_note(plan: AutoPlan) -> str: """The workload a refusal was measured at, for the receipt. diff --git a/flash_rt/structures/impls/__init__.py b/flash_rt/structures/impls/__init__.py index 2c294d1e..c1d35893 100644 --- a/flash_rt/structures/impls/__init__.py +++ b/flash_rt/structures/impls/__init__.py @@ -208,7 +208,12 @@ def hub_kernel(repo: str, version: str): # serve the same call site. _LOADED[key] = (get_kernel(repo, revision=rev) if rev else get_kernel(repo)) # pre-semver band - except (OSError, RuntimeError, ValueError) as unavailable: + except (ImportError, OSError, RuntimeError, + ValueError) as unavailable: + # ImportError: the kernels library validates a package's own + # python-dependency declarations before serving it; a host + # missing one cannot supply the package, same as any other + # unavailability _record_unavailable(repo, version, unavailable) raise KernelUnavailable( f"kernel package {repo!r} ({version}) is unavailable on " diff --git a/flash_rt/structures/impls/attention_core/__init__.py b/flash_rt/structures/impls/attention_core/__init__.py index dbad81c7..d9c441c2 100644 --- a/flash_rt/structures/impls/attention_core/__init__.py +++ b/flash_rt/structures/impls/attention_core/__init__.py @@ -6,7 +6,7 @@ from .two_way_fa2 import FactoredTwoWayAttention, bind_two_way_attention -def bind_dense_attention_best(captures): +def bind_dense_attention_best(captures, *, prefer=()): """Dense attention across the variant family, precision-descending. One structure, parallel executable forms, and no second hardware @@ -26,6 +26,15 @@ def bind_dense_attention_best(captures): package is merely absent looks identical to one where it was weighed and rejected — the two need different fixes, and only the trail tells them apart. + + ``prefer`` names forms to try ahead of that order, and is empty by + default because the order above is a precision order: the quantized + forms trade a bounded error for speed, which is a decision about the + deployment rather than about the device, and a device-shaped ladder + is the wrong place to make it. A caller that has judged the trade — + a host binding, a qualification run — names the form it wants and + gets the same qualification walk, speed gate and trail as any other + rung. An unknown name is a caller error, not a silent no-op. """ # a package can refuse a device three ways: its arch declaration # (ValueError from the loader's metadata check), the kernels @@ -47,10 +56,26 @@ def _fa4_fp8(caps): from . import fa4_fp8 return fa4_fp8.bind_dense_attention(caps) + def _sage2(caps): + from . import sage2_blackwell + return sage2_blackwell.bind_dense_attention(caps) + + def _sage3(caps): + from . import sage3_blackwell + return sage3_blackwell.bind_dense_attention(caps) + + order = [("fa2", _fa2), ("fa4_cute", _fa4_cute), + ("masked_mha", _masked_mha), ("fa4_fp8", _fa4_fp8)] + by_name = dict(order, sage2=_sage2, sage3=_sage3) + for name in reversed(tuple(prefer)): + if name not in by_name: + raise ValueError( + f"attention_core: unknown preferred form {name!r}; " + f"available: {sorted(by_name)}") + order.insert(0, (name, by_name[name])) + refusals, declined = [], 0 - for name, binder in (("fa2", _fa2), ("fa4_cute", _fa4_cute), - ("masked_mha", _masked_mha), - ("fa4_fp8", _fa4_fp8)): + for name, binder in order: try: core = binder(captures) except (ValueError, RuntimeError, OSError) as refusal: diff --git a/flash_rt/structures/impls/attention_core/sage2_blackwell.py b/flash_rt/structures/impls/attention_core/sage2_blackwell.py new file mode 100644 index 00000000..4891cb8a --- /dev/null +++ b/flash_rt/structures/impls/attention_core/sage2_blackwell.py @@ -0,0 +1,223 @@ +"""attention_core — the sage2 (Blackwell INT8-QK) dense form. + +The consumer-Blackwell quantized twin of the dense attention family: the +same stateless seam — complete Q/K/V every call, host SDPA layout — executed +by the ``flashrt/sageattention2-blackwell`` kernel: per-warp (or per-thread) +INT8 quantization of Q and K, FP8 per-channel (or FP16) V, one fused +attention, BF16 out. Where the bf16 forms preserve numerics exactly, this +form trades a bounded quantization error for roughly half the attention +time on long unmasked sites; the parity gates downstream judge that trade +on real captures like any other impl's. + +Two precision variants, both from the same artifact: + +- ``pv_fp8`` — INT8 QK / FP8 V. The speed point of the family. +- ``pv_fp16`` — INT8 QK / FP16 V. Recovers most of the quantization error + at ~60% more attention time; the option for hosts whose gates reject the + FP8 point. + +Qualification, decided from the artifact and the captures, refusal 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 has no form this kernel accepts, and + the packed-KV plan of the FA2 form does not transfer (the quantizers + consume dense NHD), so any allowed-ranges request refuses here, +- the workspace (INT8/FP8 staging, scales, output) is caller-owned and + allocated once per shape at bind — call sequences are pointer-stable and + the artifact declares itself CUDA-graph safe. +""" + +from __future__ import annotations + +from typing import NamedTuple + +import torch + +from .. import hub_kernel +from ...guard import PROCEED, GuardedSeam + +KERNEL_DEP = { + "provider": "huggingface_kernels", + "repo": "flashrt/sageattention2-blackwell", + "version": ">=1", +} + +_VARIANTS = ("pv_fp8", "pv_fp16") + + +class _Staging(NamedTuple): + q: torch.Tensor + k: torch.Tensor + v: torch.Tensor + out: torch.Tensor + workspace: object + + +#: One staging set and workspace per distinct (shape, dtype, device, +#: variant, granularity). A transformer reaches this seam once per block +#: and per attention site, all with the same shapes, and the scratch is +#: large: at 24576 tokens over 32 heads a single set is about 800MB, so +#: forty-eight blocks owning their own would not fit on any consumer part. +#: Sharing is safe because the scratch holds nothing between calls -- it +#: is filled at the top of every forward and read before returning -- and +#: because these calls are sequential on one stream by construction: the +#: host runs its blocks in order. The pool is keyed, never emptied, and +#: pointer-stable, which is also what a captured graph needs. +_STAGING: dict[tuple, _Staging] = {} + + +def _staging_for(q_shape, kv_shape, dtype, device, variant, granularity): + key = (tuple(q_shape), tuple(kv_shape), dtype, + str(torch.device(device)), variant, granularity) + staging = _STAGING.get(key) + if staging is None: + art = _artifact() + q = torch.empty(*q_shape, dtype=dtype, device=device) + k = torch.empty(*kv_shape, dtype=dtype, device=device) + v = torch.empty_like(k) + out = torch.empty_like(q) + staging = _Staging(q, k, v, out, art.allocate_workspace( + q, k, v, fp8v=(variant == "pv_fp8"), + qk_quant_granularity=granularity)) + _STAGING[key] = staging + return staging + + +def _artifact(): + return hub_kernel(KERNEL_DEP["repo"], KERNEL_DEP["version"]) + + +def supported_head_dims() -> tuple[int, ...]: + """Executable envelope, read from the artifact — never duplicated here.""" + caps = _artifact().capabilities() + dims = tuple(sorted(int(d) for d in caps["head_dims"])) + if not dims: + raise ValueError( + "attention_core sage2: artifact advertised no head dims") + return dims + + +class DenseAttentionSage2(GuardedSeam, torch.nn.Module): + """sage2 replacement for an ordinary dense unmasked SDPA call. + + Inputs and outputs use the host SDPA layout ``[B, H, S, D]``; the kernel + consumes NHD ``[B, S, H, D]``. Quantization runs per call inside the + artifact against the caller-owned workspace, so repeated calls launch an + identical sequence on identical pointers. + """ + + def __init__(self, q_shape, kv_shape, dtype: torch.dtype, device, + variant: str = "pv_fp8", + qk_quant_granularity: str = "per_warp"): + super().__init__() + if variant not in _VARIANTS: + raise ValueError( + f"attention_core sage2: unknown variant {variant!r} " + f"(expected one of {_VARIANTS})") + b, heads, seq_q, head_dim = q_shape + kb, kv_heads, seq_kv, kv_dim = kv_shape + if kb != b or kv_dim != head_dim: + raise ValueError( + "attention_core sage2: Q and KV batch/head dims differ") + if kv_heads != heads: + raise ValueError( + "attention_core sage2: GQA sites are not claimed by this " + "form yet; query and KV head counts must match") + if head_dim not in supported_head_dims(): + raise ValueError( + f"attention_core sage2: head_dim {head_dim} outside the " + f"artifact envelope {supported_head_dims()}") + if dtype != torch.bfloat16: + raise ValueError( + "attention_core sage2: the artifact consumes bf16 inputs") + self.q_shape = tuple(q_shape) + self.kv_shape = tuple(kv_shape) + self.variant = variant + self.granularity = qk_quant_granularity + + art = _artifact() + self._fn = (art.sage2_prefill_fp8v_bf16_d128 if variant == "pv_fp8" + else art.sage2_prefill_f16_bf16_d128) + staging = _staging_for( + (b, seq_q, heads, head_dim), (b, seq_kv, kv_heads, head_dim), + dtype, device, variant, qk_quant_granularity) + # Held as plain attributes, not buffers: this scratch belongs to + # the shared pool, and registering it would make every bound seam + # claim the same storage in its own state_dict. + self._q_nhd, self._k_nhd = staging.q, staging.k + self._v_nhd, self._out_nhd = staging.v, staging.out + self._workspace = staging.workspace + self._frt_arm( + dtypes=(dtype,), device=torch.device(device), + k=int(head_dim), rows=int(b * heads * seq_q)) + + def forward(self, query, key, value, *, scale=None): + admitted = self._frt_admit(query) + if admitted is not PROCEED: + return admitted + # BHSD -> NHD staging copies (fused away once the host adopts the + # NHD projection layout; kept explicit and pointer-stable here). + self._q_nhd.copy_(query.transpose(1, 2)) + self._k_nhd.copy_(key.transpose(1, 2)) + self._v_nhd.copy_(value.transpose(1, 2)) + self._fn( + self._q_nhd, self._k_nhd, self._v_nhd, + softmax_scale=scale, out=self._out_nhd, + workspace=self._workspace, + qk_quant_granularity=self.granularity) + return self._out_nhd.transpose(1, 2) + + +def bind_dense_attention(captures, *, variant: str = "pv_fp8", + qk_quant_granularity: str = "per_warp"): + """Bind one stateless dense sage2 core from repeated host captures. + + ``captures`` is the family's own convention -- a sequence of per-call + dicts holding ``q``, ``key``, ``value`` and ``mask`` in host layout -- + so this form qualifies a site the same way its BF16 siblings do: + the shape, dtype and mask must not move across the calibration call, + an unsupported shape returns ``None`` for the caller to keep its own + path, and a device the package does not serve raises so the family + binder can record it and move to the next rung. + """ + if not captures: + raise ValueError("attention_core sage2: no captures") + first = captures[0] + query, key, value = first["q"], first["key"], first["value"] + if first.get("mask") is not None: + # No packed-KV plan transfers here: the quantizers consume dense + # NHD, so a masked site is not this form's, and saying so is what + # keeps the host's own attention on it. + return None + if query.shape[-1] not in supported_head_dims(): + return None + if tuple(key.shape) != tuple(value.shape): + return None + if key.shape[0] != query.shape[0] or key.shape[1] != query.shape[1]: + # A grouped-query site is a shape this form does not claim, which + # is an answer the caller acts on by keeping its own attention -- + # not an error. The constructor still raises on it, because + # reaching it with such a shape would be this module's bug. + return None + expected = (tuple(query.shape), tuple(key.shape), tuple(value.shape), + query.dtype, key.dtype, value.dtype) + for capture in captures[1:]: + got = (tuple(capture["q"].shape), tuple(capture["key"].shape), + tuple(capture["value"].shape), capture["q"].dtype, + capture["key"].dtype, capture["value"].dtype) + if got != expected: + raise ValueError( + "attention_core sage2: shape or dtype moved within one " + f"calibration call: {expected} -> {got}") + if capture.get("mask") is not None: + raise ValueError( + "attention_core sage2: a mask appeared within one " + "calibration call") + if not (query.dtype == key.dtype == value.dtype): + raise ValueError("attention_core sage2: Q/K/V dtypes differ") + if query.dtype != torch.bfloat16: + return None + return DenseAttentionSage2( + query.shape, key.shape, query.dtype, query.device, + variant=variant, qk_quant_granularity=qk_quant_granularity) diff --git a/flash_rt/structures/impls/attention_core/sage3_blackwell.py b/flash_rt/structures/impls/attention_core/sage3_blackwell.py new file mode 100644 index 00000000..de8393b9 --- /dev/null +++ b/flash_rt/structures/impls/attention_core/sage3_blackwell.py @@ -0,0 +1,164 @@ +"""attention_core — the sage3 (Blackwell FP4-QK/PV) dense form. + +The speed point below the sage2 family: the same stateless seam — complete +Q/K/V every call, host SDPA layout — executed by the +``flashrt/sageattention3-blackwell`` artifact's fused entry: K centered and +FP4-quantized in one kernel, Q group-mean-centered and quantized likewise, +the mean contribution restored exactly through a small delta GEMM, then one +blockscaled E2M1 attention, BF16 out. Centering is a mathematical +invariant; the accuracy trade lives entirely in the FP4 quantization of +QK and PV, and the artifact declares it: ``accuracy_profile`` is read at +bind and carried on the module, so downstream gates judge this form +against the speed-first band it claims rather than the band the INT8 +forms occupy. Family order by accuracy is sage2 pv_fp16, sage2 pv_fp8, +then this form; by speed the reverse. + +Qualification, decided from the artifact and the captures, refusal legible: + +- the artifact's attention is self-attention: one shape for Q, K, and V. + Cross-attention and GQA sites return no binding, +- masked sites are not claimed, same reasoning as the sage2 form: the + quantizers consume dense NHD and a packed-KV plan does not transfer, +- head_dim must be advertised by the artifact; other dims return no + binding so the host keeps its own attention, +- sequence padding to the artifact's token alignment happens inside the + fused workspace — any length binds, the padded tail never escapes, +- the workspace (FP4 staging, scales, centered-K, group means, delta, + output) is caller-owned and allocated once at bind — call sequences + are pointer-stable and the artifact declares itself CUDA-graph safe. +""" + +from __future__ import annotations + +import torch + +from .. import hub_kernel +from ...guard import PROCEED, GuardedSeam + +KERNEL_DEP = { + "provider": "huggingface_kernels", + "repo": "flashrt/sageattention3-blackwell", + "version": ">=1", +} + + +def _artifact(): + return hub_kernel(KERNEL_DEP["repo"], KERNEL_DEP["version"]) + + +def supported_head_dims() -> tuple[int, ...]: + """Executable envelope, read from the artifact — never duplicated here.""" + caps = _artifact().capabilities() + dims = tuple(sorted(int(d) for d in caps["head_dims"])) + if not dims: + raise ValueError( + "attention_core sage3: artifact advertised no head dims") + return dims + + +class DenseAttentionSage3(GuardedSeam, torch.nn.Module): + """sage3 replacement for an ordinary dense unmasked self-attention call. + + Inputs and outputs use the host SDPA layout ``[B, H, S, D]``; the + artifact consumes NHD ``[B, S, H, D]``. Centering, quantization, the + delta correction, and attention all run inside the artifact's fused + entry against the caller-owned workspace, so repeated calls launch an + identical sequence on identical pointers. + """ + + def __init__(self, q_shape, kv_shape, dtype: torch.dtype, device): + super().__init__() + b, heads, seq_q, head_dim = q_shape + if tuple(kv_shape) != tuple(q_shape): + raise ValueError( + "attention_core sage3: the artifact is a self-attention " + "form; Q and KV shapes must match (cross-attention and " + "GQA sites are not claimed)") + if head_dim not in supported_head_dims(): + raise ValueError( + f"attention_core sage3: head_dim {head_dim} outside the " + f"artifact envelope {supported_head_dims()}") + if dtype != torch.bfloat16: + raise ValueError( + "attention_core sage3: the artifact consumes bf16 inputs") + art = _artifact() + caps = art.capabilities() + if not caps.get("fused_prep"): + raise ValueError( + "attention_core sage3: installed artifact predates the " + "fused-prep entry this form is built on") + self.q_shape = tuple(q_shape) + self.kv_shape = tuple(kv_shape) + #: accuracy band the artifact claims for itself; gates judge this + #: form against the band it declares, not the INT8 family's. + self.accuracy_profile = str(caps.get("accuracy_profile", "")) + self._fn = art.sage3_prefill_fp4_bf16 + # NHD staging + the fused caller-owned workspace, one set per + # bound shape. Padding to the artifact's token alignment lives + # inside the workspace; the entry returns the unpadded view. + self.register_buffer("_q_nhd", torch.empty( + b, seq_q, heads, head_dim, dtype=dtype, device=device)) + self.register_buffer("_k_nhd", torch.empty_like(self._q_nhd)) + self.register_buffer("_v_nhd", torch.empty_like(self._q_nhd)) + self._workspace = art.allocate_fused_workspace( + self._q_nhd, self._k_nhd, self._v_nhd) + self._frt_arm( + dtypes=(dtype,), device=torch.device(device), + k=int(head_dim), rows=int(b * heads * seq_q)) + + def forward(self, query, key, value, *, scale=None): + admitted = self._frt_admit(query) + if admitted is not PROCEED: + return admitted + # BHSD -> NHD staging copies (fused away once the host adopts the + # NHD projection layout; kept explicit and pointer-stable here). + self._q_nhd.copy_(query.transpose(1, 2)) + self._k_nhd.copy_(key.transpose(1, 2)) + self._v_nhd.copy_(value.transpose(1, 2)) + out = self._fn( + self._q_nhd, self._k_nhd, self._v_nhd, + softmax_scale=scale, workspace=self._workspace) + return out.transpose(1, 2) + + +def bind_dense_attention(captures): + """Bind one stateless dense sage3 core from repeated host captures. + + Same family convention as its siblings -- a sequence of per-call dicts + in host layout -- and the same division of answers: ``None`` for a site + this form does not claim, an exception when the calibration itself is + inconsistent or the package will not serve the device. + + This form claims less than the INT8 one. The artifact's attention is + self-attention, so a cross-attention site (K/V shaped differently from + Q) is not its shape, and neither is a GQA site. + """ + if not captures: + raise ValueError("attention_core sage3: no captures") + first = captures[0] + query, key, value = first["q"], first["key"], first["value"] + if first.get("mask") is not None: + return None + if query.shape[-1] not in supported_head_dims(): + return None + if tuple(key.shape) != tuple(query.shape) or \ + tuple(value.shape) != tuple(query.shape): + return None + expected = (tuple(query.shape), query.dtype, key.dtype, value.dtype) + for capture in captures[1:]: + got = (tuple(capture["q"].shape), capture["q"].dtype, + capture["key"].dtype, capture["value"].dtype) + if got != expected: + raise ValueError( + "attention_core sage3: shape or dtype moved within one " + f"calibration call: {expected} -> {got}") + if capture.get("mask") is not None: + raise ValueError( + "attention_core sage3: a mask appeared within one " + "calibration call") + if not (query.dtype == key.dtype == value.dtype): + raise ValueError("attention_core sage3: Q/K/V dtypes differ") + if query.dtype != torch.bfloat16: + return None + return DenseAttentionSage3( + query.shape, key.shape, query.dtype, query.device) diff --git a/flash_rt/structures/schemes.py b/flash_rt/structures/schemes.py index e684b39e..bd870076 100644 --- a/flash_rt/structures/schemes.py +++ b/flash_rt/structures/schemes.py @@ -121,6 +121,18 @@ class QuantScheme: #: quantising those projections is a precision decision. gdn_projection_format: str | None = None + #: Executable forms of ``attention_core`` to weigh ahead of the + #: family's published order, most preferred first. Empty keeps that + #: order, which is precision-descending: the BF16 forms preserve the + #: host's numerics and every existing receipt was measured against + #: them. A quantized attention form trades a bounded error for + #: speed, so it belongs to the same axis as the two attributes above + #: and arrives the same way — named by a scheme the caller selected, + #: never by a device check or a host binding. The names are impl + #: variants, exactly as in ``Decision.formats``; the family still + #: qualifies and speed-gates whatever is named. + attention_forms: tuple[str, ...] = () + def statistics(self, points: Sequence) -> dict[str, PointStat]: """Per point key (``"path|name"``): what to measure there.""" return {f"{p.path}|{p.name}": PointStat() for p in points} @@ -298,9 +310,20 @@ class Nvfp4Balance(QuantScheme): def __init__(self, alpha: float = 0.5, clamp: tuple[float, float] = (0.25, 4.0), - fuse_ffn_wire: bool = False) -> None: + fuse_ffn_wire: bool = False, + attention_forms: tuple[str, ...] = ()) -> None: self.alpha = float(alpha) self.clamp = (float(clamp[0]), float(clamp[1])) + # Naming quantized attention forms is the same kind of statement + # this scheme already makes about GEMM seams, on the same axis: + # the family still qualifies and speed-gates each one, and an + # unserved device falls through to the published order. It is a + # separate registered profile rather than a default here because + # it changes the numerics of a seam this scheme otherwise leaves + # at host precision. + self.attention_forms = tuple(attention_forms) + if attention_forms: + self.name = "nvfp4_balance_sage" # the FFN's FP4-wire chain (GEMM emits bias+GELU re-quantized, # the second GEMM consumes it) drops fc2's input-side balance — # a numerics change, so it is a scheme decision the receipt @@ -488,3 +511,5 @@ def resolve_auto() -> str: register("nvfp4_awq", Nvfp4Awq()) register("nvfp4_balance", Nvfp4Balance()) register("nvfp4_balance_wire", Nvfp4Balance(fuse_ffn_wire=True)) +register("nvfp4_balance_sage", Nvfp4Balance( + attention_forms=("sage2", "sage3"))) diff --git a/tests/test_gate_records_shape_and_memory.py b/tests/test_gate_records_shape_and_memory.py new file mode 100644 index 00000000..ac53297e --- /dev/null +++ b/tests/test_gate_records_shape_and_memory.py @@ -0,0 +1,170 @@ +"""What a gate verdict records: the shape it holds for, and its memory. + +Both are recording, not deciding. ``min_speedup`` is untouched, no new +refusal path exists, and every verdict this produces is the verdict the +gate produced before — with two facts attached that were previously lost: +the shape the measurement was taken at, and what the unit does to resident +memory. A refusal that reads as free may be the most expensive line in a +run, and until it says so that cannot even be discussed. +""" + +import pytest + +pytest.importorskip("torch") + +import torch # noqa: E402 + +from flash_rt.structures import frontdoor # noqa: E402 + + +# -------------------------------------------------------------------- +# the per-shape rule the specs have always declared +# -------------------------------------------------------------------- + +def test_specs_declaring_per_shape_are_read(): + """Thirteen specs declare it; the code has never looked at one.""" + from flash_rt.structures.registry import list_structures, load + + declared = [n for n in list_structures() + if ((load(n).gates or {}).get("latency") or {}).get("per_shape")] + assert declared, "no spec declares latency.per_shape" + for name in declared: + assert frontdoor._per_shape_rule(name) is True, name + + +def test_a_routed_unit_resolves_to_its_structure(): + """The gate's unit name is not always the structure's name.""" + assert frontdoor._per_shape_rule("attention_core_routed") is True + + +def test_an_unknown_unit_claims_nothing(): + assert frontdoor._per_shape_rule("not_a_structure") is False + + +def test_refusal_reasons_carry_the_shape_and_the_scope(): + """A verdict measured at one shape must not read as universal.""" + import inspect + + source = inspect.getsource(frontdoor.attach) + assert "per_shape_rule" in source + assert "holds for that shape and no other" in source + + +# -------------------------------------------------------------------- +# the shape a verdict was measured at +# -------------------------------------------------------------------- + +class _Guarded(torch.nn.Module): + def __init__(self, rows=None, capacity=None, bytes_=0): + super().__init__() + self._frt_guard = type("G", (), {"rows": rows, + "row_capacity": capacity})() + if bytes_: + self.register_buffer("w", torch.empty(bytes_, dtype=torch.uint8)) + + +def test_shape_comes_from_the_guards_when_the_binding_is_silent(): + """``m_profile`` is empty for hosts whose binding does not declare it, + and the refusal then read "rows unrecorded" — a shape-scoped verdict + with no shape on it. The guards always knew.""" + from flash_rt.structures.autobuild import AutoPlan + + paths = {"a": _Guarded(rows=4032), "b": _Guarded(rows=786432)} + note = frontdoor._measured_shape(paths, AutoPlan()) + assert "4032" in note and "786432" in note + + +def test_a_capacity_armed_seam_says_so(): + from flash_rt.structures.autobuild import AutoPlan + + note = frontdoor._measured_shape({"a": _Guarded(capacity=2048)}, + AutoPlan()) + assert note == "rows<=2048" + + +def test_shape_falls_back_to_the_binding_profile(): + from flash_rt.structures.autobuild import AutoPlan + from flash_rt.structures.discover import Seam + + plan = AutoPlan() + plan.seams.append(Seam(structure="vision_ffn", path="p", parent_path="", + norm_attr=None, dims={}, variant={}, + m_profile={"decode": {}})) + assert "decode" in frontdoor._measured_shape({}, plan) + + +# -------------------------------------------------------------------- +# what a unit does to resident memory +# -------------------------------------------------------------------- + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a device") +def test_a_replacement_reports_what_it_saves(): + """A quantized form replacing host weights must read as a saving. + + The retained host lives inside its replacement for fallback, so a walk + that counts it on both sides reports a form that halves the weights as + costing nothing. + """ + host = _Guarded(bytes_=8 << 20).cuda() + bound = _Guarded(rows=8, bytes_=2 << 20).cuda() + bound.host = host # retained for fallback + bound._frt_host = lambda: host + + memory = frontdoor._memory_delta({"p": bound}) + assert memory["host_bytes"] == 8 << 20 + assert memory["bound_bytes"] == 2 << 20 + assert memory["delta_bytes"] == -(6 << 20) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a device") +def test_a_form_that_brings_a_working_set_reports_it(): + """Nothing replaced, so everything it holds is new resident memory.""" + core = _Guarded(rows=8, bytes_=4 << 20).cuda() + memory = frontdoor._memory_delta({"p": core}) + assert memory["delta_bytes"] == 4 << 20 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a device") +def test_a_pooled_working_set_is_counted_once(): + """Sites sharing one staging set hold it once, not once each.""" + shared = torch.empty(4 << 20, dtype=torch.uint8, device="cuda") + + cores = {} + for name in ("a", "b", "c"): + core = _Guarded(rows=8).cuda() + core.staging = shared # a plain attribute, as pooled + cores[name] = core + memory = frontdoor._memory_delta(cores) + assert memory["bound_bytes"] == 4 << 20 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a device") +def test_tensors_held_inside_a_workspace_object_are_found(): + """A form's largest holding is often an artifact's workspace object, + which is neither a parameter nor a buffer.""" + class _Workspace: + def __init__(self): + self.packed = torch.empty(3 << 20, dtype=torch.uint8, + device="cuda") + + core = _Guarded(rows=8).cuda() + core.workspace = _Workspace() + assert frontdoor._memory_delta({"p": core})["bound_bytes"] == 3 << 20 + + +def test_small_memory_moves_stay_out_of_the_reason(): + """The note exists to make a large consequence visible, not to add + noise to every refusal.""" + assert frontdoor._memory_note({"delta_bytes": 1 << 20}) == "" + assert "GiB" in frontdoor._memory_note({"delta_bytes": 3 << 30}) + assert frontdoor._memory_note({}) == "" + + +def test_recording_does_not_change_any_threshold(): + """The decision is the same scalar comparison it always was.""" + import inspect + + source = inspect.getsource(frontdoor.attach) + assert 'timing["speedup"] < min_speedup' in source + assert "delta_bytes" not in source.split("# ---- per unit")[1].split( + "continue")[0].replace("_memory_delta(paths)", "") diff --git a/tests/test_sage_attention_forms.py b/tests/test_sage_attention_forms.py new file mode 100644 index 00000000..39261891 --- /dev/null +++ b/tests/test_sage_attention_forms.py @@ -0,0 +1,597 @@ +"""Contracts of the sage executable forms of ``attention_core``. + +Both forms qualify a site from the family's own captures and answer the +family's way: ``None`` for a site they do not claim, an exception when the +calibration is inconsistent or the package will not serve the device. These +cases need no device and no kernel package -- what they check is the +qualification walk and the selection semantics, which is where the two forms +join the public family and therefore where a mistake reaches other hosts. + +Numerical and latency evidence for the forms themselves is a hardware +qualification and lives with the model's benchmark runs, not here. +""" + +import pytest + +pytest.importorskip("torch") + +import torch # noqa: E402 + +from flash_rt.structures.impls import attention_core # noqa: E402 +from flash_rt.structures.impls.attention_core import ( # noqa: E402 + sage2_blackwell, sage3_blackwell) + +FORMS = (("sage2", sage2_blackwell), ("sage3", sage3_blackwell)) + + +def _captures(b=1, heads=32, seq=2688, dim=128, kv_seq=None, + kv_heads=None, dtype=torch.bfloat16, mask=None, n=1): + """Family-shaped captures: per-call dicts in host SDPA layout.""" + kv_seq = seq if kv_seq is None else kv_seq + kv_heads = heads if kv_heads is None else kv_heads + row = { + "q": torch.empty(b, heads, seq, dim, dtype=dtype, device="meta"), + "key": torch.empty(b, kv_heads, kv_seq, dim, dtype=dtype, + device="meta"), + "value": torch.empty(b, kv_heads, kv_seq, dim, dtype=dtype, + device="meta"), + "mask": mask, + } + return [dict(row) for _ in range(n)] + + +# -------------------------------------------------------------------- +# the family contract: what a binder is handed and what it answers +# -------------------------------------------------------------------- + +@pytest.mark.parametrize("name,module", FORMS) +def test_binder_takes_the_family_capture_convention(name, module): + """A sequence of capture dicts, not a shape-carrying object. + + This is the whole reason these forms are reachable from the family at + all, and it is the kind of mismatch that stays invisible until a host + actually routes through it: the earlier version of these binders read + attributes off a single object and raised AttributeError the first time + the family called them. + """ + captures = _captures(dim=96) # a dim no artifact advertises + assert module.bind_dense_attention(captures) is None + + +@pytest.mark.parametrize("name,module", FORMS) +def test_empty_captures_are_a_caller_error(name, module): + with pytest.raises(ValueError, match="no captures"): + module.bind_dense_attention([]) + + +@pytest.mark.parametrize("name,module", FORMS) +def test_masked_sites_are_not_claimed(name, module): + """A mask has no form these kernels accept, and the packed-KV plan of + the BF16 form does not transfer -- the quantizers consume dense NHD.""" + mask = torch.zeros(1, 1, 8, 8, device="meta") + assert module.bind_dense_attention(_captures(mask=mask)) is None + + +@pytest.mark.parametrize("name,module", FORMS) +def test_non_bf16_sites_are_not_claimed(name, module): + assert module.bind_dense_attention( + _captures(dtype=torch.float16)) is None + + +@pytest.mark.parametrize("name,module", FORMS) +def test_moving_shape_within_one_calibration_raises(name, module): + captures = _captures(n=2) + captures[1]["q"] = torch.empty(1, 32, 1344, 128, dtype=torch.bfloat16, + device="meta") + with pytest.raises(ValueError, match="moved within one"): + module.bind_dense_attention(captures) + + +@pytest.mark.parametrize("name,module", FORMS) +def test_a_mask_appearing_mid_calibration_raises(name, module): + captures = _captures(n=2) + captures[1]["mask"] = torch.zeros(1, 1, 8, 8, device="meta") + with pytest.raises(ValueError, match="mask appeared"): + module.bind_dense_attention(captures) + + +def test_sage3_does_not_claim_cross_attention(): + """The artifact's attention is self-attention: one shape for Q/K/V.""" + assert sage3_blackwell.bind_dense_attention( + _captures(kv_seq=1024)) is None + + +def test_sage3_does_not_claim_grouped_query_sites(): + assert sage3_blackwell.bind_dense_attention( + _captures(kv_heads=8)) is None + + +def test_sage2_does_not_claim_mismatched_kv(): + assert sage2_blackwell.bind_dense_attention( + _captures(kv_seq=1024, kv_heads=8)) is None + + +# -------------------------------------------------------------------- +# selection: preferring a form is a decision, never a default +# -------------------------------------------------------------------- + +def test_quantized_forms_are_absent_from_the_default_order(monkeypatch): + """The default ladder must stay precision-first for every host. + + These forms trade a bounded numerical error for speed. That is a + deployment decision, so it cannot ride in on a family default: a host + that never asked for it must keep the numerics-preserving order it has + receipts for. + """ + tried = [] + + def spy(name): + def binder(captures): + tried.append(name) + return None + return binder + + monkeypatch.setattr(attention_core, "bind_dense_attention", spy("fa2")) + for name, module in FORMS: + monkeypatch.setattr(module, "bind_dense_attention", spy(name)) + attention_core.bind_dense_attention_best(_captures()) + assert "fa2" in tried + assert "sage2" not in tried and "sage3" not in tried + + +@pytest.mark.parametrize("name", ["sage2", "sage3"]) +def test_a_preferred_form_is_tried_before_the_order(name, monkeypatch): + tried = [] + + def spy(label): + def binder(captures): + tried.append(label) + return None + return binder + + monkeypatch.setattr(attention_core, "bind_dense_attention", spy("fa2")) + for form_name, module in FORMS: + monkeypatch.setattr(module, "bind_dense_attention", spy(form_name)) + attention_core.bind_dense_attention_best(_captures(), prefer=(name,)) + assert tried[0] == name, tried + + +def test_preference_order_is_the_callers_order(monkeypatch): + tried = [] + + def spy(label): + def binder(captures): + tried.append(label) + return None + return binder + + monkeypatch.setattr(attention_core, "bind_dense_attention", spy("fa2")) + for form_name, module in FORMS: + monkeypatch.setattr(module, "bind_dense_attention", spy(form_name)) + attention_core.bind_dense_attention_best( + _captures(), prefer=("sage3", "sage2")) + assert tried[:2] == ["sage3", "sage2"], tried + + +def test_an_unknown_preferred_form_is_an_error(): + """Silently ignoring it would report the default ladder's result as + though the caller's choice had been honoured.""" + with pytest.raises(ValueError, match="unknown preferred form"): + attention_core.bind_dense_attention_best( + _captures(), prefer=("sage9",)) + + +# -------------------------------------------------------------------- +# the host-family adapter +# -------------------------------------------------------------------- + +def _attention_module(rope_type="interleaved", gate=True, heads=32): + class _Attn(torch.nn.Module): + def __init__(self): + super().__init__() + self.to_q = torch.nn.Linear(4, 4) + self.to_k = torch.nn.Linear(4, 4) + self.to_v = torch.nn.Linear(4, 4) + self.norm_q = torch.nn.LayerNorm(4) + self.norm_k = torch.nn.LayerNorm(4) + self.to_out = torch.nn.ModuleList( + [torch.nn.Linear(4, 4), torch.nn.Dropout(0.0)]) + self.heads = heads + self.rope_type = rope_type + self.to_gate_logits = torch.nn.Linear(4, 4) if gate else None + + return _Attn() + + +class _GatedRotaryProcessor: + def __call__(self, attn, hidden_states, encoder_hidden_states=None, + attention_mask=None, query_rotary_emb=None, + key_rotary_emb=None): + return hidden_states + + +class _SingleRotaryProcessor: + def __call__(self, attn, hidden_states, encoder_hidden_states=None, + attention_mask=None, rotary_emb=None): + return hidden_states + + +def test_adapter_recognises_the_gated_dual_rotary_contract(): + from flash_rt.structures.adapters.diffusers_gated_rotary_attention import ( + _compatible_site) + + module = _attention_module() + ok, reason = _compatible_site(module, _GatedRotaryProcessor()) + assert ok, reason + + +@pytest.mark.parametrize("mutate,expected", [ + (lambda m: setattr(m, "to_gate_logits", None) or m, True), + (lambda m: delattr(m, "to_gate_logits") or m, False), + (lambda m: setattr(m, "rope_type", "unknown") or m, False), + (lambda m: setattr(m, "heads", 0) or m, False), + (lambda m: setattr(m, "norm_q", None) or m, False), +]) +def test_adapter_negative_recognition(mutate, expected): + """Recognition is structural, and each missing slot refuses by name.""" + from flash_rt.structures.adapters.diffusers_gated_rotary_attention import ( + _compatible_site) + + module = mutate(_attention_module()) + ok, reason = _compatible_site(module, _GatedRotaryProcessor()) + assert ok is expected + if not ok: + assert reason + + +def test_adapter_declines_a_single_rotary_processor(): + """The sibling rotary family has its own adapter; this one must not + claim its sites just because the module slots look alike.""" + from flash_rt.structures.adapters.diffusers_gated_rotary_attention import ( + _compatible_site) + + ok, reason = _compatible_site(_attention_module(), + _SingleRotaryProcessor()) + assert not ok + assert "rotary_emb" in reason, reason + + +def test_adapter_returns_nothing_on_a_host_without_such_sites(): + """A host that is not this family gets no adapter result at all.""" + from flash_rt.structures.adapters import ( + DiffusersGatedRotaryAttentionAdapter) + + model = torch.nn.Sequential(torch.nn.Linear(4, 4)) + assert DiffusersGatedRotaryAttentionAdapter()( + model, lambda: None) is None + + +def test_adapter_passes_the_preference_to_the_family(monkeypatch): + from flash_rt.structures.adapters import ( + DiffusersGatedRotaryAttentionAdapter) + import flash_rt.structures.adapters.diffusers_gated_rotary_attention \ + as adapter_module + + seen = {} + + def fake_bind(captures, *, prefer=()): + seen["prefer"] = prefer + return None + + monkeypatch.setattr(adapter_module, "bind_dense_attention_best", + fake_bind) + + module = _attention_module(heads=2) # head_dim 2 over the 4-wide stub + module.processor = _GatedRotaryProcessor() + model = torch.nn.Module() + model.attn = module + + def forward(): + module.processor(module, torch.zeros(1, 2, 4)) + + DiffusersGatedRotaryAttentionAdapter(("sage2",))(model, forward) + assert seen["prefer"] == ("sage2",) + + +def test_adapter_default_prefers_nothing(): + from flash_rt.structures.adapters import ( + DiffusersGatedRotaryAttentionAdapter) + + assert DiffusersGatedRotaryAttentionAdapter().prefer == () + + +# -------------------------------------------------------------------- +# the precision axis: a scheme names the forms, nothing else does +# -------------------------------------------------------------------- + +def test_default_schemes_name_no_attention_forms(): + """Every existing profile must keep the published order. + + This is the property that makes the new attribute safe to add: a host + that selected any scheme before this change gets exactly the ladder it + was measured against. + """ + from flash_rt.structures import schemes + + for name in schemes.names(): + if name == "nvfp4_balance_sage": + continue + scheme = schemes.get(name) + assert getattr(scheme, "attention_forms", ()) == (), name + + +def test_the_quantized_profile_names_them(): + from flash_rt.structures import schemes + + scheme = schemes.get("nvfp4_balance_sage") + assert scheme.attention_forms == ("sage2", "sage3") + assert scheme.name == "nvfp4_balance_sage" + + +def test_adapter_takes_the_forms_from_the_scheme(monkeypatch): + from flash_rt.structures import schemes + from flash_rt.structures.adapters import ( + DiffusersGatedRotaryAttentionAdapter) + import flash_rt.structures.adapters.diffusers_gated_rotary_attention \ + as adapter_module + + seen = {} + + def fake_bind(captures, *, prefer=()): + seen["prefer"] = prefer + return None + + monkeypatch.setattr(adapter_module, "bind_dense_attention_best", + fake_bind) + module = _attention_module(heads=2) + module.processor = _GatedRotaryProcessor() + model = torch.nn.Module() + model.attn = module + + def forward(): + module.processor(module, torch.zeros(1, 2, 4)) + + DiffusersGatedRotaryAttentionAdapter()( + model, forward, scheme=schemes.get("nvfp4_balance_sage")) + assert seen["prefer"] == ("sage2", "sage3") + + +def test_a_scheme_without_forms_leaves_the_order_alone(monkeypatch): + from flash_rt.structures import schemes + from flash_rt.structures.adapters import ( + DiffusersGatedRotaryAttentionAdapter) + import flash_rt.structures.adapters.diffusers_gated_rotary_attention \ + as adapter_module + + seen = {} + + def fake_bind(captures, *, prefer=()): + seen["prefer"] = prefer + return None + + monkeypatch.setattr(adapter_module, "bind_dense_attention_best", + fake_bind) + module = _attention_module(heads=2) + module.processor = _GatedRotaryProcessor() + model = torch.nn.Module() + model.attn = module + + def forward(): + module.processor(module, torch.zeros(1, 2, 4)) + + DiffusersGatedRotaryAttentionAdapter()( + model, forward, scheme=schemes.get("nvfp4_balance")) + assert seen["prefer"] == () + + +# -------------------------------------------------------------------- +# the shared staging pool +# -------------------------------------------------------------------- + +def test_staging_is_shared_per_shape(monkeypatch): + """Seams with the same form and shape must share one scratch set. + + The scratch is the size of Q/K/V at the bound shape, and a transformer + reaches this seam once per block: at long sequence lengths, private + scratch per seam does not fit on a consumer part at all. It is safe to + share because nothing survives a call -- every buffer is written at the + top of the forward and read before it returns. + """ + from flash_rt.structures.impls.attention_core import sage2_blackwell as s2 + + made = [] + + class _Art: + def capabilities(self): + return {"head_dims": (128,)} + + def allocate_workspace(self, q, k, v, **kw): + made.append((tuple(q.shape), kw.get("fp8v"))) + return object() + + monkeypatch.setattr(s2, "_artifact", lambda: _Art()) + monkeypatch.setattr(s2, "_STAGING", {}) + + a = s2._staging_for((1, 8, 2, 128), (1, 8, 2, 128), torch.bfloat16, + "meta", "pv_fp8", "per_warp") + b = s2._staging_for((1, 8, 2, 128), (1, 8, 2, 128), torch.bfloat16, + "meta", "pv_fp8", "per_warp") + assert a is b + assert len(made) == 1, "the second seam must not allocate again" + + +@pytest.mark.parametrize("differs", [ + {"q_shape": (2, 8, 2, 128)}, + {"variant": "pv_fp16"}, + {"granularity": "per_thread"}, + {"dtype": torch.float16}, +]) +def test_a_different_form_or_shape_gets_its_own_staging(differs, monkeypatch): + """Sharing is keyed, not global: nothing may hand back a buffer of the + wrong size, precision, or layout.""" + from flash_rt.structures.impls.attention_core import sage2_blackwell as s2 + + class _Art: + def capabilities(self): + return {"head_dims": (128,)} + + def allocate_workspace(self, q, k, v, **kw): + return object() + + monkeypatch.setattr(s2, "_artifact", lambda: _Art()) + monkeypatch.setattr(s2, "_STAGING", {}) + + base = dict(q_shape=(1, 8, 2, 128), kv_shape=(1, 8, 2, 128), + dtype=torch.bfloat16, device="meta", variant="pv_fp8", + granularity="per_warp") + first = s2._staging_for(**base) + second = s2._staging_for(**{**base, **differs, + **({"kv_shape": differs["q_shape"]} + if "q_shape" in differs else {})}) + assert first is not second + + +def test_staging_is_not_registered_as_module_state(monkeypatch): + """Shared scratch must not appear in a seam's state_dict. + + Registering it would have every bound seam claim the same storage as + its own parameter state, which is both wrong and silently divergent + once two seams share one pool. + """ + from flash_rt.structures.impls.attention_core import sage2_blackwell as s2 + + class _Art: + def capabilities(self): + return {"head_dims": (128,)} + + def allocate_workspace(self, q, k, v, **kw): + return object() + + sage2_prefill_fp8v_bf16_d128 = staticmethod(lambda *a, **k: None) + sage2_prefill_f16_bf16_d128 = staticmethod(lambda *a, **k: None) + + monkeypatch.setattr(s2, "_artifact", lambda: _Art()) + monkeypatch.setattr(s2, "_STAGING", {}) + core = s2.DenseAttentionSage2((1, 8, 2, 128), (1, 8, 2, 128), + torch.bfloat16, "meta") + assert core.state_dict() == {} + + +def test_an_explicit_argument_outranks_the_profile(monkeypatch): + """A caller tuning one seam should not have to register a scheme. + + The override is the same statement the profile makes, so it reaches + adapters the same way; what it must not do is silently lose to the + profile the caller left at its default. + """ + from flash_rt.structures import autobuild, schemes + from flash_rt.structures.adapters import ( + DiffusersGatedRotaryAttentionAdapter) + import flash_rt.structures.adapters.diffusers_gated_rotary_attention \ + as adapter_module + + override = autobuild._AttentionOverride(["sage3"]) + assert override.attention_forms == ("sage3",) + + seen = {} + + def fake_bind(captures, *, prefer=()): + seen["prefer"] = prefer + return None + + monkeypatch.setattr(adapter_module, "bind_dense_attention_best", + fake_bind) + module = _attention_module(heads=2) + module.processor = _GatedRotaryProcessor() + model = torch.nn.Module() + model.attn = module + + def forward(): + module.processor(module, torch.zeros(1, 2, 4)) + + DiffusersGatedRotaryAttentionAdapter()(model, forward, scheme=override) + assert seen["prefer"] == ("sage3",) + + +def test_the_front_door_takes_attention_forms(): + """The argument has to exist where a user actually calls in.""" + import inspect + + from flash_rt import structures + from flash_rt.structures import autobuild, frontdoor + + for fn in (frontdoor.attach, autobuild.auto_swaps): + assert "attention_forms" in inspect.signature(fn).parameters, fn + # the public name is a lazy passthrough; what matters is that it does + # not filter the argument out on the way + params = inspect.signature(structures.attach).parameters + assert any(p.kind is inspect.Parameter.VAR_KEYWORD + for p in params.values()) + + +# -------------------------------------------------------------------- +# a declined form must not keep its working set +# -------------------------------------------------------------------- + +def test_release_routed_drops_the_forms_and_keeps_revert_callable(): + """Reverting a routed seam is not releasing it. + + The host processor goes back either way; what differs is whether the + bound form is still reachable. It has to be while the gate might yet + activate it, and must not be once the gate has declined -- otherwise an + attachment costs more memory than the host it replaced, for a form that + never runs. + """ + from flash_rt.structures.autobuild import AutoPlan + + routes = [("module", "original", "routed")] + reverted = [] + + plan = AutoPlan() + plan.observed["site.processor"] = object() + plan.toggles.append((lambda: None, lambda: reverted.append("off"))) + plan.revert.append(lambda: reverted.append("revert")) + plan.releases.append(routes.clear) + + plan.release_routed() + assert routes == [], "the adapter's own hold must be dropped" + assert plan.observed == {} and plan.toggles == [] + assert reverted == ["off"], "release disables, it does not revert" + + plan.release_routed() # idempotent + assert plan.releases == [] + for undo in plan.revert: # still callable, nothing left to undo + undo() + assert reverted == ["off", "revert"] + + +def test_adapter_publishes_a_release(monkeypatch): + """The adapter has to offer the hold it wants dropped.""" + from flash_rt.structures.adapters import ( + DiffusersGatedRotaryAttentionAdapter) + import flash_rt.structures.adapters.diffusers_gated_rotary_attention \ + as adapter_module + + class _Core(torch.nn.Module): + def forward(self, q, k, v, **kw): + return q + + monkeypatch.setattr(adapter_module, "bind_dense_attention_best", + lambda captures, **kw: _Core()) + module = _attention_module(heads=2) + module.processor = _GatedRotaryProcessor() + model = torch.nn.Module() + model.attn = module + + def forward(): + module.processor(module, torch.zeros(1, 2, 4)) + + _, _, extras = DiffusersGatedRotaryAttentionAdapter()(model, forward) + assert extras.get("release"), "no way to give the bound forms back" + assert extras.get("observed") + for release in extras["release"]: + release() + # the toggles still work after a release: they close over the list, + # not over what was in it + enable, disable = extras["toggle"] + enable(); disable()