From d0a53d26b08a45c424937939e9fba77310a4541c Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 14 Aug 2026 04:12:46 -0400 Subject: [PATCH 01/11] docs(ltx25): structures-layer promotion design Map the LTX-2.5 integration onto catalog structures: a sage2 qk-int8/pv-fp8 attention_core backend, a W4A4 NVFP4 vision_ffn backend, the quantization site list as a quantize_on_adopt binding attribute, and a video_generation_pipeline binding. Records the measured qualification context and the sequencing. --- docs/ltx25_structures_design.md | 125 ++++++++++++++++++++++++++++++++ 1 file changed, 125 insertions(+) create mode 100644 docs/ltx25_structures_design.md 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) From 3a79e90ce057ea076d3046a73dd018c6eaf4adb9 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 14 Aug 2026 11:03:52 -0400 Subject: [PATCH 02/11] feat(structures): sage2 Blackwell dense attention form attention_core backend on the flashrt/sageattention2-blackwell artifact: INT8 per-warp/per-thread Q/K with FP8 per-channel or FP16 V, bf16 out. Caller-owned workspace allocated once per bound shape, capability envelope read from the artifact, masked and non-128 head-dim sites refuse with the reason on the binder. Measured against the host SDPA at the qualification shapes: 2.3x at the long site with per-call cosine 0.9992 (fp8 variant) / 0.9999 (fp16 variant); real-capture parity sits inside the noise floor established for same-precision kernel substitution. --- .../impls/attention_core/sage2_blackwell.py | 168 ++++++++++++++++++ 1 file changed, 168 insertions(+) create mode 100644 flash_rt/structures/impls/attention_core/sage2_blackwell.py 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..da2fa2ab --- /dev/null +++ b/flash_rt/structures/impls/attention_core/sage2_blackwell.py @@ -0,0 +1,168 @@ +"""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 + +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") + + +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) + # NHD staging + caller-owned workspace, one set per bound shape. + self.register_buffer("_q_nhd", torch.empty( + b, seq_q, heads, head_dim, dtype=dtype, device=device)) + self.register_buffer("_k_nhd", torch.empty( + b, seq_kv, kv_heads, head_dim, dtype=dtype, device=device)) + self.register_buffer("_v_nhd", torch.empty_like(self._k_nhd)) + self.register_buffer("_out_nhd", torch.empty_like(self._q_nhd)) + self._workspace = art.allocate_workspace( + self._q_nhd, self._k_nhd, self._v_nhd, + fp8v=(variant == "pv_fp8"), + qk_quant_granularity=qk_quant_granularity) + 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_sage2_dense_attention(captures, *, variant: str = "pv_fp8", + qk_quant_granularity: str = "per_warp"): + """Bind the sage2 dense form from one real capture set, or refuse. + + ``captures`` follows the family convention: an object carrying + ``q_shape``, ``kv_shape``, ``dtype``, ``device``, and optionally + ``allowed_ranges`` / ``mask``. Returns ``None`` (with the reason as an + attribute on the function, mirroring the family's refusal trail) when + the site is outside this form's envelope. + """ + def refuse(reason: str): + bind_sage2_dense_attention.last_refusal = reason + return None + + mask = getattr(captures, "mask", None) + ranges = tuple(getattr(captures, "allowed_ranges", ()) or ()) + if mask is not None or ranges: + return refuse("masked/packed sites are not claimed by sage2") + b, heads, seq_q, head_dim = captures.q_shape + try: + dims = supported_head_dims() + except (ValueError, OSError, RuntimeError) as exc: + return refuse(f"artifact unavailable: {exc}") + if head_dim not in dims: + return refuse( + f"head_dim {head_dim} outside artifact envelope {dims}") + if captures.dtype != torch.bfloat16: + return refuse(f"dtype {captures.dtype} outside envelope (bf16)") + try: + return DenseAttentionSage2( + captures.q_shape, captures.kv_shape, captures.dtype, + captures.device, variant=variant, + qk_quant_granularity=qk_quant_granularity) + except (ValueError, RuntimeError) as exc: + return refuse(str(exc)) From 4434a080bc0bb45da4cde525a5274ddb93fde755 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 14 Aug 2026 13:40:21 -0400 Subject: [PATCH 03/11] feat(structures): sage3 Blackwell dense attention form FP4 speed point of the attention_core family, built on the artifact's fused-prep entry: centering, quantization, delta correction, and attention in one caller-owned workspace. Self-attention only; masked, cross-shape, GQA, non-bf16, and out-of-envelope head dims refuse with the reason on the binder. Accuracy profile is read from the artifact and carried on the module for downstream gates. Pointer-stable call sequences; capture-verified. --- .../impls/attention_core/sage3_blackwell.py | 159 ++++++++++++++++++ 1 file changed, 159 insertions(+) create mode 100644 flash_rt/structures/impls/attention_core/sage3_blackwell.py 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..f9663e78 --- /dev/null +++ b/flash_rt/structures/impls/attention_core/sage3_blackwell.py @@ -0,0 +1,159 @@ +"""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_sage3_dense_attention(captures): + """Bind the sage3 dense form from one real capture set, or refuse. + + ``captures`` follows the family convention: an object carrying + ``q_shape``, ``kv_shape``, ``dtype``, ``device``, and optionally + ``allowed_ranges`` / ``mask``. Returns ``None`` (with the reason as an + attribute on the function, mirroring the family's refusal trail) when + the site is outside this form's envelope. + """ + def refuse(reason: str): + bind_sage3_dense_attention.last_refusal = reason + return None + + mask = getattr(captures, "mask", None) + ranges = tuple(getattr(captures, "allowed_ranges", ()) or ()) + if mask is not None or ranges: + return refuse("masked/packed sites are not claimed by sage3") + if tuple(captures.kv_shape) != tuple(captures.q_shape): + return refuse( + "self-attention form: Q and KV shapes differ at this site") + b, heads, seq_q, head_dim = captures.q_shape + try: + dims = supported_head_dims() + except (ValueError, OSError, RuntimeError) as exc: + return refuse(f"artifact unavailable: {exc}") + if head_dim not in dims: + return refuse( + f"head_dim {head_dim} outside artifact envelope {dims}") + if captures.dtype != torch.bfloat16: + return refuse(f"dtype {captures.dtype} outside envelope (bf16)") + try: + return DenseAttentionSage3( + captures.q_shape, captures.kv_shape, captures.dtype, + captures.device) + except (ValueError, RuntimeError) as exc: + return refuse(str(exc)) From 893c9531b47dd2768f6db24b5aa969e305936bfd Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 14 Aug 2026 13:45:55 -0400 Subject: [PATCH 04/11] feat(structures): LTX-2.5 host bindings Region binding maps the video feed-forward slots on the diffusers host family (zero-bias slots where the host has none; adaptive scale/shift stays outside the seam). Pipeline binding classifies the complete joint audio+video hot path: attention, feed-forward, and whitelisted projections as catalog structures; audio-branch and small-M sites retained on the host with the measured reason; the distilled single-pass recipe, encoder residency, and resident-aware decode tiling declared as binding attributes. Validates under the complete-hot-path contract. --- flash_rt/structures/bindings/ltx25_dit.yaml | 27 +++ .../bindings/ltx25_video_pipeline.yaml | 155 ++++++++++++++++++ 2 files changed, 182 insertions(+) create mode 100644 flash_rt/structures/bindings/ltx25_dit.yaml create mode 100644 flash_rt/structures/bindings/ltx25_video_pipeline.yaml 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" From c04705e81ef6281f8ab9b9ed9f9dc07bd042be91 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Fri, 14 Aug 2026 14:18:28 -0400 Subject: [PATCH 05/11] feat(structures): gated dual-rotary Diffusers attention adapter Host family whose processors take separate query/key rotary boundaries, RMS-normalise Q/K after projection, and gate the attention output with per-head sigmoid logits computed from the pre-attention hidden states. The stock Diffusers adapter's processor-state contract does not exist on these modules, so the family gets its own adapter: reproduce the projection half, capture real Q/K/V per called site, bind through the dense attention family, replace only the dispatch. Verified on a joint audio+video block at deployment shapes: attention unit refused by the net-win gate at short sequences and activated 1.257x at long ones, with the audio head-dim refusals legible on the trail. Also treat a hub package's unmet python-dependency declaration (ImportError from the kernels validator) as KernelUnavailable: a host missing a dependency cannot supply the package, and the family ladder falls through instead of aborting the bind. --- flash_rt/structures/adapters/__init__.py | 5 + .../diffusers_gated_rotary_attention.py | 222 ++++++++++++++++++ flash_rt/structures/impls/__init__.py | 7 +- 3 files changed, 233 insertions(+), 1 deletion(-) create mode 100644 flash_rt/structures/adapters/diffusers_gated_rotary_attention.py 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..fc9670bc --- /dev/null +++ b/flash_rt/structures/adapters/diffusers_gated_rotary_attention.py @@ -0,0 +1,222 @@ +"""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.""" + + __name__ = "diffusers_gated_rotary_attention" + + def __call__(self, model, forward, *, prefix_cadence: bool = False): + del prefix_cadence + 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) + 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 + + enable() + return {}, None, { + "revert": [disable], + "observed": observed, + "toggle": (enable, disable), + "refused": refused, + "attention_variants": variants, + } 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 " From 0b9e0923a3b678d573e6aabed5521cb2e81fb073 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sun, 16 Aug 2026 08:03:03 -0400 Subject: [PATCH 06/11] fix(structures): sage forms join the family through its own contract The two forms bound from a shape-carrying object while the family passes a sequence of per-call capture dicts, and neither was reachable from the family binder -- so they were unreachable code that would have raised the first time a host routed through them. They now qualify a site the way their BF16 siblings do: shape, dtype and mask must hold across the calibration call, an unclaimed shape answers None so the caller keeps its own attention, and an inconsistent calibration or an unserved device raises for the family to record. That walk found a real one: a grouped-query site reached the constructor and raised, where the contract says an unclaimed shape is an answer, not an error. Selection is explicit. bind_dense_attention_best gains a 'prefer' argument, empty by default, so the published order stays precision-first for every host that has receipts against it: these forms trade a bounded numerical error for speed, and that is a decision about a deployment rather than about a device. A caller who has judged the trade names the form and gets the same qualification walk, speed gate and refusal trail as any other rung; an unknown name is an error rather than a silent fallback to the default order. Measured on one transformer block with real weights and real captured inputs, paired inside the gate: at S=24576 the default order reaches 1.15x with the attention family bound but declined at 1.006x, and with sage2 preferred the block reaches 1.49x with the attention unit at 1.257x. Peak memory falls from 12.2GB to 8.2GB when the projections are quantized, and rises to the ceiling of a 32GB part when four attention sites each own a workspace -- pooling those is the open item before the preferred configuration is usable at that size. Cosine against the host is 0.99999 throughout and detach restores it bit-exactly. --- docs/ltx25_usage.md | 69 ++++ .../diffusers_gated_rotary_attention.py | 14 +- .../impls/attention_core/__init__.py | 33 +- .../impls/attention_core/sage2_blackwell.py | 83 +++-- .../impls/attention_core/sage3_blackwell.py | 73 +++-- tests/test_sage_attention_forms.py | 300 ++++++++++++++++++ 6 files changed, 499 insertions(+), 73 deletions(-) create mode 100644 tests/test_sage_attention_forms.py diff --git a/docs/ltx25_usage.md b/docs/ltx25_usage.md index d4f80633..f9877cae 100644 --- a/docs/ltx25_usage.md +++ b/docs/ltx25_usage.md @@ -121,3 +121,72 @@ 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. + +### Measured on one transformer block + +Real checkpoint weights, real captured deployment inputs, paired alternating +timing inside the gate, on a 5090. "Attention" is the gate's verdict for the +attention family; the projections are the `nvfp4_balance` W4A4 form. + +| Site shape | Configuration | Block latency | Attention family | Peak memory | +|---|---|---|---|---| +| S=24576 (1536×1024×121f) | host | 134.3 ms | — | 12.2 GB | +| | attach, default order | 117.1 ms (1.15×) | bound, declined at 1.006× | 8.2 GB | +| | attach, sage2 preferred | **90.1 ms (1.49×)** | activated, 1.257× | at the 32GB ceiling | +| S=2688 (768×512×49f) | host | 10.3 ms | — | 2.3 GB | +| | attach, default order | 8.2 ms (1.25×) | declined | 1.7 GB | +| | attach, sage2 preferred | 8.0 ms (1.28×) | activated | 4.8 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). Two results are worth +reading carefully rather than skipping: + +- **The default order does not use the quantized attention forms.** They trade + a bounded numerical error for speed, which is a deployment decision, so a + caller asks for one explicitly. Without that, the family's BF16 form binds, + and at these shapes the net-win gate measures it at 1.006× and keeps the + host's attention — the projections carry the whole win. +- **Peak memory falls when the projections are quantized** (12.2 → 8.2 GB) and + rises when quantized attention is preferred, because each attention site + owns its staging and quantization workspace. At S=24576 across four sites + that reaches the ceiling of a 32GB part; pooling those workspaces is the + open item before this configuration is usable at full size. + +### Whole-model attach + +Attaching all 48 blocks and rendering end to end at 768×512×49f: **6.0 s** +(median of three warm runs) against 99.8 s for the unmodified host with +weight offloading, peak 29.9 GB. Quality is frame-inspection equivalent. Two +qualifications on that figure: the blocks are attached one at a time because +the bf16 checkpoint does not fit resident on a 32GB part, and the +feed-forward seams are bound explicitly, 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. + +### 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/diffusers_gated_rotary_attention.py b/flash_rt/structures/adapters/diffusers_gated_rotary_attention.py index fc9670bc..f4ac9da8 100644 --- a/flash_rt/structures/adapters/diffusers_gated_rotary_attention.py +++ b/flash_rt/structures/adapters/diffusers_gated_rotary_attention.py @@ -145,10 +145,20 @@ def __call__(self, attn, hidden_states, encoder_hidden_states=None, class DiffusersGatedRotaryAttentionAdapter: - """Route gated dual-rotary Diffusers processors through the family.""" + """Route gated dual-rotary Diffusers processors through the family. + + ``prefer`` is passed through to the family binder unchanged and is + empty by default: which executable form serves this seam is the + family's decision, and preferring a quantized one is a deployment + decision that belongs to whoever assembled this adapter, not to the + adapter itself. + """ __name__ = "diffusers_gated_rotary_attention" + def __init__(self, prefer=()): + self.prefer = tuple(prefer) + def __call__(self, model, forward, *, prefix_cadence: bool = False): del prefix_cadence sites = [] @@ -183,7 +193,7 @@ def __call__(self, model, forward, *, prefix_cadence: bool = False): )) continue try: - core = bind_dense_attention_best(rows) + core = bind_dense_attention_best(rows, prefer=self.prefer) except ValueError as exc: refused.append((f"{path}.processor", str(exc)[:160])) continue 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 index da2fa2ab..d23dcdf9 100644 --- a/flash_rt/structures/impls/attention_core/sage2_blackwell.py +++ b/flash_rt/structures/impls/attention_core/sage2_blackwell.py @@ -131,38 +131,55 @@ def forward(self, query, key, value, *, scale=None): return self._out_nhd.transpose(1, 2) -def bind_sage2_dense_attention(captures, *, variant: str = "pv_fp8", - qk_quant_granularity: str = "per_warp"): - """Bind the sage2 dense form from one real capture set, or refuse. - - ``captures`` follows the family convention: an object carrying - ``q_shape``, ``kv_shape``, ``dtype``, ``device``, and optionally - ``allowed_ranges`` / ``mask``. Returns ``None`` (with the reason as an - attribute on the function, mirroring the family's refusal trail) when - the site is outside this form's envelope. +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. """ - def refuse(reason: str): - bind_sage2_dense_attention.last_refusal = reason + 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 - - mask = getattr(captures, "mask", None) - ranges = tuple(getattr(captures, "allowed_ranges", ()) or ()) - if mask is not None or ranges: - return refuse("masked/packed sites are not claimed by sage2") - b, heads, seq_q, head_dim = captures.q_shape - try: - dims = supported_head_dims() - except (ValueError, OSError, RuntimeError) as exc: - return refuse(f"artifact unavailable: {exc}") - if head_dim not in dims: - return refuse( - f"head_dim {head_dim} outside artifact envelope {dims}") - if captures.dtype != torch.bfloat16: - return refuse(f"dtype {captures.dtype} outside envelope (bf16)") - try: - return DenseAttentionSage2( - captures.q_shape, captures.kv_shape, captures.dtype, - captures.device, variant=variant, - qk_quant_granularity=qk_quant_granularity) - except (ValueError, RuntimeError) as exc: - return refuse(str(exc)) + 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 index f9663e78..de8393b9 100644 --- a/flash_rt/structures/impls/attention_core/sage3_blackwell.py +++ b/flash_rt/structures/impls/attention_core/sage3_blackwell.py @@ -121,39 +121,44 @@ def forward(self, query, key, value, *, scale=None): return out.transpose(1, 2) -def bind_sage3_dense_attention(captures): - """Bind the sage3 dense form from one real capture set, or refuse. - - ``captures`` follows the family convention: an object carrying - ``q_shape``, ``kv_shape``, ``dtype``, ``device``, and optionally - ``allowed_ranges`` / ``mask``. Returns ``None`` (with the reason as an - attribute on the function, mirroring the family's refusal trail) when - the site is outside this form's envelope. +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. """ - def refuse(reason: str): - bind_sage3_dense_attention.last_refusal = reason + 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 - - mask = getattr(captures, "mask", None) - ranges = tuple(getattr(captures, "allowed_ranges", ()) or ()) - if mask is not None or ranges: - return refuse("masked/packed sites are not claimed by sage3") - if tuple(captures.kv_shape) != tuple(captures.q_shape): - return refuse( - "self-attention form: Q and KV shapes differ at this site") - b, heads, seq_q, head_dim = captures.q_shape - try: - dims = supported_head_dims() - except (ValueError, OSError, RuntimeError) as exc: - return refuse(f"artifact unavailable: {exc}") - if head_dim not in dims: - return refuse( - f"head_dim {head_dim} outside artifact envelope {dims}") - if captures.dtype != torch.bfloat16: - return refuse(f"dtype {captures.dtype} outside envelope (bf16)") - try: - return DenseAttentionSage3( - captures.q_shape, captures.kv_shape, captures.dtype, - captures.device) - except (ValueError, RuntimeError) as exc: - return refuse(str(exc)) + 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/tests/test_sage_attention_forms.py b/tests/test_sage_attention_forms.py new file mode 100644 index 00000000..0b596345 --- /dev/null +++ b/tests/test_sage_attention_forms.py @@ -0,0 +1,300 @@ +"""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 == () From 3283c49aedb3ab6a6d7732c0ff4e21abe899857c Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sun, 16 Aug 2026 10:26:41 -0400 Subject: [PATCH 07/11] feat(structures): quantized attention is a precision profile, not a default The published attention order is precision-first, so a host that never asked for a quantized form keeps the numerics it has receipts for. That left the forms reachable only by hand, which is the wrong shape for the decision: which executable form may serve a seam, when the trade is a bounded numerical error for speed, is the same kind of statement the scheme already makes about GEMM seams and about the gated-delta and MTP projections. So it arrives the same way. QuantScheme gains attention_forms, empty on every existing profile; the attention adapters that declare scheme awareness read it, exactly as the gated-delta adapter reads its projection format; and 'nvfp4_balance_sage' registers the profile that names them. 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. Measured on one transformer block, real weights and captured inputs, paired inside the gate: at S=24576 the projections-only profile reaches 1.15x with the BF16 attention form bound and declined at 1.006x, and the sage profile reaches 1.49x with the attention unit at 1.259x. The difference is the attention kernel's own 45.9 -> 19.6 ms, which is worth naming because a block ratio and a kernel ratio are different numbers for the same 27 ms. --- docs/ltx25_usage.md | 57 ++++++++----- .../diffusers_gated_rotary_attention.py | 19 +++-- flash_rt/structures/autobuild.py | 14 +++- flash_rt/structures/schemes.py | 27 +++++- tests/test_sage_attention_forms.py | 84 +++++++++++++++++++ 5 files changed, 171 insertions(+), 30 deletions(-) diff --git a/docs/ltx25_usage.md b/docs/ltx25_usage.md index f9877cae..7343d823 100644 --- a/docs/ltx25_usage.md +++ b/docs/ltx25_usage.md @@ -142,35 +142,52 @@ 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. + ### Measured on one transformer block Real checkpoint weights, real captured deployment inputs, paired alternating timing inside the gate, on a 5090. "Attention" is the gate's verdict for the attention family; the projections are the `nvfp4_balance` W4A4 form. -| Site shape | Configuration | Block latency | Attention family | Peak memory | +| Site shape | `scheme=` | Block latency | Attention family | Peak memory | |---|---|---|---|---| -| S=24576 (1536×1024×121f) | host | 134.3 ms | — | 12.2 GB | -| | attach, default order | 117.1 ms (1.15×) | bound, declined at 1.006× | 8.2 GB | -| | attach, sage2 preferred | **90.1 ms (1.49×)** | activated, 1.257× | at the 32GB ceiling | -| S=2688 (768×512×49f) | host | 10.3 ms | — | 2.3 GB | -| | attach, default order | 8.2 ms (1.25×) | declined | 1.7 GB | -| | attach, sage2 preferred | 8.0 ms (1.28×) | activated | 4.8 GB | +| 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). Two results are worth -reading carefully rather than skipping: - -- **The default order does not use the quantized attention forms.** They trade - a bounded numerical error for speed, which is a deployment decision, so a - caller asks for one explicitly. Without that, the family's BF16 form binds, - and at these shapes the net-win gate measures it at 1.006× and keeps the - host's attention — the projections carry the whole win. -- **Peak memory falls when the projections are quantized** (12.2 → 8.2 GB) and - rises when quantized attention is preferred, because each attention site - owns its staging and quantization workspace. At S=24576 across four sites - that reaches the ceiling of a 32GB part; pooling those workspaces is the - open item before this configuration is usable at full size. +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. ### Whole-model attach diff --git a/flash_rt/structures/adapters/diffusers_gated_rotary_attention.py b/flash_rt/structures/adapters/diffusers_gated_rotary_attention.py index f4ac9da8..41133d15 100644 --- a/flash_rt/structures/adapters/diffusers_gated_rotary_attention.py +++ b/flash_rt/structures/adapters/diffusers_gated_rotary_attention.py @@ -147,20 +147,25 @@ def __call__(self, attn, hidden_states, encoder_hidden_states=None, class DiffusersGatedRotaryAttentionAdapter: """Route gated dual-rotary Diffusers processors through the family. - ``prefer`` is passed through to the family binder unchanged and is - empty by default: which executable form serves this seam is the - family's decision, and preferring a quantized one is a deployment - decision that belongs to whoever assembled this adapter, not to the - adapter itself. + 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): + 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) @@ -193,7 +198,7 @@ def __call__(self, model, forward, *, prefix_cadence: bool = False): )) continue try: - core = bind_dense_attention_best(rows, prefer=self.prefer) + core = bind_dense_attention_best(rows, prefer=prefer) except ValueError as exc: refused.append((f"{path}.processor", str(exc)[:160])) continue diff --git a/flash_rt/structures/autobuild.py b/flash_rt/structures/autobuild.py index 10b07e3a..72aec048 100644 --- a/flash_rt/structures/autobuild.py +++ b/flash_rt/structures/autobuild.py @@ -1263,8 +1263,18 @@ 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): + result = adapter(model, thunks[0], + prefix_cadence=prefix_cadence, + scheme=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])) 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_sage_attention_forms.py b/tests/test_sage_attention_forms.py index 0b596345..d4d8152c 100644 --- a/tests/test_sage_attention_forms.py +++ b/tests/test_sage_attention_forms.py @@ -298,3 +298,87 @@ def test_adapter_default_prefers_nothing(): 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"] == () From be7c6ff9aa0fe5fdceb324977b3b6cc67d878f90 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sun, 16 Aug 2026 11:23:19 -0400 Subject: [PATCH 08/11] feat(structures): a caller can name attention forms, and the sage staging is pooled Two things stood between the quantized attention form and the person who would decide whether to use it. It could only be selected by registering a scheme. The profile is the right home for the decision, but a deployment tuning one seam should not have to author a profile to try it, so attach and auto_swaps take attention_forms directly and it outranks the profile -- the same statement, made by the caller. The answer still comes back measured: the family qualifies the shape, the gate times both arms and prints the accuracy band, and detach reverses it, which is what makes trying it the cheap way to find out. And it did not fit. Each bound seam allocated its own staging and workspace, about 800MB at 24576 tokens over 32 heads, so a transformer whose every block reaches this seam ran out of memory before it could be measured -- which is why the form had only ever been judged one block at a time. The staging is now pooled per (shape, dtype, device, variant, granularity). Sharing is sound because nothing survives a call: every buffer is written at the top of the forward and read before it returns, and the host runs its blocks in sequence on one stream. Held as plain attributes rather than buffers, so a pooled set does not appear in each seam's state_dict. Measured across all 48 blocks at 1536x1024x121f: memory grows by the quantized weights alone, 0.41GB per block, where before each block also took its own workspace. Per-block parity is unchanged with the pool in place -- attention unit cosine 0.999998, whole block 0.999845, band pass throughout. --- docs/ltx25_usage.md | 53 +++++-- flash_rt/structures/autobuild.py | 28 +++- flash_rt/structures/frontdoor.py | 5 +- .../impls/attention_core/sage2_blackwell.py | 60 ++++++-- tests/test_sage_attention_forms.py | 145 ++++++++++++++++++ 5 files changed, 262 insertions(+), 29 deletions(-) diff --git a/docs/ltx25_usage.md b/docs/ltx25_usage.md index 7343d823..169dbd6c 100644 --- a/docs/ltx25_usage.md +++ b/docs/ltx25_usage.md @@ -156,11 +156,39 @@ 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. -### Measured on one transformer block +To tune one seam without registering a profile, name the forms directly: -Real checkpoint weights, real captured deployment inputs, paired alternating -timing inside the gate, on a 5090. "Attention" is the gate's verdict for the -attention family; the projections are the `nvfp4_balance` W4A4 form. +```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×) | **5.7 s (17.5×)**, peak 26.8 GB | + +Medians of three warm runs. Frames are inspection-equivalent to the host's +own output at both settings. + +### 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 | |---|---|---|---|---| @@ -189,16 +217,15 @@ are easy to misread, so they are worth stating: 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. -### Whole-model attach +### How the whole-model figures were produced -Attaching all 48 blocks and rendering end to end at 768×512×49f: **6.0 s** -(median of three warm runs) against 99.8 s for the unmodified host with -weight offloading, peak 29.9 GB. Quality is frame-inspection equivalent. Two -qualifications on that figure: the blocks are attached one at a time because -the bf16 checkpoint does not fit resident on a 32GB part, and the -feed-forward seams are bound explicitly, 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. +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 diff --git a/flash_rt/structures/autobuild.py b/flash_rt/structures/autobuild.py index 72aec048..02258697 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) @@ -544,6 +558,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: @@ -1269,9 +1284,16 @@ def _stream_window_undo(): # serve its seam is a precision decision and the # scheme is where those are stated if getattr(adapter, "scheme_aware", False): - result = adapter(model, thunks[0], - prefix_cadence=prefix_cadence, - scheme=scheme_obj) + # 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) diff --git a/flash_rt/structures/frontdoor.py b/flash_rt/structures/frontdoor.py index aa64a1e8..7174695b 100644 --- a/flash_rt/structures/frontdoor.py +++ b/flash_rt/structures/frontdoor.py @@ -294,6 +294,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 +349,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, diff --git a/flash_rt/structures/impls/attention_core/sage2_blackwell.py b/flash_rt/structures/impls/attention_core/sage2_blackwell.py index d23dcdf9..4891cb8a 100644 --- a/flash_rt/structures/impls/attention_core/sage2_blackwell.py +++ b/flash_rt/structures/impls/attention_core/sage2_blackwell.py @@ -30,6 +30,8 @@ from __future__ import annotations +from typing import NamedTuple + import torch from .. import hub_kernel @@ -44,6 +46,44 @@ _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"]) @@ -99,17 +139,15 @@ def __init__(self, q_shape, kv_shape, dtype: torch.dtype, device, art = _artifact() self._fn = (art.sage2_prefill_fp8v_bf16_d128 if variant == "pv_fp8" else art.sage2_prefill_f16_bf16_d128) - # NHD staging + caller-owned workspace, one set per bound shape. - self.register_buffer("_q_nhd", torch.empty( - b, seq_q, heads, head_dim, dtype=dtype, device=device)) - self.register_buffer("_k_nhd", torch.empty( - b, seq_kv, kv_heads, head_dim, dtype=dtype, device=device)) - self.register_buffer("_v_nhd", torch.empty_like(self._k_nhd)) - self.register_buffer("_out_nhd", torch.empty_like(self._q_nhd)) - self._workspace = art.allocate_workspace( - self._q_nhd, self._k_nhd, self._v_nhd, - fp8v=(variant == "pv_fp8"), - qk_quant_granularity=qk_quant_granularity) + 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)) diff --git a/tests/test_sage_attention_forms.py b/tests/test_sage_attention_forms.py index d4d8152c..63c42a2e 100644 --- a/tests/test_sage_attention_forms.py +++ b/tests/test_sage_attention_forms.py @@ -382,3 +382,148 @@ def forward(): 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()) From aad383129ddd6adcf2a9146ee097a436edab2f13 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sun, 16 Aug 2026 12:13:51 -0400 Subject: [PATCH 09/11] docs(ltx25): lead with what a request costs, not what a block costs The measurements a reader needs first are wall clock for one request against the host they would otherwise run, at both sizes. The per-block table stays, relabelled as what it is: a diagnostic that says which family earned which part of that time. The full-size row carries its own qualifications rather than a single ratio: it is eager, the quantized attention profile does not fit at that size yet, and the audio feed-forward stays at host precision because its 126-row calls sit outside the fused chain's alignment. Those three are the distance between this number and a hand-assembled configuration on the same card, and a reader deciding whether to use this should see them next to the number, not after adopting it. --- docs/ltx25_usage.md | 25 +++++++++++++++++++++---- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/docs/ltx25_usage.md b/docs/ltx25_usage.md index 169dbd6c..2c77ff7d 100644 --- a/docs/ltx25_usage.md +++ b/docs/ltx25_usage.md @@ -179,10 +179,27 @@ actually starts from. | Request | Host (offload) | `"nvfp4_balance"` | `"nvfp4_balance_sage"` | |---|---|---|---| -| 768×512×49f | 99.8 s | 6.0 s (16.6×) | **5.7 s (17.5×)**, peak 26.8 GB | - -Medians of three warm runs. Frames are inspection-equivalent to the host's -own output at both settings. +| 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 From f8906c9f2c8e41489410d66bfd3988c8f4ea470b Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sun, 16 Aug 2026 13:06:55 -0400 Subject: [PATCH 10/11] fix(structures): a declined form must give its working set back Reverting a routed seam puts the host processor back. It does not free anything: the bound form stays reachable through the plan's observed map and through the very closures that reverted it, and the plan is what a caller holds in order to detach later. So a form the gate declined kept its whole working set for the lifetime of the attachment. Measured on one transformer block at 24576 tokens: the block's own weights fall from 0.720GB to 0.428GB when the projections are quantized, but 0.401GB appeared elsewhere and the attachment came out 0.109GB *heavier* than the host it replaced. Excluding the attention family from the same run leaves 0.009GB unaccounted, which is what identified it. Adapters can now publish a release alongside their revert, dropping their own hold on what they bound; the front door calls it once the gate has settled and no routed unit won. The revert callables survive and stay correct, because an adapter's release empties the route list its closures were built over rather than the closures themselves. The whole-host refusal path already did this through revert_all; what was missing was the mixed outcome, where some units win and the routed ones do not -- which is the ordinary case for this host, since the projections win at every shape and the attention family only wins at long ones. After: the same block attaches at 0.437GB, 0.283GB below the host, with nothing left to reclaim when the plan is dropped. A form the gate activates is untouched, which the sage profile's unchanged 2.045GB confirms. --- .../diffusers_gated_rotary_attention.py | 7 ++ flash_rt/structures/autobuild.py | 26 +++++++ flash_rt/structures/frontdoor.py | 8 +++ tests/test_sage_attention_forms.py | 68 +++++++++++++++++++ 4 files changed, 109 insertions(+) diff --git a/flash_rt/structures/adapters/diffusers_gated_rotary_attention.py b/flash_rt/structures/adapters/diffusers_gated_rotary_attention.py index 41133d15..72e34afc 100644 --- a/flash_rt/structures/adapters/diffusers_gated_rotary_attention.py +++ b/flash_rt/structures/adapters/diffusers_gated_rotary_attention.py @@ -227,9 +227,16 @@ 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, diff --git a/flash_rt/structures/autobuild.py b/flash_rt/structures/autobuild.py index 02258697..7a9d6870 100644 --- a/flash_rt/structures/autobuild.py +++ b/flash_rt/structures/autobuild.py @@ -129,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`` @@ -142,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. @@ -1326,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/frontdoor.py b/flash_rt/structures/frontdoor.py index 7174695b..f165da67 100644 --- a/flash_rt/structures/frontdoor.py +++ b/flash_rt/structures/frontdoor.py @@ -485,6 +485,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: diff --git a/tests/test_sage_attention_forms.py b/tests/test_sage_attention_forms.py index 63c42a2e..39261891 100644 --- a/tests/test_sage_attention_forms.py +++ b/tests/test_sage_attention_forms.py @@ -527,3 +527,71 @@ def test_the_front_door_takes_attention_forms(): 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() From a67f4b102776319999657149abe870d2109db045 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sun, 16 Aug 2026 13:29:30 -0400 Subject: [PATCH 11/11] fix(structures): a verdict now says which shape it holds for, and what it costs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Thirteen catalog specs declare latency.per_shape: true. Nothing has ever read it. The gate takes one measurement, compares one scalar, and records the outcome as though it were universal — while writing the shape it measured into the refusal text, which is the same program stating both that the verdict is local and that it is not. Two recordings, no decision change. min_speedup is untouched, no refusal path is added or tightened, and every verdict this produces is the verdict it produced before. The unit's spec is read, so a refusal from a per-shape rule says it holds for that shape and no other. The shape itself comes from the guards the bound seams were armed with, because m_profile is empty for hosts whose bindings do not declare it and those refusals read "rows unrecorded" — shape-scoped verdicts with no shape on them. And the receipt carries what a unit does to resident memory: what its forms hold, what the host modules they replace hold, and the difference. The walk finds tensors wherever a form keeps them, including inside an artifact's workspace object, counts pooled storage once across sites, and excludes the retained host from its replacement's total. On this host the projections read -0.292 GiB and the declined attention family +0.382 GiB, both agreeing with independent measurement to within 10 MiB. That second number is the point. A refusal has always read as free; this one now reads "no net win (1.005x), +0.38 GiB resident at rows=786432". Whether memory should ever enter the decision is a separate question, and one nobody could argue either way while the data did not exist. --- flash_rt/structures/frontdoor.py | 149 ++++++++++++++++- tests/test_gate_records_shape_and_memory.py | 170 ++++++++++++++++++++ 2 files changed, 317 insertions(+), 2 deletions(-) create mode 100644 tests/test_gate_records_shape_and_memory.py diff --git a/flash_rt/structures/frontdoor.py b/flash_rt/structures/frontdoor.py index f165da67..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, @@ -412,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 @@ -419,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: @@ -571,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/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)", "")