LTX-2.5 attention forms and the gated dual-rotary Diffusers adapter - #1
Open
LiangSu8899 wants to merge 11 commits into
Open
LTX-2.5 attention forms and the gated dual-rotary Diffusers adapter#1LiangSu8899 wants to merge 11 commits into
LiangSu8899 wants to merge 11 commits into
Conversation
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.
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.
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.
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.
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.
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.
…efault 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.
…ging 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.
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.
…t it costs 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.
The LTX-2.5 DiT and video-pipeline bindings land in flash_rt/catalog/bindings (FlashRT PR catalog/ltx25-bindings); this package reads them from there.
18 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Migrated from flashrt-project/FlashRT#174 with history (10 commits, paths and imports rewritten to this package). The LTX-2.5 DiT and video-pipeline bindings are not carried here: they land in FlashRT's catalog (flashrt-project/FlashRT#198), which this package reads through the flash-rt distribution.
Verified against that catalog branch: every module imports; the CPU test set runs 295 passed, 14 skipped, 6 failed. The 6 failures are in
tests/test_sage_attention_forms.pyand are identical on the original FlashRT branch before migration (kernel-hub version resolutionVersion >=0 not foundfor the sage package, andcutlass.cute.corelackingThrMmain this environment's CuTe-DSL); they are the PR's own open items, not the move.