[perf] FastH3 two-Spark: gather packed FP4 inputs before head projections

Opt-in h3_vsa_fp4.sp_input_gather (default false; requires sp_projection_overlap) gathers the
already-packed FP4 inputs and rotary tables, then computes each rank's assigned Q/K/V heads locally;
the gate projection overlaps the smaller transfer. Supported on the measured GB10/cuDNN shape only;
other shapes keep the existing exchange. Enabled in the Spark-pair example.
T2AV on 2 Sparks: 19.35 -> 18.76 s/clip (merge gate check, same SP reduction order rules).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

.agents/lessons/2026-10-23_asynchronous-gemm-parity.md

@@ -0,0 +1,27 @@
+---
+date: 2026-10-23
+category: evaluation
+severity: important
+---
+
+# Synchronized GEMM parity does not establish request determinism
+
+A fixed cuDNN FP4 GEMM plan matched isolated captured outputs, a complete step,
+and three scored requests on one GB10. Repeated requests on another GB10 later
+produced different latents. Restoring the prior row path reproduced the known
+latents on three fresh requests. The replacement was withdrawn.
+
+Comparing every GEMM output with `torch.equal` made all 400 calls and the complete
+request match. That comparison synchronizes execution; it can conceal an
+ordering or lifetime problem. Repeated component tests also passed. The exact
+cause remains unresolved, so these observations do not establish a vendor bug.
+
+A separate larger workspace happened to make two requests match, but cuDNN
+reported zero required workspace for the selected plan. Treat that as an
+allocation/timing perturbation, not evidence of a workspace fix.
+
+After component parity, repeat complete requests without additional per-kernel
+synchronization, including another device when the change is intended to be
+portable. Compare latent and decoded-output hashes. Preserve failing evidence
+and restore the stable path when the proposed gain is small and the cause is
+unresolved; passing a finite scored seed set does not establish determinism.

.agents/lessons/2026-10-23_compiled-forward-deepcopy.md

@@ -0,0 +1,38 @@
+---
+date: 2026-10-23
+category: evaluation
+severity: important
+---
+
+# Copying an instance-compiled module can benchmark the original module
+
+## What happened
+
+A decoder quantization experiment replaced every intended linear in a copied
+module, yet measured exactly the original output and almost the original time.
+The replacement layers were present in `named_modules()`, which was insufficient
+evidence that the new computation actually ran.
+
+## Root cause
+
+`ComposedPipelineBase._compile_with_conditions` assigns
+`submodule.forward = torch.compile(submodule.forward)`. The returned function
+closes over a bound method of the original module. Python's `copy.deepcopy`
+copies module state but retains function objects and their closures. The copied
+instance therefore still has a compiled `forward` that calls the original
+instance. Unwrapping `_orig_mod` alone does not handle this form of compilation.
+
+## Fix
+
+Prefer constructing a fresh module and loading the intended state. If copying
+this known instance-compiled form for an experiment, remove the copied instance's
+`forward` override before modifying children, then compile its newly bound class
+method. Do not remove arbitrary user-defined forward overrides blindly.
+
+## Prevention
+
+Inspect the callable as well as the module tree. Verify a replacement executes
+with a temporary hook or profiler trace, compare its output against the intended
+numerical oracle, and run final timing in a fresh process. An unchanged output
+from a supposedly lossy quantization is a reason to investigate dispatch before
+accepting the measurement.

.agents/lessons/2026-10-23_triton-affine-gather.md

@@ -0,0 +1,26 @@
+---
+date: 2026-10-23
+category: performance
+severity: important
+---
+
+# Wide Triton gathers can turn an affine slice into expensive selection code
+
+A decoder RMSNorm prototype kept a 2,048-element row in registers, then used
+`tl.gather` with indices `512 * chunk + 4 * lane + component` to extract the
+128 lanes needed by its original reduction order. Although every index was
+an affine expression, this lowered poorly on the Spark's Triton/SM121 toolchain.
+The eight-row variant emitted PTX with 168,062 virtual 32-bit registers and
+spent minutes in ptxas. Earlier variants were numerically exact but took
+392–1,475 ms for work the existing kernel completed in about 2.4 ms.
+
+For this layout, reshaping the row into four 512-element chunks and extracting
+them with `tl.trans` and `tl.split` avoided the selection expansion. The same
+ordered FMAs then compiled promptly and matched the original outputs. Keeping
+the row in registers still gave no meaningful speedup, so the rewrite was dropped.
+
+Inspect generated IR and benchmark one small launch before compiling a whole
+parameter sweep of wide gathers. Prefer static reshape/split operations for
+known affine partitions; reserve gathers for indexing that actually needs them.
+After changing an in-place residual path, also zero invalid padded rows before
+quantization: the old two-pass kernel obtained those zeros from its masked reload.

.agents/skills/ci-runner/SKILL.md

@@ -1,15 +1,15 @@
 ---
 name: ci-runner
-description: Work on FastVideo's Spark 4 and Slinky Slurm GPU CI lanes, static Buildkite graph, trusted ci-runner policy, lane scripts, and hardware validation.
+description: Work on FastVideo's Slurm-only, change-aware GPU CI lanes, static Buildkite graph, trusted ci-runner policy, lane scripts, and GB200 validation.
 ---
 
-# GPU CI lanes
+# Slinky Slurm CI lanes
 
 FastVideo's `ci-runner` Buildkite queue is the control plane for all active
-GPU CI. The trusted uploader runs on the Slinky login plane. Slinky agents
-dispatch cluster lanes through Slurm into isolated Enroot containers. The
-Spark 4 agent dispatches eligible single-GPU lanes through local Slurm into
-isolated Docker containers. Both execute the immutable PR SHA.
+GPU CI. A private host-owned dispatcher leases GPUs from the Slinky Slurm tray
+and runs the immutable PR SHA inside an isolated Enroot container. Buildkite
+pipeline upload and Slurm submission occur on the login plane; every test
+payload executes on Slurm compute.
 
 The files under `fastvideo/tests/modal/` and `.buildkite/scripts/pr_test.sh`
 are dormant rollback code. Never add an active Buildkite or slash-command
@@ -33,14 +33,6 @@ and the coordination contract with that bundle.
   A Spark runner sets `FASTVIDEO_CI_BACKEND=spark4`; gate cluster-only tests on
   that flag, never on memory heuristics. Keep
   `SPARK_CAPABLE_LANES` in `test_ci_test_collection.py` in sync.
-- The trusted uploader fills `FV_SPARK_BACKEND`: `spark4` for `fastcheck`,
-  `merge` and `direct` builds, including unset scope treated as Fastcheck,
-  while an agent tagged `queue=ci-runner,backend=spark4` is connected and
-  not paused, `slurm`
-  otherwise. Selection happens at upload time. Legacy graphs without backend
-  tags need explicit normalization; setting the variable alone does not pin
-  those jobs to Slurm. It rejects any other `backend` value. The upload step
-  in each entry pipeline requires `backend=slurm` so it never lands on the Spark runner.
 - Active CI contains no `pr_test.sh` command, Modal invocation, default queue,
   Buildkite plugin, `soft_fail`, or job-controlled artifact glob.
 - The six Fastcheck lanes use `:microscope:` labels. Full-Suite-only lanes use
@@ -72,7 +64,7 @@ and the coordination contract with that bundle.
 - Tests must preserve an inherited `MASTER_PORT`. Packed containers share the
   tray network namespace, so the private runner assigns a distinct port range
   per GPU lease and the SSIM scheduler assigns task offsets within its range.
-- The Slinky ARM64 runner image includes the pinned FA4 CuTe overlay validated on
+- The ARM64 runner image includes the pinned FA4 CuTe overlay validated on
   GB200. Keep SSIM at `FASTVIDEO_FA4=1` because its references were seeded with
   FA4; keep lanes with FA2 baselines at `FASTVIDEO_FA4=0`. A runner image change
   must revalidate both the FA4 import and an actual GB200 forward kernel.

.github/mergify.yml

@@ -9,28 +9,6 @@ merge_protections:
       - check-success=fastcheck-passed
       - check-success=full-suite-passed
 
-  # Mirrors the `paths` filter of .github/workflows/ci-macos-mlx.yml so the
-  # required checks always exist whenever this protection applies.
-  - name: MLX smoke requirements
-    if:
-      - base = main
-      - or:
-          - files~=^\.github/workflows/ci-macos-mlx\.yml$
-          - files~=^fastvideo/mlx_runtime/
-          - files~=^fastvideo/tests/mlx/
-          - files~=^fastvideo/tests/platforms/test_mps_vsa_error\.py$
-          - files~=^fastvideo/tests/platforms/test_cpu_sdpa\.py$
-          - files~=^fastvideo/platforms/cpu\.py$
-          - files~=^fastvideo/platforms/mps\.py$
-          - files~=^fastvideo/platforms/__init__\.py$
-          - files~=^fastvideo/__init__\.py$
-          - files~=^examples/inference/basic/mlx_[^/]*\.py$
-          - files~=^fastvideo/benchmarks/mlx_[^/]*\.py$
-          - files~=^pyproject\.toml$
-    success_conditions:
-      - check-success=mlx-smoke
-      - check-success=mlx-smoke-linux-cpu
-
 pull_request_rules:
 
   # ============================================================

.github/scripts/plan_merge_ci.py

@@ -131,11 +131,6 @@ class FamilyCoverage:
 
 
 FAMILY_COVERAGE = (
-    FamilyCoverage(
-        re.compile(r"(^|[/_.-])qwen[_-]?image[_-]?2[._-]?1([/_.-]|$)"),
-        ("test_qwen_image21.py", ),
-        ("test_qwen_image21_similarity.py", ),
-    ),
     FamilyCoverage(
         re.compile(r"(^|[/_.-])flux[_-]?2([/_.-]|$)"),
         ("test_flux2_klein.py", ),

.github/workflows/ci-macos-mlx.yml

@@ -2,9 +2,6 @@ name: macOS MLX Smoke
 
 on:
   pull_request:
-    # The defaults plus ready_for_review: both jobs skip drafts, and Mergify requires them
-    # for MLX changes, so marking a draft ready must run them.
-    types: [opened, synchronize, reopened, ready_for_review]
     branches: [main]
     paths:
       - ".github/workflows/ci-macos-mlx.yml"
@@ -88,10 +85,7 @@ jobs:
             fastvideo/tests/mlx/test_mlx_affine_dq_gemm.py \
             fastvideo/tests/mlx/test_mlx_minimax_h3_parity.py \
             fastvideo/tests/mlx/test_mlx_minimax_h3_vsa.py \
-            fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_sdpa.py \
-            fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_nax.py \
             fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_regressions.py \
-            fastvideo/tests/mlx/test_mlx_minimax_h3_prefix_kv.py \
             fastvideo/tests/mlx/test_mlx_minimax_h3_fast_mode.py \
             fastvideo/tests/mlx/test_mlx_minimax_h3_fast_spatial.py \
             fastvideo/tests/mlx/test_mlx_minimax_h3_ref_layout.py \
@@ -163,10 +157,7 @@ jobs:
             fastvideo/tests/mlx/test_mlx_affine_dq_gemm.py \
             fastvideo/tests/mlx/test_mlx_minimax_h3_parity.py \
             fastvideo/tests/mlx/test_mlx_minimax_h3_vsa.py \
-            fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_sdpa.py \
-            fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_nax.py \
             fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_regressions.py \
-            fastvideo/tests/mlx/test_mlx_minimax_h3_prefix_kv.py \
             fastvideo/tests/mlx/test_mlx_minimax_h3_fast_mode.py \
             fastvideo/tests/mlx/test_mlx_minimax_h3_fast_spatial.py \
             fastvideo/tests/mlx/test_mlx_minimax_h3_ref_layout.py \

docker/Dockerfile.spark-ci

@@ -1,11 +0,0 @@
-# syntax=docker/dockerfile:1.7
-# CI-only GB10 image. The regular Spark image keeps its native SDPA defaults.
-ARG SPARK_BASE_IMAGE=ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:py3.12-cuda13.0.0-latest
-FROM ${SPARK_BASE_IMAGE}
-
-# VMoBA and the Wan causal gate require FA2. Install only this wheel so the
-# validated torch/flashinfer/CUTLASS stack from the base image stays intact.
-ARG FLASH_ATTN_WHEEL=https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.22/flash_attn-2.8.3+cu130torch2.12-cp312-cp312-linux_aarch64.whl
-RUN /root/.local/bin/uv pip install --python /opt/venv/bin/python --no-deps "${FLASH_ATTN_WHEEL}" && \
-    /opt/venv/bin/python -c "import glob, shutil; [shutil.rmtree(p) for p in glob.glob('/opt/venv/lib/python*/site-packages/flash_attn/cute')]" && \
-    /opt/venv/bin/python -c "import flash_attn, flashinfer; assert hasattr(flash_attn, 'flash_attn_varlen_func'); print('FA2', flash_attn.__version__, 'flashinfer', flashinfer.__version__)"

docs-site/lib/cookbook-pages.ts

@@ -111,7 +111,6 @@ const familyInfo: Record<string, { title: string; path: string; logo: string; ca
   kandinsky5: { title: "Kandinsky 5", path: "kandinsky5", logo: "kandinsky", category: "Video" },
   kandinsky6: { title: "Kandinsky 6", path: "kandinsky6", logo: "kandinsky", category: "Video" },
   flux: { title: "FLUX", path: "flux", logo: "black-forest-labs", category: "Image" },
-  qwen_image21: { title: "Qwen-Image-2.1", path: "qwen-image21", logo: "qwen", category: "Image" },
   sd35: { title: "Stable Diffusion", path: "stable-diffusion", logo: "stabilityai", category: "Image" },
   stable_audio: { title: "Stable Audio", path: "stable-audio", logo: "stabilityai", category: "Audio" },
   mmaudio: { title: "MMAudio", path: "mmaudio", logo: "fastvideo", category: "Audio" },
@@ -133,7 +132,6 @@ const TASK_GROUPS: Array<[RegExp, string]> = [
   [/video\/text to audio|video.*to audio/, "Video to audio"],
   [/text to audio/, "Text to audio"],
   [/text to image/, "Text to image"],
-  [/image to image/, "Image to image"],
   [/reference media to video/, "Reference to video"],
   [/first\/last frame to video/, "Image to video"],
   [/text or image to video/, "Text/image to video"],
@@ -163,7 +161,6 @@ export const modelDirectory = [...new Set(sourceRecipes.map((recipe) => recipe.m
       "tencent-hunyuan": "Tencent",
       kandinsky: "Kandinsky Lab",
       "black-forest-labs": "Black Forest Labs",
-      qwen: "Qwen",
       stabilityai: "Stability AI",
       nvidia: "NVIDIA",
     };

docs-site/navigation.json

@@ -253,10 +253,6 @@
         "label": "FLUX",
         "href": "/cookbook/flux"
       },
-      {
-        "label": "Qwen-Image-2.1",
-        "href": "/cookbook/qwen-image21"
-      },
       {
         "label": "Stable Diffusion",
         "href": "/cookbook/stable-diffusion"

docs-site/tests/cookbook.spec.ts

@@ -223,7 +223,7 @@ test("inherited values, explicit false, zero and null retain their meanings", as
 test("catalog discovery opens dedicated model recipes and filters by family and task", async ({ page }) => {
   await page.goto("/FastVideo/cookbook");
   const cards = page.locator(".recipe-card");
-  await expect(cards).toHaveCount(37);
+  await expect(cards).toHaveCount(36);
   await expect(cards.first()).toContainText("FastH3 OmniRef");
   await page.getByLabel("Model family", { exact: true }).selectOption("Wan");
   await page.getByLabel("Search recipes").fill("Wan2.1 I2V");
@@ -277,7 +277,7 @@ test("output tabs support keyboard navigation and precision chips use the backen
   await page.getByLabel("Search recipes").fill("does not exist");
   await expect(page.getByRole("heading", { name: "No matching recipes" })).toBeVisible();
   await page.getByRole("button", { name: "Clear filters" }).click();
-  await expect(page.locator(".recipe-card")).toHaveCount(37);
+  await expect(page.locator(".recipe-card")).toHaveCount(36);
 });
 
 test("cookbook cards and controls remain readable in dark mode and reflow at 320px", async ({ page }) => {
@@ -297,7 +297,7 @@ test("cookbook cards and controls remain readable in dark mode and reflow at 320
 test("image and audio models have logos and dedicated source recipe pages", async ({ page }) => {
   await page.goto("/FastVideo/cookbook");
   await page.getByRole("button", { name: "Image", exact: true }).click();
-  await expect(page.locator(".recipe-card")).toHaveCount(5);
+  await expect(page.locator(".recipe-card")).toHaveCount(4);
   const card = page.locator(".recipe-card").filter({ hasText: "FLUX.2 Klein" });
   await expect(card.locator("img")).toBeVisible();
   await card.click();
@@ -309,26 +309,6 @@ test("image and audio models have logos and dedicated source recipe pages", asyn
   await expect(page.getByRole("button", { name: "Serving", exact: true })).toHaveCount(0);
   await page.getByLabel("Seed", { exact: true }).fill("0");
   await expect(page.getByRole("button", { name: "Copy example + command", exact: true })).toBeEnabled();
-  await page.goto("/FastVideo/cookbook/qwen-image21");
-  const qwen = page.locator(".recipe-card").filter({ hasText: "Qwen-Image-2.1" });
-  await expect(qwen.locator("img")).toHaveAttribute("src", /\/qwen\.webp$/);
-  await qwen.click();
-  await expect(page.getByRole("heading", { level: 1 })).toHaveText("Qwen-Image-2.1");
-  await expect(page.getByRole("button", { name: "Serving", exact: true })).toHaveCount(0);
-  await page.getByLabel("Height", { exact: true }).fill("512");
-  await page.getByRole("tab", { name: "config.yaml", exact: true }).click();
-  const qwenConfig = page.getByLabel("Source example", { exact: true });
-  await expect(qwenConfig).toContainText("model_path: Qwen/Qwen-Image-2.1");
-  await expect(qwenConfig).toContainText("height: 512");
-  await expect(qwenConfig).toContainText("dit_layerwise: true");
-  await expect(page.getByRole("button", { name: "Download YAML", exact: true })).toBeEnabled();
-  await page
-    .getByRole("group", { name: "Recipe variant", exact: true })
-    .getByRole("button", { name: "Qwen-Image-2.1 references and editing", exact: true })
-    .click();
-  await page.getByRole("tab", { name: "Commands", exact: true }).click();
-  await expect(page.getByLabel("Inference command", { exact: true })).toContainText("--reference source.png");
-  await expect(page.getByRole("button", { name: "Copy example + command", exact: true })).toBeEnabled();
   await page.goto("/FastVideo/cookbook");
   await page.getByRole("button", { name: "Audio", exact: true }).click();
   await expect(page.locator(".recipe-card")).toHaveCount(3);
@@ -604,11 +584,11 @@ test("directory filters round-trip through the URL", async ({ page }) => {
   await page.goto("/FastVideo/cookbook?category=Image&runtime=cuda");
   await expect(page.getByRole("button", { name: "Image", exact: true })).toHaveAttribute("aria-pressed", "true");
   await expect(page.getByLabel("Runtime", { exact: true })).toHaveValue("cuda");
-  await expect(page.locator(".recipe-card")).toHaveCount(5);
+  await expect(page.locator(".recipe-card")).toHaveCount(4);
   await page.reload();
   await expect(page.getByRole("button", { name: "Image", exact: true })).toHaveAttribute("aria-pressed", "true");
   await expect(page.getByLabel("Runtime", { exact: true })).toHaveValue("cuda");
-  await expect(page.locator(".recipe-card")).toHaveCount(5);
+  await expect(page.locator(".recipe-card")).toHaveCount(4);
 });
 
 test("a saved target missing from the catalog stays visible and clearable", async ({ page }) => {

docs/assets/cookbook-recipes.json

@@ -1,5 +1,5 @@
 {
-  "version": 19,
+  "version": 18,
   "recipes": [
     {
       "id": "fastwan21-t2v",
@@ -429,44 +429,6 @@
       "expected_artifact": "PNG images under outputs/flux_dev/samples/",
       "limitations": ["FLUX.1 is loadable by ID but registers no model_family in fastvideo/registry.py; it is grouped under FLUX for documentation only."]
     },
-    {
-      "id": "qwen-image21-t2i",
-      "family": "qwen_image21",
-      "stage": "inference",
-      "task": "Text to image",
-      "label": "Qwen-Image-2.1",
-      "summary": "Generate a 1024 × 1024 image in 40 steps on one CUDA GPU, with the CPU offload settings recorded in the maintained YAML.",
-      "model": "Qwen/Qwen-Image-2.1",
-      "source": "examples/inference/basic/qwen_image21_t2i.yaml",
-      "command": "fastvideo generate --config examples/inference/basic/qwen_image21_t2i.yaml",
-      "gpu_types": ["NVIDIA"],
-      "hardware": {"platform": "cuda", "gpu_count": 1, "evidence": "source-configured"},
-      "evidence": "Source-backed",
-      "expected_artifact": "RGBA PNG image at outputs/qwen_image21/teapot.png",
-      "limitations": ["The source uses DiT layerwise offload, text-encoder and VAE offload, and lazy loading. GPU memory requirements have not been established for this recipe."],
-      "related": ["qwen-image21-i2i"]
-    },
-    {
-      "id": "qwen-image21-i2i",
-      "family": "qwen_image21",
-      "stage": "inference",
-      "task": "Image to image",
-      "label": "Qwen-Image-2.1 references and editing",
-      "summary": "Generate or edit an image from a prompt and up to ten ordered reference images, using the maintained single-GPU Python example.",
-      "model": "Qwen/Qwen-Image-2.1",
-      "source": "examples/inference/basic/basic_qwen_image21.py",
-      "command": "python examples/inference/basic/basic_qwen_image21.py --reference source.png --prompt 'Change the shirt to blue; retain the subject and composition.' --output outputs/qwen_image21/edit.png",
-      "gpu_types": ["NVIDIA"],
-      "hardware": {"platform": "cuda", "gpu_count": 1, "evidence": "source-configured"},
-      "evidence": "Source-backed",
-      "expected_artifact": "RGBA PNG image at outputs/qwen_image21/edit.png",
-      "limitations": [
-        "Replace source.png with an existing image. Repeat --reference for additional images and describe their order in the prompt.",
-        "Mask images use the same ordered-reference path; they do not impose hard pixel-preservation constraints.",
-        "The source enables CPU offload and lazy loading. This entry does not establish a minimum GPU memory requirement."
-      ],
-      "related": ["qwen-image21-t2i"]
-    },
     {
       "id": "sd35-medium-t2i",
       "family": "sd35",

docs/assets/logos/SOURCES.md

@@ -13,7 +13,6 @@ plain typographic tile instead — that tile is a UI placeholder, not a logo.
 | `nvidia.webp` | [Official NVIDIA Hugging Face organization avatar](https://huggingface.co/nvidia) | Cosmos and GEN3C cards, Cosmos page header |
 | `kandinsky.webp` | [Official Kandinsky Lab Hugging Face organization avatar](https://huggingface.co/kandinskylab) | Kandinsky 5 and Kandinsky 6 cards and page headers |
 | `black-forest-labs.webp` | [Official Black Forest Labs Hugging Face organization avatar](https://huggingface.co/black-forest-labs) | FLUX family card and page header |
-| `qwen.webp` | [Official Qwen Hugging Face organization avatar](https://huggingface.co/Qwen) | Qwen-Image-2.1 family card and page header |
 | `minimax.webp` | [Official MiniMax Hugging Face organization avatar](https://huggingface.co/MiniMaxAI) | MiniMax H3 family card and page header |
 | `stabilityai.webp` | [Official Stability AI Hugging Face organization avatar](https://huggingface.co/stabilityai) | Stable Diffusion and Stable Audio cards and page headers |
 | `fastvideo.webp` | [Official FastVideo Hugging Face organization avatar](https://huggingface.co/FastVideo) | Matrix Game and MMAudio cards (converted weights published by this org) |

docs/contributing/ci_architecture.md

@@ -17,7 +17,7 @@ PR opened or updated
   |     style, lint, type, spelling, Markdown, workflow syntax, filenames
   |
   |-- Tier 2: Fastcheck
-  |     Buildkite schedules six lanes across Spark 4 and Slinky Slurm
+  |     Buildkite schedules six lanes on the Slinky Slurm cluster
   |     encoder, VAE, transformer, kernel, unit, DreamVerse
   |
   |-- /merge or ready label
@@ -25,7 +25,7 @@ PR opened or updated
         `-- Tier 3: Change-aware merge gate
               trusted base-branch planner classifies the complete PR diff
               Buildkite adds only relevant integration, quality, training,
-              API, or performance lanes on Spark 4 or Slinky Slurm
+              API, or performance lanes on Slinky Slurm
               |
               pass -> Mergify squash-merges when all merge conditions pass
               fail -> fix, push, and re-run
@@ -43,12 +43,11 @@ CI is not one monolithic job:
   updates, docs deployment, image builds, package publishing, and community
   automations.
 - Buildkite owns the GPU test graph, statuses, and trusted dispatch control
-  plane. Slinky agents dispatch tests from the login plane; the Spark 4 agent
-  dispatches local Slurm jobs.
-- Slinky Slurm runs cluster lanes in isolated Enroot containers using a
-  persistent four-GPU allocation. Spark 4 runs eligible single-GPU lanes in
-  isolated Docker containers through its own Slurm queue. Both test the
-  immutable PR SHA.
+  plane. Its agent runs on the Slurm login plane; it does not execute test
+  payloads.
+- Slinky Slurm is the only active CI compute backend. A host-owned dispatcher
+  leases GPUs from a persistent four-GPU allocation and runs each lane in an
+  isolated Enroot container at the immutable PR SHA.
 - Mergify owns merge protection, labeling, and the final squash merge.
 
 The old files under `fastvideo/tests/modal/` are retained as dormant manual
@@ -100,7 +99,7 @@ debugging a hook implementation.
 | Attribute | Value |
 |---|---|
 | Triggered by | Buildkite PR builds with `TEST_SCOPE=fastcheck` or unset |
-| Compute | Spark 4 when connected and unpaused at upload time; otherwise Slinky Slurm. Both use `ci-runner`. |
+| Compute | Slinky Slurm (`ci-runner` queue) |
 | Definition | `.buildkite/pipeline.yml` |
 | Entrypoint | Trusted host driver -> `.buildkite/scripts/unit_test.sh` or `.buildkite/scripts/lanes/*.sh` |
 
@@ -123,61 +122,15 @@ Setting it to `spark4` moves only the Spark-capable lanes to a GB10 runner.
 That runner exports `FASTVIDEO_CI_BACKEND=spark4`, which the lane scripts and
 a few tests honor to skip GB200-only coverage (for example the LTX-2 finetune
 test, which needs more than 128 GB). GB10 golden references are keyed by
-`NVIDIA_GB10`; see `fastvideo/tests/golden_gate`.
-
-The trusted uploader on the Slurm login plane picks the value at upload time.
-It sets `FV_SPARK_BACKEND=spark4` only when both hold:
-
-- the build's `TEST_SCOPE` is `fastcheck`, `merge` or `direct`, or is unset;
-- the Buildkite agents API reports an agent tagged
-  `queue=ci-runner,backend=spark4` that is connected and not paused.
-
-Otherwise it sets `slurm`. For pipelines with the backend tags above,
-`/test full`, the weekly Full Suite, a paused or disconnected Spark agent, and
-any failure to read the agents API select Slurm for every lane. This choice
-happens at upload time; it does not move jobs already queued for Spark.
-Older graphs without backend tags need the normalization described below.
-
-The uploader rejects any `backend` value other than the two forms in the
-table, so a pull request cannot pin a lane to the Spark runner directly. The upload step in each entry pipeline requires `backend=slurm`, so
-the Spark runner never takes it and this policy always runs on the login plane.
-
-The regular GB10 inference image intentionally omits FlashAttention to match
-native Spark output. CI uses `docker/Dockerfile.spark-ci` to add the pinned
-FA2 wheel required by VMoBA and Wan causal tests. This image keeps the base
-Torch, FlashInfer, and CUTLASS packages and removes the wheel's unsupported
-FA4 CuTe implementation. Build it separately from the inference image:
-
-```bash
-docker build -f docker/Dockerfile.spark-ci \
-  --build-arg SPARK_BASE_IMAGE=<validated-GB10-image> \
-  -t fastvideo-ci-gb10:<version> .
-```
-
-Golden references must match this CI image, including its FA2 version. Seed
-only from unchanged upstream main, verify in separate processes, and review
-before publishing new `golden_gates/NVIDIA_GB10/` files. Preserve existing
-references for other GPUs. A successful import alone is insufficient; validate
-an actual FA2 forward, the VMoBA tests, and each lane's complete payload before
-selecting the new image in the host-owned runner configuration.
-
-The host runner must allowlist Spark-capable lane types before allocating the
-GPU. Label each Docker container with its job token and remove it from the
-outer runner on completion or cancellation. Slurm killing the Docker client
-does not stop a container owned by the Docker daemon.
-
-The trusted uploader must preserve each literal `backend: "slurm"` even when
-it selects `FV_SPARK_BACKEND=spark4`. Older PR revisions without backend tags
-must be normalized to `backend: "slurm"` before upload; `queue: "ci-runner"`
-alone also matches the Spark agent. Update the private uploader's lane policy
-with routing changes and verify the resolved agent queries on a canary build.
+`NVIDIA_GB10`; see `fastvideo/tests/golden_gate`. The private uploader policy
+must allow the `backend` agent key before this routing is active.
 
 ### Tier 3: Change-Aware Merge Gate
 
 | Attribute | Value |
 |---|---|
 | Triggered by | `/merge`, adding `ready`, or a new push to a PR that already has `ready` |
-| Compute | Spark 4 or Slinky Slurm, selected per lane (`ci-runner` queue) |
+| Compute | Slinky Slurm only (`ci-runner` queue) |
 | Definition | `.buildkite/pipeline.yml` |
 | Entrypoint | `/opt/fastvideo-ci-runner/run-ci` (`run-unit` is a compatibility wrapper) |
 
@@ -443,7 +396,7 @@ Valid direct test names:
 | `/test eval` | `eval` |
 
 The temporary `<name>-ci` spellings remain accepted as compatibility aliases;
-they select the same test lane and do not identify a second backend.
+they select the same Slurm lane and do not identify a second backend.
 
 When a direct test completes successfully, Buildkite posts
 `direct-test-completed`. `.github/workflows/ci-aggregate-status.yml` then reads
@@ -523,7 +476,7 @@ Every active test selection lives in `.buildkite/scripts/unit_test.sh` or a
 focused `.buildkite/scripts/lanes/<lane>.sh`. The private, agent-owned lane
 table binds each internal `*_ci` type to that script, its GPU count, wall-clock
 limit, dependency extras, kernel-build policy, secrets, and artifacts. The
-internal suffix is an implementation detail; backend tags select the compute host.
+internal suffix is an implementation detail; there is only one active backend.
 
 SSIM uses `fastvideo/tests/ssim/ci_runner.py` inside a single four-GPU lease.
 It discovers `REQUIRED_GPUS` and `*_MODEL_TO_PARAMS` with AST parsing, then

docs/contributing/env_vars.md

@@ -215,7 +215,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry.
 | `FASTVIDEO_MLX_COMPILE`                                    | bool                                                               | `0`                                                   | performance | Compile the MLX DiT forward with mx.compile.                                                                                                                                                                                                                                                                                                                                                                                                |
 | `FASTVIDEO_MLX_FAST_NORM`                                  | bool                                                               | `0`                                                   | performance | Use MLX fast normalization kernels.                                                                                                                                                                                                                                                                                                                                                                                                         |
 | `FASTVIDEO_MLX_DQ_GEMM`                                    | str                                                                | `1`                                                   | performance | MLX dequantized GEMM for affine-quantized weights: 0 turns it off, 1 uses the measured minimum row count, and an integer sets the minimum row count.                                                                                                                                                                                                                                                                                        |
-| `FASTVIDEO_MLX_QUANT_GEMM`                                 | str                                                                | `auto`                                                | performance | MLX H3 affine-quantized GEMM path: auto keeps FASTVIDEO_MLX_DQ_GEMM behaviour, dequant always uses dequantize + dense GEMM for wide inputs, qmm always uses mx.quantized_matmul.                                                                                                                                                                                                                                                            |
+| `FASTVIDEO_MLX_QUANT_GEMM`                                 | str                                                                | `auto`                                                | performance | MLX H3 affine-quantized GEMM path: `auto` keeps `FASTVIDEO_MLX_DQ_GEMM` behaviour, `dequant` always uses dequantize + dense GEMM for wide inputs, `qmm` always uses `mx.quantized_matmul`.                                                                                                                                                                                                                                                  |
 | `FASTVIDEO_LTX2_VAE_CHANNELS_LAST_3D`                      | bool                                                               | `1`                                                   | performance | LTX-2 VAE uses the channels_last_3d memory format.                                                                                                                                                                                                                                                                                                                                                                                          |
 | `FASTVIDEO_LTX2_DISABLE_AUDIO_AUTOCAST`                    | bool                                                               | `1`                                                   | performance | LTX-2 audio decoding runs without CUDA autocast. Deprecated names: `LTX2_DISABLE_AUDIO_AUTOCAST`.                                                                                                                                                                                                                                                                                                                                           |
 | `FASTVIDEO_FLUX2_DISABLE_BF16_REDUCED_PRECISION_REDUCTION` | bool                                                               | `0`                                                   | performance | Flux denoising disables reduced-precision reductions in bf16 matmuls, which tightens accumulation for the 4-step Klein model.                                                                                                                                                                                                                                                                                                               |

docs/cookbook/index.md

@@ -7,7 +7,6 @@ The interactive catalog groups models by family and generates YAML for supported
 - [Wan and FastMetal](wan.md)
 - [LTX](ltx.md)
 - [FLUX](flux.md)
-- [Qwen-Image-2.1](qwen-image21.md)
 - [Stable Diffusion](stable-diffusion.md)
 - [Stable Audio](stable-audio.md)
 - [MMAudio](mmaudio.md)

docs/cookbook/qwen-image21.md

@@ -1,108 +0,0 @@
-# Qwen-Image-2.1 recipes
-
-<RecipeFamily family="qwen_image21" />
-
-Generate images, edit an image, or combine up to ten ordered references with
-`Qwen/Qwen-Image-2.1`. Outputs are RGBA PNGs. Request a transparent background
-in the prompt when needed.
-
-## Run on a GPU
-
-Follow the [CUDA installation guide](../getting_started/installation/gpu.md), then
-run from the repository root on one CUDA GPU (`tp_size=sp_size=1`). The examples
-use CPU offload and VAE tiling. The loader downloads the complete checkpoint;
-use `--model-path /path/to/Qwen-Image-2.1` for local weights. No conversion is needed.
-
-```bash
-python examples/inference/basic/basic_qwen_image21.py \
-  --prompt 'A red ceramic teapot, isolated on a transparent background.' \
-  --output outputs/teapot.png
-```
-
-This saves a 1024×1024 PNG after 40 denoising steps. For instruction editing,
-replace `source.png` with your image:
-
-```bash
-python examples/inference/basic/basic_qwen_image21.py \
-  --reference source.png \
-  --prompt 'Change the shirt to blue; retain the subject and composition.' \
-  --output outputs/edit.png
-```
-
-Repeat `--reference` for reference-to-image generation, annotations or masks:
-
-```bash
-python examples/inference/basic/basic_qwen_image21.py \
-  --reference subject.png --reference style.png \
-  --prompt 'Paint the subject from image 1 in the style of image 2.'
-```
-
-A mask is another reference interpreted through the prompt; pixels outside it
-are not guaranteed to remain identical.
-
-## Configuration
-
-- Set `--height`, `--width`, `--steps` and `--seed` as needed. Dimensions must be
-  positive multiples of 32; use `--height 2048 --width 2048` for 2K output.
-- `--reference-resolution` defaults to 1024 and must be a positive multiple of
-  32. References retain their aspect ratio and are resized to approximately
-  that resolution squared in pixels, independently of output dimensions.
-- Prefix KV caching avoids recomputing text/reference tokens each step. The
-  Python example caches on CPU; `--no-use-kv-cache` disables reuse.
-- For two-branch CFG, set `--true-cfg-scale` above 1 and supply
-  `--negative-prompt` (an empty string is valid). Embedded guidance is not used.
-
-For the CLI, use the checked-in [T2I YAML](../../examples/inference/basic/qwen_image21_t2i.yaml)
-or [editing YAML](../../examples/inference/basic/qwen_image21_edit.yaml) with
-`fastvideo generate --config <file>`. The T2I YAML keeps prefix KV caches on CUDA;
-the Python example and editing YAML default to CPU caches. See
-[Configuration](../inference/configuration.md) for the Python API and YAML options.
-
-## Attention
-
-Dense image attention uses FastVideo's shared automatic backend selector, which
-prefers FlashAttention on supported CUDA hardware. Masked text attention uses
-SDPA. To request FA4 with its dependencies installed on supported hardware:
-
-```bash
-export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
-export FASTVIDEO_FA4=1
-```
-
-Use `FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA` and `FASTVIDEO_FA4=0` for upstream
-parity comparisons. Different attention kernels can produce different pixels.
-See [Optimizations](../inference/optimizations.md) for backend and offload settings.
-
-## Measured memory
-
-Peaks below are in **GiB**, measured on one **NVIDIA GB200** with BF16,
-1024×1024 output, seed 42, compilation disabled and two 40-step requests per
-process. Image cases use one reference. These measurements do not establish
-minimum hardware capacity; an actual 24 GB GPU has not been verified.
-
-| Configuration | KV cache | GPU allocated | GPU reserved | Device VRAM including driver | Host process-tree RSS |
-| --- | --- | ---: | ---: | ---: | ---: |
-| T2I YAML, SDPA, offloaded, tiled | CUDA | 13.28 | 22.54 | 24.19 | 43.87 |
-| Plain I2I Python, SDPA, offloaded, tiled | CPU | 13.29 | 16.11 | 17.76 | 48.29 |
-| Editing Python, SDPA, offloaded, tiled | CPU | 13.29 | 16.11 | 17.76 | 49.25 |
-| T2I, FA4, resident models, untiled VAE | CUDA | 32.77 | 35.50 | 37.09 | 13.12 |
-
-Measurements cover loading, cold/warm requests, PNG export and shutdown.
-Allocated/reserved values use PyTorch counters; device VRAM and simultaneous
-process-tree RSS were sampled every 20 ms. Device VRAM includes the GB200's
-0.72 GiB driver reservation. Allow host RAM headroom beyond process RSS.
-Offloaded rows explicitly use SDPA; the resident FA4 row disables offload and
-tiling and is a separate configuration, not a controlled latency comparison.
-
-## Validation and scope
-
-On GB200, 40-step upstream SDPA parity passed for eleven 1024×1024 cases:
-T2I, plain I2I, editing, reference-to-image, RGBA generation/editing, annotations,
-mask images, true CFG, uncached execution and ten references. Separate 2048×2048
-T2I also passed. Denoising trajectories and output pixels matched exactly.
-FA4 block and T2I/edit image regressions passed against their matching references.
-See the [validation runbook](../../tests/local_tests/qwen_image21/README.md)
-for reproducible commands, pinned revisions and numerical diagnostics.
-
-The supported scope is single GPU, batch one and unquantized inference.
-SP/TP, LoRA, prompt rewriting and HTTP serving are not included.

docs/design/inference_schema_parity_inventory.yaml

@@ -89,6 +89,7 @@ surfaces:
       taeh3_checkpoint: "Optional local TAEH3 safetensors path; model-specific optimization not yet represented in the typed public schema."
       taeh3_chunk_size: "TAEH3 temporal chunk length; model-specific optimization not yet represented in the typed public schema."
       h3_vae_nvfp4_decoder: "Exported NVFP4 MiniMax-H3 video decoder file that replaces the VAE decoder at load; model-specific optimization not yet represented in the typed public schema."
+      h3_vae_nvfp4_quantize_dense: "Quantize the loaded MiniMax-H3 decoder to NVFP4 in memory without changing its depth; carried through generator.pipeline.experimental."
       vae_parallel_encode: "MiniMax-H3 sequence-parallel reference VAE encode opt-in; model-specific optimization not yet represented in the typed public schema."
       vae_parallel_decode_strategy: "Chunk-transport collective for vae_parallel_decode; model-specific optimization not yet represented in the typed public schema."
       attention_backend: "Process-wide default attention-backend request applied per component at load time; kernel-selection knob not yet represented in the typed public schema."
@@ -149,7 +150,6 @@ surfaces:
       preprocess_text_funcs: "Internal text preprocessing hooks."
       postprocess_text_funcs: "Internal text postprocessing hooks."
       scheduler_step_in_fp32: "Runtime scheduler precision toggle; not part of the public typed inference API."
-      output_channels: "Decoded pixel channel count fixed by the model; used to allocate the output buffer."
 
   pipeline_config_extensions:
     preset_owned:
@@ -443,15 +443,6 @@ surfaces:
         sources: [fastvideo.configs.pipelines.kandinsky6.Kandinsky6TI2VAConfig]
       sample_fps:
         sources: [fastvideo.configs.pipelines.kandinsky6.Kandinsky6TI2VAConfig]
-      reference_resolution:
-        target: generator.pipeline.experimental.reference_resolution
-        sources: [fastvideo.configs.pipelines.qwen_image21.QwenImage21PipelineConfig]
-      use_kv_cache:
-        target: generator.pipeline.experimental.use_kv_cache
-        sources: [fastvideo.configs.pipelines.qwen_image21.QwenImage21PipelineConfig]
-      kv_cache_device:
-        target: generator.pipeline.experimental.kv_cache_device
-        sources: [fastvideo.configs.pipelines.qwen_image21.QwenImage21PipelineConfig]
     compatibility_only:
       batch_size: "Gen3C inference-only tuning field pending typed batching design."
       gradient_checkpointing: "Gen3C inference-only compatibility field pending typed batching design."
@@ -470,8 +461,6 @@ surfaces:
       text_offset: "MagiHuman internal data-proxy text alignment offset."
       vae_stride: "MagiHuman internal VAE/data-proxy stride setting."
       vsa_ref_keep_rate: "MiniMax-H3 PDD reference-video VSA keep rate; defaults to the checkpoint's trained value."
-      video_scheduler_shift: "MiniMax-H3 training-time video noise clock; inference samples with the checkpoint's scheduler config."
-      audio_scheduler_shift: "MiniMax-H3 training-time audio noise clock; inference samples with the checkpoint's scheduler config."
       z_dim: "MagiHuman internal VAE latent channel setting."
       vocoder_config: "Legacy internal component config object."
       vocoder_precision: "Precision override pending dedicated component precision design."

docs/generate_examples.py

@@ -36,7 +36,6 @@ COOKBOOK_FAMILIES = {
     "kandinsky5",
     "kandinsky6",
     "flux",
-    "qwen_image21",
     "sd35",
     "minimax_h3",
     "stable_audio",

docs/getting_started/installation/spark_pair.md

@@ -149,6 +149,39 @@ fastvideo generate --config examples/inference/basic/basic_fasth3_spark_pair.yam
 
 Stop the cluster when you are done: `ray stop` on both nodes.
 
+For eager FastH3 FP4 attention, `generator.pipeline.experimental.h3_vsa_fp4.sp_head_chunks`
+overlaps head-group communication with independent attention work. It defaults to `1`.
+The head count must divide evenly by `sp_size * sp_head_chunks`; on a two-Spark pair FastH3's
+56 heads allow 2, 4, 7, 14 or 28 chunks. The pair example config uses 7 (four heads per group). The grouped path packs directly into each exchange buffer and writes
+received heads into their original output positions. Changing the attention batch shape can change
+floating-point reductions, so compare generated quality as well as wall time for your checkpoint.
+
+The optional `h3_vsa_fp4.sp_projection_overlap: true` starts each Q/K/V exchange as soon as its
+projection is packed, overlapping traffic with the remaining projection GEMMs. It defaults to
+`false` and requires two ranks, head dimension 128, eager BF16 inference and compatible NVFP4
+QKV/gate projections with shared input quantization. Other configurations use the ordinary
+head-group exchange. Packing arithmetic and FP8 transport precision are unchanged.
+
+The optional `h3_vsa_fp4.sp_input_gather: true` gathers the existing packed FP4 inputs and
+rotary tables before computing each rank's assigned Q/K/V heads. The gate projection overlaps
+that smaller input transfer. It defaults to `false`, requires `sp_projection_overlap: true`,
+and currently supports the measured GB10/cuDNN shape: 7430 local rows, 5376 input channels,
+56 heads and bias-free NVFP4 QKV projections. Other shapes use the existing exchange path.
+All projections, normalization, rotary operations and FP8 packet rounding remain present;
+no steps or tokens are omitted. Validate generated output when changing this GEMM shape.
+
+The optional `h3_vsa_fp4.sp_pooled_key_mean: true` reuses tile statistics for key smoothing
+on the grouped tile-64 path with FP8 prefix attention and `smooth_k: true`. It defaults to
+`false`. All key rows still contribute, including zero padding in the normalization length,
+but the changed reduction order can alter generated samples. Validate quality for the checkpoint
+and workload before enabling it.
+
+The optional `h3_vsa_fp4.sp_direct_keys: true` avoids the intermediate BF16 key tensor on
+that pooled-mean path. It defaults to `false`. Key preparation reads the current exchange's
+FP8 packets and applies the same FP32 multiply, BF16 rounding and smoothing as the ordinary
+receive path. All keys and their tile statistics still contribute; other attention configurations
+retain ordinary key stores.
+
 ## Released V2, FastH3-QAD-v2 and Trim NVFP4 stacks
 
 The public eight-forward V2 and Trim stacks include the NVFP4 encoder and

docs/getting_started/installation/spark_performance.md

@@ -390,10 +390,9 @@ off on other GPUs unless a matched benchmark validates the split there.
 `FASTVIDEO_H3_VAE_NVFP4_DECODER=/path/to/decoder.pt` swaps in an NVFP4
 decoder distilled with quantization-aware training and exported by
 `scripts/distill/minimax_h3_nvfp4_decoder/export_deploy.py`. At load the
-INT8 overlay is skipped, the VAE is depth-cut to the decoder's block count
-(the 26-layer light VAE hosts an 8-block decoder; a 36-block decoder needs
-the full `MiniMaxAI/MiniMax-H3` `vae/` in the stack's `vae/` folder), and
-only the packed FP4 weights stay resident (0.31 GiB for 8 blocks instead of
+INT8 overlay is skipped, the VAE uses the exported decoder's block count
+(a complete export can also replace a shallower light decoder when its
+head/width geometry matches), and only the packed FP4 weights stay resident (0.31 GiB for 8 blocks instead of
 2.3 GiB). The decode stage runs bf16 autocast for these decoders, the
 precision they were trained in. The decoder GEMMs use the cuDNN backend on
 GB10; outputs match CUTLASS bit for bit.
@@ -406,6 +405,16 @@ read straight onto the GPU and never materializes the full-precision weights
 of the quantized linears; it decodes bit-identically to packing the bf16 file
 at load.
 
+To quantize the original decoder without replacing it with a distilled export,
+set `generator.pipeline.experimental.h3_vae_nvfp4_quantize_dense: true` and
+`generator.engine.compile.vae_enabled: false`. This experimental option keeps
+every decoder block, quantizes its block linears to NVFP4 with unit activation
+scale during loading, and uses the existing fused BF16 decoder path. It leaves
+the checkpoint files unchanged and skips the INT8 overlay. Clear
+`FASTVIDEO_H3_VAE_NVFP4_DECODER` and `h3_vae_nvfp4_decoder`: exported and runtime
+quantization selections are mutually exclusive. The output is lossy and needs
+quality validation for the selected checkpoint. Blackwell FP4 hardware is required.
+
 V2 at 832x480, 124 frames, seed 1234, one Spark, warm calls:
 
 | Video decoder | End to end | Denoise | Video decode | Peak allocated |

docs/inference/fasth3_omniref_mlx.md

@@ -73,34 +73,6 @@ Quantized checkpoints pick a linear path per layer. `--quant-gemm` (or
 
 The same flag is on the example: `--quant-gemm {auto,dequant,qmm}`.
 
-### Opt-in M5 attention backend
-
-`--vsa-impl sdpa` evaluates the trained VSA tile mask through fused MLX SDPA,
-sharing each attention head's K/V across its query tiles. On M5, the built-in
-matrix operations can be faster than the custom SIMD kernel even though they
-visit masked keys. The default `auto` selection is unchanged.
-
-This backend keeps the reference resolution, selected tiles, ragged-row masks,
-and compression gate. Accumulation differs from the SIMD backend, so outputs
-are not bit-identical. It supports tiles 64/128 with head dimension 128 and
-raises if fused SDPA is unavailable. Prefix K/V reuse is not supported; do not
-combine it with `--reuse-prefix-kv`. Query tiles are processed in chunks of at
-most 128, capped at 256 MiB per boolean-mask chunk with two chunks in flight.
-
-`--vsa-impl nax` is a separate opt-in backend for Apple M5 GPUs. It uses Metal
-TensorOps to visit only the selected key blocks and their valid rows, avoiding
-the dense mask and masked-key traversal of the SDPA backend. Dense prefix
-queries and the trained compression branch keep their existing implementations.
-The selected reference tokens, output resolution and denoising steps are unchanged.
-
-The NAX backend uses private helper headers from the installed MLX package;
-it accepts only the header hashes validated with MLX 0.32.4. Changed headers need
-numerical and generation validation before their hashes are accepted. Missing headers, unsupported GPUs
-and compilation failures raise errors rather than silently selecting another
-backend. Tiles 64/128 and head dimension 128 are supported. Prefix K/V reuse is
-not supported. Floating-point accumulation differs from both SIMD and SDPA,
-so the generated clip can differ even with the same seed. `auto` remains unchanged.
-
 ## Generate
 
 ```bash

docs/inference/open_h3_regenerate_2k.md

@@ -1,196 +0,0 @@
-# Open-H3-Regenerate-2K
-
-`OpenH3Regenerate2KPipeline` regenerates a 768P [MiniMax-H3](https://huggingface.co/MiniMaxAI/MiniMax-H3) draft video
-(1344x768, 121 frames, 24 fps) at **2560x1440** with three refine steps. It is video to video: the draft is the input,
-and the output is a silent 2K MP4. FastVideo ports NVlabs'
-[Open-H3-Regenerate-2K](https://huggingface.co/Efficient-Large-Model/Open-H3-Regenerate-2K)
-([project page](https://nvlabs.github.io/Sana/Open-H3-Regenerate-2K/),
-[reference code](https://github.com/NVlabs/Sana/tree/sol-engine/models/Open-H3-Regenerate-2K)); the dense BF16
-`standard_fwd` path of that repo is the parity reference. The accelerated `sol_engine_fwd` path (Sol-Attn, MXFP8, CUDA
-graphs) and the fast 24-layer VAE decoder are not ported.
-
-```text
-draft video -> H3 VAE encode -> latent upscaler (2x, 3D conv) -> 3 Euler refine steps -> H3 VAE decode -> 2560x1440
-                                                 H3 DiT + merged LoRA + Sharp curve AdaLN branches
-```
-
-## What runs
-
-| Step | Details |
-|---|---|
-| Draft encode | `ffmpeg` resamples to 24 fps, cover-scales (bilinear) and centre-crops to 1344x768 (`draft_size: 768p`, or 1024x576 for `576p`), holds the last frame to 121 frames plus 3 copies. The FP32 H3 VAE posterior is sampled under seed 42. |
-| Latent upscale | The author's 3D-conv upscaler (`fastvideo/models/upsamplers/minimax_h3_latent_upscaler.py`), BF16, to the 160x90 latent grid. |
-| Refine | Euler on sigmas `[0.909375, 0.857143, 0.666667, 0]`, CFG 1, model time `1 - sigma`, initial state `(1 - 0.909375) * upscaled + 0.909375 * noise`. `num_inference_steps: 2` uses `[0.7, 0.494118, 0]` (softer, about 30% faster). |
-| Transformer | The base MiniMax-H3 DiT (the same module as the other H3 pipelines) with the released merged LoRA (FP32, rank 576 on 312 linears, scale 1) and 50 LMS Sharp curve AdaLN branches. |
-| Decode | The released H3 VAE, or an NVFP4 decoder on Blackwell (see below); the output is trimmed to 121 frames. |
-
-The merged LoRA loads through the standard H3 community-LoRA path as a constructor-time adapter: the pipeline sets
-`lora_path` to `<model>/lora/open_h3_regenerate_2k_lora.safetensors` when you do not pass one. The Sharp branches take an
-8-D timestep curve (`lora/adaln_t_table.bin`, 1025x8) instead of the layer input, so they cannot be folded into weights;
-`sharp_curve.py` adds them to the AdaLN projection outputs with forward hooks.
-
-**Audio.** The reference feeds an encoded-silence audio stream to the DiT as conditioning and writes the video with no
-audio track (`-an`); it neither regenerates nor passes through the draft's audio. FastVideo does the same: the draft's
-soundtrack is ignored and the output MP4 has no audio.
-
-## Weights
-
-The release repo keeps the unmodified MiniMax-H3 components (a copy of `MiniMaxAI/MiniMax-H3` at `6818f6c`) under
-`base_model/` next to `lora/` and `latent_upsampler/`. Point `model_path` at the Hub id or at a local copy of that layout.
-`adaln_t_table.bin` is base-derived and is not redistributed in the release repo: put the 32,800-byte file (SHA-256
-`c02a6c11888297688c1e6278185ea1f947023acfc69f9003bbcdcec9a229a8e7`) at `lora/adaln_t_table.bin`. The reference
-`download.py` fetches it by byte range from `Comfy-Org/MiniMax-H3` at revision
-`e5eb578a89295337b8ff433a035929ce0279e0b6` (tensor `adaln_t_table` of
-`diffusion_models/minimax_h3_ref2va_pruned_int8_convrot.safetensors`) and verifies the hash.
-
-A GB200/B200-class GPU is needed (about 150 GB with the DiT, the 2K activations and the VAE).
-
-## Usage
-
-```bash
-python examples/inference/basic/basic_open_h3_regenerate_2k.py \
-    --model-path Efficient-Large-Model/Open-H3-Regenerate-2K \
-    --video draft_768p.mp4 --prompt "A red fox trots through fresh snow" --output outputs/regen2k
-```
-
-The prompt conditions the refine steps (use the draft's own prompt). Drafts from other video models are better encoded at
-576P (`--draft-size 576p`). `--steps 2` runs the two-step variant. The output is always 2560x1440; the first 121 frames of
-the draft are used.
-
-## Serving
-
-```bash
-fastvideo serve --config examples/serving/openai_open_h3_regenerate_2k.yaml
-```
-
-`examples/serving/openai_open_h3_regenerate_2k.yaml` loads the release on one GPU, sets `media_roots`, and pins the 2K
-defaults. The draft arrives as the typed `video_reference` of `POST /v1/videos/sync` (the response body is the MP4) or of
-`POST /v1/videos` (an async job; poll `GET /v1/videos/{id}` and download `GET /v1/videos/{id}/content`):
-
-```bash
-curl -sS -X POST http://localhost:8000/v1/videos/sync \
-  -H 'Content-Type: application/json' \
-  -d '{
-        "prompt": "A red fox trots through fresh snow",
-        "video_reference": {"video_url": "https://example.com/draft_768p.mp4"},
-        "seed": 1766580932
-      }' \
-  -o regen_2k.mp4 -D -
-```
-
-The response headers carry `X-Request-Id`, `X-Model`, `X-Inference-Time-S` and `X-Stage-Durations`. Rules:
-
-- Exactly one `video_reference` (or the legacy `video_path` / `video_url`) and no image or audio references.
-- `http(s)` URLs must resolve to public addresses unless `server.allow_private_media_urls` is set. A local path is accepted
-  only inside a directory listed in `server.media_roots`. `data:` URLs are refused for video references; this policy is
-  unchanged.
-- `size`, `width` and `height` may be omitted or `2560x1440`; `num_frames` may be omitted or 124 (the frame count of the
-  VAE grid; the video has 121 frames); `fps` must be 24. `seconds` is not accepted.
-- `num_inference_steps` is 3 (default) or 2; `seed` sets the refine noise.
-
-## NVFP4 video decoder
-
-On Blackwell GPUs the released 36-block video decoder can be replaced by a pre-packed NVFP4 one; the latent space is
-unchanged, so the draft encode, the upscale and the refine steps are the same and only the decode changes. Both tiers of
-[FastVideo/FastH3-NVFP4-VAE-Decoder](https://huggingface.co/FastVideo/FastH3-NVFP4-VAE-Decoder) apply here, because the
-release carries MiniMax-H3's full 36-block VAE (`base_model/vae`):
-
-| Tier | Decoder blocks | File | Pick it when |
-|---|---|---|---|
-| `quality-36L` | 36 | `quality-36L/decoder-nvfp4.safetensors` | you want the released detail back; the default choice |
-| `fast-8L` | 8 (evenly spaced; the load drops 28 blocks) | `fast-8L/decoder-nvfp4.safetensors` | decode time or memory dominate and some fine detail can go |
-
-```bash
-hf download FastVideo/FastH3-NVFP4-VAE-Decoder --include 'quality-36L/decoder-nvfp4.safetensors' \
-    'fast-8L/decoder-nvfp4.safetensors' --local-dir fasth3-nvfp4-vae
-```
-
-Select it through the surface you drive the pipeline with; the field, the CLI flag and the environment variable all
-set `fastvideo_args.h3_vae_nvfp4_decoder`, and the serve-config field wins over the environment.
-
-```bash
-# example script
-python examples/inference/basic/basic_open_h3_regenerate_2k.py --video draft.mp4 --prompt "..." \
-    --vae-decoder fasth3-nvfp4-vae/quality-36L/decoder-nvfp4.safetensors
-# `fastvideo generate` and the serve config use the same field
-fastvideo generate --model-path Efficient-Large-Model/Open-H3-Regenerate-2K --video draft.mp4 --prompt "..." \
-    --h3-vae-nvfp4-decoder fasth3-nvfp4-vae/quality-36L/decoder-nvfp4.safetensors
-```
-
-```yaml
-generator:
-  pipeline:
-    experimental:
-      h3_vae_nvfp4_decoder: ./fasth3-nvfp4-vae/quality-36L/decoder-nvfp4.safetensors
-```
-
-`FASTVIDEO_H3_VAE_NVFP4_DECODER` stays as the fallback when neither the serve config nor the CLI sets the field. The
-swap happens in the VAE loader, so the dense decoder never reaches the GPU: it is replaced on the CPU by the exported
-checkpoint (2.0 GiB of packed weights instead of 9.7 GiB of FP32), moved to the device, and decoded under bf16 autocast.
-The decoder keeps the pipeline's tiling (at 2560x1440: 8x13 tiles of 256 px with 64 px overlap) and the fused Triton
-block kernels, which are on by default (`FASTVIDEO_H3_VAE_NVFP4_FUSED`); `FASTVIDEO_H3_VAE_TILE_BATCH` decodes several
-tiles per decoder call (4 is bit-identical and faster here), and `FASTVIDEO_H3_VAE_NVFP4_CUDA_GRAPH` replays the small
-fused calls. The MP4 is unaffected: the same 121 frames are trimmed and written.
-
-Warm decode of one refined latent at 2560x1440 (124 frames) on one GB200, median of 3, with the released decode as the
-reference:
-
-| Decoder | Decode (s) | Peak GPU memory (GiB) | PSNR (dB) | SSIM |
-|---|---:|---:|---:|---:|
-| released H3 VAE (dense, fp16 autocast) | 23.3 | 18.4 | - | - |
-| `quality-36L` (fused) | 6.6 | 6.2 | 41.3 | 0.979 |
-| `quality-36L` (fused, tile batch 4) | 5.0 | 6.3 | 41.3 | 0.979 |
-| `quality-36L` (eager NVFP4 GEMM) | 67.0 | 6.2 | 41.3 | 0.979 |
-| `fast-8L` (fused) | 1.7 | 5.3 | 32.8 | 0.892 |
-| `fast-8L` (eager NVFP4 GEMM) | 16.1 | 5.2 | 32.8 | 0.892 |
-
-The fused kernels are an exact speedup: the eager and fused rows decode bit-identical frames, and so does batching
-tiles, so the PSNR/SSIM columns measure the tier and nothing else. The eager NVFP4 GEMM path loses to the dense fp16
-decoder at these tile sizes (67 s against 23 s), so keep `FASTVIDEO_H3_VAE_NVFP4_FUSED` on. PSNR against the released
-decode is bounded by the released decode itself: the refined latent already sits 45.2 dB from the NVlabs reference (see
-below), so 41.3 dB is a decoder-only difference of the same order. `fast-8L` softens fine texture (fur, foliage) while
-keeping structure; use it only when the decode budget is tight.
-
-## Parity with the reference
-
-Measured on one GB200 (tray hpc-rack-1-5, BF16, dense attention) against the reference `standard_fwd` with the released
-VAE, same draft, seed (1766580932), prompt and VAE posterior seed (42). The reference's noise is a CUDA generator draw
-that FastVideo reproduces (the same `torch.Generator(device="cuda")` seed), so no injected noise was needed.
-
-| Stage / quantity | Fox draft (FastVideo H3 t2v, 768P) | Dance draft (576P, project page) |
-|---|---|---|
-| Draft latent, upscaled latent, DiT input | bit-identical | bit-identical |
-| One DiT forward (LoRA + Sharp), relative L2 | 1.1e-2 | 2.6e-2 |
-| Refined latent, relative L2 | 1.3e-2 | 3.7e-2 |
-| Decoded video PSNR / SSIM (21 frames) | 45.2 dB / 0.993 | 42.1 dB / 0.992 |
-
-The differences come from text-encoder rounding (relative L2 1e-3 on the prompt embeddings) and BF16 kernel choices;
-the reference reports the same order (1.4% relative L2 on the refined latent between its own runs).
-
-Warm timings, one GB200 (seconds):
-
-| | Upscale | Refine (3 steps) | VAE decode | Total |
-|---|---:|---:|---:|---:|
-| Reference `standard_fwd` | 0.87 | 71.6 | 15.0 | 87.5 |
-| FastVideo (end to end, incl. draft encode 5-10 s and decode to host) | 5.9-10.7 (encode + upscale + silence) | 69.5 | 25.0 | 107.5 |
-
-The first request builds nothing extra but loads weights (about 70 s once). The decode stage is slower than the
-reference because it also copies 121 frames to host memory and the MP4 encode is outside both columns.
-
-## Licences
-
-The release is not covered by one licence; its `README.md` says so, and the runtime code's licence does not replace the
-model licences. FastVideo only loads these weights and does not redistribute them.
-
-- The refiner LoRA, the fast H3 VAE decoder and `base_model/` (a copy of `MiniMaxAI/MiniMax-H3` at `6818f6c`: DiT, VAEs,
-  text encoder) follow the
-  [MiniMax H3 Community License](https://huggingface.co/MiniMaxAI/MiniMax-H3/blob/6818f6c32d12b210915e44ad56a4228c2608f160/LICENSE),
-  including its territory, use and redistribution terms (`base_model/LICENSE`, `NOTICE`). The Qwen3-VL text encoder's
-  Apache-2.0 text ships as `LICENSE-Qwen.txt`.
-- The latent upscaler follows its own licence:
-  [LBH-123-AI/Comfyui_Minimax_h3_latent_Upscaler](https://github.com/LBH-123-AI/Comfyui_Minimax_h3_latent_Upscaler)
-  (MIT, per the reference code) and its [weights](https://huggingface.co/LBH-123-AI/Minimax_h3_latent_Upscaler)
-  (`latent_upsampler/LICENSE`, Apache-2.0). `minimax_h3_latent_upscaler.py` is a port of that node's network.
-- The merged LoRA includes [LightX2V FL2VA Turbo](https://huggingface.co/lightx2v/Minimax-h3-Turbo) and
-  [LMS Sharp](https://huggingface.co/Alissonerdx/Minimax-H3-ComfyUI), which keep their own licences; the Sharp
-  `adaln_t_table` is read from [Comfy-Org/MiniMax-H3](https://huggingface.co/Comfy-Org/MiniMax-H3).

docs/inference/support_matrix.md

@@ -68,9 +68,7 @@ column links a runnable script in `examples/inference/basic/` where one exists.
 | matrixgame | `FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers`<br />`FastVideo/Matrix-Game-2.0-GTA-Distilled-Diffusers`<br />`FastVideo/Matrix-Game-2.0-TempleRun-Distilled-Diffusers`<br />`FastVideo/Matrix-Game-2.0-Base-Diffusers`<br />`FastVideo/Matrix-Game-2.0-GTA-Diffusers`<br />`FastVideo/Matrix-Game-2.0-TempleRun-Diffusers`<br />`mignonjia/mg_longtuning_distilled_zelda`<br />`mignonjia/mg_sf_distilled_zelda_1k_steps`<br />`mignonjia/mg_sf_distilled_zelda`<br />`mignonjia/mg_causal_zelda`<br />`mignonjia/mg_bidirectional_zelda` | I2V | [basic_matrixgame2.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_matrixgame2.py) |
 | matrixgame | `FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers` | I2V | [basic_matrixgame3.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_matrixgame3.py) |
 | minimax_h3 | `MiniMaxAI/MiniMax-H3` | T2V, I2V | [T2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_t2v.py)<br />[FL2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_fl2va.py)<br />[Ref2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_ref2va.py) |
-| open_h3_regenerate_2k | `Efficient-Large-Model/Open-H3-Regenerate-2K` | — | [basic_open_h3_regenerate_2k.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_open_h3_regenerate_2k.py) |
 | sd35 | `stabilityai/stable-diffusion-3.5-medium` | T2I | [basic_sd35_t2i.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_sd35_t2i.py) |
-| qwen_image21 | `Qwen/Qwen-Image-2.1` | T2I, I2I | [basic_qwen_image21.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_qwen_image21.py) |
 | stable_audio | `FastVideo/stable-audio-open-1.0-Diffusers` | T2V | [basic_stable_audio.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_stable_audio.py) |
 | stable_audio | `FastVideo/stable-audio-open-small-Diffusers` | T2V | [basic_stable_audio_small.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_stable_audio_small.py) |
 | wan | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | T2V | [basic.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic.py) |
@@ -122,9 +120,6 @@ and distilled) generate video with audio from text, optionally plus an image (se
 `kandinsky6_sr` IDs upscale an existing video (see the
 [Video SR guide](kandinsky6_sr.md)).
 
-**Note (Open-H3-Regenerate-2K)**: `open_h3_regenerate_2k` regenerates an existing 768P MiniMax-H3 draft video at
-2560x1440 (video to video, silent output); see the [Open-H3-Regenerate-2K guide](open_h3_regenerate_2k.md).
-
 **Note (LTX-2.5)**: T2V and I2V produce synchronized video and audio. The
 official `Lightricks/LTX-2.5` repository uses gated, split component files, so
 it must first be converted to FastVideo's component layout. See the

examples/inference/basic/basic_fasth3_spark_pair_qad_v2_nvfp4.yaml

@@ -1,5 +1,6 @@
 # Start Ray on both Sparks and source spark_pair_env.sh first; see spark_pair.md.
 # FastH3-QAD-v2 eight-forward video+audio on two DGX Sparks over QSFP RoCE.
+# Uses the published fast-8L NVFP4 decoder on both workers (FASTVIDEO_H3_VAE_NVFP4_DECODER) with parallel temporal chunks.
 # Download the complete release checkpoint as described in
 # docs/getting_started/installation/spark_performance.md.
 #
@@ -9,8 +10,8 @@
 # settings below are an inference choice, not its linears-only QAD recipe.
 # The NVFP4 encoder and lightweight VAE ship in the same repository.
 #
-# FASTVIDEO_H3_VSA_FP4=1 FASTVIDEO_MINIMAX_H3_FUSIONS=all FASTVIDEO_NVFP4_MM_BACKEND=cutlass \
-# FASTVIDEO_H3_VAE_TILE_BATCH=1 \
+# FASTVIDEO_H3_VSA_FP4=1 FASTVIDEO_MINIMAX_H3_FUSIONS=all FASTVIDEO_NVFP4_MM_BACKEND=cudnn FASTVIDEO_NVFP4_MM_ROW_SPLIT=table \
+# FASTVIDEO_H3_VAE_TILE_BATCH=12 FASTVIDEO_H3_VAE_INT8_OVERLAY=0 \
 # FASTVIDEO_VSA_TRITON=1 FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0 \
 # FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 FASTVIDEO_STAGE_LOGGING=1 \
 # fastvideo generate --config examples/inference/basic/basic_fasth3_spark_pair_qad_v2_nvfp4.yaml
@@ -36,7 +37,7 @@ generator:
       lazy_module_load: false
     compile:
       enabled: false
-      vae_enabled: false
+      vae_enabled: false  # Published NVFP4 decoder uses its fused eager inference path.
   pipeline:
     workload_type: t2v
     vae_tiling: true
@@ -47,6 +48,11 @@ generator:
         prefix_precision: fp8
         two_level_p: true
         smooth_k: true
+        sp_head_chunks: 7
+        sp_pooled_key_mean: true
+        sp_direct_keys: true
+        sp_projection_overlap: true
+        sp_input_gather: true
       VSA_sparsity: 0.8
       VSA_tile_size: 64
       h3_sequential_load: false

examples/inference/basic/basic_fasth3_spark_qad_v2_nvfp4.yaml

@@ -1,4 +1,5 @@
 # FastH3-QAD-v2 eight-forward video+audio on one DGX Spark.
+# Uses the published fast-8L NVFP4 decoder selected with FASTVIDEO_H3_VAE_NVFP4_DECODER (see spark_performance.md).
 # Download the complete release checkpoint as described in
 # docs/getting_started/installation/spark_performance.md.
 #
@@ -34,7 +35,7 @@ generator:
       lazy_module_load: false
     compile:
       enabled: false
-      vae_enabled: true
+      vae_enabled: false  # Packed NVFP4 decoder has its own fused inference path.
   pipeline:
     workload_type: t2v
     vae_tiling: true

examples/inference/basic/basic_open_h3_regenerate_2k.py

@@ -1,93 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Regenerate a 768P MiniMax-H3 draft video at 2560x1440 with Open-H3-Regenerate-2K.
-
-The draft is VAE-encoded (the first 121 frames at 24 fps, 1344x768), its latent is upscaled, and three refine steps
-of the H3 transformer with the released LoRA run at 2560x1440. The output video is silent. See
-docs/inference/open_h3_regenerate_2k.md.
-
-    python examples/inference/basic/basic_open_h3_regenerate_2k.py --video draft.mp4 --prompt "..." --output out/
-"""
-
-from __future__ import annotations
-
-import argparse
-import time
-from pathlib import Path
-
-from fastvideo import VideoGenerator
-from fastvideo.api import (
-    ComponentConfig,
-    EngineConfig,
-    GenerationRequest,
-    GeneratorConfig,
-    InputConfig,
-    OffloadConfig,
-    OutputConfig,
-    PipelineSelection,
-    SamplingConfig,
-)
-
-DEFAULT_MODEL = "Efficient-Large-Model/Open-H3-Regenerate-2K"
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
-    parser.add_argument("--model-path", default=DEFAULT_MODEL, help="release repo id or a local copy of it")
-    parser.add_argument("--video", required=True, help="768P draft video (URL-free local path)")
-    parser.add_argument("--prompt", required=True, help="text the refine steps are conditioned on")
-    parser.add_argument("--output", type=Path, default=Path("outputs/open_h3_regenerate_2k"))
-    parser.add_argument("--steps", type=int, choices=(2, 3), default=3, help="refine steps (3 is the released setting)")
-    parser.add_argument("--seed", type=int, default=1766580932)
-    parser.add_argument("--attention-backend", default="FLASH_ATTN")
-    parser.add_argument("--draft-size", choices=("768p", "576p"), default="768p",
-                        help="size the draft is encoded at; 576p for drafts from other video models")
-    parser.add_argument("--vae-decoder", default=None,
-                        help="NVFP4 decoder-nvfp4.safetensors that replaces the VAE decoder at load "
-                        "(Blackwell only, see docs/inference/open_h3_regenerate_2k.md)")
-    parser.add_argument("--runs", type=int, default=1, help="repeat the request to measure warm timings")
-    return parser.parse_args()
-
-
-def main() -> None:
-    args = parse_args()
-    args.output.mkdir(parents=True, exist_ok=True)
-    experimental = {
-        "attention_backend": args.attention_backend,
-        "draft_size": args.draft_size,
-    }
-    if args.vae_decoder:
-        experimental["h3_vae_nvfp4_decoder"] = args.vae_decoder
-    generator = VideoGenerator.from_config(
-        GeneratorConfig(
-            model_path=args.model_path,
-            engine=EngineConfig(num_gpus=1, offload=OffloadConfig(dit=False, dit_layerwise=False, text_encoder=True, vae=False)),
-            pipeline=PipelineSelection(
-                components=ComponentConfig(),
-                experimental=experimental,
-            ),
-        ))
-    try:
-        for run in range(args.runs):
-            started = time.perf_counter()
-            result = generator.generate(
-                GenerationRequest(
-                    prompt=args.prompt,
-                    inputs=InputConfig(video_path=args.video),
-                    sampling=SamplingConfig(height=1440,
-                                            width=2560,
-                                            num_frames=124,
-                                            fps=24,
-                                            num_inference_steps=args.steps,
-                                            guidance_scale=1.0,
-                                            seed=args.seed),
-                    output=OutputConfig(output_path=str(args.output / f"regen2k_{run}.mp4"),
-                                        save_video=True,
-                                        return_frames=False),
-                ))
-            print(f"run {run}: {time.perf_counter() - started:.1f}s -> {result.video_path}")
-    finally:
-        generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()

examples/inference/basic/basic_qwen_image21.py

@@ -1,73 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Qwen-Image-2.1 generation, editing, references and transparent PNG output."""
-
-import argparse
-
-from fastvideo import VideoGenerator
-from fastvideo.api import (
-    EngineConfig,
-    GenerationRequest,
-    GeneratorConfig,
-    InputConfig,
-    OffloadConfig,
-    OutputConfig,
-    ParallelismConfig,
-    PipelineSelection,
-    SamplingConfig,
-)
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--model-path", default="Qwen/Qwen-Image-2.1")
-    parser.add_argument("--revision")
-    parser.add_argument("--prompt", default="A red ceramic teapot, isolated on a transparent background.")
-    parser.add_argument("--reference", action="append", default=[], help="Ordered reference image; repeat up to ten times")
-    parser.add_argument("--negative-prompt", default=None)
-    parser.add_argument("--true-cfg-scale", type=float, default=1.0)
-    parser.add_argument("--height", type=int, default=1024)
-    parser.add_argument("--width", type=int, default=1024)
-    parser.add_argument("--reference-resolution", type=int, default=1024)
-    parser.add_argument("--steps", type=int, default=40)
-    parser.add_argument("--seed", type=int, default=42)
-    parser.add_argument("--use-kv-cache", action=argparse.BooleanOptionalAction, default=True)
-    parser.add_argument("--output", default="outputs/qwen_image21/image.png")
-    args = parser.parse_args()
-    if len(args.reference) > 10:
-        parser.error("At most ten ordered reference images are supported")
-    workload = "i2i" if args.reference else "t2i"
-    generator = VideoGenerator.from_config(
-        GeneratorConfig(
-            model_path=args.model_path,
-            revision=args.revision,
-            engine=EngineConfig(
-                num_gpus=1,
-                parallelism=ParallelismConfig(tp_size=1, sp_size=1),
-                offload=OffloadConfig(dit=True, dit_layerwise=True, text_encoder=True, vae=True,
-                                      lazy_module_load=True),
-            ),
-            pipeline=PipelineSelection(
-                workload_type=workload,
-                experimental={
-                    "reference_resolution": args.reference_resolution,
-                    "use_kv_cache": args.use_kv_cache,
-                },
-            ),
-        ))
-    try:
-        generator.generate(
-            GenerationRequest(
-                prompt=args.prompt,
-                negative_prompt=args.negative_prompt,
-                inputs=InputConfig(references=args.reference or None),
-                sampling=SamplingConfig(height=args.height, width=args.width, num_frames=1, fps=1,
-                                        num_inference_steps=args.steps, guidance_scale=1.0,
-                                        true_cfg_scale=args.true_cfg_scale, seed=args.seed),
-                output=OutputConfig(output_path=args.output, save_video=True, return_frames=False),
-            ))
-    finally:
-        generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()

examples/inference/basic/mlx_fasth3_omniref.py

@@ -114,10 +114,9 @@ def build_parser() -> argparse.ArgumentParser:
     parser.add_argument("--vsa-ref-keep-rate", type=float, default=None,
                         help="fraction of each reference video's tiles a query keeps (p2_multi_region); "
                         "default is the contract's trained value")
-    parser.add_argument("--vsa-impl", choices=("auto", "reference", "simd", "sdpa", "nax"), default="auto",
+    parser.add_argument("--vsa-impl", choices=("auto", "reference", "simd"), default="auto",
                         help="sparse attention backend; auto uses the SIMD Metal kernel for 128-token tiles and "
-                        "gather+SDPA otherwise, reference forces gather+SDPA; sdpa shares K/V with a tile mask "
-                        "and nax uses M5 TensorOps over selected blocks (both opt-in, incompatible with --reuse-prefix-kv)")
+                        "gather+SDPA otherwise, reference forces gather+SDPA")
     parser.add_argument("--reuse-prefix-kv", action="store_true",
                         help="opt-in, lossy: record the text/reference rows' per-block K/V on step 1 and run later "
                         "steps on the target audio/video rows only (faster, slightly different output)")
@@ -146,8 +145,6 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
         parser.error("Ref2VA needs at least one image or video reference")
     if args.model_root is None and (args.export is None or args.base_model_path is None):
         parser.error("give --model-root, or --export together with --base-model-path")
-    if args.vsa_impl in ("sdpa", "nax") and args.reuse_prefix_kv:
-        parser.error(f"--vsa-impl {args.vsa_impl} does not support --reuse-prefix-kv")
     return args
 
 

examples/inference/basic/qwen_image21_edit.yaml

@@ -1,39 +0,0 @@
-# Replace source.png and mask.png with your own images. The mask is an ordered
-# reference interpreted by the model, not a hard pixel-preservation constraint.
-generator:
-  model_path: Qwen/Qwen-Image-2.1
-  engine:
-    num_gpus: 1
-    parallelism:
-      tp_size: 1
-      sp_size: 1
-    offload:
-      dit: true
-      dit_layerwise: true
-      text_encoder: true
-      vae: true
-      lazy_module_load: true
-  pipeline:
-    workload_type: i2i
-    experimental:
-      reference_resolution: 1024
-      use_kv_cache: true
-request:
-  prompt: In image 1, turn the object highlighted by the mask in image 2 blue.
-  inputs:
-    references:
-      - source.png
-      - mask.png
-  sampling:
-    height: 1024
-    width: 1024
-    num_frames: 1
-    fps: 1
-    num_inference_steps: 40
-    guidance_scale: 1.0
-    true_cfg_scale: 1.0
-    seed: 42
-  output:
-    output_path: outputs/qwen_image21/edit.png
-    save_video: true
-    return_frames: false

examples/inference/basic/qwen_image21_t2i.yaml

@@ -1,40 +0,0 @@
-# Run on an allocated CUDA GPU after installing FastVideo and providing weights.
-# fastvideo generate --config examples/inference/basic/qwen_image21_t2i.yaml
-generator:
-  model_path: Qwen/Qwen-Image-2.1
-  engine:
-    num_gpus: 1
-    compile:
-      enabled: false
-      backend: inductor
-      fullgraph: false
-    parallelism:
-      tp_size: 1
-      sp_size: 1
-    offload:
-      dit: true
-      dit_layerwise: true
-      text_encoder: true
-      vae: true
-      lazy_module_load: true
-  pipeline:
-    workload_type: t2i
-    experimental:
-      kv_cache_device: cuda
-      reference_resolution: 1024
-      use_kv_cache: true
-request:
-  prompt: A red ceramic teapot, isolated on a transparent background.
-  sampling:
-    height: 1024
-    width: 1024
-    num_frames: 1
-    fps: 1
-    num_inference_steps: 40
-    guidance_scale: 1.0
-    true_cfg_scale: 1.0
-    seed: 42
-  output:
-    output_path: outputs/qwen_image21/teapot.png
-    save_video: true
-    return_frames: false

examples/inference/optimizations/spark_pair_env.sh

@@ -19,3 +19,7 @@ export NCCL_NVLS_ENABLE="${NCCL_NVLS_ENABLE:-0}"
 # runaway and SIGTERMs the worker around shard 11/14.
 export RAY_memory_monitor_refresh_ms="${RAY_memory_monitor_refresh_ms:-0}"
 export RAY_memory_usage_threshold="${RAY_memory_usage_threshold:-1.0}"
+# Measured on a two-Spark QSFP pair: eight NCCL channels per peer and high-priority
+# NCCL streams let the sequence-parallel exchanges overlap the projection GEMMs.
+export NCCL_NCHANNELS_PER_NET_PEER="${NCCL_NCHANNELS_PER_NET_PEER:-8}"
+export TORCH_NCCL_HIGH_PRIORITY="${TORCH_NCCL_HIGH_PRIORITY:-1}"

examples/serving/openai_open_h3_regenerate_2k.yaml

@@ -1,46 +0,0 @@
-# OpenAI-compatible Open-H3-Regenerate-2K server (one GPU with about 150 GB free, e.g. B200 / GB200).
-# POST /v1/videos/sync (or /v1/videos) with a 768P H3 draft as video_reference; the response is the 2560x1440 MP4.
-# Local draft paths are accepted only inside server.media_roots; http(s) URLs must resolve to public addresses unless
-# server.allow_private_media_urls is set. data: URLs are always refused. See docs/inference/open_h3_regenerate_2k.md.
-generator:
-  model_path: Efficient-Large-Model/Open-H3-Regenerate-2K
-  engine:
-    num_gpus: 1
-    offload:
-      dit: false
-      dit_layerwise: false
-      text_encoder: true
-      vae: false
-      pin_cpu_memory: true
-    compile:
-      enabled: false
-      vae_enabled: false
-  pipeline:
-    experimental:
-      attention_backend: FLASH_ATTN
-      # 768p for MiniMax-H3 drafts; 576p encodes drafts from other video models at 1024x576.
-      draft_size: 768p
-      # NVFP4 video decoder (Blackwell only). Download one tier from FastVideo/FastH3-NVFP4-VAE-Decoder and point
-      # this at it; the release VAE has the full 36 blocks, so both tiers apply. `fast-8L` is the cheap option.
-      # h3_vae_nvfp4_decoder: ./fasth3-nvfp4-vae/quality-36L/decoder-nvfp4.safetensors
-
-server:
-  host: 0.0.0.0
-  port: 8000
-  output_dir: outputs/openai_open_h3_regenerate_2k
-  served_model_name: open-h3-regenerate-2k
-  media_roots:
-    - /data/drafts
-  media_reference_max_bytes: 536870912
-
-default_request:
-  sampling:
-    height: 1440
-    width: 2560
-    num_frames: 124
-    fps: 24
-    num_inference_steps: 3
-    guidance_scale: 1.0
-    seed: 1766580932
-  output:
-    return_frames: false

fastvideo-kernel/attn_qat_infer/blackwell/launch.h

@@ -89,7 +89,11 @@ void run_flash_fwd(Flash_fwd_params &params, cudaStream_t stream) {
     static constexpr int ctaSize = Kernel_traits::kNWarps * 32;
     params.m_block_divmod = cutlass::FastDivmod(num_blocks_m);
     params.total_blocks = num_blocks_m * params.h * params.b;
-    dim3 grid_dims = Scheduler::get_grid_dim(scheduler_args, 170);
+    // One persistent worker per SM. A fixed 170-worker grid leaves a partial
+    // final wave on smaller GPUs such as GB10, while every worker still walks
+    // its full static share of query tiles.
+    const int num_sms = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
+    dim3 grid_dims = Scheduler::get_grid_dim(scheduler_args, num_sms);
     dim3 block_dims(ctaSize);
     dim3 cluster_dims(size<0>(ClusterShape{}), size<1>(ClusterShape{}), size<2>(ClusterShape{}));
     cutlass::ClusterLaunchParams launch_params{grid_dims, block_dims, cluster_dims, smem_size, stream};

fastvideo-kernel/csrc/attention/minimax_h3_prefix.cu

@@ -0,0 +1,84 @@
+// Exact Omni prefix attention: Torch-pinned FlashAttention, larger query tile only.
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <cmath>
+#include "flash_fwd_launch_template.h"
+
+using namespace FLASH_NAMESPACE;
+
+template<int M, int W, bool QRegs = false, bool Share = false>
+struct Variant {
+    using Traits = Flash_fwd_kernel_traits<128, M, 64, W, QRegs, Share, cutlass::bfloat16_t>;
+    static auto kernel() {
+        // Keep the original K tile and reduction order. All masking stays enabled.
+        return flash_fwd_kernel<Traits, false, false, false, false, false, true, false, false>;
+    }
+    static void setup() {
+        static const bool initialized = []() {
+            if (Traits::kSmemSize >= 48 * 1024) {
+                C10_CUDA_CHECK(cudaFuncSetAttribute(kernel(), cudaFuncAttributeMaxDynamicSharedMemorySize,
+                                                   Traits::kSmemSize));
+            }
+            return true;
+        }();
+        (void) initialized;
+    }
+    static void launch(Flash_fwd_params &params, cudaStream_t stream) {
+        setup();
+        dim3 grid((params.seqlen_q + M - 1) / M, params.b, params.h);
+        auto function = kernel();
+        function<<<grid, W * 32, Traits::kSmemSize, stream>>>(params);
+        C10_CUDA_KERNEL_LAUNCH_CHECK();
+    }
+    static std::vector<int64_t> metadata() {
+        setup();
+        cudaFuncAttributes attributes;
+        C10_CUDA_CHECK(cudaFuncGetAttributes(&attributes, kernel()));
+        int blocks = 0;
+        C10_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+            &blocks, kernel(), W * 32, Traits::kSmemSize));
+        return {M, W, Traits::kSmemSize, attributes.numRegs,
+                static_cast<int64_t>(attributes.localSizeBytes), blocks};
+    }
+};
+
+torch::Tensor prefix_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v) {
+    TORCH_CHECK(q.is_cuda() && k.device() == q.device() && v.device() == q.device(), "CUDA operands required");
+    TORCH_CHECK(q.dim() == 4 && k.dim() == 4 && v.sizes() == k.sizes(), "BHSD operands required");
+    TORCH_CHECK(q.scalar_type() == at::kBFloat16 && k.scalar_type() == at::kBFloat16 &&
+                v.scalar_type() == at::kBFloat16, "BF16 operands required");
+    TORCH_CHECK(q.size(3) == 128 && k.size(3) == 128 && q.stride(3) == 1 &&
+                k.stride(3) == 1 && v.stride(3) == 1, "Contiguous 128-wide heads required");
+    TORCH_CHECK(q.size(0) == k.size(0) && q.size(1) == k.size(1), "Batch/head mismatch");
+    TORCH_CHECK(q.size(0) > 0 && q.size(1) > 0 && q.size(2) > 0 && k.size(2) > 0,
+                "Nonempty attention required");
+    c10::cuda::CUDAGuard guard(q.device());
+    auto out = torch::empty({q.size(0), q.size(2), q.size(1), 128}, q.options()).transpose(1, 2);
+    auto lse = torch::empty({q.size(0), q.size(1), q.size(2)}, q.options().dtype(at::kFloat));
+    Flash_fwd_params params{};
+    params.q_ptr = q.data_ptr(); params.k_ptr = k.data_ptr(); params.v_ptr = v.data_ptr();
+    params.o_ptr = out.data_ptr(); params.softmax_lse_ptr = lse.data_ptr();
+    params.q_batch_stride = q.stride(0); params.q_head_stride = q.stride(1); params.q_row_stride = q.stride(2);
+    params.k_batch_stride = k.stride(0); params.k_head_stride = k.stride(1); params.k_row_stride = k.stride(2);
+    params.v_batch_stride = v.stride(0); params.v_head_stride = v.stride(1); params.v_row_stride = v.stride(2);
+    params.o_batch_stride = out.stride(0); params.o_head_stride = out.stride(1); params.o_row_stride = out.stride(2);
+    params.b = q.size(0); params.h = q.size(1); params.h_k = k.size(1); params.h_h_k_ratio = 1;
+    params.seqlen_q = q.size(2); params.seqlen_k = k.size(2); params.total_q = params.b * params.seqlen_q;
+    params.seqlen_q_rounded = (params.seqlen_q + 127) / 128 * 128;
+    params.seqlen_k_rounded = (params.seqlen_k + 127) / 128 * 128;
+    params.d = 128; params.d_rounded = 128;
+    params.scale_softmax = static_cast<float>(1.0 / std::sqrt(128.0));
+    params.scale_softmax_log2 = params.scale_softmax * M_LOG2E;
+    params.p_dropout = 1.0f; params.p_dropout_in_uint8_t = 255; params.rp_dropout = 1.0f;
+    params.scale_softmax_rp_dropout = params.scale_softmax;
+    params.window_size_left = -1; params.window_size_right = -1;
+    params.is_bf16 = true; params.is_seqlens_k_cumulative = true; params.num_splits = 1;
+    auto stream = at::cuda::getCurrentCUDAStream();
+    Variant<256, 8>::launch(params, stream);
+    return out;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, binding) {
+    binding.def("forward", &prefix_forward);
+}

fastvideo-kernel/csrc/quantization/minimax_h3_swiglu_gemm.cu

@@ -0,0 +1,203 @@
+// SPDX-License-Identifier: Apache-2.0
+// Experimental full-K SM120 GEMM with value/gate-paired columns and FP4 output.
+#include "minimax_h3_swiglu_quant.cuh"
+#include <cutlass/cutlass.h>
+#include <cutlass/epilogue/collective/collective_builder.hpp>
+#include <cutlass/gemm/collective/collective_builder.hpp>
+#include <cutlass/gemm/device/gemm_universal_adapter.h>
+#include <cutlass/gemm/kernel/tile_scheduler.hpp>
+#include <cutlass/util/packed_stride.hpp>
+
+namespace h3_epilogue {
+using namespace cute;
+using namespace cutlass::epilogue::fusion;
+struct Operation : LinearCombination<cutlass::bfloat16_t, float, void> {
+  using ElementAux = cutlass::bfloat16_t;
+  using GmemLayoutTagAux = cutlass::layout::RowMajor;
+  static constexpr int AlignmentAux = 8;
+};
+
+template<class CtaTile, class EpiTile>
+struct Store {
+  static constexpr int EM = size<0>(EpiTile{});
+  static constexpr int EN = size<1>(EpiTile{});
+  static_assert(EN % 32 == 0);
+  struct SharedStorage { cutlass::array_aligned<__nv_bfloat16, EM * EN / 2> values; };
+  struct Arguments {
+    const float* alpha_ptr = nullptr;
+    const float* scale_ptr = nullptr;
+    uint8_t* q = nullptr;
+    uint8_t* scales = nullptr;
+    float* amax = nullptr;
+    __nv_bfloat16* y = nullptr;
+  };
+  using Params = Arguments;
+  Params const* params;
+  __nv_bfloat16* shared;
+  CUTLASS_HOST_DEVICE Store() = default;
+  CUTLASS_HOST_DEVICE Store(Params const& p, SharedStorage const& s)
+      : params(&p), shared(const_cast<__nv_bfloat16*>(s.values.data())) {}
+  template<class Shape>
+  static Params to_underlying_arguments(Shape const&, Arguments const& a, void*) { return a; }
+  template<class Shape> static bool can_implement(Shape const& s, Arguments const&) {
+    return get<1>(s) % 128 == 0;
+  }
+  template<class Shape> static size_t get_workspace_size(Shape const&, Arguments const&) { return 0; }
+  template<class Shape>
+  static cutlass::Status initialize_workspace(Shape const&, Arguments const&, void*, cudaStream_t,
+                                               cutlass::CudaHostAdapter* = nullptr) {
+    return cutlass::Status::kSuccess;
+  }
+  CUTLASS_DEVICE bool is_producer_load_needed() const { return false; }
+  CUTLASS_DEVICE bool is_C_load_needed() const { return false; }
+  template<class... A> CUTLASS_DEVICE auto get_producer_load_callbacks(ProducerLoadArgs<A...> const&) {
+    return EmptyProducerLoadCallbacks{};
+  }
+
+  template<class Coords, int Threads>
+  struct Consumer : EmptyConsumerStoreCallbacks {
+    Coords coords;
+    Params p;
+    __nv_bfloat16* shared;
+    int tid, rows, width, row_base, col_base;
+    float alpha, global_sf;
+    CUTLASS_DEVICE Consumer(Coords c, Params p_, __nv_bfloat16* sm, int t, int m, int n, int rb, int cb)
+        : coords(c), p(p_), shared(sm), tid(t), rows(m), width(n/2), row_base(rb), col_base(cb),
+          alpha(*p_.alpha_ptr), global_sf(*p_.scale_ptr) {}
+    template<class Acc, int F>
+    CUTLASS_DEVICE auto visit(cutlass::Array<Acc, F> const& acc, int v, int em, int en) {
+      static_assert(F % 2 == 0);
+      auto c = coords(_,_,_,em,en);
+      CUTLASS_PRAGMA_UNROLL
+      for (int i=0; i<F; i+=2) {
+        auto xy = c(v*F+i);
+        const int row = get<0>(xy), col = get<1>(xy);
+        const float value = __bfloat162float(__float2bfloat16_rn(__fmul_rn(alpha, float(acc[i]))));
+        const float gate = __bfloat162float(__float2bfloat16_rn(__fmul_rn(alpha, float(acc[i+1]))));
+        const float y = row < rows ? swiglu_bf16(value,gate) : 0;
+        shared[(row % EM) * (EN/2) + (col % EN)/2] = __float2bfloat16_rn(y);
+      }
+      return acc;
+    }
+    template<class Buffer, class Sync, class Values>
+    CUTLASS_DEVICE void reduce(Buffer&&, Sync const& sync, int em, int en, bool, Values) {
+      sync();
+      CUTLASS_PRAGMA_UNROLL
+      for (int i=tid; i<EM*EN/2; i+=Threads) {
+        const int row = row_base + em*EM + i/(EN/2);
+        const int col = col_base/2 + en*(EN/2) + i%(EN/2);
+        const float y = __bfloat162float(shared[i]);
+        float amax = fabsf(y);
+        CUTLASS_PRAGMA_UNROLL
+        for (int d=1; d<16; d*=2)
+          amax=fmaxf(amax,__shfl_xor_sync(0xffffffff,amax,d,16));
+        const uint8_t sf=scale_fp8(__fmul_rn(global_sf,__fmul_rn(amax,rcp_approx(6.0f))));
+        const float inv=amax!=0 ? rcp_approx(__fmul_rn(scale_float(sf),rcp_approx(global_sf))) : 0;
+        const float code=__fmul_rn(y,inv);
+        const float hi=__shfl_xor_sync(0xffffffff,code,1);
+        if (row<rows && col<width) {
+          if ((col&1)==0) p.q[int64_t(row)*(width/2)+col/2]=fp4_pair(code,hi);
+          if (p.y) p.y[int64_t(row)*width+col]=__float2bfloat16_rn(y);
+        }
+        if (row < (rows + 127) / 128 * 128 && col<width && (col&15)==0) {
+          const int group=col/16;
+          const int64_t idx=(((((int64_t(row)/128)*(width/64)+group/4)*32+row%32)*4+(row/32)%4)*4+group%4);
+          p.scales[idx]=sf;
+          if (p.amax) p.amax[int64_t(row)*(width/16)+group]=amax;
+        }
+      }
+      sync();
+    }
+  };
+  template<bool ReferenceSrc, class... A>
+  CUTLASS_DEVICE auto get_consumer_store_callbacks(ConsumerStoreArgs<A...> const& args) {
+    auto [M,N,K,L]=args.problem_shape_mnkl;
+    auto [m,n,k,l]=args.tile_coord_mnkl;
+    auto identity=make_identity_tensor(make_shape(M,N));
+    auto tile=local_tile(identity,take<0,2>(args.tile_shape_mnk),make_coord(m,n));
+    auto coords=sm90_partition_for_epilogue<ReferenceSrc>(tile,args.epi_tile,args.tiled_copy,args.thread_idx);
+    constexpr int threads=decltype(size(args.tiled_copy))::value;
+    return Consumer<decltype(coords),threads>(coords,*params,shared,args.thread_idx,M,N,
+                                              m*size<0>(CtaTile{}),n*size<1>(CtaTile{}));
+  }
+};
+} // namespace h3_epilogue
+
+namespace cutlass::epilogue::fusion {
+template<int SC,int SD,int F,bool R,bool D,class CT,class ET,class... Extra>
+struct FusionCallbacks<cutlass::epilogue::Sm90TmaWarpSpecialized<SC,SD,F,R,D>,
+                       h3_epilogue::Operation,CT,ET,Extra...> : h3_epilogue::Store<CT,ET> {
+  using Impl=h3_epilogue::Store<CT,ET>;
+  using Impl::Impl;
+  using Operation=h3_epilogue::Operation;
+};
+}
+
+namespace h3_epilogue {
+using Shape=cute::Shape<int,int,int,int>;
+#ifndef H3_TILE_M
+#define H3_TILE_M 128
+#define H3_TILE_N 128
+#define H3_TILE_K 256
+#endif
+#ifndef H3_PINGPONG
+#define H3_PINGPONG 0
+#endif
+#ifndef H3_STAGES
+#define H3_STAGES 0
+#endif
+using Tile=cute::Shape<cute::Int<H3_TILE_M>,cute::Int<H3_TILE_N>,cute::Int<H3_TILE_K>>;
+using Cluster=cute::Shape<cute::_1,cute::_1,cute::_1>;
+using Epilogue=typename cutlass::epilogue::collective::CollectiveBuilder<
+  cutlass::arch::Sm120,cutlass::arch::OpClassTensorOp,Tile,Cluster,
+  cutlass::epilogue::collective::EpilogueTileAuto,float,float,
+  void,cutlass::layout::RowMajor,8,void,cutlass::layout::RowMajor,8,
+  cutlass::epilogue::TmaWarpSpecialized,Operation>::CollectiveOp;
+using Mainloop=typename cutlass::gemm::collective::CollectiveBuilder<
+  cutlass::arch::Sm120,cutlass::arch::OpClassBlockScaledTensorOp,
+  cutlass::nv_float4_t<cutlass::float_e2m1_t>,cutlass::layout::RowMajor,32,
+  cutlass::nv_float4_t<cutlass::float_e2m1_t>,cutlass::layout::ColumnMajor,32,
+  float,Tile,Cluster,
+  cute::conditional_t<H3_STAGES == 0,
+    cutlass::gemm::collective::StageCountAutoCarveout<sizeof(typename Epilogue::SharedStorage)>,
+    cutlass::gemm::collective::StageCount<H3_STAGES>>, 
+  cute::conditional_t<H3_PINGPONG != 0, cutlass::gemm::KernelTmaWarpSpecializedPingpong,
+                      cutlass::gemm::KernelTmaWarpSpecializedCooperative>>::CollectiveOp;
+using Kernel=cutlass::gemm::kernel::GemmUniversal<Shape,Mainloop,Epilogue,cutlass::gemm::StaticPersistentScheduler>;
+using Gemm=cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
+}
+extern "C" int launch_gemm_swiglu(const void* a,const void* b,const void* sa,const void* sb,
+                                 const float* alpha,const float* sf,void* q,void* scales,
+                                 float* amax,void* y,int m,int n,int k,void* workspace,
+                                 uint64_t workspace_bytes,int raster,int swizzle,void* raw_stream) {
+  using namespace h3_epilogue;
+  auto stream=static_cast<cudaStream_t>(raw_stream);
+  Gemm::Arguments args{};
+  args.mode=cutlass::gemm::GemmUniversalMode::kGemm;
+  args.problem_shape=cute::make_shape(m,n,k,1);
+  args.mainloop.ptr_A=static_cast<const cutlass::float_e2m1_t*>(a);
+  args.mainloop.ptr_B=static_cast<const cutlass::float_e2m1_t*>(b);
+  args.mainloop.ptr_SFA=static_cast<const cutlass::float_ue4m3_t*>(sa);
+  args.mainloop.ptr_SFB=static_cast<const cutlass::float_ue4m3_t*>(sb);
+  args.mainloop.dA=cutlass::make_cute_packed_stride(Kernel::StrideA{}, {m,k,1});
+  args.mainloop.dB=cutlass::make_cute_packed_stride(Kernel::StrideB{}, {n,k,1});
+  using Config=Mainloop::Sm1xxBlkScaledConfig;
+  args.mainloop.layout_SFA=Config::tile_atom_to_shape_SFA(args.problem_shape);
+  args.mainloop.layout_SFB=Config::tile_atom_to_shape_SFB(args.problem_shape);
+  args.epilogue.dC=cutlass::make_cute_packed_stride(Kernel::StrideC{}, {m,n,1});
+  args.epilogue.dD=args.epilogue.dC;
+  args.epilogue.thread={alpha,sf,static_cast<uint8_t*>(q),static_cast<uint8_t*>(scales),amax,
+                       static_cast<__nv_bfloat16*>(y)};
+  args.scheduler.max_swizzle_size=swizzle;
+  args.scheduler.raster_order=static_cast<decltype(args.scheduler.raster_order)>(raster);
+  args.hw_info.cluster_shape=dim3(1,1,1);
+  args.hw_info.cluster_shape_fallback=dim3(1,1,1);
+  Gemm gemm;
+  if (gemm.get_workspace_size(args)>workspace_bytes) return -1;
+  auto status=gemm.can_implement(args);
+  if (status!=cutlass::Status::kSuccess) return 100+int(status);
+  status=gemm.initialize(args,workspace,stream);
+  if (status!=cutlass::Status::kSuccess) return 200+int(status);
+  status=gemm.run(stream);
+  return status==cutlass::Status::kSuccess ? int(cudaGetLastError()) : 300+int(status);
+}

fastvideo-kernel/csrc/quantization/minimax_h3_swiglu_quant.cuh

@@ -0,0 +1,88 @@
+// SPDX-License-Identifier: Apache-2.0
+// Private building block for a SwiGLU/NVFP4 GEMM epilogue. No runtime dispatch.
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <cstdint>
+
+__device__ __forceinline__ float rcp_approx(float x) {
+  float y;
+  asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
+  return y;
+}
+
+__device__ __forceinline__ float swiglu_bf16(float value, float gate) {
+  float e;
+  const float log2e = -1.4426950408889634f;
+  const float exponent = __fmul_rn(gate, log2e);
+  asm("ex2.approx.f32 %0, %1;" : "=f"(e) : "f"(exponent));
+  float sigmoid;
+  const float denominator = __fadd_rn(1.0f, e);
+  // Match Triton sigmoid, including subnormal reciprocals for large negative gates.
+  asm("div.full.f32 %0, %1, %2;" : "=f"(sigmoid) : "f"(1.0f), "f"(denominator));
+  return __bfloat162float(__float2bfloat16_rn(__fmul_rn(value, __fmul_rn(gate, sigmoid))));
+}
+
+__device__ __forceinline__ uint8_t scale_fp8(float x) {
+  uint32_t packed;
+  asm("{ .reg .b16 s; cvt.rn.satfinite.e4m3x2.f32 s, %1, %1; mov.b32 %0, {s,s}; }"
+      : "=r"(packed) : "f"(x));
+  return static_cast<uint8_t>(packed);
+}
+
+__device__ __forceinline__ float scale_float(uint8_t x) {
+  // E4M3 positive finite values; use an exact bit construction including subnormals.
+  const int exponent = (x >> 3) & 15;
+  const int mantissa = x & 7;
+  return exponent == 0 ? __fmul_rn(float(mantissa), 0.001953125f)
+                       : __uint_as_float(((exponent + 120) << 23) | (mantissa << 20));
+}
+
+__device__ __forceinline__ uint8_t fp4_pair(float lo, float hi) {
+  uint32_t packed;
+  asm("{ .reg .b8 q; cvt.rn.satfinite.e2m1x2.f32 q, %1, %2; mov.b32 %0, {q,q,q,q}; }"
+      : "=r"(packed) : "f"(hi), "f"(lo));
+  return static_cast<uint8_t>(packed);
+}
+
+__global__ void quant_probe(const __nv_bfloat16* x, const float* sf, uint8_t* q,
+                            uint8_t* scales, float* amax_out, __nv_bfloat16* y_out,
+                            int rows, int width, bool paired) {
+  const int row = blockIdx.x;
+  const int col = blockIdx.y * blockDim.x + threadIdx.x;
+  float value = 0, gate = 0;
+  if (row < rows && col < width) {
+    const int64_t base = int64_t(row) * width * 2;
+    value = __bfloat162float(x[base + (paired ? 2 * col : col)]);
+    gate = __bfloat162float(x[base + (paired ? 2 * col + 1 : width + col)]);
+  }
+  const float y = swiglu_bf16(value, gate);
+  float amax = fabsf(y);
+  for (int delta = 1; delta < 16; delta *= 2)
+    amax = fmaxf(amax, __shfl_xor_sync(0xffffffff, amax, delta, 16));
+  const float global_sf = *sf;
+  const uint8_t scale = scale_fp8(__fmul_rn(global_sf, __fmul_rn(amax, rcp_approx(6.0f))));
+  const float inv = amax != 0 ? rcp_approx(__fmul_rn(scale_float(scale), rcp_approx(global_sf))) : 0;
+  const float code = __fmul_rn(y, inv);
+  const float hi = __shfl_xor_sync(0xffffffff, code, 1);
+  if (row < rows && col < width) {
+    if ((col & 1) == 0) q[int64_t(row) * (width / 2) + col / 2] = fp4_pair(code, hi);
+    if (y_out) y_out[int64_t(row) * width + col] = __float2bfloat16_rn(y);
+  }
+  if (col < width && (col & 15) == 0) {
+    const int group = col / 16;
+    const int64_t index = (((((int64_t(row) / 128) * (width / 64) + group / 4) * 32
+                           + row % 32) * 4 + (row / 32) % 4) * 4 + group % 4);
+    scales[index] = scale;
+    if (amax_out) amax_out[int64_t(row) * (width / 16) + group] = amax;
+  }
+}
+
+extern "C" int launch_quant_probe(const void* x, const float* sf, void* q, void* scales,
+                                  float* amax_out, void* y_out, int rows, int width,
+                                  int paired, void* stream) {
+  quant_probe<<<dim3((rows + 127) / 128 * 128, (width + 255) / 256), 256, 0,
+                static_cast<cudaStream_t>(stream)>>>(static_cast<const __nv_bfloat16*>(x), sf,
+      static_cast<uint8_t*>(q), static_cast<uint8_t*>(scales), amax_out,
+      static_cast<__nv_bfloat16*>(y_out), rows, width, paired);
+  return static_cast<int>(cudaGetLastError());
+}

fastvideo-kernel/python/fastvideo_kernel/h3_swiglu_gemm_binding.py

@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Private eager binding for the SM121 H3 SwiGLU/FP4 GEMM epilogue.
+
+The Spark prebuild installs this module beside its fingerprinted native library.
+Inputs are prequantized, with value/gate weight columns paired in memory.
+"""
+import ctypes
+from pathlib import Path
+
+import torch
+
+_library = ctypes.CDLL(str(Path(__file__).with_name('lib_fastvideo_h3_swiglu_gemm.so')))
+_launch = _library.launch_gemm_swiglu
+_launch.argtypes = ([ctypes.c_void_p] * 10 + [ctypes.c_int] * 3
+                    + [ctypes.c_void_p, ctypes.c_uint64, ctypes.c_int, ctypes.c_int, ctypes.c_void_p])
+_launch.restype = ctypes.c_int
+_workspaces: dict[torch.device, torch.Tensor] = {}
+
+
+def run(a, b, a_scale, b_scale, alpha, global_scale, q, scales, amax=None, values=None,
+        raster=0, swizzle=1):
+    """Write packed SwiGLU outputs; b and b_scale have contiguous transposes."""
+    workspace = _workspaces.get(a.device)
+    if workspace is None:
+        workspace = torch.empty(16 * 1024 * 1024, device=a.device, dtype=torch.uint8)
+        _workspaces[a.device] = workspace
+    rc = _launch(
+        a.data_ptr(), b.data_ptr(), a_scale.data_ptr(), b_scale.data_ptr(), alpha.data_ptr(),
+        global_scale.data_ptr(), q.data_ptr(), scales.data_ptr(),
+        0 if amax is None else amax.data_ptr(), 0 if values is None else values.data_ptr(),
+        a.shape[0], b.shape[1], a.shape[1] * 2, workspace.data_ptr(), workspace.numel(),
+        raster, swizzle, torch.cuda.current_stream(a.device).cuda_stream,
+    )
+    if rc:
+        raise RuntimeError(f'H3 SwiGLU GEMM returned CUDA/CUTLASS error {rc}')

fastvideo/configs/models/dits/__init__.py

@@ -18,7 +18,6 @@ from fastvideo.configs.models.dits.mmaudio import (
 )
 from fastvideo.configs.models.dits.stable_audio import StableAudioConfig
 from fastvideo.models.wan.config import WanVideoConfig
-from fastvideo.configs.models.dits.qwen_image21 import QwenImage21ArchConfig, QwenImage21Config
 from fastvideo.configs.models.dits.hyworld import HYWorldConfig
 from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig
 from fastvideo.configs.models.dits.kandinsky6 import Kandinsky6ArchConfig, Kandinsky6VideoAudioConfig
@@ -34,5 +33,5 @@ __all__ = [
     "StableAudioConfig", "HeliosConfig", "LingBotWorld2CausalFastVideoConfig", "LingBotWorldFastVideoConfig",
     "LingBotVideoConfig", "MiniMaxH3Config", "MMAudioArchConfig", "MMAudioTransformerConfig",
     "MMAUDIO_VARIANT_ARCHITECTURES", "MMAUDIO_TRAINING_VARIANTS", "MMAUDIO_44K_TRAINING_VARIANTS",
-    "get_mmaudio_transformer_config", "QwenImage21ArchConfig", "QwenImage21Config"
+    "get_mmaudio_transformer_config"
 ]

fastvideo/configs/models/dits/qwen_image21.py

@@ -1,55 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Architecture and checkpoint loading configuration for Qwen-Image-2.1."""
-
-from dataclasses import dataclass, field
-
-from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
-from fastvideo.platforms import AttentionBackendEnum
-
-
-def _is_transformer_block(name: str, module) -> bool:
-    parts = name.split(".")
-    return len(parts) >= 2 and parts[-2] == "transformer_blocks" and parts[-1].isdigit()
-
-
-@dataclass
-class QwenImage21ArchConfig(DiTArchConfig):
-    patch_size: int = 1
-    in_channels: int = 64
-    out_channels: int = 64
-    num_layers: int = 32
-    attention_head_dim: int = 128
-    num_attention_heads: int = 32
-    context_in_dim: int = 4096
-    mlp_ratio: int = 3
-    axes_dims_rope: tuple[int, int, int] = (16, 56, 56)
-    eps: float = 1e-6
-    causal_condition: bool = True
-    cast_prompt_embeds_to_dit_dtype: bool = True
-
-    _fsdp_shard_conditions: list = field(default_factory=lambda: [_is_transformer_block])
-    _supported_attention_backends: tuple[AttentionBackendEnum,
-                                         ...] = (AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.FLASH_ATTN,
-                                                 AttentionBackendEnum.SAGE_ATTN)
-    param_names_mapping: dict = field(default_factory=lambda: {r"^(.*)$": r"\1"})
-    reverse_param_names_mapping: dict = field(default_factory=lambda: {r"^(.*)$": r"\1"})
-
-    def __post_init__(self) -> None:
-        axes_dims_rope = tuple(self.axes_dims_rope)
-        if len(axes_dims_rope) != 3 or any(dim <= 0 or dim % 2 for dim in axes_dims_rope):
-            raise ValueError("axes_dims_rope must contain three positive even dimensions")
-        self.axes_dims_rope = (axes_dims_rope[0], axes_dims_rope[1], axes_dims_rope[2])
-        if sum(self.axes_dims_rope) != self.attention_head_dim:
-            raise ValueError("axes_dims_rope must sum to attention_head_dim")
-        if self.num_layers <= 0 or self.num_attention_heads <= 0:
-            raise ValueError("num_layers and num_attention_heads must be positive")
-        super().__post_init__()
-        self.out_channels = self.out_channels or self.in_channels
-        self.hidden_size = self.num_attention_heads * self.attention_head_dim
-        self.num_channels_latents = self.out_channels
-
-
-@dataclass
-class QwenImage21Config(DiTConfig):
-    arch_config: DiTArchConfig = field(default_factory=QwenImage21ArchConfig)
-    prefix: str = "QwenImage21"

fastvideo/configs/models/encoders/__init__.py

@@ -15,7 +15,6 @@ from fastvideo.configs.models.encoders.gemma import LTX2GemmaConfig
 from fastvideo.configs.models.encoders.mistral3 import Mistral3TextConfig
 from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import (MiniMaxH3Qwen3VLArchConfig, MiniMaxH3Qwen3VLConfig)
 from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig
-from fastvideo.configs.models.encoders.qwen_image21 import QwenImage21Qwen3VLArchConfig, QwenImage21Qwen3VLConfig
 from fastvideo.configs.models.encoders.lingbot_video import LingBotVideoQwen3VLTextConfig
 from fastvideo.configs.models.encoders.stable_audio_conditioner import (
     StableAudioConditionerArchConfig,
@@ -50,8 +49,6 @@ __all__ = [
     "StableAudioConditionerConfig",
     "T5GemmaEncoderConfig",
     "Qwen3TextConfig",
-    "QwenImage21Qwen3VLArchConfig",
-    "QwenImage21Qwen3VLConfig",
     "Mistral3TextConfig",
     "LingBotWorld2UMT5ArchConfig",
     "LingBotWorld2UMT5Config",

fastvideo/configs/models/encoders/qwen_image21.py

@@ -1,79 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Qwen3-VL architecture used for Qwen-Image-2.1 conditioning."""
-
-from dataclasses import dataclass, field
-from typing import Any
-
-from fastvideo.configs.models.base import ModelConfig
-from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import (
-    MiniMaxH3Qwen3VLArchConfig,
-    MiniMaxH3Qwen3VLConfig,
-    _OFFICIAL_ARCHITECTURES,
-    _VISION_CONFIG_MAPPING,
-)
-
-
-@dataclass
-class QwenImage21Qwen3VLArchConfig(MiniMaxH3Qwen3VLArchConfig):
-    architectures: list[str] = field(default_factory=lambda: ["QwenImage21Qwen3VLConditioner"])
-    hidden_size: int = 4096
-    intermediate_size: int = 12288
-    num_hidden_layers: int = 36
-    output_hidden_state_index: int = 36
-    num_hidden_layers_override: int | None = None
-    num_attention_heads: int = 32
-    vision_out_hidden_size: int = 4096
-    # The image pipeline's processor owns prompt lengths, including reference
-    # image tokens; the H3 text-only truncation limit is inappropriate here.
-    text_len: int = 262144
-    require_processor: bool = True
-
-    def __post_init__(self) -> None:
-        super().__post_init__()
-        if self.output_hidden_state_index != self.num_hidden_layers:
-            raise ValueError("Qwen-Image-2.1 conditioning requires the last decoder layer before final normalization")
-        if self.num_hidden_layers_override is not None:
-            raise ValueError("Qwen-Image-2.1 conditioning requires all text decoder layers")
-        if self.hidden_act != "silu" or self.vision_hidden_act != "gelu_pytorch_tanh":
-            raise ValueError("Qwen-Image-2.1 requires SiLU text MLPs and tanh-approximate GELU vision MLPs")
-        if (self.rope_scaling or {}).get("rope_type", "default") != "default":
-            raise ValueError("Qwen-Image-2.1 conditioning supports default multimodal RoPE")
-        self.tokenizer_kwargs = {
-            "padding": True,
-            "padding_side": "left",
-            "truncation": False,
-            "return_tensors": "pt",
-        }
-
-
-@dataclass
-class QwenImage21Qwen3VLConfig(MiniMaxH3Qwen3VLConfig):
-    arch_config: QwenImage21Qwen3VLArchConfig = field(default_factory=QwenImage21Qwen3VLArchConfig)
-    prefix: str = "qwen_image21_qwen3_vl"
-
-    def update_model_arch(self, source_model_dict: dict[str, Any]) -> None:
-        flattened = dict(source_model_dict)
-        flattened.update(dict(flattened.pop("text_config", {})))
-        vision = flattened.pop("vision_config", {})
-        for source_name, target_name in _VISION_CONFIG_MAPPING.items():
-            if source_name in vision:
-                flattened[target_name] = vision[source_name]
-
-        architectures = flattened.get("architectures", [])
-        if isinstance(architectures, str):
-            architectures = [architectures]
-        if any(architecture in _OFFICIAL_ARCHITECTURES for architecture in architectures):
-            flattened["architectures"] = ["QwenImage21Qwen3VLConditioner"]
-        if "num_hidden_layers" in flattened:
-            flattened["output_hidden_state_index"] = flattened["num_hidden_layers"]
-        for name in ("mrope_section", "vision_deepstack_visual_indexes"):
-            if isinstance(flattened.get(name), list):
-                flattened[name] = tuple(flattened[name])
-        # Transformers 5 serializes the same default RoPE parameters under a
-        # renamed field. Keep the checkpoint's frequencies in either format.
-        if "rope_parameters" in flattened:
-            parameters = dict(flattened["rope_parameters"])
-            flattened["rope_scaling"] = parameters
-            if "rope_theta" in parameters:
-                flattened["rope_theta"] = parameters["rope_theta"]
-        ModelConfig.update_model_arch(self, flattened)

fastvideo/configs/models/vaes/__init__.py

@@ -12,7 +12,6 @@ from fastvideo.configs.models.vaes.minimax_h3_video import (
 )
 from fastvideo.configs.models.vaes.oobleck import OobleckVAEArchConfig, OobleckVAEConfig
 from fastvideo.configs.models.vaes.flux2vae import Flux2VAEConfig
-from fastvideo.configs.models.vaes.qwen_image21 import QwenImage21VAEArchConfig, QwenImage21VAEConfig
 from fastvideo.models.wan.vae_config import WanVAEConfig
 
 __all__ = [
@@ -32,6 +31,4 @@ __all__ = [
     "OobleckVAEArchConfig",
     "OobleckVAEConfig",
     "Flux2VAEConfig",
-    "QwenImage21VAEArchConfig",
-    "QwenImage21VAEConfig",
 ]

fastvideo/configs/models/vaes/qwen_image21.py

@@ -1,187 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Qwen-Image-2.1 RGBA autoencoder configuration."""
-
-from dataclasses import dataclass, field
-
-from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig
-
-
-@dataclass
-class QwenImage21VAEArchConfig(VAEArchConfig):
-    _class_name: str = "AutoencoderKLQwenImage21"
-    base_dim: int = 96
-    decoder_base_dim: int | None = 144
-    z_dim: int = 64
-    dim_mult: tuple[int, ...] = (1, 2, 4, 8, 8)
-    num_res_blocks: int = 2
-    attn_scales: tuple[float, ...] = ()
-    temperal_downsample: tuple[bool, ...] = (False, True, True, True)
-    dropout: float = 0.0
-    is_residual: bool = True
-    in_channels: int = 4
-    out_channels: int = 4
-    patch_size: int | None = None
-    scale_factor_spatial: int = 16
-    scale_factor_temporal: int = 8
-    latents_mean: tuple[float, ...] = (
-        0.5126,
-        0.7721,
-        -0.0631,
-        1.3506,
-        -0.7855,
-        -2.1025,
-        -0.3458,
-        1.3722,
-        1.8873,
-        -1.7177,
-        -0.6510,
-        0.2732,
-        0.7562,
-        -0.6163,
-        -1.0277,
-        3.8363,
-        2.0210,
-        0.0472,
-        0.9320,
-        2.0087,
-        2.4954,
-        -0.1391,
-        -1.4249,
-        1.8464,
-        -0.5236,
-        1.2826,
-        3.7046,
-        -1.3035,
-        2.7286,
-        -1.4518,
-        -1.9036,
-        -1.9955,
-        -0.0342,
-        -1.0265,
-        -0.7636,
-        3.0555,
-        0.0746,
-        -3.0751,
-        -0.1076,
-        1.7376,
-        -1.0914,
-        -1.9435,
-        -0.2784,
-        -1.3680,
-        0.4809,
-        -0.4433,
-        0.3764,
-        0.5729,
-        -2.0595,
-        1.0960,
-        -1.3260,
-        -2.0211,
-        -5.0179,
-        0.5275,
-        4.0162,
-        1.8505,
-        0.3026,
-        1.9373,
-        1.4937,
-        0.2632,
-        0.5547,
-        -1.7121,
-        -0.1562,
-        0.0304,
-    )
-    latents_std: tuple[float, ...] = (
-        3.2001,
-        3.2936,
-        3.4321,
-        3.0091,
-        3.1061,
-        4.0379,
-        4.0705,
-        3.7910,
-        3.0785,
-        3.6500,
-        3.9308,
-        3.0904,
-        2.8778,
-        3.7675,
-        3.7320,
-        5.0756,
-        3.2864,
-        4.0397,
-        3.1317,
-        4.0443,
-        2.9249,
-        3.9454,
-        3.0988,
-        4.2489,
-        3.4896,
-        3.8513,
-        3.9323,
-        3.4719,
-        3.7498,
-        4.2830,
-        3.5694,
-        4.2467,
-        3.9037,
-        3.2947,
-        5.0770,
-        3.5075,
-        3.2700,
-        3.4767,
-        2.8063,
-        5.1125,
-        3.5327,
-        4.7833,
-        3.1286,
-        4.1819,
-        3.8527,
-        3.8312,
-        3.5605,
-        4.3875,
-        3.9624,
-        4.0168,
-        3.5643,
-        4.0550,
-        5.5614,
-        4.2963,
-        4.4080,
-        3.4959,
-        3.8747,
-        3.7608,
-        3.5735,
-        3.1490,
-        3.7662,
-        3.6746,
-        3.4563,
-        3.8161,
-    )
-    scaling_factor: float = 1.0
-    temporal_compression_ratio: int = 1
-    spatial_compression_ratio: int = 16
-    param_names_mapping: dict[str, str] = field(default_factory=lambda: {r"^(.*)$": r"\1"})
-
-    @property
-    def latent_channels(self) -> int:
-        return self.z_dim
-
-    def __post_init__(self) -> None:
-        if len(self.dim_mult) < 2 or len(self.temperal_downsample) != len(self.dim_mult) - 1:
-            raise ValueError("`temperal_downsample` must have one entry per downsampling stage.")
-        if len(self.latents_mean) != self.z_dim or len(self.latents_std) != self.z_dim:
-            raise ValueError("Qwen-Image-2.1 latent mean/std must have one entry per latent channel.")
-        if any(value <= 0 for value in self.latents_std):
-            raise ValueError("Qwen-Image-2.1 latent standard deviations must be positive.")
-        if self.patch_size is not None and self.patch_size < 1:
-            raise ValueError("`patch_size` must be positive or None.")
-        actual_spatial_scale = 2**(len(self.dim_mult) - 1) * (self.patch_size or 1)
-        if self.scale_factor_spatial != actual_spatial_scale:
-            raise ValueError("`scale_factor_spatial` must match the VAE downsampling stages and patch size.")
-        self.spatial_compression_ratio = self.scale_factor_spatial
-
-
-@dataclass
-class QwenImage21VAEConfig(VAEConfig):
-    arch_config: QwenImage21VAEArchConfig = field(default_factory=QwenImage21VAEArchConfig)
-    use_tiling: bool = False
-    use_temporal_tiling: bool = False
-    use_parallel_tiling: bool = False

fastvideo/configs/pipelines/__init__.py

@@ -14,7 +14,6 @@ from fastvideo.configs.pipelines.lingbot_video import LingBotVideoT2VConfig
 from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig
 from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig
 from fastvideo.configs.pipelines.mmaudio import MMAudioV2AConfig
-from fastvideo.configs.pipelines.qwen_image21 import QwenImage21PipelineConfig
 from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig
 from fastvideo.registry import get_pipeline_config_cls_from_name
 from fastvideo.models.wan.pipeline_config import (LucyEditDevConfig, SelfForcingWanT2V480PConfig, WanI2V480PConfig,
@@ -27,5 +26,5 @@ __all__ = [
     "Kandinsky5T2VConfig", "Kandinsky5I2VConfig", "Kandinsky5DMDConfig", "Kandinsky6TI2VAConfig",
     "Kandinsky6SRPipelineConfig", "LingBotWorld2CausalFastI2V480PConfig", "LingBotWorldFastI2V480PConfig",
     "LingBotVideoT2VConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "MMAudioV2AConfig",
-    "QwenImage21PipelineConfig", "get_pipeline_config_cls_from_name"
+    "get_pipeline_config_cls_from_name"
 ]

fastvideo/configs/pipelines/base.py

@@ -64,7 +64,6 @@ class PipelineConfig:
     vae_decode_precision: str | None = None
     vae_tiling: bool = True
     vae_sp: bool = True
-    output_channels: int = 3
 
     # Image encoder configuration
     image_encoder_config: EncoderConfig = field(default_factory=EncoderConfig)

fastvideo/configs/pipelines/open_h3_regenerate_2k.py

@@ -1,30 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Pipeline configuration for Open-H3-Regenerate-2K (768P H3 draft -> 2560x1440)."""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-
-from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig
-
-DRAFT_SIZE_CHOICES = ("768p", "576p")
-
-
-@dataclass
-class OpenH3Regenerate2KPipelineConfig(MiniMaxH3PipelineConfig):
-    """MiniMax-H3 policy plus the refine-stage knobs.
-
-    ``draft_size`` is the size the draft video is VAE-encoded at: 768p (1344x768) for MiniMax-H3 drafts, 576p
-    (1024x576) for drafts from other video models. ``refine_start_sigma`` re-noises the upscaled draft to another
-    level; the later steps then follow the shift-6 rule (None keeps the released grid).
-    """
-
-    draft_size: str = "768p"
-    refine_start_sigma: float | None = None
-
-    def check_pipeline_config(self) -> None:
-        super().check_pipeline_config()
-        if self.draft_size not in DRAFT_SIZE_CHOICES:
-            raise ValueError(f"draft_size must be one of {DRAFT_SIZE_CHOICES}, got {self.draft_size!r}.")
-        if self.refine_start_sigma is not None and not 0.0 < self.refine_start_sigma < 1.0:
-            raise ValueError(f"refine_start_sigma must be in (0, 1), got {self.refine_start_sigma}.")

fastvideo/configs/pipelines/qwen_image21.py

@@ -1,30 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Native Qwen-Image-2.1 inference configuration."""
-
-from dataclasses import dataclass, field
-
-from fastvideo.configs.models.dits.qwen_image21 import QwenImage21Config
-from fastvideo.configs.models.encoders.qwen_image21 import QwenImage21Qwen3VLConfig
-from fastvideo.configs.models.vaes.qwen_image21 import QwenImage21VAEConfig
-from fastvideo.configs.pipelines.base import PipelineConfig
-
-
-@dataclass
-class QwenImage21PipelineConfig(PipelineConfig):
-    dit_config: QwenImage21Config = field(default_factory=QwenImage21Config)
-    vae_config: QwenImage21VAEConfig = field(default_factory=QwenImage21VAEConfig)
-    text_encoder_configs: tuple[QwenImage21Qwen3VLConfig,
-                                ...] = field(default_factory=lambda: (QwenImage21Qwen3VLConfig(), ))
-    text_encoder_precisions: tuple[str, ...] = ("bf16", )
-    dit_precision: str = "bf16"
-    vae_precision: str = "bf16"
-    vae_tiling: bool = True
-    vae_sp: bool = False
-    # Four output channels are required to preserve transparency through PNG export.
-    output_channels: int = 4
-    flow_shift: float | None = None
-    scheduler_step_in_fp32: bool = True
-    reference_resolution: int = 1024
-    use_kv_cache: bool = True
-    # Runtime placement policy. CPU keeps ten reference prefixes out of VRAM.
-    kv_cache_device: str = "cpu"

fastvideo/distributed/async_collectives.py

@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tensor collectives with explicit input ownership until their consumers wait."""
+
+from dataclasses import dataclass
+import math
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+
+
+@dataclass
+class PendingTensor:
+    output: torch.Tensor
+    source: torch.Tensor
+    work: dist.Work
+    logical_shape: torch.Size | None = None
+
+    def wait(self) -> torch.Tensor:
+        """Order consumers on the current CUDA stream after the collective."""
+        self.work.wait()
+        if self.logical_shape is None:
+            return self.output
+        peer_elements = math.prod(self.logical_shape) // self.output.shape[0]
+        return self.output[:, :peer_elements].contiguous().view(self.logical_shape)
+
+
+def all_to_all_tensor_async(source: torch.Tensor, group: dist.ProcessGroup, *, align_bytes: int = 1) -> PendingTensor:
+    """Exchange equal peer slices, optionally aligning their NCCL buffer offsets.
+
+    Small FP32 scale slices with an odd head count can otherwise start eight
+    bytes off a 16-byte boundary and defeat NCCL's vectorized copies. Padding
+    only the transport buffer retains the caller's logical tensor layout.
+    """
+    logical_shape = None
+    if align_bytes > 1:
+        if align_bytes % source.element_size():
+            raise ValueError("Collective alignment must be a multiple of the element size")
+        world = dist.get_world_size(group)
+        if source.numel() % world:
+            raise ValueError("Equal all-to-all slices require a tensor divisible by the group size")
+        elements = source.numel() // world
+        padding = -elements % (align_bytes // source.element_size())
+        if padding:
+            logical_shape = source.shape
+            source = F.pad(source.reshape(world, elements), (0, padding))
+    output = torch.empty_like(source)
+    work = dist.all_to_all_single(output, source, group=group, async_op=True)
+    assert work is not None
+    return PendingTensor(output, source, work, logical_shape)
+
+
+def all_gather_tensor_async(source: torch.Tensor, group: dist.ProcessGroup, world: int) -> PendingTensor:
+    output = torch.empty((world, *source.shape), device=source.device, dtype=source.dtype)
+    work = dist.all_gather_into_tensor(output, source, group=group, async_op=True)
+    assert work is not None
+    return PendingTensor(output, source, work)
+
+
+@dataclass
+class PeerSlices:
+    """Two-rank receive whose local slice stays in the producer's storage."""
+    local: torch.Tensor
+    remote: torch.Tensor
+    template: torch.Tensor
+    rank: int
+
+    def materialize(self) -> torch.Tensor:
+        """Restore source-rank order for consumers without a two-pointer kernel."""
+        slices = (self.local, self.remote) if self.rank == 0 else (self.remote, self.local)
+        return torch.stack([part.view(torch.uint8) for part in slices]).view(self.template.dtype)
+
+
+@dataclass
+class PendingPeerSlices:
+    source: torch.Tensor
+    remote: torch.Tensor
+    rank: int
+    works: list[dist.Work]
+
+    def wait(self) -> PeerSlices:
+        for work in self.works:
+            work.wait()
+        return PeerSlices(self.source[self.rank], self.remote.view(self.source.dtype), self.source, self.rank)
+
+
+def exchange_peer_tensor_async(source: torch.Tensor, group: dist.ProcessGroup) -> PendingPeerSlices:
+    """Exchange only the remote half of a contiguous two-rank payload.
+
+    Every rank must enqueue calls in the same order. The producer is retained
+    through completion and consumption, including the local slice that bypasses
+    NCCL. Receivers index the two slices directly to avoid a concatenation copy.
+    """
+    if dist.get_world_size(group) != 2 or source.shape[0] != 2 or not source.is_contiguous():
+        raise ValueError("Peer-slice exchange requires two contiguous destination slices")
+    rank = dist.get_rank(group)
+    peer = dist.get_global_rank(group, 1 - rank)
+    send = source[1 - rank].view(torch.uint8)
+    remote = torch.empty_like(send)
+    works = dist.batch_isend_irecv(
+        [dist.P2POp(dist.irecv, remote, peer, group),
+         dist.P2POp(dist.isend, send, peer, group)])
+    return PendingPeerSlices(source, remote, rank, works)

fastvideo/distributed/projection_exchange.py

@@ -0,0 +1,120 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Two-rank QKV traffic that can start before all projections are ready."""
+
+from dataclasses import dataclass
+
+import torch
+import torch.distributed as dist
+
+from fastvideo.distributed.async_collectives import PeerSlices
+
+
+class Completion:
+
+    def __init__(self, device: torch.device) -> None:
+        self.device = device
+        self.works: list[dist.Work] = []
+        self.waited_streams: set[int] = set()
+        self.ready = torch.cuda.Event()
+        self.producer = torch.cuda.current_stream(device).cuda_stream
+
+    def wait(self) -> None:
+        stream = torch.cuda.current_stream(self.device).cuda_stream
+        if stream not in self.waited_streams:
+            if stream != self.producer:
+                torch.cuda.current_stream(self.device).wait_event(self.ready)
+            for work in self.works:
+                work.wait()
+            self.waited_streams.add(stream)
+
+
+@dataclass
+class PendingInput:
+    source: torch.Tensor
+    remote: torch.Tensor
+    rank: int
+    completion: Completion
+
+    def wait(self) -> PeerSlices:
+        self.completion.wait()
+        return PeerSlices(self.source[self.rank], self.remote.view(self.source.dtype), self.source, self.rank)
+
+
+@dataclass
+class PendingScale:
+    source: torch.Tensor
+    output: torch.Tensor
+    completion: Completion
+
+    def wait(self) -> torch.Tensor:
+        self.completion.wait()
+        return self.output
+
+
+class ProjectionExchange:
+    """Exchange Q/K/V as each projection becomes ready, retaining their storage.
+
+    Input is [head groups, destinations, QKV, rows, heads, 128]. Call start
+    in Q/K/V order on the producing CUDA stream, then packets on that stream.
+    Input payload and scales must remain unmodified until all consumers wait.
+    The local payload aliases its producer; only the remote slice travels.
+    """
+
+    def __init__(self, payload: torch.Tensor, scale: torch.Tensor, group: dist.ProcessGroup) -> None:
+        if (dist.get_world_size(group) != 2 or payload.ndim != 6 or payload.shape[1:3] != (2, 3)
+                or payload.shape[-1] != 128 or not payload.is_cuda or not payload.is_contiguous()
+                or payload.dtype != torch.float8_e4m3fn or scale.shape != payload.shape[:-1]
+                or scale.device != payload.device or scale.dtype != torch.float32 or not scale.is_contiguous()):
+            raise ValueError("Projection exchange requires contiguous two-rank FP8 QKV and FP32 scales")
+        self.payload, self.scale, self.group = payload, scale, group
+        self.rank = dist.get_rank(group)
+        self.peer = dist.get_global_rank(group, 1 - self.rank)
+        self.remote = torch.empty((payload.shape[0], *payload.shape[2:]), device=payload.device, dtype=torch.uint8)
+        self.received_scale = torch.empty_like(scale)
+        self.completion = Completion(payload.device)
+        self.next_projection = 0
+
+    def start(self, projection: int) -> None:
+        if projection != self.next_projection or projection not in (0, 1, 2):
+            raise ValueError("Every rank must start projections in Q, K, V order")
+        self.next_projection += 1
+        ops: list[dist.P2POp] = []
+        for chunk in range(self.payload.shape[0]):
+            send = self.payload[chunk, 1 - self.rank, projection].view(torch.uint8)
+            remote = self.remote[chunk, projection]
+            send_scale = self.scale[chunk, 1 - self.rank, projection]
+            remote_scale = self.received_scale[chunk, 1 - self.rank, projection]
+            ops.extend((dist.P2POp(dist.irecv, remote, self.peer,
+                                   self.group), dist.P2POp(dist.isend, send, self.peer, self.group),
+                        dist.P2POp(dist.irecv, remote_scale, self.peer,
+                                   self.group), dist.P2POp(dist.isend, send_scale, self.peer, self.group)))
+        self.completion.works.extend(dist.batch_isend_irecv(ops))
+
+    def packets(self) -> list[tuple[PendingInput, PendingScale]]:
+        """Return consumers after all projections have been queued on this stream."""
+        if self.next_projection != 3:
+            raise RuntimeError("All three projection sends must start before consumption")
+        self.received_scale[:, self.rank].copy_(self.scale[:, self.rank])
+        self.completion.ready.record()
+        return [(PendingInput(self.payload[c], self.remote[c], self.rank,
+                              self.completion), PendingScale(self.scale[c], self.received_scale[c], self.completion))
+                for c in range(self.payload.shape[0])]
+
+
+def completed_projection_packets(payload: torch.Tensor, scale: torch.Tensor,
+                                 rank: int) -> list[tuple[PendingInput, PendingScale]]:
+    """Expose already source-ordered QKV packets through the usual consumers.
+
+    Input is [head groups, source ranks, QKV, rows, heads, 128]. An event orders
+    consumers on other CUDA streams after production; handles retain both
+    source slices and scales throughout their consumption.
+    """
+    if (payload.ndim != 6 or payload.shape[1:3] != (2, 3) or payload.shape[-1] != 128 or rank not in (0, 1)
+            or not payload.is_cuda or not payload.is_contiguous() or payload.dtype != torch.float8_e4m3fn
+            or scale.shape != payload.shape[:-1] or scale.device != payload.device or scale.dtype != torch.float32
+            or not scale.is_contiguous()):
+        raise ValueError("Completed projection packets require two contiguous source-rank slices")
+    completion = Completion(payload.device)
+    completion.ready.record()
+    return [(PendingInput(payload[c], payload[c, 1 - rank].view(torch.uint8), rank,
+                          completion), PendingScale(scale[c], scale[c], completion)) for c in range(payload.shape[0])]

fastvideo/entrypoints/openai/request_adapter.py

@@ -41,10 +41,6 @@ class RequestAdaptationError(ValueError):
     """The transport request cannot be represented by the loaded pipeline."""
 
 
-OPEN_H3_REGENERATE_2K_FAMILY = "open_h3_regenerate_2k"
-OPEN_H3_REGENERATE_2K_SIZE = (2560, 1440)
-OPEN_H3_REGENERATE_2K_NUM_FRAMES = 124
-
 _VIDEO_SUFFIXES = frozenset({".avi", ".m4v", ".mkv", ".mov", ".mp4", ".mpeg", ".mpg", ".webm"})
 
 
@@ -385,13 +381,6 @@ def _apply_reference_inputs(
     if request.task is not None and model_family != "minimax_h3":
         raise RequestAdaptationError("The task selector is only defined for MiniMax-H3 servers.")
 
-    if model_family == OPEN_H3_REGENERATE_2K_FAMILY:
-        if len(videos) != 1 or images or audios:
-            raise RequestAdaptationError("Open-H3-Regenerate-2K takes exactly one video_reference (the 768P draft) "
-                                         "and no image or audio references.")
-        kwargs["video_path"] = videos[0]
-        return
-
     if model_family == "minimax_h3" and (videos or audios):
         raise RequestAdaptationError("MiniMax-H3 video/audio references require a server configured with "
                                      "override_pipeline_cls_name=MiniMaxH3Ref2VAModularPipeline.")
@@ -533,16 +522,6 @@ def build_generation_request(
     if height is not None and (not isinstance(height, int) or height <= 0):
         raise RequestAdaptationError(f"height must be a positive integer, got {height!r}")
 
-    if model_family == OPEN_H3_REGENERATE_2K_FAMILY:
-        if (width, height) not in ((None, None), OPEN_H3_REGENERATE_2K_SIZE):
-            raise RequestAdaptationError("Open-H3-Regenerate-2K always renders 2560x1440; omit size/width/height "
-                                         f"or pass 2560x1440, got {width}x{height}.")
-        if kwargs.get("num_frames", OPEN_H3_REGENERATE_2K_NUM_FRAMES) != OPEN_H3_REGENERATE_2K_NUM_FRAMES:
-            raise RequestAdaptationError("Open-H3-Regenerate-2K regenerates the draft's first 121 frames (24 fps); "
-                                         "omit seconds and num_frames.")
-        if kwargs["fps"] != 24:
-            raise RequestAdaptationError(f"Open-H3-Regenerate-2K requires fps=24, got {kwargs['fps']}.")
-
     if model_family == "minimax_h3":
         from fastvideo.pipelines.basic.minimax_h3.packing import (
             MINIMAX_H3_CANVAS_MULTIPLE,

fastvideo/entrypoints/video_generator.py

@@ -949,8 +949,7 @@ class VideoGenerator:
             samples = torch.empty(0, device='cpu')
         else:
             samples = allocate_cpu_tensor_with_pin_fallback(
-                (latent_batch_size, getattr(fastvideo_args.pipeline_config, "output_channels",
-                                            3), sampling_param.num_frames, sampling_param.height, sampling_param.width),
+                (latent_batch_size, 3, sampling_param.num_frames, sampling_param.height, sampling_param.width),
                 pin_memory=fastvideo_args.pin_cpu_memory)
         try:
             thread.join()

fastvideo/fastvideo_args.py

@@ -234,6 +234,9 @@ class FastVideoArgs:
     # Exported NVFP4 MiniMax-H3 decoder file (decoder-nvfp4.safetensors or decoder.safetensors) that replaces
     # the VAE decoder at load. Overrides FASTVIDEO_H3_VAE_NVFP4_DECODER. Needs Blackwell FP4 hardware.
     h3_vae_nvfp4_decoder: str | None = None
+    # Quantize the loaded dense decoder in memory, preserving its full depth and original weights file.
+    # Mutually exclusive with an exported NVFP4 decoder; skips the INT8 overlay and uses BF16 decode autocast.
+    h3_vae_nvfp4_quantize_dense: bool = False
 
     # Load each heavy component on first use and free it once the last stage
     # that holds it has run, instead of keeping every component resident from
@@ -443,12 +446,19 @@ class FastVideoArgs:
             env_backend = envs.FASTVIDEO_ATTENTION_BACKEND.get()
             if env_backend is not None and backend_name_to_enum(env_backend) is not None:
                 self.attention_backend = env_backend
-        allowed_fp4 = {"enabled", "prefix_precision", "two_level_p", "smooth_k"}
+        allowed_fp4 = {
+            "enabled", "prefix_precision", "two_level_p", "smooth_k", "sp_head_chunks", "sp_pooled_key_mean",
+            "sp_projection_overlap", "sp_direct_keys", "sp_input_gather"
+        }
         if not isinstance(self.h3_vsa_fp4, dict) or self.h3_vsa_fp4.keys() - allowed_fp4:
             raise ValueError(f"h3_vsa_fp4 must be a mapping with keys from {sorted(allowed_fp4)}")
-        for key in ("enabled", "two_level_p", "smooth_k"):
+        for key in ("enabled", "two_level_p", "smooth_k", "sp_pooled_key_mean", "sp_projection_overlap",
+                    "sp_direct_keys", "sp_input_gather"):
             if key in self.h3_vsa_fp4 and not isinstance(self.h3_vsa_fp4[key], bool):
                 raise ValueError(f"h3_vsa_fp4.{key} must be a boolean")
+        head_chunks = self.h3_vsa_fp4.get("sp_head_chunks", 1)
+        if type(head_chunks) is not int or head_chunks < 1:
+            raise ValueError("h3_vsa_fp4.sp_head_chunks must be a positive integer")
         if self.h3_vsa_fp4.get("prefix_precision", "") not in ("", "fp8", "bf16"):
             raise ValueError("h3_vsa_fp4.prefix_precision must be empty, fp8 or bf16")
         if self.h3_vsa_fp4.get("enabled") and (self.inference_torch_compile or self.enable_torch_compile):
@@ -881,13 +891,6 @@ class FastVideoArgs:
             default=FastVideoArgs.taeh3_chunk_size,
             help="TAEH3 latent frames per execution chunk.",
         )
-        parser.add_argument(
-            "--h3-vae-nvfp4-decoder",
-            type=str,
-            default=FastVideoArgs.h3_vae_nvfp4_decoder,
-            help="Exported NVFP4 MiniMax-H3 decoder (decoder-nvfp4.safetensors) that replaces the VAE decoder at "
-            "load. Blackwell only. Overrides FASTVIDEO_H3_VAE_NVFP4_DECODER.",
-        )
         parser.add_argument(
             "--vae-parallel-decode",
             action=StoreBoolean,

fastvideo/layers/lora/linear.py

@@ -155,43 +155,37 @@ class BaseLayerWithLoRA(nn.Module):
         """The wrapped layer's bias, or ``None`` when it has none."""
         return getattr(self.base_layer, "bias", None)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # The wrapper adds no compiled region around the base layer, so it runs exactly as it would unwrapped: eagerly,
-        # or inside a caller's compiled block. Compiling a packed NVFP4/FP8 linear on its own changes its numerics, and
-        # a wrapped layer with no adapter or strength 0 must equal the unwrapped layer exactly.
-        if self.lora_A is None or (self.side_branch and self._strength_is_zero):
-            # No adapter, or strength 0 on a side branch, is exactly the unadapted layer, down to the output dtype.
-            # A merge-mode wrapper reaches this branch while its first adapter is still loading or after a load
-            # failed: the pipeline wraps every linear before it reads the weights, so an unadapted request must
-            # not depend on an adapter that is not there.
-            return self.base_layer(x)
-
-        out, output_bias = self.base_layer(x)
-        if self.merged or self.disable_lora:
-            return out.to(x), output_bias
-        return self._add_adapter(out, x), output_bias
-
     @torch.compile()
-    def _add_adapter(self, out: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         lora_A = self.lora_A
         lora_B = self.lora_B
-        if isinstance(lora_B, DTensor):
-            lora_B = lora_B.to_local()
-            lora_A = lora_A.to_local()
-        lora_A_sliced = self.slice_lora_a_weights(lora_A.to(x, non_blocking=True))
-        lora_B_sliced = self.slice_lora_b_weights(lora_B.to(x, non_blocking=True))
-        alpha_over_rank = 1.0
-        if self.lora_alpha != self.lora_rank:
-            alpha_over_rank = float(self.lora_alpha) / float(self.lora_rank)  # type: ignore[arg-type]
-        if self.side_branch:
-            # Keep chunk temporaries opaque to Inductor so fullgraph capture cannot retain a full-size delta.
-            _add_low_rank_inplace(out, x, lora_A_sliced, lora_B_sliced, self._side_branch_strength, alpha_over_rank)
-            return out
-        scale = float(self.lora_strength) * alpha_over_rank
-        delta = x @ lora_A_sliced.transpose(0, 1) @ lora_B_sliced.transpose(0, 1)
-        if scale != 1.0:
-            delta = delta * scale
-        return out + delta
+        if isinstance(self.lora_B, DTensor):
+            lora_B = self.lora_B.to_local()
+            lora_A = self.lora_A.to_local()
+
+        if self.side_branch and (self.lora_A is None or self._strength_is_zero):
+            # No adapter, or strength 0, is exactly the unadapted layer, down to the output dtype.
+            return self.base_layer(x)
+
+        if not self.merged and not self.disable_lora:
+            lora_A_sliced = self.slice_lora_a_weights(lora_A.to(x, non_blocking=True))
+            lora_B_sliced = self.slice_lora_b_weights(lora_B.to(x, non_blocking=True))
+            alpha_over_rank = 1.0
+            if self.lora_alpha != self.lora_rank:
+                alpha_over_rank = float(self.lora_alpha) / float(self.lora_rank)  # type: ignore[arg-type]
+            out, output_bias = self.base_layer(x)
+            if self.side_branch:
+                # Keep chunk temporaries opaque to Inductor so fullgraph capture cannot retain a full-size delta.
+                _add_low_rank_inplace(out, x, lora_A_sliced, lora_B_sliced, self._side_branch_strength, alpha_over_rank)
+                return out, output_bias
+            scale = float(self.lora_strength) * alpha_over_rank
+            delta = x @ lora_A_sliced.transpose(0, 1) @ lora_B_sliced.transpose(0, 1)
+            if scale != 1.0:
+                delta = delta * scale
+            return out + delta, output_bias
+        else:
+            out, output_bias = self.base_layer(x)
+            return out.to(x), output_bias
 
     def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor:
         return A
@@ -277,14 +271,12 @@ class BaseLayerWithLoRA(nn.Module):
             current_device = self.base_layer.weight.data.device
             data = self.base_layer.weight.data.to(get_local_torch_device())
 
-            # Apply LoRA with alpha scaling. An adapter stored wider than the base (an FP32 adapter on a BF16 DiT)
-            # is multiplied and added at its own precision and rounded once, as a merge-then-cast reference does.
-            accumulate_dtype = torch.promote_types(data.dtype, self.lora_A.dtype)
-            lora_delta = (self.slice_lora_b_weights(self.lora_B.to(accumulate_dtype)) @ self.slice_lora_a_weights(
-                self.lora_A.to(accumulate_dtype)))
+            # Apply LoRA with alpha scaling
+            lora_delta = (
+                self.slice_lora_b_weights(self.lora_B.to(data)) @ self.slice_lora_a_weights(self.lora_A.to(data)))
             scale = ((self.lora_alpha / self.lora_rank) if (self.lora_alpha and self.lora_rank) else 1.0)
             lora_delta *= scale * self.lora_strength
-            data = (data.to(accumulate_dtype) + lora_delta).to(data.dtype)
+            data += lora_delta
             self.base_layer.weight.data = data.to(current_device, non_blocking=True)
 
         self.merged = True

fastvideo/layers/quantization/nvfp4_config.py

@@ -418,6 +418,14 @@ def _mm_fp4(
         raise TypeError(f"Unsupported kwargs for _mm_fp4: {sorted(kwargs)}")
     if alpha is not None and not torch.is_tensor(alpha):
         alpha = torch.tensor(alpha, device=a.device, dtype=torch.float32)
+    # Compact H3 FFN-up: direct output slices avoid the concatenation that
+    # otherwise erases cuDNN's gain from these smaller row plans on GB10.
+    if (out is None and backend == "cudnn" and block_size == 16 and use_nvfp4 and not use_8x4_sf_layout
+            and (a.shape[0], a.shape[1], b.shape[1]) == (14859, 2688, 28672) and a.dtype == torch.uint8
+            and a_scale.ndim == 2 and a_scale.shape[0] >= a.shape[0] and a.is_cuda
+            and not torch.compiler.is_compiling() and envs.FASTVIDEO_NVFP4_MM_ROW_SPLIT.get() != "off"
+            and _is_dgx_spark(a.device.index or 0)):
+        return _mm_fp4_direct_rows(a, b, a_scale, b_scale, alpha, out_dtype, 8)
     if (out is None and backend == "cudnn" and a.shape[0] >= _SPLIT_MIN_ROWS
             and envs.FASTVIDEO_NVFP4_MM_ROW_SPLIT.get() != "off" and not torch.compiler.is_compiling()):
         return _mm_fp4_row_split(a, b, a_scale, b_scale, alpha, out_dtype, block_size, use_8x4_sf_layout, backend,
@@ -462,6 +470,28 @@ def _mm_fp4_row_chunks(a: torch.Tensor, a_scale: torch.Tensor, parts: int) -> li
     return chunks
 
 
+def _mm_fp4_direct_rows(a, b, a_scale, b_scale, alpha, out_dtype, parts: int) -> torch.Tensor:
+    """Run aligned cuDNN row chunks into one contiguous destination, including the final partial tile."""
+    _, mm_fp4, _ = _require_flashinfer()
+    out = torch.empty((a.shape[0], b.shape[1]), device=a.device, dtype=out_dtype)
+    offset = 0
+    for chunk, scale in _mm_fp4_row_chunks(a, a_scale, parts):
+        rows = chunk.shape[0]
+        mm_fp4(chunk,
+               b,
+               scale,
+               b_scale,
+               alpha,
+               out_dtype,
+               out[offset:offset + rows],
+               block_size=16,
+               use_8x4_sf_layout=False,
+               backend="cudnn",
+               use_nvfp4=True)
+        offset += rows
+    return out
+
+
 def _mm_fp4_chunked(a, b, a_scale, b_scale, alpha, out_dtype, block_size, use_8x4_sf_layout, backend, use_nvfp4,
                     parts: int) -> torch.Tensor:
     chunks = _mm_fp4_row_chunks(a, a_scale, parts) if parts > 1 else [(a, a_scale)]
@@ -920,7 +950,7 @@ def load_minimax_h3_nvfp4_dit_export(
         for name, module in modules.items() if isinstance(getattr(module, "quant_method", None), NVFP4QuantizeMethod)
     }
     groups: dict[str, dict[str, str]] = {}
-    with safe_open(path, framework="pt", device="cpu") as reader:
+    with safe_open(path, framework="pt", device=str(device)) as reader:
         for key in reader.keys():  # noqa: SIM118
             if H3_NVFP4_DIT_KEY_SEP not in key:
                 raise ValueError(f"MiniMax-H3 NVFP4 DiT export key {key!r} is missing {H3_NVFP4_DIT_KEY_SEP!r}")
@@ -944,7 +974,7 @@ def load_minimax_h3_nvfp4_dit_export(
             for buffer_name in H3_NVFP4_DIT_BUFFER_NAMES + (H3_NVFP4_DIT_INPUT_SF_NAME, ):
                 if buffer_name not in buffers:
                     continue
-                tensor = reader.get_tensor(buffers[buffer_name]).to(device=device)
+                tensor = reader.get_tensor(buffers[buffer_name])
                 module.register_buffer(buffer_name, tensor, persistent=False)
             module.register_parameter("weight", None)
             loaded_names.add(next(name for name, candidate in modules.items() if candidate is module))

fastvideo/mlx_runtime/minimax_h3.py

@@ -1603,7 +1603,6 @@ class MLXMiniMaxH3DiT:
         pdd_block: tuple[int, int] | None = None,
         pdd_integration_weights: dict[str, np.ndarray] | None = None,
         record_prefix_kv: Any | None = None,
-        precomputed_text_embeds: Any | None = None,
     ):
         """Denoise-step forward served entirely from the AdaLN cache.
 
@@ -1617,11 +1616,6 @@ class MLXMiniMaxH3DiT:
         ``step_timesteps``. With a PDD checkpoint pass the step's
         ``pdd_block`` and ``pdd_integration_weights`` so the widened output
         heads are fused into the block's weighted mean.
-
-        ``precomputed_text_embeds`` may hold ``refine_text(text_rows)`` for
-        this request. The refiner is independent of the denoising timestep;
-        callers own its lifetime and must recompute it for changed text or
-        weights. No embeddings are retained on the model.
         """
         import mlx.core as mx
 
@@ -1646,16 +1640,7 @@ class MLXMiniMaxH3DiT:
             self.weights["audio_proj_in.weight"],
             self.weights["audio_proj_in.bias"],
         )
-        if precomputed_text_embeds is None:
-            text_embeds = self.refine_text(text_rows)
-        else:
-            expected_shape = (len(layout.text_indices), self.hidden_size)
-            if precomputed_text_embeds.shape != expected_shape:
-                raise ValueError(f"Precomputed text embeddings have shape {precomputed_text_embeds.shape}; "
-                                 f"expected {expected_shape}.")
-            if precomputed_text_embeds.dtype not in (mx.float32, mx.float16, mx.bfloat16):
-                raise ValueError("Precomputed text embeddings must have a floating-point dtype.")
-            text_embeds = precomputed_text_embeds
+        text_embeds = self.refine_text(text_rows)
 
         packed = mx.zeros((sequence_length, self.hidden_size), dtype=text_embeds.dtype)
         packed[mx.array(layout.text_indices)] = text_embeds

fastvideo/mlx_runtime/minimax_h3_pipeline.py

@@ -937,10 +937,6 @@ class MiniMaxH3MLXPipeline:
         elif layout.num_condition_audio_rows:
             raise ValueError("A reference-prefixed layout needs its encoded audio condition rows.")
         text = mx.array(text_rows.astype(np.float32))
-        # The token refiner depends only on this request's conditioning, not
-        # the changing latents or timesteps. A local value also prevents a
-        # resident DiT from reusing embeddings from an earlier request.
-        text_embeds = None
         dit_forward_s = 0.0
         prefix_kv = None
         if reuse_prefix_kv:
@@ -1002,9 +998,6 @@ class MiniMaxH3MLXPipeline:
                 if record_step:
                     pdd_forward_kwargs["record_prefix_kv"] = prefix_kv
                     recorded_steps.append(step_index)
-                if text_embeds is None:
-                    text_embeds = dit.refine_text(text)
-                    mx.eval(text_embeds)
                 video_velocity, audio_velocity = dit.forward_with_cache(
                     x_v,
                     x_a,
@@ -1013,7 +1006,6 @@ class MiniMaxH3MLXPipeline:
                     step_timesteps=unique,
                     row_timestep_inverse=inverse,
                     step_index=step_index,
-                    precomputed_text_embeds=text_embeds,
                     **pdd_forward_kwargs,
                 )
                 video_velocity = video_velocity[layout.num_condition_video_rows:]
@@ -1334,9 +1326,6 @@ class MiniMaxH3MLXPipeline:
         timings: dict[str, float] = {}
         peaks: dict[str, float] = {}
 
-        if vsa_impl in ("sdpa", "nax") and reuse_prefix_kv:
-            raise ValueError(
-                f"{vsa_impl.upper()} VSA does not support prefix K/V reuse; disable reuse or select another backend")
         if fast_sharpen < 0:
             raise ValueError(f"fast_sharpen must be non-negative, got {fast_sharpen}.")
         if inter_step_cooldown_s < 0:

fastvideo/mlx_runtime/minimax_h3_vsa.py

@@ -23,10 +23,6 @@ Execution backends:
 - **simd** — flash-style SIMD-group 8x8 matrix operations over the reference tile
   map for tile sizes 64 and 128 and head dimension 128. Unsupported shapes and kernel
   failures fall back to the reference backend.
-- **sdpa** — opt-in fused SDPA with the same tile mask and shared K/V. Intended
-  for M5; visits masked keys but can use faster matrix operations. ``auto`` is unchanged.
-- **nax** — opt-in M5 TensorOps kernel that visits only selected blocks and valid
-  key rows. Uses the installed MLX NAX helpers; ``auto`` is unchanged.
 
 Dense fused SDPA remains the default when VSA is disabled, the geometry is
 unsupported, or a dense-only checkpoint is loaded.
@@ -54,7 +50,7 @@ VSA_H3_TILE_SHAPES: dict[int, tuple[int, int, int]] = {
 VSA_REF_POLICIES = ("p2_multi_region", )
 VSA_GATE_KEY_SUFFIX = "attn.to_gate_compress.weight"
 VSA_PREFIX_MODES = ("exempt", "compete")
-PUBLIC_VSA_IMPLS = ("auto", "reference", "simd", "sdpa", "nax")
+PUBLIC_VSA_IMPLS = ("auto", "reference", "simd")
 VSA_IMPLS = PUBLIC_VSA_IMPLS
 # Tile sizes where ``auto`` picks the SIMD kernel (trained Omni tiles); other tiles stay opt-in.
 _AUTO_SIMD_TILES: tuple[int, ...] = (128, )
@@ -62,7 +58,7 @@ _REFERENCE_FULL_MASK_TILE_LIMIT = 24
 _REFERENCE_FULL_MASK_MAX_ELEMENTS = 256 * 1024**2
 
 PrefixMode = Literal["exempt", "compete"]
-VSAImpl = Literal["auto", "reference", "simd", "sdpa", "nax"]
+VSAImpl = Literal["auto", "reference", "simd"]
 
 
 def validate_ref_policy(ref_policy: str | None, ref_keep_rate: float | None) -> None:
@@ -877,20 +873,6 @@ def _gate_compress_output(scores, v_tiled, gate_tiled, geometry: MiniMaxH3VSAGeo
 
 
 def resolve_impl(requested: VSAImpl, head_dim: int, tile_elems: int) -> str:
-    if requested == "nax":
-        from fastvideo.mlx_runtime.minimax_h3_vsa_nax import NAX_HEAD_DIM, NAX_TILE_SIZES
-
-        if tile_elems not in NAX_TILE_SIZES or head_dim != NAX_HEAD_DIM:
-            raise ValueError("NAX VSA requires tile 64 or 128 and head dimension 128 "
-                             f"(got tile={tile_elems}, head_dim={head_dim})")
-        return "nax"
-    if requested == "sdpa":
-        from fastvideo.mlx_runtime.minimax_h3_vsa_sdpa import SDPA_HEAD_DIM, SDPA_TILE_SIZES
-
-        if tile_elems not in SDPA_TILE_SIZES or head_dim != SDPA_HEAD_DIM:
-            raise ValueError("SDPA VSA requires tile 64 or 128 and head dimension 128 "
-                             f"(got tile={tile_elems}, head_dim={head_dim})")
-        return "sdpa"
     if requested == "reference" or (requested == "auto" and tile_elems not in _AUTO_SIMD_TILES):
         return "reference"
 
@@ -1005,16 +987,6 @@ def h3_vsa_attention(
                 if stats is not None:
                     stats.dense_fallback_reason = f"simd kernel failed: {error}"
                 video_tiled = _reference_gather_sdpa(q_tiled, k_tiled, v_tiled, block_idx, geometry, scale)
-        elif chosen == "nax":
-            from fastvideo.mlx_runtime.minimax_h3_vsa_nax import nax_block_sparse
-
-            video_tiled = nax_block_sparse(q_tiled, k_tiled, v_tiled, block_idx, block_num, geometry,
-                                           scale)[n_prefix_pad:]
-        elif chosen == "sdpa":
-            from fastvideo.mlx_runtime.minimax_h3_vsa_sdpa import sdpa_block_sparse
-
-            video_tiled = sdpa_block_sparse(q_tiled, k_tiled, v_tiled, block_idx, block_num, geometry,
-                                            scale)[n_prefix_pad:]
         elif _reference_full_mask_fits(geometry, heads):
             scores_np = np.array(scores, dtype=np.float32)
             mask = build_block_mask(

fastvideo/mlx_runtime/minimax_h3_vsa_nax.py

@@ -1,388 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0 AND MIT
-# mypy: disable-error-code=no-untyped-call
-#
-# This file contains a kernel adapted from MLX. Its MIT notice follows.
-# MIT License
-#
-# Copyright © 2023 Apple Inc.
-#
-# 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 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.
-"""Opt-in M5 block-sparse TensorOps attention using installed MLX NAX helpers.
-
-Only selected blocks and their valid rows are visited, without K/V duplication
-or a dense mask. Prefix and padded query outputs remain zero. Online softmax
-and matrix accumulation can differ numerically from SIMD and fused SDPA.
-
-This backend depends on private headers from the installed MLX package, tested
-with MLX 0.32.4 and the Metal 4 SDK. Only the validated header hashes are accepted.
-Missing headers, incompatible APIs and
-unsupported hardware raise explicitly; no fallback changes the implementation.
-Importing this module performs no MLX operations. Header copyrights are retained
-and get_nax_provenance() provides hashes for reproducible benchmark receipts.
-"""
-from __future__ import annotations
-
-import hashlib
-import math
-import re
-from functools import lru_cache
-from pathlib import Path
-
-NAX_TILE_SIZES = (64, 128)
-NAX_HEAD_DIM = 128
-DEFAULT_QUERY_ROWS = 32
-_NAX_HEADER = "mlx/backend/metal/kernels/steel/attn/nax.h"
-# Private fragment layouts can change without renaming their APIs. Accept only
-# the headers used by the numerical and full-generation validation on MLX 0.32.4.
-_VALIDATED_HEADER_SHA256 = {
-    _NAX_HEADER:
-    "fb10bcae44095e9fed5c81213d03527dd5fe5ac11be11ec7784bd357c6f317ec",
-    "mlx/backend/metal/kernels/steel/defines.h":
-    "b03cea6a7d5cfe814e6838d8536e2aacdcafd2744cb408e954fda336ce21759a",
-    "mlx/backend/metal/kernels/steel/utils/integral_constant.h":
-    "886eccaac4f8f00eab5a10551fdb5315df7ab7462aba13d69bb95765c1a4735e",
-    "mlx/backend/metal/kernels/steel/utils/type_traits.h":
-    "4766c5a3809d3cf7740e63acc5703ef76b6b4e85d8baf16602c9d517a2fd426c",
-}
-
-_OPERATORS = r"""
-namespace h3_sparse_nax {
-struct MaxOp {
-  template <typename V>
-  METAL_FUNC static constexpr V apply(V x, V y) { return metal::max(x, y); }
-};
-struct SumOp {
-  template <typename V>
-  METAL_FUNC static constexpr V apply(V x, V y) { return x + y; }
-};
-struct MulOp {
-  template <typename V>
-  METAL_FUNC static constexpr V apply(V x, V y) { return x * y; }
-};
-struct ExpSubOp {
-  template <typename V>
-  METAL_FUNC static constexpr V apply(V x, V y) { return metal::fast::exp2(x - y); }
-};
-}
-"""
-
-# Adapted from MLX steel_attention_nax.h (Copyright 2024-26 Apple Inc.).
-# The fragment math/online softmax follows that implementation; this launch
-# addresses the existing H3 tiled layout and walks selected key blocks directly.
-_SOURCE = r"""
-    using namespace mlx::steel;
-    using namespace h3_sparse_nax;
-    constexpr int D = 128;
-    constexpr int BK = 32;
-    constexpr int TK = BK / 16;
-    constexpr int TD = D / 16;
-    constexpr float NEG = -3.402823466e38f;
-
-    const int qsub = int(threadgroup_position_in_grid.x);
-    const int query_tile = int(threadgroup_position_in_grid.y);
-    const int head = int(threadgroup_position_in_grid.z);
-    const int sid = int(simdgroup_index_in_threadgroup);
-    const int k_max = meta_i[0];
-    const int n_prefix = meta_i[1];
-    const int n_video = meta_i[2];
-    const int heads = meta_i[3];
-    const int qt = n_prefix + query_tile;
-    const int q_valid = vbs[qt];
-    const int row_base = qsub * BQ;
-    const int nsel = block_num[head * n_video + query_tile];
-
-    // Uniform across the threadgroup: later barriers cannot be skipped by only
-    // some SIMD groups. Inactive SIMD groups use guarded loads and stores.
-    if (row_base >= q_valid || nsel <= 0) { return; }
-
-    const int row_stride = heads * D;
-    const int qrow = row_base + sid * 16;
-    const short q_remaining = short(q_valid - qrow);
-    const device T* Q = q + (long(qt * TILE + qrow) * heads + head) * D;
-    device T* O = out + (long(qt * TILE + qrow) * heads + head) * D;
-    const float scale2 = meta_f[0] * 1.44269504088896341f;
-
-    using otile_t = NAXTile<float, 1, TD>;
-    using stile_t = NAXTile<float, 1, TK>;
-    otile_t Otile;
-    Otile.clear();
-    constexpr short kRowsPT = otile_t::kRowsPerThread;
-    metal::vec<float, kRowsPT> max_score;
-    metal::vec<float, kRowsPT> sum_score{0};
-    STEEL_PRAGMA_UNROLL
-    for (short i = 0; i < kRowsPT; ++i) { max_score[i] = NEG; }
-    const short sn = otile_t::NAXFrag_t::get_coord().x;
-
-    for (int selected = 0; selected < nsel; ++selected) {
-      const int kt = block_idx[(head * n_video + query_tile) * k_max + selected];
-      const int k_valid = vbs[kt];
-      const long kv_offset = (long(kt) * TILE * heads + head) * D;
-      for (int krow = 0; krow < k_valid; krow += BK) {
-        const short k_remaining = short(metal::min(BK, k_valid - krow));
-        const device T* K = k + kv_offset + long(krow) * row_stride;
-        const device T* V = v + kv_offset + long(krow) * row_stride;
-        stile_t Stile;
-        Stile.clear();
-
-        // QK: retain MLX's four-way partial unroll for the 128-wide head.
-        STEEL_PRAGMA_UNROLL
-        for (short ik = 0; ik < TK; ik += 2) {
-#pragma clang loop unroll_count(4)
-          for (short id = 0; id < TD; ++id) {
-            NAXTile<T, 1, 1> Qtile;
-            NAXTile<T, 2, 1> Ktile;
-            if (q_remaining < 16) {
-              Qtile.load_rows(Q + id * 16, row_stride, q_remaining);
-            } else {
-              Qtile.load(Q + id * 16, row_stride);
-            }
-            if (k_remaining < BK) {
-              Ktile.load_rows(K + ik * 16 * row_stride + id * 16,
-                              row_stride, k_remaining - ik * 16);
-            } else {
-              Ktile.load(K + ik * 16 * row_stride + id * 16, row_stride);
-            }
-            stile_t::NAXFrag_t::mma(
-                Stile.frag_at(0, ik), Stile.frag_at(0, ik + 1),
-                Qtile.frag_at(0, 0), metal::false_type{},
-                Ktile.frag_at(0, 0), Ktile.frag_at(1, 0), metal::true_type{});
-          }
-        }
-        STEEL_PRAGMA_UNROLL
-        for (short i = 0; i < stile_t::kElemsPerTile; ++i) {
-          Stile.elems()[i] *= scale2;
-        }
-        if (k_remaining < BK) {
-          STEEL_PRAGMA_UNROLL
-          for (short ik = 0; ik < TK; ++ik) {
-            thread auto& frag = Stile.frag_at(0, ik);
-            STEEL_PRAGMA_UNROLL
-            for (short i = 0; i < stile_t::kFragThrRows; ++i) {
-              STEEL_PRAGMA_UNROLL
-              for (short j = 0; j < stile_t::kFragThrCols; ++j) {
-                if (ik * 16 + sn + j >= k_remaining) {
-                  frag[i * stile_t::kFragThrCols + j] = NEG;
-                }
-              }
-            }
-          }
-        }
-
-        metal::vec<float, kRowsPT> new_max = max_score;
-        metal::vec<float, kRowsPT> factor;
-        Stile.template row_reduce<MaxOp>(new_max);
-        Stile.template row_bin_op<ExpSubOp>(new_max);
-        STEEL_PRAGMA_UNROLL
-        for (short i = 0; i < kRowsPT; ++i) {
-          factor[i] = metal::fast::exp2(max_score[i] - new_max[i]);
-          max_score[i] = new_max[i];
-          sum_score[i] *= factor[i];
-        }
-        Stile.template row_reduce<SumOp>(sum_score);
-        Otile.template row_bin_op<MulOp>(factor);
-        simdgroup_barrier(mem_flags::mem_none);
-
-        // PV retains float32 probabilities, matching MLX's NAX SDPA kernel.
-        STEEL_PRAGMA_UNROLL
-        for (short id = 0; id < TD; id += 2) {
-          if (id == 4) { threadgroup_barrier(mem_flags::mem_none); }
-          STEEL_PRAGMA_UNROLL
-          for (short ik = 0; ik < TK; ++ik) {
-            NAXTile<T, 1, 2> Vtile;
-            if (k_remaining < BK) {
-              Vtile.load_rows(V + ik * 16 * row_stride + id * 16,
-                              row_stride, k_remaining - ik * 16);
-            } else {
-              Vtile.load(V + ik * 16 * row_stride + id * 16, row_stride);
-            }
-            otile_t::NAXFrag_t::mma(
-                Otile.frag_at(0, id), Otile.frag_at(0, id + 1),
-                Stile.frag_at(0, ik), metal::false_type{},
-                Vtile.frag_at(0, 0), Vtile.frag_at(0, 1), metal::false_type{});
-          }
-        }
-      }
-    }
-    threadgroup_barrier(mem_flags::mem_none);
-    metal::vec<float, kRowsPT> rcp;
-    STEEL_PRAGMA_UNROLL
-    for (short i = 0; i < kRowsPT; ++i) {
-      rcp[i] = sum_score[i] > 0.f ? 1.f / sum_score[i] : 0.f;
-    }
-    Otile.template row_bin_op<MulOp>(rcp);
-    if (q_remaining > 0) {
-      if (q_remaining < 16) { Otile.store_rows(O, row_stride, q_remaining); }
-      else { Otile.store(O, row_stride); }
-    }
-"""
-
-
-def _mlx_include_root() -> Path:
-    """Resolve headers beside the imported package, including namespace installs."""
-    import mlx
-
-    for package_path in mlx.__path__:
-        include_root = Path(package_path) / "include"
-        if (include_root / _NAX_HEADER).is_file():
-            return include_root
-    raise RuntimeError("NAX VSA requires installed MLX NAX headers (tested with MLX 0.32.4); "
-                       "the mlx package has no include/mlx/backend/metal/kernels/steel/attn/nax.h")
-
-
-@lru_cache(maxsize=1)
-def _header():
-    """Inline quoted MLX includes; system Metal/MPP includes stay intact."""
-    include_root = _mlx_include_root()
-    visited = {}
-
-    def expand(relative):
-        if relative in visited:
-            return f"// Already included: {relative}\n"
-        path = include_root / relative
-        try:
-            original_bytes = path.read_bytes()
-        except OSError:
-            raise RuntimeError(f"NAX VSA cannot read required installed MLX header: {relative}") from None
-        digest = hashlib.sha256(original_bytes).hexdigest()
-        if digest != _VALIDATED_HEADER_SHA256.get(relative):
-            raise RuntimeError(f"Unvalidated MLX private header: {relative}; NAX VSA requires the headers "
-                               "validated with MLX 0.32.4. Revalidate this backend before accepting new headers.")
-        original = original_bytes.decode("utf-8")
-        visited[relative] = digest
-        content = re.sub(r"^\s*#pragma once\s*$", "", original, flags=re.MULTILINE)
-        content = re.sub(r'^\s*#include "(mlx/[^"\n]+)"\s*$',
-                         lambda match: expand(match.group(1)),
-                         content,
-                         flags=re.MULTILINE)
-        if relative == _NAX_HEADER:
-            old = "transpose_b,\n        true,\n        mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate"
-            expected_api = ("struct BaseNAXFrag", "struct NAXTile", "get_coord()", "row_reduce(", "row_bin_op(",
-                            "load_rows(", "store_rows(")
-            if content.count(old) != 2 or any(marker not in content for marker in expected_api):
-                raise RuntimeError("Installed MLX NAX private API changed; this backend was tested with MLX 0.32.4 "
-                                   "and needs compatibility review")
-        return f"\n// Begin installed MLX header: {relative}\n{content}\n// End: {relative}\n"
-
-    header = "#ifndef METAL_FUNC\n#define METAL_FUNC inline __attribute__((always_inline))\n#endif\n"
-    header += expand(_NAX_HEADER)
-    header += _OPERATORS
-    return header, visited
-
-
-def get_nax_provenance():
-    """CPU-only metadata, containing no machine-specific absolute paths."""
-    header, dependencies = _header()
-    return {
-        "kernel": "h3-true-sparse-nax-v1",
-        "relaxed_precision": True,
-        "math_mode": "safe",
-        "key_rows": 32,
-        "source_sha256": hashlib.sha256(_SOURCE.encode()).hexdigest(),
-        "expanded_header_sha256": hashlib.sha256(header.encode()).hexdigest(),
-        "installed_header_sha256": dependencies
-    }
-
-
-def _check_nax_device() -> None:
-    """The TensorOps backend is deliberately restricted to validated M5 GPUs."""
-    import mlx.core as mx
-
-    if not mx.metal.is_available():
-        raise RuntimeError("NAX VSA requires an Apple M5 GPU and Metal 4")
-    device_name = str(mx.device_info().get("device_name", ""))
-    if re.match(r"^Apple M5(?:\s|$)", device_name) is None:
-        raise RuntimeError(f"NAX VSA supports Apple M5 GPUs only, got {device_name or 'unknown GPU'}")
-
-
-@lru_cache(maxsize=1)
-def _nax_kernel():
-    import mlx.core as mx
-
-    _check_nax_device()
-    header, _ = _header()
-    if not hasattr(mx.fast, "metal_kernel"):
-        raise RuntimeError("NAX VSA requires mx.fast.metal_kernel with safe math compile_options")
-    try:
-        return mx.fast.metal_kernel(name="h3_vsa_sparse_nax",
-                                    input_names=["q", "k", "v", "block_idx", "block_num", "vbs", "meta_i", "meta_f"],
-                                    output_names=["out"],
-                                    source=_SOURCE,
-                                    header=header,
-                                    compile_options={"math_mode": "safe"})
-    except (TypeError, ValueError, RuntimeError) as error:
-        raise RuntimeError("NAX VSA kernel construction failed; this backend requires compatible MLX NAX headers "
-                           "and Metal 4 TensorOps support") from error
-
-
-def nax_block_sparse(q_tiled,
-                     k_tiled,
-                     v_tiled,
-                     block_idx,
-                     block_num,
-                     geometry,
-                     scale: float,
-                     *,
-                     query_rows: int = DEFAULT_QUERY_ROWS):
-    """Same [S_pad,H,128] contract as the existing H3 sparse implementations.
-
-    query_rows controls SIMD groups per threadgroup.
-    Active indices must be valid and unique, as emitted by the H3 router. Values
-    in inactive index slots are never read. Prefix/padding outputs are zero.
-    """
-    import mlx.core as mx
-
-    if query_rows not in (16, 32, 64, 128):
-        raise ValueError("query_rows must be 16, 32, 64 or 128")
-    if q_tiled.ndim != 3 or q_tiled.shape != k_tiled.shape or q_tiled.shape != v_tiled.shape:
-        raise ValueError("NAX sparse attention requires matching [S_pad,H,D] Q/K/V")
-    padded, heads, dim = q_tiled.shape
-    tile = geometry.tile_elems
-    if tile not in NAX_TILE_SIZES or dim != NAX_HEAD_DIM or query_rows > tile:
-        raise ValueError("NAX sparse attention requires tile 64/128, head dim 128, query_rows <= tile")
-    if padded != geometry.padded_length or heads < 1:
-        raise ValueError("Q/K/V shape does not match geometry")
-    if q_tiled.dtype not in (mx.float32, mx.float16, mx.bfloat16):
-        raise ValueError("NAX sparse attention supports float32, float16 or bfloat16")
-    if q_tiled.dtype != k_tiled.dtype or q_tiled.dtype != v_tiled.dtype:
-        raise ValueError("Q/K/V dtypes must match")
-    if (block_idx.ndim != 3 or block_idx.shape[:2] != (heads, geometry.num_video_tiles)
-            or block_num.shape != (heads, geometry.num_video_tiles)):
-        raise ValueError("Routing shapes do not match head/tile geometry")
-    if not math.isfinite(scale):
-        raise ValueError("scale must be finite")
-    if geometry.num_video_tiles == 0:
-        return mx.zeros_like(q_tiled)
-    threads = (query_rows // 16) * 32
-    result, = _nax_kernel()(inputs=[
-        mx.contiguous(q_tiled),
-        mx.contiguous(k_tiled),
-        mx.contiguous(v_tiled),
-        mx.contiguous(block_idx.astype(mx.int32)),
-        mx.contiguous(block_num.astype(mx.int32)),
-        mx.array(geometry.variable_block_sizes, dtype=mx.int32),
-        mx.array([block_idx.shape[-1], geometry.num_prefix_tiles, geometry.num_video_tiles, heads], dtype=mx.int32),
-        mx.array([scale], dtype=mx.float32)
-    ],
-                            template=[("T", q_tiled.dtype), ("TILE", tile), ("BQ", query_rows)],
-                            grid=(threads * (tile // query_rows), geometry.num_video_tiles, heads),
-                            threadgroup=(threads, 1, 1),
-                            output_shapes=[q_tiled.shape],
-                            output_dtypes=[q_tiled.dtype],
-                            init_value=0)
-    return result

fastvideo/mlx_runtime/minimax_h3_vsa_sdpa.py

@@ -1,97 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# mypy: disable-error-code=no-untyped-call
-"""VSA tile masks through fused SDPA, sharing K/V with grouped-query attention.
-
-This opt-in backend preserves selected blocks, ragged-row masks and the trained
-compression branch. Unlike the SIMD kernel, fused SDPA visits masked key blocks;
-on M5 its faster matrix operations can outweigh that additional work. Floating-
-point accumulation differs, so this is not a bit-identical implementation.
-"""
-
-from __future__ import annotations
-
-SDPA_TILE_SIZES = (64, 128)
-SDPA_HEAD_DIM = 128
-DEFAULT_QUERY_CHUNK = 128
-_MASK_CHUNK_MAX_BYTES = 256 * 1024**2
-
-
-def sdpa_block_sparse(q_tiled,
-                      k_tiled,
-                      v_tiled,
-                      block_idx,
-                      block_num,
-                      geometry,
-                      scale: float,
-                      *,
-                      chunk_size: int = DEFAULT_QUERY_CHUNK):
-    """Return tiled [S_pad,H,D], with zero prefix rows and zero padded queries.
-
-    Query heads are ordered (original head, query tile), so GQA maps each group
-    of query tiles to its original K/V head. A tile's key mask broadcasts across
-    its query rows. At most two mask chunks are in flight. Unsupported fused
-    operations raise instead of silently allocating unfused attention scores.
-    """
-    import mlx.core as mx
-
-    if chunk_size < 1:
-        raise ValueError("SDPA VSA chunk_size must be positive")
-    if q_tiled.ndim != 3 or k_tiled.shape != q_tiled.shape or v_tiled.shape != q_tiled.shape:
-        raise ValueError("SDPA VSA requires matching [S_pad,H,D] Q/K/V shapes")
-    padded, heads, dim = q_tiled.shape
-    tile = geometry.tile_elems
-    if tile not in SDPA_TILE_SIZES or dim != SDPA_HEAD_DIM:
-        raise ValueError(f"SDPA VSA requires tile 64 or 128 and head dimension 128, got tile={tile}, dim={dim}")
-    if padded != geometry.padded_length or heads < 1:
-        raise ValueError("SDPA VSA Q/K/V shape does not match the tile geometry")
-    if q_tiled.dtype != k_tiled.dtype or q_tiled.dtype != v_tiled.dtype:
-        raise ValueError("SDPA VSA requires matching Q/K/V dtypes")
-    n_tiles = geometry.num_tiles
-    n_video = geometry.num_video_tiles
-    n_prefix = geometry.num_prefix_tiles
-    if (block_idx.ndim != 3 or block_idx.shape[:2] != (heads, n_video) or block_num.shape != (heads, n_video)):
-        raise ValueError("SDPA VSA routing shape does not match the query head/tile geometry")
-
-    query_chunk = min(chunk_size, max(1, _MASK_CHUNK_MAX_BYTES // max(heads * padded, 1)))
-    q_video = q_tiled.reshape(n_tiles, tile, heads, dim).transpose(2, 0, 1, 3)[:, n_prefix:]
-    # Fused full attention accepts these strided views: the D axis is contiguous.
-    k_shared = k_tiled.transpose(1, 0, 2)[None]
-    v_shared = v_tiled.transpose(1, 0, 2)[None]
-    sizes = mx.array(geometry.variable_block_sizes, dtype=mx.int32)
-    row = mx.arange(tile, dtype=mx.int32)
-    key_valid = row[None, :] < sizes[:, None]
-    head_index = mx.arange(heads, dtype=mx.int32)[:, None, None]
-    selected_slot = mx.arange(block_idx.shape[-1], dtype=mx.int32)[None, None, :]
-    outputs = []
-    for start in range(0, n_video, query_chunk):
-        end = min(start + query_chunk, n_video)
-        count = end - start
-        indices = block_idx[:, start:end].astype(mx.int32)
-        counts = block_num[:, start:end]
-        # Send inactive slots to a sentinel column that is discarded. Every
-        # scatter writes True, so duplicate indices cannot erase a valid key.
-        indices = mx.where(selected_slot < counts[..., None], indices, n_tiles)
-        keep = mx.zeros((heads, count, n_tiles + 1), dtype=mx.bool_)
-        tile_index = mx.arange(count, dtype=mx.int32)[None, :, None]
-        keep[head_index, tile_index, indices] = True
-        mask = (keep[:, :, :n_tiles, None] & key_valid[None, None]).reshape(1, heads * count, 1, padded)
-        query = mx.contiguous(q_video[:, start:end].reshape(1, heads * count, tile, dim))
-        output = mx.fast.scaled_dot_product_attention(query,
-                                                      k_shared,
-                                                      v_shared,
-                                                      scale=scale,
-                                                      mask=mask,
-                                                      force_fused=True)
-        output = output.reshape(heads, count, tile, dim)
-        valid_query = row[None, :] < sizes[n_prefix + start:n_prefix + end, None]
-        output = mx.where(valid_query[None, :, :, None] & (counts[..., None, None] > 0), output, 0)
-        mx.async_eval(output)
-        outputs.append(output)
-        if len(outputs) > 1:
-            mx.eval(outputs[-2])
-    if not outputs:
-        return mx.zeros_like(q_tiled)
-    mx.eval(outputs[-1])
-    video = mx.concatenate(outputs, axis=1) if len(outputs) > 1 else outputs[0]
-    video = video.transpose(1, 2, 0, 3).reshape(n_video * tile, heads, dim)
-    return mx.concatenate((mx.zeros((n_prefix * tile, heads, dim), dtype=q_tiled.dtype), video), axis=0)

fastvideo/mlx_runtime/minimax_h3_vsa_target.py

@@ -43,9 +43,6 @@ def h3_vsa_attention_target(
     ref_keep_rate: float | None = None,
 ):
     """query/gate_compress: [T, H, D] target rows; key/value: [S, H, D] packed. Returns [T, H, D]."""
-    if impl in ("sdpa", "nax"):
-        raise ValueError(
-            f"{impl.upper()} VSA does not support prefix K/V reuse; disable reuse or select another backend")
     import mlx.core as mx
 
     _, heads, dim = query.shape

fastvideo/mlx_runtime/tests/test_minimax_h3_text_refiner_cache.py

@@ -1,120 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Request-local text refinement must preserve denoising and refresh per request."""
-from types import SimpleNamespace
-
-import numpy as np
-import pytest
-
-mx = pytest.importorskip("mlx.core")
-
-from fastvideo.mlx_runtime import minimax_h3 as h3
-from fastvideo.mlx_runtime import minimax_h3_pipeline as pipeline_module
-
-
-def _toy_dit():
-    rng = np.random.default_rng(11)
-
-    def array(shape):
-        return mx.array(rng.normal(0, 0.05, shape).astype(np.float32))
-
-    hidden, ffn = 16, 24
-    config = dict(hidden_size=hidden, num_attention_heads=1, attention_head_dim=hidden,
-                  ffn_dim=ffn, in_channels=4, audio_in_channels=4, patch_size=[1, 1, 1],
-                  text_dim=8, freq_dim=hidden, time_embed_dim=hidden, rope_freq_dim=2,
-                  rope_theta=10000., norm_eps=1e-5, qk_norm_eps=1e-5, final_norm_eps=1e-5)
-    weights = {}
-    for name, source, target in (("context_embedder", 8, hidden), ("proj_in", 4, hidden),
-                                 ("audio_proj_in", 4, hidden), ("proj_out", hidden, 4),
-                                 ("audio_proj_out", hidden, 4),
-                                 ("time_embedder.linear_1", hidden, hidden),
-                                 ("time_embedder.linear_2", hidden, hidden),
-                                 ("norm_out.linear", hidden, 2 * hidden)):
-        weights[name + ".weight"] = array((target, source))
-        weights[name + ".bias"] = array((target,))
-    weights["token_refiner.final_norm.weight"] = mx.ones((hidden,))
-    weights["norm_out.norm.weight"] = mx.ones((hidden,))
-
-    def block():
-        result = {name: mx.ones((hidden,)) for name in
-                  ("norm1.weight", "norm2.weight", "attn.norm_q.weight", "attn.norm_k.weight")}
-        for name in ("attn.to_q.weight", "attn.to_k.weight", "attn.to_v.weight", "attn.to_out.0.weight"):
-            result[name] = array((hidden, hidden))
-        result["ff.net.0.proj.weight"] = array((2 * ffn, hidden))
-        result["ff.net.2.weight"] = array((hidden, ffn))
-        result["adaln_proj.linear.weight"] = array((18 * hidden, hidden))
-        result["adaln_proj.linear.bias"] = array((18 * hidden,))
-        return result
-
-    return h3.MLXMiniMaxH3DiT(weights, [block()], [block()], config)
-
-
-def test_precomputed_refiner_is_exact_and_runs_once_for_eight_steps(monkeypatch):
-    dit = _toy_dit()
-    layout = h3.build_packed_layout(3, 1, 2, 2, 2, patch_size=dit.patch_size)
-    timesteps, inverse = h3.build_row_timesteps(layout, .7, .4)
-    dit.precompute_adaln(timesteps)
-    text = mx.array(np.arange(24, dtype=np.float32).reshape(3, 8) / 24)
-    video = mx.ones((len(layout.video_indices), 4))
-    audio = mx.ones((len(layout.audio_indices), 4))
-    calls = []
-    original = dit.refine_text
-
-    def counted(rows):
-        calls.append(rows)
-        return original(rows)
-
-    monkeypatch.setattr(dit, "refine_text", counted)
-    kwargs = dict(layout=layout, step_timesteps=timesteps, row_timestep_inverse=inverse)
-    expected = [dit.forward_with_cache(video * (step + 1), audio, text, **kwargs) for step in range(8)]
-    mx.eval(expected)
-    assert len(calls) == 8
-    calls.clear()
-    embeds = dit.refine_text(text)
-    mx.eval(embeds)
-    actual = [dit.forward_with_cache(video * (step + 1), audio, text,
-                                    precomputed_text_embeds=embeds, **kwargs) for step in range(8)]
-    mx.eval(actual)
-    assert len(calls) == 1
-    for reference, cached in zip(expected, actual, strict=True):
-        for left, right in zip(reference, cached, strict=True):
-            np.testing.assert_array_equal(np.asarray(left), np.asarray(right))
-    with pytest.raises(ValueError, match="shape"):
-        dit.forward_with_cache(video, audio, text, precomputed_text_embeds=embeds[:1], **kwargs)
-    with pytest.raises(ValueError, match="floating-point"):
-        dit.forward_with_cache(video, audio, text, precomputed_text_embeds=embeds.astype(mx.int32), **kwargs)
-
-
-def test_pipeline_refines_again_for_next_request(monkeypatch, tmp_path):
-    dit = _toy_dit()
-    video_schedule = h3.MiniMaxH3SchedulerState.create(10, 8)
-    audio_schedule = h3.MiniMaxH3SchedulerState.create(3, 8)
-    union = h3.adaln_timestep_union(video_schedule, audio_schedule)
-    schedule = SimpleNamespace(video=video_schedule, audio=audio_schedule, adaln_timesteps=union,
-                               num_steps=8, pdd_plan=None)
-    monkeypatch.setattr(pipeline_module, "resolve_h3_denoise_schedule", lambda *a: schedule)
-    monkeypatch.setattr(pipeline_module, "_cleanup_mlx", lambda: None)
-    pipeline = object.__new__(pipeline_module.MiniMaxH3MLXPipeline)
-    pipeline.resident = False
-    pipeline.model_root = tmp_path
-    pipeline.resolve_geometry = lambda *a, **k: dict(num_frames=124, latent_frame_count=1,
-                                                   latent_height=2, latent_width=2)
-    calls = []
-    original = dit.refine_text
-
-    def counted(rows):
-        calls.append(np.asarray(rows))
-        return original(rows)
-
-    monkeypatch.setattr(dit, "refine_text", counted)
-    first_text = np.arange(24, dtype=np.float32).reshape(3, 8) / 24
-    second_text = first_text[:, ::-1].copy()
-    kwargs = dict(height=32, width=32, num_frames=124, seed=0, num_steps=8, dit=dit)
-    first = pipeline.denoise(first_text, np.ones(3, dtype=np.int64), **kwargs)
-    second = pipeline.denoise(second_text, np.ones(3, dtype=np.int64), **kwargs)
-    repeated = pipeline.denoise(first_text, np.ones(3, dtype=np.int64), **kwargs)
-    assert len(calls) == 3
-    np.testing.assert_array_equal(calls[0], first_text)
-    np.testing.assert_array_equal(calls[1], second_text)
-    assert not np.array_equal(first[0], second[0])
-    np.testing.assert_array_equal(first[0], repeated[0])
-    np.testing.assert_array_equal(first[1], repeated[1])

fastvideo/mlx_runtime/tests/test_minimax_h3_vsa_gather_budget.py

@@ -50,10 +50,5 @@ def test_compiled_rotary_matches_uncompiled_fp32_math():
     rotated = mx.concatenate([-second, first], axis=-1)
     out = (x_rot * cos[:, None, :] + rotated * sin[:, None, :]).astype(x.dtype)
     expected = mx.concatenate([out, x[..., rotary_dim:]], axis=-1)
-    # The compiled kernel may fuse the multiply-add (FMA) where the eager graph rounds each op,
-    # which can move a result by one bf16 ULP on some CPUs. Allow exactly that: rtol 2**-7 is one
-    # bf16 ULP at the bottom of each binade, so a real math error still fails.
-    np.testing.assert_allclose(np.array(apply_h3_rotary(x, cos, sin).astype(mx.float32)),
-                               np.array(expected.astype(mx.float32)),
-                               rtol=2**-7,
-                               atol=0)
+    np.testing.assert_array_equal(np.array(apply_h3_rotary(x, cos, sin).astype(mx.float32)),
+                                  np.array(expected.astype(mx.float32)))

fastvideo/models/dits/minimax_h3.py

@@ -29,6 +29,7 @@ from fastvideo.distributed.parallel_state import (get_sp_group, get_sp_world_siz
 from fastvideo.forward_context import get_forward_context
 from fastvideo.layers.linear import ReplicatedLinear
 from fastvideo.layers.quantization.mxfp8_config import MXFP8QuantizeMethod
+from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod
 from fastvideo.layers.mlp import MLP
 from fastvideo.layers.pdd import PDDReplicatedLinear, fuse_pdd_heads
 from fastvideo.layers.quantization import QuantizationConfig
@@ -163,7 +164,10 @@ class MiniMaxH3FeedForward(nn.Module):
         # (4.5 GiB at 78k tokens), so chunks bound the activation peak on 24-32 GB GPUs.
         self.chunk_tokens = envs.FASTVIDEO_H3_FFN_CHUNK_TOKENS.get()
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def forward(self, hidden_states: torch.Tensor,
+                pre_quantized: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None) -> torch.Tensor:
+        if pre_quantized is not None:
+            return self._forward(hidden_states, pre_quantized)
         tokens = hidden_states.shape[-2] if hidden_states.dim() > 1 else 0
         if (self.chunk_tokens and tokens > self.chunk_tokens and not torch.is_grad_enabled()
                 and not torch.compiler.is_compiling() and hidden_states.numel() == tokens * hidden_states.shape[-1]):
@@ -174,12 +178,29 @@ class MiniMaxH3FeedForward(nn.Module):
             return out
         return self._forward(hidden_states)
 
-    def _forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def _forward(self, hidden_states: torch.Tensor,
+                 pre_quantized: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None) -> torch.Tensor:
         if self.use_mxfp8:
             from fastvideo.layers.mxfp8linear import mxfp8_swiglu_feed_forward
 
             return mxfp8_swiglu_feed_forward(hidden_states, self.fc_in, self.fc_out)
-        hidden_states, _ = self.fc_in(hidden_states)
+        if self.fuse_swiglu and pre_quantized is not None and not torch.compiler.is_compiling():
+            from fastvideo.models.dits.minimax_h3_fusions.swiglu_gemm import try_swiglu_gemm
+
+            fused = try_swiglu_gemm(self, hidden_states, pre_quantized)
+            if fused is not None:
+                return fused
+        if pre_quantized is None:
+            hidden_states, _ = self.fc_in(hidden_states)
+        else:
+            hidden_states = self.fc_in.quant_method.apply(
+                self.fc_in, hidden_states, self.fc_in.bias, pre_quantized=pre_quantized,
+            )
+        if (self.fuse_swiglu and _can_run_minimax_h3_fusion(hidden_states)
+                and not torch.compiler.is_compiling()):
+            quantized = self._quantized_swiglu_output(hidden_states)
+            if quantized is not None:
+                return quantized
         if self.fuse_swiglu and _can_run_minimax_h3_fusion(hidden_states):
             hidden_states = minimax_h3_swiglu(hidden_states)
         elif self.exact_swiglu and exact_kernels.supports_rowwise(hidden_states):
@@ -190,6 +211,43 @@ class MiniMaxH3FeedForward(nn.Module):
         hidden_states, _ = self.fc_out(hidden_states)
         return hidden_states
 
+    def input_quantization_scale(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
+        """Static scale for a fused producer, or None when ordinary preparation is needed."""
+        method = getattr(self.fc_in, "quant_method", None)
+        if (self.chunk_tokens or type(method) is not NVFP4QuantizeMethod
+                or not method.wants_prequantized_input() or hidden_states.dtype != torch.bfloat16
+                or hidden_states.shape[-1] % 64 or torch.cuda.get_device_capability(hidden_states.device)[0] < 10):
+            return None
+        scale = method._static_activation_global_sf()
+        if scale is None:
+            scale = getattr(self.fc_in, "_nvfp4_input_global_sf", None)
+        if scale is None and not method._dynamic_activation_scale():
+            scale = method.x_global_sf
+        return scale
+
+    def _quantized_swiglu_output(self, packed: torch.Tensor) -> torch.Tensor | None:
+        """Fuse BF16 SwiGLU with compatible static/unit-scale NVFP4 preparation."""
+        layer = self.fc_out
+        method = getattr(layer, "quant_method", None)
+        if (type(method) is not NVFP4QuantizeMethod or not method.wants_prequantized_input()
+                or packed.dtype != torch.bfloat16 or packed.shape[-1] % 128 or not packed.is_contiguous()
+                or torch.cuda.get_device_capability(packed.device)[0] < 10):
+            return None
+        scale = method._static_activation_global_sf()
+        if scale is None:
+            scale = getattr(layer, "_nvfp4_input_global_sf", None)
+        if scale is None:
+            if method._dynamic_activation_scale():
+                return None
+            scale = method.x_global_sf
+        from fastvideo.models.dits.minimax_h3_fusions.swiglu_quant import swiglu_quantize_nvfp4
+
+        pre = swiglu_quantize_nvfp4(packed, scale)
+        # apply(pre_quantized=...) reads only this view's shape. The BF16
+        # SwiGLU activation itself is never allocated or written.
+        shape_view = packed[..., :packed.shape[-1] // 2]
+        return method.apply(layer, shape_view, layer.bias, pre_quantized=pre)
+
 
 class MiniMaxH3Attention(nn.Module):
     """Full self-attention over one sequence-parallel packed document."""
@@ -351,12 +409,30 @@ class MiniMaxH3Attention(nn.Module):
             return exact_kernels.rope_prefix(hidden_states, rotary_emb)
         return self._apply_rotary_emb(hidden_states, rotary_emb)
 
+    def input_quantization_scale(
+        self, hidden_states: torch.Tensor, rotary_emb: tuple[torch.Tensor, torch.Tensor] | None,
+    ) -> torch.Tensor | None:
+        """Allow producer quantization only when this call will use the shared FP4 route."""
+        fine_override = getattr(getattr(self.distributed_attention, "attn_impl", None), "fine_attention_override", None)
+        if (not self._vsa_fp4 or fine_override is not None or rotary_emb is None
+                or torch.is_grad_enabled() or torch.compiler.is_compiling()
+                or not hidden_states.is_cuda or hidden_states.dtype != torch.bfloat16
+                or torch.cuda.get_device_capability(hidden_states.device)[0] < 10):
+            return None
+        meta = get_forward_context().attn_metadata
+        if not isinstance(meta, MiniMaxH3VSAMetadata) or not meta.exempt or meta.pad_tokens:
+            return None
+        from fastvideo.models.dits.minimax_h3_vsa_fp4 import _shared_qkv_gate_scale
+
+        return _shared_qkv_gate_scale(self, hidden_states)
+
     def forward(
         self,
         hidden_states: torch.Tensor,
         rotary_emb: tuple[torch.Tensor, torch.Tensor] | None,
         original_seq_len: int,
         attn_metadata: AttentionMetadata | None = None,
+        pre_quantized: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
     ) -> torch.Tensor:
         fine_override = getattr(getattr(self.distributed_attention, "attn_impl", None), "fine_attention_override", None)
         if (self._vsa_fp4 and fine_override is None and rotary_emb is not None and not torch.is_grad_enabled()
@@ -367,13 +443,22 @@ class MiniMaxH3Attention(nn.Module):
             if isinstance(meta, MiniMaxH3VSAMetadata) and meta.exempt and not meta.pad_tokens:
                 use_fused_rope = self.fuse_qknorm_rope and _can_run_minimax_h3_fusion(hidden_states)
                 if not model_parallel_is_initialized() or get_sp_world_size() == 1:
-                    hidden_states = vsa_fp4_attention(self, hidden_states, rotary_emb, meta, use_fused_rope)
+                    hidden_states = vsa_fp4_attention(self, hidden_states, rotary_emb, meta, use_fused_rope,
+                                                       pre_quantized=pre_quantized)
                 else:
                     hidden_states = vsa_fp4_attention_sp(self, hidden_states, rotary_emb, meta, use_fused_rope,
-                                                         get_sp_group())
+                                                         get_sp_group(), pre_quantized=pre_quantized)
                 with STAGES.span("out_proj"):
-                    hidden_states, _ = self.to_out(hidden_states)
+                    if isinstance(hidden_states, tuple):
+                        shape_view, prepared = hidden_states
+                        hidden_states = self.to_out.quant_method.apply(
+                            self.to_out, shape_view, self.to_out.bias, pre_quantized=prepared,
+                        )
+                    else:
+                        hidden_states, _ = self.to_out(hidden_states)
                 return hidden_states
+        if pre_quantized is not None:
+            raise RuntimeError("Prequantized H3 attention input requires the eager FP4 VSA route")
         if (self._vsa_tile_first and hidden_states.is_cuda and rotary_emb is not None
                 and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
                 and (not model_parallel_is_initialized() or get_sp_world_size() == 1)):
@@ -548,7 +633,7 @@ class MiniMaxH3AdaLayerNormModulation(nn.Module):
             prefix=f"{prefix}.linear",
         )
 
-    def enable_host_cache(self, table: dict | None = None) -> None:
+    def enable_host_cache(self, table: dict | None = None, *, keep_resident: bool = False) -> None:
         """Keep the projection in pinned host memory and cache its output per timestep set.
 
         The modulation is a pure function of the timestep embedding, and few-step checkpoints sample a fixed
@@ -563,7 +648,11 @@ class MiniMaxH3AdaLayerNormModulation(nn.Module):
             self._cache_key = None
             return
         weight, bias = weight.data, bias.data
-        if table is None:
+        if table is None and keep_resident:
+            # Unified-memory GPUs share the host pool. Reuse the original allocation
+            # instead of making a pinned copy and copying it back on every cache miss.
+            self._host_weight, self._host_bias = weight, bias
+        elif table is None:
             self._host_weight = weight.to("cpu").pin_memory()
             self._host_bias = bias.to("cpu").pin_memory()
         else:
@@ -700,7 +789,10 @@ class MiniMaxH3TransformerBlock(nn.Module):
         adaln_indices: torch.Tensor,
         rotary_emb: tuple[torch.Tensor, torch.Tensor],
         original_seq_len: int,
-    ) -> torch.Tensor:
+        *,
+        _attention_prepared: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
+        _next_block: MiniMaxH3TransformerBlock | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
         with nvtx_range("minimax_h3.transformer_block.adaln_projection"):
             shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
                 t.to(hidden_states.dtype) for t in self.adaln_proj(temb))
@@ -710,7 +802,21 @@ class MiniMaxH3TransformerBlock(nn.Module):
         use_exact_modulate = (not use_modulate_fusion and self.exact_modulate and adaln_indices.is_contiguous()
                               and exact_kernels.supports_rowwise(hidden_states, shift_msa, scale_msa, gate_msa,
                                                                  shift_mlp, scale_mlp, gate_mlp))
-        if use_modulate_fusion:
+        attn_quantized = _attention_prepared
+        attn_scale = (self.attn.input_quantization_scale(hidden_states, rotary_emb)
+                      if attn_quantized is None and use_modulate_fusion and not torch.compiler.is_compiling() else None)
+        if attn_quantized is not None:
+            norm_hidden_states = hidden_states
+        elif attn_scale is not None:
+            from fastvideo.models.dits.minimax_h3_fusions.modulation_quant import modulate_quantize_nvfp4
+
+            with nvtx_range("minimax_h3.transformer_block.modulate_quant_fusion"):
+                attn_quantized = modulate_quantize_nvfp4(
+                    hidden_states, self.norm1.weight, scale_msa, shift_msa, adaln_indices,
+                    self.norm1.eps, attn_scale,
+                )
+            norm_hidden_states = hidden_states  # Shape only; QKV and the gate consume attn_quantized.
+        elif use_modulate_fusion:
             with nvtx_range("minimax_h3.transformer_block.modulate_fusion"):
                 norm_hidden_states = fused_rmsnorm_modulate(
                     hidden_states,
@@ -730,11 +836,26 @@ class MiniMaxH3TransformerBlock(nn.Module):
                 norm_hidden_states = norm_hidden_states * (
                     1.0 + scale_msa.index_select(0, adaln_indices)) + shift_msa.index_select(0, adaln_indices)
         with nvtx_range("minimax_h3.transformer_block.self_attention"):
-            attention_output = self.attn(norm_hidden_states, rotary_emb, original_seq_len)
+            attention_output = (self.attn(norm_hidden_states, rotary_emb, original_seq_len,
+                                          pre_quantized=attn_quantized) if attn_quantized is not None
+                                else self.attn(norm_hidden_states, rotary_emb, original_seq_len))
         # The attention input is dead now. Keeping it until assignment below
         # overlaps three full-width activations during the residual fusion.
         del norm_hidden_states
-        if use_modulate_fusion:
+        ff_quantized = None
+        ff_scale = (self.ff.input_quantization_scale(hidden_states)
+                    if use_modulate_fusion and not torch.compiler.is_compiling() else None)
+        if ff_scale is not None:
+            from fastvideo.models.dits.minimax_h3_fusions.modulation_quant import residual_modulate_quantize_nvfp4
+
+            with nvtx_range("minimax_h3.transformer_block.modulate_quant_fusion"):
+                hidden_states, ff_quantized = residual_modulate_quantize_nvfp4(
+                    hidden_states, attention_output, gate_msa, self.norm2.weight, scale_mlp,
+                    shift_mlp, adaln_indices, self.norm2.eps, ff_scale,
+                )
+            # The following NVFP4 projection consumes only this tensor's shape.
+            norm_hidden_states = hidden_states
+        elif use_modulate_fusion:
             with nvtx_range("minimax_h3.transformer_block.modulate_fusion"):
                 hidden_states, norm_hidden_states = fused_residual_gate_rmsnorm_modulate(
                     hidden_states,
@@ -759,8 +880,23 @@ class MiniMaxH3TransformerBlock(nn.Module):
                     1.0 + scale_mlp.index_select(0, adaln_indices)) + shift_mlp.index_select(0, adaln_indices)
         del attention_output
         with nvtx_range("minimax_h3.transformer_block.feed_forward"), STAGES.span("feed_forward"):
-            feed_forward_output = self.ff(norm_hidden_states)
+            feed_forward_output = (self.ff(norm_hidden_states, pre_quantized=ff_quantized)
+                                   if ff_quantized is not None else self.ff(norm_hidden_states))
         if use_modulate_fusion and not torch.compiler.is_compiling():
+            if _next_block is not None and _next_block.fuse_modulate:
+                next_quant_scale = _next_block.attn.input_quantization_scale(hidden_states, rotary_emb)
+                if next_quant_scale is not None:
+                    from fastvideo.models.dits.minimax_h3_fusions.modulation_quant import (
+                        residual_modulate_quantize_nvfp4,
+                    )
+
+                    next_shift, next_scale, *_ = _next_block.adaln_proj(temb)
+                    return residual_modulate_quantize_nvfp4(
+                        hidden_states, feed_forward_output, gate_mlp, _next_block.norm1.weight,
+                        next_scale.to(hidden_states.dtype), next_shift.to(hidden_states.dtype),
+                        adaln_indices, _next_block.norm1.eps, next_quant_scale,
+                        round_residual_before_norm=True,
+                    )
             return _gated_residual(hidden_states, gate_mlp, adaln_indices, feed_forward_output)
         if use_exact_modulate:
             return exact_kernels.gate_residual(hidden_states, gate_mlp, feed_forward_output, adaln_indices)
@@ -1181,7 +1317,7 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
         object.__setattr__(self, "_splice_late", late)
         self._splice_from_step = int(from_step)
 
-    def enable_adaln_host_cache(self, table_path: str | None = None) -> None:
+    def enable_adaln_host_cache(self, table_path: str | None = None, *, keep_resident: bool = False) -> None:
         """Move every block's AdaLN projection to pinned host memory behind a per-timestep cache.
 
         With ``table_path`` (written by FASTVIDEO_H3_ADALN_DUMP), the cache is prefilled from precomputed
@@ -1193,7 +1329,7 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
             raw = torch.load(table_path, map_location="cpu")
             tables = {int(i): {ast.literal_eval(k): v for k, v in blk.items()} for i, blk in raw.items()}
         for index, block in enumerate(self.transformer_blocks):
-            block.adaln_proj.enable_host_cache(None if tables is None else tables[index])
+            block.adaln_proj.enable_host_cache(None if tables is None else tables[index], keep_resident=keep_resident)
         self._adaln_host_cache = True
         self._adaln_dumped_entries = -1
 
@@ -1331,6 +1467,18 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
                 block.adaln_proj._cache_key = key
             self._move_adaln_tables(temb.device)
 
+        # Keep cross-block preparation local to this forward. Compiled/wrapped
+        # blocks retain their normal tensor-only interface.
+        cross_block = (not torch.compiler.is_compiling() and getattr(self, "_adaln_host_cache", False)
+                       and _can_run_minimax_h3_fusion(packed_hidden_states)
+                       and packed_hidden_states.dtype == torch.bfloat16
+                       and all(type(b) is MiniMaxH3TransformerBlock
+                               and getattr(b.forward, "__func__", None) is MiniMaxH3TransformerBlock.forward
+                               and getattr(b, "_compiled_call_impl", None) is None
+                               and not b._forward_pre_hooks and not b._forward_hooks
+                               and b.norm1.weight.device == packed_hidden_states.device
+                               for b in self.transformer_blocks))
+        attention_prepared = None
         # The eager driver owns profiling markers while each block's compiled
         # forward owns the graph that the marker surrounds.
         for block_index, block in enumerate(self.transformer_blocks):
@@ -1338,13 +1486,21 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
                 logger.info("H3_MEMORY_BLOCK %d allocated=%.3f GiB reserved=%.3f GiB", block_index,
                             torch.cuda.memory_allocated() / 2**30, torch.cuda.memory_reserved() / 2**30)
             with nvtx_range(f"minimax_h3.transformer_block.{block_index}"), STAGES.span("block_total"):
-                packed_hidden_states = block(
-                    packed_hidden_states,
-                    temb,
-                    adaln_indices,
-                    rotary_emb,
-                    original_seq_len,
-                )
+                if cross_block:
+                    next_block = (self.transformer_blocks[block_index + 1]
+                                  if block_index + 1 < len(self.transformer_blocks) else None)
+                    result = block(
+                        packed_hidden_states, temb, adaln_indices, rotary_emb, original_seq_len,
+                        _attention_prepared=attention_prepared, _next_block=next_block,
+                    )
+                    if isinstance(result, tuple):
+                        packed_hidden_states, attention_prepared = result
+                    else:
+                        packed_hidden_states, attention_prepared = result, None
+                else:
+                    packed_hidden_states = block(
+                        packed_hidden_states, temb, adaln_indices, rotary_emb, original_seq_len,
+                    )
 
         packed_hidden_states = self.norm_out(
             packed_hidden_states,

fastvideo/models/dits/minimax_h3_fusions/block_lists.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Compact VSA tile masks into the native FP4 kernel's descending block lists."""
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _block_lists(MASK, VALID, IDX, NUM, KV, QUAD,
+                 MB: tl.constexpr, MH: tl.constexpr, MR: tl.constexpr, MC: tl.constexpr,
+                 H: tl.constexpr, T: tl.constexpr, HT: tl.constexpr,
+                 NB: tl.constexpr, BLOCK: tl.constexpr, WRITE_QUAD: tl.constexpr):
+    query = tl.program_id(0)
+    head = tl.program_id(1)
+    batch = tl.program_id(2)
+    lane = tl.arange(0, BLOCK)
+    col = NB - 1 - lane
+    valid_col = lane < NB
+    a, b = col * 2, col * 2 + 1
+    va = tl.load(VALID + a // HT, (a < T * HT) & valid_col, other=0).to(tl.int32) - a % HT * 64
+    vb = tl.load(VALID + b // HT, (b < T * HT) & valid_col, other=0).to(tl.int32) - b % HT * 64
+    va = tl.minimum(tl.maximum(va, 0), 64)
+    vb = tl.minimum(tl.maximum(vb, 0), 64)
+    q0, q1 = query * 2, query * 2 + 1
+    base = batch * MB + head * MH
+    m00 = tl.load(MASK + base + q0 // HT * MR + a // HT * MC,
+                  (q0 < T * HT) & (a < T * HT) & valid_col, other=0)
+    m01 = tl.load(MASK + base + q0 // HT * MR + b // HT * MC,
+                  (q0 < T * HT) & (b < T * HT) & valid_col, other=0)
+    m10 = tl.load(MASK + base + q1 // HT * MR + a // HT * MC,
+                  (q1 < T * HT) & (a < T * HT) & valid_col, other=0)
+    m11 = tl.load(MASK + base + q1 // HT * MR + b // HT * MC,
+                  (q1 < T * HT) & (b < T * HT) & valid_col, other=0)
+    # The native builder pads an odd final query half with a key-zero entry.
+    m10 = tl.where(q1 >= T * HT, a == 0, m10)
+    m11 = tl.where(q1 >= T * HT, b == 0, m11)
+    bits = ((m00 & (va > 0)).to(tl.uint8) + 2 * (m01 & (vb > 0)).to(tl.uint8)
+            + 4 * (m10 & (va > 0)).to(tl.uint8) + 8 * (m11 & (vb > 0)).to(tl.uint8))
+    keep = (bits != 0) & valid_col
+    pos = tl.cumsum(keep.to(tl.int32), 0) - 1
+    offset = ((batch * H + head) * NB + query) * NB
+    tl.store(IDX + offset + lane, 0, valid_col)
+    if WRITE_QUAD:
+        tl.store(QUAD + offset + lane, 0, valid_col)
+    tl.debug_barrier()
+    tl.store(IDX + offset + pos, col, keep)
+    if WRITE_QUAD:
+        tl.store(QUAD + offset + pos, bits, keep)
+    tl.store(NUM + (batch * H + head) * NB + query, tl.sum(keep.to(tl.int32), 0))
+    if batch == 0 and head == 0 and query == 0:
+        tl.store(KV + a, va, valid_col)
+        tl.store(KV + b, vb, valid_col)
+
+
+def tile_mask_to_fp4_blocks(mask: torch.Tensor, tile: int, valid: torch.Tensor):
+    """Match Sage's indices, counts, half validity, and optional quadrant bytes.
+
+    Inputs come from H3's mask builder: square CUDA bool masks and contiguous
+    integer validity counts, with positive tile sizes divisible by 64.
+    """
+    batch, heads, tiles, _ = mask.shape
+    halves = tile // 64
+    blocks = triton.cdiv(tiles * halves, 2)
+    indices = torch.empty((batch, heads, blocks, blocks), device=mask.device, dtype=torch.int32)
+    counts = torch.empty((batch, heads, blocks), device=mask.device, dtype=torch.int32)
+    kv_valid = torch.empty((blocks * 2,), device=mask.device, dtype=torch.int32)
+    quadrants = torch.empty_like(indices, dtype=torch.uint8) if halves == 1 else None
+    _block_lists[(blocks, heads, batch)](
+        mask, valid, indices, counts, kv_valid, quadrants if quadrants is not None else indices,
+        *mask.stride(), heads, tiles, halves, blocks, triton.next_power_of_2(blocks), halves == 1,
+        num_warps=4,
+    )
+    return indices, counts, kv_valid, quadrants

fastvideo/models/dits/minimax_h3_fusions/key_norm_mean.py

@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fuse compact key normalization with the GB10 ATen mean reduction tree.
+
+For 18,432 rows and 56 head-128 outputs, ATen uses four accumulator lanes,
+four row-reduction warps and 18 CTAs per output. Preserve those additions and
+the BF16 normalization boundary while eliminating the separate full-key read.
+The caller restricts this specialization to the measured one-Spark geometry.
+"""
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _gather_normalize_partials(
+    X, W, C, S, Index, Y, Partial, ROWS: tl.constexpr, HEADS: tl.constexpr,
+    INPUT_ROWS: tl.constexpr, ROTARY: tl.constexpr, EPS: tl.constexpr, CTAS: tl.constexpr,
+):
+    cta, head, batch = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    lanes = tl.arange(0, 16)[:, None]
+    cols = tl.arange(0, 128)[None, :]
+    accumulators = tl.full((16, 128), 0, tl.float32)
+    for iteration in range(tl.cdiv(ROWS, 16 * CTAS)):
+        # ATen input_idx = warp + cta * 4, step_input = 4 * CTAS;
+        # its four independent accumulators advance by step_input * 4.
+        row = lanes % 4 + cta * 4 + (lanes // 4) * (4 * CTAS) + iteration * (16 * CTAS)
+        source = tl.load(Index + row, row < ROWS, other=-1)
+        offset = ((batch * INPUT_ROWS + source) * HEADS + head) * 128 + cols
+        x = tl.load(X + offset, (row < ROWS) & (source >= 0), other=0).to(tl.float32)
+        squared = x * x
+        # Original four-warp normalization order, independent of this launch's layout.
+        for shift in tl.static_range(5):
+            other = tl.gather(squared, tl.broadcast_to(cols ^ (16 >> shift), (16, 128)), axis=1)
+            squared = squared + other
+        for shift in tl.static_range(2):
+            other = tl.gather(squared, tl.broadcast_to(cols ^ (64 >> shift), (16, 128)), axis=1)
+            squared = squared + other
+        variance = tl.gather(squared, tl.full((16, 1), 0, tl.int32), axis=1) / 128
+        normalized = x * tl.rsqrt(variance + EPS) * tl.load(W + cols).to(tl.float32)
+        partner = tl.where(cols < ROTARY // 2, cols + ROTARY // 2, cols - ROTARY // 2)
+        partner_normalized = tl.gather(normalized, tl.broadcast_to(partner, (16, 128)), axis=1)
+        rotated = tl.where(cols < ROTARY // 2, -partner_normalized, partner_normalized)
+        rotary_mask = (row < ROWS) & (cols < ROTARY)
+        cos = tl.load(C + row * ROTARY + cols, rotary_mask, other=1).to(tl.float32)
+        sin = tl.load(S + row * ROTARY + cols, rotary_mask, other=0).to(tl.float32)
+        output = tl.where(cols < ROTARY, normalized * cos + rotated * sin, normalized).to(tl.bfloat16)
+        tl.store(Y + ((batch * ROWS + row) * HEADS + head) * 128 + cols, output, row < ROWS)
+        # Mean consumes the stored BF16 values, not the normalization accumulator.
+        accumulators = accumulators + tl.where(row < ROWS, output.to(tl.float32), 0.0)
+
+    warp = tl.arange(0, 4)[:, None]
+    sums = tl.gather(accumulators, tl.broadcast_to(warp, (4, 128)), axis=0)
+    for accumulator in tl.static_range(1, 4):
+        sums = sums + tl.gather(accumulators, tl.broadcast_to(warp + accumulator * 4, (4, 128)), axis=0)
+    s0 = tl.gather(sums, tl.full((1, 128), 0, tl.int32), axis=0)
+    s1 = tl.gather(sums, tl.full((1, 128), 1, tl.int32), axis=0)
+    s2 = tl.gather(sums, tl.full((1, 128), 2, tl.int32), axis=0)
+    s3 = tl.gather(sums, tl.full((1, 128), 3, tl.int32), axis=0)
+    tl.store(Partial + ((batch * CTAS + cta) * HEADS + head) * 128 + cols, (s0 + s2) + (s1 + s3))
+
+
+@triton.jit
+def _finish_mean(Partial, Mean, HEADS: tl.constexpr, ROWS: tl.constexpr, CTAS: tl.constexpr):
+    head, batch = tl.program_id(0), tl.program_id(1)
+    cols = tl.arange(0, 128)
+    s0 = tl.full((128,), 0, tl.float32)
+    s1, s2, s3 = s0, s0, s0
+    for i in tl.static_range(tl.cdiv(CTAS, 4)):
+        s0 = s0 + tl.load(Partial + ((batch * CTAS + i * 4) * HEADS + head) * 128 + cols,
+                          i * 4 < CTAS, other=0)
+        s1 = s1 + tl.load(Partial + ((batch * CTAS + i * 4 + 1) * HEADS + head) * 128 + cols,
+                          i * 4 + 1 < CTAS, other=0)
+        s2 = s2 + tl.load(Partial + ((batch * CTAS + i * 4 + 2) * HEADS + head) * 128 + cols,
+                          i * 4 + 2 < CTAS, other=0)
+        s3 = s3 + tl.load(Partial + ((batch * CTAS + i * 4 + 3) * HEADS + head) * 128 + cols,
+                          i * 4 + 3 < CTAS, other=0)
+    tl.store(Mean + (batch * HEADS + head) * 128 + cols, ((s0 + s2) + (s1 + s3)) * (1.0 / ROWS))
+
+
+def gather_key_norm_mean(
+    x: torch.Tensor, weight: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
+    eps: float, input_index: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Return original BF16 normalized keys and their ATen-equivalent BF16 mean."""
+    if (x.ndim != 4 or x.shape[0] != 1 or x.shape[2:] != (56, 128)
+            or not x.is_cuda or x.dtype != torch.bfloat16 or input_index.numel() != 18432
+            or weight.shape != (128,) or cos.shape != (18432, 96) or sin.shape != cos.shape
+            or input_index.dtype not in (torch.int32, torch.int64)):
+        raise ValueError("Fused key mean requires CUDA BF16 batch-1/head-56, 18432 rows and rotary width 96")
+    if torch.is_grad_enabled() or torch.compiler.is_compiling():
+        raise RuntimeError("Fused key mean requires eager inference")
+    x, weight, cos, sin, input_index = (t.contiguous() for t in (x, weight, cos, sin, input_index))
+    rows, heads, dim, ctas = 18432, 56, 128, 18
+    out = torch.empty((1, rows, heads, dim), dtype=x.dtype, device=x.device)
+    partial = torch.empty((1, ctas, heads, dim), dtype=torch.float32, device=x.device)
+    mean = torch.empty((1, 1, heads, dim), dtype=x.dtype, device=x.device)
+    _gather_normalize_partials[(ctas, heads, 1)](
+        x, weight, cos, sin, input_index, out, partial, rows, heads, x.shape[1], 96, eps, ctas, num_warps=4,
+    )
+    _finish_mean[(heads, 1)](partial, mean, heads, rows, ctas)
+    return out, mean

fastvideo/models/dits/minimax_h3_fusions/key_quant.py

@@ -0,0 +1,231 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fuse BF16 key smoothing, Sage FP4 packing, and valid prefix K/V preparation.
+
+The FP4 token permutation and 64-row scale swizzle match fastvideo-kernel's
+scaled_fp4_quant_permute. Both FP4 and FP8 consume the rounded BF16 key.
+"""
+import torch
+import triton
+import triton.language as tl
+
+from fastvideo.distributed.async_collectives import PeerSlices
+
+from .swiglu_quant import _rcp
+from .value_quant import _left_four
+
+
+@triton.jit
+def _smooth_pack(K, V, MEAN, INV, K4, S4, K8, V8, KB: tl.constexpr, KS: tl.constexpr, KH: tl.constexpr,
+                 VB: tl.constexpr, VS: tl.constexpr, VH: tl.constexpr,
+                 B: tl.constexpr, N: tl.constexpr, H: tl.constexpr, VALID: tl.constexpr,
+                 R: tl.constexpr, VALUE_INDEX=None, GATHER_VALUE: tl.constexpr = False,
+                 PREPARE_PREFIX: tl.constexpr = True, COPY_VALUES: tl.constexpr = True,
+                 POOL=None, TILE_VALID=None, POOL_KEYS: tl.constexpr = False,
+                 PACKED_SCALE=None, TILE_INVERSE=None, Remote=None, LOCAL_ROWS: tl.constexpr = 0,
+                 PAIR: tl.constexpr = False, RANK: tl.constexpr = 0, RECEIVED_KEYS: tl.constexpr = False):
+    row = tl.program_id(0).to(tl.int64) * R + tl.arange(0, R)
+    batch = tl.program_id(1)
+    head = tl.program_id(2)
+    col = tl.arange(0, 128)
+    rem = row % 32
+    permuted = row // 32 * 32 + (rem // 8) * 2 + (rem % 8 // 2) * 8 + rem % 2
+    source = row if POOL_KEYS else permuted
+    destination4 = permuted if POOL_KEYS else row
+    mask = row[:, None] < N
+    if RECEIVED_KEYS:
+        packed_row = tl.load(TILE_INVERSE + source, row < N, other=-1)
+        valid_key = mask & (packed_row[:, None] >= 0)
+        peer = packed_row // LOCAL_ROWS
+        packed = (peer * 3 + 1) * LOCAL_ROWS + packed_row % LOCAL_ROWS
+        if PAIR:
+            offset = (LOCAL_ROWS + packed_row % LOCAL_ROWS)[:, None] * H * 128 + head * 128 + col[None, :]
+            pointer = tl.where(peer[:, None] == RANK, K + offset, Remote + offset)
+        else:
+            pointer = K + packed[:, None] * H * 128 + head * 128 + col[None, :]
+        code = tl.load(pointer, valid_key, other=0.0).to(tl.float32)
+        transport_scale = tl.load(PACKED_SCALE + packed * H + head, (row < N) & (packed_row >= 0), other=0)
+        key = (code * transport_scale[:, None]).to(tl.bfloat16).to(tl.float32)
+    else:
+        key = tl.load(K + batch * KB + source[:, None] * KS + head * KH + col[None, :],
+                      mask, other=0).to(tl.float32)
+    if POOL_KEYS:
+        raw = tl.trans(tl.reshape(key, (R // 16, 16, 128)), (1, 2, 0))
+        if R == 64:
+            even, odd = tl.split(tl.reshape(raw, (16, 128, 2, 2)))
+            x0, x2 = tl.split(even)
+            x1, x3 = tl.split(odd)
+            acc = ((x0 + x1) + x2) + x3
+        else:
+            even, odd = tl.split(tl.reshape(raw, (16, 128, 4, 2)))
+            x04, x26 = tl.split(tl.reshape(even, (16, 128, 2, 2)))
+            x15, x37 = tl.split(tl.reshape(odd, (16, 128, 2, 2)))
+            x0, x4 = tl.split(x04)
+            x1, x5 = tl.split(x15)
+            x2, x6 = tl.split(x26)
+            x3, x7 = tl.split(x37)
+            acc = ((((((x0 + x1) + x2) + x3) + x4) + x5) + x6) + x7
+        even, odd = tl.split(tl.reshape(tl.trans(acc), (128, 4, 2, 2)))
+        y0, y2 = tl.split(even)
+        y1, y3 = tl.split(odd)
+        y0, y1 = _left_four(y0, 128), _left_four(y1, 128)
+        y2, y3 = _left_four(y2, 128), _left_four(y3, 128)
+        total = (y0 + y2) + (y1 + y3)
+        tile = tl.program_id(0)
+        count = tl.load(TILE_VALID + tile).to(tl.float32)
+        tl.store(POOL + ((batch * (N // R) + tile) * H + head) * 128 + col, tl.div_rn(total, count))
+    mean = tl.load(MEAN + (batch * H + head) * 128 + col).to(tl.float32)
+    key = (key - mean[None, :]).to(tl.bfloat16).to(tl.float32)
+    groups = tl.reshape(key, (R, 8, 16))
+    maximum = tl.max(tl.abs(groups), 2)
+    sf = (maximum * (1.0 / 6.0)).to(tl.float8e4nv)
+    inv4 = tl.where(sf.to(tl.float32) == 0, 0, _rcp(sf.to(tl.float32)))
+    q = tl.reshape(groups * inv4[:, :, None], (R, 64, 2))
+    lo, hi = tl.split(q)
+    packed = tl.inline_asm_elementwise('''{
+        .reg .b8 q;
+        cvt.rn.satfinite.e2m1x2.f32 q,$1,$2;
+        mov.b32 $0,{q,q,q,q};
+    }''', constraints='=r,f,f', args=[hi, lo], dtype=tl.uint8, is_pure=True, pack=1)
+    c4 = tl.arange(0, 64)
+    tl.store(K4 + ((batch * H + head) * N + destination4[:, None]) * 64 + c4[None, :], packed, mask)
+    g = tl.arange(0, 8)
+    si = ((destination4[:, None] // 64) * 512 + (g[None, :] // 4) * 256 + g[None, :] % 4
+          + (destination4[:, None] % 64 // 16) * 4 + (destination4[:, None] % 16) * 16)
+    tl.store(S4 + (batch * H + head) * N * 8 + si, sf, mask)
+    if PREPARE_PREFIX:
+        destination = tl.load(INV + source, row < N, other=-1).to(tl.int64)
+        valid = mask & (destination[:, None] >= 0)
+        scale8 = tl.max(tl.abs(key), 1) / 448.0
+        inv8 = tl.where(scale8 == 0, 0, 1.0 / scale8)
+        key8 = tl.minimum(tl.maximum(key * inv8[:, None], -448.0), 448.0).to(tl.float8e4nv).to(tl.float32)
+        offset = ((batch * VALID + destination[:, None]) * H + head) * 128 + col[None, :]
+        tl.store(K8 + offset, key8 * scale8[:, None], valid)
+        if COPY_VALUES:
+            value_source = tl.load(VALUE_INDEX + source, row < N, other=-1) if GATHER_VALUE else source
+            value = tl.load(V + batch * VB + value_source[:, None] * VS + head * VH + col[None, :], valid, other=0)
+            tl.store(V8 + offset, value, valid)
+
+
+def _prepare_keys(
+    k: torch.Tensor, v: torch.Tensor, mean: torch.Tensor, inverse: torch.Tensor, n_valid: int,
+    value_index: torch.Tensor | None = None, pool_valid: torch.Tensor | None = None, *, tile_elems: int = 64,
+    prefix_value: torch.Tensor | None = None,
+    received_keys: tuple[torch.Tensor | PeerSlices, torch.Tensor, torch.Tensor] | None = None,
+) -> tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor, torch.Tensor, torch.Tensor | None]:
+    """Return Sage-packed K and FP8-rounded K/BF16 V gathered into prefix order."""
+    if (not k.is_cuda or k.dtype != torch.bfloat16 or v.dtype != k.dtype or v.device != k.device
+            or k.ndim != 4 or v.ndim != 4 or k.shape[0] != v.shape[0] or k.shape[2:] != v.shape[2:]
+            or (value_index is None and k.shape != v.shape) or k.shape[-1] != 128 or k.shape[1] % 128
+            or k.stride(-1) != 1 or v.stride(-1) != 1):
+        raise ValueError("Key quantization requires CUDA BF16 BSHD with 128-wide heads and padded token rows")
+    if (mean.shape != (k.shape[0], 1, k.shape[2], 128) or mean.dtype != k.dtype
+            or mean.device != k.device or not mean.is_contiguous()):
+        raise ValueError("Key quantization requires a contiguous BF16 token mean")
+    if (inverse.shape != (k.shape[1],) or inverse.device != k.device
+            or inverse.dtype not in (torch.int32, torch.int64)):
+        raise ValueError("Key quantization requires an integer inverse row map on the input device")
+    if value_index is not None and (value_index.shape != (k.shape[1],) or value_index.device != v.device
+                                    or value_index.dtype not in (torch.int32, torch.int64)
+                                    or not value_index.is_contiguous()):
+        raise ValueError("Compact prefix values require an integer tile-to-packed row map")
+    if torch.is_grad_enabled():
+        raise RuntimeError("Key quantization is inference-only")
+    batch, rows, heads, _ = k.shape
+    k4 = torch.empty((batch, heads, rows, 64), device=k.device, dtype=torch.uint8)
+    s4 = torch.empty((batch, heads, rows, 8), device=k.device, dtype=torch.float8_e4m3fn)
+    k8 = torch.empty((batch, n_valid, heads, 128), device=k.device, dtype=torch.bfloat16)
+    if prefix_value is not None and (prefix_value.shape != k8.shape or prefix_value.dtype != k8.dtype
+                                     or prefix_value.device != k8.device or not prefix_value.is_contiguous()):
+        raise ValueError("Prepared prefix values must match the contiguous BF16 prefix-key output")
+    v8 = torch.empty_like(k8) if prefix_value is None else prefix_value
+    pooled = None
+    if pool_valid is not None:
+        if tile_elems not in (64, 128):
+            raise ValueError("Key pooling supports 64- and 128-token tiles")
+        if (pool_valid.shape != (rows // tile_elems,) or pool_valid.device != k.device
+                or pool_valid.dtype not in (torch.int32, torch.int64) or not pool_valid.is_contiguous()):
+            raise ValueError("Key pooling requires contiguous integer tile validity counts")
+        pooled = torch.empty((batch, rows // tile_elems, heads, 128), device=k.device, dtype=torch.float32)
+    key_input, remote = k, k
+    transport_scale, tile_inverse, local_rows, rank, pair = None, None, 0, 0, None
+    if received_keys is not None:
+        payload, transport_scale, tile_inverse = received_keys
+        pair = payload if isinstance(payload, PeerSlices) else None
+        template = pair.template if pair is not None else payload
+        if (batch != 1 or pool_valid is not None or prefix_value is None
+                or template.ndim != 5 or template.shape[1] != 3 or template.shape[2] < 1
+                or template.shape[3:] != (heads, 128) or template.dtype != torch.float8_e4m3fn
+                or template.device != k.device or not template.is_contiguous()
+                or transport_scale.shape != template.shape[:-1] or transport_scale.dtype != torch.float32
+                or transport_scale.device != k.device or not transport_scale.is_contiguous()
+                or tile_inverse.shape != (rows,) or tile_inverse.dtype not in (torch.int32, torch.int64)
+                or tile_inverse.device != k.device or not tile_inverse.is_contiguous()):
+            raise ValueError("Received keys require FP8 QKV, FP32 scales, an inverse row map and prepared prefix values")
+        if pair is not None and (template.shape[0] != 2 or pair.rank not in (0, 1)
+                or any(t.shape != template.shape[1:] or t.dtype != template.dtype
+                       or t.device != template.device or not t.is_contiguous() for t in (pair.local, pair.remote))):
+            raise ValueError("Direct-local received keys require matching contiguous two-rank peer slices")
+        key_input = pair.local if pair is not None else template
+        remote = pair.remote if pair is not None else template
+        local_rows, rank = template.shape[2], pair.rank if pair is not None else 0
+    block = tile_elems if pooled is not None else 4
+    _smooth_pack[(triton.cdiv(rows, block), batch, heads)](
+        key_input, v, mean, inverse, k4, s4, k8, v8, *k.stride()[:3], *v.stride()[:3],
+        batch, rows, heads, n_valid, block, VALUE_INDEX=value_index, GATHER_VALUE=value_index is not None,
+        POOL=pooled, TILE_VALID=pool_valid, POOL_KEYS=pooled is not None, COPY_VALUES=prefix_value is None,
+        PACKED_SCALE=transport_scale, TILE_INVERSE=tile_inverse, Remote=remote, LOCAL_ROWS=local_rows,
+        PAIR=pair is not None, RANK=rank, RECEIVED_KEYS=received_keys is not None,
+        num_warps=8 if pooled is not None else 4,
+    )
+    return (k4, s4), k8, v8, pooled
+
+
+def quantize_smoothed_keys(
+    k: torch.Tensor, v: torch.Tensor, mean: torch.Tensor, inverse: torch.Tensor, n_valid: int,
+    value_index: torch.Tensor | None = None, *, prefix_value: torch.Tensor | None = None,
+    received_keys: tuple[torch.Tensor | PeerSlices, torch.Tensor, torch.Tensor] | None = None,
+) -> tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor, torch.Tensor]:
+    """Prepare keys and optionally consume values already in dense-prefix order.
+
+    When ``prefix_value`` is supplied, V is used only for metadata validation;
+    its contents are not read. The supplied buffer is returned without a copy.
+    With ``received_keys=(payload, scale, tile_inverse)``, K is metadata only;
+    keys are expanded from the current call's FP8 packets with the same BF16
+    rounding before smoothing. This requires prepared prefix values and batch one.
+    """
+    packed, prefix_key, prefix_value, _ = _prepare_keys(
+        k, v, mean, inverse, n_valid, value_index, prefix_value=prefix_value, received_keys=received_keys,
+    )
+    return packed, prefix_key, prefix_value
+
+
+def pool_quantize_smoothed_keys(
+    k: torch.Tensor, v: torch.Tensor, mean: torch.Tensor, inverse: torch.Tensor, n_valid: int,
+    valid: torch.Tensor, value_index: torch.Tensor | None = None, *, tile_elems: int = 64,
+) -> tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Also return exact ATen-order FP32 means of unsmoothed 64- or 128-token tiles."""
+    packed, prefix_key, prefix_value, pooled = _prepare_keys(
+        k, v, mean, inverse, n_valid, value_index, valid, tile_elems=tile_elems,
+    )
+    assert pooled is not None
+    return packed, prefix_key, prefix_value, pooled.permute(0, 2, 1, 3)
+
+
+def quantize_smoothed_keys_only(k: torch.Tensor, mean: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pack the rounded BF16 key shift without allocating dense-prefix outputs."""
+    if (not k.is_cuda or k.dtype != torch.bfloat16 or k.ndim != 4 or k.shape[-1] != 128
+            or k.shape[1] % 128 or k.stride(-1) != 1):
+        raise ValueError("Key-only preparation requires CUDA BF16 head-128 with padded token rows")
+    if (mean.shape != (k.shape[0], 1, k.shape[2], 128) or mean.dtype != k.dtype
+            or mean.device != k.device or not mean.is_contiguous()):
+        raise ValueError("Key-only preparation requires a contiguous BF16 token mean")
+    if torch.is_grad_enabled() or torch.compiler.is_compiling():
+        raise RuntimeError("Key-only preparation requires eager inference")
+    batch, rows, heads, _ = k.shape
+    codes = torch.empty((batch, heads, rows, 64), device=k.device, dtype=torch.uint8)
+    scales = torch.empty((batch, heads, rows, 8), device=k.device, dtype=torch.float8_e4m3fn)
+    _smooth_pack[(triton.cdiv(rows, 4), batch, heads)](
+        k, k, mean, k, codes, scales, k, k, *k.stride()[:3], *k.stride()[:3],
+        batch, rows, heads, 0, 4, PREPARE_PREFIX=False, num_warps=4,
+    )
+    return codes, scales

fastvideo/models/dits/minimax_h3_fusions/kv_pool.py

@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Share the compact value read across prefix attention, pooling and native FP4 packing."""
+import torch
+import triton
+import triton.language as tl
+
+from .swiglu_quant import _rcp
+from .tile_pool import _tile_sum
+
+
+@triton.jit
+def _smooth_pack(K, V, MEAN, INV, K4, S4, K8, V8, KB: tl.constexpr, KS: tl.constexpr, KH: tl.constexpr,
+                 VB: tl.constexpr, VS: tl.constexpr, VH: tl.constexpr,
+                 B: tl.constexpr, N: tl.constexpr, H: tl.constexpr, VALID: tl.constexpr,
+                 R: tl.constexpr, VALUE_INDEX=None, GATHER_VALUE: tl.constexpr = False,
+                 PREPARE_PREFIX: tl.constexpr = True,
+                 POOL=None, TILE_VALID=None, POOL_KEYS: tl.constexpr = False,
+                 VP=None, V4=None, VS4=None, HEAD_FIRST: tl.constexpr = False):
+    row = tl.program_id(0).to(tl.int64) * R + tl.arange(0, R)
+    batch = tl.program_id(1)
+    head = tl.program_id(2)
+    col = tl.arange(0, 128)
+    rem = row % 32
+    permuted = row // 32 * 32 + (rem // 8) * 2 + (rem % 8 // 2) * 8 + rem % 2
+    source = row if POOL_KEYS else permuted
+    destination4 = permuted if POOL_KEYS else row
+    mask = row[:, None] < N
+    key = tl.load(K + batch * KB + source[:, None] * KS + head * KH + col[None, :],
+                  mask, other=0).to(tl.float32)
+    if POOL_KEYS:
+        total = _tile_sum(key, R)
+        tile = tl.program_id(0)
+        count = tl.load(TILE_VALID + tile).to(tl.float32)
+        tl.store(POOL + ((batch * (N // R) + tile) * H + head) * 128 + col, tl.div_rn(total, count))
+    mean = tl.load(MEAN + (batch * H + head) * 128 + col).to(tl.float32)
+    key = (key - mean[None, :]).to(tl.bfloat16).to(tl.float32)
+    groups = tl.reshape(key, (R, 8, 16))
+    maximum = tl.max(tl.abs(groups), 2)
+    sf = (maximum * (1.0 / 6.0)).to(tl.float8e4nv)
+    inv4 = tl.where(sf.to(tl.float32) == 0, 0, _rcp(sf.to(tl.float32)))
+    q = tl.reshape(groups * inv4[:, :, None], (R, 64, 2))
+    lo, hi = tl.split(q)
+    packed = tl.inline_asm_elementwise('''{
+        .reg .b8 q;
+        cvt.rn.satfinite.e2m1x2.f32 q,$1,$2;
+        mov.b32 $0,{q,q,q,q};
+    }''', constraints='=r,f,f', args=[hi, lo], dtype=tl.uint8, is_pure=True, pack=1)
+    c4 = tl.arange(0, 64)
+    tl.store(K4 + ((batch * H + head) * N + destination4[:, None]) * 64 + c4[None, :], packed, mask)
+    g = tl.arange(0, 8)
+    si = ((destination4[:, None] // 64) * 512 + (g[None, :] // 4) * 256 + g[None, :] % 4
+          + (destination4[:, None] % 64 // 16) * 4 + (destination4[:, None] % 16) * 16)
+    tl.store(S4 + (batch * H + head) * N * 8 + si, sf, mask)
+    if PREPARE_PREFIX:
+        destination = tl.load(INV + source, row < N, other=-1).to(tl.int64)
+        valid = mask & (destination[:, None] >= 0)
+        scale8 = tl.max(tl.abs(key), 1) / 448.0
+        inv8 = tl.where(scale8 == 0, 0, 1.0 / scale8)
+        key8 = tl.minimum(tl.maximum(key * inv8[:, None], -448.0), 448.0).to(tl.float8e4nv).to(tl.float32)
+        offset = ((batch * VALID + destination[:, None]) * H + head) * 128 + col[None, :]
+        if HEAD_FIRST:
+            offset = ((batch * H + head) * VALID + destination[:, None]) * 128 + col[None, :]
+        tl.store(K8 + offset, key8 * scale8[:, None], valid)
+        value_source = tl.load(VALUE_INDEX + source, row < N, other=-1) if GATHER_VALUE else source
+        value = tl.load(V + batch * VB + value_source[:, None] * VS + head * VH + col[None, :],
+                        mask & (value_source[:, None] >= 0), other=0)
+        tl.store(V8 + offset, value, valid)
+
+
+    # V was already gathered for dense prefix attention; keep its rounded bits.
+    vf = value.to(tl.float32)
+    vg = tl.reshape(vf, (R // 16, 16, 128))
+    maximum_v = tl.max(tl.abs(vg), 1)
+    sfv = (maximum_v * (1.0 / 6.0)).to(tl.float8e4nv)
+    invv = tl.where(sfv.to(tl.float32) == 0, 0, _rcp(sfv.to(tl.float32)))
+    scaled_v = tl.reshape(vg * invv[:, None, :], (R // 2, 2, 128))
+    vlo, vhi = tl.split(tl.trans(scaled_v, (0, 2, 1)))
+    vcodes = tl.inline_asm_elementwise("""{
+        .reg .b8 q;
+        cvt.rn.satfinite.e2m1x2.f32 q,$1,$2;
+        mov.b32 $0,{q,q,q,q};
+    }""", constraints='=r,f,f', args=[vhi, vlo], dtype=tl.uint8, is_pure=True, pack=1)
+    vr = tl.arange(0, R // 2)
+    voffset = (((batch * H + head) * 128 + col[None, :]) * (N // 2)
+               + tile * (R // 2) + vr[:, None])
+    tl.store(V4 + voffset, vcodes)
+    group_v = tile * (R // 16) + tl.arange(0, R // 16)
+    soffset = ((batch * H + head) * 128 * (N // 16)
+               + (col[None, :] // 64) * 64 * (N // 16)
+               + (group_v[:, None] // 4) * 256 + group_v[:, None] % 4
+               + (col[None, :] % 64 // 16) * 4 + (col[None, :] % 16) * 16)
+    tl.store(VS4 + soffset, sfv)
+    vtotal = _tile_sum(vf, R)
+    tl.store(VP + ((batch * (N // R) + tile) * H + head) * 128 + col, tl.div_rn(vtotal, count))
+
+
+def pool_quantize_prefix_kv(
+    k: torch.Tensor, v: torch.Tensor, mean: torch.Tensor, inverse: torch.Tensor,
+    n_valid: int, valid: torch.Tensor, value_index: torch.Tensor | None, *, tile_elems: int = 64,
+) -> tuple[
+    tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor, torch.Tensor, torch.Tensor],
+    tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],
+]:
+    """Return exact pooled/native KV plus dense-prefix KV with contiguous head sequences.
+
+    Keys are already normalized in tile order. Values remain in compact projection
+    order with a row map, or padded projection order without one. A row map
+    zero-fills missing source rows; direct inputs retain every padded value. Each value is read
+    once for the prefix BF16 tensor, FP32 tile mean, and transposed FP4 codes/scales.
+    """
+    if (not k.is_cuda or k.ndim != 4 or v.ndim != 4 or k.dtype != torch.bfloat16
+            or v.dtype != k.dtype or v.device != k.device or k.shape[0] != v.shape[0]
+            or k.shape[2:] != v.shape[2:] or k.shape[-1] != 128 or k.shape[1] % 128
+            or tile_elems not in (64, 128) or (value_index is None and v.shape[1] != k.shape[1])
+            or k.stride(-1) != 1 or v.stride(-1) != 1):
+        raise ValueError("Combined KV preparation requires CUDA BF16 BSHD with aligned head-128 keys")
+    batch, rows, heads, _ = k.shape
+    if (mean.shape != (batch, 1, heads, 128) or mean.dtype != k.dtype or mean.device != k.device
+            or not mean.is_contiguous() or not 0 < n_valid <= rows):
+        raise ValueError("Combined KV preparation requires a BF16 key mean and valid prefix extent")
+    for index, size in ((inverse, rows), (value_index, rows), (valid, rows // tile_elems)):
+        if index is None:
+            continue
+        if (index.shape != (size,) or index.device != k.device or not index.is_contiguous()
+                or index.dtype not in (torch.int32, torch.int64)):
+            raise ValueError("Combined KV preparation requires contiguous integer row maps and tile counts")
+    if torch.is_grad_enabled() or torch.compiler.is_compiling():
+        raise RuntimeError("Combined KV preparation requires eager inference")
+    k4 = torch.empty((batch, heads, rows, 64), device=k.device, dtype=torch.uint8)
+    s4 = torch.empty((batch, heads, rows, 8), device=k.device, dtype=torch.float8_e4m3fn)
+    k8 = torch.empty((batch, heads, n_valid, 128), device=k.device, dtype=torch.bfloat16).transpose(1, 2)
+    v8 = torch.empty_like(k8)
+    kp = torch.empty((batch, rows // tile_elems, heads, 128), device=k.device, dtype=torch.float32)
+    vp = torch.empty_like(kp)
+    v4 = torch.empty((batch, heads, 128, rows // 2), device=k.device, dtype=torch.uint8)
+    vs4 = torch.empty((batch, heads, 128, rows // 16), device=k.device, dtype=torch.float8_e4m3fn)
+    _smooth_pack[(rows // tile_elems, batch, heads)](
+        k, v, mean, inverse, k4, s4, k8, v8, *k.stride()[:3], *v.stride()[:3],
+        batch, rows, heads, n_valid, tile_elems, VALUE_INDEX=value_index, GATHER_VALUE=value_index is not None,
+        POOL=kp, TILE_VALID=valid, POOL_KEYS=True, VP=vp, V4=v4, VS4=vs4, HEAD_FIRST=True, num_warps=4 if tile_elems == 64 else 8,
+    )
+    return ((k4, s4), k8, v8, kp.permute(0, 2, 1, 3)), (vp.permute(0, 2, 1, 3), (v4, vs4))

fastvideo/models/dits/minimax_h3_fusions/modulation_quant.py

@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Exact residual modulation fused with NVFP4 activation preparation."""
+import torch
+import triton
+import triton.language as tl
+
+from .modulation import (
+    _num_warps,
+    _require_forward_only,
+    _require_triton_cuda,
+    _row_addressable,
+    _validate_contract,
+    _validate_residual_branch,
+)
+from .swiglu_quant import _rcp as rcp
+
+
+@triton.jit
+def _modulate_quant(
+    hidden_out_ptr,
+    Q, S, SF,
+    residual_ptr,
+    branch_ptr,
+    gate_ptr,
+    weight_ptr,
+    scale_ptr,
+    shift_ptr,
+    index_ptr,
+    n_cols,
+    n_index,
+    eps,
+    stride_input_row,
+    stride_gate_row,
+    stride_scale_row,
+    stride_shift_row,
+    BLOCK: tl.constexpr,
+    WIDTH: tl.constexpr,
+    ADD_RESIDUAL: tl.constexpr,
+    ROUND_RESIDUAL: tl.constexpr = False,
+):
+    row = tl.program_id(0).to(tl.int64)
+    cols = tl.arange(0, BLOCK)
+    mask = cols < n_cols
+    input_offsets = row * stride_input_row + cols
+    table_row = tl.load(index_ptr + row % n_index).to(tl.int64)
+
+    hidden = tl.load(residual_ptr + input_offsets, mask=mask, other=0.0).to(tl.float32)
+    if ADD_RESIDUAL:
+        branch = tl.load(branch_ptr + input_offsets, mask=mask, other=0.0).to(tl.float32)
+        gate = tl.load(
+            gate_ptr + table_row * stride_gate_row + cols,
+            mask=mask,
+            other=0.0,
+        ).to(tl.float32)
+        hidden = hidden + gate * branch
+        tl.store(hidden_out_ptr + input_offsets, hidden.to(hidden_out_ptr.dtype.element_ty), mask=mask)
+
+    if ROUND_RESIDUAL:
+        hidden = hidden.to(tl.bfloat16).to(tl.float32)
+    variance = tl.sum(hidden * hidden, axis=0) / n_cols
+    normed = hidden * tl.math.rsqrt(variance + eps)
+    weight = tl.load(weight_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    scale = tl.load(
+        scale_ptr + table_row * stride_scale_row + cols,
+        mask=mask,
+        other=0.0,
+    ).to(tl.float32)
+    shift = tl.load(
+        shift_ptr + table_row * stride_shift_row + cols,
+        mask=mask,
+        other=0.0,
+    ).to(tl.float32)
+    output = normed * weight * (1.0 + scale) + shift
+    y = tl.where(mask, output.to(tl.bfloat16).to(tl.float32), 0.0)
+    grouped = tl.reshape(y, (BLOCK // 16, 16))
+    amax = tl.max(tl.abs(grouped), 1)
+    global_sf = tl.load(SF)
+    sf8 = (global_sf * (amax * rcp(6.0))).to(tl.float8e4nv)
+    inv = tl.where(amax != 0, rcp(sf8.to(tl.float32) * rcp(global_sf)), 0.0)
+    quant = tl.reshape(grouped * inv[:, None], (BLOCK // 2, 2))
+    lo, hi = tl.split(quant)
+    packed = tl.inline_asm_elementwise('''{
+        .reg .b8 q;
+        cvt.rn.satfinite.e2m1x2.f32 q, $1, $2;
+        mov.b32 $0, {q,q,q,q};
+    }''', constraints="=r,f,f", args=[hi, lo], dtype=tl.uint8, is_pure=True, pack=1)
+    qc = tl.arange(0, BLOCK // 2)
+    tl.store(Q + row * (WIDTH // 2) + qc, packed, qc < WIDTH // 2)
+    group = tl.arange(0, BLOCK // 16)
+    scale_tile = (row // 128) * (WIDTH // 64) + group // 4
+    si = ((scale_tile * 32 + row % 32) * 4 + (row // 32) % 4) * 4 + group % 4
+    tl.store(S + si, sf8.to(tl.uint8, bitcast=True), group < WIDTH // 16)
+
+
+def residual_modulate_quantize_nvfp4(
+    residual: torch.Tensor, branch: torch.Tensor, gate: torch.Tensor, weight: torch.Tensor,
+    scale: torch.Tensor, shift: torch.Tensor, index: torch.Tensor, eps: float,
+    global_scale: torch.Tensor,
+    *, round_residual_before_norm: bool = False,
+) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
+    """Return the updated residual and packed FP4 modulated activation.
+
+    The FP32 residual participates in RMSNorm before its BF16 output cast,
+    matching the existing fused residual/modulation kernel exactly. The
+    modulated activation is rounded to BF16 before FP4 preparation.
+    Across transformer blocks, round the residual to BF16 before RMSNorm too,
+    preserving the separate block output's rounding boundary.
+    """
+    _validate_residual_branch(residual, branch)
+    _validate_contract(residual, weight, (gate, scale, shift), index, eps)
+    _require_forward_only(residual, branch, gate, weight, scale, shift)
+    _require_triton_cuda(residual)
+    width = residual.shape[-1]
+    if residual.dtype != torch.bfloat16 or width % 64:
+        raise ValueError("NVFP4 modulation requires BF16 with a width divisible by 64")
+    if (global_scale.numel() != 1 or global_scale.dtype != torch.float32
+            or global_scale.device != residual.device):
+        raise ValueError("NVFP4 modulation requires a scalar FP32 scale on the input device")
+    flat = residual.reshape(-1, width).contiguous()
+    branch = branch.reshape(-1, width).contiguous()
+    gate, scale, shift = (_row_addressable(t) for t in (gate, scale, shift))
+    rows = flat.shape[0]
+    hidden = torch.empty_like(flat)
+    codes = torch.empty((rows, width // 2), dtype=torch.uint8, device=residual.device)
+    scales = torch.zeros((triton.cdiv(rows, 128) * 128, width // 16), dtype=torch.uint8,
+                         device=residual.device)
+    block = triton.next_power_of_2(width)
+    _modulate_quant[(rows,)](
+        hidden, codes, scales, global_scale, flat, branch, gate, weight.contiguous(),
+        scale, shift, index.contiguous(), width, index.numel(), eps, flat.stride(0),
+        gate.stride(0), scale.stride(0), shift.stride(0), BLOCK=block, WIDTH=width,
+        ADD_RESIDUAL=True, ROUND_RESIDUAL=round_residual_before_norm, num_warps=_num_warps(block),
+    )
+    return hidden.view_as(residual), (codes, scales, global_scale)
+
+
+def modulate_quantize_nvfp4(
+    x: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor,
+    index: torch.Tensor, eps: float, global_scale: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """RMSNorm and row-indexed modulation followed by exact BF16-to-NVFP4 preparation."""
+    _validate_contract(x, weight, (scale, shift), index, eps)
+    _require_forward_only(x, weight, scale, shift)
+    _require_triton_cuda(x)
+    width = x.shape[-1]
+    if x.dtype != torch.bfloat16 or width % 64:
+        raise ValueError("NVFP4 modulation requires BF16 with a width divisible by 64")
+    if global_scale.numel() != 1 or global_scale.dtype != torch.float32 or global_scale.device != x.device:
+        raise ValueError("NVFP4 modulation requires a scalar FP32 scale on the input device")
+    flat = x.reshape(-1, width).contiguous()
+    scale, shift = (_row_addressable(t) for t in (scale, shift))
+    rows = flat.shape[0]
+    codes = torch.empty((rows, width // 2), dtype=torch.uint8, device=x.device)
+    scales = torch.zeros((triton.cdiv(rows, 128) * 128, width // 16), dtype=torch.uint8, device=x.device)
+    block = triton.next_power_of_2(width)
+    _modulate_quant[(rows,)](
+        flat, codes, scales, global_scale, flat, flat, scale, weight.contiguous(),
+        scale, shift, index.contiguous(), width, index.numel(), eps, flat.stride(0),
+        scale.stride(0), scale.stride(0), shift.stride(0), BLOCK=block, WIDTH=width,
+        ADD_RESIDUAL=False, num_warps=_num_warps(block),
+    )
+    return codes, scales, global_scale

fastvideo/models/dits/minimax_h3_fusions/prefix_attention.py

@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Optional exact query-tile specialization for the measured GB10 Omni prefix."""
+from functools import lru_cache
+
+import torch
+
+
+@lru_cache(maxsize=1)
+def _native():
+    try:
+        import _fastvideo_h3_prefix
+    except ImportError:
+        return None
+    return _fastvideo_h3_prefix
+
+
+def try_prefix_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor | None:
+    """BHSD operands; keep Torch SDPA outside the validated eager configuration."""
+    if (torch.is_grad_enabled() or torch.compiler.is_compiling()
+            or q.shape != (1, 56, 15360, 128) or k.shape != (1, 56, 29678, 128)
+            or v.shape != k.shape or not q.is_cuda or k.device != q.device or v.device != q.device
+            or q.dtype != torch.bfloat16 or k.dtype != q.dtype or v.dtype != q.dtype
+            or q.stride(-1) != 1 or k.stride(-1) != 1 or v.stride(-1) != 1
+            or torch.cuda.get_device_capability(q.device) != (12, 1)
+            or torch.version.git_version != '7661cd9c6b841b62b7f411aa52ec51f05457263b'):
+        return None
+    native = _native()
+    return None if native is None else native.forward(q, k, v)

fastvideo/models/dits/minimax_h3_fusions/prefix_quant.py

@@ -0,0 +1,55 @@
+# SPDX-License-Identifier: Apache-2.0
+"""FP8 prefix-key preparation fused with valid-key/value gathering."""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _prefix_kv(K, V, IDX, KO, VO, KB: tl.constexpr, KS: tl.constexpr, KH: tl.constexpr,
+               VB: tl.constexpr, VS: tl.constexpr, VH: tl.constexpr,
+               B: tl.constexpr, N: tl.constexpr, H: tl.constexpr, D: tl.constexpr, R: tl.constexpr):
+    row = tl.program_id(0).to(tl.int64) * R + tl.arange(0, R)
+    col = tl.arange(0, D)
+    batch = row // (N * H)
+    token = row // H % N
+    head = row % H
+    source = tl.load(IDX + token, row < B * N * H, other=0).to(tl.int64)
+    mask = row[:, None] < B * N * H
+    k = tl.load(K + batch[:, None] * KB + source[:, None] * KS + head[:, None] * KH + col[None, :],
+                mask, other=0).to(tl.float32)
+    scale = tl.max(tl.abs(k), 1) / 448.0
+    inv = tl.where(scale == 0, 0, 1.0 / scale)
+    rounded = tl.minimum(tl.maximum(k * inv[:, None], -448.0), 448.0).to(tl.float8e4nv).to(tl.float32)
+    tl.store(KO + row[:, None] * D + col[None, :], rounded * scale[:, None], mask)
+    v = tl.load(V + batch[:, None] * VB + source[:, None] * VS + head[:, None] * VH + col[None, :],
+                mask, other=0)
+    tl.store(VO + row[:, None] * D + col[None, :], v, mask)
+
+
+def quantize_gather_prefix_kv(
+    k: torch.Tensor, v: torch.Tensor, index: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Apply the existing per-head FP8 round trip to K and gather valid K/V rows.
+
+    K remains BF16 at the SDPA boundary. The output keeps the existing BSHD
+    layout, and each head's amax is independent of the other token rows.
+    """
+    if (not k.is_cuda or k.dtype != torch.bfloat16 or v.dtype != k.dtype or v.device != k.device
+            or k.shape != v.shape or k.ndim != 4 or k.shape[-1] != 128
+            or k.stride(-1) != 1 or v.stride(-1) != 1):
+        raise ValueError("Prefix quantization requires CUDA BF16 BSHD tensors with contiguous 128-wide heads")
+    if index.ndim != 1 or index.device != k.device or index.dtype not in (torch.int32, torch.int64):
+        raise ValueError("Prefix quantization requires a CUDA integer row index")
+    if torch.is_grad_enabled():
+        raise RuntimeError("Prefix quantization is inference-only")
+    batch, _, heads, dim = k.shape
+    rows = index.numel()
+    ko = torch.empty((batch, rows, heads, dim), device=k.device, dtype=k.dtype)
+    vo = torch.empty_like(ko)
+    if rows:
+        _prefix_kv[(triton.cdiv(batch * rows * heads, 4),)](
+            k, v, index, ko, vo, *k.stride()[:3], *v.stride()[:3], batch, rows, heads, dim, 4, num_warps=4,
+        )
+    return ko, vo

fastvideo/models/dits/minimax_h3_fusions/qknorm_rope.py

@@ -19,6 +19,7 @@ except ImportError:  # pragma: no cover - exercised only in environments without
 
 
 if HAVE_TRITON:
+    from .swiglu_quant import _rcp
 
     @triton.jit
     def _qknorm_partial_rope_kernel(
@@ -34,42 +35,84 @@ if HAVE_TRITON:
         seq_len,
         eps,
         BLOCK_SIZE: tl.constexpr,
+        ROWS: tl.constexpr,
+        ROWS_PER_PROGRAM: tl.constexpr,
+        codes_ptr=None,
+        scales_ptr=None,
+        PACK_FP4: tl.constexpr = False,
+        START: tl.constexpr = 0,
+        input_index_ptr=None,
+        INPUT_ROWS: tl.constexpr = 0,
+        GATHER_INPUT: tl.constexpr = False,
     ):
         # int64, like the sibling kernels: with int32 program ids,
         # ``row * head_dim`` wraps once the flattened input reaches 2**31
         # elements (H3's 56 heads x 128 head_dim crosses that at
         # batch*seq >= 299_593 tokens per rank) and the loads/stores below
         # become out-of-bounds. ``seq_index`` inherits int64 from ``row``.
-        row = tl.program_id(0).to(tl.int64)
-        seq_index = (row // num_heads) % seq_len
-        cols = tl.arange(0, BLOCK_SIZE)
-        head_mask = cols < head_dim
-        row_offset = row * head_dim
+        # Keep each row's four-warp reduction order while amortizing CTA scheduling.
+        for i in tl.static_range(ROWS_PER_PROGRAM):
+            row = tl.program_id(0).to(tl.int64) * ROWS_PER_PROGRAM + i
+            seq_index = (row // num_heads) % seq_len
+            cols = tl.arange(0, BLOCK_SIZE)
+            head_mask = (cols < head_dim) & (row < ROWS)
+            row_offset = row * head_dim
+            if GATHER_INPUT:
+                source = tl.load(input_index_ptr + seq_index, row < ROWS, other=-1).to(tl.int64)
+                input_offset = (((row // (seq_len * num_heads)) * INPUT_ROWS + source) * num_heads
+                                + row % num_heads) * head_dim
+                input_valid = source >= 0
+            else:
+                input_offset, input_valid = row_offset, True
+            x = tl.load(x_ptr + input_offset + cols, mask=head_mask & input_valid, other=0.0).to(tl.float32)
+            variance = tl.sum(x * x, axis=0) / head_dim
+            inv_rms = tl.math.rsqrt(variance + eps)
+            weight = tl.load(weight_ptr + cols, mask=head_mask, other=0.0).to(tl.float32)
+            normalized = x * inv_rms * weight
 
-        x = tl.load(x_ptr + row_offset + cols, mask=head_mask, other=0.0).to(tl.float32)
-        variance = tl.sum(x * x, axis=0) / head_dim
-        inv_rms = tl.math.rsqrt(variance + eps)
-        weight = tl.load(weight_ptr + cols, mask=head_mask, other=0.0).to(tl.float32)
-        normalized = x * inv_rms * weight
+            rotary_mask = (cols < rotary_dim) & (row < ROWS)
+            first_half = cols < half_rotary_dim
+            partner_col = tl.where(first_half, cols + half_rotary_dim, cols - half_rotary_dim)
+            partner_x = tl.load(
+                x_ptr + input_offset + partner_col,
+                mask=rotary_mask & input_valid,
+                other=0.0,
+            ).to(tl.float32)
+            partner_weight = tl.load(weight_ptr + partner_col, mask=rotary_mask, other=0.0).to(tl.float32)
+            partner_normalized = partner_x * inv_rms * partner_weight
+            rotated = tl.where(first_half, -partner_normalized, partner_normalized)
 
-        rotary_mask = cols < rotary_dim
-        first_half = cols < half_rotary_dim
-        partner_col = tl.where(first_half, cols + half_rotary_dim, cols - half_rotary_dim)
-        partner_x = tl.load(
-            x_ptr + row_offset + partner_col,
-            mask=rotary_mask,
-            other=0.0,
-        ).to(tl.float32)
-        partner_weight = tl.load(weight_ptr + partner_col, mask=rotary_mask, other=0.0).to(tl.float32)
-        partner_normalized = partner_x * inv_rms * partner_weight
-        rotated = tl.where(first_half, -partner_normalized, partner_normalized)
-
-        table_offset = seq_index * rotary_dim + cols
-        cos = tl.load(cos_ptr + table_offset, mask=rotary_mask, other=1.0).to(tl.float32)
-        sin = tl.load(sin_ptr + table_offset, mask=rotary_mask, other=0.0).to(tl.float32)
-        rotary_output = normalized * cos + rotated * sin
-        output = tl.where(rotary_mask, rotary_output, normalized)
-        tl.store(out_ptr + row_offset + cols, output.to(out_ptr.dtype.element_ty), mask=head_mask)
+            table_offset = seq_index * rotary_dim + cols
+            cos = tl.load(cos_ptr + table_offset, mask=rotary_mask, other=1.0).to(tl.float32)
+            sin = tl.load(sin_ptr + table_offset, mask=rotary_mask, other=0.0).to(tl.float32)
+            rotary_output = normalized * cos + rotated * sin
+            output = tl.where(rotary_mask, rotary_output, normalized)
+            tl.store(out_ptr + row_offset + cols, output.to(out_ptr.dtype.element_ty), mask=head_mask)
+            if PACK_FP4:
+                rounded = output.to(tl.bfloat16).to(tl.float32)
+                grouped = tl.reshape(rounded, (BLOCK_SIZE // 16, 16))
+                maximum = tl.max(tl.abs(grouped), 1)
+                sf = (maximum * (1.0 / 6.0)).to(tl.float8e4nv)
+                inv = tl.where(sf.to(tl.float32) == 0, 0, _rcp(sf.to(tl.float32)))
+                values = tl.reshape(grouped * inv[:, None], (BLOCK_SIZE // 2, 2))
+                lo, hi = tl.split(values)
+                packed = tl.inline_asm_elementwise('''{
+                    .reg .b8 q;
+                    cvt.rn.satfinite.e2m1x2.f32 q, $1, $2;
+                    mov.b32 $0, {q,q,q,q};
+                }''', constraints='=r,f,f', args=[hi, lo], dtype=tl.uint8, is_pure=True, pack=1)
+                batch = row // (seq_len * num_heads)
+                head = row % num_heads
+                c4 = tl.arange(0, BLOCK_SIZE // 2)
+                offset = ((batch * num_heads + head) * (seq_len - START) + seq_index - START) * (head_dim // 2)
+                tl.store(codes_ptr + offset + c4, packed, (row < ROWS) & (seq_index >= START) & (c4 < head_dim // 2))
+                group = tl.arange(0, BLOCK_SIZE // 16)
+                qi = seq_index - START
+                si = ((qi // 64) * (64 * head_dim // 16) + (group // 4) * 256 + group % 4
+                      + (qi % 64 // 16) * 4 + (qi % 16) * 16)
+                scale_offset = (batch * num_heads + head) * (seq_len - START) * (head_dim // 16) + si
+                tl.store(scales_ptr + scale_offset, sf,
+                         (row < ROWS) & (seq_index >= START) & (group < head_dim // 16))
 
 
 _SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
@@ -143,7 +186,9 @@ def _fused_qknorm_rope_impl(
     flat_x = x.reshape(-1, head_dim).contiguous()
     flat_out = torch.empty_like(flat_x)
     block_size = 1 << (head_dim - 1).bit_length()
-    _qknorm_partial_rope_kernel[(flat_x.shape[0], )](
+    rows_per_program = (8 if head_dim == 128 and x.dtype == torch.bfloat16
+                        and torch.cuda.get_device_capability(x.device)[0] >= 10 else 1)
+    _qknorm_partial_rope_kernel[(triton.cdiv(flat_x.shape[0], rows_per_program), )](
         flat_out,
         flat_x,
         weight,
@@ -156,6 +201,8 @@ def _fused_qknorm_rope_impl(
         seq_len,
         eps,
         BLOCK_SIZE=block_size,
+        ROWS=flat_x.shape[0],
+        ROWS_PER_PROGRAM=rows_per_program,
         num_warps=4,
     )
     return flat_out.view(batch, seq_len, num_heads, head_dim)
@@ -214,3 +261,74 @@ def fused_qknorm_rope(
 
 
 __all__ = ["HAVE_TRITON", "fused_qknorm_rope"]
+
+
+def qknorm_rope_quantize_fp4(
+    x: torch.Tensor, weight: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
+    eps: float, start: int = 0,
+) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+    """Keep the BF16 norm/RoPE result and pack its suffix for native Sage FP4.
+
+    The suffix begins on a 128-token boundary. Its codes and 64-row scale
+    swizzle are allocated contiguously, as required by the native attention
+    kernel. BF16 rounding precedes the per-16-element scale calculation.
+    """
+    batch, seq_len, heads, dim, rotary_dim = _validate_inputs(x, weight, cos, sin, eps)
+    if (not x.is_cuda or x.dtype != torch.bfloat16 or dim != 128 or seq_len % 128
+            or start < 0 or start >= seq_len or start % 128):
+        raise ValueError("FP4 Q packing requires CUDA BF16 head-128 and aligned sequence/suffix lengths")
+    if torch.is_grad_enabled() or torch.compiler.is_compiling():
+        raise RuntimeError("FP4 Q packing requires eager inference")
+    x, weight, cos, sin = (t.contiguous() for t in (x, weight, cos, sin))
+    out = torch.empty_like(x)
+    codes = torch.empty((batch, heads, seq_len - start, 64), device=x.device, dtype=torch.uint8)
+    scales = torch.empty((batch, heads, seq_len - start, 8), device=x.device, dtype=torch.float8_e4m3fn)
+    _qknorm_partial_rope_kernel[(triton.cdiv(batch * seq_len * heads, 4),)](
+        out, x, weight, cos, sin, dim, rotary_dim, rotary_dim // 2, heads, seq_len, eps,
+        BLOCK_SIZE=128, ROWS=batch * seq_len * heads, ROWS_PER_PROGRAM=4,
+        codes_ptr=codes, scales_ptr=scales, PACK_FP4=True, START=start, num_warps=4,
+    )
+    return out, (codes, scales)
+
+
+def gather_qknorm_rope(
+    x: torch.Tensor, weight: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, eps: float,
+    input_index: torch.Tensor, *, fp4_start: int | None = None,
+) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:
+    """Normalize compact projected rows while gathering them into padded tile order.
+
+    Negative indices denote zero projection rows. Pooling and sparse attention
+    still receive every original tile row; only the linear projections omit
+    padding. An optional aligned suffix is packed directly for native FP4 Q.
+    """
+    if x.ndim != 4 or input_index.ndim != 1:
+        raise ValueError("Gathered normalization requires BSHD input and a one-dimensional row map")
+    rows = input_index.numel()
+    shape_view = x[:, :1].expand(-1, rows, -1, -1)
+    batch, _, heads, dim, rotary_dim = _validate_inputs(shape_view, weight, cos, sin, eps)
+    if (not x.is_cuda or x.dtype != torch.bfloat16 or dim != 128 or rows % 128
+            or input_index.device != x.device or input_index.dtype not in (torch.int32, torch.int64)
+            or not input_index.is_contiguous()):
+        raise ValueError("Gathered normalization requires CUDA BF16 head-128 and an aligned integer row map")
+    if torch.is_grad_enabled() or torch.compiler.is_compiling():
+        raise RuntimeError("Gathered normalization requires eager inference")
+    quantized = fp4_start is not None
+    start = 0 if fp4_start is None else fp4_start
+    if start < 0 or start >= rows or start % 128:
+        raise ValueError("The FP4 query suffix must begin at a valid 128-token boundary")
+    x, weight, cos, sin = (t.contiguous() for t in (x, weight, cos, sin))
+    out = torch.empty((batch, rows, heads, dim), device=x.device, dtype=x.dtype)
+    codes, scales = out, out
+    if quantized:
+        codes = torch.empty((batch, heads, rows - start, 64), device=x.device, dtype=torch.uint8)
+        scales = torch.empty((batch, heads, rows - start, 8), device=x.device, dtype=torch.float8_e4m3fn)
+    from .qknorm_rope_vector import qknorm_rope_vector_kernel
+
+    unroll = 16 if quantized else 8
+    qknorm_rope_vector_kernel[(triton.cdiv(batch * rows * heads, unroll),)](
+        out, x, weight, cos, sin, dim, rotary_dim, rotary_dim // 2, heads, rows, eps,
+        BLOCK_SIZE=128, ROWS=batch * rows * heads, ROWS_PER_PROGRAM=unroll,
+        codes_ptr=codes, scales_ptr=scales, PACK_FP4=quantized, START=start,
+        input_index_ptr=input_index, INPUT_ROWS=x.shape[1], GATHER_INPUT=True, num_warps=4,
+    )
+    return out, (codes, scales) if quantized else None

fastvideo/models/dits/minimax_h3_fusions/qknorm_rope_vector.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Vectorized BF16 head-128 normalization with the original four-warp sum order.
+
+The scalar-row kernel reduces within each 32-lane warp first, then combines
+its four warp sums. A plain vectorized ``tl.sum`` chooses another tree and
+changes BF16 rounding at a few inputs. Explicit butterfly permutations keep
+that original order while processing multiple rows in registers together.
+"""
+import triton
+import triton.language as tl
+
+from .swiglu_quant import _rcp
+
+
+@triton.jit
+def qknorm_rope_vector_kernel(
+    out_ptr,
+    x_ptr,
+    weight_ptr,
+    cos_ptr,
+    sin_ptr,
+    head_dim: tl.constexpr,
+    rotary_dim: tl.constexpr,
+    half_rotary_dim: tl.constexpr,
+    num_heads: tl.constexpr,
+    seq_len: tl.constexpr,
+    eps,
+    BLOCK_SIZE: tl.constexpr,
+    ROWS: tl.constexpr,
+    ROWS_PER_PROGRAM: tl.constexpr,
+    codes_ptr=None,
+    scales_ptr=None,
+    PACK_FP4: tl.constexpr = False,
+    START: tl.constexpr = 0,
+    input_index_ptr=None,
+    INPUT_ROWS: tl.constexpr = 0,
+    GATHER_INPUT: tl.constexpr = False,
+):
+    # int64, like the sibling kernels: with int32 program ids,
+    # ``row * head_dim`` wraps once the flattened input reaches 2**31
+    # elements (H3's 56 heads x 128 head_dim crosses that at
+    # batch*seq >= 299_593 tokens per rank) and the loads/stores below
+    # become out-of-bounds. ``seq_index`` inherits int64 from ``row``.
+    tl.static_assert(BLOCK_SIZE == 128 and head_dim == 128)
+    row = tl.program_id(0).to(tl.int64) * ROWS_PER_PROGRAM + tl.arange(0, ROWS_PER_PROGRAM)[:, None]
+    seq_index = (row // num_heads) % seq_len
+    cols = tl.arange(0, BLOCK_SIZE)[None, :]
+    head_mask = (cols < head_dim) & (row < ROWS)
+    row_offset = row * head_dim
+    if GATHER_INPUT:
+        source = tl.load(input_index_ptr + seq_index, row < ROWS, other=-1).to(tl.int64)
+        input_offset = (((row // (seq_len * num_heads)) * INPUT_ROWS + source) * num_heads
+                        + row % num_heads) * head_dim
+        input_valid = source >= 0
+    else:
+        input_offset, input_valid = row_offset, True
+    x = tl.load(x_ptr + input_offset + cols, mask=head_mask & input_valid, other=0.0).to(tl.float32)
+    # Warp-local offsets 16..1, then the four-warp tree at offsets 64, 32.
+    squared = x * x
+    for shift in tl.static_range(5):
+        other = tl.gather(squared, tl.broadcast_to(cols ^ (16 >> shift), (ROWS_PER_PROGRAM, BLOCK_SIZE)), axis=1)
+        squared = squared + other
+    for shift in tl.static_range(2):
+        other = tl.gather(squared, tl.broadcast_to(cols ^ (64 >> shift), (ROWS_PER_PROGRAM, BLOCK_SIZE)), axis=1)
+        squared = squared + other
+    variance = tl.gather(squared, tl.full((ROWS_PER_PROGRAM, 1), 0, tl.int32), axis=1) / head_dim
+    inv_rms = tl.math.rsqrt(variance + eps)
+    weight = tl.load(weight_ptr + cols, mask=head_mask, other=0.0).to(tl.float32)
+    normalized = x * inv_rms * weight
+
+    rotary_mask = (cols < rotary_dim) & (row < ROWS)
+    first_half = cols < half_rotary_dim
+    partner_col = tl.where(first_half, cols + half_rotary_dim, cols - half_rotary_dim)
+    # Reuse the identical normalized partner instead of loading and normalizing it again.
+    partner_normalized = tl.gather(
+        normalized, tl.broadcast_to(partner_col, (ROWS_PER_PROGRAM, BLOCK_SIZE)), axis=1,
+    )
+    rotated = tl.where(first_half, -partner_normalized, partner_normalized)
+
+    table_offset = seq_index * rotary_dim + cols
+    cos = tl.load(cos_ptr + table_offset, mask=rotary_mask, other=1.0).to(tl.float32)
+    sin = tl.load(sin_ptr + table_offset, mask=rotary_mask, other=0.0).to(tl.float32)
+    rotary_output = normalized * cos + rotated * sin
+    output = tl.where(rotary_mask, rotary_output, normalized)
+    tl.store(out_ptr + row_offset + cols, output.to(out_ptr.dtype.element_ty), mask=head_mask)
+    if PACK_FP4:
+        rounded = output.to(tl.bfloat16).to(tl.float32)
+        grouped = tl.reshape(rounded, (ROWS_PER_PROGRAM, BLOCK_SIZE // 16, 16))
+        maximum = tl.max(tl.abs(grouped), 2)
+        sf = (maximum * (1.0 / 6.0)).to(tl.float8e4nv)
+        inv = tl.where(sf.to(tl.float32) == 0, 0, _rcp(sf.to(tl.float32)))
+        values = tl.reshape(grouped * inv[:, :, None], (ROWS_PER_PROGRAM, BLOCK_SIZE // 2, 2))
+        lo, hi = tl.split(values)
+        packed = tl.inline_asm_elementwise('''{
+            .reg .b8 q;
+            cvt.rn.satfinite.e2m1x2.f32 q, $1, $2;
+            mov.b32 $0, {q,q,q,q};
+        }''', constraints='=r,f,f', args=[hi, lo], dtype=tl.uint8, is_pure=True, pack=1)
+        batch = row // (seq_len * num_heads)
+        head = row % num_heads
+        c4 = tl.arange(0, BLOCK_SIZE // 2)[None, :]
+        offset = ((batch * num_heads + head) * (seq_len - START) + seq_index - START) * (head_dim // 2)
+        tl.store(codes_ptr + offset + c4, packed, (row < ROWS) & (seq_index >= START) & (c4 < head_dim // 2))
+        group = tl.arange(0, BLOCK_SIZE // 16)[None, :]
+        qi = seq_index - START
+        si = ((qi // 64) * (64 * head_dim // 16) + (group // 4) * 256 + group % 4
+              + (qi % 64 // 16) * 4 + (qi % 16) * 16)
+        scale_offset = (batch * num_heads + head) * (seq_len - START) * (head_dim // 16) + si
+        tl.store(scales_ptr + scale_offset, sf,
+                 (row < ROWS) & (seq_index >= START) & (group < head_dim // 16))
+
+

fastvideo/models/dits/minimax_h3_fusions/quant_tile.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Reuse one packed-input quantization for tile-order QKV and packed-order gate."""
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _tile_quantized(Q, S, Index, OQ, OS, WIDTH: tl.constexpr, ROWS: tl.constexpr, BLOCK: tl.constexpr):
+    row = tl.program_id(0).to(tl.int64)
+    col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
+    src = tl.load(Index + row).to(tl.int64)
+    val = tl.load(Q + src * (WIDTH // 2) + col, (src >= 0) & (col < WIDTH // 2), other=0)
+    tl.store(OQ + row * (WIDTH // 2) + col, val, col < WIDTH // 2)
+    group = tl.program_id(1) * (BLOCK // 8) + tl.arange(0, BLOCK // 8)
+    src_scale = ((((src // 128) * (WIDTH // 64) + group // 4) * 32 + src % 32) * 4 + (src // 32) % 4) * 4 + group % 4
+    dst_scale = ((((row // 128) * (WIDTH // 64) + group // 4) * 32 + row % 32) * 4 + (row // 32) % 4) * 4 + group % 4
+    scale = tl.load(S + src_scale, (src >= 0) & (group < WIDTH // 16), other=0)
+    tl.store(OS + dst_scale, scale, group < WIDTH // 16)
+
+
+def tile_quantized(q: torch.Tensor, s: torch.Tensor, inverse: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Gather NVFP4 rows, including 128x4-swizzled scales; -1 indices produce zero padding."""
+    rows, width = inverse.numel(), q.shape[-1] * 2
+    oq = torch.empty((rows, width // 2), device=q.device, dtype=q.dtype)
+    os = torch.empty((rows, width // 16), device=s.device, dtype=s.dtype)
+    _tile_quantized[(rows, triton.cdiv(width // 2, 512))](q, s, inverse, oq, os, width, rows, 512)
+    return oq, os
+

fastvideo/models/dits/minimax_h3_fusions/query_pool.py

@@ -0,0 +1,194 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Prepare exact pooled/FP4 queries while retaining BF16 only for dense prefix attention."""
+import torch
+import triton
+import triton.language as tl
+
+from .swiglu_quant import _rcp
+from .tile_pool import _tile_sum
+
+
+@triton.jit
+def _query_pool_kernel(
+    out_ptr,
+    x_ptr,
+    weight_ptr,
+    cos_ptr,
+    sin_ptr,
+    head_dim: tl.constexpr,
+    rotary_dim: tl.constexpr,
+    half_rotary_dim: tl.constexpr,
+    num_heads: tl.constexpr,
+    seq_len: tl.constexpr,
+    eps,
+    BLOCK_SIZE: tl.constexpr,
+    ROWS: tl.constexpr,
+    ROWS_PER_PROGRAM: tl.constexpr,
+    TILE_START: tl.constexpr = 0,
+    codes_ptr=None,
+    scales_ptr=None,
+    PACK_FP4: tl.constexpr = False,
+    START: tl.constexpr = 0,
+    input_index_ptr=None,
+    INPUT_ROWS: tl.constexpr = 0,
+    GATHER_INPUT: tl.constexpr = False,
+    pool_ptr=None,
+    valid_ptr=None,
+    PREFIX: tl.constexpr = 0,
+    PREFIX_ONLY: tl.constexpr = False,
+    PREFIX_FP8: tl.constexpr = False,
+):
+    # int64, like the sibling kernels: with int32 program ids,
+    # ``row * head_dim`` wraps once the flattened input reaches 2**31
+    # elements (H3's 56 heads x 128 head_dim crosses that at
+    # batch*seq >= 299_593 tokens per rank) and the loads/stores below
+    # become out-of-bounds. ``seq_index`` inherits int64 from ``row``.
+    tl.static_assert(BLOCK_SIZE == 128 and head_dim == 128)
+    tile = tl.program_id(0).to(tl.int64) + TILE_START
+    current_head = tl.program_id(1)
+    current_batch = tl.program_id(2)
+    seq_index = tile * ROWS_PER_PROGRAM + tl.arange(0, ROWS_PER_PROGRAM)[:, None]
+    row = (current_batch * seq_len + seq_index) * num_heads + current_head
+    cols = tl.arange(0, BLOCK_SIZE)[None, :]
+    head_mask = (cols < head_dim) & (row < ROWS)
+    row_offset = row * head_dim
+    if GATHER_INPUT:
+        source = tl.load(input_index_ptr + seq_index, row < ROWS, other=-1).to(tl.int64)
+        input_offset = (((row // (seq_len * num_heads)) * INPUT_ROWS + source) * num_heads
+                        + row % num_heads) * head_dim
+        input_valid = source >= 0
+    else:
+        input_offset, input_valid = row_offset, True
+    x = tl.load(x_ptr + input_offset + cols, mask=head_mask & input_valid, other=0.0).to(tl.float32)
+    # Warp-local offsets 16..1, then the four-warp tree at offsets 64, 32.
+    squared = x * x
+    for shift in tl.static_range(5):
+        other = tl.gather(squared, tl.broadcast_to(cols ^ (16 >> shift), (ROWS_PER_PROGRAM, BLOCK_SIZE)), axis=1)
+        squared = squared + other
+    for shift in tl.static_range(2):
+        other = tl.gather(squared, tl.broadcast_to(cols ^ (64 >> shift), (ROWS_PER_PROGRAM, BLOCK_SIZE)), axis=1)
+        squared = squared + other
+    variance = tl.gather(squared, tl.full((ROWS_PER_PROGRAM, 1), 0, tl.int32), axis=1) / head_dim
+    inv_rms = tl.math.rsqrt(variance + eps)
+    weight = tl.load(weight_ptr + cols, mask=head_mask, other=0.0).to(tl.float32)
+    normalized = x * inv_rms * weight
+
+    rotary_mask = (cols < rotary_dim) & (row < ROWS)
+    first_half = cols < half_rotary_dim
+    partner_col = tl.where(first_half, cols + half_rotary_dim, cols - half_rotary_dim)
+    # Reuse the identical normalized partner instead of loading and normalizing it again.
+    partner_normalized = tl.gather(
+        normalized, tl.broadcast_to(partner_col, (ROWS_PER_PROGRAM, BLOCK_SIZE)), axis=1,
+    )
+    rotated = tl.where(first_half, -partner_normalized, partner_normalized)
+
+    table_offset = seq_index * rotary_dim + cols
+    cos = tl.load(cos_ptr + table_offset, mask=rotary_mask, other=1.0).to(tl.float32)
+    sin = tl.load(sin_ptr + table_offset, mask=rotary_mask, other=0.0).to(tl.float32)
+    rotary_output = normalized * cos + rotated * sin
+    output = tl.where(rotary_mask, rotary_output, normalized)
+    if PREFIX_ONLY:
+        output_offset = ((current_batch * PREFIX + seq_index) * num_heads + current_head) * head_dim
+        if PREFIX_FP8:
+            if tile * ROWS_PER_PROGRAM < PREFIX:
+                prefix_value = output.to(tl.bfloat16).to(tl.float32)
+                scale8 = tl.max(tl.abs(prefix_value), 1) / 448.0
+                inv8 = tl.where(scale8 == 0, 0, 1.0 / scale8)
+                prefix_value = tl.minimum(tl.maximum(prefix_value * inv8[:, None], -448.0), 448.0)
+                prefix_value = prefix_value.to(tl.float8e4nv).to(tl.float32) * scale8[:, None]
+                tl.store(out_ptr + output_offset + cols, prefix_value, mask=head_mask & (seq_index < PREFIX))
+        else:
+            tl.store(out_ptr + output_offset + cols, output.to(out_ptr.dtype.element_ty),
+                     mask=head_mask & (seq_index < PREFIX))
+    else:
+        tl.store(out_ptr + row_offset + cols, output.to(out_ptr.dtype.element_ty), mask=head_mask)
+    rounded_pool = output.to(tl.bfloat16).to(tl.float32)
+    total = _tile_sum(rounded_pool, ROWS_PER_PROGRAM)
+    count = tl.load(valid_ptr + tile).to(tl.float32)
+    pool_offset = ((current_batch * (seq_len // ROWS_PER_PROGRAM) + tile) * num_heads + current_head) * 128
+    tl.store(pool_ptr + pool_offset + tl.arange(0, 128), tl.div_rn(total, count))
+    if PACK_FP4:
+        rounded = output.to(tl.bfloat16).to(tl.float32)
+        grouped = tl.reshape(rounded, (ROWS_PER_PROGRAM, BLOCK_SIZE // 16, 16))
+        maximum = tl.max(tl.abs(grouped), 2)
+        sf = (maximum * (1.0 / 6.0)).to(tl.float8e4nv)
+        inv = tl.where(sf.to(tl.float32) == 0, 0, _rcp(sf.to(tl.float32)))
+        values = tl.reshape(grouped * inv[:, :, None], (ROWS_PER_PROGRAM, BLOCK_SIZE // 2, 2))
+        lo, hi = tl.split(values)
+        packed = tl.inline_asm_elementwise('''{
+            .reg .b8 q;
+            cvt.rn.satfinite.e2m1x2.f32 q, $1, $2;
+            mov.b32 $0, {q,q,q,q};
+        }''', constraints='=r,f,f', args=[hi, lo], dtype=tl.uint8, is_pure=True, pack=1)
+        batch = row // (seq_len * num_heads)
+        head = row % num_heads
+        c4 = tl.arange(0, BLOCK_SIZE // 2)[None, :]
+        offset = ((batch * num_heads + head) * (seq_len - START) + seq_index - START) * (head_dim // 2)
+        tl.store(codes_ptr + offset + c4, packed, (row < ROWS) & (seq_index >= START) & (c4 < head_dim // 2))
+        group = tl.arange(0, BLOCK_SIZE // 16)[None, :]
+        qi = seq_index - START
+        si = ((qi // 64) * (64 * head_dim // 16) + (group // 4) * 256 + group % 4
+              + (qi % 64 // 16) * 4 + (qi % 16) * 16)
+        scale_offset = (batch * num_heads + head) * (seq_len - START) * (head_dim // 16) + si
+        tl.store(scales_ptr + scale_offset, sf,
+                 (row < ROWS) & (seq_index >= START) & (group < head_dim // 16))
+
+
+
+
+def prepare_pooled_queries(
+    x: torch.Tensor, weight: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, eps: float,
+    input_index: torch.Tensor | None, valid: torch.Tensor, prefix_rows: int, *, tile_elems: int = 64,
+    prefix_fp8: bool = False, split_prefix: bool = False,
+) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
+    """Return BF16 prefix queries, FP4 video queries, and exact FP32 tile means.
+
+    The native suffix starts at the preceding 128-token boundary, so its FP4
+    data also includes a half-prefix block when the prefix ends on tile 64.
+    An absent row map reads padded projection rows directly.
+    Omni tile-128 callers may also fuse the existing prefix FP8 round trip.
+    Splitting the prefix and suffix launches specializes away unused packing
+    and prefix stores while retaining every consumed value and tile sum.
+    Pooling and native codes always use the original BF16-normalized queries.
+    No BF16 video-query intermediate is needed by either attention branch.
+    """
+    from .qknorm_rope import _validate_inputs
+
+    if x.ndim != 4 or (input_index is not None and input_index.ndim != 1):
+        raise ValueError("Pooled queries require BSHD input and a one-dimensional row map")
+    rows = x.shape[1] if input_index is None else input_index.numel()
+    shape_view = x[:, :1].expand(-1, rows, -1, -1)
+    batch, _, heads, dim, rotary_dim = _validate_inputs(shape_view, weight, cos, sin, eps)
+    if (not x.is_cuda or x.dtype != torch.bfloat16 or dim != 128 or rows % 128
+            or tile_elems not in (64, 128) or (prefix_fp8 and tile_elems != 128)
+            or (input_index is not None and (input_index.device != x.device
+                or input_index.dtype not in (torch.int32, torch.int64) or not input_index.is_contiguous()))):
+        raise ValueError("Pooled queries require CUDA BF16 head-128 and an aligned integer row map")
+    if (valid.shape != (rows // tile_elems,) or valid.device != x.device
+            or valid.dtype not in (torch.int32, torch.int64) or not valid.is_contiguous()
+            or prefix_rows <= 0 or prefix_rows >= rows or prefix_rows % tile_elems):
+        raise ValueError("Pooled queries require tile validity and a nonempty aligned prefix")
+    if torch.is_grad_enabled() or torch.compiler.is_compiling():
+        raise RuntimeError("Pooled query preparation requires eager inference")
+    if split_prefix and (tile_elems != 128 or not prefix_fp8):
+        raise ValueError("Split query preparation requires tile-128 FP8 prefix queries")
+    x, weight, cos, sin = (t.contiguous() for t in (x, weight, cos, sin))
+    start = prefix_rows // 128 * 128
+    prefix = torch.empty((batch, prefix_rows, heads, dim), device=x.device, dtype=x.dtype)
+    codes = torch.empty((batch, heads, rows - start, 64), device=x.device, dtype=torch.uint8)
+    scales = torch.empty((batch, heads, rows - start, 8), device=x.device, dtype=torch.float8_e4m3fn)
+    pooled = torch.empty((batch, rows // tile_elems, heads, dim), device=x.device, dtype=torch.float32)
+    launches = ((0, rows // tile_elems, True, prefix_fp8, prefix_rows),)
+    if split_prefix:
+        prefix_tiles = prefix_rows // tile_elems
+        launches = ((0, prefix_tiles, False, True, prefix_rows),
+                    (prefix_tiles, rows // tile_elems - prefix_tiles, True, False, 0))
+    for first_tile, tiles, pack, fp8, keep_prefix in launches:
+        _query_pool_kernel[(tiles, heads, batch)](
+            prefix, x, weight, cos, sin, dim, rotary_dim, rotary_dim // 2, heads, rows, eps,
+            BLOCK_SIZE=128, ROWS=batch * rows * heads, ROWS_PER_PROGRAM=tile_elems, TILE_START=first_tile,
+            codes_ptr=codes, scales_ptr=scales, PACK_FP4=pack, START=start,
+            input_index_ptr=input_index, INPUT_ROWS=x.shape[1], GATHER_INPUT=input_index is not None,
+            pool_ptr=pooled, valid_ptr=valid, PREFIX=keep_prefix, PREFIX_ONLY=True, PREFIX_FP8=fp8, num_warps=8,
+        )
+    return prefix, (codes, scales), pooled.permute(0, 2, 1, 3)

fastvideo/models/dits/minimax_h3_fusions/sp_exchange.py

@@ -0,0 +1,194 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Dequantize an H3 Ulysses exchange directly into the VSA tile layout."""
+
+import torch
+import triton
+import triton.language as tl
+
+from fastvideo.distributed.async_collectives import PeerSlices
+
+from .swiglu_quant import _rcp as rcp
+
+
+@triton.jit
+def _unpack_fp8_tiles(Payload, Scale, Inverse, Out, LOCAL_ROWS: tl.constexpr,
+                      TILE_ROWS: tl.constexpr, HEADS: tl.constexpr, DIM: tl.constexpr, BLOCK: tl.constexpr,
+                      Codes=None, QueryScales=None, PACK_QUERY: tl.constexpr = False, START: tl.constexpr = 0):
+    tile_row = tl.program_id(0).to(tl.int64)
+    projection = tl.program_id(1)
+    column = tl.arange(0, BLOCK)
+    packed_row = tl.load(Inverse + tile_row).to(tl.int64)
+    source = packed_row // LOCAL_ROWS
+    local_row = packed_row % LOCAL_ROWS
+    source_row = (source * 3 + projection) * LOCAL_ROWS + local_row
+    valid = (packed_row >= 0) & (column < HEADS * DIM)
+    value = tl.load(Payload + source_row * HEADS * DIM + column, valid, other=0.0).to(tl.float32)
+    scale = tl.load(Scale + source_row * HEADS + column // DIM, valid, other=0)
+    decoded = (value * scale).to(tl.bfloat16)
+    destination = (projection * TILE_ROWS + tile_row) * HEADS * DIM + column
+    tl.store(Out + destination, decoded, column < HEADS * DIM)
+    if PACK_QUERY and projection == 0 and tile_row >= START:
+        grouped = tl.reshape(decoded.to(tl.float32), (BLOCK // 16, 16))
+        maximum = tl.max(tl.abs(grouped), 1)
+        sf = (maximum * (1.0 / 6.0)).to(tl.float8e4nv)
+        inverse = tl.where(sf.to(tl.float32) == 0, 0, rcp(sf.to(tl.float32)))
+        values = tl.reshape(grouped * inverse[:, None], (BLOCK // 2, 2))
+        lo, hi = tl.split(values)
+        packed = tl.inline_asm_elementwise('''{
+            .reg .b8 q;
+            cvt.rn.satfinite.e2m1x2.f32 q, $1, $2;
+            mov.b32 $0, {q,q,q,q};
+        }''', constraints="=r,f,f", args=[hi, lo], dtype=tl.uint8, is_pure=True, pack=1)
+        qrow = tile_row - START
+        qrows: tl.constexpr = TILE_ROWS - START
+        code_col = tl.arange(0, BLOCK // 2)
+        head, local_code = code_col // (DIM // 2), code_col % (DIM // 2)
+        tl.store(Codes + (head * qrows + qrow) * (DIM // 2) + local_code,
+                  packed, code_col < HEADS * DIM // 2)
+        group = tl.arange(0, BLOCK // 16)
+        head, local_group = group // (DIM // 16), group % (DIM // 16)
+        swizzled = ((qrow // 64) * (64 * DIM // 16) + (local_group // 4) * 256 + local_group % 4
+                    + (qrow % 64 // 16) * 4 + (qrow % 16) * 16)
+        tl.store(QueryScales + head * qrows * (DIM // 16) + swizzled, sf, group < HEADS * DIM // 16)
+
+
+def unpack_fp8_tiles(payload: torch.Tensor, scale: torch.Tensor, inverse: torch.Tensor,
+                     output: torch.Tensor, *, query_start: int | None = None
+                     ) -> torch.Tensor | tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+    """Map `[W, 3, local_rows, Hs, D]` FP8 receives to `[3, tile_rows, Hs, D]` BF16.
+
+    `inverse[tile_row]` is the global packed row, or -1 for VSA padding. Every
+    output element is written, including padding, so a reused buffer cannot
+    retain stale values. The float32 multiply and BF16 rounding match the
+    separate unpack and index-copy operations. An optional aligned query suffix
+    also receives native Sage FP4 codes and scales after that BF16 rounding.
+    """
+    _, projections, local_rows, heads, dim = payload.shape
+    assert projections == 3 and payload.dtype == torch.float8_e4m3fn
+    assert scale.dtype == torch.float32 and output.dtype == torch.bfloat16
+    assert payload.is_contiguous() and scale.is_contiguous() and output.is_contiguous()
+    assert output.shape == (3, inverse.numel(), heads, dim)
+    if query_start is not None:
+        if dim != 128 or inverse.numel() % 128 or query_start % 128 or not 0 <= query_start < inverse.numel():
+            raise ValueError("Fused query packing requires head-128 and an aligned token suffix")
+        query_rows = inverse.numel() - query_start
+        codes = torch.empty((1, heads, query_rows, dim // 2), device=payload.device, dtype=torch.uint8)
+        query_scales = torch.empty((1, heads, query_rows, dim // 16), device=payload.device, dtype=torch.float8_e4m3fn)
+    else:
+        codes, query_scales = output, output
+    block = triton.next_power_of_2(heads * dim)
+    _unpack_fp8_tiles[(inverse.numel(), 3)](
+        payload, scale, inverse, output, local_rows, inverse.numel(), heads, dim, block,
+        Codes=codes, QueryScales=query_scales, PACK_QUERY=query_start is not None,
+        START=query_start if query_start is not None else 0, num_warps=8 if block >= 2048 else 4,
+    )
+    return output if query_start is None else (output, (codes, query_scales))
+
+
+@triton.jit
+def _unpack_fp8_gate(Payload, Scale, Compressed, Gate, RowTile, Out, Codes, Scales, GlobalScale,
+                     ROWS: tl.constexpr, HEADS: tl.constexpr, WORLD: tl.constexpr,
+                     TILES: tl.constexpr, DIM: tl.constexpr, BLOCK: tl.constexpr, QUANTIZED: tl.constexpr, CHUNK: tl.constexpr, CHUNKS: tl.constexpr,
+                     Remote=None, PAIR: tl.constexpr = False, RANK: tl.constexpr = 0):
+    row = tl.program_id(0).to(tl.int64)
+    column = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
+    head, dim = column // DIM, column % DIM
+    source, local_head = head // HEADS, head % HEADS
+    source_head = (source * ROWS + row) * HEADS + local_head
+    valid = column < WORLD * HEADS * DIM
+    if PAIR:
+        local_offset = (row * HEADS + local_head) * DIM + dim
+        pointer = tl.where(source == RANK, Payload + local_offset, Remote + local_offset)
+    else:
+        pointer = Payload + source_head * DIM + dim
+    value = tl.load(pointer, valid, other=0.0).to(tl.float32)
+    scale = tl.load(Scale + source_head, valid, other=0.0)
+    # Retain the materialized unpack's BF16 rounding before applying the gate.
+    decoded = (value * scale).to(tl.bfloat16).to(tl.float32)
+    tile = tl.load(RowTile + row).to(tl.int64)
+    compressed = tl.load(Compressed + (head * TILES + tile) * DIM + dim, valid, other=0.0).to(tl.float32)
+    full_head = source * HEADS * CHUNKS + CHUNK * HEADS + local_head
+    destination = row * WORLD * HEADS * CHUNKS * DIM + full_head * DIM + dim
+    gate = tl.load(Gate + destination, valid, other=0.0).to(tl.float32)
+    if QUANTIZED:
+        width: tl.constexpr = WORLD * HEADS * CHUNKS * DIM
+        rounded = tl.where(valid, tl.fma(compressed, gate, decoded).to(tl.bfloat16).to(tl.float32), 0.0)
+        grouped = tl.reshape(rounded, (BLOCK // 16, 16))
+        maximum = tl.max(tl.abs(grouped), 1)
+        global_scale = tl.load(GlobalScale)
+        fp4_scale = (global_scale * (maximum * rcp(6.0))).to(tl.float8e4nv)
+        inverse = tl.where(maximum != 0, rcp(fp4_scale.to(tl.float32) * rcp(global_scale)), 0.0)
+        quantized = tl.reshape(grouped * inverse[:, None], (BLOCK // 2, 2))
+        lo, hi = tl.split(quantized)
+        packed = tl.inline_asm_elementwise('''{
+            .reg .b8 q;
+            cvt.rn.satfinite.e2m1x2.f32 q, $1, $2;
+            mov.b32 $0, {q,q,q,q};
+        }''', constraints="=r,f,f", args=[hi, lo], dtype=tl.uint8, is_pure=True, pack=1)
+        code_column = tl.program_id(1) * (BLOCK // 2) + tl.arange(0, BLOCK // 2)
+        code_head = code_column // (DIM // 2)
+        full_code_head = (code_head // HEADS) * HEADS * CHUNKS + CHUNK * HEADS + code_head % HEADS
+        full_code_column = full_code_head * (DIM // 2) + code_column % (DIM // 2)
+        tl.store(Codes + row * (width // 2) + full_code_column, packed, code_column < WORLD * HEADS * DIM // 2)
+        group = tl.program_id(1) * (BLOCK // 16) + tl.arange(0, BLOCK // 16)
+        group_head = group // (DIM // 16)
+        full_group_head = (group_head // HEADS) * HEADS * CHUNKS + CHUNK * HEADS + group_head % HEADS
+        full_group = full_group_head * (DIM // 16) + group % (DIM // 16)
+        scale_tile = (row // 128) * (width // 64) + full_group // 4
+        scale_index = ((scale_tile * 32 + row % 32) * 4 + (row // 32) % 4) * 4 + full_group % 4
+        tl.store(Scales + scale_index, fp4_scale.to(tl.uint8, bitcast=True), group < WORLD * HEADS * DIM // 16)
+    else:
+        tl.store(Out + destination, tl.fma(compressed, gate, decoded), valid)
+
+
+def unpack_fp8_gate(payload: torch.Tensor | PeerSlices, scale: torch.Tensor, compressed: torch.Tensor,
+                     gate: torch.Tensor, row_tile: torch.Tensor, global_scale: torch.Tensor | None = None,
+                     *, head_chunk: int = 0, head_chunks: int = 1, output: torch.Tensor | None = None,
+                     fp4_output: tuple[torch.Tensor, torch.Tensor] | None = None,
+                     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Unpack head groups and apply the local gate without a BF16 intermediate.
+
+    Receives are `[W, rows, Hs, D]`; the gathered compression branch is
+    `[W, Hs, tiles, D]`, and the local gate/result are `[rows, W*Hs, D]`.
+    A chunked exchange writes only its corresponding heads in a shared full
+    output, preserving rank-major then chunk-major head order without copies.
+    """
+    pair = payload if isinstance(payload, PeerSlices) else None
+    payload = pair.template if pair is not None else payload
+    world, rows, heads, dim = payload.shape
+    assert payload.dtype == torch.float8_e4m3fn and scale.dtype == torch.float32
+    assert compressed.dtype == gate.dtype == torch.bfloat16
+    assert all(t.is_contiguous() for t in (payload, scale, compressed, gate, row_tile))
+    assert head_chunks > 0 and 0 <= head_chunk < head_chunks
+    assert gate.shape == (rows, world * heads * head_chunks, dim)
+    assert compressed.shape[:2] == (world, heads) and compressed.shape[-1] == dim
+    quantized = global_scale is not None
+    if quantized:
+        if output is not None or (head_chunks > 1 and fp4_output is None):
+            raise ValueError("Chunked FP4 preparation requires shared codes/scales and no BF16 output")
+        width = world * heads * head_chunks * dim
+        if width % 64 or dim % 16:
+            raise ValueError("Fused SP gate quantization requires a width divisible by 64")
+        if global_scale.numel() != 1 or global_scale.dtype != torch.float32 or global_scale.device != gate.device:
+            raise ValueError("Fused SP gate quantization requires a scalar FP32 scale on the input device")
+        if fp4_output is None:
+            codes = torch.empty((rows, width // 2), device=gate.device, dtype=torch.uint8)
+            fp4_scales = torch.zeros((triton.cdiv(rows, 128) * 128, width // 16), device=gate.device, dtype=torch.uint8)
+        else:
+            codes, fp4_scales = fp4_output
+            assert codes.shape == (rows, width // 2)
+            assert fp4_scales.shape == (triton.cdiv(rows, 128) * 128, width // 16)
+            assert all(t.dtype == torch.uint8 and t.device == gate.device and t.is_contiguous()
+                       for t in fp4_output)
+        output = gate  # Unused pointer when the kernel writes FP4 codes directly.
+    else:
+        if output is None:
+            output = torch.empty_like(gate)
+        assert output.shape == gate.shape and output.dtype == gate.dtype and output.is_contiguous()
+        codes, fp4_scales, global_scale = output, output, output
+    _unpack_fp8_gate[(rows, triton.cdiv(world * heads * dim, 512))](
+        pair.local if pair is not None else payload, scale, compressed, gate, row_tile, output, codes, fp4_scales, global_scale,
+        rows, heads, world, compressed.shape[2], dim, 512, quantized, head_chunk, head_chunks,
+        Remote=pair.remote if pair is not None else payload, PAIR=pair is not None, RANK=pair.rank if pair is not None else 0,
+    )
+    return (codes, fp4_scales, global_scale) if quantized else output

fastvideo/models/dits/minimax_h3_fusions/sp_head_pipeline.py

@@ -0,0 +1,250 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Overlap independent H3 attention head groups with FP8 Ulysses exchanges."""
+
+import torch
+
+from fastvideo.distributed.async_collectives import (
+    PeerSlices, all_gather_tensor_async, all_to_all_tensor_async, exchange_peer_tensor_async,
+)
+
+from fastvideo.attention.backends.video_sparse_attn_h3 import _build_block_mask, _pool_tiles
+from fastvideo.models.dits import minimax_h3_vsa_fp4 as base
+from fastvideo.models.dits.minimax_h3_fusions.sp_exchange import unpack_fp8_gate
+from fastvideo.models.dits.minimax_h3_fusions.sp_qkv_pack import pack_normalized_qkv_fp8
+
+
+def _start_exchange(payload, scale, group):
+    if payload.shape[0] == 2 and payload.shape[-1] == 128:
+        packet = exchange_peer_tensor_async(payload, group)
+    else:
+        packet = all_to_all_tensor_async(payload.view(torch.uint8), group)
+    return packet, all_to_all_tensor_async(scale, group, align_bytes=16)
+
+
+def _finish_exchange(packet):
+    payload, scale = packet
+    received = payload.wait()
+    if not isinstance(received, PeerSlices):
+        received = received.view(torch.float8_e4m3fn)
+    return received, scale.wait()
+
+
+def _output_quant_method(attn):
+    from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod
+
+    method = getattr(attn.to_out, "quant_method", None)
+    if (type(method) is NVFP4QuantizeMethod and method.wants_prequantized_input()
+            and method.uses_unit_activation_scale(attn.to_out) and attn.attention_head_dim % 64 == 0):
+        return method
+    return None
+
+
+def _prepare_tiles(api, layout, payload, scale, query_start):
+    if (payload.shape[-1] == 128 and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
+            and torch.cuda.get_device_capability(payload.device)[0] >= 10
+            and all(hasattr(api, name) for name in ("blockscaled_fp4_attn_sparse", "_zero_delta_s",
+                                                    "scale_and_quant_fp4_permute", "scale_and_quant_fp4_transpose"))):
+        return layout.tiles_from_fp8(payload, scale, query_start=query_start)
+    return layout.tiles_from_fp8(payload, scale), None
+
+
+def _prepare_values(api, value, valid, tile, logical):
+    if (tile == 64 and logical == value.shape[1] and value.shape[-1] == 128 and value.stride(-1) == 1
+            and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
+            and torch.cuda.get_device_capability(value.device)[0] >= 10
+            and all(hasattr(api, name) for name in ("blockscaled_fp4_attn_sparse", "_zero_delta_s",
+                                                    "scale_and_quant_fp4", "scale_and_quant_fp4_permute"))):
+        from fastvideo.models.dits.minimax_h3_fusions.value_quant import pool_quantize_values
+        return pool_quantize_values(value, valid)
+    return None, None
+
+
+def _prepare_received(api, layout, payload, scale, query_start, valid, *, prefix_values=False, direct_keys=False):
+    template = payload.template if isinstance(payload, PeerSlices) else payload
+    if (layout.tile == 64 and layout.n_tiles * layout.tile == layout.rows and template.shape[-1] == 128
+            and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
+            and torch.cuda.get_device_capability(template.device)[0] >= 10
+            and all(hasattr(api, name) for name in ("blockscaled_fp4_attn_sparse", "_zero_delta_s",
+                                                    "scale_and_quant_fp4", "scale_and_quant_fp4_permute",
+                                                    "scale_and_quant_fp4_transpose"))):
+        from .sp_receive import unpack_pool_quantize_fp8
+        inverse, qkv = layout.fp8_buffers(template)
+        prefix_value, prefix = None, None
+        # The dense branch consumes gathered values and the sparse branch
+        # consumes native V codes, so raw BF16 V need not be materialized.
+        if prefix_values:
+            valid_index, prefix_inverse = layout.valid_keys_from_tiles(valid, api.HALF_N)
+            prefix_value = torch.empty((1, valid_index.numel(), *qkv.shape[2:]),
+                                       device=qkv.device, dtype=qkv.dtype)
+            prefix = (prefix_inverse, prefix_value)
+        # Native FP4 query codes supply the suffix; the raw query suffix is
+        # used only for shape/dtype metadata. Keep a full block beyond the
+        # aligned start to cover a dense prefix ending halfway through it.
+        pooled, prepared_value, prepared_query = unpack_pool_quantize_fp8(
+            payload, scale, inverse, valid, qkv, query_start,
+            query_rows=min(query_start + base._BLOCK, layout.rows), prefix_values=prefix, store_keys=not direct_keys,
+        )
+        return qkv, prepared_query, pooled, prepared_value, prefix_value
+    if isinstance(payload, PeerSlices):
+        payload = payload.materialize()
+    qkv, prepared_query = _prepare_tiles(api, layout, payload, scale, query_start)
+    return qkv, prepared_query, None, None, None
+
+
+def _pack_output(out, layout, world, rows):
+    if isinstance(out, tuple):
+        from .sp_output_pack import pack_split_fp8
+        return pack_split_fp8(out, layout.untile, world, rows)
+    return base._pack_seq_fp8(out, layout.untile, world, rows)
+
+
+def attention_pipelined(attn, hidden_states, rotary_emb, meta, sp_group, chunks, pre_quantized=None):
+    """Pipeline head groups so communication overlaps independent attention work."""
+    world, rank = sp_group.world_size, sp_group.rank_in_group
+    rows, heads, dim = hidden_states.shape[1], attn.num_attention_heads, attn.attention_head_dim
+    assert heads % (world * chunks) == 0
+    assert attn.to_gate_compress is not None and attn._gate_active()
+    layouts = getattr(meta, '_h3_fp4_pipeline_layouts', None)
+    if layouts is None or len(layouts) != chunks:
+        layouts = [base._SPTileLayout(meta, rank, rows) for _ in range(chunks)]
+        meta._h3_fp4_pipeline_layouts = layouts
+    options = getattr(attn, '_vsa_fp4_options', {})
+    early_sends = (options.get('sp_projection_overlap', False) and world == 2 and dim == 128
+                   and hidden_states.is_cuda and hidden_states.dtype == torch.bfloat16
+                   and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
+                   and base._shared_qkv_gate_scale(attn, hidden_states) is not None)
+    input_gather = False
+    if early_sends and options.get('sp_input_gather', False):
+        from .sp_input_gather import supports_input_gather
+        input_gather = supports_input_gather(attn, hidden_states)
+    if input_gather:
+        from .sp_input_gather import project_gathered_inputs
+        with base.STAGES.span('pipe_project_pack'):
+            packets, gate = project_gathered_inputs(attn, hidden_states, rotary_emb, sp_group, chunks, pre_quantized)
+    elif early_sends:
+        from .sp_projection import project_and_exchange
+        with base.STAGES.span('pipe_project_pack'):
+            packets, gate_pre = project_and_exchange(attn, hidden_states, rotary_emb, sp_group.device_group,
+                                                     chunks, pre_quantized)
+    else:
+        with base.STAGES.span('pipe_project_pack'):
+            shared = base._shared_qkv_gate_projections(attn, hidden_states, pre_quantized=pre_quantized)
+            if shared is None:
+                projected = base._shared_input_projections((attn.to_q, attn.to_k, attn.to_v), hidden_states)
+                gate_pre = None
+            else:
+                projected, gate_pre = shared
+            query, key, value = (t.unflatten(-1, (heads, dim)) for t in projected)
+            cos, sin = (t.to(query.dtype) for t in rotary_emb)
+            payload, scale = pack_normalized_qkv_fp8(
+                query[0],
+                key[0],
+                value[0],
+                attn.norm_q.weight,
+                attn.norm_k.weight,
+                cos,
+                sin,
+                attn.norm_q.eps,
+                attn.norm_k.eps,
+                world,
+                chunks,
+            )
+        packets = [_start_exchange(payload[chunk], scale[chunk], sp_group.device_group) for chunk in range(chunks)]
+    if not input_gather:
+        with base.STAGES.span('gate_proj'):
+            if gate_pre is None:
+                gate, _ = attn.to_gate_compress(hidden_states)
+            else:
+                linear = attn.to_gate_compress
+                gate = linear.quant_method.apply(linear, hidden_states, linear.bias, pre_quantized=gate_pre)
+    vbs = meta.variable_block_sizes
+    api = base._api()
+    pending = []
+    for chunk, layout in enumerate(layouts):
+        prefix_mode, prefix_rows = base._prefix_settings(attn, meta, layout.n_tiles)
+        direct_keys = bool(options.get('sp_direct_keys', False) and options.get('sp_pooled_key_mean', False)
+                           and options.get('smooth_k') is True and prefix_rows and prefix_mode == 'fp8')
+        with base.STAGES.span('pipe_receive'):
+            received, received_scale = _finish_exchange(packets[chunk])
+            qkv, prepared_query, pooled_qkv, prepared_value, prefix_value = _prepare_received(
+                api, layout, received, received_scale, prefix_rows // base._BLOCK * base._BLOCK, vbs,
+                prefix_values=bool(prefix_rows and prefix_mode == 'fp8' and options.get('smooth_k') is True),
+                direct_keys=direct_keys,
+            )
+        q_t, k_t, v_t = qkv[0:1], qkv[1:2], qkv[2:3]
+        logical = layout.n_tiles * layout.tile
+        with base.STAGES.span('pipe_attention'):
+            q_pool = pooled_qkv[0:1] if pooled_qkv is not None else _pool_tiles(q_t[:, :logical], vbs, layout.tile)
+            k_pool = pooled_qkv[1:2] if pooled_qkv is not None else _pool_tiles(k_t[:, :logical], vbs, layout.tile)
+            scores = torch.matmul(q_pool, k_pool.transpose(-2, -1)) / (dim**0.5)
+            sparsity = 0.0 if attn._layer_idx in meta.dense_layers else meta.VSA_sparsity
+            mask = _build_block_mask(scores, meta.num_prefix_tiles, sparsity, meta.exempt, meta.video_tile_spans,
+                                     meta.span_sparsities)
+            idx, num, valid, quad = base._tile_mask_to_fp4_blocks(api, mask, layout.tile, vbs)
+            if pooled_qkv is None:
+                pooled, prepared_value = _prepare_values(api, v_t, vbs, layout.tile, logical)
+            else:
+                pooled = pooled_qkv[2:3]
+            prepared_keys = None
+            if ((options.get('sp_pooled_key_mean', False) or prefix_value is not None)
+                    and options.get('smooth_k') is True
+                    and pooled_qkv is not None and prefix_rows and prefix_mode == 'fp8'):
+                from .key_quant import quantize_smoothed_keys
+                from .sp_key_mean import key_mean_from_tiles
+                valid_index = layout.valid_key_index(valid, api.HALF_N)
+                inverse = layout.valid_key_inverse(valid, api.HALF_N)
+                mean = (key_mean_from_tiles(pooled_qkv[1:2], vbs, logical)
+                        if options.get('sp_pooled_key_mean', False) else k_t.mean(dim=1, keepdim=True))
+                key_packet = None
+                if direct_keys and prefix_value is not None:
+                    template = received.template if isinstance(received, PeerSlices) else received
+                    tile_inverse, _ = layout.fp8_buffers(template)
+                    key_packet = (received, received_scale, tile_inverse)
+                prepared_keys = quantize_smoothed_keys(k_t, v_t, mean, inverse, valid_index.numel(),
+                                                       prefix_value=prefix_value, received_keys=key_packet)
+            out = base._fp4_fine_attention(api,
+                                           q_t,
+                                           k_t,
+                                           v_t,
+                                           idx,
+                                           num,
+                                           valid,
+                                           quad,
+                                           logical=logical,
+                                           prefix_rows=prefix_rows,
+                                           precision=prefix_mode,
+                                           options=options,
+                                           layout=layout, prepared_query=prepared_query, prepared_value=prepared_value,
+                                           prepared_keys=prepared_keys, return_split=True)
+        with base.STAGES.span('pipe_output_pack'):
+            out_payload, out_scale = _pack_output(out, layout, world, rows)
+        # Queue this output before computing the next head group. The pending
+        # handles retain their inputs until the final wait and gate application.
+        out_packet = _start_exchange(out_payload, out_scale, sp_group.device_group)
+        with base.STAGES.span('gate_compress'):
+            if pooled is None:
+                pooled = _pool_tiles(v_t[:, :logical], vbs, layout.tile)
+            compressed = torch.matmul(torch.softmax(scores, dim=-1), pooled)[0].to(qkv.dtype).contiguous()
+        gathered = all_gather_tensor_async(compressed, sp_group.device_group, world)
+        pending.append((out_packet, gathered))
+    full_gate = gate[0].view(rows, heads, dim)
+    method = _output_quant_method(attn)
+    if method is None:
+        full_output = torch.empty_like(full_gate)
+    else:
+        width = heads * dim
+        codes = torch.empty((rows, width // 2), device=gate.device, dtype=torch.uint8)
+        scales = torch.zeros((((rows + 127) // 128) * 128, width // 16), device=gate.device, dtype=torch.uint8)
+    for chunk, (packet, gathered) in enumerate(pending):
+        with base.STAGES.span('pipe_output_wait_gate'):
+            received, received_scale = _finish_exchange(packet)
+            compression = gathered.wait()
+            if method is None:
+                unpack_fp8_gate(received, received_scale, compression, full_gate, layouts[chunk].local_row_tile,
+                                 head_chunk=chunk, head_chunks=chunks, output=full_output)
+            else:
+                unpack_fp8_gate(received, received_scale, compression, full_gate, layouts[chunk].local_row_tile,
+                                 method.x_global_sf, head_chunk=chunk, head_chunks=chunks, fp4_output=(codes, scales))
+    if method is not None:
+        return gate, (codes, scales, method.x_global_sf)
+    return full_output.flatten(1, 2).unsqueeze(0)

fastvideo/models/dits/minimax_h3_fusions/sp_input_gather.py

@@ -0,0 +1,141 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Gather exact packed inputs before head-sharded two-rank H3 projections.
+
+The packed inputs already carry the shared unit activation scale. Only weight
+views are sharded; each row/head computes the same Q/K/V projection, norm/RoPE,
+and transport FP8 rounding as the sequence-sharded producer. Scale reassembly
+handles the unaligned join between independently padded 128-row scale tiles.
+"""
+import torch
+import triton
+import triton.language as tl
+
+from fastvideo.distributed.async_collectives import all_gather_tensor_async
+from fastvideo.distributed.projection_exchange import completed_projection_packets
+from fastvideo.layers.quantization.nvfp4_config import _mm_fp4, _mm_fp4_backend
+
+from .sp_projection import allocate
+from .sp_qkv_ordered import _normalize
+
+
+@triton.jit
+def _join_scale_rows(Source, Out, ROWS: tl.constexpr, PADDED: tl.constexpr, GROUPS: tl.constexpr,
+                     TOTAL_PADDED: tl.constexpr, BLOCK: tl.constexpr):
+    index = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    row, group = index // GROUPS, index % GROUPS
+    rank, local_row = row // ROWS, row % ROWS
+    local_tile = (local_row // 128) * (GROUPS // 4) + group // 4
+    local_index = ((local_tile * 32 + local_row % 32) * 4 + (local_row // 32) % 4) * 4 + group % 4
+    value = tl.load(Source + rank * PADDED * GROUPS + local_index, row < 2 * ROWS, other=0)
+    tile = (row // 128) * (GROUPS // 4) + group // 4
+    destination = ((tile * 32 + row % 32) * 4 + (row // 32) % 4) * 4 + group % 4
+    tl.store(Out + destination, value, index < TOTAL_PADDED * GROUPS)
+
+
+def join_scale_rows(source: torch.Tensor, rows: int) -> torch.Tensor:
+    assert source.ndim == 3 and source.shape[0] == 2 and source.is_contiguous()
+    assert source.dtype == torch.uint8 and source.shape[2] % 4 == 0
+    padded, groups = source.shape[1:]
+    assert padded == triton.cdiv(rows, 128) * 128
+    total_padded = triton.cdiv(rows * 2, 128) * 128
+    output = torch.empty((total_padded, groups), device=source.device, dtype=source.dtype)
+    _join_scale_rows[(triton.cdiv(total_padded * groups, 512), )](
+        source, output, rows, padded, groups, total_padded, 512, num_warps=4)
+    return output
+
+
+@triton.jit
+def _pack_gathered(X, Weight, Cos, Sin, Payload, Scale, LOCAL_ROWS: tl.constexpr, HEADS: tl.constexpr,
+                   CHUNKS: tl.constexpr, ROTARY: tl.constexpr, EPS: tl.constexpr, PROJECTION: tl.constexpr,
+                   BLOCK_ROWS: tl.constexpr):
+    row_head = (tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)).to(tl.int64)
+    row, head = row_head // HEADS, row_head % HEADS
+    column = tl.arange(0, 128)
+    valid = row[:, None] < LOCAL_ROWS * 2
+    if PROJECTION == 2:
+        values = tl.load(X + row_head[:, None] * 128 + column[None, :], valid, other=0.0).to(tl.float32)
+    else:
+        values = _normalize(X, Weight, Cos, Sin, row, head, HEADS, LOCAL_ROWS * 2, 128, ROTARY, EPS, 128)
+    maximum = tl.max(tl.abs(values), 1)
+    scale = tl.maximum(maximum * (1.0 / 448.0), 1e-12)
+    quantized = (values / scale[:, None]).to(tl.float8e4nv)
+    chunk_heads: tl.constexpr = HEADS // CHUNKS
+    chunk = head // chunk_heads
+    source = row // LOCAL_ROWS
+    destination = ((
+        (chunk * 2 + source) * 3 + PROJECTION) * LOCAL_ROWS + row % LOCAL_ROWS) * chunk_heads + head % chunk_heads
+    tl.store(Payload + destination[:, None] * 128 + column[None, :], quantized, valid)
+    tl.store(Scale + destination, scale, row < LOCAL_ROWS * 2)
+
+
+def pack_gathered(x, payload, scale, projection, weight=None, cosine=None, sine=None, epsilon=0.0) -> None:
+    chunks, world, projections, rows, heads, dim = payload.shape
+    assert world == 2 and projections == 3 and dim == 128
+    assert x.shape == (2 * rows, chunks * heads * dim) and x.is_contiguous()
+    assert payload.is_contiguous() and scale.is_contiguous()
+    if projection == 2:
+        weight, cosine, sine = x, x, x
+        rotary = 0
+    else:
+        assert weight.shape == (128, ) and cosine.shape == sine.shape
+        assert cosine.shape[0] == rows * 2 and cosine.is_contiguous() and sine.is_contiguous()
+        rotary = cosine.shape[1]
+    block_rows = 16 if projection == 2 else 8
+    _pack_gathered[(triton.cdiv(2 * rows * chunks * heads, block_rows), )](
+        x, weight, cosine, sine, payload, scale, rows, chunks * heads, chunks,
+        rotary, epsilon, projection, block_rows, num_warps=4 if projection == 2 else 8)
+
+
+def supports_input_gather(attn, hidden: torch.Tensor) -> bool:
+    """Keep other shapes/backends on the established projection exchange."""
+    if (not hidden.is_cuda or hidden.dtype != torch.bfloat16
+            or hidden.shape != (1, 7430, 5376) or attn.num_attention_heads != 56
+            or torch.cuda.get_device_capability(hidden.device) != (12, 1)
+            or _mm_fp4_backend() != "cudnn"):
+        return False
+    for linear in (attn.to_q, attn.to_k, attn.to_v):
+        weight = getattr(linear, "_nvfp4_weight", None)
+        scale = getattr(linear, "_nvfp4_weight_scale", None)
+        if (weight is None or scale is None or weight.shape != (7168, 2688)
+                or scale.shape != (7168, 336) or linear.bias is not None
+                or weight.dtype != torch.uint8 or scale.dtype != torch.uint8
+                or weight.device != hidden.device or scale.device != hidden.device
+                or not weight.is_contiguous() or not scale.is_contiguous()):
+            return False
+    return True
+
+
+def project_gathered_inputs(attn, hidden, rotary, sp_group, chunks, pre_quantized=None):
+    """Return completed head packets and the gate computed while inputs travel.
+
+    The caller checks eager, two-rank, head-128 inference and the common unit
+    activation-scale contract through _shared_qkv_gate_scale. No activations
+    or packets survive this call except through the returned consumer handles.
+    """
+    linears = (attn.to_q, attn.to_k, attn.to_v)
+    pre = pre_quantized if pre_quantized is not None else linears[0].quant_method.quantize_input(hidden)
+    a, sa, input_scale = pre
+    a, sa = a.reshape(-1, a.shape[-1]), sa.reshape(-1, sa.shape[-1])
+    cosine, sine = (item.to(hidden.dtype) for item in rotary)
+    pending = [all_gather_tensor_async(value, sp_group.device_group, 2) for value in (a, sa, cosine, sine)]
+    linear = attn.to_gate_compress
+    gate = linear.quant_method.apply(linear, hidden, linear.bias, pre_quantized=pre)
+    gathered_a, gathered_sa, gathered_cosine, gathered_sine = (item.wait() for item in pending)
+    full_a = gathered_a.flatten(0, 1)
+    full_sa = join_scale_rows(gathered_sa, a.shape[0])
+    full_cosine, full_sine = gathered_cosine.flatten(0, 1), gathered_sine.flatten(0, 1)
+    rows, heads, rank = hidden.shape[1], attn.num_attention_heads, sp_group.rank_in_group
+    payload, scale = allocate(rows, heads, 2, chunks, hidden.device)
+    for projection, linear in enumerate(linears):
+        half = linear._nvfp4_weight.shape[0] // 2
+        weight = linear._nvfp4_weight.narrow(0, rank * half, half)
+        weight_scale = linear._nvfp4_weight_scale.narrow(0, rank * half, half)
+        alpha = (linear._nvfp4_alpha / input_scale if hasattr(linear, "_nvfp4_alpha")
+                 else 1.0 / (input_scale * linear._weight_global_sf))
+        value = _mm_fp4(full_a, weight.T, full_sa, weight_scale.T, alpha, torch.bfloat16, None, backend="cudnn")
+        if projection == 2:
+            pack_gathered(value, payload, scale, projection)
+        else:
+            norm = attn.norm_q if projection == 0 else attn.norm_k
+            pack_gathered(value, payload, scale, projection, norm.weight, full_cosine, full_sine, norm.eps)
+    return completed_projection_packets(payload, scale, rank), gate

fastvideo/models/dits/minimax_h3_fusions/sp_key_mean.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Recover global key means from the SP receive's complete tile statistics.
+
+Every tile contributes its FP32 mean times its valid count, and the total is
+normalized by the padded sequence length, matching the zero-padded key tensor.
+This changes floating-point reduction order and may alter the generated sample.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _key_mean(Pool, Counts, Mean, HEAD_STRIDE: tl.constexpr, TILE_STRIDE: tl.constexpr,
+              TILES: tl.constexpr, SEQ: tl.constexpr, ROWS: tl.constexpr, COLS: tl.constexpr):
+    head = tl.program_id(0)
+    rows = tl.arange(0, ROWS)
+    cols = tl.program_id(1) * COLS + tl.arange(0, COLS)
+    pooled = tl.load(Pool + head * HEAD_STRIDE + rows[:, None] * TILE_STRIDE + cols[None, :],
+                     rows[:, None] < TILES, other=0.0)
+    count = tl.load(Counts + rows, rows < TILES, other=0).to(tl.float32)
+    pooled = tl.where(count[:, None] > 0, pooled, 0.0)
+    total = tl.sum(pooled * count[:, None], 0)
+    mean = total * (1.0 / SEQ)
+    tl.store(Mean + head * 128 + cols, mean)
+
+
+def key_mean_from_tiles(pool: torch.Tensor, counts: torch.Tensor, sequence: int) -> torch.Tensor:
+    """Return BF16 ``[1, 1, H, 128]`` means from FP32 ``[1, H, tiles, 128]``.
+
+    Counts must be the same divisors used to produce ``pool``. Padding is
+    included in ``sequence`` and contributes zero, just as in ``key.mean``.
+    """
+    batch, heads, tiles, dim = pool.shape
+    if (batch != 1 or dim != 128 or sequence != tiles * 64 or pool.stride(-1) != 1
+            or pool.dtype != torch.float32 or counts.shape != (tiles,)
+            or counts.dtype not in (torch.int32, torch.int64) or not pool.is_cuda
+            or counts.device != pool.device):
+        raise ValueError("Pooled SP key means require CUDA tile-64 FP32 statistics with 128 channels")
+    result = torch.empty((1, 1, heads, dim), device=pool.device, dtype=torch.bfloat16)
+    _key_mean[(heads, 4)](
+        pool, counts, result, pool.stride(1), pool.stride(2), tiles, sequence,
+        triton.next_power_of_2(tiles), 32, num_warps=4,
+    )
+    return result

fastvideo/models/dits/minimax_h3_fusions/sp_output_pack.py

@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pack SP attention's split prefix/video results without a concatenation buffer."""
+import torch
+import triton
+import triton.language as tl
+
+@triton.jit
+def _pack_split(Prefix, Video, Untile, Payload, Scale, LENGTH: tl.constexpr, ROWS: tl.constexpr, H: tl.constexpr,
+                D: tl.constexpr, PREFIX: tl.constexpr, VIDEO: tl.constexpr, START: tl.constexpr, BLOCK_H: tl.constexpr,
+                BLOCK_D: tl.constexpr, PSH: tl.constexpr, PSS: tl.constexpr, VSH: tl.constexpr, VSS: tl.constexpr):
+    row_head = tl.program_id(0) * BLOCK_H + tl.arange(0, BLOCK_H)
+    column = tl.arange(0, BLOCK_D)
+    row, head = row_head // H, row_head % H
+    tile_row = tl.load(Untile + row, row < LENGTH, other=-1)
+    valid = (row_head[:, None] < ROWS * H) & (tile_row[:, None] >= 0) & (column[None, :] < D)
+    is_prefix = tile_row < PREFIX
+    prefix = tl.load(Prefix + head[:, None] * PSH + tile_row[:, None] * PSS + column[None, :],
+                     valid & is_prefix[:, None],
+                     other=0).to(tl.float32)
+    video = tl.load(Video + head[:, None] * VSH + (tile_row[:, None] - START) * VSS + column[None, :],
+                    valid & ~is_prefix[:, None],
+                    other=0).to(tl.float32)
+    value = tl.where(is_prefix[:, None], prefix, video)
+    scale = tl.maximum(tl.max(tl.abs(value), 1) / 448.0, 1e-12)
+    quant = (value / scale[:, None]).to(tl.float8e4nv)
+    tl.store(Payload + row_head[:, None] * D + column[None, :], quant,
+             (row_head[:, None] < ROWS * H) & (column[None, :] < D))
+    tl.store(Scale + row_head, scale, row_head < ROWS * H)
+
+
+def pack_split_fp8(parts, untile, world, rows, block=4) -> tuple[torch.Tensor, torch.Tensor]:
+    prefix, video, start = parts
+    _, heads, _, dim = video.shape
+    assert prefix.dtype == video.dtype == torch.bfloat16 and prefix.shape[0] == video.shape[0] == 1
+    assert prefix.shape[1] == heads and prefix.shape[3] == dim and prefix.stride(-1) == video.stride(-1) == 1
+    assert untile.is_contiguous() and untile.dtype in (torch.int32, torch.int64)
+    assert prefix.device == video.device == untile.device and untile.numel() <= world * rows
+    assert 0 <= start <= prefix.shape[2] and video.is_cuda
+    payload = torch.empty((world, rows, heads, dim), device=video.device, dtype=torch.float8_e4m3fn)
+    scale = torch.empty(payload.shape[:-1], device=video.device, dtype=torch.float32)
+    _pack_split[(triton.cdiv(world * rows * heads, block), )](prefix,
+                                                              video,
+                                                              untile,
+                                                              payload,
+                                                              scale,
+                                                              untile.numel(),
+                                                              world * rows,
+                                                              heads,
+                                                              dim,
+                                                              prefix.shape[2],
+                                                              video.shape[2],
+                                                              start,
+                                                              block,
+                                                              triton.next_power_of_2(dim),
+                                                              prefix.stride(1),
+                                                              prefix.stride(2),
+                                                              video.stride(1),
+                                                              video.stride(2),
+                                                              num_warps=4)
+    return payload, scale
+
+

fastvideo/models/dits/minimax_h3_fusions/sp_projection.py

@@ -0,0 +1,105 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Exact per-projection QKV packing for overlap with two-rank SP traffic."""
+
+import torch
+import triton
+import triton.language as tl
+
+from .sp_qkv_ordered import _normalize
+from .sp_qkv_pack import _pack_value_fp8
+
+
+@triton.jit
+def _pack_one(X, Weight, Cos, Sin, Payload, Scale, ROWS: tl.constexpr, HEADS: tl.constexpr, DIM: tl.constexpr,
+              ROTARY: tl.constexpr, WORLD: tl.constexpr, CHUNKS: tl.constexpr, EPS: tl.constexpr,
+              PROJECTION: tl.constexpr, BLOCK_ROWS: tl.constexpr, BLOCK_DIM: tl.constexpr):
+    row_head = (tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)).to(tl.int64)
+    row, head = row_head // HEADS, row_head % HEADS
+    col = tl.arange(0, BLOCK_DIM)
+    values = _normalize(X, Weight, Cos, Sin, row, head, HEADS, ROWS, DIM, ROTARY, EPS, BLOCK_DIM)
+    maximum = tl.max(tl.abs(values), 1)
+    scale = tl.maximum(maximum * (1.0 / 448.0), 1e-12)
+    quantized = (values / scale[:, None]).to(tl.float8e4nv)
+    local_heads: tl.constexpr = HEADS // WORLD
+    chunk_heads: tl.constexpr = local_heads // CHUNKS
+    chunk = head % local_heads // chunk_heads
+    destination = ((
+        (chunk * WORLD + head // local_heads) * 3 + PROJECTION) * ROWS + row) * chunk_heads + head % chunk_heads
+    tl.store(Payload + destination[:, None] * DIM + col[None, :], quantized,
+             (row[:, None] < ROWS) & (col[None, :] < DIM))
+    tl.store(Scale + destination, scale, row < ROWS)
+
+
+def allocate(rows, heads, world, chunks, device) -> tuple[torch.Tensor, torch.Tensor]:
+    payload = torch.empty((chunks, world, 3, rows, heads // (world * chunks), 128),
+                          device=device,
+                          dtype=torch.float8_e4m3fn)
+    return payload, torch.empty(payload.shape[:-1], device=device, dtype=torch.float32)
+
+
+def pack_one(x, weight, cosine, sine, payload, scale, epsilon, projection) -> None:
+    chunks, world, _, rows, local_heads, dim = payload.shape
+    heads = local_heads * world * chunks
+    _pack_one[(triton.cdiv(rows * heads, 8), )](
+        x,
+        weight,
+        cosine,
+        sine,
+        payload,
+        scale,
+        rows,
+        heads,
+        dim,
+        cosine.shape[1],
+        world,
+        chunks,
+        epsilon,
+        projection,
+        8,
+        128,
+        num_warps=8,
+    )
+
+
+def pack_value(x, payload, scale) -> None:
+    chunks, world, _, rows, local_heads, dim = payload.shape
+    heads = local_heads * world * chunks
+    _pack_value_fp8[(triton.cdiv(rows * heads, 16), )](
+        x,
+        payload,
+        scale,
+        rows,
+        heads,
+        dim,
+        world,
+        chunks,
+        16,
+        128,
+        num_warps=4,
+    )
+
+
+def project_and_exchange(attn, hidden, rotary, group, chunks, pre_quantized):
+    """Start Q traffic before K's GEMM, and K traffic before V's GEMM.
+
+    The caller checks that the three NVFP4 projections and gate can share
+    packed input. Normalization, RoPE and FP8 division preserve the ordinary
+    packer's arithmetic and reduction order. Nothing is reused across calls.
+    """
+    from fastvideo.distributed.projection_exchange import ProjectionExchange
+
+    linears = (attn.to_q, attn.to_k, attn.to_v)
+    pre = pre_quantized if pre_quantized is not None else linears[0].quant_method.quantize_input(hidden)
+    rows, heads = hidden.shape[1], attn.num_attention_heads
+    cosine, sine = (item.to(hidden.dtype) for item in rotary)
+    payload, scale = allocate(rows, heads, 2, chunks, hidden.device)
+    exchange = ProjectionExchange(payload, scale, group)
+    for projection, linear in enumerate(linears):
+        value = linear.quant_method.apply(linear, hidden, linear.bias, pre_quantized=pre)
+        if projection == 2:
+            pack_value(value, payload, scale)
+        else:
+            norm = attn.norm_q if projection == 0 else attn.norm_k
+            pack_one(value, norm.weight, cosine, sine, payload, scale, norm.eps, projection)
+        exchange.start(projection)
+    return exchange.packets(), pre

fastvideo/models/dits/minimax_h3_fusions/sp_qkv_ordered.py

@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Parallel SP QK rows with the original four-warp RMS addition order.
+
+The standalone head-128 normalization reduces 32 lanes in each warp first,
+then combines four warp sums. Explicit XOR gathers retain that addition tree
+when multiple heads share a CTA; a plain two-dimensional reduction changes it.
+"""
+import triton
+import triton.language as tl
+
+@triton.jit
+def _normalize(X, Weight, Cos, Sin, row, head, HEADS: tl.constexpr, ROWS: tl.constexpr, DIM: tl.constexpr,
+               ROTARY: tl.constexpr, EPS: tl.constexpr, BLOCK: tl.constexpr):
+    col = tl.arange(0, BLOCK)
+    offset = (row[:, None] * HEADS + head[:, None]) * DIM
+    mask = (row[:, None] < ROWS) & (col[None, :] < DIM)
+    value = tl.load(X + offset + col[None, :], mask, other=0.0).to(tl.float32)
+    squared = value * value
+    indices = tl.broadcast_to(col[None, :], squared.shape)
+    for shift in tl.static_range(5):
+        squared = squared + tl.gather(squared, indices ^ (16 >> shift), 1)
+    for shift in tl.static_range(2):
+        squared = squared + tl.gather(squared, indices ^ (64 >> shift), 1)
+    variance = tl.sum(tl.where(col[None, :] == 0, squared, 0.0), 1) / DIM
+    inv = tl.math.rsqrt(variance + EPS)
+    weight = tl.load(Weight + col, col < DIM, other=0.0).to(tl.float32)
+    normalized = value * inv[:, None] * weight[None, :]
+    partner = tl.where(col < ROTARY // 2, col + ROTARY // 2, col - ROTARY // 2)
+    rot_mask = (row[:, None] < ROWS) & (col[None, :] < ROTARY)
+    pv = tl.load(X + offset + partner[None, :], rot_mask, other=0.0).to(tl.float32)
+    pw = tl.load(Weight + partner, col < ROTARY, other=0.0).to(tl.float32)
+    pn = pv * inv[:, None] * pw[None, :]
+    rotated = tl.where(col[None, :] < ROTARY // 2, -pn, pn)
+    cosine = tl.load(Cos + row[:, None] * ROTARY + col[None, :], rot_mask, other=1.0).to(tl.float32)
+    sine = tl.load(Sin + row[:, None] * ROTARY + col[None, :], rot_mask, other=0.0).to(tl.float32)
+    return tl.where(col[None, :] < ROTARY, normalized * cosine + rotated * sine,
+                    normalized).to(tl.bfloat16).to(tl.float32)
+
+
+@triton.jit
+def _pack_qk(Q, K, QW, KW, Cos, Sin, Payload, Scale, ROWS: tl.constexpr, HEADS: tl.constexpr, DIM: tl.constexpr,
+             ROTARY: tl.constexpr, WORLD: tl.constexpr, CHUNKS: tl.constexpr, QEPS: tl.constexpr, KEPS: tl.constexpr,
+             BLOCK_ROWS: tl.constexpr, BLOCK_DIM: tl.constexpr):
+    row_head = (tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)).to(tl.int64)
+    row, head = row_head // HEADS, row_head % HEADS
+    col = tl.arange(0, BLOCK_DIM)
+    projection = tl.program_id(1)
+    if projection == 0:
+        values = _normalize(Q, QW, Cos, Sin, row, head, HEADS, ROWS, DIM, ROTARY, QEPS, BLOCK_DIM)
+    else:
+        values = _normalize(K, KW, Cos, Sin, row, head, HEADS, ROWS, DIM, ROTARY, KEPS, BLOCK_DIM)
+    maximum = tl.max(tl.abs(values), 1)
+    scale = tl.maximum(maximum * (1.0 / 448.0), 1e-12)
+    quantized = (values / scale[:, None]).to(tl.float8e4nv)
+    local_heads: tl.constexpr = HEADS // WORLD
+    chunk_heads: tl.constexpr = local_heads // CHUNKS
+    chunk = head % local_heads // chunk_heads
+    destination = ((
+        (chunk * WORLD + head // local_heads) * 3 + projection) * ROWS + row) * chunk_heads + head % chunk_heads
+    tl.store(Payload + destination[:, None] * DIM + col[None, :], quantized,
+             (row[:, None] < ROWS) & (col[None, :] < DIM))
+    tl.store(Scale + destination, scale, row < ROWS)
+

fastvideo/models/dits/minimax_h3_fusions/sp_qkv_pack.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fuse H3 local Q/K normalization and RoPE with FP8 Ulysses packing."""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _normalized_rotary(X, Weight, Cos, Sin, row, head, HEADS: tl.constexpr, DIM: tl.constexpr,
+                        ROTARY: tl.constexpr, EPS: tl.constexpr, BLOCK: tl.constexpr):
+    column = tl.arange(0, BLOCK)
+    offset = (row * HEADS + head) * DIM
+    value = tl.load(X + offset + column, column < DIM, other=0.0).to(tl.float32)
+    variance = tl.sum(value * value, axis=0) / DIM
+    inverse_rms = tl.math.rsqrt(variance + EPS)
+    weight = tl.load(Weight + column, column < DIM, other=0.0).to(tl.float32)
+    normalized = value * inverse_rms * weight
+    partner = tl.where(column < ROTARY // 2, column + ROTARY // 2, column - ROTARY // 2)
+    partner_value = tl.load(X + offset + partner, column < ROTARY, other=0.0).to(tl.float32)
+    partner_weight = tl.load(Weight + partner, column < ROTARY, other=0.0).to(tl.float32)
+    partner_normalized = partner_value * inverse_rms * partner_weight
+    rotated = tl.where(column < ROTARY // 2, -partner_normalized, partner_normalized)
+    cosine = tl.load(Cos + row * ROTARY + column, column < ROTARY, other=1.0).to(tl.float32)
+    sine = tl.load(Sin + row * ROTARY + column, column < ROTARY, other=0.0).to(tl.float32)
+    result = tl.where(column < ROTARY, normalized * cosine + rotated * sine, normalized)
+    # The former standalone norm/RoPE kernel materialized BF16 here.
+    return result.to(tl.bfloat16).to(tl.float32)
+
+
+@triton.jit
+def _pack_normalized_qk(Q, K, QWeight, KWeight, Cos, Sin, Payload, Scale,
+                         ROWS: tl.constexpr, HEADS: tl.constexpr, DIM: tl.constexpr,
+                         ROTARY: tl.constexpr, WORLD: tl.constexpr, CHUNKS: tl.constexpr,
+                         QEPS: tl.constexpr, KEPS: tl.constexpr, BLOCK: tl.constexpr, ROWS_PER_PROGRAM: tl.constexpr):
+    local_heads: tl.constexpr = HEADS // WORLD
+    chunk_heads: tl.constexpr = local_heads // CHUNKS
+    # Preserve each head's four-warp reduction while amortizing CTA scheduling.
+    for index in tl.static_range(ROWS_PER_PROGRAM):
+        row_head = tl.program_id(0).to(tl.int64) * ROWS_PER_PROGRAM + index
+        if ROWS * HEADS % ROWS_PER_PROGRAM == 0 or row_head < ROWS * HEADS:
+            row, head = row_head // HEADS, row_head % HEADS
+            projection = tl.program_id(1)
+            column = tl.arange(0, BLOCK)
+            if projection == 0:
+                values = _normalized_rotary(Q, QWeight, Cos, Sin, row, head, HEADS, DIM, ROTARY, QEPS, BLOCK)
+            else:
+                values = _normalized_rotary(K, KWeight, Cos, Sin, row, head, HEADS, DIM, ROTARY, KEPS, BLOCK)
+            maximum = tl.max(tl.where(column < DIM, tl.abs(values), 0.0), axis=0)
+            scale = tl.maximum(maximum * (1.0 / 448.0), 1e-12)
+            quantized = (values / scale).to(tl.float8e4nv)
+            chunk = (head % local_heads) // chunk_heads
+            destination = (((chunk * WORLD + head // local_heads) * 3 + projection) * ROWS + row) * chunk_heads + head % chunk_heads
+            tl.store(Payload + destination * DIM + column, quantized, column < DIM)
+            tl.store(Scale + destination, scale)
+
+
+
+@triton.jit
+def _pack_value_fp8(V, Payload, Scale, ROWS: tl.constexpr, HEADS: tl.constexpr, DIM: tl.constexpr,
+                    WORLD: tl.constexpr, CHUNKS: tl.constexpr, BLOCK_HEADS: tl.constexpr, BLOCK_DIM: tl.constexpr):
+    row_head = (tl.program_id(0) * BLOCK_HEADS + tl.arange(0, BLOCK_HEADS)).to(tl.int64)
+    column = tl.arange(0, BLOCK_DIM)
+    valid = (row_head[:, None] < ROWS * HEADS) & (column[None, :] < DIM)
+    values = tl.load(V + row_head[:, None] * DIM + column[None, :], valid, other=0.0).to(tl.float32)
+    scale = tl.maximum(tl.max(tl.abs(values), axis=1) * (1.0 / 448.0), 1e-12)
+    quantized = (values / scale[:, None]).to(tl.float8e4nv)
+    row, head = row_head // HEADS, row_head % HEADS
+    local_heads: tl.constexpr = HEADS // WORLD
+    chunk_heads: tl.constexpr = local_heads // CHUNKS
+    chunk = (head % local_heads) // chunk_heads
+    destination = (((chunk * WORLD + head // local_heads) * 3 + 2) * ROWS + row) * chunk_heads + head % chunk_heads
+    tl.store(Payload + destination[:, None] * DIM + column[None, :], quantized, valid)
+    tl.store(Scale + destination, scale, row_head < ROWS * HEADS)
+
+
+def pack_normalized_qkv_fp8(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,
+                            q_weight: torch.Tensor, k_weight: torch.Tensor,
+                            cos: torch.Tensor, sin: torch.Tensor, q_eps: float, k_eps: float,
+                            world: int, head_chunks: int = 1) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pack local `[rows,H,D]` after BF16 norm/RoPE rounding.
+
+    With multiple chunks, write `[chunks,W,3,rows,H/W/chunks,D]` directly so
+    each chunk's all-to-all buffer is contiguous without another memory pass.
+    """
+    rows, heads, dim = query.shape
+    assert query.dtype == key.dtype == value.dtype == torch.bfloat16
+    assert q_weight.dtype == k_weight.dtype == cos.dtype == sin.dtype == query.dtype
+    assert query.shape == key.shape == value.shape and head_chunks > 0 and heads % (world * head_chunks) == 0
+    assert all(t.is_contiguous() for t in (query, key, value, q_weight, k_weight, cos, sin))
+    assert q_weight.shape == k_weight.shape == (dim,)
+    assert cos.shape == sin.shape and cos.shape[0] == rows and 0 < cos.shape[1] <= dim
+    assert cos.shape[1] % 2 == 0
+    payload = torch.empty((head_chunks, world, 3, rows, heads // (world * head_chunks), dim), device=query.device, dtype=torch.float8_e4m3fn)
+    scales = torch.empty((head_chunks, world, 3, rows, heads // (world * head_chunks)), device=query.device, dtype=torch.float32)
+    rows_per_program = 8 if dim == 128 and torch.cuda.get_device_capability(query.device)[0] >= 10 else 1
+    if dim == 128 and torch.cuda.get_device_capability(query.device)[0] >= 10:
+        from .sp_qkv_ordered import _pack_qk
+        _pack_qk[(triton.cdiv(rows * heads, 8), 2)](
+            query, key, q_weight, k_weight, cos, sin, payload, scales,
+            rows, heads, dim, cos.shape[1], world, head_chunks, q_eps, k_eps, 8, 128, num_warps=8,
+        )
+    else:
+        _pack_normalized_qk[(triton.cdiv(rows * heads, rows_per_program), 2)](
+            query, key, q_weight, k_weight, cos, sin, payload, scales,
+            rows, heads, dim, cos.shape[1], world, head_chunks, q_eps, k_eps,
+            triton.next_power_of_2(dim), rows_per_program, num_warps=4,
+        )
+    _pack_value_fp8[(triton.cdiv(rows * heads, 16),)](
+        value, payload, scales, rows, heads, dim, world, head_chunks, 16, triton.next_power_of_2(dim), num_warps=4,
+    )
+    return (payload[0], scales[0]) if head_chunks == 1 else (payload, scales)

fastvideo/models/dits/minimax_h3_fusions/sp_receive.py

@@ -0,0 +1,169 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Share SP FP8 receive reads across BF16 expansion, tile pooling and Q/V FP4 packing."""
+import torch
+import triton
+import triton.language as tl
+
+from fastvideo.distributed.async_collectives import PeerSlices
+
+from .value_quant import _left_four
+from .swiglu_quant import _rcp
+
+@triton.jit
+def _receive_all(Payload, Scale, Inverse, Valid, Out, Pool, Codes, Scales, QCodes, QScales, ROWS: tl.constexpr,
+                 S: tl.constexpr, H: tl.constexpr, D: tl.constexpr, BLOCK: tl.constexpr, START: tl.constexpr,
+                 Remote=None, PAIR: tl.constexpr = False, RANK: tl.constexpr = 0,
+                 QUERY_ROWS: tl.constexpr = -1, PrefixOut=None, PrefixInverse=None,
+                 PREFIX_VALUES: tl.constexpr = False, STORE_KEYS: tl.constexpr = True):
+    tile = tl.program_id(1)
+    projection = tl.program_id(2)
+    col = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    row = tl.arange(0, 16)
+    head = col // D
+    acc = tl.full((16, BLOCK), 0, tl.float32)
+    for i in tl.static_range(4):
+        token = tile * 64 + i * 16 + row
+        source = tl.load(Inverse + token)
+        packed = (source // ROWS * 3 + projection) * ROWS + source % ROWS
+        valid = (col[None, :] < H * D) & (source[:, None] >= 0)
+        if PAIR:
+            local_offset = (projection * ROWS + source % ROWS)[:, None] * H * D + col[None, :]
+            pointer = tl.where((source // ROWS)[:, None] == RANK, Payload + local_offset, Remote + local_offset)
+        else:
+            pointer = Payload + packed[:, None] * H * D + col[None, :]
+        value = tl.load(pointer, valid, other=0.0).to(tl.float32)
+        scale = tl.load(Scale + packed[:, None] * H + head[None, :], valid, other=0)
+        x = (value * scale).to(tl.bfloat16).to(tl.float32)
+        store_mask = (col[None, :] < H * D) & (STORE_KEYS | (projection != 1))
+        if QUERY_ROWS >= 0:
+            store_mask = store_mask & ((projection != 0) | (token[:, None] < QUERY_ROWS))
+        if PREFIX_VALUES and projection == 2:
+            destination = tl.load(PrefixInverse + token)
+            tl.store(PrefixOut + destination[:, None] * H * D + col[None, :], x,
+                     (col[None, :] < H * D) & (destination[:, None] >= 0))
+        else:
+            tl.store(Out + (projection * S + token[:, None]) * H * D + col[None, :], x, store_mask)
+        acc = acc + x
+        if projection == 2:
+            maximum = tl.max(tl.abs(x), 0)
+            sf = (maximum * (1.0 / 6.0)).to(tl.float8e4nv)
+            inverse = tl.where(sf.to(tl.float32) == 0, 0, _rcp(sf.to(tl.float32)))
+            scaled = x * inverse[None, :]
+            pair = tl.trans(tl.reshape(scaled, (8, 2, BLOCK)), (0, 2, 1))
+            lo, hi = tl.split(pair)
+            codes = tl.inline_asm_elementwise('''{
+                .reg .b8 q;
+                cvt.rn.satfinite.e2m1x2.f32 q, $1, $2;
+                mov.b32 $0, {q,q,q,q};
+            }''',
+                                              constraints='=r,f,f',
+                                              args=[hi, lo],
+                                              dtype=tl.uint8,
+                                              is_pure=True,
+                                              pack=1)
+            qrow = tl.arange(0, 8)
+            offset = (col[None, :] * (S // 2) + tile * 32 + i * 8 + qrow[:, None])
+            tl.store(Codes + offset, codes, col[None, :] < H * D)
+            group = tile * 4 + i
+            dim = col % D
+            offset = (head * D * (S // 16) + (dim // 64) * 64 * (S // 16) + (group // 4) * 256 + group % 4 +
+                      (dim % 64 // 16) * 4 + (dim % 16) * 16)
+            tl.store(Scales + offset, sf, col < H * D)
+        elif projection == 0:
+            groups = tl.reshape(x, (16, BLOCK // 16, 16))
+            maximum = tl.max(tl.abs(groups), 2)
+            sf = (maximum * (1.0 / 6.0)).to(tl.float8e4nv)
+            inverse = tl.where(sf.to(tl.float32) == 0, 0, _rcp(sf.to(tl.float32)))
+            quant = tl.reshape(groups * inverse[:, :, None], (16, BLOCK // 2, 2))
+            lo, hi = tl.split(quant)
+            packed_q = tl.inline_asm_elementwise('''{
+                .reg .b8 q;
+                cvt.rn.satfinite.e2m1x2.f32 q, $1, $2;
+                mov.b32 $0, {q,q,q,q};
+            }''',
+                                                 constraints='=r,f,f',
+                                                 args=[hi, lo],
+                                                 dtype=tl.uint8,
+                                                 is_pure=True,
+                                                 pack=1)
+            qrow = token - START
+            qseq: tl.constexpr = S - START
+            qcol = tl.program_id(0) * (BLOCK // 2) + tl.arange(0, BLOCK // 2)
+            qhead, qdim = qcol // (D // 2), qcol % (D // 2)
+            tl.store(QCodes + (qhead[None, :] * qseq + qrow[:, None]) * (D // 2) + qdim[None, :], packed_q,
+                     (qrow[:, None] >= 0) & (qcol[None, :] < H * D // 2))
+            group = tl.program_id(0) * (BLOCK // 16) + tl.arange(0, BLOCK // 16)
+            qhead, qdim = group // (D // 16), group % (D // 16)
+            offset = ((qrow[:, None] // 64) * 512 + (qdim[None, :] // 4) * 256 + qdim[None, :] % 4 +
+                      (qrow[:, None] % 64 // 16) * 4 + (qrow[:, None] % 16) * 16)
+            tl.store(QScales + qhead[None, :] * qseq * (D // 16) + offset, sf,
+                     (qrow[:, None] >= 0) & (group[None, :] < H * D // 16))
+    even, odd = tl.split(tl.reshape(tl.trans(acc), (BLOCK, 4, 2, 2)))
+    y0, y2 = tl.split(even)
+    y1, y3 = tl.split(odd)
+    y0, y1 = _left_four(y0, BLOCK), _left_four(y1, BLOCK)
+    y2, y3 = _left_four(y2, BLOCK), _left_four(y3, BLOCK)
+    total = (y0 + y2) + (y1 + y3)
+    count = tl.load(Valid + tile).to(tl.float32)
+    mean = tl.div_rn(total, count)
+    tl.store(Pool + (projection * (S // 64) + tile) * H * D + col, mean, col < H * D)
+
+
+def unpack_pool_quantize_fp8(
+    payload: torch.Tensor | PeerSlices, scale: torch.Tensor, inverse: torch.Tensor,
+    valid: torch.Tensor, output: torch.Tensor, query_start: int, *, query_rows: int | None = None,
+    prefix_values: tuple[torch.Tensor, torch.Tensor] | None = None, store_keys: bool = True,
+) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]:
+    """Expand FP8 QKV and prepare exact tile means plus native Q/V FP4 inputs.
+
+    By default every tile row is written, including zero padding. When the
+    caller consumes native FP4 queries, ``query_rows`` may limit BF16 query
+    stores to the dense prefix. The omitted suffix is untouched and must not
+    be read; pooling and native query codes still cover every row. K/V stores
+    remain complete by default. ``prefix_values=(inverse, output)`` instead
+    writes BF16 values directly into dense-prefix order; the entire raw V
+    tensor is then untouched and must not be read. Native V codes and pooling
+    still cover every row. Pooling keeps ATen's
+    tile-64 accumulation order and reads the rounded BF16 expansion. Query
+    scales use Sage's native suffix layout; value scales use its transposed
+    layout. With ``store_keys=False``, raw K is untouched; the caller must
+    prepare keys directly from these packets. Prefix values use the
+    caller's inverse map, whose entries are -1 or in-bounds output row indices.
+    """
+    pair = payload if isinstance(payload, PeerSlices) else None
+    payload = pair.template if pair is not None else payload
+    start, block = query_start, 256
+    _, _, rows, heads, dim = payload.shape
+    assert payload.dtype == torch.float8_e4m3fn and scale.dtype == torch.float32
+    assert payload.shape[1] == 3 and dim == 128 and output.dtype == torch.bfloat16
+    assert inverse.numel() % 128 == 0 and 0 <= start < inverse.numel() and start % 128 == 0
+    assert output.shape == (3, inverse.numel(), heads, dim) and scale.shape == payload.shape[:-1]
+    assert valid.shape == (inverse.numel() // 64,)
+    assert all(t.is_contiguous() for t in (payload, scale, inverse, valid, output))
+    assert inverse.dtype in (torch.int32, torch.int64) and valid.dtype in (torch.int32, torch.int64)
+    assert all(t.device == payload.device for t in (scale, inverse, valid, output))
+    assert payload.is_cuda and not torch.is_grad_enabled()
+    prefix_inverse, prefix_output = (None, None) if prefix_values is None else prefix_values
+    if prefix_values is not None:
+        if (prefix_inverse.shape != inverse.shape or prefix_inverse.dtype not in (torch.int32, torch.int64)
+                or prefix_inverse.device != payload.device or not prefix_inverse.is_contiguous()
+                or prefix_output.ndim != 4 or prefix_output.shape[0] != 1
+                or prefix_output.shape[2:] != (heads, dim) or prefix_output.dtype != torch.bfloat16
+                or prefix_output.device != payload.device or not prefix_output.is_contiguous()):
+            raise ValueError("Prefix values require a tile inverse and contiguous BF16 [1, valid, heads, 128] output")
+    seq = inverse.numel()
+    if query_rows is not None and not 0 <= query_rows <= seq:
+        raise ValueError("BF16 query rows must lie within the tile sequence")
+    pool = torch.empty((3, seq // 64, heads, dim), device=payload.device, dtype=torch.float32)
+    codes = torch.empty((1, heads, dim, seq // 2), device=payload.device, dtype=torch.uint8)
+    scales = torch.empty((1, heads, dim, seq // 16), device=payload.device, dtype=torch.float8_e4m3fn)
+    qcodes = torch.empty((1, heads, seq - start, dim // 2), device=payload.device, dtype=torch.uint8)
+    qscales = torch.empty((1, heads, seq - start, dim // 16), device=payload.device, dtype=torch.float8_e4m3fn)
+    _receive_all[(triton.cdiv(heads * dim, block), seq // 64, 3)](
+        pair.local if pair is not None else payload, scale, inverse, valid, output, pool, codes, scales, qcodes, qscales,
+        rows, seq, heads, dim, block, start, Remote=pair.remote if pair is not None else payload,
+        PAIR=pair is not None, RANK=pair.rank if pair is not None else 0,
+        QUERY_ROWS=-1 if query_rows is None else query_rows, PrefixOut=prefix_output, PrefixInverse=prefix_inverse,
+        PREFIX_VALUES=prefix_values is not None, STORE_KEYS=store_keys, num_warps=4,
+    )
+    return pool.permute(0, 2, 1, 3), (codes, scales), (qcodes, qscales)

fastvideo/models/dits/minimax_h3_fusions/swiglu_gemm.py

@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Exact eager GB10 FFN-up GEMM, BF16 SwiGLU, and NVFP4 output preparation.
+
+Only immutable weights are cached. Their value/gate columns are paired in
+memory so each GEMM tile can produce complete SwiGLU quantization groups.
+The checkpoint layout and files are unchanged. Both original BF16 rounding
+boundaries precede the original block-16 scale/code conversion.
+"""
+from functools import lru_cache
+
+import torch
+
+
+@lru_cache(maxsize=1)
+def _backend():
+    try:
+        import _fastvideo_h3_swiglu_gemm
+    except ImportError:
+        return None
+    return _fastvideo_h3_swiglu_gemm
+
+
+def pair_weight_columns(weight: torch.Tensor, scales: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pair value/gate rows of packed weights and their 128x4-swizzled scales."""
+    n, k = weight.shape[0], weight.shape[1] * 2
+    values = torch.arange(n // 2, device=weight.device)
+    permutation = torch.stack((values, values + n // 2), dim=1).reshape(-1)
+    paired_weight = weight[permutation].contiguous().T
+    plain_scales = scales.reshape(n // 128, k // 64, 32, 4, 4).permute(0, 3, 2, 1, 4).reshape(n, k // 16)
+    paired_scales = (plain_scales[permutation].reshape(n // 128, 4, 32, k // 64, 4)
+                     .permute(0, 3, 2, 1, 4).contiguous().reshape_as(scales).T)
+    return paired_weight, paired_scales
+
+
+def _weight_key(weight: torch.Tensor, scales: torch.Tensor) -> tuple:
+    # Packed inference buffers may have no version counter. Their storage is
+    # immutable; replacement or migration changes the identity/storage key.
+    return tuple((id(t), t.data_ptr(), t.device, None if t.is_inference() else t._version)
+                 for t in (weight, scales))
+
+
+def try_swiglu_gemm(ff, hidden_states: torch.Tensor, pre_quantized) -> torch.Tensor | None:
+    """Return None outside the measured T2AV/Omni single/pair-Spark NVFP4 shapes."""
+    from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod, _mm_fp4_row_chunks
+
+    if (pre_quantized is None or torch.is_grad_enabled() or torch.compiler.is_compiling()
+            or not hidden_states.is_cuda or hidden_states.dtype != torch.bfloat16
+            or torch.cuda.get_device_capability(hidden_states.device) != (12, 1)
+            or type(getattr(ff.fc_in, 'quant_method', None)) is not NVFP4QuantizeMethod
+            or type(getattr(ff.fc_out, 'quant_method', None)) is not NVFP4QuantizeMethod
+            or ff.fc_in.bias is not None):
+        return None
+    a, a_scale, input_scale = pre_quantized
+    a = a.reshape(-1, a.shape[-1])
+    a_scale = a_scale.reshape(-1, a_scale.shape[-1])
+    weight = ff.fc_in._nvfp4_weight
+    weight_scales = ff.fc_in._nvfp4_weight_scale
+    if (a.shape not in ((14859, 2688), (7430, 2688), (29678, 2688)) or weight.shape != (28672, 2688)
+            or not ff.fc_out.quant_method.wants_prequantized_input()):
+        return None
+    backend = _backend()
+    if backend is None:
+        return None
+    method = ff.fc_out.quant_method
+    output_scale = method._static_activation_global_sf()
+    if output_scale is None:
+        output_scale = getattr(ff.fc_out, '_nvfp4_input_global_sf', None)
+    if output_scale is None:
+        if method._dynamic_activation_scale():
+            return None
+        output_scale = method.x_global_sf
+    key = _weight_key(weight, weight_scales)
+    if getattr(ff.fc_in, '_h3_swiglu_weight_key', None) != key:
+        paired_weight, paired_scales = pair_weight_columns(weight, weight_scales)
+        ff.fc_in.register_buffer('_h3_swiglu_paired_weight', paired_weight, persistent=False)
+        ff.fc_in.register_buffer('_h3_swiglu_paired_scales', paired_scales, persistent=False)
+        ff.fc_in._h3_swiglu_weight_key = key
+    paired_weight = ff.fc_in._h3_swiglu_paired_weight
+    paired_scales = ff.fc_in._h3_swiglu_paired_scales
+    if hasattr(ff.fc_in, '_nvfp4_alpha'):
+        alpha = ff.fc_in._nvfp4_alpha / input_scale
+    else:
+        alpha = 1.0 / (input_scale * ff.fc_in._weight_global_sf)
+    rows, width = a.shape[0], weight.shape[0] // 2
+    q = torch.empty((rows, width // 2), device=a.device, dtype=torch.uint8)
+    scales = torch.empty(((rows + 127) // 128 * 128, width // 16), device=a.device, dtype=torch.uint8)
+    start = 0
+    parts = 8 if rows == 29678 else (2 if rows == 7430 else 4)
+    for chunk, chunk_scale in _mm_fp4_row_chunks(a, a_scale, parts):
+        end = start + chunk.shape[0]
+        backend.run(chunk, paired_weight, chunk_scale, paired_scales, alpha, output_scale,
+                    q[start:end], scales[start:(end + 127) // 128 * 128])
+        start = end
+    # apply(pre_quantized=...) reads only hidden_states' leading shape.
+    return method.apply(ff.fc_out, hidden_states, ff.fc_out.bias, pre_quantized=(q, scales, output_scale))

fastvideo/models/dits/minimax_h3_fusions/swiglu_quant.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+"""H3 SwiGLU and exact FlashInfer-compatible NVFP4 activation quantization.
+
+The SwiGLU result is rounded to BF16 before quantization, exactly as in the
+separate kernels. Scales use FlashInfer's default fast-math recipe and 128x4
+swizzle. Fusing the two passes avoids a full-width BF16 intermediate.
+"""
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _rcp(x):
+    return tl.inline_asm_elementwise(
+        "rcp.approx.ftz.f32 $0, $1;", constraints="=f,f", args=[x],
+        dtype=tl.float32, is_pure=True, pack=1,
+    )
+
+
+@triton.jit
+def _swiglu_quant(X, SF, Q, S, ROWS: tl.constexpr, WIDTH: tl.constexpr, BLOCK: tl.constexpr):
+    row = tl.program_id(0).to(tl.int64)
+    col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
+    valid = (row < ROWS) & (col < WIDTH)
+    value = tl.load(X + row * WIDTH * 2 + col, valid, other=0).to(tl.float32)
+    gate = tl.load(X + row * WIDTH * 2 + WIDTH + col, valid, other=0).to(tl.float32)
+    y = (value * (gate * tl.sigmoid(gate))).to(tl.bfloat16).to(tl.float32)
+    grouped = tl.reshape(y, (BLOCK // 16, 16))
+    amax = tl.max(tl.abs(grouped), 1)
+    global_sf = tl.load(SF)
+    scale = (global_sf * (amax * _rcp(6.0))).to(tl.float8e4nv)
+    inv = tl.where(amax != 0, _rcp(scale.to(tl.float32) * _rcp(global_sf)), 0.0)
+    quantized = tl.reshape(grouped * inv[:, None], (BLOCK // 2, 2))
+    lo, hi = tl.split(quantized)
+    packed = tl.inline_asm_elementwise('''{
+        .reg .b8 q;
+        cvt.rn.satfinite.e2m1x2.f32 q, $1, $2;
+        mov.b32 $0, {q,q,q,q};
+    }''', constraints='=r,f,f', args=[hi,lo], dtype=tl.uint8, is_pure=True, pack=1)
+    qcol = tl.program_id(1) * (BLOCK // 2) + tl.arange(0, BLOCK // 2)
+    tl.store(Q + row * (WIDTH // 2) + qcol, packed, (row < ROWS) & (qcol < WIDTH // 2))
+    group = tl.program_id(1) * (BLOCK // 16) + tl.arange(0, BLOCK // 16)
+    sidx = (((((row // 128) * (WIDTH // 64) + group // 4) * 32 + row % 32) * 4
+             + (row // 32) % 4) * 4 + group % 4)
+    tl.store(S + sidx, scale.to(tl.uint8, bitcast=True), group < WIDTH // 16)
+
+
+
+def swiglu_quantize_nvfp4(
+    x: torch.Tensor, global_scale: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Quantize value-first SwiGLU to packed FP4 codes and padded 128x4 scales.
+
+    Requires contiguous CUDA BF16 inputs with a hidden width divisible by 64,
+    a scalar FP32 global scale, and native Blackwell FP4 conversion support.
+    The caller selects this only for eager inference with a compatible NVFP4
+    output projection; other paths retain the ordinary SwiGLU implementation.
+    """
+    if not x.is_cuda or x.dtype != torch.bfloat16 or not x.is_contiguous():
+        raise ValueError("NVFP4 SwiGLU requires contiguous CUDA BF16 input")
+    if x.ndim < 2 or x.shape[-1] == 0 or x.shape[-1] % 128:
+        raise ValueError("NVFP4 SwiGLU requires a packed width divisible by 128")
+    if global_scale.numel() != 1 or global_scale.dtype != torch.float32 or global_scale.device != x.device:
+        raise ValueError("NVFP4 SwiGLU requires a scalar FP32 scale on the input device")
+    if torch.is_grad_enabled():
+        raise RuntimeError("NVFP4 SwiGLU is inference-only")
+    rows, width = x.numel() // x.shape[-1], x.shape[-1] // 2
+    quantized = torch.empty((rows, width // 2), device=x.device, dtype=torch.uint8)
+    scales = torch.empty((triton.cdiv(rows, 128) * 128, width // 16), device=x.device, dtype=torch.uint8)
+    if rows:
+        block = 8192
+        _swiglu_quant[(scales.shape[0], triton.cdiv(width, block))](
+            x, global_scale, quantized, scales, rows, width, block, num_warps=8,
+        )
+    return quantized, scales, global_scale

fastvideo/models/dits/minimax_h3_fusions/tile_pool.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+"""ATen-compatible tile sums for the fused attention producers."""
+import triton
+import triton.language as tl
+
+from .value_quant import _left_four
+
+
+@triton.jit
+def _tile_sum(values, R: tl.constexpr):
+    """Sum 64 or 128 rows in the same order as the eager FP32 tile reduction."""
+    raw = tl.trans(tl.reshape(values, (R // 16, 16, 128)), (1, 2, 0))
+    if R == 64:
+        even, odd = tl.split(tl.reshape(raw, (16, 128, 2, 2)))
+        x0, x2 = tl.split(even)
+        x1, x3 = tl.split(odd)
+        acc = ((x0 + x1) + x2) + x3
+    else:
+        even, odd = tl.split(tl.reshape(raw, (16, 128, 4, 2)))
+        x04, x26 = tl.split(tl.reshape(even, (16, 128, 2, 2)))
+        x15, x37 = tl.split(tl.reshape(odd, (16, 128, 2, 2)))
+        x0, x4 = tl.split(x04)
+        x1, x5 = tl.split(x15)
+        x2, x6 = tl.split(x26)
+        x3, x7 = tl.split(x37)
+        acc = ((((((x0 + x1) + x2) + x3) + x4) + x5) + x6) + x7
+    even, odd = tl.split(tl.reshape(tl.trans(acc), (128, 4, 2, 2)))
+    y0, y2 = tl.split(even)
+    y1, y3 = tl.split(odd)
+    y0, y1 = _left_four(y0, 128), _left_four(y1, 128)
+    y2, y3 = _left_four(y2, 128), _left_four(y3, 128)
+    return (y0 + y2) + (y1 + y3)

fastvideo/models/dits/minimax_h3_fusions/untile_gate.py

@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fuse VSA's packed-order gather and compression-gate addcmul."""
+
+import torch
+import triton
+import triton.language as tl
+
+from .swiglu_quant import _rcp as rcp
+
+
+@triton.jit
+def _untile_gate_kernel(O, P, C, G, Untile, RowTile, Y, Q, S, SF,
+                        OB: tl.constexpr, OH: tl.constexpr, OR: tl.constexpr,
+                        PB: tl.constexpr, PH: tl.constexpr, PR: tl.constexpr,
+                        CB: tl.constexpr, CT: tl.constexpr, CH: tl.constexpr,
+                        PREFIX: tl.constexpr, START: tl.constexpr,
+                        ROWS: tl.constexpr, HD: tl.constexpr, D: tl.constexpr,
+                        BLOCK: tl.constexpr, QUANTIZED: tl.constexpr):
+    row = tl.program_id(0)
+    batch = tl.program_id(2)
+    col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
+    head, dim = col // D, col % D
+    src = tl.load(Untile + row).to(tl.int64)
+    tile = tl.load(RowTile + row).to(tl.int64)
+    if PREFIX > 0:
+        if src < PREFIX:
+            a = tl.load(P + batch * PB + head * PH + src * PR + dim, col < HD, other=0).to(tl.float32)
+        else:
+            a = tl.load(O + batch * OB + head * OH + (src - START) * OR + dim, col < HD, other=0).to(tl.float32)
+    else:
+        a = tl.load(O + batch * OB + head * OH + src * OR + dim, col < HD, other=0).to(tl.float32)
+    b = tl.load(C + batch * CB + tile * CT + head * CH + dim, col < HD, other=0).to(tl.float32)
+    offset = (batch * ROWS + row).to(tl.int64) * HD + col
+    gate = tl.load(G + offset, col < HD, other=0).to(tl.float32)
+    # ATen's BF16 addcmul uses float32 fused multiply-add and one final cast.
+    if QUANTIZED:
+        y = tl.where(col < HD, tl.fma(b, gate, a).to(tl.bfloat16).to(tl.float32), 0.0)
+        grouped = tl.reshape(y, (BLOCK // 16, 16))
+        amax = tl.max(tl.abs(grouped), 1)
+        global_sf = tl.load(SF)
+        scale = (global_sf * (amax * rcp(6.0))).to(tl.float8e4nv)
+        inv = tl.where(amax != 0, rcp(scale.to(tl.float32) * rcp(global_sf)), 0.0)
+        quant = tl.reshape(grouped * inv[:, None], (BLOCK // 2, 2))
+        lo, hi = tl.split(quant)
+        packed = tl.inline_asm_elementwise('''{
+            .reg .b8 q;
+            cvt.rn.satfinite.e2m1x2.f32 q, $1, $2;
+            mov.b32 $0, {q,q,q,q};
+        }''', constraints="=r,f,f", args=[hi, lo], dtype=tl.uint8, is_pure=True, pack=1)
+        packed_row = (batch * ROWS + row).to(tl.int64)
+        qc = tl.program_id(1) * (BLOCK // 2) + tl.arange(0, BLOCK // 2)
+        tl.store(Q + packed_row * (HD // 2) + qc, packed, qc < HD // 2)
+        group = tl.program_id(1) * (BLOCK // 16) + tl.arange(0, BLOCK // 16)
+        scale_tile = (packed_row // 128) * (HD // 64) + group // 4
+        si = ((scale_tile * 32 + packed_row % 32) * 4 + (packed_row // 32) % 4) * 4 + group % 4
+        tl.store(S + si, scale.to(tl.uint8, bitcast=True), group < HD // 16)
+    else:
+        tl.store(Y + offset, tl.fma(b, gate, a), col < HD)
+
+
+def untile_gate(out: torch.Tensor | tuple[torch.Tensor, torch.Tensor, int], compressed: torch.Tensor,
+                gate: torch.Tensor, untile: torch.Tensor, row_tile: torch.Tensor,
+                global_scale: torch.Tensor | None = None,
+                ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Combine attention and compression in packed order, optionally preparing NVFP4 directly."""
+    batch, rows, hd = gate.shape
+    quantized = global_scale is not None
+    if quantized:
+        if gate.dtype != torch.bfloat16 or hd % 64:
+            raise ValueError("Fused gate quantization requires BF16 and a width divisible by 64")
+        if (global_scale.numel() != 1 or global_scale.dtype != torch.float32
+                or global_scale.device != gate.device):
+            raise ValueError("Fused gate quantization requires a scalar FP32 scale on the input device")
+        codes = torch.empty((batch * rows, hd // 2), device=gate.device, dtype=torch.uint8)
+        scales = torch.zeros((triton.cdiv(batch * rows, 128) * 128, hd // 16),
+                             device=gate.device, dtype=torch.uint8)
+        result = gate  # Unused output pointer in the quantized specialization.
+    else:
+        result = torch.empty_like(gate)
+        codes, scales, global_scale = result, result, result
+    if isinstance(out, tuple):
+        prefix, out, start = out
+        prefix_rows = prefix.shape[2]
+    else:
+        prefix, prefix_rows, start = out, 0, 0
+    _untile_gate_kernel[(rows, triton.cdiv(hd, 512), batch)](
+        out, prefix, compressed, gate, untile, row_tile, result, codes, scales, global_scale,
+        *out.stride()[:3], *prefix.stride()[:3], *compressed.stride()[:3],
+        prefix_rows, start, rows, hd, out.shape[-1], 512, quantized,
+    )
+    return (codes, scales, global_scale) if quantized else result

fastvideo/models/dits/minimax_h3_fusions/value_quant.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Share the value read between exact tile pooling and transposed Sage FP4 packing.
+
+For BF16 tile-64/128 pooling, ATen uses four accumulators on each of four warp
+rows. Keep that addition order, including the final (0+2)+(1+3) reduction,
+instead of a tree sum. Quantization independently reduces the maximum of
+16 consecutive tokens and uses Sage's transposed 64-row scale layout.
+"""
+import torch
+import triton
+import triton.language as tl
+
+from .swiglu_quant import _rcp
+
+
+@triton.jit
+def _left_four(x, BLOCK: tl.constexpr):
+    even, odd = tl.split(tl.reshape(x, (BLOCK, 2, 2)))
+    a0, a2 = tl.split(even)
+    a1, a3 = tl.split(odd)
+    return ((a0 + a1) + a2) + a3
+
+
+@triton.jit
+def _pool_and_quantize(X, VALID, POOL, CODES, SCALES,
+                       SB: tl.constexpr, SS: tl.constexpr, SH: tl.constexpr,
+                       S: tl.constexpr, H: tl.constexpr, D: tl.constexpr, BLOCK: tl.constexpr, WIDE: tl.constexpr,
+                       INPUT_INDEX=None, GATHER_INPUT: tl.constexpr = False, TILE: tl.constexpr = 64):
+    tile, batch = tl.program_id(1), tl.program_id(2)
+    col = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    if WIDE:
+        tile, batch, col = tile.to(tl.int64), batch.to(tl.int64), col.to(tl.int64)
+    row = tl.arange(0, 16)
+    head, dim = col // D, col % D
+    acc = tl.full((16, BLOCK), 0, tl.float32)
+    for i in tl.static_range(TILE // 16):
+        token = tile * TILE + i * 16 + row
+        source = tl.load(INPUT_INDEX + token) if GATHER_INPUT else token
+        x = tl.load(X + batch * SB + source[:, None] * SS + head[None, :] * SH + dim[None, :],
+                    (col[None, :] < H * D) & (source[:, None] >= 0), other=0).to(tl.float32)
+        acc = acc + x
+        maximum = tl.max(tl.abs(x), 0)
+        sf = (maximum * (1.0 / 6.0)).to(tl.float8e4nv)
+        inverse = tl.where(sf.to(tl.float32) == 0, 0, _rcp(sf.to(tl.float32)))
+        scaled = x * inverse[None, :]
+        pair = tl.trans(tl.reshape(scaled, (8, 2, BLOCK)), (0, 2, 1))
+        lo, hi = tl.split(pair)
+        codes = tl.inline_asm_elementwise('''{
+            .reg .b8 q;
+            cvt.rn.satfinite.e2m1x2.f32 q, $1, $2;
+            mov.b32 $0, {q,q,q,q};
+        }''', constraints='=r,f,f', args=[hi, lo], dtype=tl.uint8, is_pure=True, pack=1)
+        qrow = tl.arange(0, 8)
+        offset = (((batch * H + head[None, :]) * D + dim[None, :]) * (S // 2)
+                  + tile * (TILE // 2) + i * 8 + qrow[:, None])
+        tl.store(CODES + offset, codes, col[None, :] < H * D)
+        group = tile * (TILE // 16) + i
+        offset = ((batch * H + head) * D * (S // 16) + (dim // 64) * 64 * (S // 16)
+                  + (group // 4) * 256 + group % 4 + (dim % 64 // 16) * 4 + (dim % 16) * 16)
+        tl.store(SCALES + offset, sf, col < H * D)
+    even, odd = tl.split(tl.reshape(tl.trans(acc), (BLOCK, 4, 2, 2)))
+    y0, y2 = tl.split(even)
+    y1, y3 = tl.split(odd)
+    y0, y1 = _left_four(y0, BLOCK), _left_four(y1, BLOCK)
+    y2, y3 = _left_four(y2, BLOCK), _left_four(y3, BLOCK)
+    total = (y0 + y2) + (y1 + y3)
+    count = tl.load(VALID + tile).to(tl.float32)
+    mean = tl.div_rn(total, count)
+    tl.store(POOL + ((batch * (S // TILE) + tile) * H * D) + col, mean, col < H * D)
+
+
+def pool_quantize_values(
+    value: torch.Tensor, valid: torch.Tensor, input_index: torch.Tensor | None = None, *, tile_elems: int = 64,
+) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+    """Return FP32 tile means and contiguous native FP4 V codes/scales.
+
+    The caller uses 64- or 128-token tiles, with all logical rows present. Values
+    must have unit feature stride and 128-wide heads; token rows and heads
+    may be strided. This is an eager, inference-only Blackwell path.
+    """
+    if (value.ndim != 4 or not value.is_cuda or value.dtype != torch.bfloat16
+            or value.shape[-1] != 128 or value.stride(-1) != 1
+            or (input_index is None and value.shape[1] % 128)):
+        raise ValueError("Value preparation requires CUDA BF16 BSHD, head-128 and padded token rows")
+    if torch.is_grad_enabled() or torch.compiler.is_compiling():
+        raise RuntimeError("Value preparation requires eager inference")
+    if tile_elems not in (64, 128):
+        raise ValueError("Value preparation supports 64- and 128-token tiles")
+    batch, seq_len, heads, dim = value.shape
+    if input_index is not None:
+        if (input_index.ndim != 1 or input_index.numel() % 128 or input_index.device != value.device
+                or input_index.dtype not in (torch.int32, torch.int64) or not input_index.is_contiguous()):
+            raise ValueError("Compact values require a contiguous aligned integer row map")
+        seq_len = input_index.numel()
+    if (valid.shape != (seq_len // tile_elems,) or valid.device != value.device or not valid.is_contiguous()
+            or valid.dtype not in (torch.int32, torch.int64)):
+        raise ValueError("Value preparation requires contiguous integer tile validity counts")
+    pooled = torch.empty((batch, seq_len // tile_elems, heads, dim), device=value.device, dtype=torch.float32)
+    codes = torch.empty((batch, heads, dim, seq_len // 2), device=value.device, dtype=torch.uint8)
+    scales = torch.empty((batch, heads, dim, seq_len // 16), device=value.device, dtype=torch.float8_e4m3fn)
+    # Keep wide indexing for large/strided batches without charging small clips
+    # for 64-bit vector arithmetic on every packed value.
+    span = 1 + sum((size - 1) * stride for size, stride in zip(value.shape, value.stride(), strict=True))
+    block = 128 if input_index is not None else 32
+    span = max(span, batch * seq_len * heads * dim)
+    _pool_and_quantize[(triton.cdiv(heads * dim, block), seq_len // tile_elems, batch)](
+        value, valid, pooled, codes, scales, *value.stride()[:3], seq_len, heads, dim, block, span >= 2**31,
+        INPUT_INDEX=input_index, GATHER_INPUT=input_index is not None, TILE=tile_elems, num_warps=4,
+    )
+    return pooled.permute(0, 2, 1, 3), (codes, scales)

fastvideo/models/dits/minimax_h3_vsa_fp4.py

@@ -19,6 +19,7 @@ with one gather before ``to_out``.
 from __future__ import annotations
 
 import math
+from collections.abc import Callable
 from typing import Any
 
 import torch
@@ -28,6 +29,16 @@ import fastvideo.envs as envs
 from fastvideo.attention.backends.video_sparse_attn_h3 import (MiniMaxH3VSAMetadata, _build_block_mask, _pool_tiles)
 
 
+def _tile_mask_to_fp4_blocks(api: Any, mask: torch.Tensor, tile: int, valid: torch.Tensor):
+    if (mask.is_cuda and mask.dtype == torch.bool and mask.shape[-1] == mask.shape[-2]
+            and tile > 0 and tile % 64 == 0 and valid.is_contiguous()
+            and not torch.compiler.is_compiling()):
+        from fastvideo.models.dits.minimax_h3_fusions.block_lists import tile_mask_to_fp4_blocks
+
+        return tile_mask_to_fp4_blocks(mask, tile, valid)
+    return api.vsa_tile_mask_to_fp4_blocks(mask, tile, valid)
+
+
 def _fp4_kernel_kwargs(two_level_p: bool | None = None) -> dict[str, object]:
     """Quantization options for the FP4 video kernel.
 
@@ -42,7 +53,10 @@ def _fp4_kernel_kwargs(two_level_p: bool | None = None) -> dict[str, object]:
 
 
 def _sparse_fp4_attention(api: Any, *args: Any, two_level_p: bool | None = None,
-                          smooth_k: bool | None = None) -> torch.Tensor:
+                          smooth_k: bool | None = None,
+                          prepared_key: tuple[torch.Tensor, torch.Tensor] | None = None,
+                          prepared_query: tuple[torch.Tensor, torch.Tensor] | None = None,
+                          prepared_value: tuple[torch.Tensor, torch.Tensor] | None = None) -> torch.Tensor:
     """Run the block-sparse FP4 kernel on lists built by vsa_tile_mask_to_fp4_blocks.
 
     Those lists are in range by construction, so skip the host-sync validation where the installed
@@ -52,8 +66,30 @@ def _sparse_fp4_attention(api: Any, *args: Any, two_level_p: bool | None = None,
     quantizes K. The shift is exact under softmax, so it only changes quantization error.
     """
     smooth = envs.FASTVIDEO_H3_VSA_FP4_K_SMOOTH.get() if smooth_k is None else smooth_k
+    key = args[1]
+    if (smooth and prepared_key is None and key.is_cuda and key.dtype == torch.bfloat16
+            and key.shape[-1] == 128 and key.shape[1] % 128 == 0 and key.stride(-1) == 1
+            and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
+            and torch.cuda.get_device_capability(key.device)[0] >= 10
+            and all(hasattr(api, name) for name in ("blockscaled_fp4_attn_sparse", "_zero_delta_s",
+                                                    "scale_and_quant_fp4", "scale_and_quant_fp4_transpose"))):
+        from fastvideo.models.dits.minimax_h3_fusions.key_quant import quantize_smoothed_keys_only
+
+        prepared_key = quantize_smoothed_keys_only(key, key.mean(dim=1, keepdim=True))
+        smooth = False
     if smooth:
-        args = (args[0], args[1] - args[1].mean(dim=1, keepdim=True), *args[2:])
+        args = (args[0], key - key.mean(dim=1, keepdim=True), *args[2:])
+    if prepared_key is not None or prepared_query is not None or prepared_value is not None:
+        query, key, value, indices, counts, valid, quadrants = args
+        delta = api._zero_delta_s(query.shape[0], query.shape[2], key.shape[1], query.device)
+        single_level = _fp4_kernel_kwargs(two_level_p).get("single_level_p_quant", True)
+        return api.blockscaled_fp4_attn_sparse(
+            prepared_query if prepared_query is not None else api.scale_and_quant_fp4(query.transpose(1, 2)),
+            prepared_key if prepared_key is not None else api.scale_and_quant_fp4_permute(key.transpose(1, 2)),
+            prepared_value if prepared_value is not None else api.scale_and_quant_fp4_transpose(value.transpose(1, 2)),
+            delta, key.shape[1],
+            indices, counts, valid, quadrants, False, query.dtype == torch.bfloat16, single_level,
+        )[0]
     try:
         return api.sageattn_blackwell_sparse_bshd(*args, validate=False, **_fp4_kernel_kwargs(two_level_p))
     except TypeError as error:
@@ -143,12 +179,30 @@ class _KeyIndexCache:
     tile: int
     n_tiles: int
     _valid_key_index: torch.Tensor | None = None
+    _valid_key_inverse: torch.Tensor | None = None
 
     def valid_key_index(self, kv_valid: torch.Tensor, half: int) -> torch.Tensor:
         if self._valid_key_index is None:
             self._valid_key_index = _valid_key_index(kv_valid, half, self.n_tiles * self.tile)
         return self._valid_key_index
 
+    def valid_keys_from_tiles(self, tile_valid: torch.Tensor, half: int) -> tuple[torch.Tensor, torch.Tensor]:
+        """Prepare prefix maps before mask selection, splitting 128-token tile counts into native halves."""
+        native_valid = tile_valid
+        if self._valid_key_index is None and self.tile != half:
+            offsets = torch.arange(self.tile // half, device=tile_valid.device) * half
+            native_valid = (tile_valid[:, None] - offsets[None, :]).clamp(0, half).flatten()
+        index = self.valid_key_index(native_valid, half)
+        return index, self.valid_key_inverse(native_valid, half)
+
+    def valid_key_inverse(self, kv_valid: torch.Tensor, half: int) -> torch.Tensor:
+        if self._valid_key_inverse is None:
+            index = self.valid_key_index(kv_valid, half)
+            self._valid_key_inverse = torch.full((self.n_tiles * self.tile,), -1,
+                                                 device=index.device, dtype=torch.int64)
+            self._valid_key_inverse[index] = torch.arange(index.numel(), device=index.device)
+        return self._valid_key_inverse
+
 
 class _TileLayout(_KeyIndexCache):
     """Per-step tile-order state shared by every block of one forward."""
@@ -164,6 +218,20 @@ class _TileLayout(_KeyIndexCache):
         self.cos = cos.new_zeros((self.rows, cos.shape[-1])).index_copy_(0, self.untile, cos)
         self.sin = sin.new_zeros((self.rows, sin.shape[-1])).index_copy_(0, self.untile, sin)
         self._buf: torch.Tensor | None = None
+        self._quant_inverse: torch.Tensor | None = None
+
+    def gather_quantized(self, quantized: torch.Tensor, scales: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Reorder per-row FP4 codes and scales without expanding them back to BF16."""
+        from fastvideo.models.dits.minimax_h3_fusions.quant_tile import tile_quantized
+
+        return tile_quantized(quantized, scales, self.packed_row_inverse())
+
+    def packed_row_inverse(self) -> torch.Tensor:
+        """Geometry-only map from padded tile rows to compact projection rows."""
+        if self._quant_inverse is None:
+            self._quant_inverse = torch.full((self.rows,), -1, device=self.untile.device, dtype=torch.int64)
+            self._quant_inverse[self.untile] = torch.arange(self.untile.numel(), device=self.untile.device)
+        return self._quant_inverse
 
     def gather_in(self, x: torch.Tensor) -> torch.Tensor:
         """Packed ``[B, L, C]`` -> tile-ordered ``[B, rows, C]``; pad rows stay zero.
@@ -207,6 +275,48 @@ def _shared_input_projections(linears: tuple[Any, ...], x: torch.Tensor) -> list
     return [m.apply(linear, x, linear.bias, pre_quantized=pre) for m, linear in zip(methods, linears, strict=True)]
 
 
+def _shared_qkv_gate_scale(attn: Any, x: torch.Tensor) -> torch.Tensor | None:
+    """Return the common unit scale when QKV and the gate can share packed input."""
+    from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod
+
+    linears = (attn.to_q, attn.to_k, attn.to_v, attn.to_gate_compress)
+    if attn.to_gate_compress is None or x.shape[0] != 1 or x.shape[-1] % 64 or not attn._gate_active():
+        return None
+    if not all(type(getattr(linear, "quant_method", None)) is NVFP4QuantizeMethod
+               and linear.quant_method.wants_prequantized_input()
+               and linear.quant_method.uses_unit_activation_scale(linear) for linear in linears):
+        return None
+    return linears[0].quant_method.x_global_sf
+
+
+def _shared_qkv_gate_projections(
+    attn: Any, x: torch.Tensor, layout: _TileLayout | None = None,
+    pre_quantized: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
+):
+    """Share packed activation preparation across QKV and the compression gate.
+
+    FP4 groups never cross rows, so tile reordering preserves every code and
+    scale. SP projects local packed rows directly and omits the tile layout.
+    Compatible fused modulation producers can supply the packed input.
+    """
+    if _shared_qkv_gate_scale(attn, x) is None:
+        if pre_quantized is not None:
+            raise RuntimeError("Prequantized H3 input requires compatible QKV and gate projections")
+        return None
+    linears = (attn.to_q, attn.to_k, attn.to_v, attn.to_gate_compress)
+    pre = linears[0].quant_method.quantize_input(x) if pre_quantized is None else pre_quantized
+    qkv_pre, shape_view = pre, x
+    if layout is not None:
+        q, scales = layout.gather_quantized(pre[0], pre[1])
+        qkv_pre = (q, scales, pre[2])
+        # apply(pre_quantized=...) only reads x.shape. This view supplies the
+        # tile-order shape without allocating or scattering a full BF16 tensor.
+        shape_view = x[:, :1].expand(-1, layout.rows, -1)
+    projected = [linear.quant_method.apply(linear, shape_view, linear.bias, pre_quantized=qkv_pre)
+                 for linear in linears[:3]]
+    return projected, pre
+
+
 @torch.compile(dynamic=True, fullgraph=True)
 def _fp8_per_token(x: torch.Tensor) -> torch.Tensor:
     """``[..., D]`` BF16 -> per-(row, head) FP8 E4M3 -> BF16, the prefix Q/K precision of the QK study.
@@ -232,7 +342,10 @@ def _valid_key_index(kv_valid: torch.Tensor, half: int, logical: int) -> torch.T
 
 
 def _prefix_dense_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, prefix_rows: int,
-                            logical: int, valid_index: torch.Tensor, precision: str) -> torch.Tensor:
+                            logical: int, valid_index: torch.Tensor, precision: str,
+                            key_is_smoothed: bool = False,
+                            prepared_kv: tuple[torch.Tensor, torch.Tensor] | None = None,
+                            query_is_fp8: bool = False) -> torch.Tensor:
     """Dense attention for the prefix query rows, ``[B, H, prefix_rows, D]``.
 
     VSA gives every prefix (non-video) query tile every key column, so those rows need no
@@ -246,12 +359,32 @@ def _prefix_dense_attention(query: torch.Tensor, key: torch.Tensor, value: torch
     import torch.nn.functional as F
 
     q, k, v = query[:, :prefix_rows], key[:, :logical], value[:, :logical]
-    k = k - k.mean(dim=1, keepdim=True)
-    if precision == "fp8":
-        q, k = _fp8_per_token(q), _fp8_per_token(k)
+    if not key_is_smoothed:
+        k = k - k.mean(dim=1, keepdim=True)
     # Attend only the valid keys. Gathering them keeps SDPA on its unmasked (flash) kernel; a boolean
     # mask over all rows falls back to the memory-efficient kernel, 3-4x slower at OmniRef prefix sizes.
-    k, v = k.index_select(1, valid_index), v.index_select(1, valid_index)
+    if prepared_kv is not None:
+        q = _fp8_per_token(q) if precision == "fp8" and not query_is_fp8 else q
+        k, v = prepared_kv
+    elif (precision == "fp8" and k.is_cuda and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
+            and k.dtype == torch.bfloat16 and v.dtype == k.dtype
+            and k.shape[-1] == 128 and k.stride(-1) == 1 and v.stride(-1) == 1
+            and torch.cuda.get_device_capability(k.device)[0] >= 10):
+        from fastvideo.models.dits.minimax_h3_fusions.prefix_quant import quantize_gather_prefix_kv
+
+        q = _fp8_per_token(q) if not query_is_fp8 else q
+        k, v = quantize_gather_prefix_kv(k, v, valid_index)
+    else:
+        if precision == "fp8":
+            q = _fp8_per_token(q) if not query_is_fp8 else q
+            k = _fp8_per_token(k)
+        k, v = k.index_select(1, valid_index), v.index_select(1, valid_index)
+    if prefix_rows == 15360 and precision == "fp8":
+        from fastvideo.models.dits.minimax_h3_fusions.prefix_attention import try_prefix_attention
+
+        specialized = try_prefix_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2))
+        if specialized is not None:
+            return specialized
     return F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2),
                                           scale=1.0 / math.sqrt(q.shape[-1]))
 
@@ -259,10 +392,20 @@ def _prefix_dense_attention(query: torch.Tensor, key: torch.Tensor, value: torch
 def _fp4_fine_attention(api: Any, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,
                         q2k_idx: torch.Tensor, q2k_num: torch.Tensor, kv_valid: torch.Tensor,
                         q2k_quad: torch.Tensor | None, *, logical: int, prefix_rows: int, precision: str,
-                        options: dict[str, Any], layout: _KeyIndexCache | None = None) -> torch.Tensor:
+                        options: dict[str, Any], layout: _KeyIndexCache | None = None,
+                        return_split: bool = False,
+                        prepared_query: tuple[torch.Tensor, torch.Tensor] | None = None,
+                        prepared_value: tuple[torch.Tensor, torch.Tensor] | None = None,
+                        prefix_query_is_fp8: bool = False,
+                        value_index: torch.Tensor | None = None,
+                        prepared_keys: tuple[
+                            tuple[torch.Tensor, torch.Tensor], torch.Tensor, torch.Tensor
+                        ] | None = None,
+                        ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, int]:
     numerics = {"two_level_p": options.get("two_level_p"), "smooth_k": options.get("smooth_k")}
     if not prefix_rows:
-        return _sparse_fp4_attention(api, query, key, value, q2k_idx, q2k_num, kv_valid, q2k_quad, **numerics)
+        return _sparse_fp4_attention(api, query, key, value, q2k_idx, q2k_num, kv_valid, q2k_quad,
+                                      prepared_query=prepared_query, prepared_value=prepared_value, **numerics)
     if not getattr(api, "KEY_LENGTH_SIZING", False):
         raise RuntimeError("the MiniMax-H3 FP4 prefix split needs a fastvideo-kernel build whose "
                            "sageattn_blackwell_sparse_bshd sizes by key length; rebuild fastvideo-kernel "
@@ -271,15 +414,52 @@ def _fp4_fine_attention(api: Any, query: torch.Tensor, key: torch.Tensor, value:
         valid_index = layout.valid_key_index(kv_valid, api.HALF_N)
     else:
         valid_index = _valid_key_index(kv_valid, api.HALF_N, logical)
+    smooth = numerics["smooth_k"]
+    if smooth is None:
+        smooth = envs.FASTVIDEO_H3_VSA_FP4_K_SMOOTH.get()
+    shared_smoothing = bool(smooth and logical == key.shape[1])
+    prepared_key, prepared_kv = None, None
+    if prepared_keys is not None:
+        prepared_key, prefix_key, prefix_value = prepared_keys
+        prepared_kv = (prefix_key, prefix_value)
+        numerics["smooth_k"] = False
+    elif (shared_smoothing and precision == "fp8" and layout is not None and key.is_cuda
+            and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
+            and key.dtype == torch.bfloat16 and value.dtype == key.dtype and key.shape[-1] == 128
+            and key.stride(-1) == 1 and value.stride(-1) == 1
+            and torch.cuda.get_device_capability(key.device)[0] >= 10
+            and all(hasattr(api, name) for name in ("blockscaled_fp4_attn_sparse", "_zero_delta_s",
+                                                    "scale_and_quant_fp4", "scale_and_quant_fp4_transpose"))):
+        from fastvideo.models.dits.minimax_h3_fusions.key_quant import quantize_smoothed_keys
+
+        inverse = layout.valid_key_inverse(kv_valid, api.HALF_N)
+        prepared_key, prefix_key, prefix_value = quantize_smoothed_keys(
+            key, value, key.mean(dim=1, keepdim=True), inverse, valid_index.numel(), value_index=value_index,
+        )
+        prepared_kv = (prefix_key, prefix_value)
+        numerics["smooth_k"] = False
+    elif shared_smoothing:
+        # Both branches reduce the same rows with the same strides and dtype.
+        # Keep one rounded BF16 subtraction result instead of repeating it.
+        key = key - key.mean(dim=1, keepdim=True)
+        numerics["smooth_k"] = False
     with STAGES.span("prefix_attention"):
-        prefix_out = _prefix_dense_attention(query, key, value, prefix_rows, logical, valid_index, precision)
+        prefix_out = _prefix_dense_attention(query, key, value, prefix_rows, logical, valid_index, precision,
+                                             key_is_smoothed=shared_smoothing, prepared_kv=prepared_kv,
+                                             query_is_fp8=prefix_query_is_fp8)
     # A tile-64 prefix can end halfway through a native 128-row block. Keep that whole query block and its
     # original quadrant mask, then discard the overlapping prefix half and trailing padding from its output.
     head = prefix_rows // _BLOCK
     start = head * _BLOCK
+    # With prepared FP4 queries, the raw suffix supplies only dtype/device/head
+    # metadata; query may hold just the BF16 prefix without a video intermediate.
     video_out = _sparse_fp4_attention(api, query[:, start:], key, value,
                                       q2k_idx[:, :, head:].contiguous(), q2k_num[:, :, head:].contiguous(), kv_valid,
-                                      None if q2k_quad is None else q2k_quad[:, :, head:].contiguous(), **numerics)
+                                      None if q2k_quad is None else q2k_quad[:, :, head:].contiguous(),
+                                      prepared_key=prepared_key, prepared_query=prepared_query,
+                                      prepared_value=prepared_value, **numerics)
+    if return_split:
+        return prefix_out, video_out, start
     return torch.cat((prefix_out, video_out[:, :, prefix_rows - start:logical - start]), dim=2)
 
 
@@ -353,54 +533,188 @@ def vsa_tile_first_attention(attn: Any, hidden_states: torch.Tensor,
 
 
 def vsa_fp4_attention(attn: Any, hidden_states: torch.Tensor, rotary_emb: tuple[torch.Tensor, torch.Tensor],
-                      meta: MiniMaxH3VSAMetadata, use_fused_rope: bool) -> torch.Tensor:
+                      meta: MiniMaxH3VSAMetadata, use_fused_rope: bool,
+                      pre_quantized: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
+                      ) -> torch.Tensor | tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
     """Attention core for ``MiniMaxH3Attention``; returns the pre-``to_out`` ``[B, L, H*D]``."""
     api = _api()
     layout = _layout_for(meta, rotary_emb)
     heads, dim = attn.num_attention_heads, attn.attention_head_dim
+    prefix_mode, prefix_rows = _prefix_settings(attn, meta, layout.n_tiles)
+    sim_fp8 = envs.FASTVIDEO_H3_SIM_SP_FP8.get()
+    prepared_query, q_pooled = None, None
+    prepared_key_mean = None
+    prefix_query_is_fp8 = False
+    smooth = getattr(attn, "_vsa_fp4_options", {}).get("smooth_k")
+    smooth = envs.FASTVIDEO_H3_VSA_FP4_K_SMOOTH.get() if smooth is None else smooth
+    # The measured cuDNN/GB10 route preserves every projected row while omitting
+    # zero padding from QKV GEMMs. Other backends keep the tiled projection path.
+    compact = (use_fused_rope and not sim_fp8 and layout.tile == 64
+               and layout.rows == layout.n_tiles * layout.tile and dim == 128
+               and hidden_states.is_cuda and hidden_states.dtype == torch.bfloat16
+               and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
+               and (not prefix_rows or (prefix_mode == "fp8" and smooth))
+               and envs.FASTVIDEO_NVFP4_MM_BACKEND.get() == "cudnn"
+               and torch.cuda.get_device_capability(hidden_states.device) == (12, 1)
+               and _shared_qkv_gate_scale(attn, hidden_states) is not None
+               and all(hasattr(api, name) for name in ("blockscaled_fp4_attn_sparse", "_zero_delta_s",
+                                                      "scale_and_quant_fp4", "scale_and_quant_fp4_permute",
+                                                      "scale_and_quant_fp4_transpose")))
+    input_index = layout.packed_row_inverse() if compact else None
     with STAGES.span("qkv_proj_rope"):
-        x_tiles = layout.gather_in(hidden_states)
-        query, key, value = (t.unflatten(-1, (heads, dim))
-                             for t in _shared_input_projections((attn.to_q, attn.to_k, attn.to_v), x_tiles))
-        if use_fused_rope:
+        shared = _shared_qkv_gate_projections(attn, hidden_states, None if compact else layout, pre_quantized)
+        if shared is None:
+            x_tiles = layout.gather_in(hidden_states)
+            projected = _shared_input_projections((attn.to_q, attn.to_k, attn.to_v), x_tiles)
+            gate_pre = None
+        else:
+            projected, gate_pre = shared
+        query, key, value = (t.unflatten(-1, (heads, dim)) for t in projected)
+        del projected, shared
+        if compact:
+            from fastvideo.models.dits.minimax_h3_fusions.qknorm_rope import gather_qknorm_rope
+
+            cos, sin = layout.cos.to(query.dtype), layout.sin.to(query.dtype)
+            if prefix_rows:
+                from fastvideo.models.dits.minimax_h3_fusions.query_pool import prepare_pooled_queries
+
+                query, prepared_query, q_pooled = prepare_pooled_queries(
+                    query, attn.norm_q.weight, cos, sin, attn.norm_q.eps, input_index,
+                    meta.variable_block_sizes, prefix_rows,
+                )
+            else:
+                query, prepared_query = gather_qknorm_rope(
+                    query, attn.norm_q.weight, cos, sin, attn.norm_q.eps, input_index, fp4_start=0,
+                )
+            if prefix_rows and key.shape[0] == 1 and heads == 56 and input_index.numel() == 18432 and cos.shape[1] == 96:
+                from fastvideo.models.dits.minimax_h3_fusions.key_norm_mean import gather_key_norm_mean
+
+                key, prepared_key_mean = gather_key_norm_mean(
+                    key, attn.norm_k.weight, cos, sin, attn.norm_k.eps, input_index,
+                )
+            else:
+                key, _ = gather_qknorm_rope(key, attn.norm_k.weight, cos, sin, attn.norm_k.eps, input_index)
+        elif use_fused_rope:
             from fastvideo.models.dits.minimax_h3_fusions import fused_qknorm_rope
             cos, sin = layout.cos.to(query.dtype), layout.sin.to(query.dtype)
-            query = fused_qknorm_rope(query, attn.norm_q.weight, cos, sin, attn.norm_q.eps)
+            if (not sim_fp8 and query.is_cuda and query.dtype == torch.bfloat16
+                    and dim == 128 and query.shape[1] % 128 == 0
+                    and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
+                    and torch.cuda.get_device_capability(query.device)[0] >= 10
+                    and all(hasattr(api, name) for name in ("blockscaled_fp4_attn_sparse", "_zero_delta_s",
+                                                            "scale_and_quant_fp4", "scale_and_quant_fp4_permute",
+                                                            "scale_and_quant_fp4_transpose"))):
+                from fastvideo.models.dits.minimax_h3_fusions.qknorm_rope import qknorm_rope_quantize_fp4
+
+                if (layout.tile == 128 and query.shape[1] == layout.n_tiles * layout.tile
+                        and prefix_mode == "fp8" and prefix_rows
+                        and torch.cuda.get_device_capability(query.device) == (12, 1)):
+                    from fastvideo.models.dits.minimax_h3_fusions.query_pool import prepare_pooled_queries
+
+                    query, prepared_query, q_pooled = prepare_pooled_queries(
+                        query, attn.norm_q.weight, cos, sin, attn.norm_q.eps, None,
+                        meta.variable_block_sizes, prefix_rows, tile_elems=128, prefix_fp8=True, split_prefix=True,
+                    )
+                    prefix_query_is_fp8 = True
+                else:
+                    query, prepared_query = qknorm_rope_quantize_fp4(
+                        query, attn.norm_q.weight, cos, sin, attn.norm_q.eps, prefix_rows // _BLOCK * _BLOCK,
+                    )
+            else:
+                query = fused_qknorm_rope(query, attn.norm_q.weight, cos, sin, attn.norm_q.eps)
             key = fused_qknorm_rope(key, attn.norm_k.weight, cos, sin, attn.norm_k.eps)
         else:
             rope = (layout.cos, layout.sin)
             query = attn._apply_rotary_emb(attn.norm_q(query), rope)
             key = attn._apply_rotary_emb(attn.norm_k(key), rope)
 
-    sim_fp8 = envs.FASTVIDEO_H3_SIM_SP_FP8.get()
     if sim_fp8:
         query, key, value = (_fp8_roundtrip(t) for t in (query, key, value))
 
     vbs = meta.variable_block_sizes
     logical = layout.n_tiles * layout.tile
+    prepared_keys, k_pooled = None, None
+    v_pooled, prepared_value = None, None
+    fused_gate = (attn.to_gate_compress is not None and attn._gate_active()
+                  and not sim_fp8 and query.is_cuda and query.dtype == torch.bfloat16)
+    omni_key_pool = (layout.tile == 128 and logical == key.shape[1] and prefix_mode == "fp8" and smooth
+                     and not sim_fp8 and key.is_cuda and key.dtype == value.dtype == torch.bfloat16
+                     and dim == 128 and key.stride(-1) == value.stride(-1) == 1
+                     and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
+                     and torch.cuda.get_device_capability(key.device)[0] >= 10
+                     and all(hasattr(api, name) for name in ("blockscaled_fp4_attn_sparse", "_zero_delta_s",
+                                                             "scale_and_quant_fp4", "scale_and_quant_fp4_transpose")))
+    if prefix_rows and (compact or omni_key_pool):
+        from fastvideo.models.dits.minimax_h3_fusions.key_quant import pool_quantize_smoothed_keys
+
+        valid_index, inverse = layout.valid_keys_from_tiles(vbs, api.HALF_N)
+        key_mean = key.mean(dim=1, keepdim=True) if prepared_key_mean is None else prepared_key_mean
+        if fused_gate and (compact or torch.cuda.get_device_capability(key.device) == (12, 1)):
+            from fastvideo.models.dits.minimax_h3_fusions.kv_pool import pool_quantize_prefix_kv
+
+            key_result, value_result = pool_quantize_prefix_kv(
+                key, value, key_mean, inverse, valid_index.numel(), vbs, input_index,
+                tile_elems=layout.tile,
+            )
+            packed_key, prefix_key, prefix_value, k_pooled = key_result
+            v_pooled, prepared_value = value_result
+        else:
+            packed_key, prefix_key, prefix_value, k_pooled = pool_quantize_smoothed_keys(
+                key, value, key_mean, inverse, valid_index.numel(), vbs,
+                value_index=input_index, tile_elems=layout.tile,
+            )
+        prepared_keys = (packed_key, prefix_key, prefix_value)
     with STAGES.span("select_mask"):
-        q_pooled = _pool_tiles(query[:, :logical], vbs, layout.tile)
-        k_pooled = _pool_tiles(key[:, :logical], vbs, layout.tile)
+        if q_pooled is None:
+            q_pooled = _pool_tiles(query[:, :logical], vbs, layout.tile)
+        if k_pooled is None:
+            k_pooled = _pool_tiles(key[:, :logical], vbs, layout.tile)
         scores = torch.matmul(q_pooled, k_pooled.transpose(-2, -1)) / (dim**0.5)
         sparsity = 0.0 if attn._layer_idx in meta.dense_layers else meta.VSA_sparsity
         mask = _build_block_mask(scores, meta.num_prefix_tiles, sparsity, meta.exempt, meta.video_tile_spans,
                                  meta.span_sparsities)
-        q2k_idx, q2k_num, kv_valid, q2k_quad = api.vsa_tile_mask_to_fp4_blocks(mask, layout.tile, vbs)
-    prefix_mode, prefix_rows = _prefix_settings(attn, meta, layout.n_tiles)
+        q2k_idx, q2k_num, kv_valid, q2k_quad = _tile_mask_to_fp4_blocks(api, mask, layout.tile, vbs)
+    if (prepared_value is None and fused_gate and layout.tile in (64, 128) and (compact or logical == value.shape[1]) and dim == 128
+            and value.stride(-1) == 1 and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
+            and torch.cuda.get_device_capability(value.device)[0] >= 10
+            and all(hasattr(api, name) for name in ("blockscaled_fp4_attn_sparse", "_zero_delta_s",
+                                                    "scale_and_quant_fp4", "scale_and_quant_fp4_permute"))):
+        from fastvideo.models.dits.minimax_h3_fusions.value_quant import pool_quantize_values
+
+        v_pooled, prepared_value = pool_quantize_values(value, vbs, input_index=input_index, tile_elems=layout.tile)
     with STAGES.span("fp4_attention"):
         out = _fp4_fine_attention(api, query, key, value, q2k_idx, q2k_num, kv_valid, q2k_quad,
                                   logical=logical, prefix_rows=prefix_rows, precision=prefix_mode,
-                                  options=getattr(attn, "_vsa_fp4_options", {}), layout=layout)
-    with STAGES.span("out_untile"):
-        out = out.transpose(1, 2).index_select(1, layout.untile)  # [B, L, H, D], packed order
+                                  options=getattr(attn, "_vsa_fp4_options", {}), layout=layout,
+                                  return_split=fused_gate, prepared_query=prepared_query, prepared_value=prepared_value,
+                                  prefix_query_is_fp8=prefix_query_is_fp8,
+                                  value_index=input_index, prepared_keys=prepared_keys)
+    if not fused_gate:
+        with STAGES.span("out_untile"):
+            out = out.transpose(1, 2).index_select(1, layout.untile)  # [B, L, H, D], packed order
     if sim_fp8:
         out = _fp8_roundtrip(out)
 
     if attn.to_gate_compress is not None and attn._gate_active():
         with STAGES.span("gate_compress"):
-            gate, _ = attn.to_gate_compress(hidden_states)
-            v_pooled = _pool_tiles(value[:, :logical], vbs, layout.tile)
-            out_c = torch.matmul(torch.softmax(scores, dim=-1), v_pooled).permute(0, 2, 1, 3).to(out.dtype)
+            if gate_pre is None:
+                gate, _ = attn.to_gate_compress(hidden_states)
+            else:
+                linear = attn.to_gate_compress
+                gate = linear.quant_method.apply(linear, hidden_states, linear.bias, pre_quantized=gate_pre)
+            if v_pooled is None:
+                v_pooled = _pool_tiles(value[:, :logical], vbs, layout.tile)
+            out_c = torch.matmul(torch.softmax(scores, dim=-1), v_pooled).permute(0, 2, 1, 3).to(query.dtype)
+            if fused_gate:
+                from fastvideo.models.dits.minimax_h3_fusions.untile_gate import untile_gate
+                from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod
+
+                method = getattr(attn.to_out, "quant_method", None)
+                if (type(method) is NVFP4QuantizeMethod and method.wants_prequantized_input()
+                        and method.uses_unit_activation_scale(attn.to_out) and gate.shape[-1] % 64 == 0):
+                    prepared = untile_gate(out, out_c, gate, layout.untile, layout.row_tile, method.x_global_sf)
+                    return gate, prepared
+                return untile_gate(out, out_c, gate, layout.untile, layout.row_tile)
             gate = gate.unflatten(-1, (heads, dim))
             chunk = envs.FASTVIDEO_H3_GATE_CHUNK_ROWS.get()
             if chunk > 0 and out.shape[1] > chunk:
@@ -490,12 +804,23 @@ def _apply_gate(out: torch.Tensor, out_c: torch.Tensor, row_tile: torch.Tensor,
     return out + out_c.index_select(0, row_tile) * gate
 
 
-def _all_to_all(payload: torch.Tensor, scale: torch.Tensor, group: Any) -> tuple[torch.Tensor, torch.Tensor]:
+def _all_to_all(payload: torch.Tensor, scale: torch.Tensor, group: Any,
+                overlap: Callable[[], None] | None = None) -> tuple[torch.Tensor, torch.Tensor]:
     import torch.distributed as dist
     recv = torch.empty_like(payload)
     recv_scale = torch.empty_like(scale)
-    dist.all_to_all_single(recv.view(torch.uint8), payload.view(torch.uint8), group=group)
-    dist.all_to_all_single(recv_scale, scale, group=group)
+    if overlap is None:
+        dist.all_to_all_single(recv.view(torch.uint8), payload.view(torch.uint8), group=group)
+        dist.all_to_all_single(recv_scale, scale, group=group)
+    else:
+        # Keep the current compute stream independent of NCCL until the local
+        # projection has been enqueued. wait() then orders consumers of recv.
+        payload_work = dist.all_to_all_single(recv.view(torch.uint8), payload.view(torch.uint8),
+                                              group=group, async_op=True)
+        scale_work = dist.all_to_all_single(recv_scale, scale, group=group, async_op=True)
+        overlap()
+        payload_work.wait()
+        scale_work.wait()
     return recv, recv_scale
 
 
@@ -513,6 +838,7 @@ class _SPTileLayout(_KeyIndexCache):
         # Rows past the sequence are SP padding; their outputs are discarded.
         self.local_row_tile = row_tile[local.clamp_max(self.seq_len - 1)]
         self._buf: torch.Tensor | None = None
+        self._inverse: torch.Tensor | None = None
 
     def tiles_from(self, qkv: torch.Tensor) -> torch.Tensor:
         """``[3, L, Hs, D]`` packed -> ``[3, rows, Hs, D]`` tile order with zero padding (reused buffer)."""
@@ -521,9 +847,26 @@ class _SPTileLayout(_KeyIndexCache):
             self._buf = qkv.new_zeros(shape)
         return self._buf.index_copy_(1, self.untile, qkv)
 
+    def fp8_buffers(self, payload: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Return the invariant row map and scratch tensor overwritten by each receive."""
+        if self._inverse is None:
+            self._inverse = torch.full((self.rows,), -1, device=payload.device, dtype=torch.int64)
+            self._inverse[self.untile] = torch.arange(self.seq_len, device=payload.device)
+        shape = (3, self.rows, *payload.shape[3:])
+        if self._buf is None or self._buf.shape != shape:
+            self._buf = torch.empty(shape, dtype=torch.bfloat16, device=payload.device)
+        return self._inverse, self._buf
+
+    def tiles_from_fp8(self, payload: torch.Tensor, scale: torch.Tensor, *, query_start: int | None = None):
+        from fastvideo.models.dits.minimax_h3_fusions.sp_exchange import unpack_fp8_tiles
+
+        inverse, output = self.fp8_buffers(payload)
+        return unpack_fp8_tiles(payload, scale, inverse, output, query_start=query_start)
+
 
 def vsa_fp4_attention_sp(attn: Any, hidden_states: torch.Tensor, rotary_emb: tuple[torch.Tensor, torch.Tensor],
-                         meta: MiniMaxH3VSAMetadata, use_fused_rope: bool, sp_group: Any) -> torch.Tensor:
+                         meta: MiniMaxH3VSAMetadata, use_fused_rope: bool, sp_group: Any,
+                         pre_quantized: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None) -> torch.Tensor:
     """Ulysses-SP attention core on local sequence rows ``[1, rows, C]``; returns pre-``to_out`` ``[1, rows, H*D]``."""
     import torch.distributed as dist
 
@@ -531,29 +874,65 @@ def vsa_fp4_attention_sp(attn: Any, hidden_states: torch.Tensor, rotary_emb: tup
     world, rank = sp_group.world_size, sp_group.rank_in_group
     heads, dim = attn.num_attention_heads, attn.attention_head_dim
     local_rows = hidden_states.shape[1]
+    chunks = getattr(attn, "_vsa_fp4_options", {}).get("sp_head_chunks", 1)
+    if heads % (world * chunks):
+        raise ValueError("H3 SP attention heads must divide evenly across ranks and head chunks")
+    if (chunks > 1 and use_fused_rope and hidden_states.is_cuda and hidden_states.dtype == torch.bfloat16
+            and attn.to_gate_compress is not None and attn._gate_active()):
+        from fastvideo.models.dits.minimax_h3_fusions.sp_head_pipeline import attention_pipelined
+        return attention_pipelined(attn, hidden_states, rotary_emb, meta, sp_group, chunks, pre_quantized=pre_quantized)
     layout = getattr(meta, "_h3_fp4_sp_layout", None)
     if layout is None:
         layout = _SPTileLayout(meta, rank, local_rows)
         meta._h3_fp4_sp_layout = layout  # type: ignore[attr-defined]
 
     with STAGES.span("qkv_proj_rope"):
-        query, key, value = (t.unflatten(-1, (heads, dim))
-                             for t in _shared_input_projections((attn.to_q, attn.to_k, attn.to_v), hidden_states))
+        shared = _shared_qkv_gate_projections(attn, hidden_states, pre_quantized=pre_quantized)
+        if shared is None:
+            projected = _shared_input_projections((attn.to_q, attn.to_k, attn.to_v), hidden_states)
+            gate_pre = None
+        else:
+            projected, gate_pre = shared
+        query, key, value = (t.unflatten(-1, (heads, dim)) for t in projected)
+        del projected, shared
+        fused_transport = (use_fused_rope and query.is_cuda and query.dtype == torch.bfloat16
+                           and all(t.is_contiguous() for t in (query, key, value)))
         if use_fused_rope:
-            from fastvideo.models.dits.minimax_h3_fusions import fused_qknorm_rope
             cos, sin = rotary_emb[0].to(query.dtype), rotary_emb[1].to(query.dtype)
-            query = fused_qknorm_rope(query, attn.norm_q.weight, cos, sin, attn.norm_q.eps)
-            key = fused_qknorm_rope(key, attn.norm_k.weight, cos, sin, attn.norm_k.eps)
+            if not fused_transport:
+                from fastvideo.models.dits.minimax_h3_fusions import fused_qknorm_rope
+                query = fused_qknorm_rope(query, attn.norm_q.weight, cos, sin, attn.norm_q.eps)
+                key = fused_qknorm_rope(key, attn.norm_k.weight, cos, sin, attn.norm_k.eps)
         else:
             query = attn._apply_rotary_emb(attn.norm_q(query), rotary_emb)
             key = attn._apply_rotary_emb(attn.norm_k(key), rotary_emb)
 
+    gate = None
+
+    def project_gate() -> None:
+        nonlocal gate
+        with STAGES.span("gate_proj"):
+            if gate_pre is None:
+                gate, _ = attn.to_gate_compress(hidden_states)
+            else:
+                linear = attn.to_gate_compress
+                gate = linear.quant_method.apply(linear, hidden_states, linear.bias, pre_quantized=gate_pre)
+
+    gate_active = attn.to_gate_compress is not None and attn._gate_active()
     with STAGES.span("qkv_pack"):
-        payload, scale = _pack_heads_fp8(query[0], key[0], value[0], world)
+        if fused_transport:
+            from fastvideo.models.dits.minimax_h3_fusions.sp_qkv_pack import pack_normalized_qkv_fp8
+            payload, scale = pack_normalized_qkv_fp8(
+                query[0], key[0], value[0], attn.norm_q.weight, attn.norm_k.weight,
+                cos, sin, attn.norm_q.eps, attn.norm_k.eps, world,
+            )
+        else:
+            payload, scale = _pack_heads_fp8(query[0], key[0], value[0], world)
     with STAGES.span("qkv_all_to_all"):
-        payload, scale = _all_to_all(payload, scale, sp_group.device_group)
+        payload, scale = _all_to_all(payload, scale, sp_group.device_group,
+                                     overlap=project_gate if gate_active else None)
     with STAGES.span("qkv_unpack_tile"):
-        qkv = layout.tiles_from(_unpack_seq_fp8(payload, scale, layout.seq_len))  # [3, R, Hs, D]
+        qkv = layout.tiles_from_fp8(payload, scale)  # [3, R, Hs, D]
     q_t, k_t, v_t = qkv[0:1], qkv[1:2], qkv[2:3]
 
     vbs = meta.variable_block_sizes
@@ -565,28 +944,43 @@ def vsa_fp4_attention_sp(attn: Any, hidden_states: torch.Tensor, rotary_emb: tup
         sparsity = 0.0 if attn._layer_idx in meta.dense_layers else meta.VSA_sparsity
         mask = _build_block_mask(scores, meta.num_prefix_tiles, sparsity, meta.exempt, meta.video_tile_spans,
                                  meta.span_sparsities)
-        q2k_idx, q2k_num, kv_valid, q2k_quad = api.vsa_tile_mask_to_fp4_blocks(mask, layout.tile, vbs)
+        q2k_idx, q2k_num, kv_valid, q2k_quad = _tile_mask_to_fp4_blocks(api, mask, layout.tile, vbs)
     with STAGES.span("fp4_attention"):
         out_bhsd = _fp4_fine_attention(api, q_t, k_t, v_t, q2k_idx, q2k_num, kv_valid, q2k_quad,
                                        logical=logical, prefix_rows=prefix_rows, precision=prefix_mode,
                                        options=getattr(attn, "_vsa_fp4_options", {}), layout=layout)
 
+    out_c, gathered = None, None
+
+    def compress_gate() -> None:
+        nonlocal out_c, gathered
+        with STAGES.span("gate_compress"):
+            v_pooled = _pool_tiles(v_t[:, :logical], vbs, layout.tile)
+            out_c = torch.matmul(torch.softmax(scores, dim=-1), v_pooled)[0].to(qkv.dtype)  # [Hs, n_tiles, D]
+            gathered = torch.empty((world, *out_c.shape), dtype=out_c.dtype, device=out_c.device)
+
     with STAGES.span("out_pack"):
         payload, scale = _pack_seq_fp8(out_bhsd, layout.untile, world, local_rows)
     with STAGES.span("out_all_to_all"):
-        payload, scale = _all_to_all(payload, scale, sp_group.device_group)
-    with STAGES.span("out_unpack"):
-        out = _unpack_heads_fp8(payload, scale)  # [rows, H, D]
-
-    if attn.to_gate_compress is not None and attn._gate_active():
-        with STAGES.span("gate_compress"):
-            v_pooled = _pool_tiles(v_t[:, :logical], vbs, layout.tile)
-            out_c = torch.matmul(torch.softmax(scores, dim=-1), v_pooled)[0].to(out.dtype)  # [Hs, n_tiles, D]
-            gathered = torch.empty((world, *out_c.shape), dtype=out_c.dtype, device=out_c.device)
+        payload, scale = _all_to_all(payload, scale, sp_group.device_group,
+                                     overlap=compress_gate if gate_active else None)
+    if gate_active:
+        assert out_c is not None and gathered is not None
         with STAGES.span("gate_all_gather"):
             dist.all_gather_into_tensor(gathered, out_c.contiguous(), group=sp_group.device_group)
         with STAGES.span("gate_apply"):
-            out_c_all = gathered.flatten(0, 1).transpose(0, 1)  # [n_tiles, H, D]
-            gate, _ = attn.to_gate_compress(hidden_states)
-            out = _apply_gate(out, out_c_all, layout.local_row_tile, gate[0].unflatten(-1, (heads, dim)))
+            from fastvideo.models.dits.minimax_h3_fusions.sp_exchange import unpack_fp8_gate
+            from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod
+
+            assert gate is not None
+            local_gate = gate[0].unflatten(-1, (heads, dim))
+            method = getattr(attn.to_out, "quant_method", None)
+            if (type(method) is NVFP4QuantizeMethod and method.wants_prequantized_input()
+                    and method.uses_unit_activation_scale(attn.to_out) and gate.shape[-1] % 64 == 0):
+                prepared = unpack_fp8_gate(payload, scale, gathered, local_gate, layout.local_row_tile, method.x_global_sf)
+                return gate, prepared
+            out = unpack_fp8_gate(payload, scale, gathered, local_gate, layout.local_row_tile)
+    else:
+        with STAGES.span("out_unpack"):
+            out = _unpack_heads_fp8(payload, scale)  # [rows, H, D]
     return out.flatten(1, 2).unsqueeze(0)

fastvideo/models/dits/qwen_image21.py

@@ -1,592 +0,0 @@
-# Copyright 2026 Qwen-Image Team, The HuggingFace Team. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Native Qwen-Image-2.1 single-stream diffusion transformer.
-
-The interleaved text/image block-causal stream has no matching FastVideo
-sequence-parallel attention primitive. Mask-free dense calls use FastVideo's
-shared attention selector, retaining PyTorch SDPA for causal text and padding
-masks. Tensor and sequence parallelism are not implemented.
-Prefix K/V can reside on the host, so decoding restores only one layer's cache
-at a time instead of keeping every layer's reference images on the GPU.
-"""
-
-from __future__ import annotations
-
-import math
-from typing import TYPE_CHECKING, Any
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from fastvideo.attention.selector import effective_attention_backend
-from fastvideo.configs.models.dits.qwen_image21 import QwenImage21Config
-from fastvideo.layers.linear import ReplicatedLinear
-from fastvideo.logger import init_logger
-from fastvideo.models.dits.base import BaseDiT
-
-if TYPE_CHECKING:
-    from fastvideo.platforms import AttentionBackendEnum
-
-logger = init_logger(__name__)
-
-_IMG_TOKENS_PER_SLOT = 4
-
-
-class QwenImage21KVCacheAllocationError(RuntimeError):
-    """CPU prefix-cache allocation failed; an inference caller may retry uncached."""
-
-
-class QwenImage21KVLayerCache:
-    """Post-RoPE prefix K/V in ``(batch, tokens, heads, head_dim)`` format."""
-
-    def __init__(self, storage_device: str | torch.device | None = None):
-        self.storage_device = torch.device(storage_device) if storage_device is not None else None
-        self.k: torch.Tensor | None = None
-        self.v: torch.Tensor | None = None
-
-    def store(self, k: torch.Tensor, v: torch.Tensor) -> None:
-        if k.shape != v.shape or k.ndim != 4:
-            raise ValueError("Prefix K/V must have matching (batch, tokens, heads, head_dim) shapes")
-        # A fresh allocation releases the larger prefill tensor after the layer finishes.
-        device = self.storage_device or k.device
-        try:
-            self.k = k.detach().to(device=device, copy=True)
-            self.v = v.detach().to(device=device, copy=True)
-        except (MemoryError, RuntimeError) as exc:
-            self.clear()
-            message = str(exc).lower()
-            allocation_error = isinstance(exc, MemoryError) or any(
-                marker in message for marker in (
-                    "defaultcpuallocator", "cannot allocate memory", "can't allocate memory",
-                    "not enough memory", "out of memory", "std::bad_alloc"
-                ))
-            if device.type == "cpu" and allocation_error and "cuda" not in message and "hip" not in message:
-                raise QwenImage21KVCacheAllocationError(
-                    "CPU memory exhausted while allocating the prefix KV cache") from exc
-            raise
-
-    def get(self, device: torch.device | None = None, dtype: torch.dtype | None = None):
-        if self.k is None or self.v is None:
-            raise RuntimeError("Prefix KV cache has not been populated; run an extract step first")
-        return self.k.to(device=device, dtype=dtype), self.v.to(device=device, dtype=dtype)
-
-    def clear(self) -> None:
-        self.k = self.v = None
-
-
-class QwenImage21KVCache:
-    """Per-request cache container; never share it between prompts or CFG branches."""
-
-    def __init__(self, num_layers: int, storage_device: str | torch.device | None = None):
-        if num_layers <= 0:
-            raise ValueError("num_layers must be positive")
-        self.layer_caches = [QwenImage21KVLayerCache(storage_device) for _ in range(num_layers)]
-        self.prefix_len: int | None = None
-
-    def get_layer(self, layer_idx: int) -> QwenImage21KVLayerCache:
-        return self.layer_caches[layer_idx]
-
-    def clear(self) -> None:
-        for layer in self.layer_caches:
-            layer.clear()
-        self.prefix_len = None
-
-
-def apply_rotary_emb_qwen(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
-    rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
-    return torch.view_as_real(rotated * freqs_cis.unsqueeze(1)).flatten(3).to(x.dtype)
-
-
-class QwenImage21TemporalTimesteps(nn.Module):
-
-    def __init__(self, timestep_dim: int = 256, max_period: int = 10000, time_factor: float = 1000.0):
-        super().__init__()
-        self.timestep_dim = timestep_dim
-        self.time_factor = time_factor
-        half = timestep_dim // 2
-        freqs = torch.exp(-math.log(max_period) * torch.arange(half, dtype=torch.float32, device="cpu") / half)
-        self.register_buffer("freqs", freqs, persistent=False)
-
-    def forward(self, timestep: torch.Tensor) -> torch.Tensor:
-        args = (self.time_factor * timestep.float())[:, None] * self.freqs.float()[None].to(timestep.device)
-        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
-        if self.timestep_dim % 2:
-            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
-        return embedding.to(timestep.dtype)
-
-
-class QwenImage21TimestepEmbedding(nn.Module):
-
-    def __init__(self, embedding_dim: int):
-        super().__init__()
-        self.linear_1 = ReplicatedLinear(256, embedding_dim, bias=False)
-        self.act = nn.SiLU()
-        self.linear_2 = ReplicatedLinear(embedding_dim, embedding_dim, bias=False)
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        hidden_states, _ = self.linear_1(hidden_states)
-        return self.linear_2(self.act(hidden_states))[0]
-
-
-class QwenImage21TimestepProjEmbeddings(nn.Module):
-
-    def __init__(self, embedding_dim: int):
-        super().__init__()
-        self.time_proj = QwenImage21TemporalTimesteps()
-        self.timestep_embedder = QwenImage21TimestepEmbedding(embedding_dim)
-
-    def forward(self, timestep: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor:
-        return self.timestep_embedder(self.time_proj(timestep).to(hidden_states.dtype))
-
-
-class QwenImage21ZeroCenterRMSNorm(nn.Module):
-
-    def __init__(self, dim: int, eps: float = 1e-6):
-        super().__init__()
-        self.weight = nn.Parameter(torch.zeros(dim))
-        self.eps = eps
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        input_dtype = hidden_states.dtype
-        hidden_states = hidden_states.float()
-        rrms = torch.rsqrt(torch.mean(hidden_states**2, dim=-1, keepdim=True) + self.eps)
-        return (hidden_states * rrms * (self.weight.float() + 1)).to(input_dtype)
-
-
-class QwenImage21RMSNorm(nn.Module):
-
-    def __init__(self, dim: int, eps: float):
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(dim))
-        self.eps = eps
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        variance = hidden_states.float().pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
-        if self.weight.dtype in (torch.float16, torch.bfloat16):
-            hidden_states = hidden_states.to(self.weight.dtype)
-        return hidden_states * self.weight
-
-
-class QwenImage21TextProjection(nn.Module):
-
-    def __init__(self, context_in_dim: int, hidden_size: int, eps: float):
-        super().__init__()
-        self.text_norm = QwenImage21ZeroCenterRMSNorm(context_in_dim, eps)
-        self.in_layer = ReplicatedLinear(context_in_dim, hidden_size, bias=False)
-        self.act = nn.GELU(approximate="tanh")
-        self.out_layer = ReplicatedLinear(hidden_size, hidden_size, bias=False)
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        hidden_states, _ = self.in_layer(self.text_norm(hidden_states))
-        return self.out_layer(self.act(hidden_states))[0]
-
-
-class QwenImage21SwiGLUFeedForward(nn.Module):
-
-    def __init__(self, hidden_size: int, mlp_hidden_size: int):
-        super().__init__()
-        self.proj = ReplicatedLinear(hidden_size, mlp_hidden_size, bias=False)
-        self.out = ReplicatedLinear(mlp_hidden_size, hidden_size, bias=False)
-        self.gate_layer = ReplicatedLinear(hidden_size, mlp_hidden_size, bias=False)
-        self.activation_fn = nn.SiLU()
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        gate = self.activation_fn(self.gate_layer(hidden_states)[0])
-        return self.out(gate * self.proj(hidden_states)[0])[0]
-
-
-def _select_modulation_rows(params: torch.Tensor, target_token_mask: torch.Tensor | None) -> torch.Tensor:
-    if target_token_mask is None:
-        return params.unsqueeze(1)
-    real, zero = params[:-1].unsqueeze(1), params[-1:].unsqueeze(0)
-    return torch.where(target_token_mask.view(1, -1, 1), real, zero)
-
-
-class QwenImage21AdaLayerNormContinuous(nn.Module):
-
-    def __init__(self, embedding_dim: int, eps: float):
-        super().__init__()
-        self.silu = nn.SiLU()
-        self.linear = ReplicatedLinear(embedding_dim, embedding_dim, bias=False)
-        self.norm = nn.LayerNorm(embedding_dim, eps=eps, elementwise_affine=False)
-
-    def forward(self, hidden_states, conditioning_embedding, target_token_mask=None):
-        scale = self.linear(self.silu(conditioning_embedding).to(hidden_states.dtype))[0]
-        return self.norm(hidden_states) * (1 + _select_modulation_rows(scale, target_token_mask))
-
-
-def _qwenimage21_prefix_segments(image_ids: torch.Tensor, prefix_len: int) -> list[tuple[int, int, bool]]:
-    prefix_ids = image_ids[:prefix_len].tolist()
-    segments = []
-    start = 0
-    for index in range(1, prefix_len + 1):
-        if index == prefix_len or prefix_ids[index] != prefix_ids[start]:
-            segments.append((start, index, prefix_ids[start] < 0))
-            start = index
-    return segments
-
-
-def _attention(query, key, value, mask=None):
-    # SDPA expects (batch, heads, tokens, dim); projected QKV use tokens before heads.
-    output = F.scaled_dot_product_attention(query.transpose(1, 2),
-                                            key.transpose(1, 2),
-                                            value.transpose(1, 2),
-                                            attn_mask=mask,
-                                            dropout_p=0.0)
-    if mask is not None:
-        # Older CPU SDPA returns NaNs for fully masked padding queries.
-        output = output.masked_fill(~mask.any(dim=-1, keepdim=True), 0.0)
-    return output.transpose(1, 2)
-
-
-def _make_dense_attention(heads: int, dim_head: int,
-                          supported_attention_backends: tuple[AttentionBackendEnum, ...],
-                          requested_backend: AttentionBackendEnum | None) -> nn.Module | None:
-    # Keep optional attention kernels and GPU dispatch imports off the SDPA path.
-    from fastvideo.attention.layer import LocalAttention
-    from fastvideo.attention.selector import backend_name_to_enum, get_attn_backend
-    from fastvideo.platforms import AttentionBackendEnum
-    from fastvideo.utils import get_compute_dtype
-
-    dtype = get_compute_dtype()
-    selection_kwargs = {} if requested_backend is None else {"requested": requested_backend}
-    backend_cls = get_attn_backend(dim_head, dtype,
-                                   supported_attention_backends=supported_attention_backends,
-                                   **selection_kwargs)
-    backend = backend_name_to_enum(backend_cls.get_name())
-    if backend is AttentionBackendEnum.TORCH_SDPA:
-        return None
-
-    class DenseLocalAttention(LocalAttention):
-        """Qwen owns segment lengths and has already applied its complex RoPE.
-
-        Resolve explicit component requests through the public selector above;
-        direct construction without a request retains scope/env selection.
-        LocalAttention's ordinary forward reads unrelated global metadata,
-        which must not reach our unequal Q/KV calls.
-        """
-
-        def __init__(self):
-            nn.Module.__init__(self)
-            self.softmax_scale = dim_head**-0.5
-            self.attn_impl = backend_cls.get_impl_cls()(num_heads=heads, head_size=dim_head,
-                                                        softmax_scale=self.softmax_scale,
-                                                        num_kv_heads=heads, causal=False)
-            self.num_heads = self.num_kv_heads = heads
-            self.head_size, self.backend, self.dtype = dim_head, backend, dtype
-
-        def forward(self, q, k, v):
-            return self.attn_impl.forward(q, k, v, None)
-
-    return DenseLocalAttention()
-
-
-class QwenImage21Attention(nn.Module):
-
-    def __init__(self, dim: int, heads: int, dim_head: int, eps: float, *,
-                 requested_backend: AttentionBackendEnum | None = None,
-                 supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None):
-        super().__init__()
-        self.heads = heads
-        self.inner_dim = heads * dim_head
-        self.to_q = ReplicatedLinear(dim, self.inner_dim, bias=False)
-        self.to_k = ReplicatedLinear(dim, self.inner_dim, bias=False)
-        self.to_v = ReplicatedLinear(dim, self.inner_dim, bias=False)
-        self.to_out = nn.ModuleList([ReplicatedLinear(self.inner_dim, dim, bias=False), nn.Dropout(0.0)])
-        self.norm_q = QwenImage21RMSNorm(dim_head, eps)
-        self.norm_k = QwenImage21RMSNorm(dim_head, eps)
-        self.dense_attn = None
-        if requested_backend is None or requested_backend.name != "TORCH_SDPA":
-            supported = supported_attention_backends or _DEFAULT_CONFIG.arch_config._supported_attention_backends
-            self.dense_attn = _make_dense_attention(heads, dim_head, supported, requested_backend)
-
-    def _attend(self, query, key, value, mask=None):
-        if mask is not None or self.dense_attn is None:
-            return _attention(query, key, value, mask)
-        return self.dense_attn(query, key, value)
-
-    def forward(self,
-                hidden_states: torch.Tensor,
-                rotary_emb: torch.Tensor | None = None,
-                attention_mask: torch.Tensor | None = None,
-                layer_cache: QwenImage21KVLayerCache | None = None,
-                kv_cache_mode: str | None = None,
-                cache_write_slice: slice | None = None,
-                segments: list[tuple[int, int, bool]] | None = None,
-                key_valid: torch.Tensor | None = None) -> torch.Tensor:
-        query = self.to_q(hidden_states)[0].unflatten(-1, (self.heads, -1))
-        key = self.to_k(hidden_states)[0].unflatten(-1, (self.heads, -1))
-        value = self.to_v(hidden_states)[0].unflatten(-1, (self.heads, -1))
-        query, key = self.norm_q(query).to(value.dtype), self.norm_k(key).to(value.dtype)
-        if rotary_emb is not None:
-            query = apply_rotary_emb_qwen(query, rotary_emb)
-            key = apply_rotary_emb_qwen(key, rotary_emb)
-        if layer_cache is not None:
-            if kv_cache_mode == "extract" and cache_write_slice is not None:
-                layer_cache.store(key[:, cache_write_slice], value[:, cache_write_slice])
-            elif kv_cache_mode == "cached":
-                cached_k, cached_v = layer_cache.get(key.device, key.dtype)
-                if cached_k.shape[0] != key.shape[0] or cached_k.shape[2:] != key.shape[2:]:
-                    raise ValueError("Cached prefix batch/head shape differs from the current request")
-                key = torch.cat([cached_k, key], dim=1)
-                value = torch.cat([cached_v, value], dim=1)
-                del cached_k, cached_v
-        if segments is None:
-            output = self._attend(query, key, value, attention_mask)
-        else:
-            outputs = []
-            for start, end, is_text in segments:
-                mask = None
-                if is_text:
-                    length = end - start
-                    mask = torch.cat([torch.ones(length, start, dtype=torch.bool, device=query.device),
-                                      torch.tril(torch.ones(length, length, dtype=torch.bool, device=query.device))],
-                                     dim=1)[None, None]
-                if key_valid is not None:
-                    valid = key_valid[:, None, None, :end]
-                    mask = valid if mask is None else mask & valid
-                outputs.append(self._attend(query[:, start:end], key[:, :end], value[:, :end], mask))
-            prefix_len = segments[-1][1] if segments else 0
-            mask = None if key_valid is None else key_valid[:, None, None, :]
-            outputs.append(self._attend(query[:, prefix_len:], key, value, mask))
-            output = torch.cat(outputs, dim=1)
-        output = output.flatten(2, 3).to(query.dtype)
-        return self.to_out[1](self.to_out[0](output)[0])
-
-
-class QwenImage21TransformerBlock(nn.Module):
-
-    def __init__(self, dim: int, num_attention_heads: int, attention_head_dim: int, mlp_ratio: int, eps: float, *,
-                 requested_backend: AttentionBackendEnum | None = None,
-                 supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None):
-        super().__init__()
-        self.img_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
-        self.attn = QwenImage21Attention(dim, num_attention_heads, attention_head_dim, eps,
-                                       requested_backend=requested_backend,
-                                       supported_attention_backends=supported_attention_backends)
-        self.img_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
-        self.img_mlp = QwenImage21SwiGLUFeedForward(dim, dim * mlp_ratio)
-
-    def _modulate(self, hidden_states, mod_params, target_token_mask):
-        scale, gate = mod_params.chunk(2, dim=-1)
-        scale = _select_modulation_rows(scale, target_token_mask)
-        gate = _select_modulation_rows(gate, target_token_mask)
-        return hidden_states * (1 + scale), gate
-
-    def forward(self, hidden_states, modulation, target_token_mask=None, **attention_kwargs):
-        mod1, mod2 = modulation.chunk(2, dim=-1)
-        normalized, gate = self._modulate(self.img_norm1(hidden_states), mod1, target_token_mask)
-        hidden_states = hidden_states + gate.tanh() * self.attn(normalized, **attention_kwargs)
-        normalized, gate = self._modulate(self.img_norm2(hidden_states), mod2, target_token_mask)
-        hidden_states = hidden_states + gate.tanh() * self.img_mlp(normalized)
-        if hidden_states.dtype == torch.float16:
-            hidden_states = hidden_states.clip(-65504, 65504)
-        return hidden_states
-
-
-class QwenImage21Rope(nn.Module):
-
-    def __init__(self, theta: int, axes_dim: tuple[int, int, int]):
-        super().__init__()
-        self.theta = theta
-        self.axes_dim = axes_dim
-
-    def rope_params(self, index: torch.Tensor, dim: int) -> torch.Tensor:
-        powers = torch.arange(0, dim, 2, dtype=torch.float32, device=index.device).div(dim)
-        freqs = torch.outer(index.float(), 1.0 / torch.pow(self.theta, powers))
-        return torch.polar(torch.ones_like(freqs), freqs)
-
-    def forward(self, img_shapes, image_pad_mask, device):
-        frame_index, image_height_index, image_width_index = [], [], []
-        cursor, position = 0, 0
-        is_image = image_pad_mask.tolist()
-        for _, height, width in img_shapes:
-            block_start = is_image.index(True, cursor)
-            text_len = block_start - cursor
-            frame_index.extend(range(position, position + text_len))
-            position += text_len
-            cursor = block_start + height * width
-            frame_index.extend([position] * (height * width))
-            position += max(height, width)
-            image_height_index.extend([h for h in range(-(height - height // 2), height // 2) for _ in range(width)])
-            image_width_index.extend([w for _ in range(height) for w in range(-(width - width // 2), width // 2)])
-        if cursor < len(is_image):
-            frame_index.extend(range(position, position + len(is_image) - cursor))
-        # Derive frequencies from signed coordinates, retaining Qwen's CPU fp32/complex64 numerics.
-        frame_index = torch.tensor(frame_index, dtype=torch.long, device="cpu")
-        image_pad_mask = torch.tensor(is_image, dtype=torch.bool, device="cpu")
-        height_index, width_index = frame_index.clone(), frame_index.clone()
-        height_index[image_pad_mask] = torch.tensor(image_height_index, dtype=torch.long, device="cpu")
-        width_index[image_pad_mask] = torch.tensor(image_width_index, dtype=torch.long, device="cpu")
-        return torch.cat([self.rope_params(index, dim) for index, dim in
-                          zip((frame_index, height_index, width_index), self.axes_dim, strict=True)], dim=-1).to(device)
-
-
-_DEFAULT_CONFIG = QwenImage21Config()
-
-
-class QwenImage21Transformer2DModel(BaseDiT):
-    _fsdp_shard_conditions = _DEFAULT_CONFIG.arch_config._fsdp_shard_conditions
-    _compile_conditions = _DEFAULT_CONFIG.arch_config._compile_conditions
-    _supported_attention_backends = _DEFAULT_CONFIG.arch_config._supported_attention_backends
-    param_names_mapping = _DEFAULT_CONFIG.arch_config.param_names_mapping
-    reverse_param_names_mapping = _DEFAULT_CONFIG.arch_config.reverse_param_names_mapping
-
-    def __init__(self, config: QwenImage21Config, hf_config: dict[str, Any] | None = None, **kwargs):
-        super().__init__(config, hf_config or {}, **kwargs)
-        if getattr(config, "quant_config", None) is not None:
-            raise NotImplementedError("Qwen-Image-2.1 transformer quantization has not been implemented")
-        self.inner_dim = self.hidden_size = config.hidden_size
-        self.out_channels = config.out_channels
-        self.num_attention_heads = config.num_attention_heads
-        self.num_channels_latents = config.num_channels_latents
-        self.pos_embed = QwenImage21Rope(10000, config.axes_dims_rope)
-        self.time_text_embed = QwenImage21TimestepProjEmbeddings(self.inner_dim)
-        self.txt_in = QwenImage21TextProjection(config.context_in_dim, self.inner_dim, config.eps)
-        self.img_in = ReplicatedLinear(config.in_channels * config.patch_size**2, self.inner_dim, bias=False)
-        self.modulation = nn.Sequential(nn.SiLU(), ReplicatedLinear(self.inner_dim, 4 * self.inner_dim, bias=False))
-        requested_backend = effective_attention_backend(config)
-        self.transformer_blocks = nn.ModuleList([
-            QwenImage21TransformerBlock(self.inner_dim, config.num_attention_heads, config.attention_head_dim,
-                                        config.mlp_ratio, config.eps, requested_backend=requested_backend,
-                                        supported_attention_backends=self._supported_attention_backends)
-            for _ in range(config.num_layers)
-        ])
-        self.norm_out = QwenImage21AdaLayerNormContinuous(self.inner_dim, config.eps)
-        self.proj_out = ReplicatedLinear(self.inner_dim, config.patch_size**2 * config.out_channels, bias=False)
-
-    @staticmethod
-    def build_token_metadata(image_pad_mask, img_shapes):
-        positions = image_pad_mask.nonzero(as_tuple=True)[0]
-        lengths = [math.prod(shape) for shape in img_shapes]
-        if not lengths or any(length <= 0 for length in lengths) or sum(lengths) != positions.numel():
-            raise ValueError("img_shapes token counts must match image_pad_mask")
-        image_ids = torch.full_like(image_pad_mask, -1, dtype=torch.long)
-        ids = torch.repeat_interleave(torch.arange(len(lengths), device=image_pad_mask.device),
-                                     torch.tensor(lengths, device=image_pad_mask.device))
-        image_ids[positions] = ids
-        target_mask = torch.zeros_like(image_pad_mask)
-        target_mask[positions[-lengths[-1]:]] = True
-        return image_ids, target_mask
-
-    def _validate_inputs(self, hidden_states, encoder_hidden_states, timestep, img_shapes, img_mask, text_mask):
-        batch_size = hidden_states.shape[0]
-        if hidden_states.ndim != 3 or encoder_hidden_states.ndim != 3:
-            raise ValueError("Latents and text embeddings must have (batch, tokens, channels) shapes")
-        if encoder_hidden_states.shape[0] != batch_size:
-            raise ValueError("Latent and text batch sizes must match")
-        if len(img_shapes) != batch_size or not img_shapes or not img_shapes[0]:
-            raise ValueError("img_shapes must provide a nonempty image layout for each sample")
-        if any(shapes != img_shapes[0] for shapes in img_shapes):
-            raise ValueError("All samples must share the same image layout")
-        if any(frame != 1 or height <= 0 or width <= 0 for frame, height, width in img_shapes[0]):
-            raise ValueError("Qwen-Image-2.1 expects positive single-frame image shapes")
-        target_tokens = math.prod(img_shapes[0][-1])
-        if target_tokens % _IMG_TOKENS_PER_SLOT:
-            raise ValueError("The target latent token count must be divisible by four")
-        if hidden_states.shape[1] != sum(math.prod(shape) for shape in img_shapes[0]):
-            raise ValueError("Packed latents must include all references followed by the target image")
-        expected_slots = encoder_hidden_states.shape[1] + target_tokens // _IMG_TOKENS_PER_SLOT
-        if img_mask.shape != (batch_size, expected_slots):
-            raise ValueError("img_mask must include the VLM sequence and appended target-image slots")
-        if not torch.equal(img_mask.bool(), img_mask[0:1].bool().expand_as(img_mask)):
-            raise ValueError("All samples must share the same image-slot layout")
-        if not bool(img_mask[:, -target_tokens // _IMG_TOKENS_PER_SLOT:].all()):
-            raise ValueError("Target-image slots must be appended at the end of img_mask")
-        if int(img_mask[0].count_nonzero()) * _IMG_TOKENS_PER_SLOT != hidden_states.shape[1]:
-            raise ValueError("Image-slot count must match the reference and target latent token count")
-        if timestep.numel() != batch_size:
-            raise ValueError("timestep must contain one value per sample")
-        if text_mask is not None and text_mask.shape != encoder_hidden_states.shape[:2]:
-            raise ValueError("encoder_hidden_states_mask must match the VLM embedding sequence")
-
-    def forward(self,
-                hidden_states: torch.Tensor,
-                encoder_hidden_states: torch.Tensor,
-                timestep: torch.Tensor,
-                img_shapes: list[list[tuple[int, int, int]]],
-                img_mask: torch.Tensor,
-                encoder_hidden_states_mask: torch.Tensor | None = None,
-                attention_kwargs: dict[str, Any] | None = None,
-                kv_cache: QwenImage21KVCache | None = None,
-                kv_cache_mode: str | None = None,
-                return_dict: bool | None = None,
-                **kwargs) -> torch.Tensor | tuple[torch.Tensor]:
-        """Return joint-sequence predictions; the pipeline selects the target tail.
-
-        ``kv_cache_mode='extract'`` stores timestep-independent text/reference
-        K/V. ``'cached'`` returns only the target image tokens. By default the
-        return value is a tensor; ``return_dict=False`` selects a one-item tuple.
-        """
-        self._validate_inputs(hidden_states, encoder_hidden_states, timestep, img_shapes, img_mask,
-                              encoder_hidden_states_mask)
-        if kv_cache is not None and not self.config.causal_condition:
-            raise ValueError("Prefix KV caching requires causal_condition=True")
-        if kv_cache is not None and kv_cache_mode not in ("extract", "cached"):
-            raise ValueError("kv_cache_mode must be extract or cached when a cache is provided")
-        if kv_cache is None and kv_cache_mode is not None:
-            raise ValueError("kv_cache_mode requires a cache container")
-        if kv_cache is not None and len(kv_cache.layer_caches) != len(self.transformer_blocks):
-            raise ValueError("KV cache layer count differs from the transformer")
-        if attention_kwargs:
-            raise ValueError("Custom attention processors and LoRA scaling are not implemented for Qwen-Image-2.1")
-        batch_size = hidden_states.shape[0]
-        hidden_states = self.img_in(hidden_states)[0]
-        encoder_hidden_states = self.txt_in(encoder_hidden_states)
-        img_mask = img_mask.bool().to(hidden_states.device)
-        repeats = torch.where(img_mask, _IMG_TOKENS_PER_SLOT, 1)[0]
-        image_pad_mask = torch.repeat_interleave(img_mask[0], repeats)
-        target_tokens = math.prod(img_shapes[0][-1])
-        joint = torch.cat([encoder_hidden_states,
-                           encoder_hidden_states.new_zeros(batch_size, target_tokens // 4, self.inner_dim)], dim=1)
-        joint = joint.repeat_interleave(repeats, dim=1)
-        joint[:, image_pad_mask] = hidden_states
-        rotary_emb = self.pos_embed(img_shapes[0], image_pad_mask, hidden_states.device)
-        image_ids, target_mask = self.build_token_metadata(image_pad_mask, img_shapes[0])
-        timestep = timestep.reshape(batch_size).to(device=hidden_states.device, dtype=hidden_states.dtype)
-        if self.config.causal_condition:
-            timestep = torch.cat([timestep, timestep.new_zeros(1)])
-            modulation_mask = target_mask
-        else:
-            modulation_mask = None
-        temb = self.time_text_embed(timestep, hidden_states)
-        modulation = self.modulation[1](self.modulation[0](temb))[0]
-        joint_key_valid = None
-        if encoder_hidden_states_mask is not None:
-            joint_key_valid = torch.ones(batch_size, image_pad_mask.numel(), dtype=torch.bool,
-                                         device=hidden_states.device)
-            text_positions = (~image_pad_mask).nonzero(as_tuple=True)[0]
-            vlm_text_positions = ~img_mask[0, :encoder_hidden_states.shape[1]]
-            valid_text = encoder_hidden_states_mask.bool().to(hidden_states.device)[:, vlm_text_positions]
-            joint_key_valid[:, text_positions] = valid_text
-        prefix_len = image_pad_mask.numel() - target_tokens
-        if kv_cache_mode == "cached":
-            if kv_cache.prefix_len != prefix_len:
-                raise ValueError("Cached prefix layout differs from the current request or has not been extracted")
-            joint, rotary_emb = joint[:, prefix_len:], rotary_emb[prefix_len:]
-            modulation_mask = modulation_mask[prefix_len:]
-            attention_mask = None if joint_key_valid is None else joint_key_valid[:, None, None, :]
-            segments, block_key_valid, cache_slice = None, None, None
-        else:
-            segments = _qwenimage21_prefix_segments(image_ids, prefix_len)
-            attention_mask, block_key_valid = None, joint_key_valid
-            cache_slice = slice(0, prefix_len) if kv_cache_mode == "extract" else None
-            if kv_cache is not None:
-                kv_cache.prefix_len = prefix_len
-        for index, block in enumerate(self.transformer_blocks):
-            joint = block(joint, modulation,
-                          target_token_mask=modulation_mask,
-                          rotary_emb=rotary_emb,
-                          attention_mask=attention_mask,
-                          layer_cache=kv_cache.get_layer(index) if kv_cache is not None else None,
-                          kv_cache_mode=kv_cache_mode,
-                          cache_write_slice=cache_slice,
-                          segments=segments,
-                          key_valid=block_key_valid)
-        output = self.proj_out(self.norm_out(joint, temb, modulation_mask))[0]
-        return (output,) if return_dict is False else output
-
-
-EntryClass = QwenImage21Transformer2DModel

fastvideo/models/encoders/qwen_image21.py

@@ -1,154 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Native Qwen3-VL conditioner returning Qwen-Image-2.1's pre-norm features."""
-
-from collections.abc import Iterable
-
-import torch
-
-from fastvideo.configs.models.encoders.qwen_image21 import QwenImage21Qwen3VLConfig
-from fastvideo.models.encoders.minimax_h3_qwen3_vl import MiniMaxH3Qwen3VLConditioner
-from fastvideo.models.loader.weight_utils import default_weight_loader
-
-
-class QwenImage21Qwen3VLConditioner(MiniMaxH3Qwen3VLConditioner):
-    """Qwen3-VL text and image features, without autoregressive generation.
-
-    The shared native stack stops immediately after the configured last
-    decoder layer. Its final RMSNorm is loaded for checkpoint completeness
-    but never applied to the image transformer's conditioning features.
-    Sequence parallel execution is not implemented for this conditioner.
-    """
-
-    supported_checkpoint_quantization_methods: frozenset[str] = frozenset()
-    _checkpoint_quantization_configs: dict = {}
-
-    def __init__(self, config: QwenImage21Qwen3VLConfig) -> None:
-        super().__init__(config)
-        # CPU construction of this constant avoids device-dependent power
-        # rounding accumulating across the complete 36-layer decoder stack.
-        rotary = self.language_model.rotary_emb
-        exponents = torch.arange(0, config.head_dim, 2, dtype=torch.float32, device="cpu") / config.head_dim
-        rotary.inv_freq = (1.0 / (config.rope_theta**exponents)).to(rotary.inv_freq.device)
-
-    def _visual_features(
-        self,
-        pixels: torch.Tensor,
-        grid_thw: torch.Tensor,
-    ) -> tuple[torch.Tensor, list[torch.Tensor]]:
-        if grid_thw.ndim != 2 or grid_thw.shape[1] != 3 or grid_thw.shape[0] == 0:
-            raise ValueError("image_grid_thw must contain one nonempty [time, height, width] row per image")
-        grids = grid_thw.to(device="cpu", dtype=torch.long)
-        merge = self.config.vision_spatial_merge_size
-        if bool((grids <= 0).any()) or bool((grids[:, 1:] % merge != 0).any()):
-            raise ValueError("Image grids must be positive and spatially divisible by the vision merge size")
-        patch_counts = grids.prod(dim=1).tolist()
-        if sum(patch_counts) != pixels.shape[0]:
-            raise ValueError("pixel_values patch count does not match image_grid_thw")
-
-        # Keep packed projection shapes: splitting images changes BF16 GEMM
-        # rounding before the complete text decoder amplifies the difference.
-        # Vision attention still executes separately for each image grid.
-        return self.visual(pixels.to(self.visual.patch_embed.proj.weight.dtype), grid_thw)
-
-    @torch.no_grad()
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        *,
-        attention_mask: torch.Tensor | None = None,
-        pixel_values: torch.Tensor | None = None,
-        image_grid_thw: torch.Tensor | None = None,
-        mm_token_type_ids: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        """Return ``[batch, tokens, width]``, or ``[tokens, width]`` for 1-D IDs.
-
-        Padding is zeroed and removed by the prompt stage using the original
-        attention mask before stripping its template.
-        """
-        single_sequence = input_ids.ndim == 1
-        if single_sequence:
-            input_ids = input_ids.unsqueeze(0)
-            if attention_mask is not None and attention_mask.ndim == 1:
-                attention_mask = attention_mask.unsqueeze(0)
-            if mm_token_type_ids is not None and mm_token_type_ids.ndim == 1:
-                mm_token_type_ids = mm_token_type_ids.unsqueeze(0)
-        if input_ids.ndim != 2 or input_ids.shape[1] == 0:
-            raise ValueError("Qwen-Image-2.1 input_ids must have shape [batch, nonempty sequence] or [sequence]")
-        if attention_mask is not None:
-            if attention_mask.shape != input_ids.shape:
-                raise ValueError("attention_mask must have the same shape as input_ids")
-            if bool((attention_mask.to(torch.bool).sum(dim=1) == 0).any()):
-                raise ValueError("Each conditioning sequence must contain an unmasked token")
-        if (pixel_values is None) != (image_grid_thw is None):
-            raise ValueError("pixel_values and image_grid_thw must be provided together")
-        if bool((input_ids == self.config.video_token_id).any()):
-            raise ValueError("Qwen-Image-2.1 conditioning accepts image references, not video tokens")
-        if mm_token_type_ids is not None:
-            expected = (input_ids == self.config.image_token_id).to(mm_token_type_ids.dtype)
-            if mm_token_type_ids.shape != input_ids.shape or not torch.equal(mm_token_type_ids, expected):
-                raise ValueError("mm_token_type_ids must mark image placeholders as 1 and other tokens as 0")
-
-        inputs_embeds = self.language_model.embed_tokens(input_ids)
-        visual_mask = None
-        deepstack = None
-        if pixel_values is not None:
-            features, deepstack = self._visual_features(pixel_values, image_grid_thw)
-            features = features.to(inputs_embeds.device, inputs_embeds.dtype)
-            visual_mask = self._placeholder_mask(input_ids, inputs_embeds, self.config.image_token_id,
-                                                features, "image")
-            inputs_embeds = inputs_embeds.masked_scatter(visual_mask.unsqueeze(-1), features)
-        elif bool((input_ids == self.config.image_token_id).any()):
-            raise ValueError("Image placeholder tokens require pixel_values and image_grid_thw")
-
-        # Text-only prefills retain positions in the padded sequence. A
-        # constant RoPE shift is mathematically equivalent, but changes BF16
-        # rounding; image prefills need each sample's multimodal layout.
-        positions = self._get_rope_index(input_ids, image_grid_thw, None,
-                                         attention_mask if image_grid_thw is not None else None)
-        # Unpadding changes projection and attention shapes in reduced
-        # precision. Preserve the batch through every decoder layer instead.
-        hidden_states = self.language_model(inputs_embeds, positions, attention_mask, visual_mask, deepstack)
-        if attention_mask is not None:
-            hidden_states = hidden_states.masked_fill(~attention_mask.to(torch.bool).unsqueeze(-1), 0)
-        return hidden_states[0] if single_sequence else hidden_states
-
-    @torch.no_grad()
-    def encode_ids(
-        self,
-        input_ids: torch.Tensor,
-        *,
-        attention_mask: torch.Tensor | None = None,
-        pixel_values: torch.Tensor | None = None,
-        image_grid_thw: torch.Tensor | None = None,
-        mm_token_type_ids: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        """Encode one sequence without imposing a tokenizer truncation limit."""
-        if input_ids.ndim != 1:
-            raise ValueError("encode_ids requires 1-D input_ids; use forward for a padded batch")
-        return self.forward(
-            input_ids,
-            attention_mask=attention_mask,
-            pixel_values=pixel_values,
-            image_grid_thw=image_grid_thw,
-            mm_token_type_ids=mm_token_type_ids,
-        )
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        parameters = dict(self.named_parameters())
-        loaded: set[str] = set()
-        for source_name, tensor in weights:
-            if source_name == "lm_head.weight":
-                continue
-            name = source_name.removeprefix("model.")
-            if name not in parameters:
-                raise ValueError(f"Unexpected Qwen-Image-2.1 Qwen3-VL checkpoint key: {source_name}")
-            if name in loaded:
-                raise ValueError(f"Duplicate Qwen-Image-2.1 Qwen3-VL checkpoint key: {source_name}")
-            parameter = parameters[name]
-            loader = getattr(parameter, "weight_loader", default_weight_loader)
-            loader(parameter, tensor)
-            loaded.add(name)
-        return loaded
-
-
-EntryClass = QwenImage21Qwen3VLConditioner

fastvideo/models/loader/component_loader.py

@@ -849,14 +849,57 @@ class VAELoader(ComponentLoader):
         # The swap runs where the VAE was built (the CPU, see ``_h3_nvfp4_decoder_path``); a pre-packed
         # file never allocates the quantized linears' full-precision weights, a bf16 file packs on the GPU.
         device = next(vae.parameters()).device
-        metadata = apply_nvfp4_decoder_checkpoint(vae, load_nvfp4_decoder_checkpoint(path, device), freeze=False)
+        checkpoint = load_nvfp4_decoder_checkpoint(path, device)
+        student_layers = int(checkpoint["metadata"].get("student_layers", loaded_blocks))
+        if student_layers > loaded_blocks:
+            # A complete decoder export can also replace a shallower light VAE.
+            # Build its topology from the host geometry; strict loading below
+            # requires every exported weight and rejects incompatible widths.
+            from fastvideo.models.vaes.minimax_h3_video import MiniMaxH3VideoViTDecoder3d
+
+            arch = vae.config
+            with torch.device(device):
+                vae.decoder = MiniMaxH3VideoViTDecoder3d(
+                    in_channels=arch.latent_channels,
+                    out_channels=arch.out_channels,
+                    patch_size=vae.spatial_compression_ratio,
+                    patch_size_t=vae.temporal_compression_ratio,
+                    num_layers=student_layers,
+                    num_attention_heads=arch.decoder_num_attention_heads,
+                    attention_head_dim=arch.decoder_attention_head_dim,
+                    num_register_tokens=arch.decoder_num_register_tokens,
+                    ffn_mult=arch.decoder_ffn_mult,
+                    rope_theta=arch.decoder_rope_theta,
+                    rope_dim_ratio=arch.decoder_rope_dim_ratio,
+                    norm_eps=arch.decoder_norm_eps,
+                ).eval()
+        metadata = apply_nvfp4_decoder_checkpoint(vae, checkpoint, freeze=False)
         vae.to(target_device)
         if target_device.type == "cuda":
             # Frozen layers hold only packed FP4 weights; packing needs the CUDA FP4 kernels.
             freeze_nvfp4_linears(vae.decoder)
-        logger.info("MiniMax-H3 VAE decoder: NVFP4 %s (%d of %d blocks, act_scale=%s), bf16 decode autocast", path,
+        logger.info("MiniMax-H3 VAE decoder: NVFP4 %s (%d blocks, original %d, act_scale=%s), bf16 decode autocast", path,
                     len(vae.decoder.transformer_blocks), loaded_blocks, metadata.get("act_scale"))
 
+    @staticmethod
+    def _quantize_h3_dense_decoder(vae: nn.Module) -> None:
+        """Quantize the original full-depth decoder in memory; leave the checkpoint untouched."""
+        from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import (
+            convert_decoder_to_nvfp4,
+            freeze_nvfp4_linears_on_device,
+        )
+
+        # Fused-state caching reads input_amax._version. Even an outer inference
+        # context must create ordinary version-tracked buffers during conversion.
+        with torch.inference_mode(False), torch.no_grad():
+            names = convert_decoder_to_nvfp4(vae.decoder, compute_dtype=torch.bfloat16, act_scale="unit")
+            freeze_nvfp4_linears_on_device(vae.decoder)
+        vae.requires_grad_(False)
+        vae.decode_autocast_dtype = torch.bfloat16
+        logger.info("MiniMax-H3 VAE decoder: runtime NVFP4 quantization of %d original block linears "
+                    "(%d blocks, unit activation scale), bf16 decode autocast", len(names),
+                    len(vae.decoder.transformer_blocks))
+
     def load(self, model_path: str, fastvideo_args: FastVideoArgs):
         """Load the VAE based on the model path, and inference args."""
         config = get_diffusers_config(model=model_path)
@@ -867,6 +910,10 @@ class VAELoader(ComponentLoader):
         assert class_name is not None, (
             "Model config does not contain a _class_name attribute. Only diffusers format is supported.")
         fastvideo_args.model_paths["vae"] = model_path
+        quantize_dense = (class_name == "AutoencoderKLMiniMaxH3"
+                          and getattr(fastvideo_args, "h3_vae_nvfp4_quantize_dense", False))
+        if quantize_dense and self._h3_nvfp4_decoder_path(class_name, fastvideo_args):
+            raise ValueError("h3_vae_nvfp4_quantize_dense and an exported NVFP4 decoder are mutually exclusive")
 
         from fastvideo.platforms import current_platform
 
@@ -993,9 +1040,9 @@ class VAELoader(ComponentLoader):
                 int8_convrot_overlay_to_apply,
             )
             int8_convrot_path = int8_convrot_overlay_to_apply(model_path)
-            if int8_convrot_path is not None and self._h3_nvfp4_decoder_path(class_name, fastvideo_args):
+            if int8_convrot_path is not None and (quantize_dense or self._h3_nvfp4_decoder_path(class_name, fastvideo_args)):
                 logger.info("Skipping MiniMax-H3 INT8 ConvRot VAE overlay %s: the NVFP4 decoder replaces the "
-                            "decoder", int8_convrot_path)
+                            "decoder execution", int8_convrot_path)
                 int8_convrot_path = None
             safetensors_list = dense_vae_safetensors(safetensors_list)
         if not safetensors_list:
@@ -1032,7 +1079,7 @@ class VAELoader(ComponentLoader):
         # Diffusers-format AutoencoderKL checkpoints should match exactly; load
         # strictly so missing/unexpected keys are surfaced early.
         # Kandinsky6SRVAE: a partially loaded KVAE decodes plausible-looking garbage, so it is strict as well.
-        strict_load = class_name in {"AutoencoderKL", "AutoencoderKLMiniMaxH3", "Kandinsky6SRVAE", "AutoencoderKLQwenImage21"}
+        strict_load = class_name in {"AutoencoderKL", "AutoencoderKLMiniMaxH3", "Kandinsky6SRVAE"}
         vae.load_state_dict(loaded, strict=strict_load)
         del loaded
         if class_name == "AutoencoderKLMiniMaxH3" and int8_convrot_path is not None:
@@ -1041,6 +1088,8 @@ class VAELoader(ComponentLoader):
         nvfp4_decoder = self._h3_nvfp4_decoder_path(class_name, fastvideo_args)
         if nvfp4_decoder:
             self._apply_h3_nvfp4_decoder(vae, nvfp4_decoder, target_device)
+        elif quantize_dense:
+            self._quantize_h3_dense_decoder(vae)
         if (class_name == "AutoencoderKLWan" and getattr(vae.config, "use_light_vae", False)
                 and target_device.type == "cuda" and hasattr(vae, "optimize_memory_format")):
             vae.optimize_memory_format()
@@ -1332,7 +1381,12 @@ class TransformerLoader(ComponentLoader):
             adaln_table = envs.FASTVIDEO_H3_ADALN_TABLE.get() or None
             adaln_host_cache = (fastvideo_args.inference_mode and not fastvideo_args.use_fsdp_inference
                                 and (envs.FASTVIDEO_H3_ADALN_CACHE.get() or adaln_table is not None))
-            layerwise_load = layerwise_load or adaln_host_cache
+            from fastvideo.platforms import current_platform
+            adaln_resident_cache = (adaln_host_cache and adaln_table is None
+                                    and not fastvideo_args.dit_layerwise_offload
+                                    and current_platform.has_unified_memory())
+            # Tables skip the projections; unified-memory caches reuse the device allocation.
+            layerwise_load = layerwise_load or (adaln_host_cache and adaln_table is None and not adaln_resident_cache)
             model = maybe_load_fsdp_model(
                 model_cls=model_cls,
                 init_params={
@@ -1381,8 +1435,9 @@ class TransformerLoader(ComponentLoader):
         model = model.eval()
 
         if adaln_host_cache and hasattr(model, "enable_adaln_host_cache"):
-            model.enable_adaln_host_cache(adaln_table)
+            model.enable_adaln_host_cache(adaln_table, keep_resident=adaln_resident_cache)
             logger.info("AdaLN modulation: %s", "precomputed tables from " + adaln_table if adaln_table
+                        else "resident per-timestep cache" if adaln_resident_cache
                         else "projections in pinned host memory behind a per-timestep cache")
         if layerwise_load and not fastvideo_args.dit_layerwise_offload:
             model = model.to(get_local_torch_device())
@@ -1655,8 +1710,6 @@ class PipelineComponentLoader:
             component_model_path,
         )
 
-        # A component under a subfolder (e.g. "base_model/vae") dispatches on its own name.
-        module_name = module_name.rsplit("/", 1)[-1]
         # Get the appropriate loader for this module type
         loader = ComponentLoader.for_module_type(module_name, transformers_or_diffusers)
 

fastvideo/models/registry.py

@@ -71,12 +71,7 @@ _IMAGE_TO_VIDEO_DIT_MODELS = {
     "MatrixGame3WanModel": ("dits", "matrixgame3", "MatrixGame3WanModel"),
 }
 
-_TEXT_TO_IMAGE_DIT_MODELS = {
-    "QwenImage21Transformer2DModel": ("dits", "qwen_image21", "QwenImage21Transformer2DModel"),
-}
-
 _TEXT_ENCODER_MODELS = {
-    "QwenImage21Qwen3VLConditioner": ("encoders", "qwen_image21", "QwenImage21Qwen3VLConditioner"),
     "MMAudioDFNCLIPTextEncoder": ("encoders", "mmaudio_clip", "MMAudioDFNCLIPTextEncoder"),
     "CLIPTextModel": ("encoders", "clip", "CLIPTextModel"),
     "CLIPTextModelWithProjection": ("encoders", "clip", "CLIPTextModelWithProjection"),
@@ -113,7 +108,6 @@ _VAE_MODELS = {
     "AutoencoderKLWan": ("wan", "vae", "AutoencoderKLWan"),
     "LingBotWorld2WanVAE": ("vaes", "lingbotworld2_wanvae", "LingBotWorld2WanVAE"),
     "AutoencoderKL": ("vaes", "autoencoder_kl", "AutoencoderKL"),
-    "AutoencoderKLQwenImage21": ("vaes", "qwen_image21", "AutoencoderKLQwenImage21"),
     "AutoencoderKLGen3CTokenizer": ("vaes", "gen3c_tokenizer_vae", "AutoencoderKLGen3CTokenizer"),
     "CausalVideoAutoencoder": ("vaes", "ltx2vae", "LTX2CausalVideoAutoencoder"),
     "AutoencoderKLFlux2": ("vaes", "flux2vae", "AutoencoderKLFlux2"),
@@ -160,7 +154,6 @@ _UPSAMPLERS = {
 _LEGACY_FAST_VIDEO_MODELS = {
     **_TEXT_TO_VIDEO_DIT_MODELS,
     **_IMAGE_TO_VIDEO_DIT_MODELS,
-    **_TEXT_TO_IMAGE_DIT_MODELS,
     **_TEXT_ENCODER_MODELS,
     **_IMAGE_ENCODER_MODELS,
     **_VAE_MODELS,

fastvideo/models/upsamplers/minimax_h3_latent_upscaler.py

@@ -1,161 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""MiniMax-H3 3D-convolution latent upscaler (LBH-123-AI/Comfyui_Minimax_h3_latent_Upscaler, MIT).
-
-Port of the upstream node's ``LatentResizer3D`` and its ``target dimensions`` size logic, without the ComfyUI shims,
-attention blocks (the node forces them off) and temporal chunking (Open-H3-Regenerate-2K runs without it). Module names
-match the upstream ``latent_upsampler/*.safetensors`` checkpoint, so it loads without remapping.
-
-The network works on a latent that was normalized with the statistics below, not the VAE-normalized latent the DiT
-sees: ``upscale_latent`` therefore re-normalizes its input and de-normalizes the result with the same statistics.
-"""
-from __future__ import annotations
-
-import re
-
-import torch
-import torch.nn.functional as F
-from safetensors.torch import load_file
-from torch import nn
-
-VAE_DOWNSAMPLE = 16
-PIXEL_ALIGNMENT = 32
-
-LATENTS_MEAN = (
-    0.858090341091156, -0.9606591463088989, 1.0661640167236328, -0.5090325474739075, -0.2727581858634949,
-    -1.3675414323806763, -0.2553254961967468, -0.26907554268836975, -0.5376840829849243, -0.0464097298681736,
-    0.6657370328903198, 0.19690127670764923, -0.5460608005523682, -0.4035342037677765, -0.23683024942874908,
-    0.25928452610969543, -0.30133944749832153, 0.211341992020607, -1.1206848621368408, 0.3581933379173279,
-    -0.04225143790245056, 0.2604829967021942, 0.22864092886447906, 0.7056031823158264)
-LATENTS_STD = (
-    1.2223774194717407, 1.2767263650894165, 1.6831774711608887, 1.7549455165863037, 1.5636216402053833,
-    2.194143533706665, 0.9653137922286987, 1.0569885969161987, 0.841948926448822, 0.7729952931404114,
-    1.8955937623977661, 0.946841835975647, 0.7996809482574463, 0.44988900423049927, 0.7197399735450745,
-    0.6936293244361877, 2.961095094680786, 2.7694199085235596, 3.0496184825897217, 2.1088054180145264,
-    3.276226282119751, 3.1627357006073, 2.2816812992095947, 2.6127843856811523)
-
-
-def _norm(channels: int) -> nn.GroupNorm:
-    return nn.GroupNorm(32, channels)
-
-
-class ResBlockEmb3D(nn.Module):
-
-    def __init__(self, channels: int, emb_channels: int) -> None:
-        super().__init__()
-        self.in_layers = nn.Sequential(_norm(channels), nn.SiLU(), nn.Conv3d(channels, channels, 3, padding=1))
-        self.emb_layers = nn.Sequential(nn.SiLU(), nn.Linear(emb_channels, 2 * channels))
-        self.out_norm = _norm(channels)
-        self.out_layers = nn.Sequential(nn.SiLU(), nn.Dropout(p=0.0), nn.Conv3d(channels, channels, 3, padding=1))
-
-    def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor:
-        h = self.in_layers(x)
-        scale, shift = self.emb_layers(emb).type(h.dtype)[..., None, None, None].chunk(2, dim=1)
-        h = self.out_norm(h) * (1 + scale) + shift
-        return x + self.out_layers(h)
-
-
-class TemporalConv(nn.Module):
-
-    def __init__(self, channels: int, kernel_size: int = 5) -> None:
-        super().__init__()
-        self.norm = _norm(channels)
-        self.dwconv = nn.Conv3d(channels, channels, (kernel_size, 1, 1), padding=(kernel_size // 2, 0, 0),
-                                groups=channels)
-        self.pwconv = nn.Conv3d(channels, channels, 1)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return x + self.pwconv(self.dwconv(F.silu(self.norm(x))))
-
-
-def _blocks(count: int, channels: int, emb_dim: int, temporal_every: int, kernel: int) -> nn.ModuleList:
-    blocks = nn.ModuleList()
-    for index in range(count):
-        blocks.append(ResBlockEmb3D(channels, emb_dim))
-        if temporal_every > 0 and index % temporal_every == 0:
-            blocks.append(TemporalConv(channels, kernel))
-    return blocks
-
-
-class MiniMaxH3LatentUpscaler(nn.Module):
-
-    def __init__(self,
-                 in_channels: int = 24,
-                 in_blocks: int = 12,
-                 out_blocks: int = 12,
-                 channels: int = 512,
-                 temporal_every: int = 2,
-                 temporal_kernel: int = 5) -> None:
-        super().__init__()
-        emb_dim = 64
-        self.conv_in = nn.Conv3d(in_channels, channels, 3, padding=1)
-        self.embed = nn.Sequential(nn.Linear(1, emb_dim), nn.SiLU(), nn.Linear(emb_dim, emb_dim))
-        self.in_blocks = _blocks(in_blocks, channels, emb_dim, temporal_every, temporal_kernel)
-        self.out_blocks = _blocks(out_blocks, channels, emb_dim, temporal_every, temporal_kernel)
-        self.norm_out = _norm(channels)
-        self.conv_out = nn.Conv3d(channels, in_channels, 3, padding=1)
-
-    def _run(self, blocks: nn.ModuleList, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor:
-        for block in blocks:
-            x = block(x, emb.expand(x.shape[0], -1)) if isinstance(block, ResBlockEmb3D) else block(x)
-        return x
-
-    def forward(self, x: torch.Tensor, scale: float, target_size: tuple[int, int, int]) -> torch.Tensor:
-        """``x``: [B, C, T, h, w] normalized latent; returns [B, C, T, *target_size[1:]]."""
-        emb = self.embed(torch.tensor([[scale - 1.0]], dtype=x.dtype, device=x.device))
-        x = self._run(self.in_blocks, self.conv_in(x), emb)
-        x = F.interpolate(x, size=target_size, mode="trilinear", align_corners=False)
-        x = self._run(self.out_blocks, x, emb)
-        return self.conv_out(F.silu(self.norm_out(x)))
-
-    @classmethod
-    def from_checkpoint(cls, path: str) -> MiniMaxH3LatentUpscaler:
-        state = load_file(path, device="cpu")
-        model = cls(in_channels=state["conv_in.weight"].shape[1],
-                    channels=state["conv_in.weight"].shape[0],
-                    in_blocks=_count_blocks(state, "in_blocks"),
-                    out_blocks=_count_blocks(state, "out_blocks"),
-                    temporal_kernel=_temporal_kernel(state))
-        model.load_state_dict(state, strict=True)
-        return model.eval().requires_grad_(False)
-
-
-def _count_blocks(state: dict[str, torch.Tensor], prefix: str) -> int:
-    pattern = re.compile(rf"{prefix}\.(\d+)\.in_layers\.0\.weight")
-    return len({int(m.group(1)) for key in state if (m := pattern.match(key))})
-
-
-def _temporal_kernel(state: dict[str, torch.Tensor]) -> int:
-    return next(tensor.shape[2] for key, tensor in state.items() if key.endswith("dwconv.weight"))
-
-
-def resolve_upscale_geometry(input_size: tuple[int, int], target_pixels: tuple[int, int]) -> tuple[tuple[int, int], float]:
-    """Latent (h, w) after aligning the target pixel size to 32, and the scale the network is conditioned on.
-
-    ``input_size`` is the latent (h, w); ``target_pixels`` the requested (height, width).
-    """
-    h_in, w_in = input_size
-    h_target, w_target = target_pixels
-    scale = (w_target / (w_in * VAE_DOWNSAMPLE) + h_target / (h_in * VAE_DOWNSAMPLE)) / 2.0
-    aligned = (round(round(h_target / PIXEL_ALIGNMENT) * PIXEL_ALIGNMENT / VAE_DOWNSAMPLE),
-               round(round(w_target / PIXEL_ALIGNMENT) * PIXEL_ALIGNMENT / VAE_DOWNSAMPLE))
-    return (max(1, aligned[0]), max(1, aligned[1])), scale
-
-
-@torch.inference_mode()
-def upscale_latent(model: MiniMaxH3LatentUpscaler, latent: torch.Tensor, target_pixels: tuple[int, int],
-                   dtype: torch.dtype = torch.bfloat16) -> torch.Tensor:
-    """Upscale a [C, T, h, w] latent to ``target_pixels`` (height, width); FP32 [C, T, H, W] out.
-
-    The input is returned unchanged when it already has the target latent size.
-    """
-    channels, frames, h_in, w_in = latent.shape
-    (h_out, w_out), scale = resolve_upscale_geometry((h_in, w_in), target_pixels)
-    if (h_out, w_out) == (h_in, w_in):
-        return latent.float()
-    if scale < 1.0:
-        raise ValueError("The MiniMax-H3 latent upscaler only upscales (effective scale >= 1.0).")
-    mean = torch.tensor(LATENTS_MEAN, dtype=dtype, device=latent.device).view(1, -1, 1, 1, 1)
-    std = torch.tensor(LATENTS_STD, dtype=dtype, device=latent.device).view(1, -1, 1, 1, 1)
-    source = (latent[None].to(dtype) - mean) / std
-    result = model(source, scale, (frames, h_out, w_out))
-    return (result * std + mean).float()[0]

fastvideo/models/vaes/minimax_h3_parallel.py

@@ -43,7 +43,7 @@ output is bitwise equal to serial ``decode_to_pixels``.
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Callable
 
 import torch
 
@@ -194,6 +194,8 @@ def decode_to_pixels_parallel(
     output: torch.Tensor | None,
     group: "GroupCoordinator",
     strategy: str = DEFAULT_DECODE_GATHER_STRATEGY,
+    *,
+    idle_callback: Callable[[], None] | None = None,
 ) -> torch.Tensor | None:
     """Chunk-parallel ``decode_to_pixels`` across a sequence-parallel group.
 
@@ -202,10 +204,14 @@ def decode_to_pixels_parallel(
     group's first rank, which assembles bitwise the serial
     ``decode_to_pixels`` result into ``output``. Only the first rank passes
     ``output`` (validated exactly like the serial API); other ranks pass
-    ``None`` and receive ``None``.
+    ``None`` and receive ``None``. An optional callback runs once on a rank
+    with no chunk in the final round (single-item batches only). It must not
+    issue collectives; the normal round rendezvous follows it on every rank.
     """
     if strategy not in DECODE_GATHER_STRATEGIES:
         raise ValueError(f"Unknown parallel-decode strategy {strategy!r}; expected one of {DECODE_GATHER_STRATEGIES}.")
+    if idle_callback is not None and z.shape[0] != 1:
+        raise ValueError("An idle decode callback requires a single-item batch.")
     is_leader = group.rank_in_group == 0
     if is_leader:
         if output is None:
@@ -226,7 +232,7 @@ def decode_to_pixels_parallel(
                 slice_output = output[batch_index:batch_index + 1] if output is not None else None
                 _decode_single_parallel(vae, z_slice, slice_output, group, strategy)
         else:
-            _decode_single_parallel(vae, z, output, group, strategy)
+            _decode_single_parallel(vae, z, output, group, strategy, idle_callback=idle_callback)
     finally:
         # Drain the leader's async chunk copies before the caller (or an
         # exception handler) can read or release the pinned buffer.
@@ -386,6 +392,8 @@ def _decode_single_parallel(
     output: torch.Tensor | None,
     group: "GroupCoordinator",
     strategy: str,
+    *,
+    idle_callback: Callable[[], None] | None = None,
 ) -> None:
     if _decode_single_tile_parallel(vae, z, output, group, strategy):
         return
@@ -413,6 +421,9 @@ def _decode_single_parallel(
         for round_index in range(_num_rounds(num_chunks, world_size)):
             chunk_index = round_index * world_size + rank
             if chunk_index >= num_chunks:
+                if idle_callback is not None:
+                    idle_callback()
+                    idle_callback = None
                 segment = torch.zeros(segment_shape, dtype=segment_dtype, device=z.device)
             elif round_index == 0 and first_segment is not None:
                 segment = first_segment

fastvideo/models/vaes/qwen_image21.py

@@ -1,656 +0,0 @@
-# Copyright 2026 The Qwen Team and The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Native single-frame RGBA autoencoder for Qwen-Image-2.1."""
-
-from __future__ import annotations
-
-import json
-import math
-from dataclasses import dataclass
-from pathlib import Path
-
-import torch
-from torch import nn
-from torch.nn import functional as F
-
-from fastvideo.configs.models.vaes.base import VAEConfig
-from fastvideo.configs.models.vaes.qwen_image21 import QwenImage21VAEConfig
-from fastvideo.logger import init_logger
-
-logger = init_logger(__name__)
-
-
-class DiagonalGaussianDistribution:
-    def __init__(self, parameters: torch.Tensor, deterministic: bool = False) -> None:
-        self.parameters = parameters
-        self.mean, self.logvar = parameters.chunk(2, dim=1)
-        self.logvar = self.logvar.clamp(-30.0, 20.0)
-        self.deterministic = deterministic
-        self.std = torch.exp(0.5 * self.logvar)
-        self.var = torch.exp(self.logvar)
-        if deterministic:
-            self.std = torch.zeros_like(self.mean)
-            self.var = torch.zeros_like(self.mean)
-
-    def mode(self) -> torch.Tensor:
-        return self.mean
-
-    def sample(self, generator: torch.Generator | list[torch.Generator] | None = None) -> torch.Tensor:
-        if isinstance(generator, list):
-            if len(generator) != self.mean.shape[0]:
-                raise ValueError("A generator list must contain one generator per batch element.")
-            noise = torch.cat([
-                self._noise(self.mean[index:index + 1], item) for index, item in enumerate(generator)
-            ])
-        else:
-            noise = self._noise(self.mean, generator)
-        return self.mean + self.std * noise
-
-    @staticmethod
-    def _noise(sample: torch.Tensor, generator: torch.Generator | None) -> torch.Tensor:
-        noise_device = sample.device if generator is None else generator.device
-        return torch.randn(sample.shape, generator=generator, device=noise_device, dtype=sample.dtype).to(sample.device)
-
-    def kl(self, other: DiagonalGaussianDistribution | None = None) -> torch.Tensor:
-        if self.deterministic:
-            return torch.zeros(1, device=self.mean.device, dtype=self.mean.dtype)
-        if other is None:
-            terms = self.mean.square() + self.var - 1.0 - self.logvar
-        else:
-            terms = ((self.mean - other.mean).square() + self.var) / other.var - 1.0 - self.logvar + other.logvar
-        return 0.5 * terms.sum(dim=(1, 2, 3))
-
-    def nll(self, sample: torch.Tensor, dims: tuple[int, ...] = (1, 2, 3)) -> torch.Tensor:
-        if self.deterministic:
-            return torch.zeros(1, device=self.mean.device, dtype=self.mean.dtype)
-        return 0.5 * (math.log(2 * math.pi) + self.logvar + (sample - self.mean).square() / self.var).sum(dim=dims)
-
-
-@dataclass
-class AutoencoderKLOutput:
-    latent_dist: DiagonalGaussianDistribution
-
-
-@dataclass
-class DecoderOutput:
-    sample: torch.Tensor
-
-
-class QwenImage21AvgDown3D(nn.Module):
-    def __init__(self, in_channels: int, out_channels: int, factor_t: int, factor_s: int = 1) -> None:
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = out_channels
-        self.factor_t = factor_t
-        self.factor_s = factor_s
-        self.factor = factor_t * factor_s * factor_s
-        if in_channels * self.factor % out_channels:
-            raise ValueError("Average shortcut channels must be divisible by the downsampling factor.")
-        self.group_size = in_channels * self.factor // out_channels
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t
-        x = F.pad(x, (0, 0, 0, 0, pad_t, 0))
-        batch, channels, frames, height, width = x.shape
-        x = x.view(batch, channels, frames // self.factor_t, self.factor_t,
-                   height // self.factor_s, self.factor_s, width // self.factor_s, self.factor_s)
-        x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous()
-        x = x.view(batch, channels * self.factor, frames // self.factor_t,
-                   height // self.factor_s, width // self.factor_s)
-        x = x.view(batch, self.out_channels, self.group_size, frames // self.factor_t,
-                   height // self.factor_s, width // self.factor_s)
-        return x.mean(dim=2)
-
-
-class QwenImage21DupUp3D(nn.Module):
-    def __init__(self, in_channels: int, out_channels: int, factor_t: int, factor_s: int = 1) -> None:
-        super().__init__()
-        self.out_channels = out_channels
-        self.factor_t = factor_t
-        self.factor_s = factor_s
-        factor = factor_t * factor_s * factor_s
-        if out_channels * factor % in_channels:
-            raise ValueError("Duplicate shortcut channels must be divisible by the upsampling factor.")
-        self.repeats = out_channels * factor // in_channels
-
-    def forward(self, x: torch.Tensor, first_chunk: bool = False) -> torch.Tensor:
-        x = x.repeat_interleave(self.repeats, dim=1)
-        x = x.view(x.shape[0], self.out_channels, self.factor_t, self.factor_s, self.factor_s,
-                   x.shape[2], x.shape[3], x.shape[4])
-        x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
-        x = x.view(x.shape[0], self.out_channels, x.shape[2] * self.factor_t,
-                   x.shape[4] * self.factor_s, x.shape[6] * self.factor_s)
-        return x[:, :, self.factor_t - 1:] if first_chunk else x
-
-
-class QwenImage21CausalConv3d(nn.Conv2d):
-    """Spatial convolution retaining the image VAE's singleton frame axis."""
-
-    def __init__(self, in_channels: int, out_channels: int, kernel_size, stride=1, padding=0) -> None:
-        super().__init__(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
-        self._padding = (self.padding[1], self.padding[1], self.padding[0], self.padding[0])
-        self.padding = (0, 0)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        if x.ndim != 5 or x.shape[2] != 1:
-            raise ValueError("Qwen-Image-2.1 spatial convolutions require exactly one frame.")
-        return super().forward(F.pad(x.squeeze(2), self._padding)).unsqueeze(2)
-
-
-class QwenImage21RMSNorm(nn.Module):
-    def __init__(self, dim: int, channel_first: bool = True, images: bool = True, bias: bool = False) -> None:
-        super().__init__()
-        broadcast_dims = (1, 1) if images else (1, 1, 1)
-        shape = (dim, *broadcast_dims) if channel_first else (dim,)
-        self.channel_first = channel_first
-        self.scale = dim**0.5
-        self.gamma = nn.Parameter(torch.ones(shape))
-        self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        upcast = x.dtype in (torch.float16, torch.bfloat16) or any(
-            name in str(x.dtype) for name in ("float4_", "float8_")
-        )
-        normalized = F.normalize(x.float() if upcast else x, dim=1 if self.channel_first else -1).to(x.dtype)
-        return normalized * self.scale * self.gamma + self.bias
-
-
-class QwenImage21Upsample(nn.Upsample):
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return super().forward(x.float()).type_as(x)
-
-
-class QwenImage21Resample(nn.Module):
-    def __init__(self, dim: int, mode: str, upsample_out_dim: int | None = None) -> None:
-        super().__init__()
-        self.mode = mode
-        if upsample_out_dim is None:
-            upsample_out_dim = dim // 2
-        if mode in ("upsample2d", "upsample3d"):
-            self.resample = nn.Sequential(
-                QwenImage21Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
-                nn.Conv2d(dim, upsample_out_dim, 3, padding=1),
-            )
-            if mode == "upsample3d":
-                self.time_conv = QwenImage21CausalConv3d(dim, dim * 2, (1, 1), padding=(0, 0))
-        elif mode in ("downsample2d", "downsample3d"):
-            self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)))
-            if mode == "downsample3d":
-                self.time_conv = QwenImage21CausalConv3d(dim, dim, (1, 1), stride=(1, 1), padding=(0, 0))
-        else:
-            self.resample = nn.Identity()
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        batch, channels, frames, height, width = x.shape
-        # Temporal convolutions have no preceding frame in this image model.
-        x = x.permute(0, 2, 1, 3, 4).reshape(batch * frames, channels, height, width)
-        x = self.resample(x)
-        return x.view(batch, frames, x.shape[1], x.shape[2], x.shape[3]).permute(0, 2, 1, 3, 4)
-
-
-class QwenImage21ResidualBlock(nn.Module):
-    def __init__(self, in_dim: int, out_dim: int, dropout: float = 0.0) -> None:
-        super().__init__()
-        self.norm1 = QwenImage21RMSNorm(in_dim, images=False)
-        self.conv1 = QwenImage21CausalConv3d(in_dim, out_dim, 3, padding=1)
-        self.norm2 = QwenImage21RMSNorm(out_dim, images=False)
-        self.dropout = nn.Dropout(dropout)
-        self.conv2 = QwenImage21CausalConv3d(out_dim, out_dim, 3, padding=1)
-        self.conv_shortcut = QwenImage21CausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity()
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        residual = self.conv_shortcut(x)
-        x = self.conv1(F.silu(self.norm1(x)))
-        x = self.conv2(self.dropout(F.silu(self.norm2(x))))
-        return x + residual
-
-
-class QwenImage21AttentionBlock(nn.Module):
-    def __init__(self, dim: int) -> None:
-        super().__init__()
-        self.norm = QwenImage21RMSNorm(dim)
-        self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
-        self.proj = nn.Conv2d(dim, dim, 1)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        residual = x
-        batch, channels, frames, height, width = x.shape
-        x = x.permute(0, 2, 1, 3, 4).reshape(batch * frames, channels, height, width)
-        qkv = self.to_qkv(self.norm(x)).reshape(batch * frames, 1, channels * 3, -1)
-        q, k, v = qkv.permute(0, 1, 3, 2).contiguous().chunk(3, dim=-1)
-        # The VAE's one-head spatial attention has no distributed sequence layout.
-        x = F.scaled_dot_product_attention(q, k, v)
-        x = x.squeeze(1).permute(0, 2, 1).reshape(batch * frames, channels, height, width)
-        x = self.proj(x).view(batch, frames, channels, height, width).permute(0, 2, 1, 3, 4)
-        return x + residual
-
-
-class QwenImage21MidBlock(nn.Module):
-    def __init__(self, dim: int, dropout: float = 0.0) -> None:
-        super().__init__()
-        self.resnets = nn.ModuleList([
-            QwenImage21ResidualBlock(dim, dim, dropout),
-            QwenImage21ResidualBlock(dim, dim, dropout),
-        ])
-        self.attentions = nn.ModuleList([QwenImage21AttentionBlock(dim)])
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.resnets[1](self.attentions[0](self.resnets[0](x)))
-
-
-class QwenImage21ResidualDownBlock(nn.Module):
-    def __init__(self, in_dim: int, out_dim: int, dropout: float, num_res_blocks: int,
-                 temporal_downsample: bool, down_flag: bool) -> None:
-        super().__init__()
-        self.avg_shortcut = QwenImage21AvgDown3D(in_dim, out_dim, 2 if temporal_downsample else 1,
-                                               2 if down_flag else 1)
-        self.resnets = nn.ModuleList([
-            QwenImage21ResidualBlock(in_dim if index == 0 else out_dim, out_dim, dropout)
-            for index in range(num_res_blocks)
-        ])
-        self.downsampler = QwenImage21Resample(
-            out_dim, "downsample3d" if temporal_downsample else "downsample2d"
-        ) if down_flag else None
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        residual = x.clone()
-        for resnet in self.resnets:
-            x = resnet(x)
-        if self.downsampler is not None:
-            x = self.downsampler(x)
-        return x + self.avg_shortcut(residual)
-
-
-class QwenImage21Encoder3d(nn.Module):
-    def __init__(self, arch) -> None:
-        super().__init__()
-        dims = [arch.base_dim * factor for factor in [1, *arch.dim_mult]]
-        self.conv_in = QwenImage21CausalConv3d(arch.in_channels, dims[0], 3, padding=1)
-        self.down_blocks = nn.ModuleList()
-        scale = 1.0
-        for index, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
-            down_flag = index != len(arch.dim_mult) - 1
-            temporal = arch.temperal_downsample[index] if down_flag else False
-            if arch.is_residual:
-                self.down_blocks.append(QwenImage21ResidualDownBlock(
-                    in_dim, out_dim, arch.dropout, arch.num_res_blocks, temporal, down_flag
-                ))
-            else:
-                for _ in range(arch.num_res_blocks):
-                    self.down_blocks.append(QwenImage21ResidualBlock(in_dim, out_dim, arch.dropout))
-                    if scale in arch.attn_scales:
-                        self.down_blocks.append(QwenImage21AttentionBlock(out_dim))
-                    in_dim = out_dim
-                if down_flag:
-                    mode = "downsample3d" if temporal else "downsample2d"
-                    self.down_blocks.append(QwenImage21Resample(out_dim, mode))
-                    scale /= 2.0
-        self.mid_block = QwenImage21MidBlock(out_dim, arch.dropout)
-        self.norm_out = QwenImage21RMSNorm(out_dim, images=False)
-        self.conv_out = QwenImage21CausalConv3d(out_dim, arch.z_dim * 2, 3, padding=1)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.conv_in(x)
-        for block in self.down_blocks:
-            x = block(x)
-        x = self.mid_block(x)
-        return self.conv_out(F.silu(self.norm_out(x)))
-
-
-class QwenImage21ResidualUpBlock(nn.Module):
-    def __init__(self, in_dim: int, out_dim: int, arch, temporal_upsample: bool, up_flag: bool) -> None:
-        super().__init__()
-        self.resnets = nn.ModuleList([
-            QwenImage21ResidualBlock(in_dim if index == 0 else out_dim, out_dim, arch.dropout)
-            for index in range(arch.num_res_blocks + 1)
-        ])
-        self.avg_shortcut = QwenImage21DupUp3D(
-            in_dim, out_dim, 2 if temporal_upsample else 1, 2
-        ) if up_flag else None
-        self.upsampler = QwenImage21Resample(
-            out_dim, "upsample3d" if temporal_upsample else "upsample2d", upsample_out_dim=out_dim
-        ) if up_flag else None
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        residual = x.clone()
-        for resnet in self.resnets:
-            x = resnet(x)
-        if self.upsampler is not None:
-            x = self.upsampler(x)
-        if self.avg_shortcut is not None:
-            x = x + self.avg_shortcut(residual, first_chunk=True)
-        return x
-
-
-class QwenImage21UpBlock(nn.Module):
-    def __init__(self, in_dim: int, out_dim: int, arch, upsample_mode: str | None) -> None:
-        super().__init__()
-        self.resnets = nn.ModuleList([
-            QwenImage21ResidualBlock(in_dim if index == 0 else out_dim, out_dim, arch.dropout)
-            for index in range(arch.num_res_blocks + 1)
-        ])
-        self.upsamplers = nn.ModuleList([QwenImage21Resample(out_dim, upsample_mode)]) if upsample_mode else None
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        for resnet in self.resnets:
-            x = resnet(x)
-        return self.upsamplers[0](x) if self.upsamplers is not None else x
-
-
-class QwenImage21Decoder3d(nn.Module):
-    def __init__(self, arch) -> None:
-        super().__init__()
-        dim = arch.decoder_base_dim if arch.decoder_base_dim is not None else arch.base_dim
-        dims = [dim * factor for factor in [arch.dim_mult[-1], *arch.dim_mult[::-1]]]
-        temporal_upsample = arch.temperal_downsample[::-1]
-        self.conv_in = QwenImage21CausalConv3d(arch.z_dim, dims[0], 3, padding=1)
-        self.mid_block = QwenImage21MidBlock(dims[0], arch.dropout)
-        self.up_blocks = nn.ModuleList()
-        for index, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
-            up_flag = index != len(arch.dim_mult) - 1
-            temporal = temporal_upsample[index] if up_flag else False
-            if arch.is_residual:
-                self.up_blocks.append(QwenImage21ResidualUpBlock(in_dim, out_dim, arch, temporal, up_flag))
-            else:
-                if index > 0:
-                    in_dim //= 2
-                mode = ("upsample3d" if temporal else "upsample2d") if up_flag else None
-                self.up_blocks.append(QwenImage21UpBlock(in_dim, out_dim, arch, mode))
-        self.norm_out = QwenImage21RMSNorm(out_dim, images=False)
-        self.conv_out = QwenImage21CausalConv3d(out_dim, arch.out_channels, 3, padding=1)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.mid_block(self.conv_in(x))
-        for block in self.up_blocks:
-            x = block(x)
-        return self.conv_out(F.silu(self.norm_out(x)))
-
-
-def _patchify(x: torch.Tensor, patch_size: int) -> torch.Tensor:
-    if patch_size == 1:
-        return x
-    batch, channels, frames, height, width = x.shape
-    if height % patch_size or width % patch_size:
-        raise ValueError("Image height and width must be divisible by the VAE patch size.")
-    x = x.view(batch, channels, frames, height // patch_size, patch_size, width // patch_size, patch_size)
-    x = x.permute(0, 1, 6, 4, 2, 3, 5).contiguous()
-    return x.view(batch, channels * patch_size * patch_size, frames, height // patch_size, width // patch_size)
-
-
-def _unpatchify(x: torch.Tensor, patch_size: int) -> torch.Tensor:
-    if patch_size == 1:
-        return x
-    batch, channels, frames, height, width = x.shape
-    channels //= patch_size * patch_size
-    x = x.view(batch, channels, patch_size, patch_size, frames, height, width)
-    x = x.permute(0, 1, 4, 5, 3, 6, 2).contiguous()
-    return x.view(batch, channels, frames, height * patch_size, width * patch_size)
-
-
-class AutoencoderKLQwenImage21(nn.Module):
-    _supports_gradient_checkpointing = False
-    _group_offload_block_modules = ["quant_conv", "post_quant_conv", "encoder", "decoder"]
-
-    def __init__(self, config: VAEConfig) -> None:
-        super().__init__()
-        self.config = config
-        if not config.load_encoder or not config.load_decoder:
-            raise ValueError("Qwen-Image-2.1 requires the complete encoder and decoder checkpoint.")
-        arch = config.arch_config
-        self.z_dim = arch.z_dim
-        self.latent_channels = arch.z_dim
-        self.encoder = QwenImage21Encoder3d(arch)
-        self.quant_conv = QwenImage21CausalConv3d(arch.z_dim * 2, arch.z_dim * 2, 1)
-        self.post_quant_conv = QwenImage21CausalConv3d(arch.z_dim, arch.z_dim, 1)
-        self.decoder = QwenImage21Decoder3d(arch)
-        self.spatial_compression_ratio = arch.scale_factor_spatial
-        self.temporal_compression_ratio = 1
-        self.use_slicing = False
-        self.use_tiling = config.use_tiling
-        self.tile_sample_min_height = config.tile_sample_min_height
-        self.tile_sample_min_width = config.tile_sample_min_width
-        self.tile_sample_stride_height = config.tile_sample_stride_height
-        self.tile_sample_stride_width = config.tile_sample_stride_width
-        self._validate_tiling()
-
-    @property
-    def device(self) -> torch.device:
-        return next(self.parameters()).device
-
-    @property
-    def dtype(self) -> torch.dtype:
-        return next(self.parameters()).dtype
-
-    def enable_slicing(self) -> None:
-        self.use_slicing = True
-
-    def disable_slicing(self) -> None:
-        self.use_slicing = False
-
-    def _validate_tiling(self) -> None:
-        ratio = self.spatial_compression_ratio
-        for extent, stride in ((self.tile_sample_min_height, self.tile_sample_stride_height),
-                               (self.tile_sample_min_width, self.tile_sample_stride_width)):
-            if extent < ratio or stride < ratio or extent % ratio or stride % ratio or stride > extent:
-                raise ValueError(
-                    "VAE tile size and stride must be positive spatial-scale multiples, with stride <= size."
-                )
-
-    def enable_tiling(self, tile_sample_min_height: int | None = None, tile_sample_min_width: int | None = None,
-                      tile_sample_stride_height: int | None = None, tile_sample_stride_width: int | None = None,
-                      **kwargs) -> None:
-        self.use_tiling = True
-        self.tile_sample_min_height = tile_sample_min_height or self.tile_sample_min_height
-        self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width
-        self.tile_sample_stride_height = tile_sample_stride_height or self.tile_sample_stride_height
-        self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width
-        self._validate_tiling()
-
-    def disable_tiling(self) -> None:
-        self.use_tiling = False
-
-    def clear_cache(self) -> None:
-        """Image inference retains no temporal feature cache."""
-
-    @staticmethod
-    def _validate_image(x: torch.Tensor, channels: int, name: str) -> None:
-        if x.ndim != 5 or x.shape[1] != channels or x.shape[2] != 1:
-            raise ValueError(f"`{name}` must have shape [B, {channels}, 1, H, W], got {tuple(x.shape)}.")
-        if x.shape[0] < 1 or min(x.shape[-2:]) < 1:
-            raise ValueError(f"`{name}` must contain at least one nonempty image.")
-
-    def _encode(self, x: torch.Tensor) -> torch.Tensor:
-        if self.config.patch_size is not None:
-            x = _patchify(x, self.config.patch_size)
-        if self.use_tiling and (x.shape[-2] > self.tile_sample_min_height or x.shape[-1] > self.tile_sample_min_width):
-            return self.tiled_encode(x)
-        return self.quant_conv(self.encoder(x))
-
-    def encode(
-        self, x: torch.Tensor, return_dict: bool = True
-    ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]:
-        self._validate_image(x, self.config.in_channels, "x")
-        if x.shape[-2] % self.spatial_compression_ratio or x.shape[-1] % self.spatial_compression_ratio:
-            raise ValueError("Image height and width must be divisible by the VAE spatial compression ratio.")
-        if self.use_slicing and x.shape[0] > 1:
-            moments = torch.cat([self._encode(item) for item in x.split(1)])
-        else:
-            moments = self._encode(x)
-        posterior = DiagonalGaussianDistribution(moments)
-        return AutoencoderKLOutput(posterior) if return_dict else (posterior,)
-
-    def _decode(self, z: torch.Tensor) -> torch.Tensor:
-        min_height = self.tile_sample_min_height // self.spatial_compression_ratio
-        min_width = self.tile_sample_min_width // self.spatial_compression_ratio
-        if self.use_tiling and (z.shape[-2] > min_height or z.shape[-1] > min_width):
-            return self.tiled_decode(z).sample
-        decoded = self.decoder(self.post_quant_conv(z))
-        if self.config.patch_size is not None:
-            decoded = _unpatchify(decoded, self.config.patch_size)
-        return decoded.clamp(-1.0, 1.0)
-
-    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
-        self._validate_image(z, self.z_dim, "z")
-        if self.use_slicing and z.shape[0] > 1:
-            decoded = torch.cat([self._decode(item) for item in z.split(1)])
-        else:
-            decoded = self._decode(z)
-        return DecoderOutput(decoded) if return_dict else (decoded,)
-
-    @staticmethod
-    def blend_v(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
-        blend_extent = min(a.shape[-2], b.shape[-2], blend_extent)
-        for y in range(blend_extent):
-            b[:, :, :, y, :] = (
-                a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent)
-                + b[:, :, :, y, :] * (y / blend_extent)
-            )
-        return b
-
-    @staticmethod
-    def blend_h(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
-        blend_extent = min(a.shape[-1], b.shape[-1], blend_extent)
-        for x in range(blend_extent):
-            b[:, :, :, :, x] = (
-                a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent)
-                + b[:, :, :, :, x] * (x / blend_extent)
-            )
-        return b
-
-    def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
-        height, width = x.shape[-2:]
-        ratio = self.spatial_compression_ratio // (self.config.patch_size or 1)
-        tile_height = self.tile_sample_min_height // ratio
-        tile_width = self.tile_sample_min_width // ratio
-        stride_height = self.tile_sample_stride_height // ratio
-        stride_width = self.tile_sample_stride_width // ratio
-        rows = []
-        for i in range(0, height, self.tile_sample_stride_height):
-            row = []
-            for j in range(0, width, self.tile_sample_stride_width):
-                tile = x[:, :, :, i:i + self.tile_sample_min_height, j:j + self.tile_sample_min_width]
-                row.append(self.quant_conv(self.encoder(tile)))
-            rows.append(row)
-        result_rows = []
-        for i, row in enumerate(rows):
-            result_row = []
-            for j, tile in enumerate(row):
-                if i > 0:
-                    tile = self.blend_v(rows[i - 1][j], tile, tile_height - stride_height)
-                if j > 0:
-                    tile = self.blend_h(row[j - 1], tile, tile_width - stride_width)
-                result_row.append(tile[:, :, :, :stride_height, :stride_width])
-            result_rows.append(torch.cat(result_row, dim=-1))
-        return torch.cat(result_rows, dim=3)[:, :, :, :height // ratio, :width // ratio]
-
-    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
-        height, width = z.shape[-2:]
-        ratio = self.spatial_compression_ratio
-        tile_height = self.tile_sample_min_height // ratio
-        tile_width = self.tile_sample_min_width // ratio
-        stride_height = self.tile_sample_stride_height // ratio
-        stride_width = self.tile_sample_stride_width // ratio
-        sample_height, sample_width = height * ratio, width * ratio
-        sample_stride_height = self.tile_sample_stride_height
-        sample_stride_width = self.tile_sample_stride_width
-        if self.config.patch_size is not None:
-            patch_size = self.config.patch_size
-            sample_height //= patch_size
-            sample_width //= patch_size
-            sample_stride_height //= patch_size
-            sample_stride_width //= patch_size
-            blend_height = self.tile_sample_min_height // patch_size - sample_stride_height
-            blend_width = self.tile_sample_min_width // patch_size - sample_stride_width
-        else:
-            blend_height = self.tile_sample_min_height - sample_stride_height
-            blend_width = self.tile_sample_min_width - sample_stride_width
-        rows = []
-        for i in range(0, height, stride_height):
-            row = []
-            for j in range(0, width, stride_width):
-                tile = z[:, :, :, i:i + tile_height, j:j + tile_width]
-                row.append(self.decoder(self.post_quant_conv(tile)))
-            rows.append(row)
-        result_rows = []
-        for i, row in enumerate(rows):
-            result_row = []
-            for j, tile in enumerate(row):
-                if i > 0:
-                    tile = self.blend_v(rows[i - 1][j], tile, blend_height)
-                if j > 0:
-                    tile = self.blend_h(row[j - 1], tile, blend_width)
-                result_row.append(tile[:, :, :, :sample_stride_height, :sample_stride_width])
-            result_rows.append(torch.cat(result_row, dim=-1))
-        decoded = torch.cat(result_rows, dim=3)[:, :, :, :sample_height, :sample_width]
-        if self.config.patch_size is not None:
-            decoded = _unpatchify(decoded, self.config.patch_size)
-        decoded = decoded.clamp(-1.0, 1.0)
-        return DecoderOutput(decoded) if return_dict else (decoded,)
-
-    def normalize_latents(self, latents: torch.Tensor) -> torch.Tensor:
-        mean = torch.tensor(self.config.latents_mean, device=latents.device, dtype=latents.dtype)
-        std = torch.tensor(self.config.latents_std, device=latents.device, dtype=latents.dtype)
-        mean, std = mean.view(1, self.z_dim, 1, 1, 1), std.view(1, self.z_dim, 1, 1, 1)
-        return (latents - mean) / std
-
-    def denormalize_latents(self, latents: torch.Tensor) -> torch.Tensor:
-        mean = torch.tensor(self.config.latents_mean, device=latents.device, dtype=latents.dtype)
-        std = torch.tensor(self.config.latents_std, device=latents.device, dtype=latents.dtype)
-        mean, std = mean.view(1, self.z_dim, 1, 1, 1), std.view(1, self.z_dim, 1, 1, 1)
-        return latents * std + mean
-
-    def forward(self, sample: torch.Tensor, sample_posterior: bool = False, return_dict: bool = True,
-                generator: torch.Generator | None = None) -> DecoderOutput | tuple[torch.Tensor]:
-        posterior = self.encode(sample).latent_dist
-        latents = posterior.sample(generator=generator) if sample_posterior else posterior.mode()
-        return self.decode(latents, return_dict=return_dict)
-
-    @classmethod
-    def from_pretrained(cls, model_path: str | Path, *, subfolder: str | None = None,
-                        config: VAEConfig | None = None, torch_dtype: torch.dtype | None = None,
-                        device: torch.device | str = "cpu", local_files_only: bool = True) -> AutoencoderKLQwenImage21:
-        """Strictly load a local Diffusers-layout VAE component without downloading weights."""
-        from safetensors.torch import load_file
-
-        if not local_files_only:
-            raise ValueError("Qwen-Image-2.1 VAE.from_pretrained accepts local checkpoints only.")
-        path = Path(model_path)
-        if subfolder:
-            path /= subfolder
-        if not path.is_dir():
-            raise FileNotFoundError(f"Local Qwen-Image-2.1 VAE directory does not exist: {path}")
-        with (path / "config.json").open() as handle:
-            arch_fields = json.load(handle)
-        class_name = arch_fields.get("_class_name")
-        if class_name != "AutoencoderKLQwenImage21":
-            raise ValueError(f"Expected AutoencoderKLQwenImage21 checkpoint, got {class_name!r}.")
-        config = config or QwenImage21VAEConfig()
-        config.update_model_arch(arch_fields)
-        model = cls(config)
-        files = sorted(path.glob("*.safetensors"))
-        if not files:
-            raise FileNotFoundError(f"No safetensors VAE checkpoint found in {path}.")
-        state = {}
-        for file in files:
-            shard = load_file(str(file))
-            duplicates = state.keys() & shard.keys()
-            if duplicates:
-                raise ValueError(f"Duplicate tensors in VAE checkpoint shards: {sorted(duplicates)[:5]}")
-            state.update(shard)
-        model.load_state_dict(state, strict=True)
-        logger.info("Loaded Qwen-Image-2.1 VAE strictly from %s", path)
-        return model.to(device=device, dtype=torch_dtype).eval()
-
-
-EntryClass = AutoencoderKLQwenImage21

fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py

@@ -740,8 +740,10 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase):
                 audio_scheduler=audio_scheduler,
             ),
         )
-        self.add_stage("video_decoding_stage", MiniMaxH3VideoDecodingStage(vae=None if use_taeh3 else vae))
-        self.add_stage("audio_decoding_stage", MiniMaxH3AudioDecodingStage(audio_vae=audio_vae))
+        audio_stage = MiniMaxH3AudioDecodingStage(audio_vae=audio_vae)
+        self.add_stage("video_decoding_stage",
+                       MiniMaxH3VideoDecodingStage(vae=None if use_taeh3 else vae, audio_stage=audio_stage))
+        self.add_stage("audio_decoding_stage", audio_stage)
         self._denoise_stages_ready = True
 
     def _add_stages(self, fastvideo_args: FastVideoArgs, *, ref2va: bool) -> None:

fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 from typing import Any
+from collections.abc import Callable
 
 import torch
 
@@ -28,10 +29,12 @@ from fastvideo.pipelines.stages.base import PipelineStage
 from fastvideo.pipelines.stages.validators import StageValidators as V
 from fastvideo.pipelines.stages.validators import VerificationResult
 from fastvideo.platforms import current_platform
-from fastvideo.utils import allocate_cpu_tensor_with_pin_fallback
+from fastvideo.utils import PRECISION_TO_TYPE, allocate_cpu_tensor_with_pin_fallback
 
 logger = init_logger(__name__)
 
+_IDLE_AUDIO_KEY = "minimax_h3_idle_audio"
+
 
 def _layout(batch: ForwardBatch) -> MiniMaxH3PackedLayout:
     layout = batch.extra.get(MINIMAX_H3_LAYOUT_KEY)
@@ -65,9 +68,12 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
 
     performance_component_metric = "vae_decode_time_s"
 
-    def __init__(self, vae: AutoencoderKLMiniMaxH3 | None) -> None:
+    def __init__(self,
+                 vae: AutoencoderKLMiniMaxH3 | None,
+                 audio_stage: MiniMaxH3AudioDecodingStage | None = None) -> None:
         super().__init__()
         self.vae = vae
+        self.audio_stage = audio_stage
         self._output_buffer: tuple[tuple[tuple[int, ...], bool], torch.Tensor] | None = None
         self._unified_memory: bool | None = None
 
@@ -94,6 +100,31 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
             self._output_buffer = (key, buffer)
         return self._output_buffer[1]
 
+    def _idle_audio_callback(self, batch: ForwardBatch, args: FastVideoArgs, group: Any,
+                             device: torch.device) -> Callable[[], None] | None:
+        """Use the last idle slot of an odd two-rank chunk plan for audio.
+
+        The result lives only on this request. The audio stage transfers it
+        after video assembly, keeping its normal output ownership and cleanup.
+        """
+        audio_stage, vae = self.audio_stage, self.vae
+        if audio_stage is None or vae is None or group.world_size != 2 or batch.raw_latent_shape[0] != 1:
+            return None
+        _, chunks, _ = vae._temporal_decode_plan(batch.raw_latent_shape[2])
+        if chunks <= 1 or chunks % 2 == 0:
+            return None
+        batch.extra[_IDLE_AUDIO_KEY] = None
+        if group.rank_in_group != 1:
+            return None
+        # Capture the outer stage context before video enables its own autocast.
+        enabled, dtype = torch.is_autocast_enabled(device.type), torch.get_autocast_dtype(device.type)
+
+        def decode_audio() -> None:
+            with torch.autocast(device_type=device.type, enabled=enabled, dtype=dtype):
+                batch.extra[_IDLE_AUDIO_KEY] = audio_stage.decode_audio(batch, args)
+
+        return decode_audio
+
     def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
         result = VerificationResult()
         result.add_check("layout", batch.extra.get(MINIMAX_H3_LAYOUT_KEY), V.not_none)
@@ -109,6 +140,7 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
     @torch.no_grad()
     def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
         """Decode H3 video latents into normalized CPU pixels."""
+        batch.extra.pop(_IDLE_AUDIO_KEY, None)
         placeholder = torch.empty((0, 3, 0, 0, 0), device="cpu", dtype=torch.float32)
         sp_group, is_output_rank, parallel = _decode_participation(fastvideo_args, fastvideo_args.vae_parallel_decode)
         if not is_output_rank and not parallel:
@@ -140,12 +172,15 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
                 return batch
             expected = taeh3_decoded_pixel_shape(tuple(latents.shape))
             logger.info("MiniMax-H3 video decode: TAEH3 preview (%s -> %s)", tuple(latents.shape), expected)
+            decode_precision = (fastvideo_args.pipeline_config.vae_decode_precision
+                                or fastvideo_args.pipeline_config.vae_precision)
             with nvtx_range("minimax_h3.taeh3"):
                 pixels = decode_ncthw_latents_taeh3(
                     latents,
                     device=device,
                     checkpoint_path=getattr(fastvideo_args, "taeh3_checkpoint", None),
                     chunk_size=int(getattr(fastvideo_args, "taeh3_chunk_size", 5) or 5),
+                    dtype=PRECISION_TO_TYPE[decode_precision],
                 )
             batch.output = pixels.float().cpu() if is_output_rank else placeholder
             if is_output_rank and tuple(batch.output.shape) != expected:
@@ -175,6 +210,7 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
             # Attribute the streamed decoder computation while retaining
             # per-chunk device-to-host transfer and pinned-buffer reuse.
             # fp16 for the release decoders; NVFP4 decoders were trained under bf16.
+            idle_callback = self._idle_audio_callback(batch, fastvideo_args, sp_group, device) if parallel else None
             autocast_dtype = getattr(self.vae, "decode_autocast_dtype", torch.float16)
             with (
                     nvtx_range("minimax_h3.vae"),
@@ -184,7 +220,15 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
                     strategy = fastvideo_args.vae_parallel_decode_strategy or DEFAULT_DECODE_GATHER_STRATEGY
                     logger.info("MiniMax-H3 VAE decode: sequence-parallel chunks across %d ranks (%s)",
                                 sp_group.world_size, strategy)
-                    decode_to_pixels_parallel(self.vae, latents, output, sp_group, strategy=strategy)
+                    if idle_callback is None:
+                        decode_to_pixels_parallel(self.vae, latents, output, sp_group, strategy=strategy)
+                    else:
+                        decode_to_pixels_parallel(self.vae,
+                                                  latents,
+                                                  output,
+                                                  sp_group,
+                                                  strategy=strategy,
+                                                  idle_callback=idle_callback)
                 else:
                     self.vae.decode_to_pixels(latents, output)
             if is_output_rank and output is not None and output.dtype == torch.float32:
@@ -225,15 +269,26 @@ class MiniMaxH3AudioDecodingStage(PipelineStage):
     @torch.no_grad()
     def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
         """Decode H3 audio latents into a stereo CPU waveform."""
-        # Audio decode is sub-second, so preserve the serial path's output
-        # ownership (global rank zero, or the first denoise rank under the
-        # MiniMax-H3 encoder split).
-        if model_parallel_is_initialized() and not h3_is_output_worker(fastvideo_args):
+        produced = None
+        if _IDLE_AUDIO_KEY in batch.extra:
+            group = get_sp_group()
+            produced = group.broadcast_object(batch.extra.pop(_IDLE_AUDIO_KEY), src=1)
+        is_output = not model_parallel_is_initialized() or h3_is_output_worker(fastvideo_args)
+        if not is_output:
             batch.extra["audio"] = torch.empty((0, 2), device="cpu", dtype=torch.float32)
             batch.extra["audio_sample_rate"] = self.audio_vae.sampling_rate
-            self._clear_runtime(batch)
-            return batch
+        else:
+            # A tile-parallel decoder may not use a temporal idle slot. All
+            # ranks observe None in that case and retain the serial audio path.
+            waveform, sample_rate = produced if produced is not None else self.decode_audio(batch, fastvideo_args)
+            batch.extra["audio"] = waveform
+            batch.extra["audio_sample_rate"] = sample_rate
+        self._clear_runtime(batch)
+        return batch
 
+    @torch.no_grad()
+    def decode_audio(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> tuple[torch.Tensor, int]:
+        """Synthesize this request without changing its latents or stage state."""
         layout = _layout(batch)
         if batch.audio_latents is None:
             raise ValueError("MiniMax-H3 audio latents are missing at decode.")
@@ -246,10 +301,7 @@ class MiniMaxH3AudioDecodingStage(PipelineStage):
         try:
             latents = self.audio_vae.denormalize_latents(latents.to(device=device, dtype=torch.float32))
             if fastvideo_args.output_type == "latent":
-                batch.extra["audio"] = latents.detach().float().cpu()
-                batch.extra["audio_sample_rate"] = self.audio_vae.sampling_rate
-                self._clear_runtime(batch)
-                return batch
+                return latents.detach().float().cpu(), self.audio_vae.sampling_rate
 
             # The range isolates waveform synthesis from packing and runtime
             # cleanup so the audio decoder has one stable timeline boundary.
@@ -258,10 +310,7 @@ class MiniMaxH3AudioDecodingStage(PipelineStage):
             if decoded.ndim != 3 or decoded.shape[0] != 2 or decoded.shape[1] != 1:
                 raise ValueError("MiniMax-H3 audio VAE must decode stereo channels as two mono batch items; "
                                  f"got {tuple(decoded.shape)}.")
-            batch.extra["audio"] = decoded[:, 0].transpose(0, 1).contiguous().cpu()
-            batch.extra["audio_sample_rate"] = self.audio_vae.sampling_rate
-            self._clear_runtime(batch)
-            return batch
+            return decoded[:, 0].transpose(0, 1).contiguous().cpu(), self.audio_vae.sampling_rate
         finally:
             if fastvideo_args.vae_cpu_offload:
                 pinned_offload.unload(self.audio_vae)