Prism ZeroGPU demo

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/case-1/reference_image.jpg filter=lfs diff=lfs merge=lfs -text
+examples/case-2/reference_image.jpg filter=lfs diff=lfs merge=lfs -text
+examples/case-3/reference_image.jpg filter=lfs diff=lfs merge=lfs -text
+examples/case-4/reference_image.jpg filter=lfs diff=lfs merge=lfs -text
+examples/case-5/reference_image.jpg filter=lfs diff=lfs merge=lfs -text
+examples/case-6/reference_image.jpg filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,87 @@
+Tencent is pleased to support the community by making Prism available.
+
+Copyright (C) 2026 Tencent. All rights reserved.
+
+The open-source software and/or model(s) included in this distribution may have been modified by Tencent ("Tencent Modifications"). All Tencent Modifications are Copyright (C) Tencent.
+
+Prism is licensed under MIT, except for the third-party components listed below, which remain licensed under their respective original terms. Prism does not impose any additional restrictions beyond those specified in the original licenses of these third-party components. Users are required to comply with all applicable terms and conditions of the original licenses and to ensure that the use of these third-party components conforms to all relevant laws and regulations.
+
+For the avoidance of doubt, Prism refers solely to code, parameters and weights made publicly available by Tencent in accordance with MIT.
+
+Terms of the MIT License:
+--------------------------------------------------------------------
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice (including the next paragraph) shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+
+Dependencies and Licenses:
+
+This open-source project builds upon the following open-source model(s) and/or software, each of which remains licensed under its original license(s). Certain component(s) may include modifications made by Tencent ("Tencent Modifications"), which are Copyright (C) Tencent.
+
+In case you believe there have been errors in the attribution below, you may submit the concerns to us for review and correction.
+
+
+Open Source Model(s)/Software Licensed under the Apache-2.0:
+--------------------------------------------------------------------
+1. MOVA-360p
+Copyright (c) MOVA-360p original author and authors
+Please note this model/software has been modified by Tencent in this distribution.
+Terms of the Apache-2.0:
+--------------------------------------------------------------------
+Apache License
+Version 2.0, January 2004
+http://www.apache.org/licenses/
+
+TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+1. Definitions.
+
+"License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document.
+
+"Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License.
+
+"Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity.
+
+"You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License.
+
+"Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files.
+
+"Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types.
+
+"Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below).
+
+"Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof.
+
+"Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution."
+
+"Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work.
+
+2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form.
+
+3. Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed.
+
+4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions:
+
+You must give any other recipients of the Work or Derivative Works a copy of this License; and
+You must cause any modified files to carry prominent notices stating that You changed the files; and
+You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and
+If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License.
+You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License.
+
+5. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions.
+
+6. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file.
+
+7. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License.
+
+8. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages.
+
+9. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability.
+
+END OF TERMS AND CONDITIONS
+
+==================================================
+End of the Attribution Notice of this project.

README.md

@@ -1,13 +1,20 @@
 ---
-title: Prism
-emoji: 🐠
-colorFrom: yellow
-colorTo: green
+title: Prism Video Audio
+emoji: 🎬
+colorFrom: indigo
+colorTo: pink
 sdk: gradio
-sdk_version: 6.29.1
-python_version: '3.12'
+sdk_version: 5.50.0
+python_version: "3.12"
 app_file: app.py
 pinned: false
+license: mit
+models:
+  - FrancisRing/Prism
+startup_duration_timeout: 1h
+short_description: Prism image+text to video with audio (ZeroGPU xlarge)
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+ZeroGPU demo of [FrancisRing/Prism](https://huggingface.co/FrancisRing/Prism)
+(dynamic sparse attention, joint video + audio generation on the MOVA backbone).
+Inference code is adapted from [Tencent-Hunyuan/Prism](https://github.com/Tencent-Hunyuan/Prism) (MIT).

app.py

@@ -0,0 +1,360 @@
+"""Prism (FrancisRing/Prism) joint video-audio generation on ZeroGPU (xlarge, 96GB)."""
+
+import os
+import sys
+import time
+import types
+import random
+import tempfile
+import subprocess
+
+import torch
+
+# ---------------------------------------------------------------------------
+# ZeroGPU does not support torch.compile -> turn every @torch.compile into a no-op
+# (must happen before the model code is imported).
+# ---------------------------------------------------------------------------
+def _no_compile(fn=None, *args, **kwargs):
+    if callable(fn):
+        return fn
+    return lambda f: f
+
+
+torch.compile = _no_compile
+
+# The DAC audio VAE only needs `audiotools` for training/IO helpers. Stub it so we
+# don't have to install descript-audiotools (it pins protobuf<3.20).
+_at = types.ModuleType("audiotools")
+_at_ml = types.ModuleType("audiotools.ml")
+
+
+class _AudioSignal:  # never used in the decode path
+    def __init__(self, *a, **k):
+        raise RuntimeError("audiotools.AudioSignal is not available in this Space")
+
+
+_at.AudioSignal = _AudioSignal
+_at_ml.BaseModel = torch.nn.Module
+_at.ml = _at_ml
+sys.modules.setdefault("audiotools", _at)
+sys.modules.setdefault("audiotools.ml", _at_ml)
+
+import spaces  # noqa: E402  (must be imported before CUDA is touched)
+import numpy as np  # noqa: E402
+import gradio as gr  # noqa: E402
+from PIL import Image  # noqa: E402
+from huggingface_hub import snapshot_download  # noqa: E402
+from safetensors import safe_open  # noqa: E402
+from accelerate import init_empty_weights  # noqa: E402
+from accelerate.utils import set_module_tensor_to_device  # noqa: E402
+
+from hymm.models.modules.mova import MOVABridge  # noqa: E402
+from hymm.models.modules.wan_video_dit import WanModel  # noqa: E402
+from hymm.models.modules.wan_audio_dit import WanAudioModel  # noqa: E402
+from hymm.models.modules.interactionv2 import DualTowerConditionalBridge  # noqa: E402
+from hymm.models.modules.dac_vae import DAC  # noqa: E402
+from hymm.diffusion.pipelines.mova_pipeline import MOVAPipeline  # noqa: E402
+from hymm.diffusion.schedulers.flow_match_pair import FlowMatchPairScheduler  # noqa: E402
+from diffusers.models.autoencoders import AutoencoderKLWan  # noqa: E402
+from transformers import T5TokenizerFast, UMT5EncoderModel  # noqa: E402
+
+REPO_ID = "FrancisRing/Prism"
+VARIANT = os.environ.get("PRISM_VARIANT", "preview_alpha")  # or preview_beta (motion)
+BASE = "pretrained_models/MOVA-360p"
+DTYPE = torch.bfloat16
+
+NEGATIVE_PROMPT = (
+    "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，"
+    "整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指"
+)
+
+# ---------------------------------------------------------------------------
+# Download only what inference needs: the Prism transformer checkpoint (it holds
+# both video experts, the audio tower and the bridge) + text encoder / VAEs.
+# The 57GB of base MOVA DiT weights are fully overridden, so they are skipped.
+# ---------------------------------------------------------------------------
+t0 = time.time()
+root = snapshot_download(
+    REPO_ID,
+    allow_patterns=[
+        f"{VARIANT}/*",
+        f"{BASE}/model_index.json",
+        f"{BASE}/*/config.json",
+        f"{BASE}/text_encoder/*",
+        f"{BASE}/tokenizer/*",
+        f"{BASE}/scheduler/*",
+        f"{BASE}/video_vae/*",
+        f"{BASE}/audio_vae/*",
+    ],
+)
+base = os.path.join(root, BASE)
+print(f"[prism] download done in {time.time() - t0:.0f}s -> {root}")
+
+# --- transformer: build on meta, then stream the checkpoint straight to cuda ---
+t0 = time.time()
+with init_empty_weights():
+    video_dit = WanModel.from_config(WanModel.load_config(base, subfolder="video_dit"))
+    video_dit_2 = WanModel.from_config(WanModel.load_config(base, subfolder="video_dit_2"))
+    audio_dit = WanAudioModel.from_config(WanAudioModel.load_config(base, subfolder="audio_dit"))
+    dual_bridge = DualTowerConditionalBridge.from_config(
+        DualTowerConditionalBridge.load_config(base, subfolder="dual_tower_bridge")
+    )
+    transformer = MOVABridge(
+        video_dit=video_dit,
+        video_dit_2=video_dit_2,
+        audio_dit=audio_dit,
+        dual_tower_bridge=dual_bridge,
+        boundary_ratio=0.9,
+    )
+
+ckpt_file = os.path.join(root, VARIANT, "diffusion_pytorch_model.safetensors")
+expected = set(transformer.state_dict().keys())
+loaded = set()
+with safe_open(ckpt_file, framework="pt", device="cpu") as f:
+    for k in f.keys():
+        if k not in expected:
+            print(f"[prism] unexpected key {k}")
+            continue
+        set_module_tensor_to_device(transformer, k, "cuda", value=f.get_tensor(k), dtype=DTYPE)
+        loaded.add(k)
+missing = expected - loaded
+if missing:
+    raise RuntimeError(f"[prism] {len(missing)} transformer keys missing, e.g. {sorted(missing)[:5]}")
+transformer.to("cuda")
+transformer.eval().requires_grad_(False)
+print(f"[prism] transformer ready ({len(loaded)} tensors) in {time.time() - t0:.0f}s")
+
+# --- frozen components ---
+video_vae = AutoencoderKLWan.from_pretrained(base, subfolder="video_vae", torch_dtype=DTYPE).to("cuda").eval()
+audio_vae = DAC.from_pretrained(base, subfolder="audio_vae", torch_dtype=DTYPE).to("cuda").eval()
+text_encoder = UMT5EncoderModel.from_pretrained(base, subfolder="text_encoder", torch_dtype=DTYPE).to("cuda").eval()
+tokenizer = T5TokenizerFast.from_pretrained(base, subfolder="tokenizer")
+scheduler = FlowMatchPairScheduler.from_pretrained(base, subfolder="scheduler")
+
+pipe = MOVAPipeline(
+    transformer=transformer,
+    video_vae=video_vae,
+    audio_vae=audio_vae,
+    text_encoder=text_encoder,
+    tokenizer=tokenizer,
+    scheduler=scheduler,
+    boundary_ratio=0.9,
+    device=torch.device("cuda"),
+)
+print("[prism] pipeline ready")
+
+RESOLUTIONS = {
+    "480p (848×480)": (480, 848),
+    "720p (1280×720)": (720, 1280),
+}
+
+
+def configure_sparse(sparsity, use_ivpq):
+    transformer.set_sparse_high_noise_only(False)
+    transformer.configure_bsa(
+        enable_bsa=True,
+        bsa_params={
+            "sparsity": float(sparsity),
+            "cdf_threshold": 0.20,
+            "chunk_3d_shape_q": [4, 4, 4],
+            "chunk_3d_shape_k": [4, 4, 4],
+        },
+        enable_bsa_v2a=False,
+        bsa_params_v2a=None,
+    )
+    for sa in transformer._video_self_attns():
+        sa.bsa_params["dynamic_block_lambda_a"] = 0.5
+        sa.bsa_params["dynamic_block_tau_128"] = 0.25
+        sa.bsa_params["dynamic_block_lambda_128"] = 1.0
+    transformer.configure_ivpq_dynamic_block(enable=bool(use_ivpq))
+
+
+def crop_and_resize(img, height, width):
+    w, h = img.size
+    target = width / height
+    if w / h > target:
+        nw = int(h * target)
+        img = img.crop(((w - nw) // 2, 0, (w - nw) // 2 + nw, h))
+    elif w / h < target:
+        nh = int(w / target)
+        img = img.crop((0, (h - nh) // 2, w, (h - nh) // 2 + nh))
+    return img.resize((width, height), Image.LANCZOS)
+
+
+def target_size(image, resolution):
+    h, w = RESOLUTIONS[resolution]
+    if image is not None and image.height > image.width:  # portrait input -> portrait video
+        h, w = w, h
+    return h, w
+
+
+def num_tokens(h, w, frames):
+    return ((frames - 1) // 4 + 1) * (h // 16) * (w // 16)
+
+
+# Rough GPU-time model (seconds per transformer forward), calibrated on the Space.
+SEC_PER_TOKEN = float(os.environ.get("PRISM_SEC_PER_TOKEN", "1.0e-4"))
+SEC_PER_TOKEN2 = float(os.environ.get("PRISM_SEC_PER_TOKEN2", "4.0e-9"))
+MAX_DURATION = int(os.environ.get("PRISM_MAX_DURATION", "1500"))
+
+
+def estimate_seconds(image, resolution, frames, steps, cfg, sparsity):
+    h, w = target_size(image, resolution)
+    tok = num_tokens(h, w, int(frames))
+    per_forward = SEC_PER_TOKEN * tok + SEC_PER_TOKEN2 * tok * tok * (1.0 - float(sparsity) + 0.05)
+    passes = 1 if float(cfg) == 1.0 else 2
+    return int(60 + int(steps) * passes * per_forward)
+
+
+def get_duration(image, prompt, audio_prompt, resolution, frames, steps, cfg, shift, audio_shift,
+                 seed, sparsity, use_ivpq, *args, **kwargs):
+    return min(MAX_DURATION, int(estimate_seconds(image, resolution, frames, steps, cfg, sparsity) * 1.15))
+
+
+@spaces.GPU(size="xlarge", duration=get_duration)
+def _generate_gpu(image, prompt, audio_prompt, resolution, frames, steps, cfg, shift, audio_shift,
+                  seed, sparsity, use_ivpq, progress=gr.Progress(track_tqdm=True)):
+    t_start = time.time()
+    h, w = target_size(image, resolution)
+    ref = crop_and_resize(image.convert("RGB"), h, w)
+    configure_sparse(sparsity, use_ivpq)
+    torch.manual_seed(int(seed))
+    video, audio = pipe(
+        prompt=prompt,
+        image=ref,
+        audio_prompt=audio_prompt or prompt,
+        negative_prompt=NEGATIVE_PROMPT,
+        seed=int(seed),
+        height=h,
+        width=w,
+        num_frames=int(frames),
+        video_fps=24.0,
+        num_inference_steps=int(steps),
+        visual_shift=float(shift),
+        audio_shift=float(audio_shift),
+        cfg_scale=float(cfg),
+        enable_vae_tiling=(h * w >= 720 * 1280),
+    )
+    frames_np = np.stack([np.asarray(fr) for fr in video[0]])
+    audio_np = audio[0].float().cpu().squeeze().numpy()
+    elapsed = time.time() - t_start
+    print(f"[prism] {h}x{w} {frames}f {steps} steps cfg={cfg} -> {elapsed:.1f}s on GPU "
+          f"(peak mem {torch.cuda.max_memory_allocated() / 2**30:.1f} GiB)")
+    return frames_np, audio_np, elapsed
+
+
+def save_video_with_audio(frames, audio_np, path, fps, sample_rate):
+    import wave
+    import imageio
+
+    with tempfile.TemporaryDirectory() as tmp:
+        tmp_video = os.path.join(tmp, "video.mp4")
+        tmp_audio = os.path.join(tmp, "audio.wav")
+        writer = imageio.get_writer(tmp_video, fps=fps, quality=9, macro_block_size=8)
+        for fr in frames:
+            writer.append_data(fr)
+        writer.close()
+
+        if audio_np.ndim == 1:
+            audio_np = audio_np[None, :]
+        audio_np = audio_np[:2]
+        pcm = (np.clip(audio_np, -1.0, 1.0) * 32767.0).astype(np.int16)
+        with wave.open(tmp_audio, "wb") as wf:
+            wf.setnchannels(pcm.shape[0])
+            wf.setsampwidth(2)
+            wf.setframerate(int(sample_rate))
+            wf.writeframes(pcm.T.reshape(-1).tobytes())
+
+        cmd = ["ffmpeg", "-y", "-i", tmp_video, "-i", tmp_audio, "-c:v", "copy", "-c:a", "aac",
+               "-b:a", "192k", "-movflags", "+faststart", "-shortest", path]
+        r = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        if r.returncode != 0:
+            print("[prism] ffmpeg mux failed:", r.stderr.decode(errors="ignore")[-500:])
+            import shutil
+            shutil.copyfile(tmp_video, path)
+
+
+def generate(image, prompt, audio_prompt, resolution, frames, steps, cfg, shift, audio_shift,
+             seed, randomize_seed, sparsity, use_ivpq, progress=gr.Progress(track_tqdm=True)):
+    if image is None:
+        raise gr.Error("Please upload a reference image (Prism is an image+text → video+audio model).")
+    if not prompt or not prompt.strip():
+        raise gr.Error("Please enter a prompt.")
+    if randomize_seed:
+        seed = random.randint(0, 2**31 - 1)
+    frames = int(frames)
+    frames = frames - (frames - 1) % 4
+    frames_np, audio_np, elapsed = _generate_gpu(
+        image, prompt, audio_prompt, resolution, frames, steps, cfg, shift, audio_shift,
+        seed, sparsity, use_ivpq,
+    )
+    out = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
+    save_video_with_audio(frames_np, audio_np, out, fps=24.0, sample_rate=pipe.audio_sample_rate)
+    return out, seed, f"GPU time: {elapsed:.0f}s (xlarge quota counts ×2)"
+
+
+def update_estimate(image, resolution, frames, steps, cfg, sparsity):
+    h, w = target_size(image, resolution)
+    est = estimate_seconds(image, resolution, int(frames), steps, cfg, sparsity)
+    return (f"Output {w}×{h}, {int(frames)} frames ({int(frames) / 24:.1f}s) · "
+            f"rough GPU estimate ≈ {est / 60:.1f} min → ≈ {2 * est / 60:.1f} min of daily ZeroGPU quota")
+
+
+EXAMPLES_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "examples")
+examples = []
+if os.path.isdir(EXAMPLES_DIR):
+    for name in sorted(os.listdir(EXAMPLES_DIR)):
+        d = os.path.join(EXAMPLES_DIR, name)
+        img, txt = os.path.join(d, "reference_image.jpg"), os.path.join(d, "prompt.txt")
+        if os.path.exists(img) and os.path.exists(txt):
+            with open(txt, encoding="utf-8") as fh:
+                examples.append([img, fh.read().strip()])
+
+with gr.Blocks(title="Prism · joint video + audio") as demo:
+    gr.Markdown(
+        "# Prism: image + text → video with audio\n"
+        f"[Model](https://huggingface.co/{REPO_ID}) · [Code](https://github.com/Tencent-Hunyuan/Prism) · "
+        f"[Paper](https://arxiv.org/abs/2610.05416). Checkpoint: `{VARIANT}`. Runs on one ZeroGPU **xlarge** "
+        "(RTX Pro 6000, 96GB). xlarge uses quota at 2×, so keep clips short. "
+        "Sound can be tagged in the prompt with `<music>`, `<sfx>`, `<speech>`."
+    )
+    with gr.Row():
+        with gr.Column():
+            image = gr.Image(label="Reference image (first frame)", type="pil")
+            prompt = gr.Textbox(label="Prompt (video)", lines=6)
+            audio_prompt = gr.Textbox(label="Audio prompt (optional, defaults to the video prompt)", lines=2)
+            with gr.Row():
+                resolution = gr.Radio(list(RESOLUTIONS), value="720p (1280×720)", label="Resolution")
+                frames = gr.Slider(25, 121, value=49, step=4, label="Frames (24 fps)")
+            steps = gr.Slider(10, 50, value=30, step=1, label="Denoising steps")
+            estimate = gr.Markdown()
+            with gr.Accordion("Advanced", open=False):
+                cfg = gr.Slider(1.0, 10.0, value=5.0, step=0.5, label="CFG scale (1.0 = no CFG, ~2× faster)")
+                shift = gr.Slider(1.0, 17.0, value=7.0, step=0.5, label="Visual shift (720p: 7, 480p: 5)")
+                audio_shift = gr.Slider(1.0, 15.0, value=7.0, step=0.5, label="Audio shift")
+                sparsity = gr.Dropdown([0.75, 0.85, 0.93], value=0.75, label="Block-sparse attention sparsity")
+                use_ivpq = gr.Checkbox(value=True, label="Dynamic block shape (IVPQ)")
+                with gr.Row():
+                    seed = gr.Number(value=42, precision=0, label="Seed")
+                    randomize_seed = gr.Checkbox(value=False, label="Random seed")
+            run = gr.Button("Generate", variant="primary")
+        with gr.Column():
+            out_video = gr.Video(label="Result (with audio)")
+            used_seed = gr.Number(label="Seed used", interactive=False)
+            info = gr.Markdown()
+
+    est_inputs = [image, resolution, frames, steps, cfg, sparsity]
+    for c in est_inputs:
+        c.change(update_estimate, est_inputs, estimate)
+    demo.load(update_estimate, est_inputs, estimate)
+
+    run.click(
+        generate,
+        [image, prompt, audio_prompt, resolution, frames, steps, cfg, shift, audio_shift,
+         seed, randomize_seed, sparsity, use_ivpq],
+        [out_video, used_seed, info],
+    )
+    if examples:
+        gr.Examples(examples, [image, prompt])
+
+demo.queue(max_size=20).launch()

examples/case-1/prompt.txt

@@ -0,0 +1 @@
+No background music, quiet recording environment, close-mic capture, no noticeable reverb. Close-up, eye-level view, softly lit with frontal lighting. The frame features an elderly white woman with short gray curly hair, blue eyes, deeply lined face, flushed cheeks, and red lipstick. She wears a black coat and a red scarf with a white floral pattern, a black microphone clipped to her collar. She sits in front of a brown wooden piece of furniture, with a pale yellow textured wall behind her, on which hangs a miniature model depicting a canal and Dutch-style buildings on both sides. Initially, the elderly white woman faces the camera, lips moving, speaking in a hoarse voice with a heavy expression<speech>weg haden ze nog gezegd, ze liepen weg.</speech>. Then, she briefly lowers her head, her gaze moving downward, mouth slightly open as if sighing. Immediately afterward, she raises her head to look at the camera again, continuing<speech>Nou, dan kwam de directie er bij, want de directeur moest oepen.</speech>, blinking slowly throughout, her expression solemn.

examples/case-2/prompt.txt

@@ -0,0 +1 @@
+Continuous <sfx>ambient store sounds, including people talking and footsteps</sfx> throughout. Medium close-up, eye-level view. Inside an indoor store with bright lighting. The main focus is a man and a woman conversing. On the left side of the frame is a young-to-middle-aged white woman wearing glasses, with brown wavy medium-short hair, wearing a black jacket and a dark red inner layer, holding a dark-colored wine bottle with a black label featuring gold English text <text>THE髭振酒</text>. On the right side is a middle-aged white man with a beard, receding hairline, wearing a gray plaid suit jacket and a light blue shirt, holding a green wine bottle with a white label featuring a minimalist black line pattern and text. Behind them are shelves filled with products, with multiple customers and staff walking and shopping in the background, and a row of track spotlights on the ceiling. Initially, the man in the gray suit holds up the green wine bottle, looking at the woman across from him. Then he dramatically widens his eyes, showing the bottle to her, enthusiastically saying <speech>And the second one that we have a wine.</speech>. The woman lowers her head, looking at her own dark-colored bottle, smiling. Immediately afterward, the woman in the black jacket looks up at the man, asking curiously <speech>They are also do a wine?</speech>. The man nods affirmatively, saying <speech>Yep.</speech>. The woman confirms, saying <speech>The messa.</speech>. The man then raises his left hand in a gesture of explanation, continuing to speak <speech>So when you don't know it, you're kind of afraid because maple.</speech>"}

examples/case-3/prompt.txt

@@ -0,0 +1 @@
+Continuous <music>solemn and expressive classical chamber music performance throughout, featuring two violins, one cello, one harpsichord, and one bassoon, with intertwined melodies and rich harmonies producing noticeable reverb in the stone-walled church space.</music>. High-angle wide shot, fixed frame: Gothic stone church interior. Soft, cool light from left arched window. White walls with carved memorials and shield emblems. Gray stone floor, red-and-black diamond tile path center. Eight musicians in semicircle. Center: young white male harpsichordist (black long-sleeve, facing away). Foreground left: middle-aged white female violinist (blonde, purple lace top, black dress). Beside her: balding middle-aged white male violinist. Foreground right: young white male cellist (brown short hair). Background: middle-aged white female violinist (blonde), balding middle-aged white male bassoonist, two standing male musicians. Each has black music stand. Microphones on floor. As the shot begins, all musicians start playing their instruments, gently swaying their bodies in rhythm with the music. <music>A solemn and expressive classical chamber music performance begins, featuring two violins, one cello, one harpsichord, and one bassoon, with intertwined melodies and rich harmonies producing noticeable reverb in the stone-walled church space.</music> The male and female violinists on the left and the cellist on the right are focusedly drawing their bows across the strings, fingers moving swiftly along the strings. The central harpsichordist's hands move rapidly across the keyboard while glancing sideways at the sheet music. The bassoonist in the background is playing his instrument. As the piece progresses, a cellist from the far right edge of the frame enters the shot and begins playing.

examples/case-4/prompt.txt

@@ -0,0 +1 @@
+<sfx>Constant noisy background sounds from the gym throughout, including equipment clanging, dumbbell plates hitting, and muffled conversations, with noticeable spatial reverb.</sfx>. Medium shot with a slight high-angle view, indoor gym setting, bright lighting. The main subject is a young white woman sitting on a black gym bench, with long golden ponytail, wearing white in-ear headphones, and a large black mandala tattoo on her left arm, though details like the medical giant logo beneath are somewhat blurred. She wears a white sports bra and olive-green ribbed leggings, and black-and-white Adidas sneakers. On the left side of the frame is a blue fitness equipment column with a warning label, an orange warning tag marked <text>WARNING</text>, and a white label below. The background is a white wall with blue brand text visible, partially showing <text>Nordic</text>, and <text>to</text> visible on the right. Other gym equipment and multiple members exercising are visible, including a man in a black hoodie and a woman in a light gray hoodie. Initially, the blonde woman keeps her upper body lifted with bent legs, arms hugging her knees close to her chest. Her mouth is slightly open, looking directly at the camera, beginning her explanation in a slightly hoarse but clear and strong voice in Swedish: <speech>So try to do or this can a try to move bengal。</speech>. Then, while maintaining the knee-hug position, she extends her right index finger forward, then brings it back, adjusting her left hand's position on her knee. Next, she moves her left hand to support her left hip and waist, while bringing her right knee to the side and tilting her body backward. Subsequently, she keeps her left leg bent for support, extends her right leg, and slowly moves it upward and downward in rhythm with the command, completing one snow angel-like motion. Throughout and after the movement, she continues speaking directly to the camera: <speech>Can she tell me about the room and lift it little from that? Are they just broke in them now we could。</speech>

examples/case-5/prompt.txt

@@ -0,0 +1 @@
+Continuous <sfx>sound of waves and water currents around the boat</sfx> throughout. Medium shot, eye-level view. During the day, in clear outdoor conditions, a middle-aged white man stands on a boat facing the camera. He wears a gray baseball cap and blue-tinted sunglasses, a gray short-sleeved T-shirt with a brand logo featuring <text>YETI</text> on the chest, mostly obscured by the large fish he holds. He cradles a massive dark brown grouper with a wet, shiny surface reflecting sunlight and a serrated dorsal fin. A white outboard motor is located behind his left, marked with black <text>EVINRUDE</text>, partially hidden by the fish's tail at times. On the right side of the frame, another man's profile is visible—bearded, wearing a hat, holding a fishing rod. The background shows a vast blue ocean under a sky with a few scattered clouds. Initially, the man with sunglasses beams a bright smile, enthusiastically saying <speech>I mean, so far we've caught Marlin, Cooper, I mean, a Wahoo，we got slam. Why stop now？</speech> to the camera. His expression is highly animated. Then, he briefly turns his head to the right to look at the man with the fishing rod beside him, before turning back, maintaining a joyful expression throughout. The camera gently shakes due to the boat's movement during the entire sequence.

examples/case-6/prompt.txt

@@ -0,0 +1 @@
+No background music, with slight ambient outdoor noise. Eye-level medium close-up, centered composition. On a sunny day, a young white woman and a young South Asian man sit at a round black grid table under a black patio umbrella. The woman is on the left side of the frame, with long brown hair, wearing a pink sleeveless dress, and a thin black bracelet on her left wrist. The man is on the right side, with short black hair, wearing a blue short-sleeve Polo shirt with a black microphone clipped at the collar. A glass of water sits on the table in front of them. The scene is set on a wooden deck, with open water, distant apartment buildings, and green trees in the background, and large white clouds in the sky. Initially, the man in the blue Polo shirt looks at the woman across from him, raises both hands in gesture, and says in a gentle tone <speech>That's okay. We can order。</speech>. Then, the woman in the pink dress becomes serious, lowers her hands, and gently shakes her head, responding with a disappointed tone <speech>It's not about the wine. Nothing ever works out for me.</speech>. She then lifts her right hand and touches the glass on the table. The man leans forward slightly, listening intently. The woman then asks in a curious, probing tone <speech>Why do you say that?</speech>. The woman continues to look at him seriously and explains <speech>You never smelled at me when I smelled at you back at the start, so I decided to use hope this time.</speech>

hymm/diffusion/pipelines/mova_pipeline.py

@@ -0,0 +1,485 @@
+import copy
+import gc
+import html
+import re
+from typing import Any, List, Optional, Tuple, Union
+
+import ftfy
+import torch
+import torch.distributed as dist
+from diffusers.models.autoencoders import AutoencoderKLWan
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.video_processor import VideoProcessor
+from diffusers.image_processor import PipelineImageInput
+from tqdm import tqdm
+from transformers import T5TokenizerFast, UMT5EncoderModel
+
+from hymm.models.modules.mova import MOVABridge
+from hymm.models.modules.dac_vae import DAC
+from hymm.diffusion.schedulers.flow_match_pair import FlowMatchPairScheduler
+
+
+def _basic_clean(text):
+    text = ftfy.fix_text(text)
+    text = html.unescape(html.unescape(text))
+    return text.strip()
+
+
+def _whitespace_clean(text):
+    text = re.sub(r"\s+", " ", text)
+    return text.strip()
+
+
+def _prompt_clean(text):
+    return _whitespace_clean(_basic_clean(text))
+
+
+def _retrieve_latents(encoder_output, generator=None, sample_mode="sample"):
+    if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+        return encoder_output.latent_dist.sample(generator)
+    elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+        return encoder_output.latent_dist.mode()
+    elif hasattr(encoder_output, "latents"):
+        return encoder_output.latents
+    else:
+        raise AttributeError("Could not access latents of provided encoder_output")
+
+
+class MOVAPipeline:
+    """
+    Pipeline that drives MOVA inference identically to the original
+    MOVA/mova/diffusion/pipelines/pipeline_mova.py.
+
+    Differences from the original:
+    - Does NOT inherit from DiffusionPipeline (no auto-offload magic).
+    - Uses MOVABridge as the transformer, which already holds video_dit(s),
+      audio_dit, and bridge.
+    - Compatible with the codebase's DeepSpeed / FSDP launcher.
+    """
+
+    def __init__(
+        self,
+        transformer: MOVABridge,
+        video_vae: AutoencoderKLWan,
+        audio_vae: DAC,
+        text_encoder: UMT5EncoderModel,
+        tokenizer: T5TokenizerFast,
+        scheduler: FlowMatchPairScheduler,
+        boundary_ratio: float = 0.9,
+        device=None,
+    ):
+        self.transformer = transformer
+        self.video_vae = video_vae
+        self.audio_vae = audio_vae
+        self.text_encoder = text_encoder
+        self.tokenizer = tokenizer
+        self.scheduler = scheduler
+        self.boundary_ratio = boundary_ratio
+
+        self.vae_scale_factor_spatial = self.video_vae.config.scale_factor_spatial
+        self.vae_scale_factor_temporal = self.video_vae.config.scale_factor_temporal
+        self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
+
+        self.audio_vae_scale_factor = int(self.audio_vae.hop_length)
+        self.audio_sample_rate = self.audio_vae.sample_rate
+
+        self._device = device
+        self._cpu_offload = False
+
+    @property
+    def device(self):
+        if self._device is not None:
+            return self._device
+        for p in self.transformer.parameters():
+            return p.device
+        return torch.device("cpu")
+
+    def to(self, device):
+        self._device = device
+        self.transformer.to(device)
+        self.video_vae.to(device)
+        self.audio_vae.to(device)
+        self.text_encoder.to(device)
+        return self
+
+    def enable_cpu_offload(self, gpu_device=None):
+        """Enable sequential CPU offload to reduce peak GPU memory.
+
+        Only the active component lives on GPU at any time:
+          1. text_encoder on GPU -> encode prompts -> move to CPU
+          2. video_vae on GPU -> encode condition latents -> move to CPU
+          3. transformer on GPU -> denoising loop -> move to CPU
+          4. video_vae + audio_vae on GPU -> decode -> move to CPU
+        """
+        self._cpu_offload = True
+        if gpu_device is not None:
+            self._device = gpu_device
+
+    def normalize_video_latents(self, latents):
+        mean = torch.tensor(
+            self.video_vae.config.latents_mean, device=latents.device, dtype=latents.dtype,
+        ).view(1, self.video_vae.config.z_dim, 1, 1, 1)
+        inv_std = (1.0 / torch.tensor(
+            self.video_vae.config.latents_std, device=latents.device, dtype=latents.dtype,
+        )).view(1, self.video_vae.config.z_dim, 1, 1, 1)
+        return (latents - mean) * inv_std
+
+    def denormalize_video_latents(self, latents):
+        mean = torch.tensor(
+            self.video_vae.config.latents_mean, device=latents.device, dtype=latents.dtype,
+        ).view(1, self.video_vae.config.z_dim, 1, 1, 1)
+        std = torch.tensor(
+            self.video_vae.config.latents_std, device=latents.device, dtype=latents.dtype,
+        ).view(1, self.video_vae.config.z_dim, 1, 1, 1)
+        return latents * std + mean
+
+    def check_inputs(self, height, width, num_frames):
+        target_div = self.vae_scale_factor_spatial * 2
+        if height % target_div != 0 or width % target_div != 0:
+            raise ValueError(
+                f"`height` and `width` must be divisible by {target_div}, "
+                f"got {height} and {width}."
+            )
+        if num_frames % self.vae_scale_factor_temporal != 1:
+            raise ValueError(
+                f"`num_frames - 1` must be divisible by {self.vae_scale_factor_temporal}, "
+                f"got {num_frames - 1}."
+            )
+
+    def prepare_latents(
+        self,
+        image,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        num_frames,
+        dtype,
+        device,
+        generator=None,
+        latents=None,
+        last_image=None,
+    ):
+        num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
+        latent_height = height // self.vae_scale_factor_spatial
+        latent_width = width // self.vae_scale_factor_spatial
+
+        shape = (batch_size, num_channels_latents, num_latent_frames, latent_height, latent_width)
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            latents = latents.to(device=device, dtype=dtype)
+
+        image = image.unsqueeze(2)
+
+        if last_image is None:
+            video_condition = torch.cat(
+                [image, image.new_zeros(image.shape[0], image.shape[1], num_frames - 1, height, width)], dim=2,
+            )
+        else:
+            last_image = last_image.unsqueeze(2)
+            video_condition = torch.cat(
+                [image, image.new_zeros(image.shape[0], image.shape[1], num_frames - 2, height, width), last_image],
+                dim=2,
+            )
+        video_condition = video_condition.to(device=device, dtype=self.video_vae.dtype)
+
+        if isinstance(generator, list):
+            latent_condition = [
+                _retrieve_latents(self.video_vae.encode(video_condition), sample_mode="argmax")
+                for _ in generator
+            ]
+            latent_condition = torch.cat(latent_condition)
+        else:
+            latent_condition = _retrieve_latents(
+                self.video_vae.encode(video_condition), sample_mode="argmax",
+            )
+            latent_condition = latent_condition.repeat(batch_size, 1, 1, 1, 1)
+
+        latent_condition = latent_condition.to(dtype)
+        latent_condition = self.normalize_video_latents(latent_condition)
+
+        mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height, latent_width)
+        if last_image is None:
+            mask_lat_size[:, :, list(range(1, num_frames))] = 0
+        else:
+            mask_lat_size[:, :, list(range(1, num_frames - 1))] = 0
+        first_frame_mask = mask_lat_size[:, :, 0:1]
+        first_frame_mask = torch.repeat_interleave(first_frame_mask, dim=2, repeats=self.vae_scale_factor_temporal)
+        mask_lat_size = torch.concat([first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2)
+        mask_lat_size = mask_lat_size.view(batch_size, -1, self.vae_scale_factor_temporal, latent_height, latent_width)
+        mask_lat_size = mask_lat_size.transpose(1, 2)
+        mask_lat_size = mask_lat_size.to(latent_condition.device)
+
+        return latents, torch.concat([mask_lat_size, latent_condition], dim=1)
+
+    def prepare_audio_latents(
+        self,
+        audio,
+        batch_size,
+        num_channels,
+        num_samples,
+        dtype,
+        device,
+        generator=None,
+        latents=None,
+    ):
+        latent_t = (num_samples - 1) // self.audio_vae_scale_factor + 1
+        shape = (batch_size, num_channels, latent_t)
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            latents = latents.to(device=device, dtype=dtype)
+        return latents
+
+    def _get_t5_prompt_embeds(
+        self,
+        prompt,
+        num_videos_per_prompt=1,
+        max_sequence_length=512,
+        device=None,
+        dtype=None,
+    ):
+        device = device or self.device
+        dtype = dtype or self.text_encoder.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        prompt = [_prompt_clean(u) for u in prompt]
+        batch_size = len(prompt)
+
+        text_inputs = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=max_sequence_length,
+            truncation=True,
+            add_special_tokens=True,
+            return_attention_mask=True,
+            return_tensors="pt",
+        )
+        text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask
+        seq_lens = mask.gt(0).sum(dim=1).long()
+
+        prompt_embeds = self.text_encoder(
+            text_input_ids.to(device), mask.to(device),
+        ).last_hidden_state
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+        prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)]
+        prompt_embeds = torch.stack(
+            [torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds],
+            dim=0,
+        )
+
+        _, seq_len, _ = prompt_embeds.shape
+        prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
+        return prompt_embeds
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt,
+        image,
+        audio_prompt=None,
+        negative_prompt="",
+        seed=42,
+        height=360,
+        width=640,
+        num_frames=193,
+        video_fps=24.0,
+        num_inference_steps=50,
+        visual_shift=5.0,
+        audio_shift=5.0,
+        cfg_scale=5.0,
+        remove_video_dit=False,
+        enable_vae_tiling=False,
+        vae_tile_sample_min_size=None,
+        vae_tile_sample_stride=None,
+    ):
+        """
+        Run the full MOVA denoising loop.
+        This replicates the exact logic from MOVA's __call__.
+
+        VAE tiling (Phase 4 decode only, fully isolated / default OFF):
+            enable_vae_tiling: when True, the video VAE decodes the latents tile by
+                tile (spatial), drastically cutting peak decode VRAM so that high-res
+                (e.g. 2K) videos no longer OOM. When False the decode path is byte-for-
+                byte identical to the original (VAE tiling is never touched).
+            vae_tile_sample_min_size: tile size (pixels, applied to both H and W).
+                None -> diffusers/Wan default (256), same as decode_wan_vae.py.
+            vae_tile_sample_stride: stride between tiles (pixels, H and W). overlap =
+                tile_size - stride. None -> diffusers/Wan default (192), i.e. 25% overlap,
+                same as decode_wan_vae.py.
+
+        Note on FSDP + SP: only the transformer is FSDP-sharded; the video VAE is a
+        full replica on every rank and the post-denoise latents are identical across
+        all SP ranks. VAE tiling is deterministic, so every rank produces the SAME
+        decoded video whether or not tiling is enabled.
+        """
+        self.check_inputs(height, width, num_frames)
+
+        audio_num_samples = int(self.audio_sample_rate * num_frames / video_fps)
+        device = self.device
+        offload = self._cpu_offload
+
+        # Deterministic noise generator (independent of global random state)
+        generator = torch.Generator(device=device)
+        generator.manual_seed(seed)
+
+        # ------------------------------------------------------------------
+        # Phase 1: Text encoding
+        # ------------------------------------------------------------------
+        if offload:
+            self.text_encoder.to(device)
+
+        self.scheduler.set_timesteps(num_inference_steps, shift=visual_shift, device=device)
+        self.scheduler.set_pair_postprocess_by_name(
+            "dual_sigma_shift",
+            visual_shift=float(visual_shift),
+            audio_shift=float(audio_shift),
+        )
+        paired_timesteps = self.scheduler.get_pairs()
+
+        prompt_embeds = self._get_t5_prompt_embeds(prompt)
+        audio_prompt_embeds = self._get_t5_prompt_embeds(audio_prompt) if audio_prompt is not None else None
+        negative_prompt_embeds = self._get_t5_prompt_embeds(negative_prompt)
+
+        if offload:
+            self.text_encoder.to("cpu")
+            torch.cuda.empty_cache()
+
+        # ------------------------------------------------------------------
+        # Phase 2: VAE encode (prepare latents)
+        # ------------------------------------------------------------------
+        if offload:
+            self.video_vae.to(device)
+
+        num_channels_latents = self.video_vae.config.z_dim
+        image = self.video_processor.preprocess(image, height=height, width=width).to(device, dtype=torch.float32)
+
+        latents, condition = self.prepare_latents(
+            image, 1, num_channels_latents, height, width, num_frames,
+            torch.float32, device, generator=generator, latents=None, last_image=None,
+        )
+        audio_latents = self.prepare_audio_latents(
+            None, 1, self.audio_vae.latent_dim, audio_num_samples,
+            torch.float32, device, generator=generator, latents=None,
+        )
+
+        if offload:
+            self.video_vae.to("cpu")
+            torch.cuda.empty_cache()
+
+        # ------------------------------------------------------------------
+        # Phase 3: Denoising loop (transformer on GPU)
+        # ------------------------------------------------------------------
+        if offload:
+            self.transformer.to(device)
+
+        cur_use_dit_2 = False
+        total_steps = paired_timesteps.shape[0]
+        switched = False
+        boundary_timestep = self.boundary_ratio * self.scheduler.config.num_train_timesteps
+
+        is_main = not dist.is_initialized() or dist.get_rank() == 0
+
+        for idx_step in tqdm(range(total_steps), disable=not is_main):
+            timestep, audio_timestep = paired_timesteps[idx_step]
+
+            if not switched and timestep.item() < boundary_timestep:
+                cur_use_dit_2 = True
+                if remove_video_dit:
+                    self.transformer.video_dit = None
+                    gc.collect()
+                switched = True
+
+            latent_model_input = torch.cat([latents, condition], dim=1)
+            timestep_t = timestep.unsqueeze(0).to(device=device, dtype=torch.float32)
+            audio_timestep_t = audio_timestep.unsqueeze(0).to(device=device, dtype=torch.float32)
+
+            noise_pred_posi = self.transformer(
+                visual_latents=latent_model_input,
+                audio_latents=audio_latents,
+                context=prompt_embeds,
+                timestep=timestep_t,
+                audio_context=audio_prompt_embeds,
+                audio_timestep=audio_timestep_t,
+                video_fps=video_fps,
+                num_train_timesteps=self.scheduler.config.num_train_timesteps,
+                use_video_dit_2=cur_use_dit_2,
+            )
+
+            if cfg_scale == 1.0:
+                visual_noise_pred = noise_pred_posi[0].float()
+                audio_noise_pred = noise_pred_posi[1].float()
+            else:
+                noise_pred_nega = self.transformer(
+                    visual_latents=latent_model_input,
+                    audio_latents=audio_latents,
+                    context=negative_prompt_embeds,
+                    timestep=timestep_t,
+                    audio_context=None,
+                    audio_timestep=audio_timestep_t,
+                    video_fps=video_fps,
+                    num_train_timesteps=self.scheduler.config.num_train_timesteps,
+                    use_video_dit_2=cur_use_dit_2,
+                )
+                visual_noise_pred_nega = noise_pred_nega[0].float()
+                audio_noise_pred_nega = noise_pred_nega[1].float()
+                visual_noise_pred_posi = noise_pred_posi[0].float()
+                audio_noise_pred_posi = noise_pred_posi[1].float()
+                visual_noise_pred = visual_noise_pred_nega + cfg_scale * (visual_noise_pred_posi - visual_noise_pred_nega)
+                audio_noise_pred = audio_noise_pred_nega + cfg_scale * (audio_noise_pred_posi - audio_noise_pred_nega)
+
+            next_timestep = paired_timesteps[idx_step + 1, 0] if idx_step + 1 < total_steps else None
+            next_audio_timestep = paired_timesteps[idx_step + 1, 1] if idx_step + 1 < total_steps else None
+            latents = self.scheduler.step_from_to(
+                visual_noise_pred, timestep, next_timestep, latents,
+            )
+            audio_latents = self.scheduler.step_from_to(
+                audio_noise_pred, audio_timestep, next_audio_timestep, audio_latents,
+            )
+
+        if offload:
+            self.transformer.to("cpu")
+            torch.cuda.empty_cache()
+
+        # ------------------------------------------------------------------
+        # Phase 4: VAE decode
+        # ------------------------------------------------------------------
+        if offload:
+            self.video_vae.to(device)
+
+        # Optional tiled VAE decode (default OFF -> original path untouched). When
+        # enabled, decode the latents tile-by-tile spatially to cut peak decode VRAM
+        # for high-res (e.g. 2K) videos. Mirrors decode_wan_vae.py: min_size/stride
+        # map to the AutoencoderKLWan (height & width) tiling knobs; None keeps the
+        # diffusers defaults (256 / 192).
+        if enable_vae_tiling:
+            tiling_kwargs = {}
+            if vae_tile_sample_min_size is not None:
+                tiling_kwargs["tile_sample_min_height"] = vae_tile_sample_min_size
+                tiling_kwargs["tile_sample_min_width"] = vae_tile_sample_min_size
+            if vae_tile_sample_stride is not None:
+                tiling_kwargs["tile_sample_stride_height"] = vae_tile_sample_stride
+                tiling_kwargs["tile_sample_stride_width"] = vae_tile_sample_stride
+            self.video_vae.enable_tiling(**tiling_kwargs)
+
+        video_latents = self.denormalize_video_latents(latents)
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            video = self.video_vae.decode(video_latents).sample
+        video = self.video_processor.postprocess_video(video, output_type="pil")
+
+        if offload:
+            self.video_vae.to("cpu")
+            self.audio_vae.to(device)
+            torch.cuda.empty_cache()
+
+        with torch.autocast("cuda", dtype=torch.float32):
+            audio = self.audio_vae.decode(audio_latents)
+
+        if offload:
+            self.audio_vae.to("cpu")
+            torch.cuda.empty_cache()
+
+        return video, audio

hymm/diffusion/schedulers/__init__.py

@@ -0,0 +1 @@
+from .flow_match_pair import FlowMatchPairScheduler

hymm/diffusion/schedulers/flow_match_pair.py

@@ -0,0 +1,277 @@
+import math
+import torch
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+
+
+class FlowMatchSchedulerBase:
+    def __init__(
+        self,
+        num_inference_steps=100,
+        num_train_timesteps=1000,
+        shift=3.0,
+        sigma_max=1.0,
+        sigma_min=0.003 / 1.002,
+        inverse_timesteps=False,
+        extra_one_step=False,
+        reverse_sigmas=False,
+        exponential_shift=False,
+        exponential_shift_mu=None,
+        shift_terminal=None,
+    ):
+        self.num_train_timesteps = num_train_timesteps
+        self.shift = shift
+        self.sigma_max = sigma_max
+        self.sigma_min = sigma_min
+        self.inverse_timesteps = inverse_timesteps
+        self.extra_one_step = extra_one_step
+        self.reverse_sigmas = reverse_sigmas
+        self.exponential_shift = exponential_shift
+        self.exponential_shift_mu = exponential_shift_mu
+        self.shift_terminal = shift_terminal
+        self.train_timesteps = None
+        self.train_sigmas = None
+        self.set_timesteps(num_train_timesteps)
+        self.set_timesteps(num_inference_steps)
+
+    def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, training=False, shift=None, dynamic_shift_len=None, device=None):
+        if shift is not None:
+            self.shift = shift
+        sigma_start = self.sigma_min + (self.sigma_max - self.sigma_min) * denoising_strength
+        if self.extra_one_step:
+            self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps + 1)[:-1]
+        else:
+            self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps)
+        if self.inverse_timesteps:
+            self.sigmas = torch.flip(self.sigmas, dims=[0])
+        if self.exponential_shift:
+            mu = self.calculate_shift(dynamic_shift_len) if dynamic_shift_len is not None else self.exponential_shift_mu
+            self.sigmas = math.exp(mu) / (math.exp(mu) + (1 / self.sigmas - 1))
+        else:
+            self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas)
+        if self.shift_terminal is not None:
+            one_minus_z = 1 - self.sigmas
+            scale_factor = one_minus_z[-1] / (1 - self.shift_terminal)
+            self.sigmas = 1 - (one_minus_z / scale_factor)
+        if self.reverse_sigmas:
+            self.sigmas = 1 - self.sigmas
+        self.timesteps = self.sigmas * self.num_train_timesteps
+        if self.train_timesteps is None:
+            self.train_timesteps = self.timesteps
+            self.train_sigmas = self.sigmas
+        if training:
+            x = self.timesteps
+            y = torch.exp(-2 * ((x - num_inference_steps / 2) / num_inference_steps) ** 2)
+            y_shifted = y - y.min()
+            bsmntw_weighing = y_shifted * (num_inference_steps / y_shifted.sum())
+            self.linear_timesteps_weights = bsmntw_weighing
+            self.training = True
+        else:
+            self.training = False
+
+    def step(self, model_output, timestep, sample, to_final=False, **kwargs):
+        if isinstance(timestep, torch.Tensor):
+            timestep = timestep.cpu()
+        timestep_id = torch.argmin((self.timesteps - timestep).abs())
+        sigma = self.sigmas[timestep_id]
+        if to_final or timestep_id + 1 >= len(self.timesteps):
+            sigma_ = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0
+        else:
+            sigma_ = self.sigmas[timestep_id + 1]
+        prev_sample = sample + model_output * (sigma_ - sigma)
+        return prev_sample
+
+    def add_noise(self, original_samples, noise, timestep):
+        if isinstance(timestep, torch.Tensor):
+            timestep = timestep.cpu()
+        timestep_id = torch.argmin((self.timesteps - timestep).abs())
+        sigma = self.sigmas[timestep_id]
+        sample = (1 - sigma) * original_samples + sigma * noise
+        return sample
+
+    def training_target(self, sample, noise, timestep):
+        return noise - sample
+
+    def calculate_shift(
+        self,
+        image_seq_len,
+        base_seq_len: int = 256,
+        max_seq_len: int = 8192,
+        base_shift: float = 0.5,
+        max_shift: float = 0.9,
+    ):
+        m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+        b = base_shift - m * base_seq_len
+        mu = image_seq_len * m + b
+        return mu
+
+
+class FlowMatchPairScheduler(FlowMatchSchedulerBase, SchedulerMixin, ConfigMixin):
+    """
+    Paired timestep scheduler for dual-tower (video + audio) diffusion.
+    Returns tensor with shape [num_timesteps, 2], each row is (visual_t, audio_t).
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        num_inference_steps=100,
+        num_train_timesteps=1000,
+        shift=3.0,
+        sigma_max=1.0,
+        sigma_min=0.003 / 1.002,
+        inverse_timesteps=False,
+        extra_one_step=False,
+        reverse_sigmas=False,
+        exponential_shift=False,
+        exponential_shift_mu=None,
+        shift_terminal=None,
+    ):
+        self._pair_postprocess_fn = None
+        self._pair_postprocess_requires_source = False
+        self.pair_timesteps = None
+        self.pair_sigmas = None
+        self.timesteps = None
+        self.sigmas = None
+        super().__init__(
+            num_inference_steps=num_inference_steps,
+            num_train_timesteps=num_train_timesteps,
+            shift=shift,
+            sigma_max=sigma_max,
+            sigma_min=sigma_min,
+            inverse_timesteps=inverse_timesteps,
+            extra_one_step=extra_one_step,
+            reverse_sigmas=reverse_sigmas,
+            exponential_shift=exponential_shift,
+            exponential_shift_mu=exponential_shift_mu,
+            shift_terminal=shift_terminal,
+        )
+
+    def set_pair_postprocess(self, fn):
+        if fn is not None and not callable(fn):
+            raise TypeError("pair_postprocess must be a callable object or None")
+        self._pair_postprocess_fn = fn
+        self._pair_postprocess_requires_source = False if fn is None else bool(getattr(fn, "_requires_source", False))
+        if self.timesteps is None or self.sigmas is None:
+            raise RuntimeError("Scheduler not initialized, please call set_timesteps() first")
+        self._refresh_pair_cache()
+
+    def set_pair_postprocess_by_name(self, name, **kwargs):
+        if name is None or str(name).lower() in ("none", "off", "false", "no"):
+            self.set_pair_postprocess(None)
+            return
+        if name == "dual_sigma_shift":
+            visual_shift = float(kwargs.get("visual_shift", self.shift))
+            audio_shift = float(kwargs.get("audio_shift", self.shift))
+            visual_denoising_strength = float(kwargs.get("visual_denoising_strength", 1.0))
+            audio_denoising_strength = float(kwargs.get("audio_denoising_strength", 1.0))
+            visual_mu = kwargs.get("visual_exponential_shift_mu", self.exponential_shift_mu)
+            audio_mu = kwargs.get("audio_exponential_shift_mu", self.exponential_shift_mu)
+
+            def _dual_sigma_shift(pairs, *, source):
+                num_steps = pairs.shape[0]
+                device = pairs.device
+                dtype = pairs.dtype
+
+                def _build_column(shift_value, denoising_strength, mu_override):
+                    sigma_start = self.sigma_min + (self.sigma_max - self.sigma_min) * denoising_strength
+                    if self.extra_one_step:
+                        base = torch.linspace(sigma_start, self.sigma_min, num_steps + 1, device=device, dtype=dtype)[:-1]
+                    else:
+                        base = torch.linspace(sigma_start, self.sigma_min, num_steps, device=device, dtype=dtype)
+                    if self.inverse_timesteps:
+                        base = torch.flip(base, dims=[0])
+                    if self.exponential_shift:
+                        mu_value = mu_override
+                        if mu_value is None:
+                            raise RuntimeError("exponential_shift is enabled but exponential_shift_mu is not provided")
+                        exp_mu = math.exp(float(mu_value))
+                        base = exp_mu / (exp_mu + (1 / base - 1))
+                    else:
+                        base = shift_value * base / (1 + (shift_value - 1) * base)
+                    if self.shift_terminal is not None:
+                        one_minus_z = 1 - base
+                        scale_factor = one_minus_z[-1] / (1 - self.shift_terminal)
+                        base = 1 - (one_minus_z / scale_factor)
+                    if self.reverse_sigmas:
+                        base = 1 - base
+                    if source == "timesteps":
+                        return base * self.num_train_timesteps
+                    return base
+
+                visual_column = _build_column(visual_shift, visual_denoising_strength, visual_mu)
+                audio_column = _build_column(audio_shift, audio_denoising_strength, audio_mu)
+                return torch.stack([visual_column, audio_column], dim=1)
+
+            _dual_sigma_shift._requires_source = True
+            self.set_pair_postprocess(_dual_sigma_shift)
+            return
+        raise ValueError(f"Unsupported pair postprocessing name: {name}")
+
+    def _make_pairs_from_vector(self, vec):
+        return torch.stack([vec, vec], dim=1)
+
+    def get_pairs(self, source="timesteps"):
+        if source == "timesteps":
+            pairs = self.pair_timesteps
+        elif source == "sigmas":
+            pairs = self.pair_sigmas
+        else:
+            raise ValueError("source only supports 'timesteps' or 'sigmas'")
+        if pairs is None:
+            raise RuntimeError("Scheduler not initialized, please call set_timesteps() first")
+        return pairs
+
+    @property
+    def visual_timesteps(self):
+        if self.pair_timesteps is None:
+            raise RuntimeError("Scheduler not initialized")
+        return self.pair_timesteps[:, 0]
+
+    @property
+    def audio_timesteps(self):
+        if self.pair_timesteps is None:
+            raise RuntimeError("Scheduler not initialized")
+        return self.pair_timesteps[:, 1]
+
+    def set_timesteps(self, *args, **kwargs):
+        super().set_timesteps(*args, **kwargs)
+        self._refresh_pair_cache()
+
+    def timestep_to_sigma(self, timestep):
+        t_value = float(timestep)
+        t_cpu = torch.tensor(t_value)
+        idx = torch.argmin((self.train_timesteps - t_cpu).abs())
+        return self.train_sigmas[idx]
+
+    def step_from_to(self, model_output, timestep_from, timestep_to, sample):
+        sigma_from = self.timestep_to_sigma(timestep_from)
+        if timestep_to is None:
+            sigma_to = torch.tensor(
+                1.0 if (self.inverse_timesteps or self.reverse_sigmas) else 0.0,
+                device=sigma_from.device, dtype=sigma_from.dtype,
+            )
+        else:
+            sigma_to = self.timestep_to_sigma(timestep_to)
+        prev_sample = sample + model_output * (sigma_to - sigma_from)
+        return prev_sample
+
+    def _refresh_pair_cache(self):
+        if self.timesteps is None or self.sigmas is None:
+            return
+
+        def _apply_postprocess(pairs, source):
+            if self._pair_postprocess_fn is None:
+                return pairs
+            if self._pair_postprocess_requires_source:
+                modified = self._pair_postprocess_fn(pairs, source=source)
+            else:
+                modified = self._pair_postprocess_fn(pairs)
+            return modified
+
+        base_pairs_timesteps = self._make_pairs_from_vector(self.timesteps)
+        base_pairs_sigmas = self._make_pairs_from_vector(self.sigmas)
+
+        self.pair_timesteps = _apply_postprocess(base_pairs_timesteps, "timesteps")
+        self.pair_sigmas = _apply_postprocess(base_pairs_sigmas, "sigmas")

hymm/models/modules/__init__.py

@@ -0,0 +1 @@
+from .mova import MOVABridge, FusedMOVABlock

hymm/models/modules/block_sparse_attention/__init__.py

@@ -0,0 +1,16 @@
+from .bsa_interface import (
+    flash_attn_bsa_3d,
+    flash_attn_bsa_cross,
+    flash_attn_bsa_3d_audio_guided,
+    flash_attn_bsa_3d_variance_guided,
+    compute_audio_spatial_concentration_gate,
+    compute_timestep_reliability_gate,
+    compute_block_audio_saliency,
+    audio_weighted_mean_pooling_compression,
+    compute_channel_variance_density,
+)
+
+from .bias_rectification import (
+    bsa_taylor_sparse_attn,
+    bsa_rectified_sparse_attn,
+)

hymm/models/modules/block_sparse_attention/bias_rectification.py

@@ -0,0 +1,452 @@
+
+
+from typing import Optional, Tuple
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+
+from .bsa_interface import (
+    mean_pooling_compression,
+    masked_mean_pooling_compression,
+    cal_score,
+    get_select_indices_topk_from_score,
+    get_select_indices_cdf_from_score,
+    get_select_indices_cdf_topk_from_score,
+    create_mask_from_indices_varlen,
+    attn_fwd_bsa_varlen_triton,
+    attn_bwd_bsa_varlen_triton,
+)
+
+
+# =====================================================================
+# Pure BSA Triton wrapper — only this is an autograd Function.
+# Everything else (pooling, scoring, correction) lives outside in
+# autograd-friendly PyTorch ops so gradients flow correctly.
+# =====================================================================
+class _bsa_kernel_only(torch.autograd.Function):
+    """Forward + backward through `attn_fwd_bsa_varlen_triton` only.
+
+    Returns (o, lse_natural). lse is the log-sum-exp of the *attended* keys
+    in natural log units (= LN2 * the kernel's log2 lse).
+    """
+
+    @staticmethod
+    def forward(ctx, q, k, v, sm_scale, block_indices, block_indices_lens,
+                chunk_size_q, chunk_size_k, sparsity, kv_valid_mask):
+        o, lse_natural = attn_fwd_bsa_varlen_triton(
+            q, k, v, sm_scale, block_indices, block_indices_lens,
+            chunk_size_q, chunk_size_k, sparsity,
+            kv_valid_mask=kv_valid_mask,
+        )
+        ctx.save_for_backward(q, k, v, o, lse_natural, block_indices, block_indices_lens)
+        ctx.sm_scale = sm_scale
+        ctx.chunk_size_q = chunk_size_q
+        ctx.chunk_size_k = chunk_size_k
+        ctx.sparsity = sparsity
+        ctx.kv_valid_mask = kv_valid_mask
+        return o, lse_natural
+
+    @staticmethod
+    def backward(ctx, do, dlse_unused):
+        q, k, v, o, lse, block_indices, block_indices_lens = ctx.saved_tensors
+        dq = torch.empty_like(q)
+        dk = torch.empty_like(k)
+        dv = torch.empty_like(v)
+        do = do.contiguous()
+        attn_bwd_bsa_varlen_triton(
+            do, q, k, v, o, dq, dk, dv,
+            ctx.sm_scale, lse,
+            block_indices, block_indices_lens,
+            ctx.chunk_size_q, ctx.chunk_size_k, ctx.sparsity,
+            kv_valid_mask=ctx.kv_valid_mask,
+        )
+        # 10 inputs: q, k, v, sm_scale, block_indices, block_indices_lens,
+        #            chunk_size_q, chunk_size_k, sparsity, kv_valid_mask
+        return dq, dk, dv, None, None, None, None, None, None, None
+
+
+def _bsa_kernel(q, k, v, sm_scale, block_indices, block_indices_lens,
+                chunk_size_q, chunk_size_k, sparsity, kv_valid_mask):
+    """Wrapper that exposes (o, lse) but always detaches lse.
+
+    `lse` is an attention normalisation statistic used downstream as a
+    coefficient in the Taylor merge weights. We treat it as a fixed
+    quantity (FlashAttention convention) and explicitly detach it so the
+    autograd engine does not try to backpropagate through a path we
+    intentionally do not implement in `_bsa_kernel_only.backward`.
+    """
+    o, lse = _bsa_kernel_only.apply(
+        q, k, v, sm_scale, block_indices, block_indices_lens,
+        chunk_size_q, chunk_size_k, sparsity, kv_valid_mask,
+    )
+    return o, lse.detach()
+
+
+# =====================================================================
+# Shared helpers (block-level pooling, selection, validity)
+# =====================================================================
+def _pool_qk(q, k, chunk_size_q, chunk_size_k, kv_valid_mask, q_valid_mask):
+    """Pool Q and K to block representatives, masking padded tokens.
+
+    Self-attn case (Q/K share grid): both use kv_valid_mask.
+    Cross-attn case: each side uses its own valid_mask.
+
+    All ops are autograd-tracked (gradients to q / k flow through pooling).
+    """
+    if q_valid_mask is not None:
+        q_cmp = masked_mean_pooling_compression(q, chunk_size_q, q_valid_mask)
+    elif kv_valid_mask is not None and q.shape[2] == k.shape[2]:
+        # self-attention: Q shares K's valid mask
+        q_cmp = masked_mean_pooling_compression(q, chunk_size_q, kv_valid_mask)
+    else:
+        q_cmp = mean_pooling_compression(q, chunk_size_q)
+    if kv_valid_mask is not None:
+        k_cmp = masked_mean_pooling_compression(k, chunk_size_k, kv_valid_mask)
+    else:
+        k_cmp = mean_pooling_compression(k, chunk_size_k)
+    return q_cmp, k_cmp
+
+
+def _pool_v(v, chunk_size_k, kv_valid_mask):
+    if kv_valid_mask is not None:
+        return masked_mean_pooling_compression(v, chunk_size_k, kv_valid_mask)
+    return mean_pooling_compression(v, chunk_size_k)
+
+
+def _select_blocks(selection_score, sparsity, cdf_threshold, sm_scale):
+    """topk / top-p / hybrid selection (non-differentiable)."""
+    if sparsity is not None and cdf_threshold is None:
+        return get_select_indices_topk_from_score(selection_score, sparsity)
+    if sparsity is None and cdf_threshold is not None:
+        return get_select_indices_cdf_from_score(selection_score, cdf_threshold, sm_scale)
+    if sparsity is not None and cdf_threshold is not None:
+        return get_select_indices_cdf_topk_from_score(selection_score, sparsity, cdf_threshold, sm_scale)
+    raise ValueError("Either sparsity or cdf_threshold must be provided")
+
+
+def _compute_block_valid_mask(kv_valid_mask: torch.Tensor, chunk_size: int) -> torch.Tensor:
+    """A K block is valid if any token within it is valid."""
+    S = kv_valid_mask.shape[0]
+    num_blocks = S // chunk_size
+    return kv_valid_mask.view(num_blocks, chunk_size).any(dim=-1)
+
+
+def _compute_block_valid_token_counts(kv_valid_mask: torch.Tensor, chunk_size: int) -> torch.Tensor:
+    S = kv_valid_mask.shape[0]
+    num_blocks = S // chunk_size
+    return kv_valid_mask.view(num_blocks, chunk_size).sum(dim=-1).float()
+
+
+def _masked_softmax_along_last(score: torch.Tensor, sm_scale: float,
+                               block_valid: Optional[torch.Tensor]) -> torch.Tensor:
+    """softmax(score * sm_scale, dim=-1) with -inf at padded K blocks."""
+    if block_valid is None:
+        return torch.softmax(score * sm_scale, dim=-1)
+    masked = score.masked_fill(~block_valid.view(1, 1, 1, -1), float('-inf'))
+    return torch.softmax(masked * sm_scale, dim=-1)
+
+
+def _masked_var_along_last(probs: torch.Tensor,
+                           block_valid: Optional[torch.Tensor]) -> torch.Tensor:
+    """Variance of softmax distribution restricted to valid K blocks.
+
+    Without masking, padded blocks contribute zero (since A_pool=0 there) but
+    inflate the denominator → biases variance downward → biases the LIVEditor
+    flat/sharp split. We compute Var[p] only over valid entries.
+    """
+    if block_valid is None:
+        return probs.var(dim=-1, unbiased=False)
+    bv = block_valid.view(1, 1, 1, -1).to(probs.dtype)
+    n_valid = bv.sum(dim=-1).clamp(min=1.0)
+    p_masked = probs * bv
+    mean = p_masked.sum(dim=-1, keepdim=True) / n_valid
+    var = ((p_masked - mean) ** 2 * bv).sum(dim=-1) / n_valid
+    return var
+
+
+def _global_quantile_across_sp(values: torch.Tensor, q: float) -> torch.Tensor:
+
+    # Imported at call scope but NOT guarded by try/except: swallowing an import
+    # error here would make one rank skip the all_gather below while its peers
+    # still enter it, which is an unrecoverable NCCL hang rather than a crash.
+    from hymm.utils.parallel_states import nccl_info, get_sequence_parallel_state
+
+    if (not get_sequence_parallel_state()) or nccl_info.sp_size <= 1:
+        return torch.quantile(values.float().reshape(-1), q)
+
+    sp_group = nccl_info.sp_group
+    flat = values.float().reshape(-1).contiguous()
+
+    # All ranks in SP have the same flat shape because Q/K/V are split by HEAD,
+    # not by token. So we can use a tensor all-gather.
+    world_size = nccl_info.sp_size
+    gather_list = [torch.empty_like(flat) for _ in range(world_size)]
+    dist.all_gather(gather_list, flat, group=sp_group)
+    global_flat = torch.cat(gather_list, dim=0)
+    return torch.quantile(global_flat, q)
+
+
+def _taylor_merge_pytorch(
+    q: torch.Tensor,
+    k_cmp: torch.Tensor,
+    v_cmp: torch.Tensor,
+    o_spa: torch.Tensor,
+    lse_spa: torch.Tensor,
+    nonsel_mask: torch.Tensor,
+    is_flat: torch.Tensor,
+    block_token_counts: torch.Tensor,
+    chunk_size_q: int,
+    sm_scale: float,
+) -> torch.Tensor:
+
+    B, H, S_q, D = o_spa.shape
+    N_q, N_k = nonsel_mask.shape[-2], nonsel_mask.shape[-1]
+
+    nonsel_per_tok = (
+        nonsel_mask.unsqueeze(3)
+        .expand(B, H, N_q, chunk_size_q, N_k)
+        .contiguous()
+        .reshape(B, H, S_q, N_k)
+    )
+
+    raw_scores = torch.matmul(q, k_cmp.transpose(-1, -2)) * sm_scale  # [B, H, S_q, N_k]
+    log_counts = torch.log(block_token_counts.float().clamp(min=1.0)).view(1, 1, 1, N_k)
+    scores_full = raw_scores + log_counts.to(raw_scores.dtype)
+
+    # ---- LSE: -inf-masked path; logsumexp handles all-(-inf) rows safely ----
+    scores_inf = scores_full.masked_fill(~nonsel_per_tok, float('-inf'))
+    lse_taylor = torch.logsumexp(scores_inf, dim=-1)             # [B, H, S_q]
+    has_taylor = torch.isfinite(lse_taylor)                      # bool
+
+    # ---- Softmax: use a finite sentinel to avoid NaN; gate rows w/o nonsel ----
+    LARGE_NEG = -1.0e30
+    scores_safe = scores_full.masked_fill(~nonsel_per_tok, LARGE_NEG)
+    weights = torch.softmax(scores_safe, dim=-1)                 # finite everywhere
+    # For rows with NO valid nonsel, every entry is LARGE_NEG → softmax = uniform
+    # → zero them out so o_taylor = 0 there.
+    weights = weights * has_taylor.unsqueeze(-1).to(weights.dtype)
+    o_taylor = torch.matmul(weights, v_cmp)
+
+    # ---- Stable lse merge ----
+    lse_max = torch.maximum(lse_spa, lse_taylor)
+    finite = torch.isfinite(lse_max)
+
+    # diff is undefined when both lses are -inf (rare; lse_spa is always finite
+    # in practice because BSA selects ≥1 block). Replace those NaNs by 0 inside
+    # the discarded `where` branch so backward never sees NaN.
+    diff = torch.minimum(lse_spa, lse_taylor) - lse_max
+    diff = torch.where(finite, diff, torch.zeros_like(diff))
+    lse_total = torch.where(finite, lse_max + torch.log1p(torch.exp(diff)), lse_max)
+
+    w_spa = torch.where(finite, torch.exp(lse_spa - lse_total), torch.ones_like(lse_total))
+    w_taylor = torch.where(
+        finite & has_taylor,
+        torch.exp(lse_taylor - lse_total),
+        torch.zeros_like(lse_total),
+    )
+
+    merged = w_spa.unsqueeze(-1) * o_spa + w_taylor.unsqueeze(-1) * o_taylor
+
+    # Sharp queries (block-level boolean) keep o_spa unchanged.
+    flat_tok = (
+        is_flat.unsqueeze(-1)
+        .expand(B, H, N_q, chunk_size_q)
+        .contiguous()
+        .reshape(B, H, S_q)
+    )
+    return torch.where(flat_tok.unsqueeze(-1), merged, o_spa).to(o_spa.dtype).contiguous()
+
+
+def bsa_taylor_sparse_attn(
+    q: torch.Tensor,                  # [B, H, S_q, D] in 3d-block layout
+    k: torch.Tensor,                  # [B, H, S_k, D]
+    v: torch.Tensor,                  # [B, H, S_k, D]
+    chunk_size_q: int,
+    chunk_size_k: int,
+    sparsity: Optional[float],
+    cdf_threshold: Optional[float],
+    sm_scale: float,
+    *,
+    alpha_f: float = 0.5,
+    kv_valid_mask: Optional[torch.Tensor] = None,
+    q_valid_mask: Optional[torch.Tensor] = None,
+    score_modifier: Optional[torch.Tensor] = None,
+    use_triton_inference: bool = True,
+) -> torch.Tensor:
+    # --- (a) pool Q and K (autograd-tracked) ---
+    q_cmp, k_cmp = _pool_qk(q, k, chunk_size_q, chunk_size_k, kv_valid_mask, q_valid_mask)
+
+    # --- (b) block scoring + selection (selection itself is non-differentiable) ---
+    raw_score = cal_score(q_cmp, k_cmp)  # [B, H, N_q, N_k]
+    with torch.no_grad():
+        sel_score = raw_score * score_modifier if score_modifier is not None else raw_score
+        block_indices, block_indices_lens = _select_blocks(sel_score, sparsity, cdf_threshold, sm_scale)
+
+    # --- (c) BSA Triton kernel; this is the only autograd Function call ---
+    o_spa, lse_spa = _bsa_kernel(
+        q, k, v, sm_scale, block_indices, block_indices_lens,
+        chunk_size_q, chunk_size_k, sparsity, kv_valid_mask,
+    )
+
+    # --- (d) Taylor side products (block_valid, A_pool, M_n, is_flat) ---
+    block_valid = _compute_block_valid_mask(kv_valid_mask, chunk_size_k) if kv_valid_mask is not None else None
+
+    with torch.no_grad():
+        sel_mask = create_mask_from_indices_varlen(block_indices, raw_score.shape[-1])
+        nonsel_mask = ~sel_mask
+        if block_valid is not None:
+            nonsel_mask = nonsel_mask & block_valid.view(1, 1, 1, -1)
+
+        # Block-level true probability distribution (no score_modifier, no scale boost)
+        A_pool_for_class = _masked_softmax_along_last(raw_score.detach(), sm_scale, block_valid)
+        M_n = _masked_var_along_last(A_pool_for_class, block_valid)  # [B, H, N_q]
+        tau_M = _global_quantile_across_sp(M_n, alpha_f)
+        is_flat = M_n < tau_M
+
+        if kv_valid_mask is not None:
+            btc = _compute_block_valid_token_counts(kv_valid_mask, chunk_size_k)
+        else:
+            btc = torch.full(
+                (k.shape[2] // chunk_size_k,), float(chunk_size_k),
+                device=q.device, dtype=torch.float32,
+            )
+
+    # --- (e) v_cmp (autograd-tracked: gradients to V flow through this) ---
+    v_cmp = _pool_v(v, chunk_size_k, kv_valid_mask)
+
+    # --- (f) merge (autograd-tracked) ---
+    use_triton = (
+        use_triton_inference
+        and not torch.is_grad_enabled()
+        and q.is_cuda
+    )
+    if use_triton:
+        try:
+            from .bias_rectification_triton import taylor_merge_triton
+            return taylor_merge_triton(
+                q, k_cmp, v_cmp, o_spa, lse_spa, nonsel_mask, is_flat, btc,
+                chunk_size_q, sm_scale,
+            )
+        except Exception as e:
+            # Fall through to PyTorch path on any kernel error.
+            print(f"[bias_rectification] taylor_merge_triton fallback: {e}")
+
+    return _taylor_merge_pytorch(
+        q, k_cmp, v_cmp, o_spa, lse_spa, nonsel_mask, is_flat, btc,
+        chunk_size_q, sm_scale,
+    )
+
+def _rectified_correction_pytorch(
+    o_spa: torch.Tensor,
+    R_n: torch.Tensor,
+    A_pool: torch.Tensor,
+    rectify_mask: torch.Tensor,
+    v_cmp: torch.Tensor,
+    chunk_size_q: int,
+) -> torch.Tensor:
+    """O' = R_n * O_spa + (A_pool * rectify_mask) @ V_cmp, autograd-tracked.
+
+    Memory-contiguous via expand+contiguous+reshape: each broadcast tensor is
+    materialised exactly once, never produces a non-contiguous view that the
+    next op would silently re-materialise.
+    """
+    B, H, S, D = o_spa.shape
+    N_q = R_n.shape[-1]
+
+    # Broadcast R_n [B, H, N_q] → [B, H, S, 1]: one materialisation, contiguous.
+    R_tok = (
+        R_n.unsqueeze(-1).unsqueeze(-1)
+        .expand(B, H, N_q, chunk_size_q, 1)
+        .contiguous()
+        .reshape(B, H, S, 1)
+    )
+    o_cri = R_tok * o_spa
+
+    # Non-critical contribution: weighted sum of pooled V over rectifiable blocks.
+    A_rect = A_pool * rectify_mask.to(A_pool.dtype)
+    o_ncri_block = torch.matmul(A_rect, v_cmp)              # [B, H, N_q, D]
+
+    # Broadcast block-level output to per-token, again contiguous.
+    o_ncri = (
+        o_ncri_block.unsqueeze(3)
+        .expand(B, H, N_q, chunk_size_q, D)
+        .contiguous()
+        .reshape(B, H, S, D)
+    )
+    return (o_cri + o_ncri).to(o_spa.dtype).contiguous()
+
+
+def bsa_rectified_sparse_attn(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    chunk_size_q: int,
+    chunk_size_k: int,
+    sparsity: Optional[float],
+    cdf_threshold: Optional[float],
+    sm_scale: float,
+    *,
+    kv_valid_mask: Optional[torch.Tensor] = None,
+    q_valid_mask: Optional[torch.Tensor] = None,
+    score_modifier: Optional[torch.Tensor] = None,
+    use_triton_inference: bool = True,
+) -> torch.Tensor:
+    # --- (a) pooling (autograd-tracked) ---
+    q_cmp, k_cmp = _pool_qk(q, k, chunk_size_q, chunk_size_k, kv_valid_mask, q_valid_mask)
+
+    # --- (b) raw scoring (autograd-tracked) ---
+    raw_score = cal_score(q_cmp, k_cmp)
+
+    # --- (c) selection (no grad) ---
+    with torch.no_grad():
+        sel_score = raw_score * score_modifier if score_modifier is not None else raw_score
+        block_indices, block_indices_lens = _select_blocks(sel_score, sparsity, cdf_threshold, sm_scale)
+
+    # --- (d) BSA Triton (autograd via _bsa_kernel_only) ---
+    o_spa, _ = _bsa_kernel(
+        q, k, v, sm_scale, block_indices, block_indices_lens,
+        chunk_size_q, chunk_size_k, sparsity, kv_valid_mask,
+    )
+
+    # --- (e) sel_mask / rectify_mask (no grad; bool tensors, contiguous) ---
+    block_valid = _compute_block_valid_mask(kv_valid_mask, chunk_size_k) if kv_valid_mask is not None else None
+    with torch.no_grad():
+        sel_mask = create_mask_from_indices_varlen(block_indices, raw_score.shape[-1])
+        rectify_mask = ~sel_mask                              # all non-selected blocks
+        if block_valid is not None:
+            rectify_mask = rectify_mask & block_valid.view(1, 1, 1, -1)
+        rectify_mask = rectify_mask.contiguous()
+
+    # --- (f) A_pool & R_n (autograd-tracked: ∂L/∂q,∂L/∂k flow through these) ---
+    A_pool = _masked_softmax_along_last(raw_score, sm_scale, block_valid)
+    R_n = (1.0 - (A_pool * rectify_mask.to(A_pool.dtype)).sum(dim=-1)).clamp(min=1e-6)
+
+    # --- (g) v_cmp (autograd-tracked: ∂L/∂v flows through pooling) ---
+    v_cmp = _pool_v(v, chunk_size_k, kv_valid_mask)
+
+    # --- (h) correction (fused Triton kernel for inference, PyTorch for training) ---
+    use_triton = (
+        use_triton_inference
+        and not torch.is_grad_enabled()
+        and q.is_cuda
+    )
+    if use_triton:
+        try:
+            from .bias_rectification_triton import rectified_correction_triton
+            return rectified_correction_triton(
+                o_spa, R_n.detach(), A_pool.detach(), rectify_mask, v_cmp.detach(),
+                chunk_size_q,
+            )
+        except Exception as e:
+            print(f"[bias_rectification] rectified_correction_triton fallback: {e}")
+
+    return _rectified_correction_pytorch(
+        o_spa, R_n, A_pool, rectify_mask, v_cmp, chunk_size_q,
+    )
+
+
+__all__ = [
+    "bsa_taylor_sparse_attn",
+    "bsa_rectified_sparse_attn",
+]

hymm/models/modules/block_sparse_attention/bias_rectification_triton.py

@@ -0,0 +1,363 @@
+import torch
+import triton
+import triton.language as tl
+
+
+# =====================================================================
+# Rectified-SpaAttn: fused (R_n * o_spa + (A_pool * rectify_mask) @ v_cmp)
+# =====================================================================
+@triton.jit
+def _rectified_correction_fwd_kernel(
+    O_spa, R_n, A_pool, rectify_mask, V_cmp, Out,
+    stride_oz, stride_oh, stride_om, stride_od,
+    stride_rz, stride_rh, stride_rn,                 # R_n: [B, H, N_q]
+    stride_az, stride_ah, stride_an, stride_ak,      # A_pool: [B, H, N_q, N_k]
+    stride_mz, stride_mh, stride_mn, stride_mk,      # rectify_mask: [B, H, N_q, N_k]
+    stride_vz, stride_vh, stride_vk, stride_vd,      # V_cmp: [B, H, N_k, D]
+    stride_cz, stride_ch, stride_cm, stride_cd,      # Out: [B, H, S_q, D]
+    H, N_q, N_k,
+    CHUNK_SIZE_Q: tl.constexpr,                       # tokens per Q block
+    BLOCK_M: tl.constexpr,                            # tokens per program (<= CHUNK_SIZE_Q)
+    BLOCK_N: tl.constexpr,                            # K blocks per inner step
+    HEAD_DIM: tl.constexpr,
+):
+    """Compute Out[b, h, s, :] = R_n[b, h, n_q(s)] * O_spa[b, h, s, :]
+                                + sum_{j} A_pool[b, h, n_q(s), j]
+                                         * rectify_mask[b, h, n_q(s), j]
+                                         * V_cmp[b, h, j, :]
+    where n_q(s) = s // CHUNK_SIZE_Q.
+
+    Each program handles BLOCK_M consecutive query tokens (within one Q block,
+    BLOCK_M <= CHUNK_SIZE_Q so every token in the program shares the same n_q).
+    """
+    pid_m = tl.program_id(0)
+    pid_zh = tl.program_id(1)
+    off_z = pid_zh // H
+    off_h = pid_zh % H
+
+    # All BLOCK_M rows belong to the same Q block (BLOCK_M divides CHUNK_SIZE_Q).
+    n_q = (pid_m * BLOCK_M) // CHUNK_SIZE_Q
+
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_d = tl.arange(0, HEAD_DIM)
+
+    # -------- Load R_n and the BSA output --------
+    r_ptr = R_n + off_z * stride_rz + off_h * stride_rh + n_q * stride_rn
+    r = tl.load(r_ptr).to(tl.float32)
+
+    o_spa_ptrs = (
+        O_spa
+        + off_z * stride_oz + off_h * stride_oh
+        + offs_m[:, None] * stride_om + offs_d[None, :] * stride_od
+    )
+    out_ptrs = (
+        Out
+        + off_z * stride_cz + off_h * stride_ch
+        + offs_m[:, None] * stride_cm + offs_d[None, :] * stride_cd
+    )
+
+    s_q_total = pid_m * BLOCK_M + BLOCK_M  # valid rows are those < S_q_total in caller
+    # We don't know S_q here; rely on caller to guarantee S_q is a multiple of CHUNK_SIZE_Q
+    # via padding. But we still do a row mask to be safe for the last program.
+    row_mask = offs_m[:, None] < (n_q + 1) * CHUNK_SIZE_Q
+
+    o_spa = tl.load(o_spa_ptrs, mask=row_mask, other=0.0).to(tl.float32)
+    acc = r * o_spa  # broadcast scalar * [BLOCK_M, HEAD_DIM]
+
+    # -------- Accumulate non-critical contribution: sum_j A * mask * V --------
+    a_base = (
+        A_pool
+        + off_z * stride_az + off_h * stride_ah + n_q * stride_an
+    )
+    m_base = (
+        rectify_mask
+        + off_z * stride_mz + off_h * stride_mh + n_q * stride_mn
+    )
+    v_base = V_cmp + off_z * stride_vz + off_h * stride_vh
+
+    for j_start in range(0, N_k, BLOCK_N):
+        offs_n = j_start + tl.arange(0, BLOCK_N)
+        n_mask = offs_n < N_k
+
+        a = tl.load(a_base + offs_n * stride_ak, mask=n_mask, other=0.0).to(tl.float32)
+        m = tl.load(m_base + offs_n * stride_mk, mask=n_mask, other=0).to(tl.float32)
+        coef = a * m  # [BLOCK_N], zero where mask is False
+
+        v_ptrs = v_base + offs_n[:, None] * stride_vk + offs_d[None, :] * stride_vd
+        v = tl.load(v_ptrs, mask=n_mask[:, None], other=0.0).to(tl.float32)
+
+        # acc += coef[:, None] * v summed over BLOCK_N (one Q row uses all K blocks)
+        # Contribution is the same for all Q tokens in this block, so we add a row-wise
+        # contribution: acc[i, :] += sum_j coef[j] * v[j, :]
+        contrib = tl.sum(coef[:, None] * v, axis=0)  # [HEAD_DIM]
+        acc += contrib[None, :]
+
+    tl.store(out_ptrs, acc.to(Out.type.element_ty), mask=row_mask)
+
+
+def rectified_correction_triton(
+    o_spa: torch.Tensor,            # [B, H, S_q, D]
+    R_n: torch.Tensor,              # [B, H, N_q]
+    A_pool: torch.Tensor,           # [B, H, N_q, N_k] (probabilities, fp32)
+    rectify_mask: torch.Tensor,     # [B, H, N_q, N_k] (bool)
+    v_cmp: torch.Tensor,            # [B, H, N_k, D]
+    chunk_size_q: int,
+) -> torch.Tensor:
+    """Fused Rectified-SpaAttn correction without materialising the
+    [B, H, S_q, D] o_ncri intermediate.
+
+    Output: out = R_n_token_level * o_spa + (A_pool * rectify_mask) @ v_cmp,
+    where R_n is broadcast from N_q to S_q via the chunk index.
+    """
+    B, H, S_q, D = o_spa.shape
+    N_q = R_n.shape[-1]
+    N_k = A_pool.shape[-1]
+    assert S_q == N_q * chunk_size_q, f"S_q={S_q} must equal N_q*chunk_size_q={N_q*chunk_size_q}"
+
+    # Choose BLOCK_M = chunk_size_q for chunk_size_q <= 64; else use 64.
+    if chunk_size_q <= 64:
+        BLOCK_M = chunk_size_q
+    else:
+        # When chunk_size_q is large, split each Q block into multiple programs.
+        BLOCK_M = 64
+        assert chunk_size_q % BLOCK_M == 0, (
+            f"chunk_size_q={chunk_size_q} must be divisible by BLOCK_M={BLOCK_M}")
+
+    BLOCK_N = 64 if N_k >= 64 else max(1, triton.next_power_of_2(N_k))
+
+    o_spa = o_spa.contiguous()
+    A_pool = A_pool.contiguous()
+    rectify_mask = rectify_mask.contiguous()
+    v_cmp = v_cmp.contiguous()
+    R_n = R_n.contiguous()
+    out = torch.empty_like(o_spa)
+
+    grid = (triton.cdiv(S_q, BLOCK_M), B * H, 1)
+
+    _rectified_correction_fwd_kernel[grid](
+        o_spa, R_n, A_pool, rectify_mask, v_cmp, out,
+        o_spa.stride(0), o_spa.stride(1), o_spa.stride(2), o_spa.stride(3),
+        R_n.stride(0), R_n.stride(1), R_n.stride(2),
+        A_pool.stride(0), A_pool.stride(1), A_pool.stride(2), A_pool.stride(3),
+        rectify_mask.stride(0), rectify_mask.stride(1), rectify_mask.stride(2), rectify_mask.stride(3),
+        v_cmp.stride(0), v_cmp.stride(1), v_cmp.stride(2), v_cmp.stride(3),
+        out.stride(0), out.stride(1), out.stride(2), out.stride(3),
+        H, N_q, N_k,
+        CHUNK_SIZE_Q=chunk_size_q,
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        HEAD_DIM=D,
+        num_warps=4,
+        num_stages=2,
+    )
+    return out
+
+
+@triton.jit
+def _taylor_merge_fwd_kernel(
+    Q, K_cmp, V_cmp,
+    O_spa, LSE_spa,
+    Out,
+    block_token_counts,        # [N_k]   (fp32, log already taken outside)
+    nonsel_mask,               # [B, H, N_q, N_k]   bool
+    is_flat,                   # [B, H, N_q]        bool
+    sm_scale_log2,             # scalar, in log2
+    sm_scale,                  # scalar, natural
+    stride_qz, stride_qh, stride_qm, stride_qd,
+    stride_kz, stride_kh, stride_kk, stride_kd,
+    stride_vz, stride_vh, stride_vk, stride_vd,
+    stride_oz, stride_oh, stride_om, stride_od,
+    stride_lz, stride_lh, stride_lm,                 # LSE_spa [B, H, S_q]
+    stride_cz, stride_ch, stride_cm, stride_cd,      # Out
+    stride_btc,                                       # block_token_counts [N_k]
+    stride_mz, stride_mh, stride_mn, stride_mk,
+    stride_fz, stride_fh, stride_fn,
+    H, N_q, N_k,
+    CHUNK_SIZE_Q: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+):
+    """For each query token i in a Q block n_q:
+       z_taylor[j] = q_i · K_cmp[j] * sm_scale + log_count[j]   (only for j in nonsel)
+       lse_taylor = logsumexp_j z_taylor                        (over nonsel j)
+       o_taylor   = sum_j softmax(z_taylor)[j] * V_cmp[j]
+    Then if is_flat[n_q]:
+       w_spa = exp(lse_spa - lse_total)
+       w_taylor = exp(lse_taylor - lse_total)
+       out = w_spa * o_spa + w_taylor * o_taylor
+    else:
+       out = o_spa
+    """
+    pid_m = tl.program_id(0)
+    pid_zh = tl.program_id(1)
+    off_z = pid_zh // H
+    off_h = pid_zh % H
+
+    n_q = (pid_m * BLOCK_M) // CHUNK_SIZE_Q
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_d = tl.arange(0, HEAD_DIM)
+
+    flat_ptr = is_flat + off_z * stride_fz + off_h * stride_fh + n_q * stride_fn
+    flat_val = tl.load(flat_ptr).to(tl.int32)
+
+    # Load Q rows for this program
+    q_ptrs = (
+        Q + off_z * stride_qz + off_h * stride_qh
+        + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd
+    )
+    q = tl.load(q_ptrs).to(tl.float32)
+
+    # Always read o_spa
+    o_spa_ptrs = (
+        O_spa + off_z * stride_oz + off_h * stride_oh
+        + offs_m[:, None] * stride_om + offs_d[None, :] * stride_od
+    )
+    o_spa = tl.load(o_spa_ptrs).to(tl.float32)
+
+    out_ptrs = (
+        Out + off_z * stride_cz + off_h * stride_ch
+        + offs_m[:, None] * stride_cm + offs_d[None, :] * stride_cd
+    )
+
+    if flat_val == 0:
+        # Sharp queries: keep o_spa unchanged.
+        tl.store(out_ptrs, o_spa.to(Out.type.element_ty))
+        return
+
+    # --- Streaming softmax over nonsel K blocks for Taylor ---
+    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
+    acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
+
+    mask_base = nonsel_mask + off_z * stride_mz + off_h * stride_mh + n_q * stride_mn
+    k_base = K_cmp + off_z * stride_kz + off_h * stride_kh
+    v_base = V_cmp + off_z * stride_vz + off_h * stride_vh
+
+    for j_start in range(0, N_k, BLOCK_N):
+        offs_n = j_start + tl.arange(0, BLOCK_N)
+        n_mask = offs_n < N_k
+
+        # Load K_cmp [BLOCK_N, HEAD_DIM]
+        k_ptrs = k_base + offs_n[:, None] * stride_kk + offs_d[None, :] * stride_kd
+        k_cmp_b = tl.load(k_ptrs, mask=n_mask[:, None], other=0.0).to(tl.float32)
+
+        # Score q @ k_cmp^T : [BLOCK_M, BLOCK_N]
+        qk = tl.dot(q, tl.trans(k_cmp_b)) * sm_scale
+
+        # Add log(block_token_count[j]) to make Taylor weights commensurate.
+        log_cnt = tl.load(block_token_counts + offs_n * stride_btc, mask=n_mask, other=0.0)
+        qk = qk + log_cnt[None, :]
+
+        # Mask out non-eligible blocks (selected, padded, etc.)
+        nm = tl.load(mask_base + offs_n * stride_mk, mask=n_mask, other=0).to(tl.int32)
+        valid = (nm != 0) & n_mask
+        qk = tl.where(valid[None, :], qk, float("-inf"))
+
+        # Streaming online softmax update
+        m_ij = tl.maximum(m_i, tl.max(qk, axis=1))
+        alpha = tl.math.exp2((m_i - m_ij) * 1.4426950408889634)
+        # exp in natural log: convert qk (already natural) to natural exp
+        # We want p = exp(qk - m_ij). Use tl.math.exp.
+        p = tl.math.exp(qk - m_ij[:, None])
+        p = tl.where(valid[None, :], p, 0.0)
+
+        # Load V_cmp [BLOCK_N, HEAD_DIM]
+        v_ptrs = v_base + offs_n[:, None] * stride_vk + offs_d[None, :] * stride_vd
+        v_cmp_b = tl.load(v_ptrs, mask=n_mask[:, None], other=0.0).to(tl.float32)
+
+        l_i = l_i * alpha + tl.sum(p, axis=1)
+        acc = acc * alpha[:, None] + tl.dot(p, v_cmp_b)
+        m_i = m_ij
+
+    # lse_taylor = m_i + log(l_i); guard l_i = 0 (no nonsel blocks → o_taylor = 0, lse = -inf)
+    l_safe = tl.where(l_i > 0.0, l_i, 1.0)
+    lse_taylor = tl.where(l_i > 0.0, m_i + tl.log(l_safe), tl.full([BLOCK_M], -float("inf"), tl.float32))
+    o_taylor = tl.where(l_i[:, None] > 0.0, acc / l_safe[:, None], tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32))
+
+    # Load lse_spa for these tokens
+    lse_spa_ptrs = LSE_spa + off_z * stride_lz + off_h * stride_lh + offs_m * stride_lm
+    lse_spa = tl.load(lse_spa_ptrs)
+
+    # Stable lse merge: lse_total = max + log(exp(lse_a - max) + exp(lse_b - max))
+    lse_max = tl.maximum(lse_spa, lse_taylor)
+    finite = (lse_max > -float("inf"))
+    a_ext = tl.where(finite, lse_spa - lse_max, 0.0)
+    b_ext = tl.where(finite, lse_taylor - lse_max, 0.0)
+    lse_total = tl.where(finite, lse_max + tl.log(tl.math.exp(a_ext) + tl.math.exp(b_ext)), lse_max)
+
+    w_spa = tl.where(finite, tl.math.exp(lse_spa - lse_total), 1.0)
+    w_taylor = tl.where(finite, tl.math.exp(lse_taylor - lse_total), 0.0)
+
+    out = w_spa[:, None] * o_spa + w_taylor[:, None] * o_taylor
+    tl.store(out_ptrs, out.to(Out.type.element_ty))
+
+
+def taylor_merge_triton(
+    q: torch.Tensor,                # [B, H, S_q, D]   float (bf16/fp16/fp32)
+    k_cmp: torch.Tensor,            # [B, H, N_k, D]
+    v_cmp: torch.Tensor,            # [B, H, N_k, D]
+    o_spa: torch.Tensor,            # [B, H, S_q, D]
+    lse_spa: torch.Tensor,          # [B, H, S_q]   natural-log LSE from BSA
+    nonsel_mask: torch.Tensor,      # [B, H, N_q, N_k] bool
+    is_flat: torch.Tensor,          # [B, H, N_q]      bool
+    block_token_counts: torch.Tensor,  # [N_k] fp32 — number of valid tokens per K block
+    chunk_size_q: int,
+    sm_scale: float,
+) -> torch.Tensor:
+    """Fused Taylor merge.
+
+    Materialises NONE of the [B, H, S_q, N_k] taylor_scores tensor.
+    Sharp queries (is_flat=False) skip the streaming softmax entirely.
+    """
+    B, H, S_q, D = o_spa.shape
+    N_q = nonsel_mask.shape[-2]
+    N_k = nonsel_mask.shape[-1]
+    assert S_q == N_q * chunk_size_q
+
+    if chunk_size_q <= 64:
+        BLOCK_M = chunk_size_q
+    else:
+        BLOCK_M = 64
+        assert chunk_size_q % BLOCK_M == 0
+
+    BLOCK_N = 64 if N_k >= 64 else max(1, triton.next_power_of_2(N_k))
+
+    # Take log of token counts once; pad blocks (count=0) → log(1) = 0, never reached
+    # because they're already excluded via nonsel_mask & block_valid in the caller.
+    log_counts = torch.log(block_token_counts.float().clamp(min=1.0)).contiguous()
+
+    q = q.contiguous()
+    k_cmp = k_cmp.contiguous()
+    v_cmp = v_cmp.contiguous()
+    o_spa = o_spa.contiguous()
+    lse_spa = lse_spa.contiguous()
+    nonsel_mask = nonsel_mask.contiguous()
+    is_flat = is_flat.contiguous()
+    out = torch.empty_like(o_spa)
+
+    grid = (triton.cdiv(S_q, BLOCK_M), B * H, 1)
+
+    _taylor_merge_fwd_kernel[grid](
+        q, k_cmp, v_cmp,
+        o_spa, lse_spa, out,
+        log_counts, nonsel_mask, is_flat,
+        sm_scale * 1.4426950408889634, sm_scale,
+        q.stride(0), q.stride(1), q.stride(2), q.stride(3),
+        k_cmp.stride(0), k_cmp.stride(1), k_cmp.stride(2), k_cmp.stride(3),
+        v_cmp.stride(0), v_cmp.stride(1), v_cmp.stride(2), v_cmp.stride(3),
+        o_spa.stride(0), o_spa.stride(1), o_spa.stride(2), o_spa.stride(3),
+        lse_spa.stride(0), lse_spa.stride(1), lse_spa.stride(2),
+        out.stride(0), out.stride(1), out.stride(2), out.stride(3),
+        log_counts.stride(0),
+        nonsel_mask.stride(0), nonsel_mask.stride(1), nonsel_mask.stride(2), nonsel_mask.stride(3),
+        is_flat.stride(0), is_flat.stride(1), is_flat.stride(2),
+        H, N_q, N_k,
+        CHUNK_SIZE_Q=chunk_size_q,
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        HEAD_DIM=D,
+        num_warps=4,
+        num_stages=2,
+    )
+    return out

hymm/models/modules/block_sparse_attention/bsa_interface.py

@@ -0,0 +1,1503 @@
+import os
+import torch
+import torch.nn.functional as F
+import triton
+import triton.language as tl
+import math
+
+from .common import _attn_fwd_gating, _attn_bwd_preprocess, configs_gating_preset
+from .flash_attn_bsa_varlen_mask import (
+    _attn_fwd_bsa_varlen, _attn_fwd_bsa_varlen_align, _attn_bwd_dkdv_bsa_varlen_wrapper, _attn_bwd_dq_bsa_varlen_wrapper, _attn_bwd_dq_bsa_varlen_align_wrapper,
+    configs_fwd_bsa_varlen_preset, configs_fwd_bsa_varlen_align_preset, configs_bwd_dkdv_bsa_varlen_preset, configs_bwd_dq_bsa_varlen_preset, configs_bwd_dq_bsa_varlen_align_preset
+)
+
+torch._dynamo.config.cache_size_limit = 32
+
+def is_cuda():
+    return triton.runtime.driver.active.get_current_target().backend == "cuda"
+
+def supports_tma():
+    return is_cuda() and torch.cuda.get_device_capability()[0] >= 9
+
+HAS_TMA_DESC = "nv_tma_desc_type" in dir(tl)
+
+if HAS_TMA_DESC:
+    print("TMA benchmarks will be running with experimental grid constant TMA descriptor.", )
+else:
+    print("TMA benchmarks will be running without grid constant TMA descriptor.", )
+
+
+# TmaAutoTuneHelper used in htyu's PR #5622
+class TmaAutoTuneHelper:
+
+    # duck typing wrapper to implement the same interface as TmaDescKernelParam in Triton PR #4498
+    class KernelParamWrapper:
+
+        def __init__(self, desc):
+            self.desc = desc
+
+        def tma_desc_cpu_ptr(self):
+            return self.desc.data_ptr()
+
+    TMA_SIZE = 128
+
+    def __init__(self):
+        self.fill_1d_tma_descriptor_inner = (triton.runtime.driver.active.utils.fill_1d_tma_descriptor)
+        self.fill_2d_tma_descriptor_inner = (triton.runtime.driver.active.utils.fill_2d_tma_descriptor)
+        if HAS_TMA_DESC:
+            self.descriptors = {}
+        else:
+            self.cuda_descriptors = {}
+
+    # Call this method outside of the lambda function for grid size
+    def init_tma_descriptor(self, name):
+        if HAS_TMA_DESC:
+            self.descriptors[name] = torch.empty(TmaAutoTuneHelper.TMA_SIZE, device="cpu", dtype=torch.int8)
+        else:
+            self.cuda_descriptors[name] = torch.empty(TmaAutoTuneHelper.TMA_SIZE, device="cuda", dtype=torch.int8)
+
+    # Call this method inside the lambda function for grid size
+    def fill_1d_tma_descriptor(self, name, ptr, dim, block_dim, element_size):
+        if HAS_TMA_DESC:
+            desc_x = self.descriptors[name]
+            assert desc_x.data_ptr() % 64 == 0
+            self.fill_1d_tma_descriptor_inner(ptr, dim, block_dim, element_size, desc_x.data_ptr())
+        else:
+            desc_x = self.cuda_descriptors[name]
+            buf_x = torch.empty_like(desc_x, device="cpu", pin_memory=True)
+            self.fill_1d_tma_descriptor_inner(ptr, dim, block_dim, element_size, buf_x.data_ptr())
+            desc_x.copy_(buf_x, non_blocking=True)
+
+    # Call this method inside the lambda function for grid size
+    def fill_2d_tma_descriptor(self, name, ptr, dim1, dim0, block_dim1, block_dim0, element_size):
+        if HAS_TMA_DESC:
+            desc_x = self.descriptors[name]
+            assert desc_x.data_ptr() % 64 == 0
+            self.fill_2d_tma_descriptor_inner(ptr, dim1, dim0, block_dim1, block_dim0, element_size, desc_x.data_ptr())
+        else:
+            desc_x = self.cuda_descriptors[name]
+            buf_x = torch.empty_like(desc_x, device="cpu", pin_memory=True)
+            self.fill_2d_tma_descriptor_inner(ptr, dim1, dim0, block_dim1, block_dim0, element_size, buf_x.data_ptr())
+            desc_x.copy_(buf_x, non_blocking=True)
+
+    def get_tma_descriptor_kernel_param(self, name):
+        if HAS_TMA_DESC:
+            assert self.descriptors[name] is not None
+            return self.KernelParamWrapper(self.descriptors[name])
+        else:
+            assert self.cuda_descriptors[name] is not None
+            return self.cuda_descriptors[name]
+
+
+@triton.jit
+def create_mask_from_indices_kernel(
+    block_indices,
+    block_mask,
+    stride_bz, stride_bh, stride_bm, stride_bs,
+    stride_mz, stride_mh, stride_mm, stride_mn,
+    H,
+):
+    i_zh, i_m, i_s = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_z, i_h = i_zh // H, i_zh % H
+
+    off_b = i_z.to(tl.int64) * stride_bz + i_h.to(tl.int64) * stride_bh + i_m.to(tl.int64) * stride_bm + i_s.to(tl.int64) * stride_bs
+    
+    b_i = tl.load(block_indices + off_b)
+    
+    off_m = i_z.to(tl.int64) * stride_mz + i_h.to(tl.int64) * stride_mh + i_m.to(tl.int64) * stride_mm + b_i.to(tl.int64) * stride_mn
+    
+    b_m = 1
+    tl.store(block_mask + off_m, b_m.to(block_mask.dtype.element_ty))
+
+def create_mask_from_indices_triton(
+    block_indices,
+    N_cols
+):
+    B, H, N_rows, S = block_indices.shape
+    block_mask = torch.zeros((B, H, N_rows, N_cols), dtype=torch.bool, device=block_indices.device)
+    create_mask_from_indices_kernel[(B * H, N_rows, S)](
+        block_indices,
+        block_mask,
+        block_indices.stride(0), block_indices.stride(1), block_indices.stride(2), block_indices.stride(3),
+        block_mask.stride(0), block_mask.stride(1), block_mask.stride(2), block_mask.stride(3),
+        H,
+    )
+    return block_mask
+
+@torch.compile
+def create_mask_from_indices_varlen(block_indices, N_cols_mask):
+   
+    B, H, M, _ = block_indices.shape
+    device = block_indices.device
+    
+    mask = torch.zeros((B, H, M, N_cols_mask), dtype=torch.bool, device=device)
+    
+    valid = block_indices < N_cols_mask
+    
+    b_idx = torch.arange(B, device=device)[:, None, None, None].expand_as(block_indices)
+    h_idx = torch.arange(H, device=device)[None, :, None, None].expand_as(block_indices)
+    m_idx = torch.arange(M, device=device)[None, None, :, None].expand_as(block_indices)
+    
+    valid_coords = (b_idx[valid], h_idx[valid], m_idx[valid], block_indices[valid])
+    
+    mask[valid_coords] = True
+    
+    return mask
+
+@torch.compile
+def create_indices_k_from_indices_q_varlen(
+    block_indices,
+    N_cols_mask # indicate the number of the last dimension of the bool mask, since this information cannot be determined by block_indices, which may contain invalid elements
+):
+    block_mask_qk = create_mask_from_indices_varlen(block_indices, N_cols_mask)
+    B, H, M, N = block_mask_qk.shape
+    block_mask_kq = block_mask_qk.permute(0, 1, 3, 2)
+    indices = torch.arange(M, device=block_indices.device).view(1, 1, 1, -1).expand_as(block_mask_kq)
+    block_indices_k = torch.where(block_mask_kq, indices, M)
+    block_indices_k, _ = torch.sort(block_indices_k, dim=-1)
+        
+    block_indices_k_lens = (block_indices_k < M).sum(dim=-1)
+
+    return block_indices_k, block_indices_k_lens
+
+
+@torch.compile
+def mean_pooling_compression(
+    x: torch.Tensor,
+    block_size: int
+) -> torch.Tensor:
+    B, H, S = x.shape[:3]
+    num_block = math.ceil(S / block_size)
+    if S % block_size != 0:
+        x = F.pad(x, (0, 0, 0, num_block * block_size - S))
+    x_cmp = x.view(B, H, num_block, block_size, -1).mean(dim=3)
+    return x_cmp
+
+
+@torch.compile
+def masked_mean_pooling_compression(
+    x: torch.Tensor,
+    block_size: int,
+    valid_mask: torch.Tensor,
+) -> torch.Tensor:
+    """Mean pooling that ignores padded positions marked False in valid_mask.
+
+    Handles both cases:
+      - S divisible by block_size (after 3D padding + block rearrangement)
+      - S not divisible by block_size (pads the tail and extends valid_mask)
+    """
+    B, H, S, D = x.shape
+    num_block = math.ceil(S / block_size)
+    if S % block_size != 0:
+        pad_len = num_block * block_size - S
+        x = F.pad(x, (0, 0, 0, pad_len))
+        valid_mask = F.pad(valid_mask, (0, pad_len), value=False)
+
+    x_blocks = x.view(B, H, num_block, block_size, D)
+    mask_blocks = valid_mask.view(num_block, block_size)
+    mask_f = mask_blocks.float().unsqueeze(0).unsqueeze(0).unsqueeze(-1)
+    counts = mask_f.sum(dim=3, keepdim=True).clamp(min=1.0)
+    x_cmp = (x_blocks * mask_f).sum(dim=3) / counts.squeeze(3)
+    return x_cmp
+
+@torch.compile
+def cal_score(q, k):
+    k_transposed = k.transpose(-1, -2)  # [b, h, d, s_k]
+    score = torch.matmul(q, k_transposed)  # [b, h, s_q, s_k]
+    return score
+
+def cal_score_triton(q, k):
+    B, H, s_q, D = q.shape
+    s_k = k.shape[2]
+    
+    score = torch.empty(B, H, s_q, s_k, device=q.device, dtype=q.dtype)
+    
+    kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', '0') == '1' else configs_gating_preset['default']
+    
+    grid = lambda args: (triton.cdiv(s_q, args["BLOCK_M"]), B * H, 1)
+    _attn_fwd_gating[grid](
+        q, k, score,
+        q.stride(0), q.stride(1), q.stride(2), q.stride(3),
+        k.stride(0), k.stride(1), k.stride(2), k.stride(3),
+        score.stride(0), score.stride(1), score.stride(2), score.stride(3),
+        H, s_q, s_k,
+        HEAD_DIM=D,
+        **kernel_config
+    )
+    return score
+
+@torch.compile
+def get_select_indices_topk(q, k, sparsity):
+    score = cal_score(q, k)
+    block_indices, block_indices_lens = get_select_indices_topk_from_score(score, sparsity)
+    return block_indices, block_indices_lens
+
+@torch.compile
+def get_select_indices_topk_from_score(score, sparsity):
+    num_selected = int((1 - sparsity) * score.shape[-1])
+    block_indices = torch.topk(score, num_selected)[1]
+    block_indices, _ = torch.sort(block_indices, dim=-1)
+
+    block_indices_lens = torch.full(
+        (block_indices.shape[0], block_indices.shape[1], block_indices.shape[2]), 
+        num_selected, 
+        dtype=torch.int32,
+        device=block_indices.device
+    )
+
+    return block_indices, block_indices_lens
+
+@torch.compile
+def get_select_indices_cdf(q, k, cdf_threshold):
+    score = cal_score(q, k)
+    head_dim = q.shape[-1]
+    block_indices, block_indices_lens = get_select_indices_cdf_from_score(score, cdf_threshold, 1 / head_dim**0.5)
+    return block_indices, block_indices_lens
+
+@torch.compile
+def get_select_indices_cdf_from_score(score, cdf_threshold, sm_scale):
+    weights = torch.softmax(score * sm_scale, dim=-1)
+    
+    B, H, Sq, Sk = weights.shape
+    upper_bound = min(Sk, int(cdf_threshold * Sk) + 1)
+    topk_vals, topk_idx = torch.topk(weights, k=upper_bound, dim=-1, largest=True, sorted=True)
+    cdf = torch.cumsum(topk_vals, dim=-1)
+    # Standard nucleus (top-p): pick the smallest set whose cumulative prob >= threshold,
+    # i.e. INCLUDE the block that crosses the threshold. (cdf < p).sum() counts the prefix
+    # blocks still strictly below p; +1 adds the crossing block.
+    num_selected = (cdf < cdf_threshold).to(torch.int32).sum(dim=-1, keepdim=True) + 1
+    num_selected = num_selected.clamp(min=1, max=Sk)
+    
+    block_indices = topk_idx.contiguous()
+    pos = torch.arange(upper_bound, device=block_indices.device)
+    block_indices = block_indices.masked_fill(pos >= num_selected, Sk)
+    block_indices, _ = torch.sort(block_indices, dim=-1)
+    return block_indices, num_selected.squeeze(-1).to(torch.int32)
+
+@torch.compile
+def get_select_indices_cdf_topk(q, k, sparsity, cdf_threshold):
+    score = cal_score(q, k)
+    head_dim = q.shape[-1]
+    block_indices, block_indices_lens = get_select_indices_cdf_topk_from_score(score, sparsity, cdf_threshold, 1 / head_dim**0.5)
+    return block_indices, block_indices_lens
+
+@torch.compile
+def get_select_indices_cdf_topk_from_score(score, sparsity, cdf_threshold, sm_scale):
+    weights = torch.softmax(score * sm_scale, dim=-1)
+    
+    B, H, Sq, Sk = weights.shape
+    num_selected_topk = max(1, int((1 - sparsity) * Sk))
+    upper_bound = min(Sk, max(num_selected_topk, int(cdf_threshold * Sk) + 1))
+
+    topk_vals, topk_idx = torch.topk(weights, k=upper_bound, dim=-1, largest=True, sorted=True)
+    cdf = torch.cumsum(topk_vals, dim=-1)
+    # Standard nucleus (top-p): pick the smallest set whose cumulative prob >= threshold,
+    # i.e. INCLUDE the block that crosses the threshold. (cdf < p).sum() counts the prefix
+    # blocks still strictly below p; +1 adds the crossing block.
+    num_selected = (cdf < cdf_threshold).to(torch.int32).sum(dim=-1, keepdim=True) + 1
+    num_selected = num_selected.clamp(min=num_selected_topk, max=Sk)
+
+    block_indices = topk_idx.contiguous()
+    pos = torch.arange(upper_bound, device=block_indices.device)
+    block_indices = block_indices.masked_fill(pos >= num_selected, Sk)
+    block_indices, _ = torch.sort(block_indices, dim=-1)
+    return block_indices, num_selected.squeeze(-1).to(torch.int32)
+
+def get_select_indices(q, k, sparsity, cdf_threshold):
+    if sparsity is not None and cdf_threshold is None:
+        block_indices, block_indices_lens = get_select_indices_topk(q, k, sparsity)
+    elif sparsity is None and cdf_threshold is not None:
+        block_indices, block_indices_lens = get_select_indices_cdf(q, k, cdf_threshold)
+    elif sparsity is not None and cdf_threshold is not None:
+        block_indices, block_indices_lens = get_select_indices_cdf_topk(q, k, sparsity, cdf_threshold)
+    else:
+        raise ValueError
+    return block_indices, block_indices_lens
+
+def get_select_indices_from_score(score, sparsity, cdf_threshold):
+    if sparsity is not None and cdf_threshold is None:
+        block_indices, block_indices_lens = get_select_indices_topk_from_score(score, sparsity)
+    elif sparsity is None and cdf_threshold is not None:
+        block_indices, block_indices_lens = get_select_indices_cdf_from_score(score, cdf_threshold)
+    elif sparsity is not None and cdf_threshold is not None:
+        block_indices, block_indices_lens = get_select_indices_cdf_topk_from_score(score, sparsity, cdf_threshold)
+    else:
+        raise ValueError
+    return block_indices, block_indices_lens
+
+def attn_fwd_bsa_varlen_triton(
+    q, 
+    k, 
+    v, 
+    sm_scale, 
+    block_indices, 
+    block_indices_lens, 
+    chunk_size_q, 
+    chunk_size_k,
+    sparsity,
+    kv_valid_mask=None,
+):
+    
+    B, H, Seq, D = q.shape
+
+    o = torch.empty_like(q)
+    M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32)
+    
+    grid = lambda args: (triton.cdiv(q.shape[2], args["BLOCK_M"]), q.shape[0] * q.shape[1], 1)
+
+    config_key = 'BLOCK_N_LG=64' if chunk_size_k == 64 else 'default'
+    if chunk_size_k > 128:
+        fwd_func = _attn_fwd_bsa_varlen
+        kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', '0') == '1' else configs_fwd_bsa_varlen_preset[config_key]
+    else:
+        fwd_func = _attn_fwd_bsa_varlen_align
+        kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', '0') == '1' else configs_fwd_bsa_varlen_align_preset[config_key]
+    
+    block_indices = block_indices.contiguous()
+    block_indices_lens = block_indices_lens.contiguous()
+
+    has_kv_mask = 1 if kv_valid_mask is not None else 0
+    _kv_mask = kv_valid_mask if kv_valid_mask is not None else torch.empty(0, dtype=torch.bool, device=q.device)
+    
+    fwd_func[grid](
+        q, k, v, sm_scale, M, o, 
+        block_indices,
+        block_indices_lens,
+        _kv_mask,
+        q.stride(0), q.stride(1), q.stride(2), q.stride(3),
+        k.stride(0), k.stride(1), k.stride(2), k.stride(3),
+        v.stride(0), v.stride(1), v.stride(2), v.stride(3), 
+        o.stride(0), o.stride(1), o.stride(2), o.stride(3),
+        block_indices.stride(0), block_indices.stride(1), block_indices.stride(2), block_indices.stride(3),
+        block_indices_lens.stride(0), block_indices_lens.stride(1), block_indices_lens.stride(2),
+        H, Seq, 
+        D,
+        BLOCK_M=chunk_size_q,
+        BLOCK_N_LG=chunk_size_k,
+        SPARSITY=sparsity,
+        HAS_KV_MASK=has_kv_mask,
+        **kernel_config
+    )
+    
+    LN2 = 0.6931471824645996
+    lse = M * LN2
+
+    return o, lse
+
+def attn_bwd_bsa_varlen_triton(
+    do, 
+    q, 
+    k, 
+    v, 
+    o, 
+    dq,
+    dk,
+    dv,
+    sm_scale, 
+    M, 
+    block_indices, 
+    block_indices_lens,
+    chunk_size_q, 
+    chunk_size_k,
+    sparsity,
+    kv_valid_mask=None,
+):
+    RCP_LN2 = 1.4426950408889634
+    M = M * RCP_LN2 # ln -> log2
+    
+    do = do.contiguous()
+    # assert q.stride() == k.stride() == v.stride() == o.stride() == do.stride()
+
+    BATCH, N_HEAD, N_CTX, HEAD_DIM = q.shape
+    N_CTX_KV = k.shape[-2]
+
+    RCP_LN2 = 1.4426950408889634  # = 1.0 / ln(2) # reciprocal 
+    arg_k = k
+    arg_k = arg_k * (sm_scale * RCP_LN2)
+    
+    if min(chunk_size_q, chunk_size_k) >= 128:
+        PRE_BLOCK = 128
+    else:
+        PRE_BLOCK = min(chunk_size_q, chunk_size_k)
+        
+    assert N_CTX % PRE_BLOCK == 0
+    pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD)
+    delta = torch.empty_like(M)
+    _attn_bwd_preprocess[pre_grid](
+        o, do,
+        delta,
+        N_CTX,
+        BLOCK_M=PRE_BLOCK, 
+        HEAD_DIM=HEAD_DIM
+    )
+
+    _N_cols = N_CTX_KV // chunk_size_k
+    _Hn = block_indices.shape[1]
+    _ki_list, _kl_list = [], []
+    for _h in range(_Hn):
+        _ki_h, _kl_h = create_indices_k_from_indices_q_varlen(
+            block_indices=block_indices[:, _h:_h+1],
+            N_cols_mask=_N_cols,
+        )
+        _ki_list.append(_ki_h)
+        _kl_list.append(_kl_h)
+        del _ki_h, _kl_h
+    block_indices_k = torch.cat(_ki_list, dim=1)
+    block_indices_k_lens = torch.cat(_kl_list, dim=1)
+    del _ki_list, _kl_list
+    
+    block_indices = block_indices.contiguous()
+    block_indices_lens = block_indices_lens.contiguous()
+    block_indices_k = block_indices_k.contiguous()
+    block_indices_k_lens = block_indices_k_lens.contiguous()
+
+    has_kv_mask = 1 if kv_valid_mask is not None else 0
+    _kv_mask = kv_valid_mask if kv_valid_mask is not None else torch.empty(0, dtype=torch.bool, device=q.device)
+
+    config_key = 'BLOCK_N_DQ_LG=64' if chunk_size_k == 64 else 'default'
+    kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', '0') == '1' else configs_bwd_dkdv_bsa_varlen_preset[config_key]
+    
+    grid_dkdv = lambda args: (triton.cdiv(arg_k.shape[2], args["BLOCK_N"]), 1, arg_k.shape[0] * arg_k.shape[1])
+    _attn_bwd_dkdv_bsa_varlen_wrapper[grid_dkdv](
+        q, arg_k, v, sm_scale,
+        do,
+        dk, dv,
+        M,
+        delta,
+        block_indices_k,
+        block_indices_k_lens,
+        _kv_mask,
+        q.stride(0), q.stride(1), q.stride(2), q.stride(3),
+        k.stride(0), k.stride(1), k.stride(2), k.stride(3),
+        v.stride(0), v.stride(1), v.stride(2), v.stride(3),
+        dk.stride(0), dk.stride(1), dk.stride(2), dk.stride(3),
+        dv.stride(0), dv.stride(1), dv.stride(2), dv.stride(3),
+        do.stride(0), do.stride(1), do.stride(2), do.stride(3),
+        M.stride(0), M.stride(1), M.stride(2),
+        delta.stride(0), delta.stride(1), delta.stride(2),
+        block_indices_k.stride(0), block_indices_k.stride(1), block_indices_k.stride(2), block_indices_k.stride(3), 
+        block_indices_k_lens.stride(0), block_indices_k_lens.stride(1), block_indices_k_lens.stride(2), 
+        N_HEAD, N_CTX,
+        BLOCK_M=chunk_size_q,
+        BLOCK_N_DQ_LG=chunk_size_k,
+        HEAD_DIM=HEAD_DIM,
+        SPARSITY=sparsity,
+        HAS_KV_MASK=has_kv_mask,
+        **kernel_config
+    )
+
+    config_key = 'BLOCK_N_DQ_LG=64' if chunk_size_k == 64 else 'default'
+    if chunk_size_k > 128:
+        bwd_dq_func = _attn_bwd_dq_bsa_varlen_wrapper
+        kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', '0') == '1' else configs_bwd_dq_bsa_varlen_preset[config_key]
+    else:
+        bwd_dq_func = _attn_bwd_dq_bsa_varlen_align_wrapper
+        kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', '0') == '1' else configs_bwd_dq_bsa_varlen_align_preset[config_key]
+        
+    grid_dq = lambda args: (triton.cdiv(q.shape[2], args["BLOCK_M"]), 1, q.shape[0] * q.shape[1])
+    bwd_dq_func[grid_dq](
+        q, arg_k, v,
+        do, 
+        dq,
+        M,
+        delta,
+        block_indices,
+        block_indices_lens,
+        _kv_mask,
+        q.stride(0), q.stride(1), q.stride(2), q.stride(3),
+        k.stride(0), k.stride(1), k.stride(2), k.stride(3),
+        v.stride(0), v.stride(1), v.stride(2), v.stride(3),
+        dq.stride(0), dq.stride(1), dq.stride(2), dq.stride(3),
+        do.stride(0), do.stride(1), do.stride(2), do.stride(3),
+        M.stride(0), M.stride(1), M.stride(2),
+        delta.stride(0), delta.stride(1), delta.stride(2),
+        block_indices.stride(0), block_indices.stride(1), block_indices.stride(2), block_indices.stride(3),
+        block_indices_lens.stride(0), block_indices_lens.stride(1), block_indices_lens.stride(2),
+        N_HEAD, N_CTX,
+        BLOCK_M=chunk_size_q, 
+        BLOCK_N_DQ_LG=chunk_size_k,
+        HEAD_DIM=HEAD_DIM,
+        SPARSITY=sparsity,
+        HAS_KV_MASK=has_kv_mask,
+        **kernel_config
+    )    
+
+@torch.compile
+def make_block_indices_varlen_cp_list(block_indices, cp_size, num_blocks_k_full):
+    """
+    Args:
+        block_indices: [B, H, num_blocks_q_per_cp_rank, num_blocks_k_full]
+    
+    Return:
+        a list of [block_indices, block_indices_lens] for k from each cp_rank
+            - each block_indices starts from zero
+            - block_indices_lens indicates the valid number of elements in the last dimension of block_indices
+    """
+    res = []
+    num_blocks_per_rank = num_blocks_k_full // cp_size
+    for i in range(cp_size):
+        block_indices_tmp = block_indices.clone()
+        min_block_idx = i * num_blocks_per_rank
+        block_indices_tmp -= min_block_idx
+        block_indices_tmp[block_indices_tmp < 0] = num_blocks_per_rank # block_indices_tmp < 0 indicate invalid indices, set them to num_blocks_per_rank in order to sort them to the tail, so that the first N elements of the block_indices indicated by block_indices_lens are valid
+        
+        block_indices_tmp, _ = torch.sort(block_indices_tmp, dim=-1)
+        
+        block_indices_tmp_lens = (block_indices_tmp < num_blocks_per_rank).sum(dim=-1)
+        
+        res.append([block_indices_tmp, block_indices_tmp_lens])
+
+    return res
+
+@torch.compile
+def flash_attn_fwd_softmax_lse_correction(
+    softmax_lse: torch.Tensor, 
+    softmax_lse_per_step: torch.Tensor,
+):
+    """Merge softmax stats of each step in Attention with context parallelism"""
+    max_scale = torch.max(softmax_lse, softmax_lse_per_step)
+    min_scale = torch.min(softmax_lse, softmax_lse_per_step)
+    lse_diff = min_scale - max_scale
+    lse_diff = lse_diff.nan_to_num(nan=0.) # handle cases: tensor(-inf) - tensor(-inf) = tensor(nan); In the current cp implementation, it is possible that lses of 2 cp ranks are both -inf, if no block is selected from both cp ranks. In such cases, the finally corrected lse should remain -inf.
+    new_scale = max_scale + torch.log1p(torch.exp(lse_diff)) # a + ln(1 + e^(b - a)) = ln(e^a) + ln(1 + e^(b - a)) = ln(e^a + e^b)
+    softmax_lse.copy_(new_scale)
+
+@torch.compile
+def flash_attn_fwd_out_correction_init(
+    out_init_step: torch.Tensor, # b h s d
+    softmax_lse: torch.Tensor, # b h s
+    softmax_lse_init_step: torch.Tensor,
+):
+    """Merge partial outputs of the first step in Attention with context parallelism"""
+    softmax_lse_corrected_exp = torch.exp(softmax_lse_init_step - softmax_lse)
+    softmax_lse_corrected_exp = softmax_lse_corrected_exp.unsqueeze(-1)
+    out_corrected = out_init_step * softmax_lse_corrected_exp
+    return out_corrected.to(out_init_step.dtype)
+
+
+@torch.compile
+def flash_attn_fwd_out_correction(
+    out: torch.Tensor,
+    out_per_step: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    softmax_lse_per_step: torch.Tensor,
+):
+    """Merge partial outputs of each step in Attention with context parallelism"""
+    softmax_lse_corrected_exp = torch.exp(softmax_lse_per_step - softmax_lse)
+    softmax_lse_corrected_exp = softmax_lse_corrected_exp.unsqueeze(-1)
+    out_corrected = out_per_step * softmax_lse_corrected_exp
+    out.add_(out_corrected)
+
+@torch.compile
+def topk_sort(score, num_chunks_selected):
+    block_indices = torch.topk(score, num_chunks_selected)[1]
+    block_indices, _ = torch.sort(block_indices, dim=-1)
+    return block_indices
+
+class _attention_bsa(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, q, k, v, chunk_size_q, chunk_size_k, sparsity, cdf_threshold, sm_scale,
+                kv_valid_mask=None, q_valid_mask=None, score_modifier=None):
+        HEAD_DIM_Q, HEAD_DIM_K = q.shape[-1], k.shape[-1]
+        HEAD_DIM_V = v.shape[-1]
+        assert HEAD_DIM_Q == HEAD_DIM_K and HEAD_DIM_K == HEAD_DIM_V
+        assert HEAD_DIM_K in {16, 32, 64, 128, 256}
+        
+        # ---------------------- gating ----------------------
+        if q_valid_mask is not None:
+            q_cmp = masked_mean_pooling_compression(q, chunk_size_q, q_valid_mask)
+        elif kv_valid_mask is not None and q.shape[2] == k.shape[2]:
+            q_cmp = masked_mean_pooling_compression(q, chunk_size_q, kv_valid_mask)
+        else:
+            q_cmp = mean_pooling_compression(q, chunk_size_q)
+
+        if kv_valid_mask is not None:
+            k_cmp = masked_mean_pooling_compression(k, chunk_size_k, kv_valid_mask)
+        else:
+            k_cmp = mean_pooling_compression(k, chunk_size_k)
+
+        score = cal_score(q_cmp, k_cmp)
+
+        if score_modifier is not None:
+            score = score * score_modifier
+
+        if sparsity is not None and cdf_threshold is None:
+            block_indices, block_indices_lens = get_select_indices_topk_from_score(score, sparsity)
+        elif sparsity is None and cdf_threshold is not None:
+            block_indices, block_indices_lens = get_select_indices_cdf_from_score(score, cdf_threshold, sm_scale)
+        elif sparsity is not None and cdf_threshold is not None:
+            block_indices, block_indices_lens = get_select_indices_cdf_topk_from_score(score, sparsity, cdf_threshold, sm_scale)
+        else:
+            raise ValueError("Either sparsity or cdf_threshold must be provided")
+
+        # ---------------------- bsa ----------------------
+
+        o, lse = attn_fwd_bsa_varlen_triton(
+            q, k, v, 
+            sm_scale, block_indices, block_indices_lens,
+            chunk_size_q, chunk_size_k, 
+            sparsity,
+            kv_valid_mask=kv_valid_mask,
+        )
+
+        ctx.save_for_backward(q, k, v, o, lse, block_indices, block_indices_lens)
+        ctx.sm_scale = sm_scale
+        ctx.HEAD_DIM = HEAD_DIM_K
+        ctx.chunk_size_q = chunk_size_q
+        ctx.chunk_size_k = chunk_size_k
+        ctx.sparsity = sparsity
+        ctx.kv_valid_mask = kv_valid_mask
+
+        return o
+
+    @staticmethod
+    def backward(ctx, do):
+        q, k, v, o, lse, block_indices, block_indices_lens = ctx.saved_tensors
+
+        dq = torch.empty_like(q)
+        dk = torch.empty_like(k)
+        dv = torch.empty_like(v)
+
+        attn_bwd_bsa_varlen_triton(
+            do, 
+            q, 
+            k, 
+            v, 
+            o,
+            dq,
+            dk,
+            dv,
+            ctx.sm_scale, 
+            lse, 
+            block_indices, 
+            block_indices_lens,
+            ctx.chunk_size_q, 
+            ctx.chunk_size_k,
+            ctx.sparsity,
+            kv_valid_mask=ctx.kv_valid_mask,
+        )
+
+        return dq, dk, dv, None, None, None, None, None, None, None, None
+
+flash_attn_bsa = _attention_bsa.apply
+
+
+class _attention_bsa_audio_guided(torch.autograd.Function):
+    """BSA with audio-guided score modulation (Gate 2: Audio Spatial Concentration Gate).
+
+    The audio saliency modulates the block-level QK scores BEFORE topk selection,
+    but the actual sparse attention Triton kernel remains unchanged. This ensures
+    Triton-level efficiency for the compute-heavy attention, while the lightweight
+    score modulation uses standard PyTorch ops (torch.compile friendly).
+    """
+
+    @staticmethod
+    def forward(ctx, q, k, v, chunk_size_q, chunk_size_k, sparsity, cdf_threshold, sm_scale,
+                audio_block_saliency, audio_gate_value, audio_boost_gamma,
+                kv_valid_mask=None, q_valid_mask=None,
+                audio_norms_block_order=None, audio_weighted_lambda=0.0,
+                score_modifier=None):
+        """
+        Args:
+            audio_block_saliency: [N_blocks_k] or [B, 1, N_blocks_k] — per-key-block saliency in [0,1]
+            audio_gate_value: scalar in [0,1] — Audio Spatial Concentration Gate g_a
+            audio_boost_gamma: scalar >= 0 — boost strength γ
+            audio_norms_block_order: [S] — per-token norms in block-rearranged order (for Path B)
+            audio_weighted_lambda: float — Path B weighting strength (0=standard pooling)
+            score_modifier: optional broadcastable tensor — external additive boost term
+                (e.g., γ_var * ρ from variance guidance). Combined with audio INSIDE
+                one boost bracket: score *= (1 + γ_audio * S̃ + score_modifier).
+        """
+        HEAD_DIM_Q, HEAD_DIM_K = q.shape[-1], k.shape[-1]
+        HEAD_DIM_V = v.shape[-1]
+        assert HEAD_DIM_Q == HEAD_DIM_K and HEAD_DIM_K == HEAD_DIM_V
+        assert HEAD_DIM_K in {16, 32, 64, 128, 256}
+
+        B, H, Sq, D = q.shape
+        Sk = k.shape[2]
+
+        # ---------------------- gating (block-level scoring) ----------------------
+        if q_valid_mask is not None:
+            q_cmp = masked_mean_pooling_compression(q, chunk_size_q, q_valid_mask)
+        elif kv_valid_mask is not None and q.shape[2] == k.shape[2]:
+            q_cmp = masked_mean_pooling_compression(q, chunk_size_q, kv_valid_mask)
+        else:
+            q_cmp = mean_pooling_compression(q, chunk_size_q)
+
+        use_path_b = (audio_norms_block_order is not None and audio_weighted_lambda > 0)
+        if use_path_b:
+            k_cmp = audio_weighted_mean_pooling_compression(
+                k, chunk_size_k, audio_norms_block_order,
+                valid_mask=kv_valid_mask, lambda_w=audio_weighted_lambda,
+            )
+        elif kv_valid_mask is not None:
+            k_cmp = masked_mean_pooling_compression(k, chunk_size_k, kv_valid_mask)
+        else:
+            k_cmp = mean_pooling_compression(k, chunk_size_k)
+
+        score = cal_score(q_cmp, k_cmp)  # [B, H, N_blocks_q, N_blocks_k]
+
+        # ---------------------- Unified Additive Boost (md 3.5.1.2) ----------------------
+        # I_{i,j} = P^QK * (1 + γ_audio * S̃ + additive_boost_from_external)
+        # All boost terms are ADDITIVE inside one bracket, then multiply P^QK.
+        N_blocks_k = score.shape[-1]
+        boost = torch.zeros(1, device=score.device, dtype=score.dtype)
+
+        has_audio = (audio_block_saliency is not None
+                     and audio_block_saliency.numel() > 0
+                     and audio_boost_gamma > 0)
+        if has_audio:
+            if audio_block_saliency.dim() == 1:
+                saliency = audio_block_saliency.view(1, 1, 1, -1)
+            else:
+                saliency = audio_block_saliency.unsqueeze(-2)
+            uniform_val = 1.0 / N_blocks_k
+            gated_saliency = audio_gate_value * saliency + (1.0 - audio_gate_value) * uniform_val
+            boost = boost + audio_boost_gamma * gated_saliency
+
+        if score_modifier is not None:
+            boost = boost + score_modifier
+
+        score = score * (1.0 + boost)
+
+        # ---------------------- topk selection from modulated scores ----------------------
+        if sparsity is not None and cdf_threshold is None:
+            block_indices, block_indices_lens = get_select_indices_topk_from_score(score, sparsity)
+        elif sparsity is None and cdf_threshold is not None:
+            block_indices, block_indices_lens = get_select_indices_cdf_from_score(score, cdf_threshold, sm_scale)
+        elif sparsity is not None and cdf_threshold is not None:
+            block_indices, block_indices_lens = get_select_indices_cdf_topk_from_score(score, sparsity, cdf_threshold, sm_scale)
+        else:
+            raise ValueError("Either sparsity or cdf_threshold must be provided")
+
+        # ---------------------- bsa (Triton kernel unchanged) ----------------------
+        o, lse = attn_fwd_bsa_varlen_triton(
+            q, k, v,
+            sm_scale, block_indices, block_indices_lens,
+            chunk_size_q, chunk_size_k,
+            sparsity,
+            kv_valid_mask=kv_valid_mask,
+        )
+
+        ctx.save_for_backward(q, k, v, o, lse, block_indices, block_indices_lens)
+        ctx.sm_scale = sm_scale
+        ctx.HEAD_DIM = HEAD_DIM_K
+        ctx.chunk_size_q = chunk_size_q
+        ctx.chunk_size_k = chunk_size_k
+        ctx.sparsity = sparsity
+        ctx.kv_valid_mask = kv_valid_mask
+
+        return o
+
+    @staticmethod
+    def backward(ctx, do):
+        q, k, v, o, lse, block_indices, block_indices_lens = ctx.saved_tensors
+
+        dq = torch.empty_like(q)
+        dk = torch.empty_like(k)
+        dv = torch.empty_like(v)
+
+        attn_bwd_bsa_varlen_triton(
+            do,
+            q,
+            k,
+            v,
+            o,
+            dq,
+            dk,
+            dv,
+            ctx.sm_scale,
+            lse,
+            block_indices,
+            block_indices_lens,
+            ctx.chunk_size_q,
+            ctx.chunk_size_k,
+            ctx.sparsity,
+            kv_valid_mask=ctx.kv_valid_mask,
+        )
+
+        # Order: q, k, v, chunk_size_q, chunk_size_k, sparsity, cdf_threshold, sm_scale,
+        #         audio_block_saliency, audio_gate_value, audio_boost_gamma,
+        #         kv_valid_mask, q_valid_mask, audio_norms_block_order, audio_weighted_lambda,
+        #         score_modifier
+        return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None
+
+
+flash_attn_bsa_audio_guided = _attention_bsa_audio_guided.apply
+
+
+@torch.compile
+def compute_timestep_reliability_gate(
+    timestep_ratio: torch.Tensor,
+    mode: str = "power",
+    power: float = 1.5,
+    scale: float = 10.0,
+    threshold: float = 0.7,
+) -> torch.Tensor:
+    """Compute Gate 1: Timestep Reliability Gate r_t (non-learnable).
+
+    At high noise (t→1): r_t→0 (don't trust audio guidance, bridge residual is noise-driven).
+    At low noise (t→0): r_t→1 (trust audio guidance, bridge residual reflects real saliency).
+
+    Two modes:
+      - "power": r_t = (1 - t)^power. Physically motivated by flow matching signal decay.
+        power=1.5 is recommended (moderate decay, default).
+      - "sigmoid": r_t = sigmoid(scale * (threshold - t)). Sharp transition.
+
+    Args:
+        timestep_ratio: scalar or [B] — normalized timestep in [0, 1] (0=clean, 1=noise)
+        mode: "power" (default) or "sigmoid"
+        power: exponent for power-law mode (default 1.5)
+        scale: sharpness for sigmoid mode (default 10.0)
+        threshold: midpoint for sigmoid mode (default 0.7)
+
+    Returns:
+        r_t: scalar in [0, 1]
+    """
+    if mode == "power":
+        r_t = (1.0 - timestep_ratio.clamp(0.0, 1.0)) ** power
+    else:
+        r_t = torch.sigmoid(scale * (threshold - timestep_ratio))
+    return r_t
+
+
+@torch.compile
+def compute_audio_spatial_concentration_gate(
+    token_norms: torch.Tensor,
+    valid_mask: torch.Tensor = None,
+    alpha: float = 5.0,
+    beta: float = 2.0,
+) -> torch.Tensor:
+    """Compute Gate 2: Audio Spatial Concentration Gate g_a from per-token L2 norms.
+
+    g_a = sigmoid(alpha * CV - beta)
+    where CV = std(norms) / (mean(norms) + eps)
+
+    Args:
+        token_norms: [N_tokens] — L2 norms of bridge residual per video token
+        valid_mask: [N_tokens] bool — True for real tokens, False for padded
+        alpha, beta: fixed hyperparameters (non-learnable)
+
+    Returns:
+        g_a: scalar in [0, 1]
+    """
+    eps = 1e-6
+    if valid_mask is not None:
+        mask_f = valid_mask.float()
+        count = mask_f.sum().clamp(min=1.0)
+        mean_a = (token_norms * mask_f).sum() / count
+        var_a = ((token_norms - mean_a).pow(2) * mask_f).sum() / count
+    else:
+        mean_a = token_norms.mean()
+        var_a = token_norms.var(unbiased=False)
+
+    std_a = var_a.sqrt()
+    cv = std_a / (mean_a + eps)
+    g_a = torch.sigmoid(alpha * cv - beta)
+    return g_a
+
+
+@torch.compile
+def audio_weighted_mean_pooling_compression(
+    x: torch.Tensor,
+    block_size: int,
+    token_norms: torch.Tensor,
+    valid_mask: torch.Tensor = None,
+    lambda_w: float = 1.0,
+) -> torch.Tensor:
+    """Audio-weighted mean pooling for K-side block representatives (Path B).
+
+    w_m = 1 + lambda_w * (a_m / max(a_m' in block) + eps)
+    K_bar_j = sum(w_m * K_m) / sum(w_m)
+
+    Automatically degrades to standard mean pooling when:
+    - silence: all a_m ≈ 0 → w_m ≈ 1
+    - high noise: all a_m similar → a_m/max ≈ 1 → uniform weighting
+    - lambda_w = 0: w_m = 1 exactly
+
+    Args:
+        x: [B, H, S, D] — K or V tokens (already in 3D-block rearranged order)
+        block_size: number of tokens per block (e.g., 64)
+        token_norms: [S] — per-token L2 norms of bridge residual (block-rearranged order)
+        valid_mask: [S] bool — True for real tokens, False for padded
+        lambda_w: weighting strength (default 1.0; 0=standard mean pooling)
+
+    Returns:
+        x_cmp: [B, H, N_blocks, D] — weighted block representatives
+    """
+    B, H, S, D = x.shape
+    eps = 1e-6
+    num_blocks = math.ceil(S / block_size)
+    if S % block_size != 0:
+        pad_len = num_blocks * block_size - S
+        x = F.pad(x, (0, 0, 0, pad_len))
+        token_norms = F.pad(token_norms, (0, pad_len), value=0.0)
+        if valid_mask is not None:
+            valid_mask = F.pad(valid_mask, (0, pad_len), value=False)
+
+    # Reshape to [num_blocks, block_size]
+    norms_blocks = token_norms.view(num_blocks, block_size)
+
+    # Intra-block normalized weighting: w_m = 1 + λ * (a_m / max_in_block)
+    if valid_mask is not None:
+        mask_blocks = valid_mask.view(num_blocks, block_size).float()
+        # Mask out padded positions for max computation
+        norms_masked = norms_blocks * mask_blocks + (-1e9) * (1 - mask_blocks)
+        block_max = norms_masked.max(dim=1, keepdim=True).values.clamp(min=eps)
+        # Reset negatives from masking
+        norms_for_weight = norms_blocks * mask_blocks
+    else:
+        block_max = norms_blocks.max(dim=1, keepdim=True).values.clamp(min=eps)
+        norms_for_weight = norms_blocks
+        mask_blocks = None
+
+    # w_m = 1 + λ * (a_m / max)
+    weights = 1.0 + lambda_w * (norms_for_weight / block_max)  # [num_blocks, block_size]
+
+    if mask_blocks is not None:
+        weights = weights * mask_blocks  # zero out padded positions
+
+    # Reshape x for block-wise weighted mean: [B, H, num_blocks, block_size, D]
+    x_blocks = x.view(B, H, num_blocks, block_size, D)
+
+    # weights: [num_blocks, block_size] → [1, 1, num_blocks, block_size, 1] for broadcasting
+    w = weights.unsqueeze(0).unsqueeze(0).unsqueeze(-1)  # [1, 1, num_blocks, block_size, 1]
+    w_sum = w.sum(dim=3, keepdim=True).clamp(min=1.0)  # [1, 1, num_blocks, 1, 1]
+
+    x_cmp = (x_blocks * w).sum(dim=3) / w_sum.squeeze(3)  # [B, H, num_blocks, D]
+    return x_cmp
+
+
+@torch.compile
+def compute_channel_variance_density(
+    x: torch.Tensor,
+    block_size: int,
+    bt: int, bh: int, bw: int,
+    valid_mask: torch.Tensor = None,
+) -> torch.Tensor:
+    """Compute per-block channel-variance information density (non-learnable).
+
+    Averages across all (local) heads for head-agnostic block-level statistics
+    (different heads project into different subspaces — single-head would bias
+    selection toward that head's variance pattern). Mean pooling stays in native
+    dtype (bf16) to avoid large float32 copies; only the tiny [num_blocks, D]
+    result converts to float32 for numerically precise variance computation.
+
+    Note (SP): under sequence parallelism each rank owns a head shard, so ρ is
+    computed from local heads only. This keeps the per-layer cost minimal (no
+    extra collective). ρ is a coarse per-block density signal whose estimate is
+    stable across head subsets, and each head's attention is computed on its own
+    rank, so the per-rank ρ is correct for that rank's heads.
+
+    Args:
+        x: [B, H, S, D] V tensor in 3D-block-rearranged order (V = x @ W_V, no RMSNorm/RoPE).
+        block_size: bt * bh * bw tokens per block.
+        bt, bh, bw: block shape in temporal, height, width.
+        valid_mask: [S] bool — True for real tokens, False for padding.
+
+    Returns:
+        rho: [num_blocks] ∈ [0, 1] — normalized variance density per block.
+    """
+    S, D = x.shape[2], x.shape[3]
+    num_blocks = S // block_size
+    spatial_size = bh * bw
+
+    # Average across heads in native dtype (bf16) — no large float32 copy.
+    # x[0]: [H, S, D] → .mean(dim=0): [S, D]. bf16 mean is numerically stable.
+    x_avg = x[0].mean(dim=0)                               # [S, D], bf16, head-agnostic
+    x_blocks = x_avg.view(num_blocks, block_size, D)        # [num_blocks, block_size, D]
+
+    if valid_mask is not None:
+        mask_blocks = valid_mask.view(num_blocks, block_size)
+        mask_f = mask_blocks.unsqueeze(-1).to(x_blocks.dtype)  # [num_blocks, block_size, 1]
+        counts = mask_f.sum(dim=1).clamp(min=1.0)              # [num_blocks, 1]
+    else:
+        mask_f = None
+        counts = None
+
+    # --- Spatial Channel Variance σ²_ch-s ---
+    # Per-block mean per channel (bf16), then variance across channels (float32)
+    if mask_f is not None:
+        channel_means = (x_blocks * mask_f).sum(dim=1) / counts  # [num_blocks, D], bf16
+    else:
+        channel_means = x_blocks.mean(dim=1)                      # [num_blocks, D], bf16
+
+    cm_f = channel_means.float()                                  # [num_blocks, D] — only this is float32 (~2.5MB)
+    global_cm = cm_f.mean(dim=-1, keepdim=True)                   # [num_blocks, 1]
+    spatial_var = ((cm_f - global_cm) ** 2).mean(dim=-1)          # [num_blocks]
+
+    # --- Temporal Channel Variance σ²_ch-t ---
+    if bt > 1:
+        x_temporal = x_blocks.view(num_blocks, bt, spatial_size, D)
+
+        if valid_mask is not None:
+            mask_t_3d = mask_blocks.view(num_blocks, bt, spatial_size)
+            mask_t = mask_t_3d.unsqueeze(-1).to(x_blocks.dtype)   # [num_blocks, bt, spatial_size, 1]
+            fc = mask_t.sum(dim=2).clamp(min=1.0)                 # [num_blocks, bt, 1]
+            frame_means = (x_temporal * mask_t).sum(dim=2) / fc    # [num_blocks, bt, D], bf16
+
+            # Frame-level validity: only diff between pairs where BOTH frames have ≥1 valid token
+            frame_valid = mask_t_3d.any(dim=2)                     # [num_blocks, bt]
+            both_valid = frame_valid[:, 1:] & frame_valid[:, :-1]  # [num_blocks, bt-1]
+            bv_f = both_valid.float().unsqueeze(-1)                # [num_blocks, bt-1, 1]
+            valid_pairs = bv_f.sum(dim=1).clamp(min=1.0)           # [num_blocks, 1]
+
+            fm_f = frame_means.float()
+            frame_diffs = fm_f[:, 1:, :] - fm_f[:, :-1, :]        # [num_blocks, bt-1, D]
+            temporal_var = ((frame_diffs.pow(2) * bv_f).sum(dim=1) / valid_pairs).mean(dim=-1)
+        else:
+            frame_means = x_temporal.mean(dim=2)                    # [num_blocks, bt, D], bf16
+            fm_f = frame_means.float()
+            frame_diffs = fm_f[:, 1:, :] - fm_f[:, :-1, :]         # [num_blocks, bt-1, D]
+            temporal_var = frame_diffs.pow(2).mean(dim=1).mean(dim=-1)
+    else:
+        temporal_var = torch.zeros(num_blocks, device=x.device, dtype=torch.float32)
+
+    # --- Combine and normalize to [0, 1] ---
+    combined = spatial_var + temporal_var
+    rho = combined / combined.max().clamp(min=1e-6)
+
+    return rho
+
+
+@torch.compile
+def compute_block_audio_saliency(
+    token_norms: torch.Tensor,
+    chunk_size: int,
+    valid_mask: torch.Tensor = None,
+) -> torch.Tensor:
+    """Compute per-block audio saliency from per-token L2 norms.
+
+    For each block j: S_j = mean(norms in block j) / max(mean norms across all blocks)
+    Then clamp to [0, 1].
+
+    Args:
+        token_norms: [S] — per-token L2 norms (already in 3D-block rearranged order)
+        chunk_size: block size (e.g., 64 for 4x4x4)
+        valid_mask: [S] bool — True for real tokens, False for padded
+
+    Returns:
+        block_saliency: [N_blocks] in [0, 1]
+    """
+    S = token_norms.shape[0]
+    num_blocks = math.ceil(S / chunk_size)
+    if S % chunk_size != 0:
+        pad_len = num_blocks * chunk_size - S
+        token_norms = F.pad(token_norms, (0, pad_len), value=0.0)
+        if valid_mask is not None:
+            valid_mask = F.pad(valid_mask, (0, pad_len), value=False)
+
+    norms_blocks = token_norms.view(num_blocks, chunk_size)
+
+    if valid_mask is not None:
+        mask_blocks = valid_mask.view(num_blocks, chunk_size).float()
+        counts = mask_blocks.sum(dim=1).clamp(min=1.0)
+        block_means = (norms_blocks * mask_blocks).sum(dim=1) / counts
+    else:
+        block_means = norms_blocks.mean(dim=1)
+
+    max_val = block_means.max().clamp(min=1e-6)
+    block_saliency = block_means / max_val
+    return block_saliency
+
+
+def rearrange_THW_to_3d_block(x, Nt, Nh, Nw, t, h, w, D):
+    B, H, _, D = x.shape
+    x = x.view(B, H, Nt, t, Nh, h, Nw, w, D)
+    x = x.permute(0, 1, 2, 4, 6, 3, 5, 7, 8)  # B H Nt Nh Nw t h w D
+    return x.contiguous().view(B, H, Nt * Nh * Nw * t * h * w, D)
+
+def rearrange_3d_block_to_THW(x, Nt, Nh, Nw, t, h, w, D):
+    B, H, _, D = x.shape
+    x = x.view(B, H, Nt, Nh, Nw, t, h, w, D)
+    x = x.permute(0, 1, 2, 5, 3, 6, 4, 7, 8)  # B H Nt t Nh h Nw w D
+    return x.contiguous().view(B, H, Nt * t * Nh * h * Nw * w, D)
+
+def rearrange_THW_to_3d_block_1d(mask, Nt, Nh, Nw, t, h, w):
+    """Rearrange a 1D bool mask from THW order to 3D block order (same permutation as tokens)."""
+    mask = mask.view(Nt, t, Nh, h, Nw, w)
+    mask = mask.permute(0, 2, 4, 1, 3, 5)  # Nt Nh Nw t h w
+    return mask.contiguous().view(-1)
+
+def flash_attn_bsa_3d(
+    q: torch.Tensor, # [B, H, Sq, D]
+    k: torch.Tensor, # [B, H, Skv, D]
+    v: torch.Tensor, # [B, H, Skv, D]
+    latent_shape_q,
+    latent_shape_k,
+    # bsa_params
+    sparsity=0.875,
+    cdf_threshold=None,
+    chunk_3d_shape_q=[4, 4, 8],
+    chunk_3d_shape_k=[4, 4, 8],
+    valid_mask=None,  # [Sq] bool, True=valid token, False=padded
+) -> torch.Tensor:
+    _, _, Sq, head_dim_q = q.shape
+    _, _, Sk, head_dim_k = k.shape
+    
+    assert head_dim_q == head_dim_k
+    head_dim = head_dim_q
+    
+    Tq, Hq, Wq = latent_shape_q
+    Tk, Hk, Wk = latent_shape_k
+    
+    assert Tq * Hq * Wq == Sq
+    assert Tk * Hk * Wk == Sk
+    
+    tq, hq, wq = chunk_3d_shape_q
+    tk, hk, wk = chunk_3d_shape_k
+    
+    assert Tq % tq == 0 and Hq % hq == 0 and Wq % wq == 0
+    assert Tk % tk == 0 and Hk % hk == 0 and Wk % wk == 0
+    
+    Ntq = Tq // tq
+    Nhq = Hq // hq
+    Nwq = Wq // wq
+
+    Ntk = Tk // tk
+    Nhk = Hk // hk
+    Nwk = Wk // wk
+
+    q = rearrange_THW_to_3d_block(q, Ntq, Nhq, Nwq, tq, hq, wq, q.shape[-1])
+    k = rearrange_THW_to_3d_block(k, Ntk, Nhk, Nwk, tk, hk, wk, k.shape[-1])
+    v = rearrange_THW_to_3d_block(v, Ntk, Nhk, Nwk, tk, hk, wk, v.shape[-1])
+
+    kv_valid_mask = None
+    if valid_mask is not None:
+        kv_valid_mask = rearrange_THW_to_3d_block_1d(valid_mask, Ntk, Nhk, Nwk, tk, hk, wk)
+
+    chunk_size_q = tq * hq * wq
+    chunk_size_k = tk * hk * wk
+
+    output = flash_attn_bsa(q, k, v, chunk_size_q, chunk_size_k, sparsity, cdf_threshold, 1 / head_dim**0.5, kv_valid_mask)
+
+    output = rearrange_3d_block_to_THW(output, Ntq, Nhq, Nwq, tq, hq, wq, output.shape[-1])
+    return output
+
+
+def flash_attn_bsa_3d_variance_guided(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    latent_shape_q,
+    latent_shape_k,
+    variance_boost_gamma: float = 1.0,
+    sparsity=0.875,
+    cdf_threshold=None,
+    chunk_3d_shape_q=[4, 4, 4],
+    chunk_3d_shape_k=[4, 4, 4],
+    valid_mask=None,
+) -> torch.Tensor:
+    """BSA with channel-variance information density score modulation.
+
+    Completely independent from audio guidance — no shared state or code.
+
+    Computes ρ_j on V blocks (key/value side only), then applies:
+      score_{i,j} *= (1 + γ * ρ_j)
+
+    Query-awareness comes from P^QK being per-query (zero-gate effect):
+      - Sky query → lip key: P^QK ≈ 0, boost = 1.85, final ≈ 0 → not selected
+      - Face query → lip key: P^QK ≈ 0.12, boost = 1.85, final = 0.222 → selected
+    Same ρ_j, different P^QK → different top-k per query block.
+
+    Why V (not K or hidden states)?
+      V = x @ W_V — raw linear projection, no RMSNorm, no RoPE distortion.
+      V channel variance directly predicts information contribution I(V_block; Output):
+      low variance → V rows near-parallel → output gain ≈ 0 regardless of attention weight.
+    """
+    _, _, Sq, head_dim_q = q.shape
+    _, _, Sk, head_dim_k = k.shape
+
+    assert head_dim_q == head_dim_k
+    head_dim = head_dim_q
+
+    Tq, Hq, Wq = latent_shape_q
+    Tk, Hk, Wk = latent_shape_k
+
+    assert Tq * Hq * Wq == Sq
+    assert Tk * Hk * Wk == Sk
+
+    tq, hq, wq = chunk_3d_shape_q
+    tk, hk, wk = chunk_3d_shape_k
+
+    assert Tq % tq == 0 and Hq % hq == 0 and Wq % wq == 0
+    assert Tk % tk == 0 and Hk % hk == 0 and Wk % wk == 0
+
+    Ntq = Tq // tq
+    Nhq = Hq // hq
+    Nwq = Wq // wq
+
+    Ntk = Tk // tk
+    Nhk = Hk // hk
+    Nwk = Wk // wk
+
+    q = rearrange_THW_to_3d_block(q, Ntq, Nhq, Nwq, tq, hq, wq, q.shape[-1])
+    k = rearrange_THW_to_3d_block(k, Ntk, Nhk, Nwk, tk, hk, wk, k.shape[-1])
+    v = rearrange_THW_to_3d_block(v, Ntk, Nhk, Nwk, tk, hk, wk, v.shape[-1])
+
+    kv_valid_mask = None
+    if valid_mask is not None:
+        kv_valid_mask = rearrange_THW_to_3d_block_1d(valid_mask, Ntk, Nhk, Nwk, tk, hk, wk)
+
+    chunk_size_q = tq * hq * wq
+    chunk_size_k = tk * hk * wk
+
+    with torch.no_grad():
+        rho = compute_channel_variance_density(v, chunk_size_k, tk, hk, wk, kv_valid_mask)
+
+    # Multiplicative modifier for variance-only path: score *= (1 + γ * ρ)
+    score_modifier = (
+        1.0 + variance_boost_gamma * rho.to(q.dtype).view(1, 1, 1, -1)
+    ).contiguous()
+
+    output = flash_attn_bsa(
+        q, k, v, chunk_size_q, chunk_size_k, sparsity, cdf_threshold,
+        1 / head_dim**0.5, kv_valid_mask, None, score_modifier,
+    )
+
+    output = rearrange_3d_block_to_THW(output, Ntq, Nhq, Nwq, tq, hq, wq, output.shape[-1])
+    return output
+
+
+def _pad_1d_for_bsa(x, chunk_size):
+    """Pad a 1D sequence [B, H, S, D] to be divisible by chunk_size. Returns (padded_x, valid_mask_or_None, orig_len)."""
+    B, H, S, D = x.shape
+    remainder = S % chunk_size
+    if remainder == 0:
+        return x, None, S
+    pad_len = chunk_size - remainder
+    x_padded = F.pad(x, (0, 0, 0, pad_len)).contiguous()
+    valid_mask = torch.zeros(S + pad_len, dtype=torch.bool, device=x.device)
+    valid_mask[:S] = True
+    return x_padded, valid_mask, S
+
+
+def flash_attn_bsa_cross(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    q_structure: str,
+    k_structure: str,
+    q_grid_size=None,
+    k_grid_size=None,
+    sparsity=0.875,
+    cdf_threshold=None,
+    chunk_size_q=64,
+    chunk_size_k=64,
+    chunk_3d_shape_q=None,
+    chunk_3d_shape_k=None,
+) -> torch.Tensor:
+    """Cross-modal BSA supporting asymmetric Q/K structures (3D video vs 1D audio).
+
+    Args:
+        q: [B, H, Sq, D]
+        k: [B, H, Sk, D]
+        v: [B, H, Sk, D]
+        q_structure: "3d" (video with grid T,H,W) or "1d" (audio sequential)
+        k_structure: "3d" or "1d"
+        q_grid_size: (T, H, W) when q_structure="3d"
+        k_grid_size: (T, H, W) when k_structure="3d"
+        sparsity: fraction of K blocks to skip
+        chunk_size_q: block size for 1D Q (must be multiple of 64)
+        chunk_size_k: block size for 1D K (must be multiple of 64)
+        chunk_3d_shape_q: [t,h,w] for 3D Q blocking
+        chunk_3d_shape_k: [t,h,w] for 3D K blocking
+    """
+    _, _, Sq, head_dim = q.shape
+    _, _, Sk, _ = k.shape
+
+    q_valid_mask = None
+    k_valid_mask = None
+    q_orig_len = Sq
+    k_orig_len = Sk
+
+    if q_structure == "3d":
+        assert q_grid_size is not None
+        Tq, Hq, Wq = q_grid_size
+        tq, hq, wq = chunk_3d_shape_q or [4, 4, 4]
+        pad_tq = (tq - Tq % tq) % tq
+        pad_hq = (hq - Hq % hq) % hq
+        pad_wq = (wq - Wq % wq) % wq
+
+        if pad_tq > 0 or pad_hq > 0 or pad_wq > 0:
+            B, H, _, D = q.shape
+            Tq_p, Hq_p, Wq_p = Tq + pad_tq, Hq + pad_hq, Wq + pad_wq
+            idx_t = torch.arange(Tq_p, device=q.device)
+            idx_h = torch.arange(Hq_p, device=q.device)
+            idx_w = torch.arange(Wq_p, device=q.device)
+            q_valid_mask = ((idx_t[:, None, None] < Tq) & (idx_h[None, :, None] < Hq) & (idx_w[None, None, :] < Wq)).reshape(-1).contiguous()
+            q_3d = q.view(B, H, Tq, Hq, Wq, D)
+            q = F.pad(q_3d, (0, 0, 0, pad_wq, 0, pad_hq, 0, pad_tq)).reshape(B, H, -1, D).contiguous()
+            Tq, Hq, Wq = Tq_p, Hq_p, Wq_p
+
+        Ntq, Nhq, Nwq = Tq // tq, Hq // hq, Wq // wq
+        q = rearrange_THW_to_3d_block(q, Ntq, Nhq, Nwq, tq, hq, wq, q.shape[-1])
+        if q_valid_mask is not None:
+            q_valid_mask = rearrange_THW_to_3d_block_1d(q_valid_mask, Ntq, Nhq, Nwq, tq, hq, wq)
+        actual_chunk_size_q = tq * hq * wq
+    else:
+        q, q_valid_mask, q_orig_len = _pad_1d_for_bsa(q, chunk_size_q)
+        actual_chunk_size_q = chunk_size_q
+        Ntq = Nhq = Nwq = tq = hq = wq = None
+
+    if k_structure == "3d":
+        assert k_grid_size is not None
+        Tk, Hk, Wk = k_grid_size
+        tk, hk, wk = chunk_3d_shape_k or [4, 4, 4]
+        pad_tk = (tk - Tk % tk) % tk
+        pad_hk = (hk - Hk % hk) % hk
+        pad_wk = (wk - Wk % wk) % wk
+
+        if pad_tk > 0 or pad_hk > 0 or pad_wk > 0:
+            B, H, _, D = k.shape
+            Tk_p, Hk_p, Wk_p = Tk + pad_tk, Hk + pad_hk, Wk + pad_wk
+            idx_t = torch.arange(Tk_p, device=k.device)
+            idx_h = torch.arange(Hk_p, device=k.device)
+            idx_w = torch.arange(Wk_p, device=k.device)
+            k_valid_mask = ((idx_t[:, None, None] < Tk) & (idx_h[None, :, None] < Hk) & (idx_w[None, None, :] < Wk)).reshape(-1).contiguous()
+            B, H, _, D = k.shape
+            k_3d = k.view(B, H, Tk, Hk, Wk, D)
+            k = F.pad(k_3d, (0, 0, 0, pad_wk, 0, pad_hk, 0, pad_tk)).reshape(B, H, -1, D).contiguous()
+            v_3d = v.view(B, H, Tk, Hk, Wk, D)
+            v = F.pad(v_3d, (0, 0, 0, pad_wk, 0, pad_hk, 0, pad_tk)).reshape(B, H, -1, D).contiguous()
+            Tk, Hk, Wk = Tk_p, Hk_p, Wk_p
+
+        Ntk, Nhk, Nwk = Tk // tk, Hk // hk, Wk // wk
+        k = rearrange_THW_to_3d_block(k, Ntk, Nhk, Nwk, tk, hk, wk, k.shape[-1])
+        v = rearrange_THW_to_3d_block(v, Ntk, Nhk, Nwk, tk, hk, wk, v.shape[-1])
+        if k_valid_mask is not None:
+            k_valid_mask = rearrange_THW_to_3d_block_1d(k_valid_mask, Ntk, Nhk, Nwk, tk, hk, wk)
+        actual_chunk_size_k = tk * hk * wk
+    else:
+        k, k_valid_mask_k, k_orig_len = _pad_1d_for_bsa(k, chunk_size_k)
+        v, _, _ = _pad_1d_for_bsa(v, chunk_size_k)
+        if k_valid_mask_k is not None:
+            k_valid_mask = k_valid_mask_k
+        actual_chunk_size_k = chunk_size_k
+        Ntk = Nhk = Nwk = tk = hk = wk = None
+
+    output = flash_attn_bsa(q, k, v, actual_chunk_size_q, actual_chunk_size_k, sparsity, cdf_threshold, 1 / head_dim**0.5, k_valid_mask, q_valid_mask)
+
+    if q_structure == "3d":
+        output = rearrange_3d_block_to_THW(output, Ntq, Nhq, Nwq, tq, hq, wq, output.shape[-1])
+        if pad_tq > 0 or pad_hq > 0 or pad_wq > 0:
+            B, H, _, D = output.shape
+            Tq_orig, Hq_orig, Wq_orig = q_grid_size
+            output = output.view(B, H, Tq, Hq, Wq, D)[:, :, :Tq_orig, :Hq_orig, :Wq_orig, :].contiguous().reshape(B, H, q_orig_len, D)
+    else:
+        if q_orig_len != output.shape[2]:
+            output = output[:, :, :q_orig_len, :].contiguous()
+
+    return output
+
+
+def flash_attn_bsa_3d_audio_guided(
+    q: torch.Tensor,  # [B, H, Sq, D]
+    k: torch.Tensor,  # [B, H, Skv, D]
+    v: torch.Tensor,  # [B, H, Skv, D]
+    latent_shape_q,
+    latent_shape_k,
+    audio_token_norms: torch.Tensor = None,  # [Sq] per-token L2 norms of A->V bridge residual
+    audio_gate_value: torch.Tensor = None,  # scalar gate g_a in [0,1]
+    audio_boost_gamma: torch.Tensor = None,  # scalar boost γ >= 0
+    sparsity=0.875,
+    cdf_threshold=None,
+    chunk_3d_shape_q=[4, 4, 4],
+    chunk_3d_shape_k=[4, 4, 4],
+    valid_mask=None,  # [Sq] bool, True=valid, False=padded
+    audio_weighted_lambda: float = 0.0,  # Path B: weighted pooling strength (0=disabled)
+    score_modifier: torch.Tensor = None,  # external additive boost (e.g., γ_var * ρ from variance guidance)
+    variance_boost_gamma: float = 0.0,  # if > 0, compute ρ INSIDE on rearranged V (avoids double rearrange)
+) -> torch.Tensor:
+    """BSA with audio saliency score modulation.
+
+    Audio-specific: gated saliency boost + optional audio-weighted K pooling.
+
+    Combine modes for additional boosts (all additive inside one bracket):
+      score *= (1 + γ_audio * S̃ + γ_var * ρ + score_modifier)
+
+    For audio + variance combination, prefer passing `variance_boost_gamma > 0`
+    instead of `score_modifier`: that lets us compute ρ on the **already
+    rearranged** V tensor, avoiding a second `rearrange_THW_to_3d_block(v)`
+    on the caller side. Both produce identical results, but the in-function
+    path saves one V permute+contiguous (≈30-80 MB memory bandwidth on 2K).
+    """
+    _, _, Sq, head_dim_q = q.shape
+    _, _, Sk, head_dim_k = k.shape
+
+    assert head_dim_q == head_dim_k
+    head_dim = head_dim_q
+
+    Tq, Hq, Wq = latent_shape_q
+    Tk, Hk, Wk = latent_shape_k
+
+    assert Tq * Hq * Wq == Sq
+    assert Tk * Hk * Wk == Sk
+
+    tq, hq, wq = chunk_3d_shape_q
+    tk, hk, wk = chunk_3d_shape_k
+
+    assert Tq % tq == 0 and Hq % hq == 0 and Wq % wq == 0
+    assert Tk % tk == 0 and Hk % hk == 0 and Wk % wk == 0
+
+    Ntq = Tq // tq
+    Nhq = Hq // hq
+    Nwq = Wq // wq
+
+    Ntk = Tk // tk
+    Nhk = Hk // hk
+    Nwk = Wk // wk
+
+    q = rearrange_THW_to_3d_block(q, Ntq, Nhq, Nwq, tq, hq, wq, q.shape[-1])
+    k = rearrange_THW_to_3d_block(k, Ntk, Nhk, Nwk, tk, hk, wk, k.shape[-1])
+    v = rearrange_THW_to_3d_block(v, Ntk, Nhk, Nwk, tk, hk, wk, v.shape[-1])
+
+    kv_valid_mask = None
+    if valid_mask is not None:
+        kv_valid_mask = rearrange_THW_to_3d_block_1d(valid_mask, Ntk, Nhk, Nwk, tk, hk, wk)
+
+    chunk_size_q = tq * hq * wq
+    chunk_size_k = tk * hk * wk
+
+    # Optional in-function variance density (avoids the caller re-rearranging V).
+    # Folded additively into score_modifier so audio + variance share one boost
+    # bracket — no cross-term, exactly matches the standalone variance path.
+    if variance_boost_gamma is not None and variance_boost_gamma > 0:
+        with torch.no_grad():
+            _rho = compute_channel_variance_density(v, chunk_size_k, tk, hk, wk, kv_valid_mask)
+        _var_term = (variance_boost_gamma * _rho.to(q.dtype).view(1, 1, 1, -1)).contiguous()
+        score_modifier = (score_modifier + _var_term) if score_modifier is not None else _var_term
+
+    # Audio saliency computation (only when audio norms are provided)
+    has_audio_norms = (audio_token_norms is not None and audio_token_norms.numel() > 0)
+    if has_audio_norms:
+        audio_norms_block_order = rearrange_THW_to_3d_block_1d(
+            audio_token_norms, Ntk, Nhk, Nwk, tk, hk, wk
+        )
+        block_saliency = compute_block_audio_saliency(
+            audio_norms_block_order, chunk_size_k, valid_mask=kv_valid_mask
+        )
+    else:
+        block_saliency = None
+        audio_norms_block_order = None
+        if audio_gate_value is None:
+            audio_gate_value = torch.ones(1, device=q.device, dtype=q.dtype)
+        if audio_boost_gamma is None:
+            audio_boost_gamma = torch.zeros(1, device=q.device, dtype=q.dtype)
+
+    output = flash_attn_bsa_audio_guided(
+        q, k, v, chunk_size_q, chunk_size_k, sparsity, cdf_threshold,
+        1 / head_dim**0.5,
+        block_saliency, audio_gate_value, audio_boost_gamma,
+        kv_valid_mask,
+        None,  # q_valid_mask (same grid in self-attn)
+        audio_norms_block_order if audio_weighted_lambda > 0 else None,
+        audio_weighted_lambda,
+        score_modifier,
+    )
+
+    output = rearrange_3d_block_to_THW(output, Ntq, Nhq, Nwq, tq, hq, wq, output.shape[-1])
+    return output
+
+

hymm/models/modules/block_sparse_attention/common.py

@@ -0,0 +1,111 @@
+import triton
+import triton.language as tl
+import os
+
+if os.environ.get('TRITON_AUTOTUNE_ENBALE', '0') == '1':
+    autotune = triton.autotune
+else:
+    def autotune(*args, **kwargs):
+        def decorator(func):
+            return func
+        return decorator
+
+configs_gating_preset = {
+    'default': {
+        'BLOCK_M': 64,
+        'BLOCK_N': 64,
+        'num_stages': 3,
+        'num_warps': 8,
+    }
+}
+
+configs_gating = [
+    triton.Config({'BLOCK_M': BM, 'BLOCK_N': BN}, num_stages=s, num_warps=w) \
+    for BM in [64, 128] \
+    for BN in [32, 64] \
+    for s in [2, 3, 4, 5] \
+    for w in [4, 8] \
+]
+
+gating_reevaluate_keys = ["M", "N"] if os.environ.get('TRITON_REEVALUATE_KEY', '0') == '1' else []
+@autotune(configs_gating, key=gating_reevaluate_keys)
+@triton.jit
+def _attn_fwd_gating(
+    Q, K, Out, 
+    stride_qz, stride_qh, stride_qm, stride_qk, 
+    stride_kz, stride_kh, stride_kn, stride_kk, 
+    stride_oz, stride_oh, stride_om, stride_on, 
+    H, M, N, 
+    HEAD_DIM: tl.constexpr, 
+    BLOCK_M: tl.constexpr, 
+    BLOCK_N: tl.constexpr, 
+    ):
+    
+    tl.static_assert(BLOCK_N <= HEAD_DIM)
+    start_m = tl.program_id(0)
+    off_hz = tl.program_id(1)
+    off_z = off_hz // H
+    off_h = off_hz % H
+    q_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh
+    k_offset = off_z.to(tl.int64) * stride_kz + off_h.to(tl.int64) * stride_kh
+    o_offset = off_z.to(tl.int64) * stride_oz + off_h.to(tl.int64) * stride_oh
+
+    # block pointers
+    Q_block_ptr = tl.make_block_ptr(
+        base=Q + q_offset,
+        shape=(M, HEAD_DIM),
+        strides=(stride_qm, stride_qk),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, HEAD_DIM),
+        order=(1, 0),
+    )
+
+    K_block_ptr = tl.make_block_ptr(
+        base=K + k_offset,
+        shape=(HEAD_DIM, N),
+        strides=(stride_kk, stride_kn),
+        offsets=(0, 0),
+        block_shape=(HEAD_DIM, BLOCK_N),
+        order=(0, 1),
+    )
+    O_block_ptr = tl.make_block_ptr(
+        base=Out + o_offset,
+        shape=(M, N),
+        strides=(stride_om, stride_on),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, BLOCK_N),
+        order=(1, 0),
+    )
+
+    # load q: it will stay in SRAM throughout
+    q = tl.load(Q_block_ptr, boundary_check=(0,)) 
+    for start_n in range(0, N, BLOCK_N):
+        start_n = tl.multiple_of(start_n, BLOCK_N)
+        # -- compute qk ----
+        k = tl.load(K_block_ptr, boundary_check=(1,))
+        qk = tl.dot(q, k)
+
+        tl.store(O_block_ptr, qk.to(Out.type.element_ty), boundary_check=(0, 1))
+        
+        K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
+        O_block_ptr = tl.advance(O_block_ptr, (0, BLOCK_N))
+
+
+@triton.jit
+def _attn_bwd_preprocess(
+    O, DO,
+    Delta, # output
+    N_CTX,
+    BLOCK_M: tl.constexpr, 
+    HEAD_DIM: tl.constexpr
+    ):
+    
+    off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
+    off_hz = tl.program_id(1)
+    off_n = tl.arange(0, HEAD_DIM)
+    # load
+    o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :])
+    do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :]).to(tl.float32)
+    delta = tl.sum(o * do, axis=1)
+    # write-back
+    tl.store(Delta + off_hz * N_CTX + off_m, delta)

hymm/models/modules/block_sparse_attention/dynamic_block_attention.py

@@ -0,0 +1,335 @@
+
+from typing import Optional, Tuple
+
+import torch
+
+# NOTE: `bsa_interface` pulls in Triton. We import its symbols lazily (inside the
+# functions that launch kernels) so that the pure-PyTorch geometry helpers in
+# this module remain importable / unit-testable on machines without Triton.
+from .dynamic_block_shape import (
+    ALL_CANDIDATE_SHAPES,
+    ZONE_SIZE,
+    select_zone_shapes,
+)
+
+_TILE = 64  # kernel tile size (Section 8.7: unified BLOCK_N=64)
+
+
+# =====================================================================
+# Per-shape local permutation precompute (constant; cached per device).
+# =====================================================================
+def _precompute_local_orders(device):
+    """For each candidate shape, the [512] permutation mapping a rearranged
+    within-zone position -> local flat coord (lt*64 + lh*8 + lw).
+
+    Also returns, per shape, (product, tiles_per_block). All cached.
+    """
+    cache = _precompute_local_orders._cache
+    if device in cache:
+        return cache[device]
+
+    zs = ZONE_SIZE
+    n = zs * zs * zs  # 512
+    lf = torch.arange(n, device=device)
+    lt = lf // (zs * zs)
+    lh = (lf // zs) % zs
+    lw = lf % zs
+
+    orders = []
+    prods = []
+    tpb = []
+    for (bt, bh, bw) in ALL_CANDIDATE_SHAPES:
+        bsz = bt * bh * bw
+        Nh, Nw = zs // bh, zs // bw
+        nt, it = lt // bt, lt % bt
+        nh, ih = lh // bh, lh % bh
+        nw, iw = lw // bw, lw % bw
+        blk = (nt * Nh + nh) * Nw + nw                 # block index within zone
+        off = (it * bh + ih) * bw + iw                 # offset within block
+        rearr_pos = blk * bsz + off                    # [512]: local_flat -> rearr pos
+        order = torch.empty(n, dtype=torch.long, device=device)
+        order[rearr_pos] = lf                          # invert -> rearr pos -> local_flat
+        orders.append(order)
+        prods.append(bsz)
+        tpb.append(bsz // _TILE)
+
+    local_orders = torch.stack(orders, dim=0)                                   # [C, 512]
+    products = torch.tensor(prods, dtype=torch.long, device=device)            # [C]
+    tiles_per_block = torch.tensor(tpb, dtype=torch.long, device=device)       # [C]
+    cache[device] = (local_orders, products, tiles_per_block)
+    return cache[device]
+
+
+_precompute_local_orders._cache = {}
+
+
+# =====================================================================
+# Pure BSA Triton wrapper (forward + backward through the kernel only).
+# Mirrors bias_rectification._bsa_kernel_only but lives here to keep the
+# dynamic path fully isolated. Selection / pooling stay in PyTorch outside.
+# =====================================================================
+class _dyn_bsa_kernel(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, q, k, v, sm_scale, block_indices, block_indices_lens,
+                chunk_size_q, chunk_size_k, sparsity, kv_valid_mask):
+        from .bsa_interface import attn_fwd_bsa_varlen_triton
+        o, lse = attn_fwd_bsa_varlen_triton(
+            q, k, v, sm_scale, block_indices, block_indices_lens,
+            chunk_size_q, chunk_size_k, sparsity, kv_valid_mask=kv_valid_mask,
+        )
+        ctx.save_for_backward(q, k, v, o, lse, block_indices, block_indices_lens)
+        ctx.sm_scale = sm_scale
+        ctx.chunk_size_q = chunk_size_q
+        ctx.chunk_size_k = chunk_size_k
+        ctx.sparsity = sparsity
+        ctx.kv_valid_mask = kv_valid_mask
+        return o
+
+    @staticmethod
+    def backward(ctx, do):
+        from .bsa_interface import attn_bwd_bsa_varlen_triton
+        q, k, v, o, lse, block_indices, block_indices_lens = ctx.saved_tensors
+        dq = torch.empty_like(q)
+        dk = torch.empty_like(k)
+        dv = torch.empty_like(v)
+        attn_bwd_bsa_varlen_triton(
+            do.contiguous(), q, k, v, o, dq, dk, dv,
+            ctx.sm_scale, lse, block_indices, block_indices_lens,
+            ctx.chunk_size_q, ctx.chunk_size_k, ctx.sparsity,
+            kv_valid_mask=ctx.kv_valid_mask,
+        )
+        return dq, dk, dv, None, None, None, None, None, None, None
+
+
+# =====================================================================
+# Geometry helpers (no grad).
+# =====================================================================
+def _build_rearrange_index(shape_ids, ZT, ZH, ZW, H, W, device):
+    """Return src_idx[S], inv_idx[S] for the per-zone anisotropic rearrange.
+
+    src_idx[rearr_pos] = original_thw_index ; inv_idx = its inverse permutation.
+    Both contiguous int64.
+    """
+    zs = ZONE_SIZE
+    Nz = ZT * ZH * ZW
+    local_orders, _, _ = _precompute_local_orders(device)
+    lo = local_orders[shape_ids]                       # [Nz, 512] local flat per rearr pos
+    lt = lo // (zs * zs)
+    lh = (lo // zs) % zs
+    lw = lo % zs
+
+    zidx = torch.arange(Nz, device=device)
+    zt = zidx // (ZH * ZW)
+    zh = (zidx // ZW) % ZH
+    zw = zidx % ZW
+    gt = zt[:, None] * zs + lt
+    gh = zh[:, None] * zs + lh
+    gw = zw[:, None] * zs + lw
+    global_thw = (gt * H + gh) * W + gw                # [Nz, 512]
+    src_idx = global_thw.reshape(-1).contiguous()      # [S]
+    # src_idx is a permutation, so its inverse is an O(S) scatter (cheaper than
+    # an O(S log S) argsort): inv_idx[src_idx[r]] = r.
+    S = src_idx.shape[0]
+    inv_idx = torch.empty_like(src_idx)
+    inv_idx[src_idx] = torch.arange(S, device=device)
+    return src_idx, inv_idx
+
+
+def _build_block_tile_maps(shape_ids, ZT, ZH, ZW, device):
+    """Logical-block <-> tile bookkeeping for the rearranged sequence.
+
+    Returns:
+      block_id_of_tile [N_tiles]   : logical-block id of each 64-tile
+      kblock_tiles     [N_b+1, 2]  : the (up-to-2) tile ids of each logical block
+                                     (row N_b is the sentinel = N_tiles, used for
+                                     invalid/sentinel selection slots)
+      N_b (int), N_tiles (int)
+    """
+    Nz = ZT * ZH * ZW
+    _, products, tiles_per_block = _precompute_local_orders(device)
+    tpb = tiles_per_block[shape_ids]                   # [Nz] in {1,2}
+    nbz = (512 // products[shape_ids])                 # blocks per zone [Nz] in {8,4}
+    N_tiles = Nz * (512 // _TILE)                      # = Nz * 8
+    N_b = int(nbz.sum().item())
+
+    # exclusive cumsum of blocks-per-zone -> per-zone block offset
+    block_offset = torch.zeros(Nz, dtype=torch.long, device=device)
+    if Nz > 1:
+        block_offset[1:] = torch.cumsum(nbz, dim=0)[:-1]
+
+    tile_g = torch.arange(N_tiles, device=device)
+    z_of_tile = tile_g // 8
+    tl_local = tile_g % 8
+    tpb_t = tpb[z_of_tile]
+    block_local = tl_local // tpb_t
+    pos_in_block = tl_local % tpb_t
+    block_id_of_tile = block_offset[z_of_tile] + block_local   # [N_tiles]
+
+    kblock_tiles = torch.full((N_b + 1, 2), N_tiles, dtype=torch.long, device=device)
+    kblock_tiles[block_id_of_tile, pos_in_block] = tile_g
+    return block_id_of_tile, kblock_tiles, N_b, N_tiles
+
+
+def _logical_block_means(x_tile, tile_counts, block_id_of_tile, N_b):
+    """Aggregate masked tile-means -> logical-block masked means.
+
+    block_mean = sum(tile_mean * tile_valid_count) / sum(tile_valid_count),
+    correctly averaging only valid tokens across the 1-or-2 tiles of a block.
+
+    x_tile: [B,H,N_tiles,D]; tile_counts: [N_tiles]; returns ([B,H,N_b,D], block_count[N_b]).
+    """
+    B, Hn, N_tiles, D = x_tile.shape
+    w = x_tile * tile_counts.view(1, 1, N_tiles, 1).to(x_tile.dtype)
+    block_sum = x_tile.new_zeros(B, Hn, N_b, D)
+    block_sum.index_add_(2, block_id_of_tile, w)
+    block_count = tile_counts.new_zeros(N_b)
+    block_count.index_add_(0, block_id_of_tile, tile_counts.to(block_count.dtype))
+    block_mean = block_sum / block_count.clamp(min=1.0).view(1, 1, N_b, 1).to(x_tile.dtype)
+    return block_mean, block_count
+
+
+# =====================================================================
+# Public entry point.
+# =====================================================================
+def flash_attn_bsa_3d_dynamic(
+    q: torch.Tensor,           # [B, H, S, D] THW order (already padded to mult of ZONE_SIZE)
+    k: torch.Tensor,
+    v: torch.Tensor,
+    grid_padded: Tuple[int, int, int],
+    *,
+    method: str,               # "ivpq" or "penalty"
+    sparsity: Optional[float] = 0.9375,
+    cdf_threshold: Optional[float] = None,
+    audio_token_norms: Optional[torch.Tensor] = None,  # [S] THW order, or None
+    valid_mask: Optional[torch.Tensor] = None,         # [S] bool THW order, or None
+    lambda_a: float = 0.5,
+    tau_128: float = 0.15,
+    lambda_128: float = 1.0,
+    shape_ids: Optional[torch.Tensor] = None,
+    return_shape_ids: bool = False,
+) -> torch.Tensor:
+    """Dynamic-block-shape BSA for video self-attention. Returns [B,H,S,D] THW order.
+
+    shape_ids: if given, skip the (non-deterministic) shape selection and use it
+        verbatim — used to keep the checkpoint recompute identical to the forward.
+    return_shape_ids: if True, also return the shape_ids used (for caching).
+    """
+    from .bsa_interface import (
+        mean_pooling_compression,
+        masked_mean_pooling_compression,
+        cal_score,
+        get_select_indices_topk_from_score,
+        get_select_indices_cdf_from_score,
+        get_select_indices_cdf_topk_from_score,
+    )
+    B, Hn, S, D = q.shape
+    T, H, W = grid_padded
+    assert T * H * W == S
+    assert T % ZONE_SIZE == 0 and H % ZONE_SIZE == 0 and W % ZONE_SIZE == 0
+    device = q.device
+    sm_scale = 1.0 / (D ** 0.5)
+    ZT, ZH, ZW = T // ZONE_SIZE, H // ZONE_SIZE, W // ZONE_SIZE
+
+    # ---- (1) per-zone shape decision (head-agnostic, SP-consistent) ----
+    with torch.no_grad():
+        if shape_ids is None:
+            v_headavg = v[0].mean(dim=0).view(T, H, W, D)          # [T,H,W,D]
+            shape_ids = select_zone_shapes(
+                v_headavg, audio_token_norms, grid_padded, method,
+                lambda_a=lambda_a, tau_128=tau_128, lambda_128=lambda_128,
+                valid_mask_thw=valid_mask, sp_reduce=True,
+            )                                                       # [Nz] int64
+
+        # ---- (2) build the rearrange + tile/block bookkeeping ----
+        src_idx, inv_idx = _build_rearrange_index(shape_ids, ZT, ZH, ZW, H, W, device)
+        block_id_of_tile, kblock_tiles, N_b, N_tiles = _build_block_tile_maps(
+            shape_ids, ZT, ZH, ZW, device
+        )
+
+    # ---- (3) rearrange q/k/v into the tile-contiguous layout (differentiable) ----
+    q_re = q.index_select(2, src_idx).contiguous()
+    k_re = k.index_select(2, src_idx).contiguous()
+    v_re = v.index_select(2, src_idx).contiguous()
+    valid_re = valid_mask[src_idx].contiguous() if valid_mask is not None else None
+
+    # ---- (4) logical-block scoring + selection (no grad — selection only) ----
+    with torch.no_grad():
+        if valid_re is not None:
+            q_tile = masked_mean_pooling_compression(q_re, _TILE, valid_re)
+            k_tile = masked_mean_pooling_compression(k_re, _TILE, valid_re)
+            tile_counts = valid_re.view(N_tiles, _TILE).sum(dim=-1).float()
+        else:
+            q_tile = mean_pooling_compression(q_re, _TILE)
+            k_tile = mean_pooling_compression(k_re, _TILE)
+            tile_counts = torch.full((N_tiles,), float(_TILE), device=device)
+
+        q_cmp, _ = _logical_block_means(q_tile, tile_counts, block_id_of_tile, N_b)
+        k_cmp, block_count = _logical_block_means(k_tile, tile_counts, block_id_of_tile, N_b)
+        block_valid = block_count > 0                          # [N_b]
+        del q_tile, k_tile
+
+        has_invalid = not bool(block_valid.all())
+        invalid_mask = ~block_valid.view(1, 1, 1, -1) if has_invalid else None
+        del block_valid, block_count
+
+        sel_blocks_list = []
+        for _h in range(Hn):
+            score_h = cal_score(q_cmp[:, _h:_h+1], k_cmp[:, _h:_h+1])
+            if invalid_mask is not None:
+                score_h = score_h.masked_fill(invalid_mask, float('-inf'))
+            if sparsity is not None and cdf_threshold is None:
+                sel_h, _ = get_select_indices_topk_from_score(score_h, sparsity)
+            elif sparsity is None and cdf_threshold is not None:
+                sel_h, _ = get_select_indices_cdf_from_score(score_h, cdf_threshold, sm_scale)
+            elif sparsity is not None and cdf_threshold is not None:
+                sel_h, _ = get_select_indices_cdf_topk_from_score(score_h, sparsity, cdf_threshold, sm_scale)
+            else:
+                raise ValueError("Either sparsity or cdf_threshold must be provided")
+            sel_blocks_list.append(sel_h)
+            del score_h, sel_h
+        del q_cmp, k_cmp, invalid_mask
+
+        sel_blocks = torch.cat(sel_blocks_list, dim=1)
+        del sel_blocks_list
+        # sel_blocks values in [0, N_b]; N_b is the sentinel (cdf/varlen padding).
+        sel_blocks = sel_blocks.clamp(max=N_b)                 # safety (int64 index)
+
+        # ---- (5) expand selected logical blocks -> tile-level indices ----
+        kblock_tiles = kblock_tiles.to(torch.int32)
+        t0 = kblock_tiles[:, 0][sel_blocks]                    # [B,H,N_bq,r] int32
+        t1 = kblock_tiles[:, 1][sel_blocks]
+        del sel_blocks, kblock_tiles
+        tile_sel = torch.cat([t0, t1], dim=-1)                 # [B,H,N_bq,2r] int32; values in [0, N_tiles]
+        del t0, t1
+
+        if valid_re is not None:
+            tile_is_valid = tile_counts > 0                    # [N_tiles] bool
+            in_range = tile_sel < N_tiles
+            sel_valid = in_range & tile_is_valid[tile_sel.clamp(max=N_tiles - 1)]
+            tile_sel = torch.where(sel_valid, tile_sel, torch.full_like(tile_sel, N_tiles))
+        del tile_counts
+
+        tile_sel, _ = torch.sort(tile_sel, dim=-1)             # int32 sort; valid first, sentinel(N_tiles) at tail
+        tile_lens_block = (tile_sel < N_tiles).sum(dim=-1).to(torch.int32)
+
+        # map q-block rows -> q-tile rows (a 128 q-block shares its row across its 2 tiles).
+        # tile_sel is already int32, so index_select yields the int32 block indices directly.
+        tile_block_indices = tile_sel.index_select(2, block_id_of_tile).contiguous()
+        tile_block_indices_lens = tile_lens_block.index_select(2, block_id_of_tile).contiguous()
+        del tile_sel, tile_lens_block, block_id_of_tile
+
+    # ---- (6) BSA Triton kernel at chunk=64 (autograd through q_re/k_re/v_re) ----
+    o_re = _dyn_bsa_kernel.apply(
+        q_re, k_re, v_re, sm_scale,
+        tile_block_indices, tile_block_indices_lens,
+        _TILE, _TILE, sparsity, valid_re,
+    )
+
+    # ---- (7) scatter back to original THW order ----
+    out = o_re.index_select(2, inv_idx).contiguous()
+    if return_shape_ids:
+        return out, shape_ids
+    return out
+
+
+__all__ = ["flash_attn_bsa_3d_dynamic"]

hymm/models/modules/block_sparse_attention/dynamic_block_shape.py

@@ -0,0 +1,406 @@
+
+from typing import List, Optional, Tuple
+
+import torch
+
+CANDIDATE_SHAPES_64: List[Tuple[int, int, int]] = [
+    (4, 4, 4),  # isotropic, default
+    (8, 2, 4),  # long-T, narrow-H (lip/face)
+    (8, 4, 2),  # long-T, narrow-W
+    (2, 4, 8),  # short-T, wide-W (fast motion)
+    (2, 8, 4),  # short-T, wide-H
+    (4, 2, 8),  # mid-T, narrow-H wide-W
+    (4, 8, 2),  # mid-T, wide-H narrow-W
+]
+
+CANDIDATE_SHAPES_128: List[Tuple[int, int, int]] = [
+    (2, 8, 8),  # thin-T, large spatial
+    (8, 2, 8),  # long-T, narrow-H wide-W
+    (8, 8, 2),  # long-T, wide-H narrow-W
+    (4, 4, 8),  # mid-T, wide-W
+    (4, 8, 4),  # mid-T, wide-H
+    (8, 4, 4),  # long-T, mid spatial (static sky)
+]
+
+ALL_CANDIDATE_SHAPES: List[Tuple[int, int, int]] = (
+    CANDIDATE_SHAPES_64 + CANDIDATE_SHAPES_128
+)
+NUM_CANDIDATES_64 = len(CANDIDATE_SHAPES_64)
+NUM_CANDIDATES = len(ALL_CANDIDATE_SHAPES)
+
+ZONE_SIZE = 8  # macro-zone edge length (Section 8.4.1); lcm(2,4,8)
+
+_EPS = 1e-6
+
+
+def candidate_tensors(device, dtype=torch.float32):
+    """Return (shapes[NUM_CANDIDATES,3], log2_shapes[NUM_CANDIDATES,3],
+    is_128[NUM_CANDIDATES], product[NUM_CANDIDATES]) on `device`.
+
+    Cached per (device, dtype) to avoid re-allocation each layer/step.
+    """
+    key = (device, dtype)
+    cache = candidate_tensors._cache
+    if key not in cache:
+        shapes = torch.tensor(ALL_CANDIDATE_SHAPES, device=device, dtype=dtype)  # [C,3]
+        log2_shapes = torch.log2(shapes)
+        prod = shapes.prod(dim=-1)
+        is_128 = prod > 64.5
+        cache[key] = (shapes, log2_shapes, is_128, prod)
+    return cache[key]
+
+
+candidate_tensors._cache = {}
+
+def _zone_reshape(x_thw_d: torch.Tensor, ZT: int, ZH: int, ZW: int) -> torch.Tensor:
+    """[T,H,W,D] -> [num_zones, ZS,ZS,ZS, D] with zone order zt-major.
+
+    Memory-contiguous output (one permute+contiguous). T,H,W must be multiples
+    of ZONE_SIZE.
+    """
+    T, H, W, D = x_thw_d.shape
+    zs = ZONE_SIZE
+    x = x_thw_d.reshape(ZT, zs, ZH, zs, ZW, zs, D)
+    x = x.permute(0, 2, 4, 1, 3, 5, 6).contiguous()  # [ZT,ZH,ZW, zs,zs,zs, D]
+    return x.view(ZT * ZH * ZW, zs, zs, zs, D)
+
+
+def compute_zone_aggregates(
+    v_headavg_thw_d: torch.Tensor,
+    ZT: int, ZH: int, ZW: int,
+    valid_mask_thw: Optional[torch.Tensor] = None,
+):
+    """Compute the per-zone *masked mean* aggregates that are linear in V.
+
+    Padded tokens (valid_mask == False) are excluded from every mean. Because
+    the valid mask is identical on every SP rank (it is a token-level grid
+    mask, not a head shard), the per-zone valid token *counts* are
+    rank-independent — so a masked-mean computed per rank and then AVG-all-
+    reduced across the SP group yields the exact global (head- and token-)
+    masked mean. This keeps the SP reduction a single cheap collective.
+
+    Returns (packed[num_zones, D + 3*ZS*D], frame_valid[num_zones, ZS]).
+    Packed last-dim layout:
+        [ 0        : D        ]  channel_means  (masked mean over all tokens)
+        [ D        : D + ZS*D ]  frame_means_T  (masked mean over H,W per t)
+        [ D + ZS*D : D + 2ZS*D]  axis_means_H   (masked mean over T,W per h)
+        [ D +2ZS*D : D + 3ZS*D]  axis_means_W   (masked mean over T,H per w)
+    `frame_valid[z,t]` marks t-frames that contain >=1 valid token (used to
+    exclude padded frames from the temporal-variance diff downstream).
+    """
+    zs = ZONE_SIZE
+    D = v_headavg_thw_d.shape[-1]
+    xz = _zone_reshape(v_headavg_thw_d, ZT, ZH, ZW).float()  # [Nz, t,h,w, D]
+    Nz = xz.shape[0]
+
+    if valid_mask_thw is not None:
+        mz = _zone_reshape(valid_mask_thw.reshape(*v_headavg_thw_d.shape[:3], 1).float(),
+                           ZT, ZH, ZW)                       # [Nz, t,h,w, 1]
+        c_all = mz.sum(dim=(1, 2, 3)).clamp(min=1.0)         # [Nz,1]
+        channel_means = (xz * mz).sum(dim=(1, 2, 3)) / c_all
+        cT = mz.sum(dim=(2, 3)).clamp(min=1.0)               # [Nz, t, 1]
+        frame_means_T = (xz * mz).sum(dim=(2, 3)) / cT
+        cH = mz.sum(dim=(1, 3)).clamp(min=1.0)
+        axis_means_H = (xz * mz).sum(dim=(1, 3)) / cH
+        cW = mz.sum(dim=(1, 2)).clamp(min=1.0)
+        axis_means_W = (xz * mz).sum(dim=(1, 2)) / cW
+        m = mz.squeeze(-1)
+        t_valid = m.sum(dim=(2, 3)) > 0                      # [Nz, t]
+        h_valid = m.sum(dim=(1, 3)) > 0                      # [Nz, h]
+        w_valid = m.sum(dim=(1, 2)) > 0                      # [Nz, w]
+        axis_valid = (t_valid, h_valid, w_valid)
+    else:
+        channel_means = xz.mean(dim=(1, 2, 3))               # [Nz, D]
+        frame_means_T = xz.mean(dim=(2, 3))                  # [Nz, t, D]
+        axis_means_H = xz.mean(dim=(1, 3))                   # [Nz, h, D]
+        axis_means_W = xz.mean(dim=(1, 2))                   # [Nz, w, D]
+        ones = torch.ones(Nz, zs, dtype=torch.bool, device=xz.device)
+        axis_valid = (ones, ones, ones)
+
+    packed = torch.cat(
+        [
+            channel_means,
+            frame_means_T.reshape(Nz, zs * D),
+            axis_means_H.reshape(Nz, zs * D),
+            axis_means_W.reshape(Nz, zs * D),
+        ],
+        dim=-1,
+    ).contiguous()
+    return packed, axis_valid
+
+
+def _maybe_sp_average(packed: torch.Tensor) -> torch.Tensor:
+    """AVG-all-reduce the packed aggregates across the SP group (no-op if SP off).
+
+    Mirrors the FSDP/SP convention used by bias_rectification: only the SP path
+    splits heads across ranks, so the head-agnostic shape decision needs this
+    reduction to be globally consistent. FSDP-only (no SP) sees identical data
+    on all ranks already → no collective.
+    """
+    # Not wrapped in try/except on purpose: a swallowed import error would let one
+    # rank return early while the rest block inside the all_reduce, turning a
+    # loud failure into a silent NCCL hang.
+    import torch.distributed as dist
+    from hymm.utils.parallel_states import nccl_info, get_sequence_parallel_state
+
+    if (not get_sequence_parallel_state()) or nccl_info.sp_size <= 1:
+        return packed
+    out = packed.contiguous()
+    dist.all_reduce(out, op=dist.ReduceOp.SUM, group=nccl_info.sp_group)
+    out /= float(nccl_info.sp_size)
+    return out
+
+
+def _unpack_aggregates(packed: torch.Tensor, D: int):
+    zs = ZONE_SIZE
+    Nz = packed.shape[0]
+    channel_means = packed[:, :D]
+    off = D
+    frame_means_T = packed[:, off:off + zs * D].reshape(Nz, zs, D); off += zs * D
+    axis_means_H = packed[:, off:off + zs * D].reshape(Nz, zs, D); off += zs * D
+    axis_means_W = packed[:, off:off + zs * D].reshape(Nz, zs, D); off += zs * D
+    return channel_means, frame_means_T, axis_means_H, axis_means_W
+
+
+def _axis_var(axis_means: torch.Tensor, valid: Optional[torch.Tensor] = None) -> torch.Tensor:
+    """Variance across the ZS axis positions, averaged over channels. [Nz,ZS,D]->[Nz].
+
+    If `valid` [Nz,ZS] is given, only valid positions contribute (padded
+    axis positions are excluded so they cannot fabricate variance).
+    """
+    if valid is None:
+        var_per_channel = axis_means.var(dim=1, unbiased=False)   # [Nz, D]
+        return var_per_channel.mean(dim=-1)                       # [Nz]
+    vf = valid.unsqueeze(-1).to(axis_means.dtype)                 # [Nz,ZS,1]
+    n = vf.sum(dim=1).clamp(min=1.0)                              # [Nz,1]
+    mean = (axis_means * vf).sum(dim=1, keepdim=True) / n.unsqueeze(1)
+    var_per_channel = ((axis_means - mean) ** 2 * vf).sum(dim=1) / n
+    return var_per_channel.mean(dim=-1)                           # [Nz]
+
+
+def _masked_var_1d(vals: torch.Tensor, valid: torch.Tensor) -> torch.Tensor:
+    """Variance over the last axis restricted to valid positions. [Nz,L]->[Nz]."""
+    vf = valid.to(vals.dtype)
+    n = vf.sum(dim=-1).clamp(min=1.0)
+    mean = (vals * vf).sum(dim=-1) / n
+    return ((vals - mean.unsqueeze(-1)) ** 2 * vf).sum(dim=-1) / n
+
+
+def compute_audio_directional(
+    audio_norms_thw: Optional[torch.Tensor],
+    ZT: int, ZH: int, ZW: int,
+    valid_mask_thw: Optional[torch.Tensor] = None,
+    device: Optional[torch.device] = None,
+):
+    """Per-zone audio descriptor (Section 8.4.3).
+
+    Returns (a_bar[Nz], v_hat_T[Nz], v_hat_H[Nz], v_hat_W[Nz]) where a_bar is
+    the (max-normalized) mean audio-absorption and v_hat_d is the normalized
+    directional variance along axis d. All in float32.
+
+    `audio_norms_thw` is identical on every SP rank (all-gathered upstream), so
+    no collective is needed here.
+    """
+    zs = ZONE_SIZE
+    Nz = ZT * ZH * ZW
+    if device is None:
+        device = (audio_norms_thw.device if audio_norms_thw is not None
+                  else (valid_mask_thw.device if valid_mask_thw is not None else 'cpu'))
+    if audio_norms_thw is None:
+        z = torch.zeros(Nz, device=device, dtype=torch.float32)
+        return z, z.clone(), z.clone(), z.clone()
+
+    a = audio_norms_thw.reshape(ZT, zs, ZH, zs, ZW, zs).float()
+    a = a.permute(0, 2, 4, 1, 3, 5).reshape(Nz, zs, zs, zs)  # [Nz,t,h,w]
+
+    if valid_mask_thw is not None:
+        m = valid_mask_thw.reshape(ZT, zs, ZH, zs, ZW, zs).float()
+        m = m.permute(0, 2, 4, 1, 3, 5).reshape(Nz, zs, zs, zs)
+        cnt = m.sum(dim=(1, 2, 3)).clamp(min=1.0)
+        a_bar = (a * m).sum(dim=(1, 2, 3)) / cnt
+        # Per-axis masked mean over the other two axes, then masked variance over
+        # the axis. Padded positions are excluded so they cannot fabricate audio
+        # directional structure (mirrors the V-feature masking).
+        cT = m.sum(dim=(2, 3)).clamp(min=1.0); aT = (a * m).sum(dim=(2, 3)) / cT  # [Nz,t]
+        cH = m.sum(dim=(1, 3)).clamp(min=1.0); aH = (a * m).sum(dim=(1, 3)) / cH
+        cW = m.sum(dim=(1, 2)).clamp(min=1.0); aW = (a * m).sum(dim=(1, 2)) / cW
+        t_valid = m.sum(dim=(2, 3)) > 0
+        h_valid = m.sum(dim=(1, 3)) > 0
+        w_valid = m.sum(dim=(1, 2)) > 0
+        var_T = _masked_var_1d(aT, t_valid)
+        var_H = _masked_var_1d(aH, h_valid)
+        var_W = _masked_var_1d(aW, w_valid)
+    else:
+        a_bar = a.mean(dim=(1, 2, 3))
+        var_T = a.mean(dim=(2, 3)).var(dim=1, unbiased=False)  # [Nz]
+        var_H = a.mean(dim=(1, 3)).var(dim=1, unbiased=False)
+        var_W = a.mean(dim=(1, 2)).var(dim=1, unbiased=False)
+
+    denom = (var_T + var_H + var_W).clamp(min=_EPS)
+    v_hat_T = var_T / denom
+    v_hat_H = var_H / denom
+    v_hat_W = var_W / denom
+
+    a_bar_norm = a_bar / a_bar.max().clamp(min=_EPS)
+    return a_bar_norm, v_hat_T, v_hat_H, v_hat_W
+
+
+def compute_effective_gradient(
+    v_headavg_thw_d: torch.Tensor,
+    audio_norms_thw: Optional[torch.Tensor],
+    ZT: int, ZH: int, ZW: int,
+    lambda_a: float = 0.5,
+    valid_mask_thw: Optional[torch.Tensor] = None,
+    sp_reduce: bool = True,
+):
+    """Compute (g[Nz,3], rho[Nz]) — Section 8.4.4 effective gradient + density.
+
+    g_d = aniso_ratio_d  +  lambda_a * a_bar_z * v_hat_d^audio        (d in T,H,W)
+    rho = normalized (spatial + temporal) channel-variance information density.
+    """
+    D = v_headavg_thw_d.shape[-1]
+    packed, axis_valid = compute_zone_aggregates(
+        v_headavg_thw_d, ZT, ZH, ZW, valid_mask_thw=valid_mask_thw
+    )
+    if sp_reduce:
+        packed = _maybe_sp_average(packed)
+    t_valid, h_valid, w_valid = axis_valid
+    channel_means, frame_means_T, axis_means_H, axis_means_W = _unpack_aggregates(packed, D)
+
+    # --- channel-variance information density rho (Section 3.3 @ zone level) ---
+    global_cm = channel_means.mean(dim=-1, keepdim=True)
+    spatial_var = ((channel_means - global_cm) ** 2).mean(dim=-1)            # [Nz]
+    # Temporal diff only between adjacent frames that are BOTH valid (mirrors
+    # compute_channel_variance_density) so padded frames cannot fabricate motion.
+    both_valid = (t_valid[:, 1:] & t_valid[:, :-1]).unsqueeze(-1).to(frame_means_T.dtype)  # [Nz,ZS-1,1]
+    n_pairs = both_valid.sum(dim=1).clamp(min=1.0)                           # [Nz,1]
+    frame_diffs = frame_means_T[:, 1:, :] - frame_means_T[:, :-1, :]         # [Nz,ZS-1,D]
+    temporal_var = ((frame_diffs.pow(2) * both_valid).sum(dim=1) / n_pairs).mean(dim=-1)  # [Nz]
+    rho_raw = spatial_var + temporal_var
+    rho = rho_raw / rho_raw.max().clamp(min=_EPS)                            # [Nz] in [0,1]
+
+    # --- anisotropy ratios (Section 8.4.2), padded axis positions excluded ---
+    var_T = _axis_var(frame_means_T, t_valid)
+    var_H = _axis_var(axis_means_H, h_valid)
+    var_W = _axis_var(axis_means_W, w_valid)
+    asum = (var_T + var_H + var_W).clamp(min=_EPS)
+    aniso_T = var_T / asum
+    aniso_H = var_H / asum
+    aniso_W = var_W / asum
+
+    # --- audio directional guidance (Section 8.4.3) ---
+    a_bar, v_hat_T, v_hat_H, v_hat_W = compute_audio_directional(
+        audio_norms_thw, ZT, ZH, ZW, valid_mask_thw,
+        device=v_headavg_thw_d.device,
+    )
+
+    g_T = aniso_T + lambda_a * a_bar * v_hat_T
+    g_H = aniso_H + lambda_a * a_bar * v_hat_H
+    g_W = aniso_W + lambda_a * a_bar * v_hat_W
+    g = torch.stack([g_T, g_H, g_W], dim=-1).clamp(min=1e-4)  # [Nz,3]
+    return g, rho
+
+def select_shapes_ivpq(
+    g: torch.Tensor,         # [Nz, 3]
+    rho: torch.Tensor,       # [Nz]
+    tau_128: float = 0.15,
+) -> torch.Tensor:
+    """IVPQ (Section 8.5): b_d* ∝ g_d^{-2/3}; nearest candidate in log2-space.
+
+    The candidate pool gating (rho vs tau_128) restricts high-info zones to the
+    7 64-token shapes and lets low-info zones also consider the 6 128-token
+    shapes (Section 8.5.3). Returns shape_id [Nz] (index into ALL_CANDIDATE_SHAPES).
+    """
+    device = g.device
+    _, log2_shapes, is_128, _ = candidate_tensors(device, dtype=g.dtype)  # [C,3], [C]
+
+    log2_ratio = (-2.0 / 3.0) * torch.log2(g)               # [Nz, 3], ratio only
+    log2_ideal = log2_ratio + (6.0 - log2_ratio.sum(dim=-1, keepdim=True)) / 3.0
+
+    # log-space squared distance to each candidate: [Nz, C]
+    dist = ((log2_shapes.unsqueeze(0) - log2_ideal.unsqueeze(1)) ** 2).sum(dim=-1)
+
+    # Pool gating: forbid 128-token candidates for high-info zones (rho >= tau_128).
+    forbid = (rho.unsqueeze(1) >= tau_128) & is_128.unsqueeze(0)  # [Nz, C]
+    dist = dist.masked_fill(forbid, float('inf'))
+
+    return dist.argmin(dim=-1).to(torch.int64)
+
+
+def select_shapes_penalty(
+    g: torch.Tensor,         # [Nz, 3]
+    rho: torch.Tensor,       # [Nz]
+    tau_128: float = 0.15,
+    lambda_128: float = 1.0,
+) -> torch.Tensor:
+    """Penalty Matching (Section 8.6): argmin_c sum_d g_d * b_d (+128 bonus).
+
+    128-token candidates receive a density bonus ``-lambda_128 * (1 - rho)`` so
+    low-info zones can fairly select them; high-info zones (rho >= tau_128) are
+    still restricted to the 64-token pool. Returns shape_id [Nz].
+    """
+    device = g.device
+    shapes, _, is_128, _ = candidate_tensors(device, dtype=g.dtype)  # [C,3], [C]
+
+    penalty = (g.unsqueeze(1) * shapes.unsqueeze(0)).sum(dim=-1)     # [Nz, C]
+
+    # 128-token density bonus (lowers penalty for low-info zones).
+    bonus = (lambda_128 * (1.0 - rho)).unsqueeze(1) * is_128.unsqueeze(0).to(g.dtype)
+    penalty = penalty - bonus
+
+    # Pool gating identical to IVPQ.
+    forbid = (rho.unsqueeze(1) >= tau_128) & is_128.unsqueeze(0)
+    penalty = penalty.masked_fill(forbid, float('inf'))
+
+    return penalty.argmin(dim=-1).to(torch.int64)
+
+
+def select_zone_shapes(
+    v_headavg_thw_d: torch.Tensor,
+    audio_norms_thw: Optional[torch.Tensor],
+    grid_padded: Tuple[int, int, int],
+    method: str,
+    lambda_a: float = 0.5,
+    tau_128: float = 0.15,
+    lambda_128: float = 1.0,
+    valid_mask_thw: Optional[torch.Tensor] = None,
+    sp_reduce: bool = True,
+) -> torch.Tensor:
+    """End-to-end: features -> g_d -> shape_id per zone. Returns int64 [num_zones].
+
+    Args:
+        v_headavg_thw_d: [T,H,W,D] head-averaged V (already padded). float ok.
+        audio_norms_thw: [T*H*W] per-token L2 norm of A->V bridge residual, or None.
+        grid_padded: (T,H,W), each a multiple of ZONE_SIZE.
+        method: "ivpq" or "penalty".
+    """
+    T, H, W = grid_padded
+    assert T % ZONE_SIZE == 0 and H % ZONE_SIZE == 0 and W % ZONE_SIZE == 0, (
+        f"grid {grid_padded} must be multiple of zone size {ZONE_SIZE}"
+    )
+    ZT, ZH, ZW = T // ZONE_SIZE, H // ZONE_SIZE, W // ZONE_SIZE
+
+    g, rho = compute_effective_gradient(
+        v_headavg_thw_d, audio_norms_thw, ZT, ZH, ZW,
+        lambda_a=lambda_a, valid_mask_thw=valid_mask_thw, sp_reduce=sp_reduce,
+    )
+    if method == "ivpq":
+        return select_shapes_ivpq(g, rho, tau_128=tau_128)
+    elif method == "penalty":
+        return select_shapes_penalty(g, rho, tau_128=tau_128, lambda_128=lambda_128)
+    raise ValueError(f"unknown dynamic-block method: {method!r}")
+
+
+__all__ = [
+    "CANDIDATE_SHAPES_64",
+    "CANDIDATE_SHAPES_128",
+    "ALL_CANDIDATE_SHAPES",
+    "NUM_CANDIDATES_64",
+    "NUM_CANDIDATES",
+    "ZONE_SIZE",
+    "candidate_tensors",
+    "compute_effective_gradient",
+    "select_shapes_ivpq",
+    "select_shapes_penalty",
+    "select_zone_shapes",
+]

hymm/models/modules/block_sparse_attention/flash_attn_bsa_varlen_mask.py

@@ -0,0 +1,1014 @@
+import triton
+import triton.language as tl
+import os
+
+from .common import autotune
+"""
+TRITON_REEVALUATE_KEY=1
+- autotune whenever params in reevaluate keys change 
+- use in benchmark script to fine the best config
+
+TRITON_AUTOTUNE_ENBALE=1
+- if set to 0, autotune will not work, and the related params must be passed to the function call.
+"""
+
+configs_fwd_bsa_varlen_preset = {
+    'default': {
+        'BLOCK_N': 64,
+        'num_stages': 3,
+        'num_warps': 8,
+    },
+    'BLOCK_N_LG=64': {
+        'BLOCK_N': 64,
+        'num_stages': 3,
+        'num_warps': 4,
+    },
+}
+configs_fwd_bsa_varlen = [
+    triton.Config({'BLOCK_N': BN}, num_stages=s, num_warps=w) \
+    for BN in [32, 64, 128] \
+    for s in [2, 3, 4, 5] \
+    for w in [4, 8] \
+]
+
+fwd_bsa_reevaluate_varlen_keys = ['N_CTX', 'BLOCK_M', 'BLOCK_N_LG', 'SPARSITY'] if os.environ.get('TRITON_REEVALUATE_KEY', '0') == '1' else []
+@autotune(list(configs_fwd_bsa_varlen), key=fwd_bsa_reevaluate_varlen_keys)
+@triton.jit
+def _attn_fwd_bsa_varlen(
+    Q, K, V, sm_scale, M, Out, 
+    block_indices, # [B, H, M_COMPRESS, S_MAX]
+    block_indices_lens, # [B, H, M_COMPRESS]
+    kv_valid_mask, # [N_CTX] bool, True=valid (shared across B,H)
+    stride_qz, stride_qh, stride_qm, stride_qk,
+    stride_kz, stride_kh, stride_kn, stride_kk,
+    stride_vz, stride_vh, stride_vn, stride_vk,
+    stride_oz, stride_oh, stride_om, stride_ok,
+    stride_bz, stride_bh, stride_bm, stride_bs,
+    stride_lz, stride_lh, stride_lm,
+    H, N_CTX,
+    HEAD_DIM: tl.constexpr, 
+    BLOCK_M: tl.constexpr, 
+    BLOCK_N_LG: tl.constexpr, 
+    BLOCK_N: tl.constexpr, 
+    SPARSITY: tl.constexpr,
+    HAS_KV_MASK: tl.constexpr,
+    ):
+    
+    start_m = tl.program_id(0)
+    off_hz = tl.program_id(1)
+    off_z = off_hz // H
+    off_h = off_hz % H
+    
+    q_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh
+    k_offset = off_z.to(tl.int64) * stride_kz + off_h.to(tl.int64) * stride_kh
+    v_offset = off_z.to(tl.int64) * stride_vz + off_h.to(tl.int64) * stride_vh
+    o_offset = off_z.to(tl.int64) * stride_oz + off_h.to(tl.int64) * stride_oh
+    b_offset = off_z.to(tl.int64) * stride_bz + off_h.to(tl.int64) * stride_bh
+    l_offset = off_z.to(tl.int64) * stride_lz + off_h.to(tl.int64) * stride_lh
+
+    # block pointers
+    Q_block_ptr = tl.make_block_ptr(
+        base=Q + q_offset,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_qm, stride_qk),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, HEAD_DIM),
+        order=(1, 0),
+    )
+    V_block_ptr = tl.make_block_ptr(
+        base=V + v_offset,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_vn, stride_vk),
+        offsets=(0, 0),
+        block_shape=(BLOCK_N, HEAD_DIM),
+        order=(1, 0),
+    )
+    KT_block_ptr = tl.make_block_ptr(
+        base=K + k_offset,
+        shape=(HEAD_DIM, N_CTX),
+        strides=(stride_kk, stride_kn),
+        offsets=(0, 0),
+        block_shape=(HEAD_DIM, BLOCK_N),
+        order=(0, 1),
+    )
+    O_block_ptr = tl.make_block_ptr(
+        base=Out + o_offset,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_om, stride_ok),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, HEAD_DIM),
+        order=(1, 0),
+    )
+    block_indices += b_offset + start_m * stride_bm
+    block_indices_lens += l_offset + start_m * stride_lm
+    # initialize offsets
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+
+    # initialize pointer to m and l
+    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
+    acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
+    # load scales
+    qk_scale = sm_scale
+    qk_scale *= 1.44269504  # 1/ln2
+    # load q: it will stay in SRAM throughout
+    q = tl.load(Q_block_ptr)
+    S = tl.load(block_indices_lens)
+    for i in range(S):
+        block_id = tl.load(block_indices + i * stride_bs).to(tl.int32)
+        lo, hi = block_id * BLOCK_N_LG, (block_id + 1) * BLOCK_N_LG
+        lo = tl.multiple_of(lo, BLOCK_N)
+        KT_block_ptr_i = tl.advance(KT_block_ptr, (0, lo))
+        V_block_ptr_i = tl.advance(V_block_ptr, (lo, 0))
+        mask_offset = lo
+
+        # loop over k, v and update accumulator
+        for start_n in range(lo, hi, BLOCK_N):
+            start_n = tl.multiple_of(start_n, BLOCK_N)
+            # -- compute qk ----
+            kT = tl.load(KT_block_ptr_i)
+            qkT = tl.dot(q, kT)
+
+            if HAS_KV_MASK:
+                offs_n = mask_offset + tl.arange(0, BLOCK_N)
+                k_valid = tl.load(kv_valid_mask + offs_n)
+                qkT = tl.where(k_valid[None, :], qkT, float('-inf'))
+
+            m_ij = tl.maximum(m_i, tl.max(qkT, 1) * qk_scale)
+            qkT = qkT * qk_scale - m_ij[:, None]
+            p = tl.math.exp2(qkT)
+
+            if HAS_KV_MASK:
+                p = tl.where(k_valid[None, :], p, 0.0)
+            
+            # -- update m_i and l_i
+            # Guard the all-(-inf) block case: if a selected block is fully padded
+            # (every key masked) AND it is the first processed block, m_i and m_ij
+            # are both -inf, so (m_i - m_ij) = -inf-(-inf) = NaN. Such a block
+            # contributes nothing (p is all-zero via the kv-mask above), so alpha
+            # must be 1.0 to leave acc / l_i untouched (alpha=0 would zero l_i and
+            # produce 0/0 = NaN in the epilogue for an all-padding-selection query).
+            # When any selected block has >=1 valid key (the common case) m_ij is
+            # finite and this is numerically identical to the original.
+            alpha = tl.where(m_ij == float("-inf"), 1.0, tl.math.exp2(m_i - m_ij))
+            l_ij = tl.sum(p, 1)
+            # -- update output accumulator --
+            acc = acc * alpha[:, None]
+            # update acc
+            v = tl.load(V_block_ptr_i)
+            acc = tl.dot(p.to(v.dtype), v, acc)
+            # update m_i and l_i
+            l_i = l_i * alpha + l_ij
+            m_i = m_ij
+            V_block_ptr_i = tl.advance(V_block_ptr_i, (BLOCK_N, 0))
+            KT_block_ptr_i = tl.advance(KT_block_ptr_i, (0, BLOCK_N))
+            mask_offset += BLOCK_N
+    
+    # epilogue
+    m_i += tl.math.log2(l_i)
+    acc = acc / l_i[:, None]
+    m_ptrs = M + off_hz * N_CTX + offs_m
+    tl.store(m_ptrs, m_i)
+    tl.store(O_block_ptr, acc.to(Out.type.element_ty))
+
+configs_fwd_bsa_varlen_align_preset = {
+    'default': {
+        'num_stages': 3,
+        'num_warps': 8,
+    },
+    'BLOCK_N_LG=64': {
+        'num_stages': 3,
+        'num_warps': 4,
+    },
+}
+configs_fwd_bsa_varlen_align = [
+    triton.Config({}, num_stages=s, num_warps=w) \
+    for s in [2, 3, 4, 5] \
+    for w in [4, 8] \
+]
+
+fwd_bsa_reevaluate_varlen_align_keys = ['N_CTX', 'BLOCK_M', 'BLOCK_N_LG', 'SPARSITY'] if os.environ.get('TRITON_REEVALUATE_KEY', '0') == '1' else []
+@autotune(list(configs_fwd_bsa_varlen_align), key=fwd_bsa_reevaluate_varlen_align_keys)
+@triton.jit
+def _attn_fwd_bsa_varlen_align(
+    Q, K, V, sm_scale, M, Out, 
+    block_indices, # [B, H, M_COMPRESS, S_MAX]
+    block_indices_lens, # [B, H, M_COMPRESS]
+    kv_valid_mask, # [N_CTX] bool
+    stride_qz, stride_qh, stride_qm, stride_qk,
+    stride_kz, stride_kh, stride_kn, stride_kk,
+    stride_vz, stride_vh, stride_vn, stride_vk,
+    stride_oz, stride_oh, stride_om, stride_on,
+    stride_bz, stride_bh, stride_bm, stride_bs,
+    stride_lz, stride_lh, stride_lm,
+    H, N_CTX,
+    HEAD_DIM: tl.constexpr, 
+    BLOCK_M: tl.constexpr, 
+    BLOCK_N_LG: tl.constexpr, 
+    SPARSITY: tl.constexpr,
+    HAS_KV_MASK: tl.constexpr,
+    ):
+    
+    start_m = tl.program_id(0)
+    off_hz = tl.program_id(1)
+    off_z = off_hz // H
+    off_h = off_hz % H
+    
+    q_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh
+    k_offset = off_z.to(tl.int64) * stride_kz + off_h.to(tl.int64) * stride_kh
+    v_offset = off_z.to(tl.int64) * stride_vz + off_h.to(tl.int64) * stride_vh
+    o_offset = off_z.to(tl.int64) * stride_oz + off_h.to(tl.int64) * stride_oh
+    b_offset = off_z.to(tl.int64) * stride_bz + off_h.to(tl.int64) * stride_bh
+    l_offset = off_z.to(tl.int64) * stride_lz + off_h.to(tl.int64) * stride_lh
+
+    # block pointers
+    Q_block_ptr = tl.make_block_ptr(
+        base=Q + q_offset,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_qm, stride_qk),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, HEAD_DIM),
+        order=(1, 0),
+    )
+    V_block_ptr = tl.make_block_ptr(
+        base=V + v_offset,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_vn, stride_vk),
+        offsets=(0, 0),
+        block_shape=(BLOCK_N_LG, HEAD_DIM),
+        order=(1, 0),
+    )
+    KT_block_ptr = tl.make_block_ptr(
+        base=K + k_offset,
+        shape=(HEAD_DIM, N_CTX),
+        strides=(stride_kk, stride_kn),
+        offsets=(0, 0),
+        block_shape=(HEAD_DIM, BLOCK_N_LG),
+        order=(0, 1),
+    )
+    O_block_ptr = tl.make_block_ptr(
+        base=Out + o_offset,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_om, stride_on),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, HEAD_DIM),
+        order=(1, 0),
+    )
+    block_indices += b_offset + start_m * stride_bm
+    block_indices_lens += l_offset + start_m * stride_lm
+    # initialize offsets
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+
+    # initialize pointer to m and l
+    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
+    acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
+    # load scales
+    qk_scale = sm_scale
+    qk_scale *= 1.44269504  # 1/ln2
+    # load q: it will stay in SRAM throughout
+    q = tl.load(Q_block_ptr)
+    S = tl.load(block_indices_lens)
+    for i in range(S):
+        block_id = tl.load(block_indices + i * stride_bs).to(tl.int32)
+        lo = block_id * BLOCK_N_LG
+        lo = tl.multiple_of(lo, BLOCK_N_LG)
+        KT_block_ptr_i = tl.advance(KT_block_ptr, (0, lo))
+        V_block_ptr_i = tl.advance(V_block_ptr, (lo, 0))
+
+        # -- compute qk ----
+        kT = tl.load(KT_block_ptr_i)
+        qkT = tl.dot(q, kT)
+
+        if HAS_KV_MASK:
+            offs_n = lo + tl.arange(0, BLOCK_N_LG)
+            k_valid = tl.load(kv_valid_mask + offs_n)
+            qkT = tl.where(k_valid[None, :], qkT, float('-inf'))
+
+        m_ij = tl.maximum(m_i, tl.max(qkT, 1) * qk_scale)
+        qkT = qkT * qk_scale - m_ij[:, None]
+        p = tl.math.exp2(qkT)
+
+        if HAS_KV_MASK:
+            p = tl.where(k_valid[None, :], p, 0.0)
+        
+        # -- update m_i and l_i
+        # Guard the all-(-inf) block case: a fully-padded selected block processed
+        # first leaves m_i == m_ij == -inf, so (m_i - m_ij) = NaN. Such a block
+        # contributes nothing (p is all-zero via the kv-mask above); alpha must be
+        # 1.0 to leave acc / l_i untouched (alpha=0 would zero l_i -> 0/0 = NaN in
+        # the epilogue for an all-padding-selection query). Identical to the
+        # original whenever any selected block has >=1 valid key (the common case).
+        alpha = tl.where(m_ij == float("-inf"), 1.0, tl.math.exp2(m_i - m_ij))
+        l_ij = tl.sum(p, 1)
+        # -- update output accumulator --
+        acc = acc * alpha[:, None]
+        # update acc
+        v = tl.load(V_block_ptr_i)
+        acc = tl.dot(p.to(v.dtype), v, acc)
+        # update m_i and l_i
+        l_i = l_i * alpha + l_ij
+        m_i = m_ij
+
+    
+    # epilogue
+    m_i += tl.math.log2(l_i)
+    acc = acc / l_i[:, None]
+    m_ptrs = M + off_hz * N_CTX + offs_m
+    tl.store(m_ptrs, m_i)
+    tl.store(O_block_ptr, acc.to(Out.type.element_ty))
+
+
+# The main inner-loop logic for computing dK and dV.
+@triton.jit
+def _attn_bwd_dkdv_bsa_varlen(
+    dk, dv,
+    k, v,
+    Q, DO,
+    M, D,
+    block_indices,
+    block_indices_lens,
+    # shared by Q/K/V/DO.
+    stride_qm, stride_qk,
+    stride_dom, stride_dok,
+    stride_mm,
+    stride_dm,
+    stride_bm,
+    N_CTX, 
+    kv_valid_mask,
+    k_start_pos,
+    BLOCK_N_K: tl.constexpr,
+    BLOCK_M1: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    HAS_KV_MASK: tl.constexpr,
+    ):
+
+    if HAS_KV_MASK:
+        offs_k = k_start_pos + tl.arange(0, BLOCK_N_K)
+        k_valid = tl.load(kv_valid_mask + offs_k)
+    
+    QT_block_ptr = tl.make_block_ptr(
+        base=Q,
+        shape=(HEAD_DIM, N_CTX),
+        strides=(stride_qk, stride_qm),
+        offsets=(0, 0),
+        block_shape=(HEAD_DIM, BLOCK_M1),
+        order=(0, 1),
+    )
+    
+    DO_block_ptr = tl.make_block_ptr(
+        base=DO,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_dom, stride_dok),
+        offsets=(0, 0),
+        block_shape=(BLOCK_M1, HEAD_DIM),
+        order=(1, 0),
+    )
+
+    S = tl.load(block_indices_lens)
+    for i in range(S):
+        block_id = tl.load(block_indices + i * stride_bm).to(tl.int32)
+        start_m = block_id * BLOCK_M1
+        start_m = tl.multiple_of(start_m, BLOCK_M1)
+        
+        QT_block_ptr_i = tl.advance(QT_block_ptr, (0, start_m))
+        DO_block_ptr_i = tl.advance(DO_block_ptr, (start_m, 0))
+        
+        qT = tl.load(QT_block_ptr_i)
+        offs_m = start_m + tl.arange(0, BLOCK_M1) * stride_mm
+        m = tl.load(M + offs_m)
+        kqT = tl.dot(k, qT)
+        if HAS_KV_MASK:
+            kqT = tl.where(k_valid[:, None], kqT, float('-inf'))
+        pT = tl.math.exp2(kqT - m[None, :])
+        if HAS_KV_MASK:
+            pT = tl.where(k_valid[:, None], pT, 0.0)
+
+        do = tl.load(DO_block_ptr_i)
+        # Compute dV.
+        ppT = pT
+        ppT = ppT.to(v.dtype)
+        dv += tl.dot(ppT, do)
+        # D (= delta) is pre-divided by ds_scale.
+        offs_d = start_m + tl.arange(0, BLOCK_M1) * stride_dm
+        Di = tl.load(D + offs_d)
+        # Compute dP and dS.
+        dpT = tl.dot(v, tl.trans(do)).to(tl.float32)
+        dsT = pT * (dpT - Di[None, :])
+        dsT = dsT.to(v.dtype)
+        dk += tl.dot(dsT, tl.trans(qT))
+        
+
+    return dk, dv
+
+
+# the main inner-loop logic for computing dQ
+@triton.jit
+def _attn_bwd_dq_bsa_varlen(
+    dq,
+    q, do, 
+    m, d,
+    K, V, 
+    N_CTX,
+    BLOCK_N2: tl.constexpr,
+    BLOCK_N_LG: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    block_indices,
+    block_indices_lens,
+    stride_bn,
+    stride_kn, stride_kk,        
+    stride_vn, stride_vk,
+    kv_valid_mask,
+    HAS_KV_MASK: tl.constexpr,
+    ):
+    
+    VT_block_ptr = tl.make_block_ptr(
+        base=V,
+        shape=(HEAD_DIM, N_CTX),
+        strides=(stride_vk, stride_vn),
+        offsets=(0, 0),
+        block_shape=(HEAD_DIM, BLOCK_N2),
+        order=(0, 1),
+    )
+    KT_block_ptr = tl.make_block_ptr(
+        base=K,
+        shape=(HEAD_DIM, N_CTX),
+        strides=(stride_kk, stride_kn),
+        offsets=(0, 0),
+        block_shape=(HEAD_DIM, BLOCK_N2),
+        order=(0, 1),
+    )
+    
+    S = tl.load(block_indices_lens)
+    for i in range(S):
+        block_id = tl.load(block_indices + i * stride_bn).to(tl.int32)
+        lo, hi = block_id * BLOCK_N_LG, (block_id + 1) * BLOCK_N_LG
+        lo = tl.multiple_of(lo, BLOCK_N2)
+        
+        KT_block_ptr_i = tl.advance(KT_block_ptr, (0, lo))
+        VT_block_ptr_i = tl.advance(VT_block_ptr, (0, lo))
+        mask_offset = lo
+    
+        for start_n in range(lo, hi, BLOCK_N2):
+            start_n = tl.multiple_of(start_n, BLOCK_N2)
+            
+            kT = tl.load(KT_block_ptr_i)
+            vT = tl.load(VT_block_ptr_i)
+            qkT = tl.dot(q, kT)
+            if HAS_KV_MASK:
+                offs_n = mask_offset + tl.arange(0, BLOCK_N2)
+                k_valid = tl.load(kv_valid_mask + offs_n)
+                qkT = tl.where(k_valid[None, :], qkT, float('-inf'))
+            p = tl.math.exp2(qkT - m)
+            if HAS_KV_MASK:
+                p = tl.where(k_valid[None, :], p, 0.0)
+            # Compute dP and dS.
+            dp = tl.dot(do, vT).to(tl.float32)
+            ds = p * (dp - d)
+            ds = ds.to(kT.dtype)
+            # Compute dQ.
+            dq += tl.dot(ds, tl.trans(kT))
+            # Increment pointers.
+            KT_block_ptr_i = tl.advance(KT_block_ptr_i, (0, BLOCK_N2))
+            VT_block_ptr_i = tl.advance(VT_block_ptr_i, (0, BLOCK_N2))
+            mask_offset += BLOCK_N2
+        
+    return dq
+
+# the main inner-loop logic for computing dQ
+@triton.jit
+def _attn_bwd_dq_bsa_varlen_align(
+    dq,
+    q, do, 
+    m, d,
+    K, V, 
+    N_CTX,
+    BLOCK_N_LG: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    block_indices,
+    block_indices_lens,
+    stride_bn,
+    stride_kn, stride_kk,        
+    stride_vn, stride_vk,
+    kv_valid_mask,
+    HAS_KV_MASK: tl.constexpr,
+    ):
+    
+    VT_block_ptr = tl.make_block_ptr(
+        base=V,
+        shape=(HEAD_DIM, N_CTX),
+        strides=(stride_vk, stride_vn),
+        offsets=(0, 0),
+        block_shape=(HEAD_DIM, BLOCK_N_LG),
+        order=(0, 1),
+    )
+    KT_block_ptr = tl.make_block_ptr(
+        base=K,
+        shape=(HEAD_DIM, N_CTX),
+        strides=(stride_kk, stride_kn),
+        offsets=(0, 0),
+        block_shape=(HEAD_DIM, BLOCK_N_LG),
+        order=(0, 1),
+    )
+    
+    S = tl.load(block_indices_lens)
+    for i in range(S):
+        block_id = tl.load(block_indices + i * stride_bn).to(tl.int32)
+        lo = block_id * BLOCK_N_LG
+        lo = tl.multiple_of(lo, BLOCK_N_LG)
+        
+        KT_block_ptr_i = tl.advance(KT_block_ptr, (0, lo))
+        VT_block_ptr_i = tl.advance(VT_block_ptr, (0, lo))
+          
+        kT = tl.load(KT_block_ptr_i)
+        vT = tl.load(VT_block_ptr_i)
+        qkT = tl.dot(q, kT)
+        if HAS_KV_MASK:
+            offs_n = lo + tl.arange(0, BLOCK_N_LG)
+            k_valid = tl.load(kv_valid_mask + offs_n)
+            qkT = tl.where(k_valid[None, :], qkT, float('-inf'))
+        p = tl.math.exp2(qkT - m)
+        if HAS_KV_MASK:
+            p = tl.where(k_valid[None, :], p, 0.0)
+        # Compute dP and dS.
+        dp = tl.dot(do, vT).to(tl.float32)
+        ds = p * (dp - d)
+        ds = ds.to(kT.dtype)
+        # Compute dQ.
+        dq += tl.dot(ds, tl.trans(kT))
+
+        
+    return dq
+
+configs_bwd_dkdv_bsa_varlen_preset = {
+    'default': {
+        'BLOCK_N': 128,
+        'num_stages': 2,
+        'num_warps': 8,
+    },
+    'BLOCK_N_DQ_LG=64': {
+        'BLOCK_N': 64,
+        'num_stages': 2,
+        'num_warps': 4,
+    }
+}
+configs_bwd_dkdv_bsa_varlen = [
+    triton.Config({'BLOCK_N': BN}, num_stages=s, num_warps=w) \
+    for BN in [32, 64, 128] \
+    for s in [2, 3, 4, 5] \
+    for w in [4, 8] \
+]
+bwd_dkdv_bsa_varlen_reevaluate_keys = ['N_CTX', 'BLOCK_M', 'BLOCK_N_DQ_LG', 'SPARSITY'] if os.environ.get('TRITON_REEVALUATE_KEY', '0') == '1' else []
+@autotune(list(configs_bwd_dkdv_bsa_varlen), key=bwd_dkdv_bsa_varlen_reevaluate_keys)
+@triton.jit
+def _attn_bwd_dkdv_bsa_varlen_wrapper(
+    Q, K, V, sm_scale,  # softmax scale
+    DO,
+    DK, DV,
+    M, # lse (log2)
+    D,
+    block_indices,
+    block_indices_lens,
+    kv_valid_mask,
+    # qkv
+    stride_qz, stride_qh, stride_qm, stride_qk,
+    stride_kz, stride_kh, stride_kn, stride_kk,
+    stride_vz, stride_vh, stride_vn, stride_vk,
+    # dk dv do
+    stride_dkz, stride_dkh, stride_dkn, stride_dkk,
+    stride_dvz, stride_dvh, stride_dvn, stride_dvk,
+    stride_doz, stride_doh, stride_dom, stride_dok,
+    # m, d
+    stride_mz, stride_mh, stride_mm,
+    stride_dz, stride_dh, stride_dm,
+    #
+    stride_bz, stride_bh, stride_bn, stride_bm, # block_indices
+    stride_lz, stride_lh, stride_ln, # block_indices_lens
+    #
+    H, N_CTX,
+    BLOCK_M: tl.constexpr, 
+    BLOCK_N: tl.constexpr, 
+    BLOCK_N_DQ_LG: tl.constexpr, # logical block size
+    HEAD_DIM: tl.constexpr,
+    SPARSITY: tl.constexpr,
+    HAS_KV_MASK: tl.constexpr,
+    ):
+
+    tl.static_assert(BLOCK_N_DQ_LG % BLOCK_N == 0)
+    start_n = tl.program_id(0)
+
+    off_hz = tl.program_id(2)
+    off_z = off_hz // H
+    off_h = off_hz % H
+    
+    off_q = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh
+    off_k = off_z.to(tl.int64) * stride_kz + off_h.to(tl.int64) * stride_kh
+    off_v = off_z.to(tl.int64) * stride_vz + off_h.to(tl.int64) * stride_vh
+
+    off_dk = off_z.to(tl.int64) * stride_dkz + off_h.to(tl.int64) * stride_dkh
+    off_dv = off_z.to(tl.int64) * stride_dvz + off_h.to(tl.int64) * stride_dvh
+    off_do = off_z.to(tl.int64) * stride_doz + off_h.to(tl.int64) * stride_doh
+    
+    off_m = off_z.to(tl.int64) * stride_mz + off_h.to(tl.int64) * stride_mh
+    off_d = off_z.to(tl.int64) * stride_dz + off_h.to(tl.int64) * stride_dh
+
+    off_block_incides =      off_z.to(tl.int64) * stride_bz + off_h.to(tl.int64) * stride_bh
+    off_block_incides_lens = off_z.to(tl.int64) * stride_lz + off_h.to(tl.int64) * stride_lh
+    
+    # offset pointers for batch/head
+    Q += off_q
+    K += off_k
+    V += off_v
+    DO += off_do
+    DK += off_dk
+    DV += off_dv
+    
+    M += off_m
+    D += off_d
+    block_indices += off_block_incides
+    block_indices_lens += off_block_incides_lens
+
+    # ---------------------------------------- [DKDV] ----------------------------------------
+
+    dv = tl.zeros([BLOCK_N, HEAD_DIM], dtype=tl.float32)
+    dk = tl.zeros([BLOCK_N, HEAD_DIM], dtype=tl.float32)
+
+    # load K and V: they stay in SRAM throughout the inner loop.
+    
+    K_block_ptr = tl.make_block_ptr(
+        base=K,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_kn, stride_kk),
+        offsets=(start_n * BLOCK_N, 0),
+        block_shape=(BLOCK_N, HEAD_DIM),
+        order=(1, 0),
+    )
+    V_block_ptr = tl.make_block_ptr(
+        base=V,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_vn, stride_vk),
+        offsets=(start_n * BLOCK_N, 0),
+        block_shape=(BLOCK_N, HEAD_DIM),
+        order=(1, 0),
+    )
+    DK_block_ptr = tl.make_block_ptr(
+        base=DK,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_dkn, stride_dkk),
+        offsets=(start_n * BLOCK_N, 0),
+        block_shape=(BLOCK_N, HEAD_DIM),
+        order=(1, 0),
+    )
+    DV_block_ptr = tl.make_block_ptr(
+        base=DV,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_dvn, stride_dvk),
+        offsets=(start_n * BLOCK_N, 0),
+        block_shape=(BLOCK_N, HEAD_DIM),
+        order=(1, 0),
+    )
+    
+    k = tl.load(K_block_ptr)
+    v = tl.load(V_block_ptr)
+    
+    k_compress_idx = start_n * BLOCK_N // BLOCK_N_DQ_LG
+    block_indices_i = block_indices + k_compress_idx * stride_bn
+    block_indices_lens_i = block_indices_lens + k_compress_idx * stride_ln
+    
+    dk, dv = _attn_bwd_dkdv_bsa_varlen(
+        dk, dv, 
+        k, v, 
+        Q, DO, 
+        M, D, 
+        block_indices_i,
+        block_indices_lens_i,
+        stride_qm, stride_qk,
+        stride_dom, stride_dok,
+        stride_mm,
+        stride_dm,
+        stride_bm,
+        N_CTX, 
+        kv_valid_mask,
+        start_n * BLOCK_N,
+        BLOCK_N,
+        BLOCK_M, 
+        HEAD_DIM,
+        HAS_KV_MASK,
+    )
+    
+    # Write back dk
+    dk *= sm_scale # S = scale * QKT; dK = scale * QdST
+    tl.store(DK_block_ptr, dk.to(k.dtype))
+    
+    # Write back dv
+    tl.store(DV_block_ptr, dv.to(v.dtype))
+
+configs_bwd_dq_bsa_varlen_preset = {
+    'default': {
+        'BLOCK_N_DQ': 64,
+        'num_stages': 2,
+        'num_warps': 8,
+    },
+    'BLOCK_N_DQ_LG=64': {
+        'BLOCK_N_DQ': 64,
+        'num_stages': 2,
+        'num_warps': 4,
+    },
+}
+configs_bwd_dq_bsa_varlen = [
+    triton.Config({'BLOCK_N_DQ': BN}, num_stages=s, num_warps=w) \
+    for BN in [32, 64, 128] \
+    for s in [2, 3, 4, 5] \
+    for w in [4, 8] \
+]
+bwd_dq_bsa_varlen_reevaluate_keys = ['N_CTX', 'BLOCK_M', 'BLOCK_N_DQ_LG', 'SPARSITY'] if os.environ.get('TRITON_REEVALUATE_KEY', '0') == '1' else []
+@autotune(list(configs_bwd_dq_bsa_varlen), key=bwd_dq_bsa_varlen_reevaluate_keys)
+@triton.jit
+def _attn_bwd_dq_bsa_varlen_wrapper(
+    Q, K, V,
+    DO, 
+    DQ,
+    M, # lse (log2)
+    D,
+    block_indices,
+    block_indices_lens,
+    kv_valid_mask,
+    # qkv
+    stride_qz, stride_qh, stride_qm, stride_qk,
+    stride_kz, stride_kh, stride_kn, stride_kk,
+    stride_vz, stride_vh, stride_vn, stride_vk,
+    # dq do
+    stride_dqz, stride_dqh, stride_dqm, stride_dqk,
+    stride_doz, stride_doh, stride_dom, stride_dok,
+    # m, d
+    stride_mz, stride_mh, stride_mm,
+    stride_dz, stride_dh, stride_dm,
+    #
+    stride_bz, stride_bh, stride_bm, stride_bn, # block_indices
+    stride_lz, stride_lh, stride_lm, # block_indices_lens
+    #
+    H, N_CTX,
+    BLOCK_M: tl.constexpr, 
+    BLOCK_N_DQ_LG: tl.constexpr, # logical block size
+    BLOCK_N_DQ: tl.constexpr, 
+    HEAD_DIM: tl.constexpr,
+    SPARSITY: tl.constexpr,
+    HAS_KV_MASK: tl.constexpr,
+    ):
+
+    tl.static_assert(BLOCK_N_DQ_LG % BLOCK_N_DQ == 0)
+    tl.static_assert(BLOCK_N_DQ_LG % BLOCK_M == 0)
+    
+    LN2: tl.constexpr = 0.6931471824645996  # = ln(2)
+    
+    start_m = tl.program_id(0)
+
+    off_hz = tl.program_id(2)
+    off_z = off_hz // H
+    off_h = off_hz % H
+    
+    off_q = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh
+    off_k = off_z.to(tl.int64) * stride_kz + off_h.to(tl.int64) * stride_kh
+    off_v = off_z.to(tl.int64) * stride_vz + off_h.to(tl.int64) * stride_vh
+
+    off_dq = off_z.to(tl.int64) * stride_dqz + off_h.to(tl.int64) * stride_dqh
+    off_do = off_z.to(tl.int64) * stride_doz + off_h.to(tl.int64) * stride_doh
+    
+    off_m = off_z.to(tl.int64) * stride_mz + off_h.to(tl.int64) * stride_mh
+    off_d = off_z.to(tl.int64) * stride_dz + off_h.to(tl.int64) * stride_dh
+
+    off_block_incides =      off_z.to(tl.int64) * stride_bz + off_h.to(tl.int64) * stride_bh
+    off_block_incides_lens = off_z.to(tl.int64) * stride_lz + off_h.to(tl.int64) * stride_lh
+
+    # offset pointers for batch/head
+    Q += off_q
+    K += off_k
+    V += off_v
+    DO += off_do
+    DQ += off_dq
+    
+    M += off_m
+    D += off_d
+    block_indices += off_block_incides
+    block_indices_lens += off_block_incides_lens
+
+    # ---------------------------------------- [DQ] ----------------------------------------
+    
+    Q_block_ptr = tl.make_block_ptr(
+        base=Q,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_qm, stride_qk),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, HEAD_DIM),
+        order=(1, 0),
+    )
+    DO_block_ptr = tl.make_block_ptr(
+        base=DO,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_dom, stride_dok),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, HEAD_DIM),
+        order=(1, 0),
+    )
+    DQ_block_ptr = tl.make_block_ptr(
+        base=DQ,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_dqm, stride_dqk),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, HEAD_DIM),
+        order=(1, 0),
+    )
+    
+    q = tl.load(Q_block_ptr)
+    do = tl.load(DO_block_ptr)
+    
+    dq = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
+    
+    start_m = start_m * BLOCK_M
+
+    offs_m = start_m + tl.arange(0, BLOCK_M) * stride_mm
+    offs_d = start_m + tl.arange(0, BLOCK_M) * stride_dm
+
+    m = tl.load(M + offs_m)
+    m = m[:, None]
+
+    d = tl.load(D + offs_d)
+    d = d[:, None]
+    
+    block_indices_m = block_indices + (start_m // BLOCK_M) * stride_bm
+    block_indices_lens_m = block_indices_lens + (start_m // BLOCK_M) * stride_lm
+
+    dq = _attn_bwd_dq_bsa_varlen(
+        dq,
+        q, do, 
+        m, d,
+        K, V, 
+        N_CTX,
+        BLOCK_N_DQ,
+        BLOCK_N_DQ_LG,   
+        HEAD_DIM,
+        block_indices_m, 
+        block_indices_lens_m, 
+        stride_bn,   
+        stride_kn, stride_kk,        
+        stride_vn, stride_vk,
+        kv_valid_mask,
+        HAS_KV_MASK,
+    )
+    
+    # Write back dQ.
+    dq *= LN2
+    tl.store(DQ_block_ptr, dq.to(q.dtype))
+
+configs_bwd_dq_bsa_varlen_align_preset = {
+    'default': {
+        'num_stages': 2,
+        'num_warps': 8,
+    },
+    'BLOCK_N_DQ_LG=64': {
+        'num_stages': 2,
+        'num_warps': 4,
+    },
+}
+configs_bwd_dq_bsa_varlen_align = [
+    triton.Config({}, num_stages=s, num_warps=w) \
+    for s in [2, 3, 4, 5] \
+    for w in [4, 8] \
+]
+bwd_dq_bsa_varlen_align_reevaluate_keys = ['N_CTX', 'BLOCK_M', 'BLOCK_N_DQ_LG', 'SPARSITY'] if os.environ.get('TRITON_REEVALUATE_KEY', '0') == '1' else []
+@autotune(list(configs_bwd_dq_bsa_varlen_align), key=bwd_dq_bsa_varlen_align_reevaluate_keys)
+@triton.jit
+def _attn_bwd_dq_bsa_varlen_align_wrapper(
+    Q, K, V,
+    DO, 
+    DQ,
+    M, # lse (log2)
+    D,
+    block_indices,
+    block_indices_lens,
+    kv_valid_mask,
+    # qkv
+    stride_qz, stride_qh, stride_qm, stride_qk,
+    stride_kz, stride_kh, stride_kn, stride_kk,
+    stride_vz, stride_vh, stride_vn, stride_vk,
+    # dq do
+    stride_dqz, stride_dqh, stride_dqm, stride_dqk,
+    stride_doz, stride_doh, stride_dom, stride_dok,
+    # m, d
+    stride_mz, stride_mh, stride_mm,
+    stride_dz, stride_dh, stride_dm,
+    #
+    stride_bz, stride_bh, stride_bm, stride_bn, # block_indices
+    stride_lz, stride_lh, stride_lm, # block_indices_lens
+    #
+    H, N_CTX,
+    BLOCK_M: tl.constexpr, 
+    BLOCK_N_DQ_LG: tl.constexpr, # logical block size
+    HEAD_DIM: tl.constexpr,
+    SPARSITY: tl.constexpr,
+    HAS_KV_MASK: tl.constexpr,
+    ):
+
+    tl.static_assert(BLOCK_N_DQ_LG % BLOCK_N_DQ_LG == 0)
+    tl.static_assert(BLOCK_N_DQ_LG % BLOCK_M == 0)
+    
+    LN2: tl.constexpr = 0.6931471824645996  # = ln(2)
+    
+    start_m = tl.program_id(0)
+    
+    off_hz = tl.program_id(2)
+    off_z = off_hz // H
+    off_h = off_hz % H
+    
+    off_q = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh
+    off_k = off_z.to(tl.int64) * stride_kz + off_h.to(tl.int64) * stride_kh
+    off_v = off_z.to(tl.int64) * stride_vz + off_h.to(tl.int64) * stride_vh
+
+    off_dq = off_z.to(tl.int64) * stride_dqz + off_h.to(tl.int64) * stride_dqh
+    off_do = off_z.to(tl.int64) * stride_doz + off_h.to(tl.int64) * stride_doh
+    
+    off_m = off_z.to(tl.int64) * stride_mz + off_h.to(tl.int64) * stride_mh
+    off_d = off_z.to(tl.int64) * stride_dz + off_h.to(tl.int64) * stride_dh
+
+    off_block_incides =      off_z.to(tl.int64) * stride_bz + off_h.to(tl.int64) * stride_bh
+    off_block_incides_lens = off_z.to(tl.int64) * stride_lz + off_h.to(tl.int64) * stride_lh
+
+    # offset pointers for batch/head
+    Q += off_q
+    K += off_k
+    V += off_v
+    DO += off_do
+    DQ += off_dq
+    
+    M += off_m
+    D += off_d
+    block_indices += off_block_incides
+    block_indices_lens += off_block_incides_lens
+
+    # ---------------------------------------- [DQ] ----------------------------------------
+    
+    Q_block_ptr = tl.make_block_ptr(
+        base=Q,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_qm, stride_qk),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, HEAD_DIM),
+        order=(1, 0),
+    )
+    DO_block_ptr = tl.make_block_ptr(
+        base=DO,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_dom, stride_dok),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, HEAD_DIM),
+        order=(1, 0),
+    )
+    DQ_block_ptr = tl.make_block_ptr(
+        base=DQ,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_dqm, stride_dqk),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, HEAD_DIM),
+        order=(1, 0),
+    )
+    
+    q = tl.load(Q_block_ptr)
+    do = tl.load(DO_block_ptr)
+    
+    dq = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
+    
+    start_m = start_m * BLOCK_M
+
+    offs_m = start_m + tl.arange(0, BLOCK_M) * stride_mm
+    offs_d = start_m + tl.arange(0, BLOCK_M) * stride_dm
+
+    m = tl.load(M + offs_m)
+    m = m[:, None]
+    
+    # D (= delta) is pre-divided by ds_scale.
+    d = tl.load(D + offs_d)
+    d = d[:, None]
+    
+    block_indices_m = block_indices + (start_m // BLOCK_M) * stride_bm
+    block_indices_lens_m = block_indices_lens + (start_m // BLOCK_M) * stride_lm
+
+    dq = _attn_bwd_dq_bsa_varlen_align(
+        dq,
+        q, do, 
+        m, d,
+        K, V, 
+        N_CTX,
+        BLOCK_N_DQ_LG,   
+        HEAD_DIM,
+        block_indices_m, 
+        block_indices_lens_m, 
+        stride_bn,    
+        stride_kn, stride_kk,        
+        stride_vn, stride_vk,
+        kv_valid_mask,
+        HAS_KV_MASK,
+    )
+    
+    # Write back dQ.    
+    dq *= LN2
+    tl.store(DQ_block_ptr, dq.to(q.dtype))

hymm/models/modules/dac_vae.py

@@ -0,0 +1,1103 @@
+
+import math
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List, Union
+
+import numpy as np
+import torch
+import tqdm
+from audiotools import AudioSignal
+from audiotools.ml import BaseModel
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils.accelerate_utils import apply_forward_hook
+from torch import nn
+from torch.nn.utils import weight_norm
+
+SUPPORTED_VERSIONS = ["1.0.0"]
+
+
+@dataclass
+class DACFile:
+    codes: torch.Tensor
+
+    # Metadata
+    chunk_length: int
+    original_length: int
+    input_db: float
+    channels: int
+    sample_rate: int
+    padding: bool
+    dac_version: str
+
+    def save(self, path):
+        artifacts = {
+            "codes": self.codes.numpy().astype(np.uint16),
+            "metadata": {
+                "input_db": self.input_db.numpy().astype(np.float32),
+                "original_length": self.original_length,
+                "sample_rate": self.sample_rate,
+                "chunk_length": self.chunk_length,
+                "channels": self.channels,
+                "padding": self.padding,
+                "dac_version": SUPPORTED_VERSIONS[-1],
+            },
+        }
+        path = Path(path).with_suffix(".dac")
+        with open(path, "wb") as f:
+            np.save(f, artifacts)
+        return path
+
+    @classmethod
+    def load(cls, path):
+        artifacts = np.load(path, allow_pickle=True)[()]
+        codes = torch.from_numpy(artifacts["codes"].astype(int))
+        if artifacts["metadata"].get("dac_version", None) not in SUPPORTED_VERSIONS:
+            raise RuntimeError(
+                f"Given file {path} can't be loaded with this version of descript-audio-codec."
+            )
+        return cls(codes=codes, **artifacts["metadata"])
+
+
+class CodecMixin:
+    @property
+    def padding(self):
+        if not hasattr(self, "_padding"):
+            self._padding = True
+        return self._padding
+
+    @padding.setter
+    def padding(self, value):
+        assert isinstance(value, bool)
+
+        layers = [
+            l for l in self.modules() if isinstance(l, (nn.Conv1d, nn.ConvTranspose1d))
+        ]
+
+        for layer in layers:
+            if value:
+                if hasattr(layer, "original_padding"):
+                    layer.padding = layer.original_padding
+            else:
+                layer.original_padding = layer.padding
+                layer.padding = tuple(0 for _ in range(len(layer.padding)))
+
+        self._padding = value
+
+    def get_delay(self):
+        # Any number works here, delay is invariant to input length
+        l_out = self.get_output_length(0)
+        L = l_out
+
+        layers = []
+        for layer in self.modules():
+            if isinstance(layer, (nn.Conv1d, nn.ConvTranspose1d)):
+                layers.append(layer)
+
+        for layer in reversed(layers):
+            d = layer.dilation[0]
+            k = layer.kernel_size[0]
+            s = layer.stride[0]
+
+            if isinstance(layer, nn.ConvTranspose1d):
+                L = ((L - d * (k - 1) - 1) / s) + 1
+            elif isinstance(layer, nn.Conv1d):
+                L = (L - 1) * s + d * (k - 1) + 1
+
+            L = math.ceil(L)
+
+        l_in = L
+
+        return (l_in - l_out) // 2
+
+    def get_output_length(self, input_length):
+        L = input_length
+        # Calculate output length
+        for layer in self.modules():
+            if isinstance(layer, (nn.Conv1d, nn.ConvTranspose1d)):
+                d = layer.dilation[0]
+                k = layer.kernel_size[0]
+                s = layer.stride[0]
+
+                if isinstance(layer, nn.Conv1d):
+                    L = ((L - d * (k - 1) - 1) / s) + 1
+                elif isinstance(layer, nn.ConvTranspose1d):
+                    L = (L - 1) * s + d * (k - 1) + 1
+
+                L = math.floor(L)
+        return L
+
+    @torch.no_grad()
+    def compress(
+        self,
+        audio_path_or_signal: Union[str, Path, AudioSignal],
+        win_duration: float = 1.0,
+        verbose: bool = False,
+        normalize_db: float = -16,
+        n_quantizers: int = None,
+    ) -> DACFile:
+        """Processes an audio signal from a file or AudioSignal object into
+        discrete codes. This function processes the signal in short windows,
+        using constant GPU memory.
+
+        Parameters
+        ----------
+        audio_path_or_signal : Union[str, Path, AudioSignal]
+            audio signal to reconstruct
+        win_duration : float, optional
+            window duration in seconds, by default 5.0
+        verbose : bool, optional
+            by default False
+        normalize_db : float, optional
+            normalize db, by default -16
+
+        Returns
+        -------
+        DACFile
+            Object containing compressed codes and metadata
+            required for decompression
+        """
+        audio_signal = audio_path_or_signal
+        if isinstance(audio_signal, (str, Path)):
+            audio_signal = AudioSignal.load_from_file_with_ffmpeg(str(audio_signal))
+
+        self.eval()
+        original_padding = self.padding
+        original_device = audio_signal.device
+
+        audio_signal = audio_signal.clone()
+        audio_signal = audio_signal.to_mono()
+        original_sr = audio_signal.sample_rate
+
+        resample_fn = audio_signal.resample
+        loudness_fn = audio_signal.loudness
+
+        # If audio is > 10 minutes long, use the ffmpeg versions
+        if audio_signal.signal_duration >= 10 * 60 * 60:
+            resample_fn = audio_signal.ffmpeg_resample
+            loudness_fn = audio_signal.ffmpeg_loudness
+
+        original_length = audio_signal.signal_length
+        resample_fn(self.sample_rate)
+        input_db = loudness_fn()
+
+        if normalize_db is not None:
+            audio_signal.normalize(normalize_db)
+        audio_signal.ensure_max_of_audio()
+
+        nb, nac, nt = audio_signal.audio_data.shape
+        audio_signal.audio_data = audio_signal.audio_data.reshape(nb * nac, 1, nt)
+        win_duration = (
+            audio_signal.signal_duration if win_duration is None else win_duration
+        )
+
+        if audio_signal.signal_duration <= win_duration:
+            # Unchunked compression (used if signal length < win duration)
+            self.padding = True
+            n_samples = nt
+            hop = nt
+        else:
+            # Chunked inference
+            self.padding = False
+            # Zero-pad signal on either side by the delay
+            audio_signal.zero_pad(self.delay, self.delay)
+            n_samples = int(win_duration * self.sample_rate)
+            # Round n_samples to nearest hop length multiple
+            n_samples = int(math.ceil(n_samples / self.hop_length) * self.hop_length)
+            hop = self.get_output_length(n_samples)
+
+        codes = []
+        range_fn = range if not verbose else tqdm.trange
+
+        for i in range_fn(0, nt, hop):
+            x = audio_signal[..., i : i + n_samples]
+            x = x.zero_pad(0, max(0, n_samples - x.shape[-1]))
+
+            audio_data = x.audio_data.to(self.device)
+            audio_data = self.preprocess(audio_data, self.sample_rate)
+            _, c, _, _, _ = self.encode(audio_data, n_quantizers)
+            codes.append(c.to(original_device))
+            chunk_length = c.shape[-1]
+
+        codes = torch.cat(codes, dim=-1)
+
+        dac_file = DACFile(
+            codes=codes,
+            chunk_length=chunk_length,
+            original_length=original_length,
+            input_db=input_db,
+            channels=nac,
+            sample_rate=original_sr,
+            padding=self.padding,
+            dac_version=SUPPORTED_VERSIONS[-1],
+        )
+
+        if n_quantizers is not None:
+            codes = codes[:, :n_quantizers, :]
+
+        self.padding = original_padding
+        return dac_file
+
+    @torch.no_grad()
+    def decompress(
+        self,
+        obj: Union[str, Path, DACFile],
+        verbose: bool = False,
+    ) -> AudioSignal:
+        """Reconstruct audio from a given .dac file
+
+        Parameters
+        ----------
+        obj : Union[str, Path, DACFile]
+            .dac file location or corresponding DACFile object.
+        verbose : bool, optional
+            Prints progress if True, by default False
+
+        Returns
+        -------
+        AudioSignal
+            Object with the reconstructed audio
+        """
+        self.eval()
+        if isinstance(obj, (str, Path)):
+            obj = DACFile.load(obj)
+
+        original_padding = self.padding
+        self.padding = obj.padding
+
+        range_fn = range if not verbose else tqdm.trange
+        codes = obj.codes
+        original_device = codes.device
+        chunk_length = obj.chunk_length
+        recons = []
+
+        for i in range_fn(0, codes.shape[-1], chunk_length):
+            c = codes[..., i : i + chunk_length].to(self.device)
+            z = self.quantizer.from_codes(c)[0]
+            r = self.decode(z)
+            recons.append(r.to(original_device))
+
+        recons = torch.cat(recons, dim=-1)
+        recons = AudioSignal(recons, self.sample_rate)
+
+        resample_fn = recons.resample
+        loudness_fn = recons.loudness
+
+        # If audio is > 10 minutes long, use the ffmpeg versions
+        if recons.signal_duration >= 10 * 60 * 60:
+            resample_fn = recons.ffmpeg_resample
+            loudness_fn = recons.ffmpeg_loudness
+
+        if obj.input_db is not None:
+            recons.normalize(obj.input_db)
+
+        resample_fn(obj.sample_rate)
+
+        if obj.original_length is not None:
+            recons = recons[..., : obj.original_length]
+            loudness_fn()
+            recons.audio_data = recons.audio_data.reshape(
+                -1, obj.channels, obj.original_length
+            )
+        else:
+            loudness_fn()
+
+        self.padding = original_padding
+        return recons
+
+
+def WNConv1d(*args, **kwargs):
+    return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+def WNConvTranspose1d(*args, **kwargs):
+    return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
+
+
+# Scripting this brings model speed up 1.4x
+@torch.jit.script
+def snake(x, alpha):
+    shape = x.shape
+    x = x.reshape(shape[0], shape[1], -1)
+    x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
+    x = x.reshape(shape)
+    return x
+
+
+class Snake1d(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.ones(1, channels, 1))
+
+    def forward(self, x):
+        return snake(x, self.alpha)
+
+
+import torch.nn.functional as F
+from einops import rearrange
+
+
+class VectorQuantize(nn.Module):
+    """
+    Implementation of VQ similar to Karpathy's repo:
+    https://github.com/karpathy/deep-vector-quantization
+    Additionally uses following tricks from Improved VQGAN
+    (https://arxiv.org/pdf/2110.04627.pdf):
+        1. Factorized codes: Perform nearest neighbor lookup in low-dimensional space
+            for improved codebook usage
+        2. l2-normalized codes: Converts euclidean distance to cosine similarity which
+            improves training stability
+    """
+
+    def __init__(self, input_dim: int, codebook_size: int, codebook_dim: int):
+        super().__init__()
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+
+        self.in_proj = WNConv1d(input_dim, codebook_dim, kernel_size=1)
+        self.out_proj = WNConv1d(codebook_dim, input_dim, kernel_size=1)
+        self.codebook = nn.Embedding(codebook_size, codebook_dim)
+
+    def forward(self, z):
+        """Quantized the input tensor using a fixed codebook and returns
+        the corresponding codebook vectors
+
+        Parameters
+        ----------
+        z : Tensor[B x D x T]
+
+        Returns
+        -------
+        Tensor[B x D x T]
+            Quantized continuous representation of input
+        Tensor[1]
+            Commitment loss to train encoder to predict vectors closer to codebook
+            entries
+        Tensor[1]
+            Codebook loss to update the codebook
+        Tensor[B x T]
+            Codebook indices (quantized discrete representation of input)
+        Tensor[B x D x T]
+            Projected latents (continuous representation of input before quantization)
+        """
+
+        # Factorized codes (ViT-VQGAN) Project input into low-dimensional space
+        z_e = self.in_proj(z)  # z_e : (B x D x T)
+        z_q, indices = self.decode_latents(z_e)
+
+        commitment_loss = F.mse_loss(z_e, z_q.detach(), reduction="none").mean([1, 2])
+        codebook_loss = F.mse_loss(z_q, z_e.detach(), reduction="none").mean([1, 2])
+
+        z_q = (
+            z_e + (z_q - z_e).detach()
+        )  # noop in forward pass, straight-through gradient estimator in backward pass
+
+        z_q = self.out_proj(z_q)
+
+        return z_q, commitment_loss, codebook_loss, indices, z_e
+
+    def embed_code(self, embed_id):
+        return F.embedding(embed_id, self.codebook.weight)
+
+    def decode_code(self, embed_id):
+        return self.embed_code(embed_id).transpose(1, 2)
+
+    def decode_latents(self, latents):
+        encodings = rearrange(latents, "b d t -> (b t) d")
+        codebook = self.codebook.weight  # codebook: (N x D)
+
+        # L2 normalize encodings and codebook (ViT-VQGAN)
+        encodings = F.normalize(encodings)
+        codebook = F.normalize(codebook)
+
+        # Compute euclidean distance with codebook
+        dist = (
+            encodings.pow(2).sum(1, keepdim=True)
+            - 2 * encodings @ codebook.t()
+            + codebook.pow(2).sum(1, keepdim=True).t()
+        )
+        indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0))
+        z_q = self.decode_code(indices)
+        return z_q, indices
+
+
+class ResidualVectorQuantize(nn.Module):
+    """
+    Introduced in SoundStream: An end2end neural audio codec
+    https://arxiv.org/abs/2107.03312
+    """
+
+    def __init__(
+        self,
+        input_dim: int = 512,
+        n_codebooks: int = 9,
+        codebook_size: int = 1024,
+        codebook_dim: Union[int, list] = 8,
+        quantizer_dropout: float = 0.0,
+    ):
+        super().__init__()
+        if isinstance(codebook_dim, int):
+            codebook_dim = [codebook_dim for _ in range(n_codebooks)]
+
+        self.n_codebooks = n_codebooks
+        self.codebook_dim = codebook_dim
+        self.codebook_size = codebook_size
+
+        self.quantizers = nn.ModuleList(
+            [
+                VectorQuantize(input_dim, codebook_size, codebook_dim[i])
+                for i in range(n_codebooks)
+            ]
+        )
+        self.quantizer_dropout = quantizer_dropout
+
+    def forward(self, z, n_quantizers: int = None):
+        """Quantized the input tensor using a fixed set of `n` codebooks and returns
+        the corresponding codebook vectors
+        Parameters
+        ----------
+        z : Tensor[B x D x T]
+        n_quantizers : int, optional
+            No. of quantizers to use
+            (n_quantizers < self.n_codebooks ex: for quantizer dropout)
+            Note: if `self.quantizer_dropout` is True, this argument is ignored
+                when in training mode, and a random number of quantizers is used.
+        Returns
+        -------
+        dict
+            A dictionary with the following keys:
+
+            "z" : Tensor[B x D x T]
+                Quantized continuous representation of input
+            "codes" : Tensor[B x N x T]
+                Codebook indices for each codebook
+                (quantized discrete representation of input)
+            "latents" : Tensor[B x N*D x T]
+                Projected latents (continuous representation of input before quantization)
+            "vq/commitment_loss" : Tensor[1]
+                Commitment loss to train encoder to predict vectors closer to codebook
+                entries
+            "vq/codebook_loss" : Tensor[1]
+                Codebook loss to update the codebook
+        """
+        z_q = 0
+        residual = z
+        commitment_loss = 0
+        codebook_loss = 0
+
+        codebook_indices = []
+        latents = []
+
+        if n_quantizers is None:
+            n_quantizers = self.n_codebooks
+        if self.training:
+            n_quantizers = torch.ones((z.shape[0],)) * self.n_codebooks + 1
+            dropout = torch.randint(1, self.n_codebooks + 1, (z.shape[0],))
+            n_dropout = int(z.shape[0] * self.quantizer_dropout)
+            n_quantizers[:n_dropout] = dropout[:n_dropout]
+            n_quantizers = n_quantizers.to(z.device)
+
+        for i, quantizer in enumerate(self.quantizers):
+            if self.training is False and i >= n_quantizers:
+                break
+
+            z_q_i, commitment_loss_i, codebook_loss_i, indices_i, z_e_i = quantizer(
+                residual
+            )
+
+            # Create mask to apply quantizer dropout
+            mask = (
+                torch.full((z.shape[0],), fill_value=i, device=z.device) < n_quantizers
+            )
+            z_q = z_q + z_q_i * mask[:, None, None]
+            residual = residual - z_q_i
+
+            # Sum losses
+            commitment_loss += (commitment_loss_i * mask).mean()
+            codebook_loss += (codebook_loss_i * mask).mean()
+
+            codebook_indices.append(indices_i)
+            latents.append(z_e_i)
+
+        codes = torch.stack(codebook_indices, dim=1)
+        latents = torch.cat(latents, dim=1)
+
+        return z_q, codes, latents, commitment_loss, codebook_loss
+
+    def from_codes(self, codes: torch.Tensor):
+        """Given the quantized codes, reconstruct the continuous representation
+        Parameters
+        ----------
+        codes : Tensor[B x N x T]
+            Quantized discrete representation of input
+        Returns
+        -------
+        Tensor[B x D x T]
+            Quantized continuous representation of input
+        """
+        z_q = 0.0
+        z_p = []
+        n_codebooks = codes.shape[1]
+        for i in range(n_codebooks):
+            z_p_i = self.quantizers[i].decode_code(codes[:, i, :])
+            z_p.append(z_p_i)
+
+            z_q_i = self.quantizers[i].out_proj(z_p_i)
+            z_q = z_q + z_q_i
+        return z_q, torch.cat(z_p, dim=1), codes
+
+    def from_latents(self, latents: torch.Tensor):
+        """Given the unquantized latents, reconstruct the
+        continuous representation after quantization.
+
+        Parameters
+        ----------
+        latents : Tensor[B x N x T]
+            Continuous representation of input after projection
+
+        Returns
+        -------
+        Tensor[B x D x T]
+            Quantized representation of full-projected space
+        Tensor[B x D x T]
+            Quantized representation of latent space
+        """
+        z_q = 0
+        z_p = []
+        codes = []
+        dims = np.cumsum([0] + [q.codebook_dim for q in self.quantizers])
+
+        n_codebooks = np.where(dims <= latents.shape[1])[0].max(axis=0, keepdims=True)[
+            0
+        ]
+        for i in range(n_codebooks):
+            j, k = dims[i], dims[i + 1]
+            z_p_i, codes_i = self.quantizers[i].decode_latents(latents[:, j:k, :])
+            z_p.append(z_p_i)
+            codes.append(codes_i)
+
+            z_q_i = self.quantizers[i].out_proj(z_p_i)
+            z_q = z_q + z_q_i
+
+        return z_q, torch.cat(z_p, dim=1), torch.stack(codes, dim=1)
+
+
+class AbstractDistribution:
+    def sample(self):
+        raise NotImplementedError()
+
+    def mode(self):
+        raise NotImplementedError()
+
+
+class DiracDistribution(AbstractDistribution):
+    def __init__(self, value):
+        self.value = value
+
+    def sample(self):
+        return self.value
+
+    def mode(self):
+        return self.value
+
+
+class DiagonalGaussianDistribution(object):
+    def __init__(self, parameters, deterministic=False):
+        self.parameters = parameters
+        self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
+        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+        self.deterministic = deterministic
+        self.std = torch.exp(0.5 * self.logvar)
+        self.var = torch.exp(self.logvar)
+        if self.deterministic:
+            self.var = self.std = torch.zeros_like(self.mean).to(device=self.parameters.device)
+
+    def sample(self):
+        x = self.mean + self.std * torch.randn(self.mean.shape).to(device=self.parameters.device)
+        return x
+
+    def kl(self, other=None):
+        if self.deterministic:
+            return torch.Tensor([0.0])
+        else:
+            if other is None:
+                return 0.5 * torch.mean(
+                    torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar,
+                    dim=[1, 2],
+                )
+            else:
+                return 0.5 * torch.mean(
+                    torch.pow(self.mean - other.mean, 2) / other.var
+                    + self.var / other.var
+                    - 1.0
+                    - self.logvar
+                    + other.logvar,
+                    dim=[1, 2],
+                )
+
+    def nll(self, sample, dims=[1, 2]):
+        if self.deterministic:
+            return torch.Tensor([0.0])
+        logtwopi = np.log(2.0 * np.pi)
+        return 0.5 * torch.sum(
+            logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
+            dim=dims,
+        )
+
+    def mode(self):
+        return self.mean
+
+
+def normal_kl(mean1, logvar1, mean2, logvar2):
+    """
+    source: https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/losses.py#L12
+    Compute the KL divergence between two gaussians.
+    Shapes are automatically broadcasted, so batches can be compared to
+    scalars, among other use cases.
+    """
+    tensor = None
+    for obj in (mean1, logvar1, mean2, logvar2):
+        if isinstance(obj, torch.Tensor):
+            tensor = obj
+            break
+    assert tensor is not None, "at least one argument must be a Tensor"
+
+    # Force variances to be Tensors. Broadcasting helps convert scalars to
+    # Tensors, but it does not work for torch.exp().
+    logvar1, logvar2 = [x if isinstance(x, torch.Tensor) else torch.tensor(x).to(tensor) for x in (logvar1, logvar2)]
+
+    return 0.5 * (
+        -1.0 + logvar2 - logvar1 + torch.exp(logvar1 - logvar2) + ((mean1 - mean2) ** 2) * torch.exp(-logvar2)
+    )
+
+
+def init_weights(m):
+    if isinstance(m, nn.Conv1d):
+        nn.init.trunc_normal_(m.weight, std=0.02)
+        nn.init.constant_(m.bias, 0)
+
+
+class ResidualUnit(nn.Module):
+    def __init__(self, dim: int = 16, dilation: int = 1):
+        super().__init__()
+        pad = ((7 - 1) * dilation) // 2
+        self.block = nn.Sequential(
+            Snake1d(dim),
+            WNConv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
+            Snake1d(dim),
+            WNConv1d(dim, dim, kernel_size=1),
+        )
+
+    def forward(self, x):
+        y = self.block(x)
+        pad = (x.shape[-1] - y.shape[-1]) // 2
+        if pad > 0:
+            x = x[..., pad:-pad]
+        return x + y
+
+
+class EncoderBlock(nn.Module):
+    def __init__(self, dim: int = 16, stride: int = 1):
+        super().__init__()
+        self.block = nn.Sequential(
+            ResidualUnit(dim // 2, dilation=1),
+            ResidualUnit(dim // 2, dilation=3),
+            ResidualUnit(dim // 2, dilation=9),
+            Snake1d(dim // 2),
+            WNConv1d(
+                dim // 2,
+                dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=math.ceil(stride / 2),
+            ),
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        d_model: int = 64,
+        strides: list = [2, 4, 8, 8],
+        d_latent: int = 64,
+    ):
+        super().__init__()
+        # Create first convolution
+        self.block = [WNConv1d(1, d_model, kernel_size=7, padding=3)]
+
+        # Create EncoderBlocks that double channels as they downsample by `stride`
+        for stride in strides:
+            d_model *= 2
+            self.block += [EncoderBlock(d_model, stride=stride)]
+
+        # Create last convolution
+        self.block += [
+            Snake1d(d_model),
+            WNConv1d(d_model, d_latent, kernel_size=3, padding=1),
+        ]
+
+        # Wrap black into nn.Sequential
+        self.block = nn.Sequential(*self.block)
+        self.enc_dim = d_model
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class DecoderBlock(nn.Module):
+    def __init__(self, input_dim: int = 16, output_dim: int = 8, stride: int = 1):
+        super().__init__()
+        self.block = nn.Sequential(
+            Snake1d(input_dim),
+            WNConvTranspose1d(
+                input_dim,
+                output_dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=math.ceil(stride / 2),
+                output_padding=stride % 2,
+            ),
+            ResidualUnit(output_dim, dilation=1),
+            ResidualUnit(output_dim, dilation=3),
+            ResidualUnit(output_dim, dilation=9),
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        input_channel,
+        channels,
+        rates,
+        d_out: int = 1,
+    ):
+        super().__init__()
+
+        # Add first conv layer
+        layers = [WNConv1d(input_channel, channels, kernel_size=7, padding=3)]
+
+        # Add upsampling + MRF blocks
+        for i, stride in enumerate(rates):
+            input_dim = channels // 2**i
+            output_dim = channels // 2 ** (i + 1)
+            layers += [DecoderBlock(input_dim, output_dim, stride)]
+
+        # Add final conv layer
+        layers += [
+            Snake1d(output_dim),
+            WNConv1d(output_dim, d_out, kernel_size=7, padding=3),
+            nn.Tanh(),
+        ]
+
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class DAC(BaseModel, CodecMixin, ModelMixin, ConfigMixin):
+
+    @register_to_config
+    def __init__(
+        self,
+        encoder_dim: int = 64,
+        encoder_rates: List[int] = [2, 4, 8, 8],
+        latent_dim: int = None,
+        decoder_dim: int = 1536,
+        decoder_rates: List[int] = [8, 8, 4, 2],
+        n_codebooks: int = 9,
+        codebook_size: int = 1024,
+        codebook_dim: Union[int, list] = 8,
+        quantizer_dropout: bool = False,
+        sample_rate: int = 44100,
+        continuous: bool = False,
+        use_weight_norm: bool = True,
+    ):
+        super().__init__()
+
+        self.encoder_dim = encoder_dim
+        self.encoder_rates = encoder_rates
+        self.decoder_dim = decoder_dim
+        self.decoder_rates = decoder_rates
+        self.sample_rate = sample_rate
+        self.continuous = continuous
+        self.use_weight_norm = use_weight_norm
+
+        if latent_dim is None:
+            latent_dim = encoder_dim * (2 ** len(encoder_rates))
+
+        self.latent_dim = latent_dim
+
+        self.hop_length = np.prod(encoder_rates)
+        self.encoder = Encoder(encoder_dim, encoder_rates, latent_dim)
+
+        if not continuous:
+            self.n_codebooks = n_codebooks
+            self.codebook_size = codebook_size
+            self.codebook_dim = codebook_dim
+            self.quantizer = ResidualVectorQuantize(
+                input_dim=latent_dim,
+                n_codebooks=n_codebooks,
+                codebook_size=codebook_size,
+                codebook_dim=codebook_dim,
+                quantizer_dropout=quantizer_dropout,
+            )
+        else:
+            self.quant_conv = torch.nn.Conv1d(latent_dim, 2 * latent_dim, 1)
+            self.post_quant_conv = torch.nn.Conv1d(latent_dim, latent_dim, 1)
+
+        self.decoder = Decoder(
+            latent_dim,
+            decoder_dim,
+            decoder_rates,
+        )
+        self.sample_rate = sample_rate
+        self.apply(init_weights)
+
+        self.delay = self.get_delay()
+
+        if not self.use_weight_norm:
+            self.remove_weight_norm()
+
+    @property
+    def dtype(self):
+        """Get the dtype of the model parameters."""
+        # Return the dtype of the first parameter found
+        for param in self.parameters():
+            return param.dtype
+        return torch.float32  # fallback
+
+    @property
+    def device(self):
+        """Get the device of the model parameters."""
+        # Return the device of the first parameter found
+        for param in self.parameters():
+            return param.device
+        return torch.device('cpu')  # fallback
+
+    def preprocess(self, audio_data, sample_rate):
+        if sample_rate is None:
+            sample_rate = self.sample_rate
+        assert sample_rate == self.sample_rate
+
+        length = audio_data.shape[-1]
+        right_pad = math.ceil(length / self.hop_length) * self.hop_length - length
+        audio_data = nn.functional.pad(audio_data, (0, right_pad))
+
+        return audio_data
+
+    @apply_forward_hook
+    def encode(
+        self,
+        audio_data: torch.Tensor,
+        n_quantizers: int = None,
+    ):
+        """Encode given audio data and return quantized latent codes
+
+        Parameters
+        ----------
+        audio_data : Tensor[B x 1 x T]
+            Audio data to encode
+        n_quantizers : int, optional
+            Number of quantizers to use, by default None
+            If None, all quantizers are used.
+
+        Returns
+        -------
+        dict
+            A dictionary with the following keys:
+            "z" : Tensor[B x D x T]
+                Quantized continuous representation of input
+            "codes" : Tensor[B x N x T]
+                Codebook indices for each codebook
+                (quantized discrete representation of input)
+            "latents" : Tensor[B x N*D x T]
+                Projected latents (continuous representation of input before quantization)
+            "vq/commitment_loss" : Tensor[1]
+                Commitment loss to train encoder to predict vectors closer to codebook
+                entries
+            "vq/codebook_loss" : Tensor[1]
+                Codebook loss to update the codebook
+            "length" : int
+                Number of samples in input audio
+        """
+        z = self.encoder(audio_data)  # [B x D x T]
+        if not self.continuous:
+            z, codes, latents, commitment_loss, codebook_loss = self.quantizer(z, n_quantizers)
+        else:
+            z = self.quant_conv(z)  # [B x 2D x T]
+            z = DiagonalGaussianDistribution(z)
+            codes, latents, commitment_loss, codebook_loss = None, None, 0, 0
+
+        return z, codes, latents, commitment_loss, codebook_loss
+
+    @apply_forward_hook
+    def decode(self, z: torch.Tensor):
+        """Decode given latent codes and return audio data
+
+        Parameters
+        ----------
+        z : Tensor[B x D x T]
+            Quantized continuous representation of input
+        length : int, optional
+            Number of samples in output audio, by default None
+
+        Returns
+        -------
+        dict
+            A dictionary with the following keys:
+            "audio" : Tensor[B x 1 x length]
+                Decoded audio data.
+        """
+        if not self.continuous:
+            audio = self.decoder(z)
+        else:
+            z = self.post_quant_conv(z)
+            audio = self.decoder(z)
+
+        return audio
+
+    def forward(
+        self,
+        audio_data: torch.Tensor,
+        sample_rate: int = None,
+        n_quantizers: int = None,
+    ):
+        """Model forward pass
+
+        Parameters
+        ----------
+        audio_data : Tensor[B x 1 x T]
+            Audio data to encode
+        sample_rate : int, optional
+            Sample rate of audio data in Hz, by default None
+            If None, defaults to `self.sample_rate`
+        n_quantizers : int, optional
+            Number of quantizers to use, by default None.
+            If None, all quantizers are used.
+
+        Returns
+        -------
+        dict
+            A dictionary with the following keys:
+            "z" : Tensor[B x D x T]
+                Quantized continuous representation of input
+            "codes" : Tensor[B x N x T]
+                Codebook indices for each codebook
+                (quantized discrete representation of input)
+            "latents" : Tensor[B x N*D x T]
+                Projected latents (continuous representation of input before quantization)
+            "vq/commitment_loss" : Tensor[1]
+                Commitment loss to train encoder to predict vectors closer to codebook
+                entries
+            "vq/codebook_loss" : Tensor[1]
+                Codebook loss to update the codebook
+            "length" : int
+                Number of samples in input audio
+            "audio" : Tensor[B x 1 x length]
+                Decoded audio data.
+        """
+        length = audio_data.shape[-1]
+        audio_data = self.preprocess(audio_data, sample_rate)
+        if not self.continuous:
+            z, codes, latents, commitment_loss, codebook_loss = self.encode(audio_data, n_quantizers)
+
+            x = self.decode(z)
+            return {
+                "audio": x[..., :length],
+                "z": z,
+                "codes": codes,
+                "latents": latents,
+                "vq/commitment_loss": commitment_loss,
+                "vq/codebook_loss": codebook_loss,
+            }
+        else:
+            posterior, _, _, _, _ = self.encode(audio_data, n_quantizers)
+            z = posterior.sample()
+            x = self.decode(z)
+
+            kl_loss = posterior.kl()
+            kl_loss = kl_loss.mean()
+
+            return {
+                "audio": x[..., :length],
+                "z": z,
+                "kl_loss": kl_loss,
+            }
+
+    def remove_weight_norm(self):
+        """
+        Remove weight_norm from all modules in the model.
+        This fuses the weight_g and weight_v parameters into a single weight parameter.
+        Should be called before inference for better performance.
+        Returns:
+            self: The model with weight_norm removed
+        """
+        from torch.nn.utils import remove_weight_norm
+        num_removed = 0
+        for name, module in list(self.named_modules()):
+            if hasattr(module, "_forward_pre_hooks"):
+                for hook_id, hook in list(module._forward_pre_hooks.items()):
+                    if "WeightNorm" in str(type(hook)):
+                        try:
+                            remove_weight_norm(module)
+                            num_removed += 1
+                            # print(f"Removed weight_norm from: {name}")
+                        except ValueError as e:
+                            print(f"Failed to remove weight_norm from {name}: {e}")
+        if num_removed > 0:
+            print(f"Successfully removed weight_norm from {num_removed} modules")
+            self.use_weight_norm = False
+            self.register_to_config(use_weight_norm=False)
+        else:
+            print("No weight_norm found in the model")
+        return self
+
+if __name__ == "__main__":
+    from functools import partial
+
+    import numpy as np
+
+    model = DAC().to("cpu")
+
+    for n, m in model.named_modules():
+        o = m.extra_repr()
+        p = sum([np.prod(p.size()) for p in m.parameters()])
+        fn = lambda o, p: o + f" {p/1e6:<.3f}M params."
+        setattr(m, "extra_repr", partial(fn, o=o, p=p))
+    print(model)
+    print("Total # of params: ", sum([np.prod(p.size()) for p in model.parameters()]))
+
+    length = 88200 * 2
+    x = torch.randn(1, 1, length).to(model.device)
+    x.requires_grad_(True)
+    x.retain_grad()
+
+    # Make a forward pass
+    out = model(x)["audio"]
+    print("Input shape:", x.shape)
+    print("Output shape:", out.shape)
+
+    # Create gradient variable
+    grad = torch.zeros_like(out)
+    grad[:, :, grad.shape[-1] // 2] = 1
+
+    # Make a backward pass
+    out.backward(grad)
+
+    # Check non-zero values
+    gradmap = x.grad.squeeze(0)
+    gradmap = (gradmap != 0).sum(0)  # sum across features
+    rf = (gradmap != 0).sum()
+
+    print(f"Receptive field: {rf.item()}")
+
+    x = AudioSignal(torch.randn(1, 1, 44100 * 60), 44100)
+    model.decompress(model.compress(x, verbose=True), verbose=True)

hymm/models/modules/interactionv2.py

@@ -0,0 +1,708 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Dict, List, Tuple, Optional
+
+from einops import rearrange
+
+from .wan_video_dit import RMSNorm, AttentionModule, flash_attention
+from hymm.utils.parallel_states import get_sequence_parallel_state, nccl_info, sp_seq_ctx
+from hymm.utils.communications import (
+    all_to_all_4D,
+    sp_pad_num_heads, sp_pad_heads, sp_num_local_real_heads,
+)
+
+class RotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor  # fix linting for `register_buffer`
+
+    def __init__(self, base: float, dim: int, device=None):
+        super().__init__()
+        self.base = base
+        self.dim = dim
+        self.attention_scaling = 1.0
+
+        inv_freq = 1.0 / (
+                    base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+
+    @torch.no_grad()
+    def forward(self, x, position_ids):
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2:]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+@torch.compile(fullgraph=True)
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`, *optional*):
+            Deprecated and unused.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+class PerFrameAttentionPooling(nn.Module):
+    """
+    Per-frame multi-head attention pooling.
+
+    Given a flattened token sequence [B, L, D] and grid size (T, H, W), perform a
+    single-query attention pooling over the H*W tokens for each time frame, producing
+    [B, T, D].
+
+    Inspired by SigLIP's Multihead Attention Pooling head (without MLP/residual stack).
+    """
+
+    def __init__(self, dim: int, num_heads: int, eps: float = 1e-6):
+        super().__init__()
+        assert dim % num_heads == 0, "dim must be divisible by num_heads"
+        self.dim = dim
+        self.num_heads = num_heads
+
+        self.probe = nn.Parameter(torch.randn(1, 1, dim))
+        nn.init.normal_(self.probe, std=0.02)
+
+        self.attention = nn.MultiheadAttention(embed_dim=dim, num_heads=num_heads, batch_first=True)
+        self.layernorm = nn.LayerNorm(dim, eps=eps)
+
+    def forward(self, x: torch.Tensor, grid_size: Tuple[int, int, int]) -> torch.Tensor:
+        """
+        Args:
+            x: [B, L, D], where L = T*H*W
+            grid_size: (T, H, W)
+        Returns:
+            pooled: [B, T, D]
+        """
+        B, L, D = x.shape
+        T, H, W = grid_size
+        assert D == self.dim, f"Channel dimension mismatch: D={D} vs dim={self.dim}"
+        assert L == T * H * W, f"Flattened length mismatch: L={L} vs T*H*W={T * H * W}"
+
+        S = H * W
+        # Re-arrange tokens grouped by frame.
+        x_bt_s_d = x.view(B, T, S, D).contiguous().view(B * T, S, D)  # [B*T, S, D]
+
+        # A learnable probe as the query (one query per frame).
+        probe = self.probe.expand(B * T, -1, -1)  # [B*T, 1, D]
+
+        # Attention pooling: query=probe, key/value=H*W tokens within the frame.
+        pooled_bt_1_d = self.attention(probe, x_bt_s_d, x_bt_s_d, need_weights=False)[0]  # [B*T, 1, D]
+        pooled_bt_d = pooled_bt_1_d.squeeze(1)  # [B*T, D]
+
+        # Restore to [B, T, D].
+        pooled = pooled_bt_d.view(B, T, D)
+        pooled = self.layernorm(pooled)
+        return pooled
+
+
+class CrossModalInteractionController:
+    """
+    Strategy class that controls interactions between two towers.
+    Manages the interaction mapping between visual DiT (e.g. 30 layers) and audio DiT (e.g. 30 layers).
+    """
+
+    def __init__(self, visual_layers: int = 30, audio_layers: int = 30):
+        self.visual_layers = visual_layers
+        self.audio_layers = audio_layers
+        self.min_layers = min(visual_layers, audio_layers)
+
+    def get_interaction_layers(self, strategy: str = "shallow_focus") -> Dict[str, List[Tuple[int, int]]]:
+        """
+        Get interaction layer mappings.
+
+        Args:
+            strategy: interaction strategy
+                - "shallow_focus": emphasize shallow layers to avoid deep-layer asymmetry
+                - "distributed": distributed interactions across the network
+                - "progressive": dense shallow interactions, sparse deeper interactions
+                - "custom": custom interaction layers
+
+        Returns:
+            A dict containing mappings for 'v2a' (visual -> audio) and 'a2v' (audio -> visual).
+        """
+
+        if strategy == "shallow_focus":
+            # Emphasize the first ~1/3 layers to avoid deep-layer asymmetry.
+            num_interact = min(10, self.min_layers // 3)
+            interact_layers = list(range(0, num_interact))
+
+        elif strategy == "distributed":
+            # Distribute interactions across the network (every few layers).
+            step = 3
+            interact_layers = list(range(0, self.min_layers, step))
+
+        elif strategy == "progressive":
+            # Progressive: dense shallow interactions, sparse deeper interactions.
+            shallow = list(range(0, min(8, self.min_layers)))  # Dense for the first 8 layers.
+            if self.min_layers > 8:
+                deep = list(range(8, self.min_layers, 3))  # Every 3 layers afterwards.
+                interact_layers = shallow + deep
+            else:
+                interact_layers = shallow
+
+        elif strategy == "custom":
+            # Custom strategy: adjust as needed.
+            interact_layers = [0, 2, 4, 6, 8, 12, 16, 20]  # Explicit layer indices.
+            interact_layers = [i for i in interact_layers if i < self.min_layers]
+
+        elif strategy == "full":
+            interact_layers = list(range(0, self.min_layers))
+
+        else:
+            raise ValueError(f"Unknown interaction strategy: {strategy}")
+
+        # Build bidirectional mapping.
+        mapping = {
+            'v2a': [(i, i) for i in interact_layers],  # visual layer i -> audio layer i
+            'a2v': [(i, i) for i in interact_layers]  # audio layer i -> visual layer i
+        }
+
+        return mapping
+
+    def should_interact(self, layer_idx: int, direction: str, interaction_mapping: Dict) -> bool:
+        """
+        Check whether a given layer should interact.
+
+        Args:
+            layer_idx: current layer index
+            direction: interaction direction ('v2a' or 'a2v')
+            interaction_mapping: interaction mapping table
+
+        Returns:
+            bool: whether to interact
+        """
+        if direction not in interaction_mapping:
+            return False
+
+        return any(src == layer_idx for src, _ in interaction_mapping[direction])
+
+
+class ConditionalCrossAttention(nn.Module):
+    def __init__(self, dim: int, kv_dim: int, num_heads: int, eps: float = 1e-6):
+        super().__init__()
+        self.q_dim = dim
+        self.kv_dim = kv_dim
+        self.num_heads = num_heads
+        self.head_dim = self.q_dim // num_heads
+
+        self.q = nn.Linear(dim, dim)
+        self.k = nn.Linear(kv_dim, dim)
+        self.v = nn.Linear(kv_dim, dim)
+        self.o = nn.Linear(dim, dim)
+        self.norm_q = RMSNorm(dim, eps=eps)
+        self.norm_k = RMSNorm(dim, eps=eps)
+
+        self.attn = AttentionModule(self.num_heads)
+
+        self.enable_bsa = False
+        self.bsa_params = {}
+
+        # Cross-modal bridge: q and k/v come from two differently-sized SP-sharded
+        # streams, so each needs its own global length for the Ulysses all-to-all.
+        # MOVABridge sets these per direction (a2v: q=visual/kv=audio, v2a: swapped).
+        self.sp_q_stream = "visual"
+        self.sp_kv_stream = "audio"
+
+    def _run_cross_bsa(self, q_bhsd, k_bhsd, v_bhsd, q_structure, k_structure,
+                       q_grid_size=None, k_grid_size=None):
+        """Run cross-modal BSA between video (3D) and audio (1D)."""
+        from hymm.models.modules.block_sparse_attention import flash_attn_bsa_cross
+        return flash_attn_bsa_cross(
+            q_bhsd, k_bhsd, v_bhsd,
+            q_structure=q_structure,
+            k_structure=k_structure,
+            q_grid_size=q_grid_size,
+            k_grid_size=k_grid_size,
+            sparsity=self.bsa_params.get('sparsity', 0.9375),
+            cdf_threshold=self.bsa_params.get('cdf_threshold', None),
+            chunk_size_q=self.bsa_params.get('chunk_size_q', 64),
+            chunk_size_k=self.bsa_params.get('chunk_size_k', 64),
+            chunk_3d_shape_q=self.bsa_params.get('chunk_3d_shape_q', [4, 4, 4]),
+            chunk_3d_shape_k=self.bsa_params.get('chunk_3d_shape_k', [4, 4, 4]),
+        )
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor, x_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+                y_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+                q_structure: str = "1d", k_structure: str = "1d",
+                q_grid_size=None, k_grid_size=None):
+        ctx = y
+        q = self.norm_q(self.q(x))
+        k = self.norm_k(self.k(ctx))
+        v = self.v(ctx)
+        if x_freqs is not None:
+            x_cos, x_sin = x_freqs
+            q_view = rearrange(q, 'b l (h d) -> b l h d', d=self.head_dim)
+            x_cos = x_cos.to(q_view.dtype).to(q_view.device)
+            x_sin = x_sin.to(q_view.dtype).to(q_view.device)
+            q_view, _ = apply_rotary_pos_emb(q_view, q_view, x_cos, x_sin, unsqueeze_dim=2)
+            q = rearrange(q_view, 'b l h d -> b l (h d)')
+        if y_freqs is not None:
+            y_cos, y_sin = y_freqs
+            k_view = rearrange(k, 'b l (h d) -> b l h d', d=self.head_dim)
+            y_cos = y_cos.to(k_view.dtype).to(k_view.device)
+            y_sin = y_sin.to(k_view.dtype).to(k_view.device)
+            _, k_view = apply_rotary_pos_emb(k_view, k_view, y_cos, y_sin, unsqueeze_dim=2)
+            k = rearrange(k_view, 'b l h d -> b l (h d)')
+
+        use_bsa = self.enable_bsa and self.bsa_params.get('sparsity', 0) > 0
+
+        if get_sequence_parallel_state():
+            sp_size = nccl_info.sp_size
+            sp_rank = nccl_info.rank_within_spgroup
+            # Same Ulysses head-padding contract as the DiT attentions. v2a bridge
+            # operates in the audio 12-head space, so with sp_size=8 the head dim is
+            # padded to 16; padded heads are sliced off before the BSA cross kernel
+            # / flash and truncated after, fully isolated from the real heads.
+            padded_heads = sp_pad_num_heads(self.num_heads, sp_size)
+
+            q, k, v = [rearrange(t, "b s (n d) -> b s n d", n=self.num_heads) for t in (q, k, v)]
+            if padded_heads != self.num_heads:
+                q = sp_pad_heads(q, padded_heads)
+                k = sp_pad_heads(k, padded_heads)
+                v = sp_pad_heads(v, padded_heads)
+
+            q_full_len = sp_seq_ctx.get(self.sp_q_stream)
+            kv_full_len = sp_seq_ctx.get(self.sp_kv_stream)
+            q = all_to_all_4D(q, nccl_info.sp_group, scatter_dim=2, gather_dim=1,
+                              full_seq_len=q_full_len)
+            k = all_to_all_4D(k, nccl_info.sp_group, scatter_dim=2, gather_dim=1,
+                              full_seq_len=kv_full_len)
+            v = all_to_all_4D(v, nccl_info.sp_group, scatter_dim=2, gather_dim=1,
+                              full_seq_len=kv_full_len)
+
+            hpr = q.shape[2]
+            local_heads = sp_num_local_real_heads(self.num_heads, padded_heads, sp_size, sp_rank)
+            if local_heads > 0:
+                if local_heads != hpr:
+                    q = q[:, :, :local_heads, :]
+                    k = k[:, :, :local_heads, :]
+                    v = v[:, :, :local_heads, :]
+                if use_bsa:
+                    q_bsa = q.permute(0, 2, 1, 3).contiguous()
+                    k_bsa = k.permute(0, 2, 1, 3).contiguous()
+                    v_bsa = v.permute(0, 2, 1, 3).contiguous()
+                    x_bsa = self._run_cross_bsa(q_bsa, k_bsa, v_bsa, q_structure, k_structure, q_grid_size, k_grid_size)
+                    x = rearrange(x_bsa, "b n s d -> b s (n d)")
+                else:
+                    x = flash_attention(q.flatten(2), k.flatten(2), v.flatten(2), num_heads=local_heads)
+                x = rearrange(x, "b s (n d) -> b s n d", n=local_heads)
+                if x.shape[2] != hpr:
+                    x = sp_pad_heads(x, hpr)
+            else:
+                # Padding-only rank: no BSA cross / flash kernel is run (padding never
+                # enters the kernel). Keep x CONNECTED to q/k/v (value forced to zero)
+                # so the q/k/v backward all_to_all collectives still fire here, keeping
+                # NCCL symmetric and routing this rank's real-token grads back correctly.
+                # Cross-attention: q and k/v have DIFFERENT seq lengths, so connect k/v
+                # via scalar reductions (broadcast-safe). Output takes q's seq shape.
+                x = q * 0.0 + (k.sum() + v.sum()) * 0.0
+            x = all_to_all_4D(x, nccl_info.sp_group, scatter_dim=1, gather_dim=2,
+                              full_seq_len=q_full_len)
+            if padded_heads != self.num_heads:
+                x = x[:, :, :self.num_heads, :]
+            x = x.flatten(2)
+        else:
+            if use_bsa:
+                q_bsa = rearrange(q, "b s (n d) -> b n s d", n=self.num_heads).contiguous()
+                k_bsa = rearrange(k, "b s (n d) -> b n s d", n=self.num_heads).contiguous()
+                v_bsa = rearrange(v, "b s (n d) -> b n s d", n=self.num_heads).contiguous()
+                x_bsa = self._run_cross_bsa(q_bsa, k_bsa, v_bsa, q_structure, k_structure, q_grid_size, k_grid_size)
+                x = rearrange(x_bsa, "b n s d -> b s (n d)")
+            else:
+                x = self.attn(q, k, v)
+
+        return self.o(x)
+
+
+# from diffusers.models.attention import AdaLayerNorm
+class AdaLayerNorm(nn.Module):
+    r"""
+    Norm layer modified to incorporate timestep embeddings.
+
+    Parameters:
+        embedding_dim (`int`): The size of each embedding vector.
+        num_embeddings (`int`, *optional*): The size of the embeddings dictionary.
+        output_dim (`int`, *optional*):
+        norm_elementwise_affine (`bool`, defaults to `False):
+        norm_eps (`bool`, defaults to `False`):
+        chunk_dim (`int`, defaults to `0`):
+    """
+
+    def __init__(
+            self,
+            embedding_dim: int,
+            num_embeddings: Optional[int] = None,
+            output_dim: Optional[int] = None,
+            norm_elementwise_affine: bool = False,
+            norm_eps: float = 1e-5,
+            chunk_dim: int = 0,
+    ):
+        super().__init__()
+
+        self.chunk_dim = chunk_dim
+        output_dim = output_dim or embedding_dim * 2
+
+        if num_embeddings is not None:
+            self.emb = nn.Embedding(num_embeddings, embedding_dim)
+        else:
+            self.emb = None
+
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(embedding_dim, output_dim)
+        self.norm = nn.LayerNorm(output_dim // 2, norm_eps, norm_elementwise_affine)
+
+    def forward(
+            self, x: torch.Tensor, timestep: Optional[torch.Tensor] = None, temb: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        if self.emb is not None:
+            temb = self.emb(timestep)
+
+        temb = self.linear(self.silu(temb))
+
+        if self.chunk_dim == 2:
+            scale, shift = temb.chunk(2, dim=2)
+            # print(f"{x.shape = }, {scale.shape = }, {shift.shape = }")
+        elif self.chunk_dim == 1:
+            # This is a bit weird why we have the order of "shift, scale" here and "scale, shift" in the
+            # other if-branch. This branch is specific to CogVideoX and OmniGen for now.
+            shift, scale = temb.chunk(2, dim=1)
+            shift = shift[:, None, :]
+            scale = scale[:, None, :]
+        else:
+            scale, shift = temb.chunk(2, dim=0)
+
+        x = self.norm(x) * (1 + scale) + shift
+        return x
+
+
+class ConditionalCrossAttentionBlock(nn.Module):
+    """
+    A thin wrapper around ConditionalCrossAttention.
+    Applies LayerNorm to the conditioning input `y` before cross-attention.
+    """
+
+    def __init__(self, dim: int, kv_dim: int, num_heads: int, eps: float = 1e-6, pooled_adaln: bool = False):
+        super().__init__()
+        self.y_norm = nn.LayerNorm(kv_dim, eps=eps)
+        self.inner = ConditionalCrossAttention(dim=dim, kv_dim=kv_dim, num_heads=num_heads, eps=eps)
+        self.pooled_adaln = pooled_adaln
+        if pooled_adaln:
+            self.per_frame_pooling = PerFrameAttentionPooling(kv_dim, num_heads=num_heads, eps=eps)
+            self.adaln = AdaLayerNorm(kv_dim, output_dim=dim * 2, chunk_dim=2)
+
+    def set_sp_streams(self, q_stream: str, kv_stream: str):
+        """Declare which SP-sharded stream feeds q and which feeds k/v."""
+        self.inner.sp_q_stream = q_stream
+        self.inner.sp_kv_stream = kv_stream
+
+    def forward(
+            self,
+            x: torch.Tensor,
+            y: torch.Tensor,
+            x_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+            y_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+            video_grid_size: Optional[Tuple[int, int, int]] = None,
+            q_structure: str = "1d",
+            k_structure: str = "1d",
+            q_grid_size=None,
+            k_grid_size=None,
+    ) -> torch.Tensor:
+        if self.pooled_adaln:
+            assert not get_sequence_parallel_state() or nccl_info.sp_size == 1, (
+                "pooled_adaln requires full-length y for per-frame pooling, "
+                "which is incompatible with SP > 1. Set sp_size=1 or disable pooled_adaln."
+            )
+            assert video_grid_size is not None, "video_grid_size must not be None"
+            pooled_y = self.per_frame_pooling(y, video_grid_size)
+            if pooled_y.shape[1] != x.shape[1]:
+                pooled_y = F.interpolate(
+                    pooled_y.permute(0, 2, 1),
+                    size=x.shape[1],
+                    mode='linear',
+                    align_corners=False,
+                ).permute(0, 2, 1)
+            x = self.adaln(x, temb=pooled_y)
+        y = self.y_norm(y)
+        return self.inner(x=x, y=y, x_freqs=x_freqs, y_freqs=y_freqs,
+                          q_structure=q_structure, k_structure=k_structure,
+                          q_grid_size=q_grid_size, k_grid_size=k_grid_size)
+
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+
+
+class DualTowerConditionalBridge(ModelMixin, ConfigMixin):
+    """
+    Dual-tower conditional bridge.
+    """
+    _repeated_blocks = ("ConditionalCrossAttentionBlock",)
+
+    @register_to_config
+    def __init__(self,
+                 visual_layers: int = 30,
+                 audio_layers: int = 30,
+                 visual_hidden_dim: int = 3072,  # visual DiT hidden state dimension
+                 audio_hidden_dim: int = 1536,  # audio DiT hidden state dimension
+                 audio_fps: float = 44100.0 / 2048.0,
+                 head_dim: int = 128,  # attention head dimension
+                 interaction_strategy: str = "shallow_focus",
+                 apply_cross_rope: bool = False,  # whether to apply RoPE in cross-attention
+                 apply_first_frame_bias_in_rope: bool = False,
+                 # whether to account for 1/video_fps bias for the first frame in RoPE alignment
+                 trainable_condition_scale: bool = False,
+                 pooled_adaln: bool = False,
+                 ):
+        super().__init__()
+
+        self.visual_hidden_dim = visual_hidden_dim
+        self.audio_hidden_dim = audio_hidden_dim
+        self.audio_fps = audio_fps
+        self.head_dim = head_dim
+        self.apply_cross_rope = apply_cross_rope
+        self.apply_first_frame_bias_in_rope = apply_first_frame_bias_in_rope
+        self.trainable_condition_scale = trainable_condition_scale
+        self.pooled_adaln = pooled_adaln
+        if self.trainable_condition_scale:
+            self.condition_scale = nn.Parameter(torch.tensor([1.0], dtype=torch.float32))
+        else:
+            self.condition_scale = 1.0
+
+        self.controller = CrossModalInteractionController(visual_layers, audio_layers)
+        self.interaction_mapping = self.controller.get_interaction_layers(interaction_strategy)
+
+        # Conditional cross-attention modules operating at the DiT hidden-state level.
+        self.audio_to_video_conditioners = nn.ModuleDict()  # audio hidden states -> visual DiT conditioning
+        self.video_to_audio_conditioners = nn.ModuleDict()  # visual hidden states -> audio DiT conditioning
+
+        # Build conditioners for layers that should interact.
+        # audio hidden states condition the visual DiT
+        self.rotary = RotaryEmbedding(base=10000.0, dim=head_dim)
+        for v_layer, _ in self.interaction_mapping['a2v']:
+            self.audio_to_video_conditioners[str(v_layer)] = ConditionalCrossAttentionBlock(
+                dim=visual_hidden_dim,  # 3072 (visual DiT hidden states)
+                kv_dim=audio_hidden_dim,  # 1536 (audio DiT hidden states)
+                num_heads=visual_hidden_dim // head_dim,  # derive number of heads from hidden dim
+                pooled_adaln=False  # a2v typically does not need pooled AdaLN
+            )
+
+        # visual hidden states condition the audio DiT
+        for a_layer, _ in self.interaction_mapping['v2a']:
+            self.video_to_audio_conditioners[str(a_layer)] = ConditionalCrossAttentionBlock(
+                dim=audio_hidden_dim,  # 1536 (audio DiT hidden states)
+                kv_dim=visual_hidden_dim,  # 3072 (visual DiT hidden states)
+                num_heads=audio_hidden_dim // head_dim,  # safe head count derivation
+                pooled_adaln=self.pooled_adaln
+            )
+
+    @torch.no_grad()
+    def build_aligned_freqs(self,
+                            video_fps: float,
+                            grid_size: Tuple[int, int, int],
+                            audio_steps: int,
+                            device: Optional[torch.device] = None,
+                            dtype: Optional[torch.dtype] = None) -> Tuple[
+        Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:
+        """
+        Build aligned RoPE (cos, sin) pairs based on video fps, video grid size (f_v, h, w),
+        and audio sequence length `audio_steps` (with fixed audio fps = 44100/2048).
+
+        Returns:
+            visual_freqs: (cos_v, sin_v), shape [1, f_v*h*w, head_dim]
+            audio_freqs:  (cos_a, sin_a), shape [1, audio_steps, head_dim]
+        """
+        f_v, h, w = grid_size
+        L_v = f_v * h * w
+        L_a = int(audio_steps)
+
+        device = device or next(self.parameters()).device
+        dtype = dtype or torch.float32
+
+        # Audio positions: 0,1,2,...,L_a-1 (audio as reference).
+        audio_pos = torch.arange(L_a, device=device, dtype=torch.float32).unsqueeze(0)
+
+        # Video positions: align video frames to audio-step units.
+        # FIXME(dhyu): hard-coded VAE temporal stride = 4
+        if self.apply_first_frame_bias_in_rope:
+            # Account for the "first frame lasts 1/video_fps" bias.
+            video_effective_fps = float(video_fps) / 4.0
+            if f_v > 0:
+                t_starts = torch.zeros((f_v,), device=device, dtype=torch.float32)
+                if f_v > 1:
+                    t_starts[1:] = (1.0 / float(video_fps)) + torch.arange(f_v - 1, device=device,
+                                                                           dtype=torch.float32) * (
+                                               1.0 / video_effective_fps)
+            else:
+                t_starts = torch.zeros((0,), device=device, dtype=torch.float32)
+            # Convert to audio-step units.
+            video_pos_per_frame = t_starts * float(self.audio_fps)
+        else:
+            # No first-frame bias: uniform alignment.
+            scale = float(self.audio_fps) / float(video_fps / 4.0)
+            video_pos_per_frame = torch.arange(f_v, device=device, dtype=torch.float32) * scale
+        # Flatten to f*h*w; tokens within the same frame share the same time position.
+        video_pos = video_pos_per_frame.repeat_interleave(h * w).unsqueeze(0)
+
+        # print(f"video fps: {video_fps}, audio fps: {self.audio_fps}, scale: {scale}")
+        # print(f"video pos: {video_pos.shape}, audio pos: {audio_pos.shape}")
+
+        # Build dummy x to produce cos/sin, dim=head_dim.
+        dummy_v = torch.zeros((1, L_v, self.head_dim), device=device, dtype=dtype)
+        dummy_a = torch.zeros((1, L_a, self.head_dim), device=device, dtype=dtype)
+
+        cos_v, sin_v = self.rotary(dummy_v, position_ids=video_pos)
+        cos_a, sin_a = self.rotary(dummy_a, position_ids=audio_pos)
+
+        return (cos_v, sin_v), (cos_a, sin_a)
+
+    def should_interact(self, layer_idx: int, direction: str) -> bool:
+        return self.controller.should_interact(layer_idx, direction, self.interaction_mapping)
+
+    def apply_conditional_control(self,
+                                  layer_idx: int,
+                                  direction: str,
+                                  primary_hidden_states: torch.Tensor,
+                                  condition_hidden_states: torch.Tensor,
+                                  x_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+                                  y_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+                                  condition_scale: Optional[float] = None,
+                                  video_grid_size: Optional[Tuple[int, int, int]] = None) -> torch.Tensor:
+        """
+        Apply conditional control (at the DiT hidden-state level).
+
+        Args:
+            layer_idx: current layer index
+            direction: conditioning direction
+                - 'a2v': audio hidden states -> visual DiT
+                - 'v2a': visual hidden states -> audio DiT
+            primary_hidden_states: primary DiT hidden states [B, L, hidden_dim]
+            condition_hidden_states: condition DiT hidden states [B, L, hidden_dim]
+            condition_scale: conditioning strength (similar to CFG scale)
+
+        Returns:
+            Conditioned primary DiT hidden states [B, L, hidden_dim]
+        """
+
+        if not self.controller.should_interact(layer_idx, direction, self.interaction_mapping):
+            return primary_hidden_states
+
+        if direction == 'a2v':
+            # audio hidden states condition the visual DiT
+            conditioner = self.audio_to_video_conditioners[str(layer_idx)]
+
+        elif direction == 'v2a':
+            # visual hidden states condition the audio DiT
+            conditioner = self.video_to_audio_conditioners[str(layer_idx)]
+        else:
+            raise ValueError(f"Invalid direction: {direction}")
+
+        conditioned_features = conditioner(
+            x=primary_hidden_states,
+            y=condition_hidden_states,
+            x_freqs=x_freqs,
+            y_freqs=y_freqs,
+            video_grid_size=video_grid_size,
+        )
+
+        if self.trainable_condition_scale and condition_scale is not None:
+            print(
+                "[WARN] This model has a trainable condition_scale, but an external "
+                f"condition_scale={condition_scale} was provided. The trainable condition_scale "
+                "will be ignored in favor of the external value."
+            )
+
+        scale = condition_scale if condition_scale is not None else self.condition_scale
+
+        primary_hidden_states = primary_hidden_states + conditioned_features * scale
+
+        return primary_hidden_states
+
+    def forward(
+            self,
+            layer_idx: int,
+            visual_hidden_states: torch.Tensor,
+            audio_hidden_states: torch.Tensor,
+            *,
+            x_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+            y_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+            a2v_condition_scale: Optional[float] = None,
+            v2a_condition_scale: Optional[float] = None,
+            condition_scale: Optional[float] = None,
+            video_grid_size: Optional[Tuple[int, int, int]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Apply bidirectional conditional control to both visual/audio towers.
+
+        Args:
+            layer_idx: current layer index
+            visual_hidden_states: visual DiT hidden states
+            audio_hidden_states: audio DiT hidden states
+            x_freqs / y_freqs: cross-modal RoPE (cos, sin) pairs.
+                If provided, x_freqs is assumed to correspond to the primary tower and y_freqs
+                to the conditioning tower.
+            a2v_condition_scale: audio->visual conditioning strength (overrides global condition_scale)
+            v2a_condition_scale: visual->audio conditioning strength (overrides global condition_scale)
+            condition_scale: fallback conditioning strength when per-direction scale is None
+            video_grid_size: (F, H, W), used on the audio side when pooled_adaln is enabled
+
+        Returns:
+            (visual_hidden_states, audio_hidden_states), both conditioned in their respective directions.
+        """
+
+        visual_conditioned = self.apply_conditional_control(
+            layer_idx=layer_idx,
+            direction="a2v",
+            primary_hidden_states=visual_hidden_states,
+            condition_hidden_states=audio_hidden_states,
+            x_freqs=x_freqs,
+            y_freqs=y_freqs,
+            condition_scale=a2v_condition_scale if a2v_condition_scale is not None else condition_scale,
+            video_grid_size=video_grid_size,
+        )
+
+        audio_conditioned = self.apply_conditional_control(
+            layer_idx=layer_idx,
+            direction="v2a",
+            primary_hidden_states=audio_hidden_states,
+            condition_hidden_states=visual_hidden_states,
+            x_freqs=y_freqs,
+            y_freqs=x_freqs,
+            condition_scale=v2a_condition_scale if v2a_condition_scale is not None else condition_scale,
+            video_grid_size=video_grid_size,
+        )
+
+        return visual_conditioned, audio_conditioned

hymm/models/modules/mova.py

@@ -0,0 +1,964 @@
+import gc
+import json
+import os
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint as torch_checkpoint
+from diffusers import ModelMixin, ConfigMixin
+
+from hymm.models.modules.wan_video_dit import (
+    DiTBlock,
+    WanModel,
+    sinusoidal_embedding_1d,
+    advance_dynamic_block_pass_id,
+)
+from hymm.models.modules.wan_audio_dit import WanAudioModel
+from hymm.models.modules.interactionv2 import (
+    ConditionalCrossAttentionBlock,
+    DualTowerConditionalBridge,
+)
+from hymm.utils.parallel_states import get_sequence_parallel_state, nccl_info, sp_seq_ctx
+from hymm.utils.communications import all_gather, sp_split, assert_sp_splittable
+
+
+def assemble_visual_freqs(freqs_tuple, t, h, w, device):
+    """Assemble 3D RoPE complex-valued freqs for the video DiT.
+
+    Matches the inline assembly in WanModel.forward() and
+    MOVA/mova/diffusion/pipelines/pipeline_mova.py.
+
+    Args:
+        freqs_tuple: (f_freqs_cis, h_freqs_cis, w_freqs_cis) from
+            precompute_freqs_cis_3d.  Each is a complex tensor.
+        t, h, w: temporal, height, width grid dimensions (after patchify).
+        device: target device.
+
+    Returns:
+        Complex tensor of shape [t*h*w, 1, head_dim//2].
+    """
+    freqs = tuple(f.to(device) for f in freqs_tuple)
+    return torch.cat([
+        freqs[0][:t].view(t, 1, 1, -1).expand(t, h, w, -1),
+        freqs[1][:h].view(1, h, 1, -1).expand(t, h, w, -1),
+        freqs[2][:w].view(1, 1, w, -1).expand(t, h, w, -1),
+    ], dim=-1).reshape(t * h * w, 1, -1)
+
+
+def assemble_audio_freqs(freqs_tuple, f, device):
+    """Assemble 1D RoPE complex-valued freqs for the audio DiT.
+
+    Matches the inline assembly in WanAudioModel.forward() and
+    MOVA/mova/diffusion/pipelines/pipeline_mova.py.
+
+    Args:
+        freqs_tuple: (f0, f1, f2) from precompute_freqs_cis.
+        f: sequence length (after audio patchify).
+        device: target device.
+
+    Returns:
+        Complex tensor of shape [f, 1, head_dim//2].
+    """
+    return torch.cat([
+        freqs_tuple[0][:f].view(f, -1).expand(f, -1),
+        freqs_tuple[1][:f].view(f, -1).expand(f, -1),
+        freqs_tuple[2][:f].view(f, -1).expand(f, -1),
+    ], dim=-1).reshape(f, 1, -1).to(device)
+
+
+
+class FusedMOVABlock(nn.Module):
+
+    def __init__(self, video_block, audio_block,
+                 a2v_conditioner=None, v2a_conditioner=None):
+        super().__init__()
+        self.video_block = video_block
+        self.audio_block = audio_block
+        self.a2v_conditioner = a2v_conditioner
+        self.v2a_conditioner = v2a_conditioner
+
+        # Mark inner DiTBlocks so apply_fsdp_checkpointing skips them,
+        # avoiding nested checkpointing (outer FusedMOVABlock already covers them).
+        self.video_block._inside_fused_block = True
+        self.audio_block._inside_fused_block = True
+
+        # Set to True by _apply_mova_fsdp_checkpointing (only when the YAML trigger fine_grained_gc is enabled) for blocks that ALSO get the outer block-level checkpoint
+        # wrapper. When on (and training), the a2v/v2a bridge submodules are
+        # additionally inner-checkpointed (deferred) so they are NOT co-resident
+        # with the video block at the video-FFN recompute peak. The video/audio
+        # DiT blocks run inline (Version B — see forward()). Pure recompute-
+        # schedule change: fwd/bwd numerics and loss are bit-identical to the
+        # default whole-block checkpointing.
+        self.fine_grained_gc = False
+
+    def forward(
+        self,
+        visual_x, audio_x,
+        visual_context, audio_context,
+        visual_t_mod, audio_t_mod,
+        visual_freqs, audio_freqs,
+        visual_rope_cos_sin=None, audio_rope_cos_sin=None,
+        a2v_scale=1.0, v2a_scale=1.0,
+        video_grid_size=None,
+        override_video_block=None,
+        timestep_ratio=None,
+    ):
+        # Fine-grained (nested) checkpointing helper. wraps ONLY the a2v/v2a bridges; the video/audio
+        # blocks run inline. All arguments are passed POSITIONALLY in
+        # the submodule's signature order, so the computation is byte-for-byte identical to the un-checkpointed calls (checkpointing only re-runs forward; there are no RNG ops here).
+        inner_gc = self.training and self.fine_grained_gc
+
+        def _maybe_ckpt(fn, *args):
+            if inner_gc:
+                return torch_checkpoint.checkpoint(fn, *args, use_reentrant=False)
+            return fn(*args)
+
+        # Bridge conditioning — must use ORIGINAL states for both directions,
+        # matching DualTowerConditionalBridge.forward which computes a2v and v2a
+        # from the same unmodified inputs before returning both conditioned outputs.
+        visual_x_orig = visual_x
+
+        a2v_bridge_residual = None
+        if self.a2v_conditioner is not None:
+            # Positional order matches ConditionalCrossAttentionBlock.forward:
+            # (x, y, x_freqs, y_freqs, video_grid_size, q_structure, k_structure,
+            #  q_grid_size, k_grid_size)
+            cond = _maybe_ckpt(
+                self.a2v_conditioner,
+                visual_x, audio_x,
+                visual_rope_cos_sin, audio_rope_cos_sin,
+                video_grid_size,
+                "3d", "1d",
+                video_grid_size, None,
+            )
+            # Store the raw bridge residual (before scaling) for audio guidance
+            a2v_bridge_residual = cond
+            visual_x = visual_x + cond * a2v_scale
+
+        if self.v2a_conditioner is not None:
+            cond = _maybe_ckpt(
+                self.v2a_conditioner,
+                audio_x, visual_x_orig,
+                audio_rope_cos_sin, visual_rope_cos_sin,
+                video_grid_size,
+                "1d", "3d",
+                None, video_grid_size,
+            )
+            audio_x = audio_x + cond * v2a_scale
+
+        # DiT blocks (allow video block override for video_dit_2).
+        # If this OOMs -> wrapping BOTH calls below in _maybe_ckpt(...) (positional order as in the bridges above).
+        # Numerically bit-identical either way (checkpointing only re-runs fwd).
+        vblock = override_video_block if override_video_block is not None else self.video_block
+
+        visual_x = vblock(
+            visual_x, visual_context, visual_t_mod, visual_freqs,
+            grid_size=video_grid_size,
+            a2v_bridge_residual=a2v_bridge_residual,
+            timestep_ratio=timestep_ratio,
+        )
+        audio_x = self.audio_block(audio_x, audio_context, audio_t_mod, audio_freqs)
+
+        return visual_x, audio_x
+
+
+# ---------------------------------------------------------------------------
+# MOVABridge — top-level wrapper
+# ---------------------------------------------------------------------------
+
+class MOVABridge(ModelMixin, ConfigMixin):
+
+    def __init__(
+        self,
+        video_dit: WanModel,
+        video_dit_2: Optional[WanModel],
+        audio_dit: WanAudioModel,
+        dual_tower_bridge: DualTowerConditionalBridge,
+        boundary_ratio: float = 0.9,
+        gradient_checkpoint: bool = False,
+    ):
+        super().__init__()
+        self.boundary_ratio = boundary_ratio
+        self.gradient_checkpoint = gradient_checkpoint
+
+        min_layers = min(len(video_dit.blocks), len(audio_dit.blocks))
+
+        # --- Build fused blocks (video + audio + bridge per layer) ---
+        self.fusion_blocks = nn.ModuleList()
+        for i in range(min_layers):
+            key = str(i)
+            a2v = dual_tower_bridge.audio_to_video_conditioners[key] if key in dual_tower_bridge.audio_to_video_conditioners else None
+            v2a = dual_tower_bridge.video_to_audio_conditioners[key] if key in dual_tower_bridge.video_to_audio_conditioners else None
+            fused = FusedMOVABlock(
+                video_block=video_dit.blocks[i],
+                audio_block=audio_dit.blocks[i],
+                a2v_conditioner=a2v,
+                v2a_conditioner=v2a,
+            )
+            self.fusion_blocks.append(fused)
+
+        # Remaining video-only blocks (primary dit, beyond min_layers)
+        self.remaining_video_blocks = nn.ModuleList(
+            [video_dit.blocks[i] for i in range(min_layers, len(video_dit.blocks))]
+        )
+
+        # --- Clear extracted blocks to avoid double registration ---
+        video_dit.blocks = nn.ModuleList()
+        audio_dit.blocks = nn.ModuleList()
+        dual_tower_bridge.audio_to_video_conditioners = nn.ModuleDict()
+        dual_tower_bridge.video_to_audio_conditioners = nn.ModuleDict()
+
+        # Keep DiTs for utility functions (patchify, unpatchify, head, embeddings, freqs)
+        self.video_dit = video_dit
+        self.audio_dit = audio_dit
+
+        # Keep bridge for rotary / build_aligned_freqs / condition_scale
+        self.dual_tower_bridge = dual_tower_bridge
+
+        # Secondary video dit (kept as-is with its own blocks, not fused).
+        # Mark override blocks (0..min_layers-1) so they are skipped by
+        # apply_fsdp_checkpointing — they run inside a checkpointed
+        # FusedMOVABlock and must not be double-wrapped.
+        if video_dit_2 is not None:
+            for i in range(min(min_layers, len(video_dit_2.blocks))):
+                video_dit_2.blocks[i]._inside_fused_block = True
+        self.video_dit_2 = video_dit_2
+
+        # When True, ALL sparse-attention features (top-k/top-p BSA, audio guidance,
+        # variance guidance, taylor/rectified bias correction, dynamic block shape)
+        # are applied ONLY to the high-noise expert (video_dit = fusion_blocks +
+        # remaining_video_blocks). The low-noise expert (video_dit_2) then keeps the
+        # backbone's original dense full attention. Default False = sparse on BOTH
+        # experts (legacy behavior). Plain python attr (not an FSDP param/buffer).
+        self.sparse_high_noise_only = False
+
+        self._tag_sp_streams()
+
+        gc.collect()
+
+    def _tag_sp_streams(self):
+        """Tell every attention which SP-sharded stream its tensors belong to.
+
+        Under sequence parallelism the visual and audio towers are chunked to
+        different lengths, and the Ulysses all-to-all has to undo the exact
+        ``torch.chunk`` split of the stream it is handling. Attentions default to
+        "visual", so only the audio tower and the two bridge directions need
+        retagging here.
+        """
+        from hymm.models.modules.wan_audio_dit import tag_audio_sp_stream
+
+        for fused in self.fusion_blocks:
+            tag_audio_sp_stream([fused.audio_block])
+            if fused.a2v_conditioner is not None:
+                # a2v: query = visual tokens, context = audio tokens
+                fused.a2v_conditioner.set_sp_streams("visual", "audio")
+            if fused.v2a_conditioner is not None:
+                # v2a: query = audio tokens, context = visual tokens
+                fused.v2a_conditioner.set_sp_streams("audio", "visual")
+
+    @classmethod
+    def from_mova_pretrained(cls, ckpt_path, torch_dtype=torch.bfloat16, logger=None):
+        """
+        Load from a MOVA HuggingFace-style checkpoint directory.
+
+        The directory should contain subfolders:
+            video_dit/, video_dit_2/, audio_dit/, dual_tower_bridge/,
+            video_vae/, audio_vae/, text_encoder/, tokenizer/, scheduler/
+
+        Returns:
+            (mova_bridge, extra_components) where extra_components is a dict
+            containing {video_vae, audio_vae, text_encoder, tokenizer, scheduler,
+            boundary_ratio, audio_vae_type}.
+        """
+        from diffusers.models.autoencoders import AutoencoderKLWan
+        from transformers import T5TokenizerFast, UMT5EncoderModel
+
+        from hymm.models.modules.dac_vae import DAC
+        from hymm.diffusion.schedulers.flow_match_pair import FlowMatchPairScheduler
+
+        _log = logger.info if logger else print
+
+        import time as _time
+        _load_start = _time.time()
+        _log(f"[MOVABridge] Loading from {ckpt_path}")
+
+        # Trainable components: load in requested dtype on CPU
+        _t0 = _time.time()
+        video_dit = WanModel.from_pretrained(
+            ckpt_path, subfolder="video_dit", torch_dtype=torch_dtype,
+        ).cpu()
+        _log(f"[MOVABridge] Loaded video_dit ({_time.time()-_t0:.1f}s)")
+
+        video_dit_2_path = os.path.join(ckpt_path, "video_dit_2")
+        if os.path.isdir(video_dit_2_path):
+            _t0 = _time.time()
+            video_dit_2 = WanModel.from_pretrained(
+                ckpt_path, subfolder="video_dit_2", torch_dtype=torch_dtype,
+            ).cpu()
+            _log(f"[MOVABridge] Loaded video_dit_2 ({_time.time()-_t0:.1f}s)")
+        else:
+            video_dit_2 = None
+            _log("[MOVABridge] No video_dit_2 found, skipping")
+
+        _t0 = _time.time()
+        audio_dit = WanAudioModel.from_pretrained(
+            ckpt_path, subfolder="audio_dit", torch_dtype=torch_dtype,
+        ).cpu()
+        _log(f"[MOVABridge] Loaded audio_dit ({_time.time()-_t0:.1f}s)")
+
+        _t0 = _time.time()
+        dual_tower_bridge = DualTowerConditionalBridge.from_pretrained(
+            ckpt_path, subfolder="dual_tower_bridge", torch_dtype=torch_dtype,
+        ).cpu()
+        _log(f"[MOVABridge] Loaded dual_tower_bridge ({_time.time()-_t0:.1f}s)")
+
+        # Frozen components: load in bf16 on CPU to save memory
+        frozen_dtype = torch.bfloat16
+
+        _t0 = _time.time()
+        video_vae = AutoencoderKLWan.from_pretrained(
+            ckpt_path, subfolder="video_vae", torch_dtype=frozen_dtype,
+        ).cpu()
+        _log(f"[MOVABridge] Loaded video_vae (bf16, frozen, {_time.time()-_t0:.1f}s)")
+
+        _t0 = _time.time()
+        audio_vae = DAC.from_pretrained(
+            ckpt_path, subfolder="audio_vae", torch_dtype=frozen_dtype,
+        ).cpu()
+        _log(f"[MOVABridge] Loaded audio_vae (bf16, frozen, {_time.time()-_t0:.1f}s)")
+
+        _t0 = _time.time()
+        text_encoder = UMT5EncoderModel.from_pretrained(
+            ckpt_path, subfolder="text_encoder", torch_dtype=frozen_dtype,
+        ).to("cpu")
+        _log(f"[MOVABridge] Loaded text_encoder (bf16, frozen, {_time.time()-_t0:.1f}s)")
+        _log(f"[MOVABridge] Total model loading time: {_time.time()-_load_start:.1f}s")
+
+        tokenizer = T5TokenizerFast.from_pretrained(
+            ckpt_path, subfolder="tokenizer",
+        )
+        _log("[MOVABridge] Loaded tokenizer")
+
+        scheduler_path = os.path.join(ckpt_path, "scheduler")
+        if os.path.isdir(scheduler_path):
+            scheduler = FlowMatchPairScheduler.from_pretrained(
+                ckpt_path, subfolder="scheduler",
+            )
+            _log("[MOVABridge] Loaded scheduler from pretrained")
+        else:
+            scheduler = FlowMatchPairScheduler(shift=5.0)
+            _log("[MOVABridge] Using default FlowMatchPairScheduler(shift=5.0)")
+
+        model_index_path = os.path.join(ckpt_path, "model_index.json")
+        boundary_ratio = 0.9
+        audio_vae_type = "dac"
+        if os.path.exists(model_index_path):
+            with open(model_index_path, "r") as f:
+                model_index = json.load(f)
+            boundary_ratio = model_index.get("boundary_ratio", 0.9)
+            audio_vae_type = model_index.get("audio_vae_type", "dac")
+
+        bridge = cls(
+            video_dit=video_dit,
+            video_dit_2=video_dit_2,
+            audio_dit=audio_dit,
+            dual_tower_bridge=dual_tower_bridge,
+            boundary_ratio=boundary_ratio,
+        )
+
+        extra_components = {
+            "video_vae": video_vae,
+            "audio_vae": audio_vae,
+            "text_encoder": text_encoder,
+            "tokenizer": tokenizer,
+            "scheduler": scheduler,
+            "boundary_ratio": boundary_ratio,
+            "audio_vae_type": audio_vae_type,
+        }
+
+        return bridge, extra_components
+
+    # -----------------------------------------------------------------------
+    # FSDP helpers
+    # -----------------------------------------------------------------------
+
+    @staticmethod
+    def get_fsdp_no_split_modules():
+        """Return module types that should not be split by FSDP."""
+        return (FusedMOVABlock, DiTBlock)
+
+    def get_fsdp_block_list(self):
+        """Return a flat list of all blocks for FSDP wrapping."""
+        blocks = list(self.fusion_blocks) + list(self.remaining_video_blocks)
+        if self.video_dit_2 is not None:
+            blocks.extend(self.video_dit_2.blocks)
+        return blocks
+
+    def get_fsdp_nested_block_list(self):
+        """Sub-blocks that need an FSDP unit of their own inside a parent block.
+
+        Each FusedMOVABlock owns the high-noise expert's video block, but on a
+        low-noise step that block is bypassed in favour of ``video_dit_2``. If it
+        stays in the fused block's parameter group it is still all-gathered on
+        every one of those steps -- twice, since activation checkpointing
+        re-runs the forward during backward -- and then discarded unused. Its own
+        group is simply never entered, the same way ``video_dit_2.blocks`` are
+        never entered on a high-noise step.
+
+        Only worth doing when a second expert exists; otherwise the video block
+        runs every step and splitting it just buys an extra collective.
+        """
+        if self.video_dit_2 is None:
+            return []
+        return [fused.video_block for fused in self.fusion_blocks]
+
+    def get_fsdp_forward_tail_blocks(self):
+        """Blocks that can be the LAST one executed in a forward pass.
+
+        FSDP re-gathers the final block almost immediately for backward, so
+        resharding it after forward is a wasted round trip. Which block ends the
+        forward depends on the expert this step routed to, so both candidates are
+        returned rather than just the last entry of ``get_fsdp_block_list()`` --
+        that entry belongs to ``video_dit_2`` and is idle on high-noise steps.
+        """
+        tails = []
+        high_noise_tail = self.remaining_video_blocks or self.fusion_blocks
+        if len(high_noise_tail) > 0:
+            tails.append(high_noise_tail[-1])
+        if self.video_dit_2 is not None and len(self.video_dit_2.blocks) > 0:
+            tails.append(self.video_dit_2.blocks[-1])
+        return tails
+
+    # -----------------------------------------------------------------------
+    # BSA configuration
+    # -----------------------------------------------------------------------
+
+    def set_sparse_high_noise_only(self, flag: bool):
+        """Restrict ALL video self-attn sparse features to the high-noise expert.
+
+        When ``flag`` is True, every subsequent ``configure_*`` call (BSA top-k/top-p,
+        audio guidance, variance guidance, taylor/rectified bias correction, IVPQ /
+        penalty dynamic block, layer-adaptive) skips the low-noise expert
+        (``video_dit_2``), leaving it on the backbone's dense full attention. The
+        high-noise expert (``fusion_blocks`` + ``remaining_video_blocks``) is
+        unaffected. Must be called BEFORE the ``configure_*`` methods. Default False
+        = sparse applied to BOTH experts (legacy behavior).
+        """
+        self.sparse_high_noise_only = bool(flag)
+
+    def configure_bsa(self, enable_bsa: bool, bsa_params: dict = None,
+                       enable_bsa_v2a: bool = False, bsa_params_v2a: dict = None):
+        """Runtime configuration of Block Sparse Attention.
+
+        Args:
+            enable_bsa: enable BSA on video self-attention
+            bsa_params: params for video self-attn BSA
+            enable_bsa_v2a: enable BSA on v2a bridge cross-attention
+                (Q=audio 1D blocked internally, K=video 3D blocked)
+                a2v does NOT use BSA — audio K is too short.
+            bsa_params_v2a: params for v2a BSA
+        """
+        self._bsa_enabled = enable_bsa
+        self._bsa_params = bsa_params or {}
+        for fused_block in self.fusion_blocks:
+            fused_block.video_block.self_attn.enable_bsa = enable_bsa
+            fused_block.video_block.self_attn.bsa_params = bsa_params or {}
+            if fused_block.v2a_conditioner is not None:
+                fused_block.v2a_conditioner.inner.enable_bsa = enable_bsa_v2a
+                fused_block.v2a_conditioner.inner.bsa_params = bsa_params_v2a or {}
+        for block in self.remaining_video_blocks:
+            block.self_attn.enable_bsa = enable_bsa
+            block.self_attn.bsa_params = bsa_params or {}
+        if self.video_dit_2 is not None and not self.sparse_high_noise_only:
+            for block in self.video_dit_2.blocks:
+                block.self_attn.enable_bsa = enable_bsa
+                block.self_attn.bsa_params = bsa_params or {}
+
+    # -----------------------------------------------------------------------
+    # Audio Guidance configuration
+    # -----------------------------------------------------------------------
+
+    def configure_audio_guidance(self, enable: bool, enable_concentration_gate: bool = False,
+                                enable_timestep_reliability_gate: bool = False,
+                                enable_audio_weighted_pooling: bool = False):
+        """Runtime configuration of audio-guided BSA on video self-attention.
+
+        Two independent paths:
+          - Path A (enable=True): Multiplicative score fusion with optional gates.
+          - Path B (enable_audio_weighted_pooling=True): Audio-weighted K pooling.
+
+        Either or both can be enabled independently.
+
+        Args:
+            enable: whether to enable Path A (score modulation).
+            enable_concentration_gate: Gate 2 for Path A.
+            enable_timestep_reliability_gate: Gate 1 for Path A.
+            enable_audio_weighted_pooling: whether to enable Path B (weighted K pooling).
+        """
+        for fused_block in self.fusion_blocks:
+            sa = fused_block.video_block.self_attn
+            sa.enable_audio_guidance = enable
+            sa.enable_audio_concentration_gate = enable_concentration_gate
+            sa.enable_timestep_reliability_gate = enable_timestep_reliability_gate
+            sa.enable_audio_weighted_pooling = enable_audio_weighted_pooling
+        for block in self.remaining_video_blocks:
+            sa = block.self_attn
+            sa.enable_audio_guidance = enable
+            sa.enable_audio_concentration_gate = enable_concentration_gate
+            sa.enable_timestep_reliability_gate = enable_timestep_reliability_gate
+            sa.enable_audio_weighted_pooling = enable_audio_weighted_pooling
+        if self.video_dit_2 is not None and not self.sparse_high_noise_only:
+            for block in self.video_dit_2.blocks:
+                sa = block.self_attn
+                sa.enable_audio_guidance = enable
+                sa.enable_audio_concentration_gate = enable_concentration_gate
+                sa.enable_timestep_reliability_gate = enable_timestep_reliability_gate
+                sa.enable_audio_weighted_pooling = enable_audio_weighted_pooling
+
+    # -----------------------------------------------------------------------
+    # Variance Guidance configuration (independent from audio guidance)
+    # -----------------------------------------------------------------------
+
+    def configure_variance_guidance(self, enable: bool):
+        """Runtime configuration of channel-variance guided BSA on video self-attention.
+
+        Non-learnable, no gate. Computes ρ_j on V blocks (key/value side only):
+          score_{i,j} *= (1 + γ * ρ_j)
+        Query-awareness via P^QK zero-gate: same ρ_j, different QK scores per query
+        → different top-k per query block. γ read from bsa_params['variance_boost_gamma'].
+        Completely independent from audio guidance — no shared code.
+
+        Args:
+            enable: whether to enable channel-variance guidance.
+        """
+        for fused_block in self.fusion_blocks:
+            fused_block.video_block.self_attn.enable_variance_guidance = enable
+        for block in self.remaining_video_blocks:
+            block.self_attn.enable_variance_guidance = enable
+        if self.video_dit_2 is not None and not self.sparse_high_noise_only:
+            for block in self.video_dit_2.blocks:
+                block.self_attn.enable_variance_guidance = enable
+
+    def configure_taylor_sparse_attn(self, enable: bool):
+        """LIVEditor: selected + non-selected (Taylor) in one unified softmax."""
+        for fused_block in self.fusion_blocks:
+            fused_block.video_block.self_attn.enable_taylor_sparse_attn = enable
+            if enable:
+                fused_block.video_block.self_attn.enable_rectified_sparse_attn = False
+        for block in self.remaining_video_blocks:
+            block.self_attn.enable_taylor_sparse_attn = enable
+            if enable:
+                block.self_attn.enable_rectified_sparse_attn = False
+        if self.video_dit_2 is not None and not self.sparse_high_noise_only:
+            for block in self.video_dit_2.blocks:
+                block.self_attn.enable_taylor_sparse_attn = enable
+                if enable:
+                    block.self_attn.enable_rectified_sparse_attn = False
+
+    def configure_rectified_sparse_attn(self, enable: bool):
+        """Rectified SpaAttn: o = R_n * o_spa + A_pool[nonsel] · V_pool."""
+        for fused_block in self.fusion_blocks:
+            fused_block.video_block.self_attn.enable_rectified_sparse_attn = enable
+            if enable:
+                fused_block.video_block.self_attn.enable_taylor_sparse_attn = False
+        for block in self.remaining_video_blocks:
+            block.self_attn.enable_rectified_sparse_attn = enable
+            if enable:
+                block.self_attn.enable_taylor_sparse_attn = False
+        if self.video_dit_2 is not None and not self.sparse_high_noise_only:
+            for block in self.video_dit_2.blocks:
+                block.self_attn.enable_rectified_sparse_attn = enable
+                if enable:
+                    block.self_attn.enable_taylor_sparse_attn = False
+
+    def _video_self_attns(self):
+        """Yield EVERY video self-attn (both experts). NOT gated by
+        sparse_high_noise_only — this enumerates all attention modules and is used
+        for setting numeric bsa_params (which are inert unless a feature is enabled).
+        Never used in the forward pass.
+        """
+        for fused_block in self.fusion_blocks:
+            yield fused_block.video_block.self_attn
+        for block in self.remaining_video_blocks:
+            yield block.self_attn
+        if self.video_dit_2 is not None:
+            for block in self.video_dit_2.blocks:
+                yield block.self_attn
+
+    def _video_self_attns_with_idx(self):
+        """Yield (global_video_layer_idx, self_attn) for every video self-attn.
+
+        Layer indexing matches the forward loop: fusion_blocks[i] -> i;
+        remaining_video_blocks[j] -> n_fused + j; video_dit_2.blocks[i] -> i
+        (the boundary model mirrors the same per-layer indices). NOT gated.
+        """
+        n_fused = len(self.fusion_blocks)
+        for i, fused_block in enumerate(self.fusion_blocks):
+            yield i, fused_block.video_block.self_attn
+        for j, block in enumerate(self.remaining_video_blocks):
+            yield n_fused + j, block.self_attn
+        if self.video_dit_2 is not None:
+            for i, block in enumerate(self.video_dit_2.blocks):
+                yield i, block.self_attn
+
+    def _sparse_cfg_video_self_attns(self):
+
+        for fused_block in self.fusion_blocks:
+            yield fused_block.video_block.self_attn
+        for block in self.remaining_video_blocks:
+            yield block.self_attn
+        if self.video_dit_2 is not None and not self.sparse_high_noise_only:
+            for block in self.video_dit_2.blocks:
+                yield block.self_attn
+
+    def _sparse_cfg_video_self_attns_with_idx(self):
+        """Indexed variant of ``_sparse_cfg_video_self_attns`` (low-noise expert
+        omitted when ``sparse_high_noise_only`` is set)."""
+        n_fused = len(self.fusion_blocks)
+        for i, fused_block in enumerate(self.fusion_blocks):
+            yield i, fused_block.video_block.self_attn
+        for j, block in enumerate(self.remaining_video_blocks):
+            yield n_fused + j, block.self_attn
+        if self.video_dit_2 is not None and not self.sparse_high_noise_only:
+            for i, block in enumerate(self.video_dit_2.blocks):
+                yield i, block.self_attn
+
+    @property
+    def _num_video_layers(self):
+        return len(self.fusion_blocks) + len(self.remaining_video_blocks)
+
+    @staticmethod
+    def _disable_non_dynamic_bsa_features(sa):
+        sa.enable_audio_guidance = False
+        sa.enable_audio_concentration_gate = False
+        sa.enable_timestep_reliability_gate = False
+        sa.enable_audio_weighted_pooling = False
+        sa.enable_variance_guidance = False
+        sa.enable_taylor_sparse_attn = False
+        sa.enable_rectified_sparse_attn = False
+
+    def configure_ivpq_dynamic_block(self, enable: bool):
+        """IVPQ anisotropic dynamic block shape (Section 8.5)."""
+        for sa in self._sparse_cfg_video_self_attns():
+            sa.enable_ivpq_dynamic_block = enable
+            if enable:
+                sa.enable_penalty_dynamic_block = False
+                self._disable_non_dynamic_bsa_features(sa)
+
+    def configure_penalty_dynamic_block(self, enable: bool):
+        """Penalty-Matching anisotropic dynamic block shape (Section 8.6)."""
+        for sa in self._sparse_cfg_video_self_attns():
+            sa.enable_penalty_dynamic_block = enable
+            if enable:
+                sa.enable_ivpq_dynamic_block = False
+                self._disable_non_dynamic_bsa_features(sa)
+
+    def configure_layer_adaptive_dynamic_block(self, enable: bool):
+        """Layer-Adaptive Dynamic Block Shape wrapper (Section 8.2).
+
+        Must be configured AFTER ``configure_ivpq/penalty_dynamic_block`` (it only
+        takes effect when one of those is active). When ``enable``:
+
+          * shallow half of the video self-attn layers (idx < split) run FIXED
+            uniform-shape BSA;
+          * deep half (idx >= split, split = num_video_layers // 2) run the
+            dynamic-block method;
+          * a deep layer that HAS its own a2v bridge uses its OWN audio; only the
+            deep layers WITHOUT a bridge (the tail beyond audio_layers) reuse the
+            cache. The cache is seeded by the DEEPEST bridge layer — its
+            full-sequence ``||Δh_v||_2`` is cached (per forward) and reused by those
+            no-bridge deep layers so the audio-directional gradient still guides
+            their 64/128 + triple split, combined with each deep layer's own video
+            channel-variance.
+
+        Returns (split, cache_source_idx) for logging. Idempotent / safe to call
+        with ``enable=False`` to turn the wrapper off (back to all-layer dynamic).
+        """
+        n_fused = len(self.fusion_blocks)
+        num_video_layers = self._num_video_layers
+        split = num_video_layers // 2
+
+
+        bridge_layers = [
+            i for i, fb in enumerate(self.fusion_blocks)
+            if getattr(fb, "a2v_conditioner", None) is not None
+        ]
+        cache_source_idx = max(bridge_layers) if bridge_layers else None
+
+        # One mutable holder shared by reference across ALL video self-attns; reset
+        # at the start of every forward. Plain python attr → not an FSDP param/buffer.
+        if getattr(self, "_la_audio_cache", None) is None:
+            self._la_audio_cache = {"norms": None}
+        cache = self._la_audio_cache
+
+        for idx, sa in self._sparse_cfg_video_self_attns_with_idx():
+            sa.enable_layer_adaptive_dynamic_block = enable
+            sa._la_layer_idx = idx
+            sa._la_split = split
+            sa._la_is_cache_source = bool(
+                enable and cache_source_idx is not None and idx == cache_source_idx
+            )
+            sa._la_audio_cache = cache if enable else None
+
+        return split, cache_source_idx
+
+    # -----------------------------------------------------------------------
+    # Forward
+    # -----------------------------------------------------------------------
+
+    def forward(
+        self,
+        visual_latents: torch.Tensor,
+        audio_latents: torch.Tensor,
+        context: torch.Tensor,
+        timestep: torch.Tensor,
+        audio_context: Optional[torch.Tensor] = None,
+        audio_timestep: Optional[torch.Tensor] = None,
+        video_fps: float = 24.0,
+        num_train_timesteps: int = 1000,
+        condition_scale: float = 1.0,
+        a2v_condition_scale: Optional[float] = None,
+        v2a_condition_scale: Optional[float] = None,
+        use_video_dit_2: bool = False,
+    ):
+        """
+        Single denoising step.  Reproduces MOVA's inference_single_step exactly.
+        """
+        advance_dynamic_block_pass_id()  # once per real forward (not on recompute)
+
+        use_dit_2 = use_video_dit_2 and self.video_dit_2 is not None
+        visual_dit = self.video_dit_2 if use_dit_2 else self.video_dit
+
+        if audio_context is None:
+            audio_context = context
+        if audio_timestep is None:
+            audio_timestep = timestep
+
+        # Timestep embeddings (force fp32 for numerical stability, same as MOVA)
+        with torch.autocast("cuda", dtype=torch.float32):
+            visual_t = visual_dit.time_embedding(
+                sinusoidal_embedding_1d(visual_dit.freq_dim, timestep)
+            )
+            visual_t_mod = visual_dit.time_projection(visual_t).unflatten(1, (6, visual_dit.dim))
+
+            audio_t = self.audio_dit.time_embedding(
+                sinusoidal_embedding_1d(self.audio_dit.freq_dim, audio_timestep)
+            )
+            audio_t_mod = self.audio_dit.time_projection(audio_t).unflatten(1, (6, self.audio_dit.dim))
+
+        model_dtype = visual_dit.dtype
+        visual_t = visual_t.to(model_dtype)
+        visual_t_mod = visual_t_mod.to(model_dtype)
+        audio_t = audio_t.to(model_dtype)
+        audio_t_mod = audio_t_mod.to(model_dtype)
+
+        visual_context_emb = visual_dit.text_embedding(context)
+        audio_context_emb = self.audio_dit.text_embedding(audio_context)
+
+        visual_x = visual_latents.to(model_dtype)
+        audio_x = audio_latents.to(model_dtype)
+
+        # Patchify
+        visual_x, (t, h, w) = visual_dit.patchify(visual_x)
+        grid_size = (t, h, w)
+        visual_freqs = assemble_visual_freqs(visual_dit.freqs, t, h, w, visual_x.device)
+
+        audio_x, (f,) = self.audio_dit.patchify(audio_x, None)
+        audio_freqs = assemble_audio_freqs(self.audio_dit.freqs, f, audio_x.device)
+
+        # Compute normalized timestep ratio for Timestep Reliability Gate
+        # In flow matching: timestep ∈ [0, 1] where 0=clean, 1=noise
+        # If timestep > 1 (e.g., in [0, 1000] convention), normalize
+        t_val = timestep.float().mean()
+        timestep_ratio = t_val if t_val <= 1.0 else t_val / float(num_train_timesteps)
+
+        # Forward through fused dual-tower blocks
+        visual_x, audio_x = self._forward_dual_tower(
+            visual_x=visual_x,
+            audio_x=audio_x,
+            visual_context=visual_context_emb,
+            audio_context=audio_context_emb,
+            visual_t_mod=visual_t_mod,
+            audio_t_mod=audio_t_mod,
+            visual_freqs=visual_freqs,
+            audio_freqs=audio_freqs,
+            grid_size=grid_size,
+            video_fps=video_fps,
+            condition_scale=condition_scale,
+            a2v_condition_scale=a2v_condition_scale,
+            v2a_condition_scale=v2a_condition_scale,
+            use_dit_2=use_dit_2,
+            timestep_ratio=timestep_ratio,
+        )
+
+        # Head projections (operate on SP-split sequence for efficiency)
+        visual_output = visual_dit.head(visual_x, visual_t)
+        audio_output = self.audio_dit.head(audio_x, audio_t)
+
+        # SP: gather back full sequences after head
+        if get_sequence_parallel_state():
+            visual_output = all_gather(visual_output, dim=1, group=nccl_info.sp_group,
+                                       full_seq_len=t * h * w)
+            audio_output = all_gather(audio_output, dim=1, group=nccl_info.sp_group,
+                                      full_seq_len=f)
+
+        visual_output = visual_dit.unpatchify(visual_output, grid_size)
+        audio_output = self.audio_dit.unpatchify(audio_output, (f,))
+
+        return visual_output, audio_output
+
+    def _forward_dual_tower(
+        self,
+        visual_x,
+        audio_x,
+        visual_context,
+        audio_context,
+        visual_t_mod,
+        audio_t_mod,
+        visual_freqs,
+        audio_freqs,
+        grid_size,
+        video_fps,
+        condition_scale=1.0,
+        a2v_condition_scale=None,
+        v2a_condition_scale=None,
+        use_dit_2=False,
+        timestep_ratio=None,
+    ):
+        """
+        Forward through fused dual-tower blocks + remaining visual-only blocks.
+
+        Ulysses SP flow:
+          - Before blocks: chunk sequences and freqs
+          - Inside each attention: all_to_all_4D scatter heads / gather seq
+          - After all blocks: all_gather to reconstruct full sequences
+        """
+        sp_enabled = get_sequence_parallel_state()
+
+        # Build bridge cross-modal RoPE
+        if self.dual_tower_bridge.apply_cross_rope:
+            (visual_rope_cos_sin, audio_rope_cos_sin) = (
+                self.dual_tower_bridge.build_aligned_freqs(
+                    video_fps=video_fps,
+                    grid_size=grid_size,
+                    audio_steps=audio_x.shape[1],
+                    device=visual_x.device,
+                    dtype=visual_x.dtype,
+                )
+            )
+        else:
+            visual_rope_cos_sin = None
+            audio_rope_cos_sin = None
+
+        # Resolve condition scales
+        bridge_scale = self.dual_tower_bridge.condition_scale
+        a2v_scale = (a2v_condition_scale if a2v_condition_scale is not None
+                     else (condition_scale if condition_scale is not None
+                           else bridge_scale))
+        v2a_scale = (v2a_condition_scale if v2a_condition_scale is not None
+                     else (condition_scale if condition_scale is not None
+                           else bridge_scale))
+
+        # --- SP: chunk sequences and freqs ---
+        if sp_enabled:
+            if getattr(nccl_info, 'use_dynamic_ring_attention', False):
+                raise NotImplementedError(
+                    "MOVABridge does not support dynamic_ring_attention. "
+                    "Please set use_dynamic_ring_attention=False."
+                )
+            sp_size = nccl_info.sp_size
+            sp_rank = nccl_info.rank_within_spgroup
+
+            # Publish the pre-shard lengths of both streams. Every Ulysses
+            # all-to-all below reads them from here instead of renegotiating the
+            # split with a per-call all_gather_object.
+            visual_seq_len = visual_x.shape[1]
+            audio_seq_len = audio_x.shape[1]
+            assert_sp_splittable(visual_seq_len, sp_size, "visual sequence")
+            assert_sp_splittable(audio_seq_len, sp_size, "audio sequence")
+            sp_seq_ctx.set(visual=visual_seq_len, audio=audio_seq_len)
+
+            visual_x = sp_split(visual_x, sp_size, sp_rank, dim=1)
+            audio_x = sp_split(audio_x, sp_size, sp_rank, dim=1)
+            visual_freqs = sp_split(visual_freqs, sp_size, sp_rank, dim=0)
+            audio_freqs = sp_split(audio_freqs, sp_size, sp_rank, dim=0)
+
+            if visual_rope_cos_sin is not None:
+                v_cos, v_sin = visual_rope_cos_sin
+                v_cos = sp_split(v_cos, sp_size, sp_rank, dim=1)
+                v_sin = sp_split(v_sin, sp_size, sp_rank, dim=1)
+                visual_rope_cos_sin = (v_cos, v_sin)
+            if audio_rope_cos_sin is not None:
+                a_cos, a_sin = audio_rope_cos_sin
+                a_cos = sp_split(a_cos, sp_size, sp_rank, dim=1)
+                a_sin = sp_split(a_sin, sp_size, sp_rank, dim=1)
+                audio_rope_cos_sin = (a_cos, a_sin)
+
+        use_gc = self.training and self.gradient_checkpoint
+
+        # Reset the layer-adaptive audio cache for this forward. The deepest-bridge
+        # cache-source layer seeds it; the no-bridge tail layers (which run AFTER it,
+        # sequentially) read it. During checkpoint recompute it still holds the
+        # forward value because the tail layers recompute first (reverse order) and
+        # the cache is detached/no-grad, so the seed→read order is preserved.
+        if getattr(self, "_la_audio_cache", None) is not None:
+            self._la_audio_cache["norms"] = None
+
+        # --- Fused dual-tower layers ---
+        for i, fused_block in enumerate(self.fusion_blocks):
+            override = None
+            if use_dit_2:
+                override = self.video_dit_2.blocks[i]
+
+            if use_gc:
+                visual_x, audio_x = torch_checkpoint.checkpoint(
+                    fused_block,
+                    visual_x, audio_x,
+                    visual_context, audio_context,
+                    visual_t_mod, audio_t_mod,
+                    visual_freqs, audio_freqs,
+                    visual_rope_cos_sin, audio_rope_cos_sin,
+                    a2v_scale, v2a_scale,
+                    grid_size,
+                    override,
+                    timestep_ratio,
+                    use_reentrant=False,
+                )
+            else:
+                visual_x, audio_x = fused_block(
+                    visual_x, audio_x,
+                    visual_context, audio_context,
+                    visual_t_mod, audio_t_mod,
+                    visual_freqs, audio_freqs,
+                    visual_rope_cos_sin=visual_rope_cos_sin,
+                    audio_rope_cos_sin=audio_rope_cos_sin,
+                    a2v_scale=a2v_scale, v2a_scale=v2a_scale,
+                    video_grid_size=grid_size,
+                    override_video_block=override,
+                    timestep_ratio=timestep_ratio,
+                )
+
+        # --- Remaining video-only layers ---
+        if use_dit_2:
+            remaining = [
+                self.video_dit_2.blocks[i]
+                for i in range(len(self.fusion_blocks), len(self.video_dit_2.blocks))
+            ]
+        else:
+            remaining = self.remaining_video_blocks
+
+        for block in remaining:
+            if use_gc:
+                visual_x = torch_checkpoint.checkpoint(
+                    block, visual_x, visual_context, visual_t_mod, visual_freqs, grid_size,
+                    None, timestep_ratio,
+                    use_reentrant=False,
+                )
+            else:
+                visual_x = block(
+                    visual_x, visual_context, visual_t_mod, visual_freqs,
+                    grid_size=grid_size, timestep_ratio=timestep_ratio,
+                )
+
+        return visual_x, audio_x

hymm/models/modules/wan_audio_dit.py

@@ -0,0 +1,271 @@
+
+import math
+from typing import Literal, Optional, Tuple
+
+import torch
+import torch.nn as nn
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from einops import rearrange
+
+from .wan_video_dit import DiTBlock
+from hymm.utils.parallel_states import get_sequence_parallel_state, nccl_info, sp_seq_ctx
+from hymm.utils.communications import all_gather, sp_split, assert_sp_splittable
+
+
+def tag_audio_sp_stream(blocks):
+    """Mark a stack of DiTBlocks as operating on the audio SP stream.
+
+    The blocks are plain video DiTBlocks, so their attentions default to the
+    "visual" stream; without this the Ulysses all-to-all would look up the wrong
+    sequence length once MOVABridge fuses the two towers.
+    """
+    for block in blocks:
+        block.self_attn.sp_stream = "audio"
+        block.cross_attn.sp_stream = "audio"
+
+
+def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor):
+    # x is fp32 after layer norm
+    # print(f"{shift.dtype = }")
+    return (x * (1 + scale) + shift).to(shift.dtype)
+
+
+def sinusoidal_embedding_1d(dim, position):
+    sinusoid = torch.outer(position.type(torch.float64), torch.pow(
+        10000, -torch.arange(dim//2, dtype=torch.float64, device=position.device).div(dim//2)))
+    x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
+    return x.to(position.dtype)
+
+
+def precompute_freqs_cis_3d(dim: int, end: int = 1024, theta: float = 10000.0):
+    # 3d rope precompute
+    f_freqs_cis = precompute_freqs_cis(dim - 2 * (dim // 3), end, theta)
+    h_freqs_cis = precompute_freqs_cis(dim // 3, end, theta)
+    w_freqs_cis = precompute_freqs_cis(dim // 3, end, theta)
+    return f_freqs_cis, h_freqs_cis, w_freqs_cis
+
+def legacy_precompute_freqs_cis_1d(dim: int, end: int = 16384, theta: float = 10000.0, base_tps=4.0, target_tps=44100/2048):
+    s = float(base_tps) / float(target_tps)
+    # 1d rope precompute
+    f_freqs_cis = precompute_freqs_cis(dim - 2 * (dim // 3), end, theta, s)
+    # Do not apply positional encoding to the remaining dimensions.
+    no_freqs_cis = precompute_freqs_cis(dim // 3, end, theta, s)
+    no_freqs_cis = torch.ones_like(no_freqs_cis)
+    return f_freqs_cis, no_freqs_cis, no_freqs_cis
+
+
+def precompute_freqs_cis_1d(dim: int, end: int = 16384, theta: float = 10000.0):
+    f_freqs_cis = precompute_freqs_cis(dim, end, theta)
+    return f_freqs_cis.chunk(3, dim=-1)
+
+
+def precompute_freqs_cis(dim: int, end: int = 16384, theta: float = 10000.0, s: float = 1.0):
+    # 1d rope precompute
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)
+                   [: (dim // 2)].double() / dim))
+    pos = torch.arange(end, dtype=torch.float64, device=freqs.device) * s
+    freqs = torch.outer(pos, freqs)
+    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
+    return freqs_cis
+
+
+class MLP(torch.nn.Module):
+    def __init__(self, in_dim, out_dim, has_pos_emb=False):
+        super().__init__()
+        self.proj = torch.nn.Sequential(
+            nn.LayerNorm(in_dim),
+            nn.Linear(in_dim, in_dim),
+            nn.GELU(),
+            nn.Linear(in_dim, out_dim),
+            nn.LayerNorm(out_dim)
+        )
+        self.has_pos_emb = has_pos_emb
+        if has_pos_emb:
+            self.emb_pos = torch.nn.Parameter(torch.zeros((1, 514, 1280)))
+
+    def forward(self, x):
+        if self.has_pos_emb:
+            x = x + self.emb_pos.to(dtype=x.dtype, device=x.device)
+        return self.proj(x)
+
+
+class Head(nn.Module):
+    def __init__(self, dim: int, out_dim: int, patch_size: Tuple[int, int, int], eps: float):
+        super().__init__()
+        self.dim = dim
+        self.patch_size = patch_size
+        self.norm = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
+        self.head = nn.Linear(dim, out_dim * math.prod(patch_size))
+        self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
+
+    def forward(self, x, t_mod):
+        # print(f"{t_mod.shape = }")
+        if len(t_mod.shape) == 3:
+            shift, scale = (self.modulation.unsqueeze(0).to(dtype=t_mod.dtype, device=t_mod.device) + t_mod.unsqueeze(2)).chunk(2, dim=2)
+            x = (self.head(self.norm(x) * (1 + scale.squeeze(2)) + shift.squeeze(2)))
+        else:
+            # NOTE: `t_mod` used to be [B, C]. When B=1 broadcasting works, but for B>1
+            # it does not align with [1, 2, C]. We therefore unsqueeze at dim=1 here.
+            shift, scale = (self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod.unsqueeze(1)).chunk(2, dim=1)
+            x = (self.head(self.norm(x) * (1 + scale) + shift))
+        return x
+
+
+class WanAudioModel(ModelMixin, ConfigMixin):
+    _repeated_blocks = ("DiTBlock",)
+
+    @register_to_config
+    def __init__(
+        self,
+        dim: int,
+        in_dim: int,
+        ffn_dim: int,
+        out_dim: int,
+        text_dim: int,
+        freq_dim: int,
+        eps: float,
+        patch_size: Tuple[int, int, int],
+        num_heads: int,
+        num_layers: int,
+        has_image_input: bool,
+        has_image_pos_emb: bool = False,
+        has_ref_conv: bool = False,
+        add_control_adapter: bool = False,
+        in_dim_control_adapter: int = 24,
+        seperated_timestep: bool = False,
+        require_vae_embedding: bool = True,
+        require_clip_embedding: bool = True,
+        fuse_vae_embedding_in_latents: bool = False,
+        vae_type: Literal["oobleck", "dac"] = "oobleck",
+    ):
+        super().__init__()
+        self.dim = dim
+        self.freq_dim = freq_dim
+        self.has_image_input = has_image_input
+        self.patch_size = patch_size
+        self.seperated_timestep = seperated_timestep
+        self.require_vae_embedding = require_vae_embedding
+        self.require_clip_embedding = require_clip_embedding
+        self.fuse_vae_embedding_in_latents = fuse_vae_embedding_in_latents
+        self.vae_type = vae_type
+        # self.patch_embedding = nn.Conv3d(
+        #     in_dim, dim, kernel_size=patch_size, stride=patch_size)
+        self.patch_embedding = nn.Conv1d(
+            in_dim, dim, kernel_size=patch_size, stride=patch_size
+        )
+        self.text_embedding = nn.Sequential(
+            nn.Linear(text_dim, dim),
+            nn.GELU(approximate='tanh'),
+            nn.Linear(dim, dim)
+        )
+        self.time_embedding = nn.Sequential(
+            nn.Linear(freq_dim, dim),
+            nn.SiLU(),
+            nn.Linear(dim, dim)
+        )
+        self.time_projection = nn.Sequential(
+            nn.SiLU(), nn.Linear(dim, dim * 6))
+        self.blocks = nn.ModuleList([
+            DiTBlock(has_image_input, dim, num_heads, ffn_dim, eps)
+            for _ in range(num_layers)
+        ])
+        tag_audio_sp_stream(self.blocks)
+        self.head = Head(dim, out_dim, patch_size, eps)
+        head_dim = dim // num_heads
+        if vae_type == "oobleck":
+            self.freqs = legacy_precompute_freqs_cis_1d(head_dim, base_tps=4.0, target_tps=44100/2048)
+        elif vae_type == "dac":
+            self.freqs = precompute_freqs_cis_1d(head_dim)
+        else:
+            raise ValueError(f"Invalid VAE type: {vae_type}")
+
+        if has_image_input:
+            self.img_emb = MLP(1280, dim, has_pos_emb=has_image_pos_emb)  # clip_feature_dim = 1280
+        if has_ref_conv:
+            self.ref_conv = nn.Conv2d(16, dim, kernel_size=(2, 2), stride=(2, 2))
+        self.has_image_pos_emb = has_image_pos_emb
+        self.has_ref_conv = has_ref_conv
+        self.control_adapter = None
+
+    def patchify(self, x: torch.Tensor, control_camera_latents_input: Optional[torch.Tensor] = None):
+        x = self.patch_embedding(x)
+        if self.control_adapter is not None and control_camera_latents_input is not None:
+            y_camera = self.control_adapter(control_camera_latents_input)
+            x = [u + v for u, v in zip(x, y_camera)]
+            x = x[0].unsqueeze(0)
+        grid_size = x.shape[2:]
+        x = rearrange(x, 'b c f -> b f c').contiguous()
+        return x, grid_size  # x, grid_size: (f)
+
+    def unpatchify(self, x: torch.Tensor, grid_size: torch.Tensor):
+        return rearrange(
+            x, 'b f (p c) -> b c (f p)',
+            f=grid_size[0],
+            p=self.patch_size[0]
+        )
+
+    def forward(self,
+                x: torch.Tensor,
+                timestep: torch.Tensor,
+                context: torch.Tensor,
+                clip_feature: Optional[torch.Tensor] = None,
+                y: Optional[torch.Tensor] = None,
+                use_gradient_checkpointing: bool = False,
+                use_gradient_checkpointing_offload: bool = False,
+                **kwargs,
+                ):
+        t = self.time_embedding(
+            sinusoidal_embedding_1d(self.freq_dim, timestep))
+        t_mod = self.time_projection(t).unflatten(1, (6, self.dim))
+        context = self.text_embedding(context)
+
+        if self.has_image_input:
+            x = torch.cat([x, y], dim=1)  # (b, c_x + c_y, f, h, w)
+            clip_embdding = self.img_emb(clip_feature)
+            context = torch.cat([clip_embdding, context], dim=1)
+
+        x, (f, ) = self.patchify(x)
+
+        freqs = torch.cat([
+            self.freqs[0][:f].view(f, -1).expand(f, -1),
+            self.freqs[1][:f].view(f, -1).expand(f, -1),
+            self.freqs[2][:f].view(f, -1).expand(f, -1),
+        ], dim=-1).reshape(f, 1, -1).to(x.device)
+
+        if get_sequence_parallel_state():
+            sp_size = nccl_info.sp_size
+            sp_rank = nccl_info.rank_within_spgroup
+            assert_sp_splittable(f, sp_size, "audio sequence")
+            sp_seq_ctx.set(audio=f)
+            x = sp_split(x, sp_size, sp_rank, dim=1)
+            freqs = sp_split(freqs, sp_size, sp_rank, dim=0)
+
+        def create_custom_forward(module):
+            def custom_forward(*inputs):
+                return module(*inputs)
+            return custom_forward
+
+        for block in self.blocks:
+            if self.training and use_gradient_checkpointing:
+                if use_gradient_checkpointing_offload:
+                    with torch.autograd.graph.save_on_cpu():
+                        x = torch.utils.checkpoint.checkpoint(
+                            create_custom_forward(block),
+                            x, context, t_mod, freqs,
+                            use_reentrant=False,
+                        )
+                else:
+                    x = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward(block),
+                        x, context, t_mod, freqs,
+                        use_reentrant=False,
+                    )
+            else:
+                x = block(x, context, t_mod, freqs)
+
+        x = self.head(x, t)
+        if get_sequence_parallel_state():
+            x = all_gather(x, dim=1, group=nccl_info.sp_group, full_seq_len=f)
+        x = self.unpatchify(x, (f, ))
+        return x

hymm/models/modules/wan_video_dit.py

@@ -0,0 +1,1324 @@
+
+import math
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from einops import rearrange
+from torch.distributed.tensor import DTensor
+from torch.nn import RMSNorm
+
+from hymm.utils.parallel_states import get_sequence_parallel_state, nccl_info, sp_seq_ctx
+from hymm.utils.communications import (
+    all_to_all_4D, all_gather,
+    sp_pad_num_heads, sp_pad_heads, sp_num_local_real_heads,
+    sp_split, assert_sp_splittable,
+)
+from hymm.models.modules.block_sparse_attention import flash_attn_bsa_3d
+
+_DYN_BLOCK_FWD_ID = 0
+
+
+def advance_dynamic_block_pass_id():
+    """Call once at the start of each real outer forward (never from recompute)."""
+    global _DYN_BLOCK_FWD_ID
+    _DYN_BLOCK_FWD_ID += 1
+    return _DYN_BLOCK_FWD_ID
+
+
+def current_dynamic_block_pass_id():
+    return _DYN_BLOCK_FWD_ID
+
+
+try:
+    from flash_attn_interface import flash_attn_func as flash_attn_3_func
+    print("flash_attn_interface loaded from flash_attn_interface")
+    FLASH_ATTN_3_AVAILABLE = True
+except ModuleNotFoundError:
+    FLASH_ATTN_3_AVAILABLE = False
+    try:
+        from kernels import get_kernel
+        flash_attn_kernel = get_kernel("kernels-community/flash-attn3")
+        flash_attn_3_func = flash_attn_kernel.flash_attn_func
+        print("flash_attn_interface loaded from kernels-community/flash-attn3")
+        FLASH_ATTN_3_AVAILABLE = True
+    except ModuleNotFoundError:
+        FLASH_ATTN_3_AVAILABLE = False
+
+try:
+    import flash_attn
+    FLASH_ATTN_2_AVAILABLE = True
+except ModuleNotFoundError:
+    FLASH_ATTN_2_AVAILABLE = False
+
+try:
+    from sageattention import sageattn
+    SAGE_ATTN_AVAILABLE = True
+except ModuleNotFoundError:
+    SAGE_ATTN_AVAILABLE = False
+
+try:
+    from yunchang import LongContextAttention
+    from yunchang.kernels import AttnType
+    LONG_CONTEXT_ATTN_AVAILABLE = True
+except Exception:
+    LONG_CONTEXT_ATTN_AVAILABLE = False
+    LongContextAttention = None
+    # Fallback so module import succeeds even when yunchang isn't available.
+    # `USPAttention` will error if instantiated without LongContextAttention.
+    class AttnType:  # type: ignore
+        FA = None
+        FA3 = None
+
+
+def flash_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, num_heads: int, compatibility_mode=False):
+    if compatibility_mode:
+        q = rearrange(q, "b s (n d) -> b n s d", n=num_heads)
+        k = rearrange(k, "b s (n d) -> b n s d", n=num_heads)
+        v = rearrange(v, "b s (n d) -> b n s d", n=num_heads)
+        x = F.scaled_dot_product_attention(q, k, v)
+        x = rearrange(x, "b n s d -> b s (n d)", n=num_heads)
+    elif FLASH_ATTN_3_AVAILABLE:
+        q = rearrange(q, "b s (n d) -> b s n d", n=num_heads)
+        k = rearrange(k, "b s (n d) -> b s n d", n=num_heads)
+        v = rearrange(v, "b s (n d) -> b s n d", n=num_heads)
+        x = flash_attn_3_func(q, k, v)
+        if isinstance(x,tuple):
+            x = x[0]
+        x = rearrange(x, "b s n d -> b s (n d)", n=num_heads)
+    elif FLASH_ATTN_2_AVAILABLE:
+        q = rearrange(q, "b s (n d) -> b s n d", n=num_heads)
+        k = rearrange(k, "b s (n d) -> b s n d", n=num_heads)
+        v = rearrange(v, "b s (n d) -> b s n d", n=num_heads)
+        x = flash_attn.flash_attn_func(q, k, v)
+        x = rearrange(x, "b s n d -> b s (n d)", n=num_heads)
+    elif SAGE_ATTN_AVAILABLE:
+        q = rearrange(q, "b s (n d) -> b s n d", n=num_heads)
+        k = rearrange(k, "b s (n d) -> b s n d", n=num_heads)
+        v = rearrange(v, "b s (n d) -> b s n d", n=num_heads)
+        x = sageattn(q, k, v, tensor_layout="NHD")
+        x = rearrange(x, "b s n d -> b s (n d)", n=num_heads)
+    else:
+        q = rearrange(q, "b s (n d) -> b n s d", n=num_heads)
+        k = rearrange(k, "b s (n d) -> b n s d", n=num_heads)
+        v = rearrange(v, "b s (n d) -> b n s d", n=num_heads)
+        x = F.scaled_dot_product_attention(q, k, v)
+        x = rearrange(x, "b n s d -> b s (n d)", n=num_heads)
+    return x
+
+
+@torch.compile(fullgraph=True)
+def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor):
+    return (x * (1 + scale) + shift)
+
+
+def sinusoidal_embedding_1d(dim, position):
+    sinusoid = torch.outer(position.type(torch.float64), torch.pow(
+        10000, -torch.arange(dim//2, dtype=torch.float64, device=position.device).div(dim//2)))
+    x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
+    return x.to(position.dtype)
+
+
+def precompute_freqs_cis_3d(dim: int, end: int = 1024, theta: float = 10000.0):
+    # 3d rope precompute
+    f_freqs_cis = precompute_freqs_cis(dim - 2 * (dim // 3), end, theta)
+    h_freqs_cis = precompute_freqs_cis(dim // 3, end, theta)
+    w_freqs_cis = precompute_freqs_cis(dim // 3, end, theta)
+    return f_freqs_cis, h_freqs_cis, w_freqs_cis
+
+
+def precompute_freqs_cis(dim: int, end: int = 1024, theta: float = 10000.0):
+    # 1d rope precompute
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)
+                   [: (dim // 2)].double() / dim))
+    freqs = torch.outer(torch.arange(end, device=freqs.device), freqs)
+    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
+    return freqs_cis
+
+
+def _rope_rotate(x_4d, freqs):
+    """RoPE via sin/cos rotation in the input dtype — avoids float64 allocation.
+
+    Complex multiply (a+bi)(c+di) = (ac-bd) + (ad+bc)i is equivalent to:
+      out_even = x_even * cos - x_odd * sin
+      out_odd  = x_even * sin + x_odd * cos
+    where freqs = cos + i*sin (complex64).
+
+    Args:
+        x_4d: [B, S, n, head_dim] in bf16/fp32.
+        freqs: [S, 1, head_dim//2] complex64.
+    Returns:
+        [B, S, n*head_dim] in same dtype as x_4d.
+    """
+    cos_f = freqs.real.unsqueeze(0).to(x_4d.dtype)   # [1, S, 1, head_dim//2]
+    sin_f = freqs.imag.unsqueeze(0).to(x_4d.dtype)   # [1, S, 1, head_dim//2]
+    x_pairs = x_4d.unflatten(-1, (-1, 2))             # [..., head_dim//2, 2]
+    x_even = x_pairs[..., 0]                          # view
+    x_odd = x_pairs[..., 1]                           # view
+    out_even = x_even * cos_f
+    out_even.sub_(x_odd * sin_f)
+    out_odd = x_even * sin_f
+    out_odd.add_(x_odd * cos_f)
+    return torch.stack([out_even, out_odd], dim=-1).flatten(-2).flatten(2)
+
+
+@torch.amp.autocast('cuda', enabled=False)
+def rope_apply(x, freqs, num_heads):
+    x = rearrange(x, "b s (n d) -> b s n d", n=num_heads)
+    x_out = torch.view_as_complex(x.to(torch.float64).reshape(x.shape[0], x.shape[1], x.shape[2], -1, 2))
+    x_out = torch.view_as_real(x_out * freqs).flatten(2)
+    return x_out.to(x.dtype)
+    # return _rope_rotate(x, freqs)
+
+
+@torch.amp.autocast('cuda', enabled=False)
+def rope_apply_head_dim(x, freqs, head_dim):
+    x = rearrange(x, "b s (n d) -> b s n d", d=head_dim)
+    x_out = torch.view_as_complex(x.to(torch.float64).reshape(x.shape[0], x.shape[1], x.shape[2], -1, 2))
+    x_out = torch.view_as_real(x_out * freqs).flatten(2)
+    return x_out.to(x.dtype)
+    # return _rope_rotate(x, freqs)
+
+
+class SlowRMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        dtype = x.dtype
+        return (self.norm(x.float()) * self.weight).to(dtype)
+
+
+class AttentionModule(nn.Module):
+    def __init__(self, num_heads):
+        super().__init__()
+        self.num_heads = num_heads
+
+    def forward(self, q, k, v):
+        x = flash_attention(q=q, k=k, v=v, num_heads=self.num_heads)
+        return x
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, eps: float = 1e-6,
+                 enable_bsa: bool = False, bsa_params: dict = None,
+                 enable_audio_guidance: bool = False,
+                 enable_audio_concentration_gate: bool = False,
+                 enable_timestep_reliability_gate: bool = False,
+                 enable_variance_guidance: bool = False):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+
+        self.q = nn.Linear(dim, dim)
+        self.k = nn.Linear(dim, dim)
+        self.v = nn.Linear(dim, dim)
+        self.o = nn.Linear(dim, dim)
+        self.norm_q = RMSNorm(dim, eps=eps)
+        self.norm_k = RMSNorm(dim, eps=eps)
+
+        self.attn = AttentionModule(self.num_heads)
+        self.enable_bsa = enable_bsa
+        self.bsa_params = bsa_params or {}
+
+        # Which SP-sharded stream this attention operates on. MOVABridge retags
+        # the audio tower to "audio"; see sp_seq_ctx in parallel_states.
+        self.sp_stream = "visual"
+
+        # --- Audio Guidance (Section 3.5.1 Path A: Post-QK Score Modulation) ---
+        self.enable_audio_guidance = enable_audio_guidance
+        self.enable_audio_concentration_gate = enable_audio_concentration_gate
+        self.enable_timestep_reliability_gate = enable_timestep_reliability_gate
+        # --- Audio Weighted Pooling (Section 3.5.2 Path B: Pre-QK Representation Enhancement) ---
+        self.enable_audio_weighted_pooling = False
+
+        # --- Channel-Variance Guidance (non-learnable, no gate) ---
+        self.enable_variance_guidance = enable_variance_guidance
+
+        # --- Bias correction (two independent methods, at most one enabled) ---
+        self.enable_taylor_sparse_attn = False    # LIVEditor method
+        self.enable_rectified_sparse_attn = False  # Rectified SpaAttn method
+
+        # --- Anisotropic Dynamic Block Shape (Section 8; two methods, at most
+        #     one enabled). When EITHER is on, ALL of audio-guidance,
+        #     variance-guidance and bias-correction are disabled (full isolation).
+        self.enable_ivpq_dynamic_block = False     # Section 8.5
+        self.enable_penalty_dynamic_block = False  # Section 8.6
+
+        # (fwd_id, grid_padded, shape_ids) cache; training-only, reused on recompute.
+        self._dyn_shape_cache = None
+
+        # --- Layer-Adaptive Dynamic Block Shape (Section 8.2) ---
+        # Optional wrapper on top of a dynamic-block method. When enabled:
+        #   * shallow layers (idx < _la_split) use FIXED uniform-shape BSA;
+        #   * deep   layers (idx >= _la_split) use the dynamic-block method;
+        #   * deep layers WITH their own a2v bridge use their own audio; deep
+        #     layers WITHOUT a bridge (tail beyond audio_layers) reuse the cached
+        #     audio norms (seeded by the deepest bridge layer = cache-source) to
+        #     drive the audio-directional gradient, combined with the deep layer's
+        #     own video channel-variance.
+        # All fields are plain attributes set by configure_layer_adaptive_dynamic_block.
+        self.enable_layer_adaptive_dynamic_block = False
+        self._la_layer_idx = None          # this self-attn's global video-layer index
+        self._la_split = None              # shallow/deep boundary (num_video_layers // 2)
+        self._la_is_cache_source = False   # deepest shallow bridge layer → seeds the cache
+        self._la_audio_cache = None        # shared {'norms': tensor|None} across all video self-attns
+
+    _bsa_fallback_logged = False
+    _path_b_in_correction_warned = False
+
+    def _check_bsa(self, grid_size):
+        """
+        Check if BSA can be applied. Returns (can_use, fallback_reason).
+        Divisibility is NOT required — internal padding + masking handles arbitrary T/H/W.
+        When sparsity=0, BSA selects ALL blocks and produces results identical to full attention.
+        """
+        if not self.enable_bsa:
+            return False, None
+        if grid_size is None:
+            return False, "grid_size is None"
+        T, H, W = grid_size
+        if T <= 1:
+            return False, f"T={T} <= 1 (image mode)"
+        return True, None
+
+    def _prepare_audio_guidance(self, audio_token_norms, q_bhsd, valid_mask,
+                                T, H, W, pad_t, pad_h, pad_w, need_pad,
+                                timestep_ratio, use_path_a, use_path_b):
+        """Prepare audio guidance parameters (isolated from variance guidance)."""
+        from hymm.models.modules.block_sparse_attention import (
+            compute_audio_spatial_concentration_gate,
+            compute_timestep_reliability_gate,
+        )
+        if need_pad:
+            norms_3d = audio_token_norms.view(T, H, W)
+            norms_padded = F.pad(norms_3d, (0, pad_w, 0, pad_h, 0, pad_t), value=0.0)
+            audio_norms_padded = norms_padded.reshape(-1).contiguous()
+        else:
+            audio_norms_padded = audio_token_norms.contiguous()
+
+        audio_gate = torch.ones(1, device=q_bhsd.device, dtype=q_bhsd.dtype)
+        audio_boost_gamma = torch.zeros(1, device=q_bhsd.device, dtype=q_bhsd.dtype)
+        audio_weighted_lambda = 0.0
+
+        if use_path_a:
+            if self.enable_timestep_reliability_gate and timestep_ratio is not None:
+                r_t = compute_timestep_reliability_gate(
+                    timestep_ratio,
+                    mode=self.bsa_params.get('timestep_gate_mode', 'power'),
+                    power=self.bsa_params.get('timestep_gate_power', 1.5),
+                    scale=self.bsa_params.get('timestep_gate_scale', 10.0),
+                    threshold=self.bsa_params.get('timestep_gate_threshold', 0.7),
+                )
+            else:
+                r_t = torch.ones(1, device=q_bhsd.device, dtype=q_bhsd.dtype)
+
+            if self.enable_audio_concentration_gate:
+                g_a = compute_audio_spatial_concentration_gate(
+                    audio_norms_padded, valid_mask=valid_mask)
+            else:
+                g_a = torch.ones(1, device=q_bhsd.device, dtype=q_bhsd.dtype)
+
+            audio_gate = r_t * g_a
+            audio_boost_gamma = torch.tensor(
+                self.bsa_params.get('audio_boost_gamma', 1.0),
+                device=q_bhsd.device, dtype=q_bhsd.dtype,
+            )
+
+        if use_path_b:
+            audio_weighted_lambda = self.bsa_params.get('audio_weighted_lambda', 1.0)
+
+        return audio_norms_padded, audio_gate, audio_boost_gamma, audio_weighted_lambda
+
+    def _run_bsa_dynamic(self, q_bhsd, k_bhsd, v_bhsd, grid_size, audio_token_norms=None):
+        """Anisotropic Dynamic Block Shape BSA (Section 8) — video self-attn only.
+
+        Fully isolated from `_run_bsa`'s uniform path. Pads T/H/W to a multiple of
+        the macro-zone size (8) using the SAME internal padding+mask+unpad
+        contract, builds a `valid_mask` so padded tokens never participate, then
+        defers all the heavy lifting (per-zone shape selection, anisotropic
+        rearrange, logical-block selection, Triton kernel reuse, scatter back) to
+        `flash_attn_bsa_3d_dynamic`.
+
+        `audio_token_norms` (per-token ||Δh_v||_2 of the A→V bridge residual, in
+        THW order, full sequence) drives the audio-directional gradient when
+        available; it is optional (silence/no-bridge → video-only shape decision).
+        """
+        from hymm.models.modules.block_sparse_attention.dynamic_block_attention import (
+            flash_attn_bsa_3d_dynamic,
+        )
+        from hymm.models.modules.block_sparse_attention.dynamic_block_shape import ZONE_SIZE
+
+        T, H, W = grid_size
+        B, n_heads, S, D = q_bhsd.shape
+
+        pad_t = (ZONE_SIZE - T % ZONE_SIZE) % ZONE_SIZE
+        pad_h = (ZONE_SIZE - H % ZONE_SIZE) % ZONE_SIZE
+        pad_w = (ZONE_SIZE - W % ZONE_SIZE) % ZONE_SIZE
+        need_pad = pad_t > 0 or pad_h > 0 or pad_w > 0
+
+        valid_mask = None
+        audio_norms_p = audio_token_norms
+        if need_pad:
+            T_p, H_p, W_p = T + pad_t, H + pad_h, W + pad_w
+            idx_t = torch.arange(T_p, device=q_bhsd.device)
+            idx_h = torch.arange(H_p, device=q_bhsd.device)
+            idx_w = torch.arange(W_p, device=q_bhsd.device)
+            valid_mask = (
+                (idx_t[:, None, None] < T)
+                & (idx_h[None, :, None] < H)
+                & (idx_w[None, None, :] < W)
+            ).reshape(-1).contiguous()
+
+            def _pad_3d(t):
+                t_3d = t.view(B, n_heads, T, H, W, D)
+                t_padded = F.pad(t_3d, (0, 0, 0, pad_w, 0, pad_h, 0, pad_t))
+                return t_padded.reshape(B, n_heads, -1, D).contiguous()
+
+            q_bhsd = _pad_3d(q_bhsd)
+            k_bhsd = _pad_3d(k_bhsd)
+            v_bhsd = _pad_3d(v_bhsd)
+            if audio_token_norms is not None:
+                norms_3d = audio_token_norms.reshape(T, H, W)
+                audio_norms_p = F.pad(
+                    norms_3d, (0, pad_w, 0, pad_h, 0, pad_t), value=0.0
+                ).reshape(-1).contiguous()
+            grid_padded = (T_p, H_p, W_p)
+        else:
+            grid_padded = grid_size
+
+        method = "ivpq" if self.enable_ivpq_dynamic_block else "penalty"
+
+        # Training-only: reuse this forward's shape decision on the recompute
+        # (cache hit when fwd_id+grid match). Eval bypasses the cache entirely.
+        use_shape_cache = self.training
+        cached_shape_ids = None
+        if use_shape_cache:
+            cur_fwd_id = current_dynamic_block_pass_id()
+            cached = self._dyn_shape_cache
+            if (cached is not None and cached[0] == cur_fwd_id
+                    and cached[1] == grid_padded):
+                cached_shape_ids = cached[2]
+
+        out, used_shape_ids = flash_attn_bsa_3d_dynamic(
+            q_bhsd, k_bhsd, v_bhsd, grid_padded,
+            method=method,
+            sparsity=self.bsa_params.get('sparsity', 0.9375),
+            cdf_threshold=self.bsa_params.get('cdf_threshold', None),
+            audio_token_norms=audio_norms_p,
+            valid_mask=valid_mask,
+            lambda_a=self.bsa_params.get('dynamic_block_lambda_a', 0.5),
+            tau_128=self.bsa_params.get('dynamic_block_tau_128', 0.15),
+            lambda_128=self.bsa_params.get('dynamic_block_lambda_128', 1.0),
+            shape_ids=cached_shape_ids,
+            return_shape_ids=True,
+        )
+
+        if use_shape_cache:
+            self._dyn_shape_cache = (cur_fwd_id, grid_padded, used_shape_ids)
+
+        if need_pad:
+            out_3d = out.view(B, n_heads, T_p, H_p, W_p, D)
+            out = out_3d[:, :, :T, :H, :W, :].contiguous().reshape(B, n_heads, T * H * W, D)
+        return out
+
+    def _run_bsa(self, q_bhsd, k_bhsd, v_bhsd, grid_size, audio_token_norms=None, timestep_ratio=None):
+        """Run BSA with internal temporal+spatial padding/masking/unpadding.
+
+        Execution order (when all features are enabled):
+          1. Padding (temporal + spatial) + valid_mask creation
+          2. Audio/Variance score modulation (multiplicative boost)
+          3. Block selection (top-k / top-p)
+          4. BSA Triton kernel → o_spa
+          5. **Bias rectification** (post-hoc, training-free, Section 3.6):
+             a. Reproduce block-level scoring (O(N_q*N_k), negligible cost)
+             b. Compute R_n (selected weight fraction), M_n (query sharpness)
+             c. Compute O_nonsel_pool (Taylor approximation from non-selected blocks)
+             d. flat:  O = R_n * O_spa + O_nonsel_pool
+                sharp: O = R_n * O_spa
+          6. Unpadding
+
+        Bias correction controlled by enable_taylor_sparse_attn / enable_rectified_sparse_attn.
+        Isolated from audio/variance guidance code paths.
+        """
+        # === Anisotropic Dynamic Block Shape (Section 8) — fully isolated path ===
+        # When a dynamic-block trigger is on we route to a dedicated method and
+        # never touch the uniform-shape / audio / variance / bias-correction code
+        # below. The two dynamic methods are mutually exclusive with each other
+        # and with all the other BSA features (enforced in configure_* + asserts).
+        use_dynamic = self.enable_ivpq_dynamic_block or self.enable_penalty_dynamic_block
+        if use_dynamic and self.enable_layer_adaptive_dynamic_block \
+                and self._la_layer_idx is not None and self._la_split is not None:
+            # Layer-adaptive (Section 8.2): shallow half → FIXED uniform-shape BSA
+            # (falls through to the uniform dispatch below); deep half → dynamic.
+            if self._la_layer_idx < self._la_split:
+                use_dynamic = False
+        if use_dynamic:
+            return self._run_bsa_dynamic(
+                q_bhsd, k_bhsd, v_bhsd, grid_size, audio_token_norms=audio_token_norms,
+            )
+
+        from hymm.models.modules.block_sparse_attention import flash_attn_bsa_3d_audio_guided
+
+        T, H, W = grid_size
+        cq = self.bsa_params.get('chunk_3d_shape_q', [4, 4, 4])
+        ck = self.bsa_params.get('chunk_3d_shape_k', [4, 4, 4])
+        chunk_t = max(cq[0], ck[0])
+        chunk_h = max(cq[1], ck[1])
+        chunk_w = max(cq[2], ck[2])
+
+        pad_t = (chunk_t - T % chunk_t) % chunk_t
+        pad_h = (chunk_h - H % chunk_h) % chunk_h
+        pad_w = (chunk_w - W % chunk_w) % chunk_w
+
+        B, n_heads, S, D = q_bhsd.shape
+        need_pad = pad_t > 0 or pad_h > 0 or pad_w > 0
+
+        if need_pad:
+            T_p, H_p, W_p = T + pad_t, H + pad_h, W + pad_w
+
+            idx_t = torch.arange(T_p, device=q_bhsd.device)
+            idx_h = torch.arange(H_p, device=q_bhsd.device)
+            idx_w = torch.arange(W_p, device=q_bhsd.device)
+            valid_mask = (
+                (idx_t[:, None, None] < T)
+                & (idx_h[None, :, None] < H)
+                & (idx_w[None, None, :] < W)
+            ).reshape(-1).contiguous()
+
+            def _pad_3d(t):
+                t_3d = t.view(B, n_heads, T, H, W, D)
+                t_padded = F.pad(t_3d, (0, 0, 0, pad_w, 0, pad_h, 0, pad_t))
+                return t_padded.reshape(B, n_heads, -1, D).contiguous()
+
+            q_bhsd = _pad_3d(q_bhsd)
+            k_bhsd = _pad_3d(k_bhsd)
+            v_bhsd = _pad_3d(v_bhsd)
+            grid_padded = (T_p, H_p, W_p)
+        else:
+            grid_padded = grid_size
+            valid_mask = None
+
+        # --- Separate guidance-specific keys from core BSA params ---
+        _GUIDANCE_ONLY_KEYS = ('audio_boost_gamma', 'audio_weighted_lambda',
+                               'timestep_gate_mode', 'timestep_gate_power',
+                               'timestep_gate_scale', 'timestep_gate_threshold',
+                               'variance_boost_gamma',
+                               'taylor_alpha_f',
+                               # dynamic-block knobs — never go to the uniform kernel
+                               # (used by shallow layers under layer-adaptive mode)
+                               'dynamic_block_lambda_a', 'dynamic_block_tau_128',
+                               'dynamic_block_lambda_128')
+        bsa_kwargs = {k: v for k, v in self.bsa_params.items() if k not in _GUIDANCE_ONLY_KEYS}
+
+        has_audio_norms = (audio_token_norms is not None and audio_token_norms.numel() > 0)
+        use_path_a = self.enable_audio_guidance and has_audio_norms
+        use_path_b = self.enable_audio_weighted_pooling and has_audio_norms
+        use_audio = use_path_a or use_path_b
+        use_variance = self.enable_variance_guidance
+
+        # =====================================================================
+        # Prepare audio guidance data (shared by both dispatch and bias_rect).
+        # Done once, used in whichever path is active.
+        # =====================================================================
+        audio_norms_padded = audio_gate = audio_boost_gamma = audio_weighted_lambda = None
+        if use_audio:
+            audio_norms_padded, audio_gate, audio_boost_gamma, audio_weighted_lambda = (
+                self._prepare_audio_guidance(
+                    audio_token_norms, q_bhsd, valid_mask, T, H, W,
+                    pad_t, pad_h, pad_w, need_pad, timestep_ratio,
+                    use_path_a, use_path_b,
+                )
+            )
+
+        # =====================================================================
+        # Dispatch: 3 mutually exclusive modes
+        #   1. Taylor sparse attn (LIVEditor) — selected + Taylor in one softmax
+        #   2. Rectified sparse attn (Rectified SpaAttn) — R_n * o_spa + o_ncri
+        #   3. Plain BSA (default) — original paths, unchanged
+        # Audio/variance guidance affects block SELECTION only (score_modifier).
+        # =====================================================================
+        use_correction = self.enable_taylor_sparse_attn or self.enable_rectified_sparse_attn
+
+        if use_correction:
+            # Bias rectification lives in a fully-isolated module so that the
+            # audio-guidance / variance-guidance code paths above are untouched.
+            from hymm.models.modules.block_sparse_attention.bias_rectification import (
+                bsa_taylor_sparse_attn, bsa_rectified_sparse_attn,
+            )
+            from hymm.models.modules.block_sparse_attention.bsa_interface import (
+                rearrange_THW_to_3d_block, rearrange_3d_block_to_THW,
+                rearrange_THW_to_3d_block_1d,
+            )
+
+            _ck = bsa_kwargs.get('chunk_3d_shape_k', [4, 4, 4])
+            _cq = bsa_kwargs.get('chunk_3d_shape_q', [4, 4, 4])
+            _tk, _hk, _wk = _ck[0], _ck[1], _ck[2]
+            _tq, _hq, _wq = _cq[0], _cq[1], _cq[2]
+            _Tp, _Hp, _Wp = grid_padded
+            _Ntk, _Nhk, _Nwk = _Tp // _tk, _Hp // _hk, _Wp // _wk
+            _Ntq, _Nhq, _Nwq = _Tp // _tq, _Hp // _hq, _Wp // _wq
+
+            # Rearrange once, reuse for variance/audio/correction
+            q_ro = rearrange_THW_to_3d_block(q_bhsd, _Ntq, _Nhq, _Nwq, _tq, _hq, _wq, D)
+            k_ro = rearrange_THW_to_3d_block(k_bhsd, _Ntk, _Nhk, _Nwk, _tk, _hk, _wk, D)
+            v_ro = rearrange_THW_to_3d_block(v_bhsd, _Ntk, _Nhk, _Nwk, _tk, _hk, _wk, D)
+            kv_vm_ro = rearrange_THW_to_3d_block_1d(valid_mask, _Ntk, _Nhk, _Nwk, _tk, _hk, _wk) if valid_mask is not None else None
+
+            score_mod = None
+            if use_variance:
+                from hymm.models.modules.block_sparse_attention.bsa_interface import compute_channel_variance_density
+                with torch.no_grad():
+                    _rho = compute_channel_variance_density(v_ro, _tk*_hk*_wk, _tk, _hk, _wk, kv_vm_ro)
+                score_mod = self.bsa_params.get('variance_boost_gamma', 1.0) * _rho.to(q_bhsd.dtype).view(1, 1, 1, -1)
+
+            if use_audio and audio_norms_padded is not None:
+                # Path A boost: γ_audio * (g_a * S̃_j + (1 - g_a) / N_bk)
+                # All ops stay on-GPU. Avoid `.item()` here — `audio_gate` and
+                # `audio_boost_gamma` are 1-element CUDA tensors, and `.item()`
+                # synchronises the GPU pipeline (each call drains the queue;
+                # 30 layers x 50 steps x 2 (checkpoint recompute) = 3000 syncs
+                # per training step). Use plain tensor broadcasting instead.
+                from hymm.models.modules.block_sparse_attention.bsa_interface import compute_block_audio_saliency
+                _anorms_bo = rearrange_THW_to_3d_block_1d(audio_norms_padded, _Ntk, _Nhk, _Nwk, _tk, _hk, _wk)
+                _bsal = compute_block_audio_saliency(_anorms_bo, _tk*_hk*_wk, valid_mask=kv_vm_ro)
+                _N_bk = _Ntk * _Nhk * _Nwk
+                # Cast scalar/Python audio_gate to a 1-elem tensor for uniform broadcast.
+                if not torch.is_tensor(audio_gate):
+                    audio_gate = torch.tensor(audio_gate, device=q_bhsd.device, dtype=q_bhsd.dtype)
+                if not torch.is_tensor(audio_boost_gamma):
+                    audio_boost_gamma = torch.tensor(audio_boost_gamma, device=q_bhsd.device, dtype=q_bhsd.dtype)
+                _gv = audio_gate.to(q_bhsd.dtype)
+                _bg = audio_boost_gamma.to(q_bhsd.dtype)
+                _gated = _gv * _bsal.view(1, 1, 1, -1).to(q_bhsd.dtype) + (1.0 - _gv) / _N_bk
+                audio_term = _bg * _gated
+                score_mod = (score_mod + audio_term) if score_mod is not None else audio_term
+
+                # Warn-once: Path B (audio-weighted K pooling) has no effect in the
+                # bias-rectification path. Taylor / Rectified methods use plain mean
+                # pooling for K (so autograd through pooling is correct). Path A
+                # boost (audio_boost_gamma) still applies.
+                if self.enable_audio_weighted_pooling and not SelfAttention._path_b_in_correction_warned:
+                    print(
+                        "[BSA] enable_audio_weighted_pooling (Path B) is silently ignored "
+                        "when enable_taylor_sparse_attn / enable_rectified_sparse_attn is "
+                        "active. Path A boost still applies."
+                    )
+                    SelfAttention._path_b_in_correction_warned = True
+
+            if score_mod is not None:
+                score_mod = (1.0 + score_mod).contiguous()
+
+            _sm = 1.0 / (D ** 0.5)
+            _sp = bsa_kwargs.get('sparsity', 0.9375)
+            _cdf = bsa_kwargs.get('cdf_threshold', None)
+            if self.enable_taylor_sparse_attn:
+                out = bsa_taylor_sparse_attn(
+                    q_ro, k_ro, v_ro,
+                    chunk_size_q=_tq * _hq * _wq,
+                    chunk_size_k=_tk * _hk * _wk,
+                    sparsity=_sp, cdf_threshold=_cdf, sm_scale=_sm,
+                    alpha_f=self.bsa_params.get('taylor_alpha_f', 0.5),
+                    kv_valid_mask=kv_vm_ro, score_modifier=score_mod,
+                )
+            else:
+                out = bsa_rectified_sparse_attn(
+                    q_ro, k_ro, v_ro,
+                    chunk_size_q=_tq * _hq * _wq,
+                    chunk_size_k=_tk * _hk * _wk,
+                    sparsity=_sp, cdf_threshold=_cdf, sm_scale=_sm,
+                    kv_valid_mask=kv_vm_ro, score_modifier=score_mod,
+                )
+            out = rearrange_3d_block_to_THW(out, _Ntq, _Nhq, _Nwq, _tq, _hq, _wq, D)
+
+        else:
+            # ---------- Original dispatch paths (completely unchanged) ----------
+            if use_variance and not use_audio:
+                from hymm.models.modules.block_sparse_attention import flash_attn_bsa_3d_variance_guided
+                out = flash_attn_bsa_3d_variance_guided(
+                    q_bhsd, k_bhsd, v_bhsd, grid_padded, grid_padded,
+                    variance_boost_gamma=self.bsa_params.get('variance_boost_gamma', 1.0),
+                    valid_mask=valid_mask, **bsa_kwargs,
+                )
+            elif use_audio and not use_variance:
+                out = flash_attn_bsa_3d_audio_guided(
+                    q_bhsd, k_bhsd, v_bhsd, grid_padded, grid_padded,
+                    audio_token_norms=audio_norms_padded,
+                    audio_gate_value=audio_gate,
+                    audio_boost_gamma=audio_boost_gamma,
+                    valid_mask=valid_mask,
+                    audio_weighted_lambda=audio_weighted_lambda,
+                    **bsa_kwargs,
+                )
+            elif use_audio and use_variance:
+                # Pass variance_boost_gamma so ρ is computed INSIDE on the already-
+                # rearranged V — saves one large `rearrange_THW_to_3d_block(v)` per
+                # layer per step (significant on 2K where V is ~600K tokens).
+                # Mathematically identical to the explicit-score_modifier route.
+                out = flash_attn_bsa_3d_audio_guided(
+                    q_bhsd, k_bhsd, v_bhsd, grid_padded, grid_padded,
+                    audio_token_norms=audio_norms_padded,
+                    audio_gate_value=audio_gate,
+                    audio_boost_gamma=audio_boost_gamma,
+                    valid_mask=valid_mask,
+                    audio_weighted_lambda=audio_weighted_lambda,
+                    variance_boost_gamma=self.bsa_params.get('variance_boost_gamma', 1.0),
+                    **bsa_kwargs,
+                )
+            else:
+                out = flash_attn_bsa_3d(
+                    q_bhsd, k_bhsd, v_bhsd, grid_padded, grid_padded,
+                    valid_mask=valid_mask, **bsa_kwargs,
+                )
+
+        if need_pad:
+            out_3d = out.view(B, n_heads, T_p, H_p, W_p, D)
+            out = out_3d[:, :, :T, :H, :W, :].contiguous().reshape(B, n_heads, T * H * W, D)
+
+        return out
+
+    def forward(self, x, freqs, grid_size=None, a2v_bridge_residual=None, timestep_ratio=None):
+        """
+        Args:
+            a2v_bridge_residual: [B, S, D] — A→V bridge cross-attention output (before scale).
+                Used for audio-guided BSA scoring. Only effective when
+                enable_audio_guidance=True and BSA is active.
+            timestep_ratio: scalar in [0,1] — normalized timestep (0=clean, 1=noise).
+                Used for Timestep Reliability Gate. Only effective when
+                enable_timestep_reliability_gate=True.
+        """
+        q = self.norm_q(self.q(x))
+        k = self.norm_k(self.k(x))
+        v = self.v(x)
+        if isinstance(freqs, DTensor):
+            freqs = freqs.to_local()
+        q = rope_apply_head_dim(q, freqs, self.head_dim)
+        k = rope_apply_head_dim(k, freqs, self.head_dim)
+
+        use_bsa, fallback_reason = self._check_bsa(grid_size)
+
+        if self.enable_bsa and not use_bsa and not SelfAttention._bsa_fallback_logged:
+            print(f"[BSA] Falling back to full attention: {fallback_reason} "
+                  f"(grid_size={grid_size})")
+            SelfAttention._bsa_fallback_logged = True
+
+        # Compute per-token audio norms for audio guidance (before SP reshuffling).
+        # Also needed by the dynamic-block path for the audio-directional gradient.
+        audio_token_norms = None
+        # Layer-adaptive dynamic block (Section 8.2): only effective with a dynamic method.
+        la_on = (self.enable_layer_adaptive_dynamic_block
+                 and (self.enable_ivpq_dynamic_block or self.enable_penalty_dynamic_block)
+                 and self._la_layer_idx is not None and self._la_split is not None)
+        la_is_deep = la_on and (self._la_layer_idx >= self._la_split)
+
+        use_any_audio = use_bsa and (
+            self.enable_audio_guidance or self.enable_audio_weighted_pooling
+            or self.enable_ivpq_dynamic_block or self.enable_penalty_dynamic_block
+        )
+        # Decide whether THIS layer needs to compute audio norms (a per-layer
+        # all-gather under SP). Under layer-adaptive mode we avoid wasted work:
+        #   * shallow layers run FIXED BSA (ignore audio) → only the designated
+        #     cache-source layer computes audio (to seed the shallow→deep cache);
+        #   * deep layers compute their own only if they actually have a bridge
+        #     (otherwise audio stays None and they reuse the cache below).
+        if la_on:
+            need_audio = True if la_is_deep else self._la_is_cache_source
+        else:
+            need_audio = use_any_audio
+        sp_full_len = sp_seq_ctx.get(self.sp_stream)
+
+        if need_audio and a2v_bridge_residual is not None:
+            # a2v_bridge_residual: [B, S_local, D] (SP-split) or [B, S, D] (no SP)
+            # Compute per-token L2 norm: ||Δh_v[m]||_2
+            with torch.no_grad():
+                local_norms = a2v_bridge_residual.norm(dim=-1)  # [B, S_local]
+
+            if get_sequence_parallel_state():
+                # All-gather norms across SP ranks to get full sequence
+                local_norms_4d = local_norms.unsqueeze(-1)  # [B, S_local, 1]
+                full_norms = all_gather(local_norms_4d, dim=1, group=nccl_info.sp_group,
+                                        full_seq_len=sp_full_len)
+                audio_token_norms = full_norms.squeeze(-1)[0]  # [S]
+            else:
+                audio_token_norms = local_norms[0]  # [S]
+
+            # Seed the shallow→deep cache (detached, full-seq THW order, SP-consistent).
+            # `audio_token_norms` is computed under no_grad above (no graph retained).
+            if la_on and self._la_is_cache_source and self._la_audio_cache is not None:
+                self._la_audio_cache['norms'] = audio_token_norms.contiguous()
+
+        # Deep layers under layer-adaptive: reuse cached shallow audio when this layer
+        # has no bridge of its own (own audio_token_norms is None for those layers).
+        if la_is_deep and audio_token_norms is None and self._la_audio_cache is not None:
+            audio_token_norms = self._la_audio_cache.get('norms', None)
+
+        if get_sequence_parallel_state():
+            sp_size = nccl_info.sp_size
+            sp_rank = nccl_info.rank_within_spgroup
+            # Ulysses requires num_heads % sp_size == 0. When it is not (e.g. audio
+            # DiT = 12 heads with sp_size=8), pad the head dim with ZERO heads to the
+            # next multiple. Padded heads are sliced off before any attention compute
+            # and truncated after the inverse all_to_all, so they never enter the
+            # BSA Triton / flash kernel and never affect the real heads' fwd/bwd.
+            # When num_heads is already divisible (e.g. video = 40 heads), padded_heads
+            # == num_heads and this path is bit-identical to the original code.
+            padded_heads = sp_pad_num_heads(self.num_heads, sp_size)
+
+            q, k, v = [rearrange(t, "b s (n d) -> b s n d", n=self.num_heads) for t in (q, k, v)]
+            if padded_heads != self.num_heads:
+                q = sp_pad_heads(q, padded_heads)
+                k = sp_pad_heads(k, padded_heads)
+                v = sp_pad_heads(v, padded_heads)
+
+            q = all_to_all_4D(q, nccl_info.sp_group, scatter_dim=2, gather_dim=1,
+                              full_seq_len=sp_full_len)
+            k = all_to_all_4D(k, nccl_info.sp_group, scatter_dim=2, gather_dim=1,
+                              full_seq_len=sp_full_len)
+            v = all_to_all_4D(v, nccl_info.sp_group, scatter_dim=2, gather_dim=1,
+                              full_seq_len=sp_full_len)
+
+            hpr = q.shape[2]  # heads-per-rank after scatter (= padded_heads // sp_size)
+            local_heads = sp_num_local_real_heads(self.num_heads, padded_heads, sp_size, sp_rank)
+            if local_heads > 0:
+                if local_heads != hpr:
+                    # Mixed rank: keep only the real-head prefix. The slice stays
+                    # connected to the all_to_all output, so its backward yields grad
+                    # for all hpr local heads (zeros for the padded tail) and the
+                    # q/k/v all_to_all collectives still fire here.
+                    q = q[:, :, :local_heads, :]
+                    k = k[:, :, :local_heads, :]
+                    v = v[:, :, :local_heads, :]
+                if use_bsa:
+                    q_bsa = q.permute(0, 2, 1, 3).contiguous()
+                    k_bsa = k.permute(0, 2, 1, 3).contiguous()
+                    v_bsa = v.permute(0, 2, 1, 3).contiguous()
+                    x_bsa = self._run_bsa(q_bsa, k_bsa, v_bsa, grid_size,
+                                          audio_token_norms=audio_token_norms,
+                                          timestep_ratio=timestep_ratio)
+                    x = rearrange(x_bsa, "b n s d -> b s (n d)")
+                else:
+                    x = flash_attention(q.flatten(2), k.flatten(2), v.flatten(2), num_heads=local_heads)
+                x = rearrange(x, "b s (n d) -> b s n d", n=local_heads)
+                if x.shape[2] != hpr:
+                    # Re-pad to hpr so the inverse all_to_all (head gather) is balanced.
+                    x = sp_pad_heads(x, hpr)
+            else:
+                # Padding-only rank: this rank holds NO real heads, so NO attention
+                # kernel is run (padding never enters the BSA Triton / flash kernel).
+                # CRITICAL: x must stay CONNECTED to q/k/v (value forced to zero) so the
+                # backward all_to_all for q/k/v still fires on this rank. This keeps the
+                # NCCL collectives symmetric across ranks (no deadlock) AND lets this
+                # rank's real tokens receive their correct q/k/v gradients, which are
+                # computed on the head-owner ranks and routed back through the all_to_all.
+                # `* 0.0` guarantees zero contribution to forward/backward of real heads.
+                x = (q + k + v) * 0.0
+            x = all_to_all_4D(x, nccl_info.sp_group, scatter_dim=1, gather_dim=2,
+                              full_seq_len=sp_full_len)
+            if padded_heads != self.num_heads:
+                x = x[:, :, :self.num_heads, :]
+            x = x.flatten(2)
+        else:
+            if use_bsa:
+                q_bsa = rearrange(q, "b s (n d) -> b n s d", n=self.num_heads).contiguous()
+                k_bsa = rearrange(k, "b s (n d) -> b n s d", n=self.num_heads).contiguous()
+                v_bsa = rearrange(v, "b s (n d) -> b n s d", n=self.num_heads).contiguous()
+                x_bsa = self._run_bsa(q_bsa, k_bsa, v_bsa, grid_size,
+                                      audio_token_norms=audio_token_norms,
+                                      timestep_ratio=timestep_ratio)
+                x = rearrange(x_bsa, "b n s d -> b s (n d)")
+            else:
+                x = self.attn(q, k, v)
+
+        return self.o(x)
+
+
+class USPAttention(nn.Module):
+    def __init__(self, num_heads: int, attn_type=AttnType.FA):
+        super().__init__()
+        if not LONG_CONTEXT_ATTN_AVAILABLE or LongContextAttention is None:
+            raise RuntimeError(
+                "USPAttention requires `yunchang` (LongContextAttention). "
+                "Please install/enable yunchang or avoid using USPAttention."
+            )
+        self.num_heads = num_heads
+        self.attn = LongContextAttention(ring_impl_type="basic", attn_type=attn_type)
+
+    def forward(self, q, k, v):
+        q = rearrange(q, "b s (n d) -> b s n d", n=self.num_heads)
+        k = rearrange(k, "b s (n d) -> b s n d", n=self.num_heads)
+        v = rearrange(v, "b s (n d) -> b s n d", n=self.num_heads)
+        x = self.attn(q, k, v)
+        return rearrange(x, "b s n d -> b s (n d)")
+
+
+class CrossAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, eps: float = 1e-6, has_image_input: bool = False):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+
+        self.q = nn.Linear(dim, dim)
+        self.k = nn.Linear(dim, dim)
+        self.v = nn.Linear(dim, dim)
+        self.o = nn.Linear(dim, dim)
+        self.norm_q = RMSNorm(dim, eps=eps)
+        self.norm_k = RMSNorm(dim, eps=eps)
+        self.has_image_input = has_image_input
+        if has_image_input:
+            self.k_img = nn.Linear(dim, dim)
+            self.v_img = nn.Linear(dim, dim)
+            self.norm_k_img = RMSNorm(dim, eps=eps)
+
+        self.attn = AttentionModule(self.num_heads)
+
+        # Query stream for SP; the text context is replicated, not sharded.
+        self.sp_stream = "visual"
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor):
+        if self.has_image_input:
+            img = y[:, :257]
+            ctx = y[:, 257:]
+        else:
+            ctx = y
+        q = self.norm_q(self.q(x))
+        k = self.norm_k(self.k(ctx))
+        v = self.v(ctx)
+
+        sp_enabled = get_sequence_parallel_state()
+
+        if sp_enabled:
+            sp_size = nccl_info.sp_size
+            sp_rank = nccl_info.rank_within_spgroup
+            # Same Ulysses head-padding contract as SelfAttention: pad to a multiple
+            # of sp_size, slice real heads before compute, truncate after. The k/v
+            # context is NOT all_to_all'd (replicated); it is padded then sliced by
+            # head_start so each rank takes the same global heads as its q.
+            padded_heads = sp_pad_num_heads(self.num_heads, sp_size)
+
+            q_4d = rearrange(q, "b s (n d) -> b s n d", n=self.num_heads)
+            k_4d = rearrange(k, "b s (n d) -> b s n d", n=self.num_heads)
+            v_4d = rearrange(v, "b s (n d) -> b s n d", n=self.num_heads)
+            if padded_heads != self.num_heads:
+                q_4d = sp_pad_heads(q_4d, padded_heads)
+                k_4d = sp_pad_heads(k_4d, padded_heads)
+                v_4d = sp_pad_heads(v_4d, padded_heads)
+
+            sp_full_len = sp_seq_ctx.get(self.sp_stream)
+            q_4d = all_to_all_4D(q_4d, nccl_info.sp_group, scatter_dim=2, gather_dim=1,
+                                 full_seq_len=sp_full_len)
+            hpr = q_4d.shape[2]
+            head_start = sp_rank * hpr
+            k_4d = k_4d[:, :, head_start:head_start + hpr, :]
+            v_4d = v_4d[:, :, head_start:head_start + hpr, :]
+
+            local_heads = sp_num_local_real_heads(self.num_heads, padded_heads, sp_size, sp_rank)
+            if local_heads > 0:
+                if local_heads != hpr:
+                    q_4d = q_4d[:, :, :local_heads, :]
+                    k_4d = k_4d[:, :, :local_heads, :]
+                    v_4d = v_4d[:, :, :local_heads, :]
+                out = flash_attention(q_4d.flatten(2), k_4d.flatten(2), v_4d.flatten(2), num_heads=local_heads)
+                out = rearrange(out, "b s (n d) -> b s n d", n=local_heads)
+                if self.has_image_input:
+                    k_img = self.norm_k_img(self.k_img(img))
+                    v_img = self.v_img(img)
+                    k_img_4d = rearrange(k_img, "b s (n d) -> b s n d", n=self.num_heads)
+                    v_img_4d = rearrange(v_img, "b s (n d) -> b s n d", n=self.num_heads)
+                    if padded_heads != self.num_heads:
+                        k_img_4d = sp_pad_heads(k_img_4d, padded_heads)
+                        v_img_4d = sp_pad_heads(v_img_4d, padded_heads)
+                    k_img_4d = k_img_4d[:, :, head_start:head_start + local_heads, :]
+                    v_img_4d = v_img_4d[:, :, head_start:head_start + local_heads, :]
+                    img_out = flash_attention(q_4d.flatten(2), k_img_4d.flatten(2), v_img_4d.flatten(2), num_heads=local_heads)
+                    img_out = rearrange(img_out, "b s (n d) -> b s n d", n=local_heads)
+                    out = out + img_out
+                if out.shape[2] != hpr:
+                    out = sp_pad_heads(out, hpr)
+            else:
+                # Padding-only rank: no kernel. Keep out CONNECTED to q_4d (the only
+                # all_to_all'd tensor here; k/v are local context slices with no
+                # collective) so q's backward all_to_all fires on this rank too.
+                # `* 0.0` forces zero so it never affects real heads' fwd/bwd.
+                out = q_4d * 0.0
+            out = all_to_all_4D(out, nccl_info.sp_group, scatter_dim=1, gather_dim=2,
+                                full_seq_len=sp_full_len)
+            if padded_heads != self.num_heads:
+                out = out[:, :, :self.num_heads, :]
+            x = out.flatten(2)
+        else:
+            x = self.attn(q, k, v)
+            if self.has_image_input:
+                k_img = self.norm_k_img(self.k_img(img))
+                v_img = self.v_img(img)
+                y = flash_attention(q, k_img, v_img, num_heads=self.num_heads)
+                x = x + y
+
+        return self.o(x)
+
+
+class GateModule(nn.Module):
+    def __init__(self,):
+        super().__init__()
+
+    def forward(self, x, gate, residual):
+        return x + gate * residual
+
+class DiTBlock(nn.Module):
+    def __init__(self, has_image_input: bool, dim: int, num_heads: int, ffn_dim: int, eps: float = 1e-6,
+                 enable_bsa: bool = False, bsa_params: dict = None,
+                 enable_audio_guidance: bool = False,
+                 enable_audio_concentration_gate: bool = False,
+                 enable_variance_guidance: bool = False):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.ffn_dim = ffn_dim
+
+        self.self_attn = SelfAttention(
+            dim, num_heads, eps,
+            enable_bsa=enable_bsa, bsa_params=bsa_params,
+            enable_audio_guidance=enable_audio_guidance,
+            enable_audio_concentration_gate=enable_audio_concentration_gate,
+            enable_variance_guidance=enable_variance_guidance,
+        )
+        self.cross_attn = CrossAttention(
+            dim, num_heads, eps, has_image_input=has_image_input)
+        self.norm1 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
+        self.norm2 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
+        self.norm3 = nn.LayerNorm(dim, eps=eps)
+        self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU(
+            approximate='tanh'), nn.Linear(ffn_dim, dim))
+        self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+        self.gate = GateModule()
+
+    def forward(self, x, context, t_mod, freqs, grid_size=None, a2v_bridge_residual=None, timestep_ratio=None):
+        has_seq = len(t_mod.shape) == 4
+        chunk_dim = 2 if has_seq else 1
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+            self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(6, dim=chunk_dim)
+        if has_seq:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+                shift_msa.squeeze(2), scale_msa.squeeze(2), gate_msa.squeeze(2),
+                shift_mlp.squeeze(2), scale_mlp.squeeze(2), gate_mlp.squeeze(2),
+            )
+        input_x = modulate(self.norm1(x), shift_msa, scale_msa)
+        x = self.gate(x, gate_msa, self.self_attn(
+            input_x, freqs, grid_size=grid_size,
+            a2v_bridge_residual=a2v_bridge_residual,
+            timestep_ratio=timestep_ratio,
+        ))
+        x = x + self.cross_attn(self.norm3(x), context)
+        input_x = modulate(self.norm2(x), shift_mlp, scale_mlp)
+        x = self.gate(x, gate_mlp, self.ffn(input_x))
+        return x
+
+
+class MLP(torch.nn.Module):
+    def __init__(self, in_dim, out_dim, has_pos_emb=False):
+        super().__init__()
+        self.proj = torch.nn.Sequential(
+            nn.LayerNorm(in_dim),
+            nn.Linear(in_dim, in_dim),
+            nn.GELU(),
+            nn.Linear(in_dim, out_dim),
+            nn.LayerNorm(out_dim)
+        )
+        self.has_pos_emb = has_pos_emb
+        if has_pos_emb:
+            self.emb_pos = torch.nn.Parameter(torch.zeros((1, 514, 1280)))
+
+    def forward(self, x):
+        if self.has_pos_emb:
+            x = x + self.emb_pos.to(dtype=x.dtype, device=x.device)
+        return self.proj(x)
+
+
+class Head(nn.Module):
+    def __init__(self, dim: int, out_dim: int, patch_size: Tuple[int, int, int], eps: float):
+        super().__init__()
+        self.dim = dim
+        self.patch_size = patch_size
+        self.norm = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
+        self.head = nn.Linear(dim, out_dim * math.prod(patch_size))
+        self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
+
+    def forward(self, x, t_mod):
+        if len(t_mod.shape) == 3:
+            shift, scale = (self.modulation.unsqueeze(0).to(dtype=t_mod.dtype, device=t_mod.device) + t_mod.unsqueeze(2)).chunk(2, dim=2)
+            x = (self.head(self.norm(x) * (1 + scale.squeeze(2)) + shift.squeeze(2)))
+        else:
+            shift, scale = (self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(2, dim=1)
+            x = (self.head(self.norm(x) * (1 + scale) + shift))
+        return x
+
+
+class WanModel(ModelMixin, ConfigMixin):
+    _repeated_blocks = ("DiTBlock",)
+
+    @register_to_config
+    def __init__(
+        self,
+        dim: int,
+        in_dim: int,
+        ffn_dim: int,
+        out_dim: int,
+        text_dim: int,
+        freq_dim: int,
+        eps: float,
+        patch_size: Tuple[int, int, int],
+        num_heads: int,
+        num_layers: int,
+        has_image_input: bool,
+        has_image_pos_emb: bool = False,
+        has_ref_conv: bool = False,
+        add_control_adapter: bool = False,
+        in_dim_control_adapter: int = 24,
+        seperated_timestep: bool = False,
+        require_vae_embedding: bool = True,
+        require_clip_embedding: bool = True,
+        fuse_vae_embedding_in_latents: bool = False,
+        enable_bsa: bool = False,
+        bsa_params: dict = None,
+        enable_variance_guidance: bool = False,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.freq_dim = freq_dim
+        self.has_image_input = has_image_input
+        self.patch_size = patch_size
+        self.seperated_timestep = seperated_timestep
+        self.require_vae_embedding = require_vae_embedding
+        self.require_clip_embedding = require_clip_embedding
+        self.fuse_vae_embedding_in_latents = fuse_vae_embedding_in_latents
+        self.enable_bsa = enable_bsa
+        self.bsa_params = bsa_params
+        self.enable_variance_guidance = enable_variance_guidance
+
+        self.patch_embedding = nn.Conv3d(
+            in_dim, dim, kernel_size=patch_size, stride=patch_size)
+        self.text_embedding = nn.Sequential(
+            nn.Linear(text_dim, dim),
+            nn.GELU(approximate='tanh'),
+            nn.Linear(dim, dim)
+        )
+        self.time_embedding = nn.Sequential(
+            nn.Linear(freq_dim, dim),
+            nn.SiLU(),
+            nn.Linear(dim, dim)
+        )
+        self.time_projection = nn.Sequential(
+            nn.SiLU(), nn.Linear(dim, dim * 6))
+        self.blocks = nn.ModuleList([
+            DiTBlock(has_image_input, dim, num_heads, ffn_dim, eps,
+                     enable_bsa=enable_bsa, bsa_params=bsa_params,
+                     enable_audio_guidance=False,
+                     enable_audio_concentration_gate=False,
+                     enable_variance_guidance=enable_variance_guidance)
+            for _ in range(num_layers)
+        ])
+        self.head = Head(dim, out_dim, patch_size, eps)
+        head_dim = dim // num_heads
+        self.freqs = precompute_freqs_cis_3d(head_dim)
+
+        if has_image_input:
+            self.img_emb = MLP(1280, dim, has_pos_emb=has_image_pos_emb)  # clip_feature_dim = 1280
+        if has_ref_conv:
+            self.ref_conv = nn.Conv2d(16, dim, kernel_size=(2, 2), stride=(2, 2))
+        self.has_image_pos_emb = has_image_pos_emb
+        self.has_ref_conv = has_ref_conv
+        self.control_adapter = None
+
+    def patchify(self, x: torch.Tensor,control_camera_latents_input: torch.Tensor = None):
+        # NOTE(dhyu): avoid slow_conv
+        x = x.contiguous(memory_format=torch.channels_last_3d)
+        x = self.patch_embedding(x)
+        if self.control_adapter is not None and control_camera_latents_input is not None:
+            y_camera = self.control_adapter(control_camera_latents_input)
+            x = [u + v for u, v in zip(x, y_camera)]
+            x = x[0].unsqueeze(0)
+        grid_size = x.shape[2:]
+        x = rearrange(x, 'b c f h w -> b (f h w) c').contiguous()
+        return x, grid_size  # x, grid_size: (f, h, w)
+
+    def unpatchify(self, x: torch.Tensor, grid_size: torch.Tensor):
+        return rearrange(
+            x, 'b (f h w) (x y z c) -> b c (f x) (h y) (w z)',
+            f=grid_size[0], h=grid_size[1], w=grid_size[2],
+            x=self.patch_size[0], y=self.patch_size[1], z=self.patch_size[2]
+        )
+
+    def forward(self,
+                x: torch.Tensor,
+                timestep: torch.Tensor,
+                context: torch.Tensor,
+                clip_feature: Optional[torch.Tensor] = None,
+                y: Optional[torch.Tensor] = None,
+                use_gradient_checkpointing: bool = False,
+                use_gradient_checkpointing_offload: bool = False,
+                **kwargs,
+                ):
+        advance_dynamic_block_pass_id()  # once per real forward (not on recompute)
+        t = self.time_embedding(
+            sinusoidal_embedding_1d(self.freq_dim, timestep))
+        t_mod = self.time_projection(t).unflatten(1, (6, self.dim))
+        context = self.text_embedding(context)
+
+        if self.has_image_input:
+            x = torch.cat([x, y], dim=1)  # (b, c_x + c_y, f, h, w)
+            clip_embdding = self.img_emb(clip_feature)
+            context = torch.cat([clip_embdding, context], dim=1)
+
+        x, (f, h, w) = self.patchify(x)
+
+        freqs = torch.cat([
+            self.freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
+            self.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+            self.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
+        ], dim=-1).reshape(f * h * w, 1, -1).to(x.device)
+
+        visual_seq_len = f * h * w
+        if get_sequence_parallel_state():
+            sp_size = nccl_info.sp_size
+            sp_rank = nccl_info.rank_within_spgroup
+            assert_sp_splittable(visual_seq_len, sp_size, "visual sequence")
+            sp_seq_ctx.set(visual=visual_seq_len)
+            x = sp_split(x, sp_size, sp_rank, dim=1)
+            freqs = sp_split(freqs, sp_size, sp_rank, dim=0)
+
+        def create_custom_forward(module):
+            def custom_forward(*inputs):
+                return module(*inputs)
+            return custom_forward
+
+        grid_size = (f, h, w) if self.enable_bsa else None
+
+        for block in self.blocks:
+            if self.training and use_gradient_checkpointing:
+                if use_gradient_checkpointing_offload:
+                    with torch.autograd.graph.save_on_cpu():
+                        x = torch.utils.checkpoint.checkpoint(
+                            create_custom_forward(block),
+                            x, context, t_mod, freqs, grid_size,
+                            use_reentrant=False,
+                        )
+                else:
+                    x = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward(block),
+                        x, context, t_mod, freqs, grid_size,
+                        use_reentrant=False,
+                    )
+            else:
+                x = block(x, context, t_mod, freqs, grid_size=grid_size)
+
+        x = self.head(x, t)
+        if get_sequence_parallel_state():
+            x = all_gather(x, dim=1, group=nccl_info.sp_group,
+                           full_seq_len=visual_seq_len)
+        x = self.unpatchify(x, (f, h, w))
+        return x
+
+    def configure_bsa(self, enable_bsa: bool, bsa_params: dict = None):
+        """Runtime configuration of Block Sparse Attention on all video self-attention blocks."""
+        self.enable_bsa = enable_bsa
+        self.bsa_params = bsa_params
+        for block in self.blocks:
+            block.self_attn.enable_bsa = enable_bsa
+            block.self_attn.bsa_params = bsa_params or {}
+
+    def configure_audio_guidance(self, enable: bool, enable_concentration_gate: bool = False):
+        """Runtime configuration of audio-guided BSA on all video self-attention blocks.
+
+        Audio guidance modulates block-level QK scores with audio saliency
+        from the A→V bridge residual. Only effective when BSA is also enabled.
+
+        The boost strength γ (audio_boost_gamma) is read from bsa_params at runtime.
+        Set bsa_params['audio_boost_gamma'] to control the modulation strength
+        (default 1.0; 0.0 disables modulation).
+
+        Args:
+            enable: whether to enable audio-guided score modulation.
+            enable_concentration_gate: whether to enable Gate 2 (Audio Spatial
+                Concentration Gate) which down-weights audio saliency when the
+                bridge residual is spatially uniform (silence/ambient scenarios).
+        """
+        for block in self.blocks:
+            sa = block.self_attn
+            sa.enable_audio_guidance = enable
+            sa.enable_audio_concentration_gate = enable_concentration_gate
+
+    def configure_variance_guidance(self, enable: bool):
+        """Runtime configuration of channel-variance guidance on all video self-attention blocks.
+
+        Non-learnable, no gate. When enabled, block-level QK scores are boosted:
+          score_{i,j} *= (1 + γ * ρ_j)
+        where ρ_j is the V-block channel-variance density ∈ [0,1] (key/value side only).
+        Query-awareness comes from P^QK being per-query (zero-gate effect):
+        each query block has its own QK scores → different top-k after boost.
+
+        Completely independent from audio guidance — no shared code or state.
+        Only effective when BSA (enable_bsa) is also enabled.
+        Boost strength γ is read from bsa_params['variance_boost_gamma'] (default 1.0).
+
+        Args:
+            enable: whether to enable channel-variance guidance.
+        """
+        self.enable_variance_guidance = enable
+        for block in self.blocks:
+            block.self_attn.enable_variance_guidance = enable
+
+    def configure_taylor_sparse_attn(self, enable: bool):
+        """LIVEditor: selected (exact) + non-selected (Taylor) in one unified softmax.
+        Weights sum to exactly 1. No R_n. Mutually exclusive with rectified."""
+        for block in self.blocks:
+            block.self_attn.enable_taylor_sparse_attn = enable
+            if enable:
+                block.self_attn.enable_rectified_sparse_attn = False
+
+    def configure_rectified_sparse_attn(self, enable: bool):
+        """Rectified SpaAttn: o = R_n * o_spa + A_pool[nonsel] · V_pool.
+        Mutually exclusive with taylor."""
+        for block in self.blocks:
+            block.self_attn.enable_rectified_sparse_attn = enable
+            if enable:
+                block.self_attn.enable_taylor_sparse_attn = False
+
+    def _disable_non_dynamic_bsa_features(self, sa):
+        """Turn off every other BSA feature on a SelfAttention module so the
+        dynamic-block path is fully isolated (Section 8 requirement)."""
+        sa.enable_audio_guidance = False
+        sa.enable_audio_concentration_gate = False
+        sa.enable_timestep_reliability_gate = False
+        sa.enable_audio_weighted_pooling = False
+        sa.enable_variance_guidance = False
+        sa.enable_taylor_sparse_attn = False
+        sa.enable_rectified_sparse_attn = False
+
+    def configure_ivpq_dynamic_block(self, enable: bool):
+        """IVPQ anisotropic dynamic block shape (Section 8.5). Video self-attn only.
+        Mutually exclusive with penalty-matching AND with all other BSA features."""
+        for block in self.blocks:
+            sa = block.self_attn
+            sa.enable_ivpq_dynamic_block = enable
+            if enable:
+                sa.enable_penalty_dynamic_block = False
+                self._disable_non_dynamic_bsa_features(sa)
+
+    def configure_penalty_dynamic_block(self, enable: bool):
+        """Penalty-Matching anisotropic dynamic block shape (Section 8.6). Video
+        self-attn only. Mutually exclusive with IVPQ AND all other BSA features."""
+        for block in self.blocks:
+            sa = block.self_attn
+            sa.enable_penalty_dynamic_block = enable
+            if enable:
+                sa.enable_ivpq_dynamic_block = False
+                self._disable_non_dynamic_bsa_features(sa)

hymm/utils/__init__.py

@@ -0,0 +1 @@
+

hymm/utils/communications.py

@@ -0,0 +1,413 @@
+
+from typing import Any, List, Optional, Tuple
+
+import torch
+import torch.distributed as dist
+from torch.nn import functional as F
+
+
+
+
+def broadcast(input_: torch.Tensor, group: dist.ProcessGroup):
+    src = dist.get_global_rank(group, 0)
+    dist.broadcast(input_, src=src, group=group)
+
+
+# ---------------------------------------------------------------------------
+# Shard bookkeeping
+#
+# Sequences are split across SP ranks with ``torch.chunk`` semantics: every rank
+# gets ``ceil(S / sp)`` tokens except the last, which gets the remainder. The
+# collectives below need those sizes, but a rank cannot infer the global length
+# S from its own shard. Rather than negotiating it at runtime with a
+# ``dist.all_gather_object`` on every call (which costs a pickle round-trip plus
+# a device sync ~800 times per training step), the towers publish S through
+# ``sp_seq_ctx`` and pass it down as ``full_seq_len``.
+# ---------------------------------------------------------------------------
+
+def sp_chunk_sizes(total: int, world_size: int) -> List[int]:
+    """Shard sizes produced by ``torch.chunk(x, world_size)`` along one dim."""
+    chunk = (total + world_size - 1) // world_size
+    return [min(chunk, max(total - i * chunk, 0)) for i in range(world_size)]
+
+
+def assert_sp_splittable(total: int, world_size: int, name: str = "sequence"):
+    """``torch.chunk`` silently returns fewer than ``world_size`` chunks when the
+    sequence is very short, which would make high SP ranks index out of range."""
+    if world_size <= 1:
+        return
+    sizes = sp_chunk_sizes(total, world_size)
+    assert sizes[-1] > 0, (
+        f"{name} length {total} is too short for sp_size={world_size}: "
+        f"the trailing SP rank(s) would receive zero tokens. Reduce sp_size or "
+        f"increase the resolution / duration bucket."
+    )
+
+
+def sp_split(x: torch.Tensor, world_size: int, rank: int, dim: int) -> torch.Tensor:
+    """``torch.chunk``-compatible shard selection that never indexes past the end."""
+    if world_size <= 1:
+        return x
+    sizes = sp_chunk_sizes(x.shape[dim], world_size)
+    return torch.split(x, sizes, dim=dim)[rank]
+
+
+# ---------------------------------------------------------------------------
+# Ulysses head padding (for num_heads not divisible by sp_size, e.g. audio
+# DiT = 12 heads with sp_size = 8).
+#
+# `all_to_all_4D` requires `num_heads % sp_size == 0` because it scatters the
+# head dim evenly across SP ranks. These helpers let a module pad its head dim
+# with ZERO heads up to the next multiple of sp_size purely at runtime, with a
+# strict isolation contract enforced by the call sites:
+#
+#   1. Padding is appended (never weight-level) -> pretrained checkpoints load
+#      unchanged.
+#   2. After the head-scatter all_to_all, each rank holds the CONTIGUOUS global
+#      head block [r*hpr, (r+1)*hpr); real heads are the prefix [0, num_heads),
+#      so the real heads on a rank are always its FIRST `local_real` local heads.
+#      Call sites slice those off BEFORE any attention compute, so padded (zero)
+#      heads NEVER enter the BSA Triton / flash kernels and cannot produce NaNs
+#      or perturb the real heads.
+#   3. Ranks that hold only padding heads (local_real == 0) skip the kernel
+#      entirely and emit a zero-width head slice; they still join the all_to_all
+#      collectives so distributed/FSDP execution stays in lock-step.
+#   4. After the inverse all_to_all the head dim is truncated back to num_heads,
+#      so forward and backward for the real heads are bit-identical to the
+#      unpadded computation (padding contributes exactly zero gradient).
+# ---------------------------------------------------------------------------
+
+def sp_pad_num_heads(num_heads: int, sp_size: int) -> int:
+    """Smallest multiple of ``sp_size`` that is >= ``num_heads``."""
+    remainder = num_heads % sp_size
+    if remainder == 0:
+        return num_heads
+    return num_heads + (sp_size - remainder)
+
+
+def sp_pad_heads(x: torch.Tensor, padded_heads: int, dim: int = 2) -> torch.Tensor:
+    """Append zero-valued heads along ``dim`` so the head count becomes
+    ``padded_heads``. Autograd-safe: gradients routed to the appended heads are
+    discarded by the caller's later truncation, leaving the real heads' grads
+    untouched. Returns ``x`` unchanged when no padding is needed."""
+    cur = x.shape[dim]
+    if cur >= padded_heads:
+        return x
+    pad_shape = list(x.shape)
+    pad_shape[dim] = padded_heads - cur
+    pad = x.new_zeros(pad_shape)
+    return torch.cat([x, pad], dim=dim)
+
+
+def sp_num_local_real_heads(num_real_heads: int, padded_heads: int,
+                            sp_size: int, sp_rank: int) -> int:
+    """Number of this rank's local heads that are REAL after ``all_to_all_4D``
+    evenly distributes ``padded_heads``. Real heads are the global prefix
+    ``[0, num_real_heads)`` and rank ``r`` owns the contiguous global block
+    ``[r*hpr, (r+1)*hpr)`` (hpr = padded_heads // sp_size), so the real heads are
+    always this rank's FIRST ``k`` local heads."""
+    hpr = padded_heads // sp_size
+    return max(0, min(num_real_heads - sp_rank * hpr, hpr))
+
+
+# ---------------------------------------------------------------------------
+# Ulysses all-to-all
+# ---------------------------------------------------------------------------
+
+_UNKNOWN_SEQLEN_WARNED = False
+
+
+def _negotiate_seq_lens(local_len: int, group) -> List[int]:
+    global _UNKNOWN_SEQLEN_WARNED
+    if not _UNKNOWN_SEQLEN_WARNED:
+        print(
+            "[SP] all_to_all_4D called without full_seq_len; falling back to a "
+            "per-call all_gather_object. This is correct but slow -- publish the "
+            "sequence length through hymm.utils.parallel_states.sp_seq_ctx."
+        )
+        _UNKNOWN_SEQLEN_WARNED = True
+    seq_lens = [None] * dist.get_world_size(group)
+    dist.all_gather_object(seq_lens, local_len, group)
+    return seq_lens
+
+
+def _all_to_all_4D(
+    input: torch.Tensor,
+    scatter_idx: int = 2,
+    gather_idx: int = 1,
+    group=None,
+    full_seq_len: Optional[int] = None,
+) -> torch.Tensor:
+    """all-to-all for QKV.
+
+    Args:
+        input: a 4D tensor sharded along the scatter dim.
+        scatter_idx / gather_idx: 2/1 scatters heads and gathers sequence,
+            1/2 does the inverse.
+        full_seq_len: global (unsharded) sequence length. Required for the 2/1
+            direction, where it is the only way to know how much padding the
+            trailing shard carries.
+    """
+    assert (
+        input.dim() == 4
+    ), f"input must be 4D tensor, got {input.dim()} and shape {input.shape}"
+
+    seq_world_size = dist.get_world_size(group)
+
+    if scatter_idx == 2 and gather_idx == 1:
+        # (bs, seqlen/P, hc, hs) -> (bs, seqlen, hc/P, hs)
+        if full_seq_len is not None:
+            sizes = sp_chunk_sizes(int(full_seq_len), seq_world_size)
+            gap = sizes[0] - sizes[-1]
+        else:
+            seq_lens = _negotiate_seq_lens(input.shape[1], group)
+            gap = seq_lens[0] - seq_lens[-1] if seq_lens[-1] != seq_lens[0] else 0
+            assert gap >= 0, "only the trailing SP rank may hold a shorter shard"
+
+        if gap > 0 and dist.get_group_rank(group, dist.get_rank()) == seq_world_size - 1:
+            input = F.pad(input, (0, 0, 0, 0, 0, gap))
+
+        bs, shard_seqlen, hc, hs = input.shape
+        seqlen = shard_seqlen * seq_world_size
+        assert hc % seq_world_size == 0, (
+            f'Invalid Head size: {hc}, which should be divisible by spsize {seq_world_size}'
+        )
+        shard_hc = hc // seq_world_size
+
+        # transpose groups of heads with the seq-len parallel dimension, so that we can scatter them!
+        # (bs, seqlen/P, hc, hs) -reshape-> (bs, seq_len/P, P, hc/P, hs) -transpose(0,2)-> (P, seq_len/P, bs, hc/P, hs)
+        input_t = (
+            input.reshape(bs, shard_seqlen, seq_world_size, shard_hc, hs)
+            .transpose(0, 2)
+            .contiguous()
+        )
+
+        if seq_world_size > 1:
+            output = torch.empty_like(input_t)
+            # Stream-ordered: NCCL enqueues on the current stream and every
+            # consumer below is on that same stream, so no host sync is needed.
+            dist.all_to_all_single(output, input_t, group=group)
+        else:
+            output = input_t
+
+        output = output.reshape(seqlen, bs, shard_hc, hs)
+        # (seq_len, bs, hc/P, hs) -reshape-> (bs, seq_len, hc/P, hs)
+        output = output.transpose(0, 1).contiguous().reshape(bs, seqlen, shard_hc, hs)
+        if gap > 0:
+            output = output[:, :-gap]
+
+        return output
+
+    elif scatter_idx == 1 and gather_idx == 2:
+        # (bs, seqlen, hc/P, hs) -> (bs, seqlen/P, hc, hs)
+        # `seqlen` here is already the global length, so the split is derived
+        # locally with no negotiation.
+        bs, seqlen, shard_hc, hs = input.shape
+
+        hc = shard_hc * seq_world_size
+        if seqlen % seq_world_size != 0:
+            new_seqlen = (seqlen // seq_world_size + 1) * seq_world_size
+            gap = new_seqlen - seqlen
+            input = F.pad(input, (0, 0, 0, 0, 0, gap))
+            bs, seqlen, shard_hc, hs = input.shape
+        else:
+            gap = 0
+
+        shard_seqlen = seqlen // seq_world_size
+
+        # (bs, seqlen, hc/P, hs) -reshape-> (bs, P, seq_len/P, hc/P, hs) -transpose(0, 3)-> (hc/P, P, seqlen/P, bs, hs) -transpose(0, 1) -> (P, hc/P, seqlen/P, bs, hs)
+        input_t = (
+            input.reshape(bs, seq_world_size, shard_seqlen, shard_hc, hs)
+            .transpose(0, 3)
+            .transpose(0, 1)
+            .contiguous()
+            .reshape(seq_world_size, shard_hc, shard_seqlen, bs, hs)
+        )
+
+        if seq_world_size > 1:
+            output = torch.empty_like(input_t)
+            dist.all_to_all_single(output, input_t, group=group)
+        else:
+            output = input_t
+
+        output = output.reshape(hc, shard_seqlen, bs, hs)
+        # (hc, seqlen/N, bs, hs) -tranpose(0,2)-> (bs, seqlen/N, hc, hs)
+        output = output.transpose(0, 2).contiguous().reshape(bs, shard_seqlen, hc, hs)
+
+        if gap > 0 and dist.get_group_rank(group, dist.get_rank()) == seq_world_size - 1:
+            output = output[:, :-gap]
+
+        return output
+    else:
+        raise RuntimeError("scatter_idx must be 1 or 2 and gather_idx must be 1 or 2")
+
+
+class SeqAllToAll4D(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx: Any,
+        group: dist.ProcessGroup,
+        input: torch.Tensor,
+        scatter_idx: int,
+        gather_idx: int,
+        full_seq_len: Optional[int] = None,
+    ) -> torch.Tensor:
+        ctx.group = group
+        ctx.scatter_idx = scatter_idx
+        ctx.gather_idx = gather_idx
+        ctx.full_seq_len = full_seq_len
+
+        return _all_to_all_4D(input, scatter_idx, gather_idx, group=group,
+                              full_seq_len=full_seq_len)
+
+    @staticmethod
+    def backward(ctx: Any, *grad_output: torch.Tensor):
+        return (
+            None,
+            SeqAllToAll4D.apply(
+                ctx.group, *grad_output, ctx.gather_idx, ctx.scatter_idx,
+                ctx.full_seq_len,
+            ),
+            None,
+            None,
+            None,
+        )
+
+
+def all_to_all_4D(
+    input_: torch.Tensor,
+    group: dist.ProcessGroup,
+    scatter_dim: int = 2,
+    gather_dim: int = 1,
+    full_seq_len: Optional[int] = None,
+):
+    return SeqAllToAll4D.apply(group, input_, scatter_dim, gather_dim, full_seq_len)
+
+
+def _all_to_all(
+    input_: torch.Tensor,
+    world_size: int,
+    group: dist.ProcessGroup,
+    scatter_dim: int,
+    gather_dim: int,
+):
+    input_list = [
+        t.contiguous() for t in torch.tensor_split(input_, world_size, scatter_dim)
+    ]
+    output_list = [torch.empty_like(input_list[0]) for _ in range(world_size)]
+    dist.all_to_all(output_list, input_list, group=group)
+    return torch.cat(output_list, dim=gather_dim).contiguous()
+
+
+class _AllToAll(torch.autograd.Function):
+    """All-to-all communication.
+
+    Args:
+        input_: input matrix
+        process_group: communication group
+        scatter_dim: scatter dimension
+        gather_dim: gather dimension
+    """
+
+    @staticmethod
+    def forward(ctx, input_, process_group, scatter_dim, gather_dim):
+        ctx.process_group = process_group
+        ctx.scatter_dim = scatter_dim
+        ctx.gather_dim = gather_dim
+        ctx.world_size = dist.get_world_size(process_group)
+        output = _all_to_all(
+            input_, ctx.world_size, process_group, scatter_dim, gather_dim
+        )
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_output = _all_to_all(
+            grad_output,
+            ctx.world_size,
+            ctx.process_group,
+            ctx.gather_dim,
+            ctx.scatter_dim,
+        )
+        return (
+            grad_output,
+            None,
+            None,
+            None,
+        )
+
+
+def all_to_all(
+    input_: torch.Tensor, group: dist.ProcessGroup, scatter_dim: int = 2, gather_dim: int = 1
+):
+    return _AllToAll.apply(input_, group, scatter_dim, gather_dim)
+
+
+class _AllGather(torch.autograd.Function):
+    """All-gather along ``dim`` with autograd support.
+
+    ``full_seq_len`` lets both the forward gather and the backward split derive
+    the per-rank shard sizes locally, so neither direction needs the shape
+    negotiation collective.
+    """
+
+    @staticmethod
+    def forward(ctx, input_, dim, group, full_seq_len=None):
+        world_size = dist.get_world_size(group)
+        rank = dist.get_group_rank(group, dist.get_rank())
+
+        if full_seq_len is not None:
+            sizes = sp_chunk_sizes(int(full_seq_len), world_size)
+        else:
+            shapes = [None] * world_size
+            dist.all_gather_object(shapes, input_.shape[dim], group)
+            sizes = [int(s) for s in shapes]
+
+        ctx.dim = dim
+        ctx.group = group
+        ctx.sizes = sizes
+        ctx.rank = rank
+
+        input_ = input_.contiguous()
+        if all(s == sizes[0] for s in sizes):
+            # Uniform shards: a single ncclAllGather instead of the
+            # broadcast-per-rank fallback PyTorch falls back to for ragged inputs.
+            if dim == 0:
+                out_shape = list(input_.shape)
+                out_shape[0] = sizes[0] * world_size
+                flat = torch.empty(out_shape, dtype=input_.dtype, device=input_.device)
+                dist.all_gather_into_tensor(flat, input_, group=group)
+                return flat
+            tensor_list = [torch.empty_like(input_) for _ in range(world_size)]
+            dist.all_gather(tensor_list, input_, group=group)
+            return torch.cat(tensor_list, dim=dim)
+
+        tensor_list = []
+        for size in sizes:
+            shape = list(input_.shape)
+            shape[dim] = size
+            tensor_list.append(
+                torch.empty(shape, dtype=input_.dtype, device=input_.device)
+            )
+        dist.all_gather(tensor_list, input_, group=group)
+        return torch.cat(tensor_list, dim=dim)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_input = torch.split(grad_output, ctx.sizes, dim=ctx.dim)[ctx.rank]
+        return grad_input.contiguous(), None, None, None
+
+
+def all_gather(input_: torch.Tensor, dim: int = 1, group=None,
+               full_seq_len: Optional[int] = None):
+    """Gather ``input_`` shards along ``dim`` across the SP group.
+
+    Args:
+        input_: this rank's shard.
+        dim: concatenation dimension.
+        group: SP process group.
+        full_seq_len: global length along ``dim``. When supplied the shard sizes
+            are computed locally, avoiding a shape-negotiation collective in both
+            the forward and the backward pass.
+    """
+    return _AllGather.apply(input_, dim, group, full_seq_len)

hymm/utils/parallel_states.py

@@ -0,0 +1,453 @@
+"""Distributed / sequence-parallel state for MOVA.
+
+The topology follows HunyuanVideo-1.5: two independent device meshes over the
+same world.
+
+    world_mesh = [dp, sp]                 -> data parallel x Ulysses sequence parallel
+    fsdp_mesh  = [dp_replicate, fsdp_shard] -> FSDP2 parameter sharding
+
+SP shards activations, FSDP shards parameters, so the two meshes are
+orthogonal and ``fsdp_shard`` deliberately spans SP ranks as well.
+
+Model code only ever touches :data:`nccl_info` and
+:func:`get_sequence_parallel_state`, which are kept as thin views over the
+active :class:`ParallelDims`, so training and inference share one init path.
+"""
+
+import os
+import random
+from dataclasses import dataclass, field
+from datetime import timedelta
+from typing import Any, Optional
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from torch.distributed.device_mesh import init_device_mesh
+
+
+# ---------------------------------------------------------------------------
+# Topology
+# ---------------------------------------------------------------------------
+
+@dataclass
+class ParallelDims:
+    sp: int = 1
+    dp_replicate: int = 1
+    world_size: int = -1
+    use_dynamic_ring_attention: bool = False
+
+    world_mesh: Any = field(default=None, init=False, repr=False)
+    fsdp_mesh: Any = field(default=None, init=False, repr=False)
+    sp_group: Any = field(default=None, init=False, repr=False)
+    sp_rank: int = field(default=0, init=False)
+    dp_rank: int = field(default=0, init=False)
+    dp_size: int = field(default=1, init=False)
+
+    def __post_init__(self):
+        if self.world_size == -1:
+            if dist.is_initialized():
+                self.world_size = dist.get_world_size()
+            else:
+                self.world_size = int(os.environ.get("WORLD_SIZE", "1"))
+        if self.dp_replicate == -1:
+            # One replica per node; FSDP shards within a node.
+            local_world = int(os.environ.get("LOCAL_WORLD_SIZE", "8"))
+            assert self.world_size % local_world == 0, (
+                f"world_size({self.world_size}) must be divisible by "
+                f"LOCAL_WORLD_SIZE({local_world}) for dp_replicate=-1"
+            )
+            self.dp_replicate = self.world_size // local_world
+        assert self.sp >= 1 and self.dp_replicate >= 1
+        assert self.world_size % self.sp == 0, (
+            f"world_size({self.world_size}) must be divisible by sp({self.sp})"
+        )
+        assert self.world_size % self.dp_replicate == 0, (
+            f"world_size({self.world_size}) must be divisible by "
+            f"dp_replicate({self.dp_replicate})"
+        )
+
+    # -- mesh -------------------------------------------------------------
+
+    def build_mesh(self, device_type: str = "cuda"):
+        self.world_mesh = init_device_mesh(
+            device_type,
+            (self.world_size // self.sp, self.sp),
+            mesh_dim_names=("dp", "sp"),
+        )
+        self.fsdp_mesh = init_device_mesh(
+            device_type,
+            (self.dp_replicate, self.world_size // self.dp_replicate),
+            mesh_dim_names=("dp_replicate", "fsdp_shard"),
+        )
+
+        self.dp_rank = self.world_mesh["dp"].get_local_rank()
+        self.dp_size = self.world_mesh["dp"].size()
+        if self.sp_enabled:
+            self.sp_rank = self.world_mesh["sp"].get_local_rank()
+            self.sp_group = self.world_mesh["sp"].get_group()
+        else:
+            self.sp_rank = 0
+            self.sp_group = None
+        return self.world_mesh
+
+    # -- properties -------------------------------------------------------
+
+    @property
+    def sp_enabled(self) -> bool:
+        return self.sp > 1
+
+    @property
+    def dp_enabled(self) -> bool:
+        return self.dp_size > 1
+
+    @property
+    def dp_replicate_enabled(self) -> bool:
+        return self.dp_replicate > 1
+
+    @property
+    def dp_shard_enabled(self) -> bool:
+        return (self.world_size // self.dp_replicate) > 1
+
+    @property
+    def pp_enabled(self) -> bool:
+        return False
+
+    @property
+    def sp_mesh(self):
+        return self.world_mesh["sp"]
+
+    @property
+    def dp_mesh(self):
+        return self.world_mesh["dp"]
+
+    @property
+    def sp_src_rank(self) -> int:
+        """Global rank of local rank 0 inside this rank's SP group."""
+        if not self.sp_enabled:
+            return dist.get_rank() if dist.is_initialized() else 0
+        return dist.get_global_rank(self.sp_group, 0)
+
+
+# ---------------------------------------------------------------------------
+# Legacy view used by the model code (nccl_info.sp_group / .sp_size / ...)
+# ---------------------------------------------------------------------------
+
+class COMM_INFO:
+    def __init__(self):
+        self.sp_group = None
+        self.sp_size = 1
+        self.global_rank = 0
+        self.rank_within_spgroup = 0
+        self.parallel_dims = None
+        self.device_mesh = None
+        self.use_dynamic_ring_attention = False
+        self.sp_stream = None
+        self.sp_rank_list = None
+
+
+nccl_info = COMM_INFO()
+_SEQUENCE_PARALLEL_STATE = False
+_PARALLEL_DIMS: Optional[ParallelDims] = None
+
+
+def set_sequence_parallel_state(state: bool):
+    global _SEQUENCE_PARALLEL_STATE
+    _SEQUENCE_PARALLEL_STATE = state
+
+
+def get_sequence_parallel_state() -> bool:
+    return _SEQUENCE_PARALLEL_STATE
+
+
+def get_parallel_state() -> ParallelDims:
+    """Active topology, or a mesh-free single-rank default if none was set up.
+
+    The fallback deliberately does NOT build a device mesh: ``init_device_mesh``
+    is collective, and creating one from an incidental accessor call would mean
+    some ranks entering a collective their peers never reach.
+    """
+    global _PARALLEL_DIMS
+    if _PARALLEL_DIMS is None:
+        _PARALLEL_DIMS = ParallelDims(sp=1, dp_replicate=1, world_size=1)
+    return _PARALLEL_DIMS
+
+
+# ---------------------------------------------------------------------------
+# Init
+# ---------------------------------------------------------------------------
+
+def init_distributed(timeout_seconds: int = 5400) -> int:
+    """Initialise the default process group and bind this rank's device.
+
+    ``device_id`` is passed on purpose: it makes PyTorch build the default NCCL
+    communicator eagerly during ``init_process_group`` instead of lazily on the
+    first collective. Lazy communicator creation is a classic source of hangs
+    when different ranks reach their first collective through different code
+    paths (e.g. the MoE expert alternating between steps).
+
+    The timeout here covers *startup*, where a rank can legitimately sit in a
+    barrier for a long time while a peer streams a 15B checkpoint off a network
+    filesystem. Shrink it with :func:`set_collective_timeout` once the training
+    loop starts, where a long wait means a genuine desync rather than slow I/O.
+    """
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    if torch.cuda.is_available():
+        torch.cuda.set_device(local_rank)
+    if not dist.is_initialized():
+        kwargs = {}
+        if torch.cuda.is_available():
+            kwargs["device_id"] = torch.device("cuda", local_rank)
+        dist.init_process_group(
+            backend="nccl",
+            timeout=timedelta(seconds=timeout_seconds),
+            **kwargs,
+        )
+    return local_rank
+
+
+def set_collective_timeout(timeout_seconds: int, logger=None):
+    """Re-arm the watchdog with a tighter per-collective deadline.
+
+    Called once the steady-state training loop begins. A short deadline is what
+    turns a desync into an abort with a stack trace instead of a job that sits
+    idle until someone notices. Applied to every group we created, since the
+    watchdog is per-communicator.
+
+    Failures are swallowed deliberately: this touches a private torch API and
+    performs no collective, so a rank that cannot re-arm simply keeps the
+    startup timeout rather than diverging from its peers.
+    """
+    if not dist.is_initialized():
+        return False
+    try:
+        from torch.distributed.distributed_c10d import _set_pg_timeout
+    except ImportError:
+        if logger is not None:
+            logger.warning("[dist] _set_pg_timeout unavailable; keeping the startup timeout")
+        return False
+
+    timeout = timedelta(seconds=int(timeout_seconds))
+    groups = [None]
+    pd = _PARALLEL_DIMS
+    if pd is not None:
+        if pd.world_mesh is not None:
+            groups += [pd.world_mesh["sp"].get_group(), pd.world_mesh["dp"].get_group()]
+        if pd.fsdp_mesh is not None:
+            groups += [pd.fsdp_mesh["fsdp_shard"].get_group(),
+                       pd.fsdp_mesh["dp_replicate"].get_group()]
+    for group in groups:
+        try:
+            _set_pg_timeout(timeout, group)
+        except Exception as exc:  # noqa: BLE001
+            if logger is not None:
+                logger.warning(f"[dist] could not re-arm timeout on {group}: {exc}")
+            return False
+    if logger is not None:
+        logger.info(f"[dist] per-collective timeout tightened to {timeout_seconds}s for training")
+    return True
+
+
+def initialize_parallel_state(
+    sp: int = 1,
+    dp_replicate: int = 1,
+    use_dynamic_ring_attention: bool = False,
+    warmup: bool = True,
+) -> ParallelDims:
+    """Build the device meshes and publish them to the model-facing globals."""
+    global _PARALLEL_DIMS
+
+    parallel_dims = ParallelDims(
+        sp=sp,
+        dp_replicate=dp_replicate,
+        use_dynamic_ring_attention=use_dynamic_ring_attention,
+    )
+    if dist.is_initialized():
+        parallel_dims.build_mesh("cuda")
+
+    nccl_info.sp_size = parallel_dims.sp
+    nccl_info.global_rank = dist.get_rank() if dist.is_initialized() else 0
+    nccl_info.rank_within_spgroup = parallel_dims.sp_rank
+    nccl_info.sp_group = parallel_dims.sp_group
+    nccl_info.parallel_dims = parallel_dims
+    nccl_info.device_mesh = parallel_dims.world_mesh
+    nccl_info.use_dynamic_ring_attention = use_dynamic_ring_attention
+    if use_dynamic_ring_attention and parallel_dims.sp_enabled:
+        nccl_info.sp_stream = torch.cuda.Stream()
+        nccl_info.sp_rank_list = parallel_dims.world_mesh["sp"].mesh.tolist()
+
+    set_sequence_parallel_state(parallel_dims.sp_enabled)
+    _PARALLEL_DIMS = parallel_dims
+
+    if warmup and dist.is_initialized():
+        warmup_process_groups(parallel_dims)
+    return parallel_dims
+
+
+def warmup_process_groups(parallel_dims: Optional[ParallelDims] = None):
+    """Force NCCL communicator creation for every group we will ever use.
+
+    Every rank walks the same group list in the same order, so all communicators
+    exist before the training loop starts. Without this, the first collective on
+    a lazily-created group can land while other ranks are inside a different
+    collective, which NCCL cannot untangle and the watchdog eventually kills.
+    """
+    if not dist.is_initialized():
+        return
+    parallel_dims = parallel_dims or get_parallel_state()
+    device = torch.device("cuda", torch.cuda.current_device())
+    probe = torch.zeros(1, device=device)
+
+    groups = [None]  # default process group
+    if parallel_dims.world_mesh is not None:
+        groups.append(parallel_dims.world_mesh["sp"].get_group())
+        groups.append(parallel_dims.world_mesh["dp"].get_group())
+    if parallel_dims.fsdp_mesh is not None:
+        groups.append(parallel_dims.fsdp_mesh["fsdp_shard"].get_group())
+        groups.append(parallel_dims.fsdp_mesh["dp_replicate"].get_group())
+
+    for group in groups:
+        dist.all_reduce(probe, group=group)
+    torch.cuda.synchronize()
+
+
+def destroy_sequence_parallel_group():
+    if dist.is_initialized():
+        dist.destroy_process_group()
+
+
+# ---------------------------------------------------------------------------
+# Sequence-parallel sequence-length context
+# ---------------------------------------------------------------------------
+
+class SPSequenceContext:
+    """Full (pre-shard) sequence length of each SP-sharded stream.
+
+    Ulysses' head-scatter all-to-all has to undo the uneven ``torch.chunk``
+    split, which means it needs the *global* sequence length -- information a
+    rank cannot recover from its own shard. The towers publish it here once per
+    forward instead of paying a ``dist.all_gather_object`` inside every single
+    attention call (~800 object collectives per step at 40 layers x 2 towers).
+
+    Values stay valid until the next forward overwrites them, so activation
+    checkpointing recompute during backward reads the correct lengths.
+    """
+
+    __slots__ = ("visual", "audio")
+
+    def __init__(self):
+        self.visual = None
+        self.audio = None
+
+    def set(self, visual=None, audio=None):
+        if visual is not None:
+            self.visual = int(visual)
+        if audio is not None:
+            self.audio = int(audio)
+
+    def get(self, stream: str):
+        return getattr(self, stream, None)
+
+    def clear(self):
+        self.visual = None
+        self.audio = None
+
+
+sp_seq_ctx = SPSequenceContext()
+
+
+# ---------------------------------------------------------------------------
+# Cross-rank data / RNG alignment (replaces hy_parallelism)
+# ---------------------------------------------------------------------------
+
+def _flatten_tensors(obj, out):
+    if isinstance(obj, torch.Tensor):
+        out.append(obj)
+    elif isinstance(obj, dict):
+        for value in obj.values():
+            _flatten_tensors(value, out)
+    elif isinstance(obj, (list, tuple)):
+        for value in obj:
+            _flatten_tensors(value, out)
+    return out
+
+
+def _rebuild_with_tensors(obj, tensors, cursor):
+    if isinstance(obj, torch.Tensor):
+        value = tensors[cursor[0]]
+        cursor[0] += 1
+        return value
+    if isinstance(obj, dict):
+        return {k: _rebuild_with_tensors(v, tensors, cursor) for k, v in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        rebuilt = [_rebuild_with_tensors(v, tensors, cursor) for v in obj]
+        return tuple(rebuilt) if isinstance(obj, tuple) else rebuilt
+    return obj
+
+
+def sync_data_for_sp(obj, parallel_dims: Optional[ParallelDims] = None, force_object: bool = False):
+    """Make ``obj`` identical on every rank of the SP group (source = local rank 0).
+
+    Containers are walked recursively, so a whole dataloader batch can be passed
+    straight in. Non-tensor leaves go through ``broadcast_object_list``; tensors
+    are broadcast on the device.
+
+    The shapes are agreed on FIRST, in a single object broadcast, and any rank
+    whose local tensor disagrees reallocates before the data broadcast. That
+    matters because ``dist.broadcast`` demands matching shapes on every
+    participant: without the agreement step a rank that silently substituted a
+    different sample (the dataset retries on an I/O error, and an error need not
+    hit every rank) would corrupt the batch or hang the group -- the exact
+    failure this call exists to prevent.
+    """
+    parallel_dims = parallel_dims or get_parallel_state()
+    if not parallel_dims.sp_enabled:
+        return obj
+
+    src = parallel_dims.sp_src_rank
+    group = parallel_dims.sp_group
+
+    if force_object or not isinstance(obj, (torch.Tensor, dict, list, tuple)):
+        holder = [obj]
+        dist.broadcast_object_list(holder, src=src, group=group)
+        return holder[0]
+
+    tensors = _flatten_tensors(obj, [])
+    if not tensors:
+        holder = [obj]
+        dist.broadcast_object_list(holder, src=src, group=group)
+        return holder[0]
+
+    specs = [(tuple(t.shape), t.dtype) for t in tensors]
+    holder = [specs]
+    dist.broadcast_object_list(holder, src=src, group=group)
+    specs = holder[0]
+
+    device = torch.device("cuda", torch.cuda.current_device())
+    synced = []
+    for tensor, (shape, dtype) in zip(tensors, specs):
+        target_device = tensor.device
+        if tuple(tensor.shape) != shape or tensor.dtype != dtype:
+            staged = torch.empty(shape, dtype=dtype, device=device)
+        else:
+            staged = tensor if tensor.is_cuda else tensor.to(device)
+        dist.broadcast(staged, src=src, group=group)
+        synced.append(staged if target_device.type == "cuda" else staged.to(target_device))
+
+    return _rebuild_with_tensors(obj, synced, [0])
+
+
+def sync_random_states(step: int, base_seed: int = 0, dp_rank: int = 0):
+    """Re-seed the RNGs so an SP group samples identical noise.
+
+    Derived purely from ``(base_seed, dp_rank, step)``: every SP rank computes
+    the same value locally, so there is no collective and no device sync, while
+    distinct DP groups still see distinct noise.
+    """
+    seed = (int(base_seed) * 1_000_003 + int(dp_rank) * 10_007 + int(step)) % (2 ** 31 - 1)
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    return seed

packages.txt

@@ -0,0 +1 @@
+ffmpeg

requirements.txt

@@ -0,0 +1,14 @@
+torch==2.8.0
+torchvision==0.23.0
+diffusers==0.33.0
+transformers==4.57.1
+accelerate>=1.10.0
+safetensors>=0.6.2
+huggingface_hub[hf_xet]
+einops>=0.8.1
+loguru
+ftfy
+imageio[ffmpeg]>=2.37.0
+numpy
+scipy
+sentencepiece