[misc]: regenerate env docs for NVFP4 FA4 Q/K transform and BF16 block variables

fp4_dense_qat routes FlashAttention layers through the FP4 training kernel with every
key block listed (Q/K NVFP4, P and V BF16, straight-through gradients), matching the
FlashAttention-4 NVFP4 deployment; qad.fp4_attention enables it in the Flex profile.
check_fp4_attention_parity.py compares it with the real FA4 NVFP4 kernel on a GPU.

docs-site/islands/RecipeFamily.tsx

@@ -5,7 +5,7 @@ export default function RecipeFamily({ family }: { family: string }) {
   const models = modelDirectory.filter(
     (model) =>
       model.recipes.some((recipe) => recipe.family === family) ||
-      (family === "minimax_h3" && model.family === "FastH3"),
+      (family === "minimax_h3" && (model.family === "FastH3" || model.familyPath === "/cookbook/minimax-h3")),
   );
   return (
     <div className="recipe-directory not-prose">

docs-site/lib/cookbook-pages.ts

@@ -105,6 +105,8 @@ export interface SourceRecipe {
 export const sourceRecipes = sourceCatalog.recipes as SourceRecipe[];
 const familyInfo: Record<string, { title: string; path: string; logo: string; category: string }> = {
   minimax_h3: { title: "MiniMax H3", path: "minimax-h3", logo: "minimax", category: "Video" },
+  // A post-trained MiniMax-H3 (Pika adapters on H3 components), listed with the H3 family.
+  pika_h3_flex: { title: "MiniMax H3", path: "minimax-h3", logo: "pika", category: "Video" },
   wan: { title: "Wan", path: "wan", logo: "wan-ai", category: "Video" },
   ltx2: { title: "LTX", path: "ltx", logo: "ltx", category: "Video" },
   hunyuan: { title: "Hunyuan", path: "hunyuan", logo: "tencent-hunyuan", category: "Video" },
@@ -158,6 +160,7 @@ export const modelDirectory = [...new Set(sourceRecipes.map((recipe) => recipe.m
     const providers: Record<string, string> = {
       fastvideo: "FastVideo",
       minimax: "MiniMax",
+      pika: "Pika",
       "wan-ai": "Wan AI",
       ltx: "Lightricks",
       "tencent-hunyuan": "Tencent",

docs-site/navigation.json

@@ -174,6 +174,10 @@
             "label": "FastH3 4-Step V1",
             "href": "/cookbook/minimax-h3/fasth3-4step"
           },
+          {
+            "label": "Pika H3 Flex",
+            "href": "/cookbook/minimax-h3/models/pika-h3-flex-t2va"
+          },
           {
             "label": "Recipes and setup",
             "href": "/cookbook/minimax-h3"
@@ -310,6 +314,10 @@
         "label": "Attn-QAT Training",
         "href": "/training/attn_qat"
       },
+      {
+        "label": "FastH3 NVFP4 QAD",
+        "href": "/training/minimax_h3_fp4_qad"
+      },
       {
         "label": "HunyuanVideo 1.5 T2V",
         "href": "/training/hunyuan15_overfit"

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(38);
   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(38);
 });
 
 test("cookbook cards and controls remain readable in dark mode and reflow at 320px", async ({ page }) => {
@@ -395,6 +395,35 @@ test("source-backed controls update the copied command without changing models",
   await expect(page.getByLabel("Inference command", { exact: true })).toContainText("--num-gpus '2'");
 });
 
+test("Pika H3 Flex knobs select the tier, bucket, duration, task, and NVFP4 decoder", async ({ page }) => {
+  await page.goto("/FastVideo/cookbook/minimax-h3");
+  const pika = page.locator(".recipe-card").filter({ hasText: "Pika H3 Flex" });
+  await expect(pika.locator("img")).toHaveAttribute("src", /\/pika\.webp$/);
+  await expect(pika).toContainText("Pika");
+  await pika.click();
+  await expect(page).toHaveURL(/cookbook\/minimax-h3\/models\/pika-h3-flex-t2va$/);
+  const command = page.getByLabel("Inference command", { exact: true });
+  await expect(command).toContainText("FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN FASTVIDEO_FA4=1");
+  await chip(page, "Resolution", "480p");
+  await chip(page, "Aspect ratio", "9:16");
+  await chip(page, "Duration", "10 s (243 frames)");
+  await chip(page, "VAE decoder", "NVFP4 36L quality");
+  await expect(command).toContainText("--resolution '480p'");
+  await expect(command).toContainText("--aspect '9:16'");
+  await expect(command).toContainText("--duration '10'");
+  await expect(command).toContainText("--nvfp4-decoder './fasth3-nvfp4-vae/quality-36L/decoder-nvfp4.safetensors'");
+  await page
+    .getByRole("group", { name: "Recipe variant", exact: true })
+    .getByRole("button", { name: "Pika H3 Flex keyframes", exact: true })
+    .click();
+  await expect(command).toContainText("--task fl2va --offload resident");
+  await chip(page, "Task", "First frame (I2VA)");
+  await chip(page, "Memory", "Lazy load + offload");
+  await expect(command).toContainText("--task 'i2va'");
+  await expect(command).toContainText("--offload 'lazy'");
+  await expect(page.getByRole("button", { name: "Copy example + command", exact: true })).toBeEnabled();
+});
+
 test("Wan I2V requires an image for CLI and SDK examples without changing server defaults", async ({ page }) => {
   await page.goto("/FastVideo/cookbook/wan/wan21-i2v");
   await chip(page, "Run mode", "CLI");

docs/assets/cookbook-recipes.json

@@ -1,5 +1,5 @@
 {
-  "version": 19,
+  "version": 20,
   "recipes": [
     {
       "id": "fastwan21-t2v",
@@ -760,6 +760,329 @@
       ],
       "limitations": ["Pass --reference-audio for an additional audio reference. The checked-in source defaults to four GPUs."]
     },
+    {
+      "id": "pika-h3-flex-t2va",
+      "family": "pika_h3_flex",
+      "stage": "inference",
+      "task": "Text to video (with audio)",
+      "label": "Pika H3 Flex",
+      "summary": "Pika's post-trained MiniMax-H3: synchronized video and stereo audio from text in six forwards, at 480p or 768p, 5-15 s, with optional NVFP4 decoder and text encoder on Blackwell.",
+      "model": "Pika-Labs/Pika-H3-Flex",
+      "source": "examples/inference/basic/basic_pika_h3_flex.py",
+      "command": "FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN FASTVIDEO_FA4=1 python examples/inference/basic/basic_pika_h3_flex.py --task t2va --offload resident --prompt 'A paper boat drifts down a rainy city street at night; rain patters on the water.'",
+      "gpu_types": ["NVIDIA"],
+      "hardware": {
+        "platform": "cuda",
+        "gpu_count": 1,
+        "accelerator": "NVIDIA GB200",
+        "peak_memory": "About 129 GiB resident (BF16, 768p, 5 s)",
+        "evidence": "validated",
+        "evidence_url": "https://github.com/hao-ai-lab/FastVideo/pull/2038"
+      },
+      "evidence": "Verified",
+      "expected_artifact": "MP4 with synchronized stereo audio at outputs/pika_h3_flex/pika_h3_flex_t2va.mp4",
+      "modes": ["T2VA"],
+      "knobs": [
+        {
+          "key": "resolution",
+          "label": "Resolution",
+          "hint": "Pika resolution tier",
+          "flag": "--resolution",
+          "options": [{"value": "768p", "label": "768p"}, {"value": "480p", "label": "480p"}],
+          "default": "768p"
+        },
+        {
+          "key": "aspect",
+          "label": "Aspect ratio",
+          "hint": "Bucket within the tier",
+          "flag": "--aspect",
+          "options": ["21:9", "16:9", "3:2", "4:3", "5:4", "1:1", "4:5", "3:4", "2:3", "9:16", "9:21"],
+          "default": "16:9"
+        },
+        {
+          "key": "duration",
+          "label": "Duration",
+          "hint": "Rounded up to the 17n+5 frame grid",
+          "flag": "--duration",
+          "options": [
+            {"value": 5, "label": "5 s (124 frames)"},
+            {"value": 10, "label": "10 s (243 frames)"},
+            {"value": 15, "label": "15 s (362 frames)"}
+          ],
+          "default": 5
+        },
+        {
+          "key": "nvfp4_decoder",
+          "label": "VAE decoder",
+          "hint": "NVFP4 needs Blackwell",
+          "flag": "--nvfp4-decoder",
+          "options": [
+            {"value": "", "label": "BF16 H3 VAE"},
+            {"value": "./fasth3-nvfp4-vae/quality-36L/decoder-nvfp4.safetensors", "label": "NVFP4 36L quality"},
+            {"value": "./fasth3-nvfp4-vae/fast-8L/decoder-nvfp4.safetensors", "label": "NVFP4 8L fast"}
+          ],
+          "default": ""
+        },
+        {
+          "key": "nvfp4_text_encoder",
+          "label": "Text encoder",
+          "hint": "NVFP4 needs Blackwell",
+          "flag": "--nvfp4-text-encoder",
+          "options": [
+            {"value": "", "label": "BF16 Qwen3-VL"},
+            {"value": "./fasth3-nvfp4-encoder/text_encoder", "label": "NVFP4 Qwen3-VL"}
+          ],
+          "default": ""
+        },
+        {
+          "key": "offload",
+          "label": "Memory",
+          "hint": "Resident is how the GB200 timings were measured",
+          "flag": "--offload",
+          "options": [
+            {"value": "resident", "label": "All resident (fastest)"},
+            {"value": "lazy", "label": "Lazy load + offload"}
+          ],
+          "default": "resident"
+        },
+        {
+          "key": "num_gpus",
+          "label": "GPUs",
+          "hint": "Sequence-parallel degree",
+          "flag": "--num-gpus",
+          "options": [1, 2, 4, 8],
+          "default": 1
+        }
+      ],
+      "limitations": [
+        "Measured on one GB200 at 768p 16:9, 5 s, warm: official pika-h3-pro CLI 59.7 s; FastVideo BF16 26.2 s; with the NVFP4 quality-36L decoder 18.0 s; with the fast-8L decoder 17.4 s. The 36L decoder measured 40.5 dB PSNR against the BF16 decoder.",
+        "Peak resident memory was about 129 GiB for the all-BF16 run and 104.6 GiB with the NVFP4 text encoder (25.49 s) at 768p / 5 s on GB200; the official CLI's sampling peak was 71.7 GiB, a different measurement. The NVFP4 text encoder's first request in a fresh process pays a one-time cost of about 183 s.",
+        "--offload resident keeps every component on the GPU (104.6-108.5 GiB peak with the NVFP4 text encoder, about 129 GiB all-BF16, on GB200). --offload lazy loads components on demand for smaller devices; on GB200 it took 48.06 s for the same T2VA request.",
+        "FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN FASTVIDEO_FA4=1 selects FlashAttention-4 on Hopper or Blackwell; FA4 is never enabled just because it is installed. Without FA4, attention uses FA3/FA2 if installed and otherwise falls back to SDPA, which is slow and can run out of memory at 768p.",
+        "The NVFP4 decoder and text encoder are Blackwell-only. Download them before selecting them: hf download FastVideo/FastH3-NVFP4-VAE-Decoder --include 'quality-36L/decoder-nvfp4.safetensors' 'fast-8L/decoder-nvfp4.safetensors' --local-dir fasth3-nvfp4-vae; hf download FastVideo/FastVideo-FastH3-8-Step-V2-NVFP4-Consumer --include 'text_encoder/*' --local-dir fasth3-nvfp4-encoder. The 8L decoder is a distilled student; validate it on your prompts.",
+        "The six-forward re-noise ladder is part of the adapter contract; the example exposes no step-count flag."
+      ],
+      "related": ["pika-h3-flex-keyframes", "pika-h3-flex-ref2va"]
+    },
+    {
+      "id": "pika-h3-flex-keyframes",
+      "family": "pika_h3_flex",
+      "stage": "inference",
+      "task": "First/last frame to video (with audio)",
+      "label": "Pika H3 Flex keyframes",
+      "summary": "Animate from a first frame (I2VA), toward a last frame (L2VA), or between both (FL2VA). Pika center-crops each keyframe to the selected bucket.",
+      "model": "Pika-Labs/Pika-H3-Flex",
+      "source": "examples/inference/basic/basic_pika_h3_flex.py",
+      "command": "FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN FASTVIDEO_FA4=1 python examples/inference/basic/basic_pika_h3_flex.py --task fl2va --offload resident --image assets/images/natural.jpg --last-image assets/images/peninsula.jpg --prompt 'The camera drifts from the misty mountain ridge down to the coastal cliffs above the sea.'",
+      "gpu_types": ["NVIDIA"],
+      "hardware": {
+        "platform": "cuda",
+        "gpu_count": 1,
+        "accelerator": "NVIDIA GB200",
+        "peak_memory": "105.2 GiB FL2VA, 104.9 GiB I2VA (NVFP4 text encoder, 768p, 5 s)",
+        "evidence": "validated",
+        "evidence_url": "https://github.com/hao-ai-lab/FastVideo/pull/2038"
+      },
+      "evidence": "Verified",
+      "expected_artifact": "MP4 with synchronized stereo audio at outputs/pika_h3_flex/pika_h3_flex_<task>.mp4",
+      "modes": ["I2VA", "L2VA", "FL2VA"],
+      "knobs": [
+        {
+          "key": "task",
+          "label": "Task",
+          "hint": "Which keyframes condition the clip",
+          "flag": "--task",
+          "options": [
+            {"value": "i2va", "label": "First frame (I2VA)"},
+            {"value": "l2va", "label": "Last frame (L2VA)"},
+            {"value": "fl2va", "label": "First and last (FL2VA)"}
+          ],
+          "default": "fl2va"
+        },
+        {
+          "key": "resolution",
+          "label": "Resolution",
+          "hint": "Pika resolution tier",
+          "flag": "--resolution",
+          "options": [{"value": "768p", "label": "768p"}, {"value": "480p", "label": "480p"}],
+          "default": "768p"
+        },
+        {
+          "key": "aspect",
+          "label": "Aspect ratio",
+          "hint": "Bucket within the tier",
+          "flag": "--aspect",
+          "options": ["21:9", "16:9", "3:2", "4:3", "5:4", "1:1", "4:5", "3:4", "2:3", "9:16", "9:21"],
+          "default": "16:9"
+        },
+        {
+          "key": "duration",
+          "label": "Duration",
+          "hint": "Rounded up to the 17n+5 frame grid",
+          "flag": "--duration",
+          "options": [
+            {"value": 5, "label": "5 s (124 frames)"},
+            {"value": 10, "label": "10 s (243 frames)"},
+            {"value": 15, "label": "15 s (362 frames)"}
+          ],
+          "default": 5
+        },
+        {
+          "key": "nvfp4_decoder",
+          "label": "VAE decoder",
+          "hint": "NVFP4 needs Blackwell",
+          "flag": "--nvfp4-decoder",
+          "options": [
+            {"value": "", "label": "BF16 H3 VAE"},
+            {"value": "./fasth3-nvfp4-vae/quality-36L/decoder-nvfp4.safetensors", "label": "NVFP4 36L quality"},
+            {"value": "./fasth3-nvfp4-vae/fast-8L/decoder-nvfp4.safetensors", "label": "NVFP4 8L fast"}
+          ],
+          "default": ""
+        },
+        {
+          "key": "nvfp4_text_encoder",
+          "label": "Text encoder",
+          "hint": "NVFP4 needs Blackwell",
+          "flag": "--nvfp4-text-encoder",
+          "options": [
+            {"value": "", "label": "BF16 Qwen3-VL"},
+            {"value": "./fasth3-nvfp4-encoder/text_encoder", "label": "NVFP4 Qwen3-VL"}
+          ],
+          "default": ""
+        },
+        {
+          "key": "offload",
+          "label": "Memory",
+          "hint": "Resident is how the GB200 timings were measured",
+          "flag": "--offload",
+          "options": [
+            {"value": "resident", "label": "All resident (fastest)"},
+            {"value": "lazy", "label": "Lazy load + offload"}
+          ],
+          "default": "resident"
+        },
+        {
+          "key": "num_gpus",
+          "label": "GPUs",
+          "hint": "Sequence-parallel degree",
+          "flag": "--num-gpus",
+          "options": [1, 2, 4, 8],
+          "default": 1
+        }
+      ],
+      "limitations": [
+        "Measured on one GB200 at 768p, 5 s, seed 42, NVFP4 text encoder, all resident: FL2VA 28.95 s (first frame 39.6 dB / SSIM 0.971 and last frame 29.1 dB / 0.819 against the inputs); I2VA 27.37 s (first frame 26.5 dB / 0.715). L2VA was not measured separately.",
+        "The command passes both keyframes; the example uses only the ones the selected task needs. Replace them with your own images.",
+        "--offload resident keeps every component on the GPU (104.6-108.5 GiB peak with the NVFP4 text encoder, about 129 GiB all-BF16, on GB200). --offload lazy loads components on demand for smaller devices; on GB200 it took 48.06 s for the same T2VA request.",
+        "FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN FASTVIDEO_FA4=1 selects FlashAttention-4 on Hopper or Blackwell; FA4 is never enabled just because it is installed. Without FA4, attention uses FA3/FA2 if installed and otherwise falls back to SDPA, which is slow and can run out of memory at 768p.",
+        "The NVFP4 decoder and text encoder are Blackwell-only. Download them before selecting them: hf download FastVideo/FastH3-NVFP4-VAE-Decoder --include 'quality-36L/decoder-nvfp4.safetensors' 'fast-8L/decoder-nvfp4.safetensors' --local-dir fasth3-nvfp4-vae; hf download FastVideo/FastVideo-FastH3-8-Step-V2-NVFP4-Consumer --include 'text_encoder/*' --local-dir fasth3-nvfp4-encoder. The 8L decoder is a distilled student; validate it on your prompts."
+      ],
+      "related": ["pika-h3-flex-t2va", "pika-h3-flex-ref2va"]
+    },
+    {
+      "id": "pika-h3-flex-ref2va",
+      "family": "pika_h3_flex",
+      "stage": "inference",
+      "task": "Reference media to video (with audio)",
+      "label": "Pika H3 Flex references",
+      "summary": "Generate from ordered image, video, and audio references through Pika's ref2va adapter: up to 9 images, 3 videos, and 3 audio clips.",
+      "model": "Pika-Labs/Pika-H3-Flex",
+      "source": "examples/inference/basic/basic_pika_h3_flex.py",
+      "command": "FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN FASTVIDEO_FA4=1 python examples/inference/basic/basic_pika_h3_flex.py --task ref2va --offload resident --reference image:assets/images/natural.jpg --reference image:assets/images/bus_terminal.jpg --prompt 'Start on the misty mountain ridge from the first reference, then cut to the night bus terminal from the second.'",
+      "gpu_types": ["NVIDIA"],
+      "hardware": {
+        "platform": "cuda",
+        "gpu_count": 1,
+        "accelerator": "NVIDIA GB200",
+        "peak_memory": "108.5 GiB with two image references (NVFP4 text encoder, 768p, 5 s)",
+        "evidence": "validated",
+        "evidence_url": "https://github.com/hao-ai-lab/FastVideo/pull/2038"
+      },
+      "evidence": "Verified",
+      "expected_artifact": "MP4 with synchronized stereo audio at outputs/pika_h3_flex/pika_h3_flex_ref2va.mp4",
+      "modes": ["Ref2VA"],
+      "knobs": [
+        {
+          "key": "resolution",
+          "label": "Resolution",
+          "hint": "Pika resolution tier",
+          "flag": "--resolution",
+          "options": [{"value": "768p", "label": "768p"}, {"value": "480p", "label": "480p"}],
+          "default": "768p"
+        },
+        {
+          "key": "aspect",
+          "label": "Aspect ratio",
+          "hint": "Bucket within the tier",
+          "flag": "--aspect",
+          "options": ["21:9", "16:9", "3:2", "4:3", "5:4", "1:1", "4:5", "3:4", "2:3", "9:16", "9:21"],
+          "default": "16:9"
+        },
+        {
+          "key": "duration",
+          "label": "Duration",
+          "hint": "Rounded up to the 17n+5 frame grid",
+          "flag": "--duration",
+          "options": [
+            {"value": 5, "label": "5 s (124 frames)"},
+            {"value": 10, "label": "10 s (243 frames)"},
+            {"value": 15, "label": "15 s (362 frames)"}
+          ],
+          "default": 5
+        },
+        {
+          "key": "nvfp4_decoder",
+          "label": "VAE decoder",
+          "hint": "NVFP4 needs Blackwell",
+          "flag": "--nvfp4-decoder",
+          "options": [
+            {"value": "", "label": "BF16 H3 VAE"},
+            {"value": "./fasth3-nvfp4-vae/quality-36L/decoder-nvfp4.safetensors", "label": "NVFP4 36L quality"},
+            {"value": "./fasth3-nvfp4-vae/fast-8L/decoder-nvfp4.safetensors", "label": "NVFP4 8L fast"}
+          ],
+          "default": ""
+        },
+        {
+          "key": "nvfp4_text_encoder",
+          "label": "Text encoder",
+          "hint": "NVFP4 needs Blackwell",
+          "flag": "--nvfp4-text-encoder",
+          "options": [
+            {"value": "", "label": "BF16 Qwen3-VL"},
+            {"value": "./fasth3-nvfp4-encoder/text_encoder", "label": "NVFP4 Qwen3-VL"}
+          ],
+          "default": ""
+        },
+        {
+          "key": "offload",
+          "label": "Memory",
+          "hint": "Resident is how the GB200 timings were measured",
+          "flag": "--offload",
+          "options": [
+            {"value": "resident", "label": "All resident (fastest)"},
+            {"value": "lazy", "label": "Lazy load + offload"}
+          ],
+          "default": "resident"
+        },
+        {
+          "key": "num_gpus",
+          "label": "GPUs",
+          "hint": "Sequence-parallel degree",
+          "flag": "--num-gpus",
+          "options": [1, 2, 4, 8],
+          "default": 1
+        }
+      ],
+      "limitations": [
+        "Measured on one GB200 at 768p, 5 s, seed 42, NVFP4 text encoder, all resident, two image references: 51.17 s (38.42 s denoise, 64,241 packed tokens); both references appear in the output.",
+        "Repeat --reference image:PATH, video:PATH, or audio:PATH in order. At least one image or video is required; a video reference needs at least 22 frames at 24 fps.",
+        "Ref2VA loads transformer_ref and the ref2va adapter, a separate 61.7 GiB trunk download. Reference images are encoded at up to a 1024 (480p) or 2048 (768p) short edge, so conditioning memory is higher than T2VA's.",
+        "--offload resident keeps every component on the GPU (104.6-108.5 GiB peak with the NVFP4 text encoder, about 129 GiB all-BF16, on GB200). --offload lazy loads components on demand for smaller devices; on GB200 it took 48.06 s for the same T2VA request.",
+        "FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN FASTVIDEO_FA4=1 selects FlashAttention-4 on Hopper or Blackwell; FA4 is never enabled just because it is installed. Without FA4, attention uses FA3/FA2 if installed and otherwise falls back to SDPA, which is slow and can run out of memory at 768p.",
+        "The NVFP4 decoder and text encoder are Blackwell-only. Download them before selecting them: hf download FastVideo/FastH3-NVFP4-VAE-Decoder --include 'quality-36L/decoder-nvfp4.safetensors' 'fast-8L/decoder-nvfp4.safetensors' --local-dir fasth3-nvfp4-vae; hf download FastVideo/FastVideo-FastH3-8-Step-V2-NVFP4-Consumer --include 'text_encoder/*' --local-dir fasth3-nvfp4-encoder. The 8L decoder is a distilled student; validate it on your prompts."
+      ],
+      "related": ["pika-h3-flex-t2va", "pika-h3-flex-keyframes"]
+    },
     {
       "id": "fasth3-lora-preview",
       "family": "minimax_h3",

docs/assets/logos/SOURCES.md

@@ -15,6 +15,7 @@ plain typographic tile instead — that tile is a UI placeholder, not a logo.
 | `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 |
+| `pika.webp` | [Official Pika-Labs Hugging Face organization avatar](https://huggingface.co/Pika-Labs) | Pika H3 Flex 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/env_vars.md

@@ -177,6 +177,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry.
 | `FASTVIDEO_CACHE_ROOT`                                     | path                                                               | computed                                              | path        | Root directory for FastVideo cache files. Defaults to ~/.cache/fastvideo, or $XDG_CACHE_HOME/fastvideo when XDG_CACHE_HOME is set.                                                                                                                                                                                                                                                                                                          |
 | `FASTVIDEO_REASON1_WEIGHTS_PATH`                           | str                                                                | unset                                                 | path        | Local path or Hugging Face id of Reason1 weights to load instead of the checkpoint's own.                                                                                                                                                                                                                                                                                                                                                   |
 | `FASTVIDEO_HOST_IP`                                        | str                                                                | `""`                                                  | distributed | IP address of this node when the node has several network interfaces. Set it on each node for multi-node inference.                                                                                                                                                                                                                                                                                                                         |
+| `FASTVIDEO_PG_TIMEOUT_MIN`                                 | int                                                                | `0`                                                   | distributed | Timeout in minutes for the training process groups' collectives (0 keeps PyTorch's 10-minute default). Raise it for large models whose rank-0 load and scatter, checkpoint saves or long evaluations outlast the default.                                                                                                                                                                                                                   |
 | `FASTVIDEO_LOOPBACK_IP`                                    | str                                                                | `""`                                                  | distributed | Loopback IP address to use instead of the detected one.                                                                                                                                                                                                                                                                                                                                                                                     |
 | `FASTVIDEO_RAY_PER_WORKER_GPUS`                            | float                                                              | `1.0`                                                 | distributed | GPUs per Ray worker. A fraction lets Ray schedule several actors on one GPU, so other actors can share the GPUs with FastVideo.                                                                                                                                                                                                                                                                                                             |
 | `FASTVIDEO_NCCL_SO_PATH`                                   | str                                                                | unset                                                 | distributed | Path to the NCCL library file. Needed because the nccl>=2.19 that PyTorch ships has a bug (https://github.com/NVIDIA/nccl/issues/1234).                                                                                                                                                                                                                                                                                                  |
@@ -197,6 +198,8 @@ longer exists also fails the test, so the fixing pull request deletes its entry.
 | `FASTVIDEO_VSA_PTX_BLACKWELL`                              | bool                                                               | `0`                                                   | attention   | VIDEO_SPARSE_ATTN_H3 sends tile-64 forwards to the data-center Blackwell (sm_100a / sm_103a) kernel; grad-tracking forwards need its autograd operator, else Triton. fastvideo-kernel reads the same variable with the same rule. Deprecated names: `FASTVIDEO_VSA_SM100A`.                                                                                                                                                                 |
 | `FASTVIDEO_VSA_TRITON`                                     | bool                                                               | `0`                                                   | attention   | Force the Triton MiniMax-H3 sparse attention kernel. fastvideo-kernel reads the same variable.                                                                                                                                                                                                                                                                                                                                              |
 | `FASTVIDEO_NVFP4_FA4`                                      | bool                                                               | `0`                                                   | attention   | FlashAttention-4 quantizes Q and K to NVFP4. An explicit nvfp4_fa4 attention implementation argument takes precedence.                                                                                                                                                                                                                                                                                                                      |
+| `FASTVIDEO_NVFP4_FA4_QK_TRANSFORM`                         | one of none, smooth_k, hadamard, smooth_k+hadamard                 | `none`                                                | attention   | Exact Q/K preconditioning applied right before NVFP4 quantization on the FlashAttention-4 NVFP4 path: smooth_k subtracts each head's token-mean of K, hadamard rotates Q and K by the head-dim Hadamard matrix. Must match the qk_transform the model was QAD-trained with. An explicit nvfp4_qk_transform attention implementation argument takes precedence.                                                                              |
+| `FASTVIDEO_NVFP4_FA4_BF16_BLOCKS`                          | str                                                                | `""`                                                  | attention   | Comma-separated MiniMax-H3 transformer_blocks indices whose attention stays in BF16 while FASTVIDEO_NVFP4_FA4 quantizes the rest, for example 45,49.                                                                                                                                                                                                                                                                                        |
 | `FASTVIDEO_FA4_PV_MODE`                                    | one of bf16, fp8                                                   | `bf16`                                                | attention   | V dtype on the NVFP4 FlashAttention-4 path (FLASH_ATTN with nvfp4_fa4, ATTN_QAT_INFER on sm_100a/sm_103a): bf16 keeps V in BF16; fp8 casts V to float8 e4m3 without scaling. An explicit fa4_pv_mode attention implementation argument takes precedence.                                                                                                                                                                                    |
 | `FASTVIDEO_DISABLE_ATTENTION_COMPILE`                      | bool                                                               | `1`                                                   | attention   | Keep attention forward out of torch.compile graphs (torch.compiler.disable). Set it to 0 to let attention constructed under that setting be traced. Setting it explicitly to true also blocks regional compile.                                                                                                                                                                                                                             |
 | `FASTVIDEO_MLX_WINDOW`                                     | int                                                                | `0`                                                   | attention   | MLX FastWan windowed attention size in tokens. 0 uses full attention.                                                                                                                                                                                                                                                                                                                                                                       |

docs/cookbook/minimax-h3.md

@@ -72,6 +72,32 @@
     </div>
   </details>
 
+<details  id="pika-h3-flex">
+  <summary>Pika H3 Flex</summary>
+  <div >
+      <p>Pika H3 Flex is Pika's post-trained MiniMax-H3: the H3 components plus a rank-128 adapter that FastVideo applies automatically, sampled with a fixed six-forward re-noise ladder. One recipe page covers T2VA, the keyframe tasks (I2VA, L2VA, FL2VA), and Ref2VA. Its builder selects the resolution tier (480p or 768p), the aspect bucket, the duration (5, 10, or 15 s), and the optional Blackwell NVFP4 decoder and text encoder. The <a href="../../inference/pika_h3_flex/">Pika H3 Flex runbook</a> has the full sampling and geometry contract.</p>
+      <pre><code>FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN FASTVIDEO_FA4=1 \
+python examples/inference/basic/basic_pika_h3_flex.py --task t2va --offload resident \
+    --resolution 768p --aspect 16:9 --duration 5 \
+    --prompt "A paper boat drifts down a rainy city street at night."</code></pre>
+      <table>
+        <thead>
+          <tr>
+            <th>T2VA on one GB200, 768p 16:9, 5 s, warm, all resident</th>
+            <th>Time</th>
+          </tr>
+        </thead>
+        <tbody>
+          <tr><td>Official <code>pika-h3-pro</code> CLI</td><td>59.7 s</td></tr>
+          <tr><td>FastVideo BF16</td><td>26.2 s</td></tr>
+          <tr><td>+ NVFP4 <code>quality-36L</code> decoder (40.5 dB PSNR vs BF16 decoder)</td><td>18.0 s</td></tr>
+          <tr><td>+ NVFP4 <code>fast-8L</code> decoder</td><td>17.4 s</td></tr>
+        </tbody>
+      </table>
+      <p>The BF16 run peaked at about 129 GiB resident on GB200; the official CLI's sampling peak was 71.7 GiB. With the NVFP4 text encoder, I2VA took 27.37 s, FL2VA 28.95 s, and Ref2VA with two image references 51.17 s, at 104.9 to 108.5 GiB peak. The example's default <code>--offload lazy</code> fits smaller devices but reloads components per request (48.06 s for the same T2VA request on GB200). FlashAttention-4 needs <code>FASTVIDEO_FA4=1</code> even when it is installed. Without it, attention uses FA3 or FA2 if one is installed and otherwise falls back to SDPA, which is slow and can run out of memory at 768p. GB10 has no FA4, so drop the prefix on DGX Spark.</p>
+  </div>
+</details>
+
 <details  id="cookbook-setup">
   <summary>Setup</summary>
   <div >

docs/design/inference_schema_parity_inventory.yaml

@@ -88,7 +88,7 @@ surfaces:
       video_decode_backend: "MiniMax-H3 video decoder selection (full VAE vs TAEH3 preview); model-specific optimization not yet represented in the typed public schema."
       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_decoder: "Exported NVFP4 MiniMax-H3 video decoder file that replaces the VAE decoder at load; represented as ComponentConfig.h3_vae_nvfp4_decoder."
       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."
@@ -452,6 +452,9 @@ surfaces:
       kv_cache_device:
         target: generator.pipeline.experimental.kv_cache_device
         sources: [fastvideo.configs.pipelines.qwen_image21.QwenImage21PipelineConfig]
+      pika_resolution_tier:
+        target: generator.pipeline.experimental.pika_resolution_tier
+        sources: [fastvideo.configs.pipelines.pika_h3_flex.PikaH3FlexPipelineConfig]
     compatibility_only:
       batch_size: "Gen3C inference-only tuning field pending typed batching design."
       gradient_checkpointing: "Gen3C inference-only compatibility field pending typed batching design."
@@ -465,6 +468,7 @@ surfaces:
       frame_receptive_field: "MagiHuman internal data-proxy receptive-field setting."
       image_conditioning: "MagiHuman preset variant marker for reference-image conditioning."
       pdd_step_indices: "MiniMax-H3 PDD fine-grid partition; resolve_checkpoint_settings sets it from fastvideo_inference.json."
+      pika_denoising_steps: "Pika H3 Flex six-forward re-noise ladder fixed by the adapter contract; check_pipeline_config rejects any other value."
       ref_audio_offset: "MagiHuman internal data-proxy audio alignment offset."
       sr_local_attn_layers: "MagiHuman SR internal sparse-attention layer selection."
       text_offset: "MagiHuman internal data-proxy text alignment offset."

docs/generate_examples.py

@@ -39,6 +39,7 @@ COOKBOOK_FAMILIES = {
     "qwen_image21",
     "sd35",
     "minimax_h3",
+    "pika_h3_flex",
     "stable_audio",
     "mmaudio",
     "matrixgame",

docs/inference/pika_h3_flex.md

@@ -0,0 +1,237 @@
+# Pika H3 Flex
+
+Pika H3 Flex is a post-trained MiniMax-H3 model that keeps H3's packed
+video/audio transformer, Qwen3-VL conditioner, and video/audio VAEs. Pika ships
+two rank-128 PEFT adapters instead of a second base checkpoint:
+
+| Adapter | Trunk | Tasks |
+|---|---|---|
+| `fl2va/adapter.safetensors` | `transformer` | T2VA, I2VA, L2VA, FL2VA |
+| `ref2va/adapter.safetensors` | `transformer_ref` | Ref2VA |
+
+FastVideo detects `Pika-Labs/Pika-H3-Flex`, downloads only the selected task's
+adapter, and applies it as an ordinary LoRA. The released adapter is selected
+automatically; pass `lora_path` only when overriding it with a compatible
+Pika adapter.
+
+## Sampling contract
+
+Pika's sampler is fixed, not a tunable step count:
+
+* 6 transformer forwards over clean-time rungs
+  `1000, 750, 500, 250, 100, 40, 0`;
+* video shift `12.0`, audio shift `3.0`;
+* no CFG, `guidance_scale=1.0`, empty negative prompt;
+* each step re-noises the predicted clean sample with fresh Gaussian noise from
+  a generator seeded `seed + 1`.
+
+`num_inference_steps` must therefore be `6`; other values are rejected. This is
+not a PDD checkpoint and it has no trained VSA gates, so PDD fused blocks and
+VSA sparsity are not used.
+
+## Tasks
+
+The commands below assume a FlashAttention-4 host; see
+[Attention backend](#attention-backend) before running them elsewhere.
+
+```bash
+export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN FASTVIDEO_FA4=1
+
+# text to video + audio
+python examples/inference/basic/basic_pika_h3_flex.py --task t2va \
+    --prompt "a paper boat drifting down a rainy street"
+
+# 480p portrait, ten seconds
+python examples/inference/basic/basic_pika_h3_flex.py --task t2va \
+    --resolution 480p --aspect 9:16 --duration 10 \
+    --prompt "a lighthouse at dusk, waves breaking below"
+
+# first frame, last frame, or both
+python examples/inference/basic/basic_pika_h3_flex.py --task i2va \
+    --image first.png --prompt "the camera slowly pushes in"
+python examples/inference/basic/basic_pika_h3_flex.py --task fl2va \
+    --image first.png --last-image last.png --prompt "the flower opens"
+
+# ordered references, including reference video and audio
+python examples/inference/basic/basic_pika_h3_flex.py --task ref2va \
+    --reference image:person.png --reference video:dance.mp4 \
+    --reference audio:voice.wav --prompt "the person dances to the music"
+```
+
+Ref2VA accepts up to 12 references: 9 images, 3 videos, and 3 audio. These
+counts are enforced locally, as is the rule that audio-only requests are
+invalid. A video reference needs at least 22 decoded frames at 24 fps. At least
+one image or video is required, and the Ref2VA pipeline must be selected with
+`PikaH3FlexRef2VAModularPipeline` (the example does this).
+
+Pika's hosted API also documents a 2–15 s limit per video/audio reference and a
+15 s total per modality. **The local path does not currently enforce those
+duration caps**; it trims each reference to the requested output duration. Treat
+the per-reference and per-modality limits as recommendations until the checks
+are added.
+
+## Canvas and frames
+
+Pika uses an exact 32-pixel bucket lattice. A request that matches a bucket is
+kept exactly; otherwise FastVideo picks the tier from the requested short edge
+(`<= 480` → p480, else p768) and snaps to the nearest aspect bucket in that tier.
+
+| Tier | Short edge | Max pixels | Reference-image short edge |
+|---|---|---|---|
+| p480 | 480 | 832×480 | 1024 |
+| p768 | 768 | 1344×768 | 2048 |
+
+Frame counts follow Pika's `17n + 5` grid at 24 fps (5–15 s, 120–362 aligned
+frames). The default 16:9 p768 canvas is 1344×768 and 124 frames.
+
+| Duration | 5 s | 10 s | 15 s |
+|---|---|---|---|
+| `num_frames` | 124 | 243 | 362 |
+
+The example script picks the bucket with `--resolution {480p,768p}` and
+`--aspect` (`21:9`, `16:9`, `3:2`, `4:3`, `5:4`, `1:1`, `4:5`, `3:4`, `2:3`,
+`9:16`, `9:21`), and the frame count with `--duration` seconds. `--height`,
+`--width`, and `--num-frames` override them.
+
+Through the API, a request that omits `height` and `width` uses the 16:9
+bucket of the configured tier:
+
+```python
+PipelineSelection(experimental={"pika_resolution_tier": "p480"})  # 832x480 when the size is omitted
+```
+
+An explicit `height`/`width` always wins and selects its own tier. "Omitted"
+means absent from a request parsed from a mapping or config file, or from a
+`generate_video(...)` call without `height`/`width`; a `GenerationRequest`
+constructed in Python counts `SamplingConfig`'s 1280×720 default as set, which
+snaps to the 768p 16:9 bucket. Pass the size explicitly in that case.
+
+Pika center-crops both keyframes to the bucket. MiniMax-H3 stretches its first
+frame instead; do not reuse that fit for Pika.
+
+## Attention backend
+
+Run Pika with FlashAttention-4 on Hopper or Blackwell:
+
+```bash
+FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN FASTVIDEO_FA4=1 \
+    python examples/inference/basic/basic_pika_h3_flex.py --task t2va --prompt "..."
+```
+
+FA4 is opt-in. Installing `flash_attn.cute` is not enough:
+`fastvideo/attention/utils/flash_attn_default.py` only selects it when
+`FASTVIDEO_FA4=1`, and neither the example's defaults nor `--lossless-accel`
+set that variable. With `FASTVIDEO_FA4=1` and an unusable FA4 install the run
+fails at import instead of silently falling back.
+
+Without FA4, the `FLASH_ATTN` backend uses FA3 or FA2 if one is installed. On a
+host with neither (common on Blackwell), backend selection falls back to
+`TORCH_SDPA`. SDPA at 768p is much slower and can run out of memory, so do not
+use it for 768p benchmarks or production runs.
+
+## Performance and memory
+
+All measurements: one GB200, 768p 16:9 (1344×768), 5 s (124 frames), six
+forwards, warm requests, FA4, every component resident (`--offload resident`).
+
+| Configuration (T2VA) | Generation time |
+|---|---|
+| Official `pika-h3-pro` CLI | 59.7 s |
+| FastVideo, BF16 | 26.2 s |
+| FastVideo, NVFP4 `quality-36L` decoder | 18.0 s |
+| FastVideo, NVFP4 `fast-8L` decoder | 17.4 s |
+
+The 36L decoder measured 40.5 dB PSNR against the BF16 decoder on the same
+latents. The 8L decoder was timed but not scored here.
+
+Task coverage, seed 42, NVFP4 text encoder, BF16 decoder, also with
+`FASTVIDEO_MINIMAX_H3_FUSIONS=all`:
+
+| Task | Wall | Denoise | Decode | Peak memory | Check |
+|---|---|---|---|---|---|
+| T2VA | 25.49 s | 16.70 s | 8.04 s | 104.6 GiB | on prompt; 21.7 dB PSNR vs the BF16-encoder output |
+| I2VA | 27.37 s | 18.02 s | 8.23 s | 104.9 GiB | first frame vs input 26.5 dB, SSIM 0.715 |
+| FL2VA | 28.95 s | 19.49 s | 8.12 s | 105.2 GiB | first frame 39.6 dB / 0.971, last frame 29.1 dB / 0.819 |
+| Ref2VA, two image references | 51.17 s | 38.42 s | 7.74 s | 108.5 GiB | 64,241 packed tokens; both references used |
+
+The NVFP4 text encoder's first request in a fresh process paid a one-time
+first-use cost of about 183 s; later requests did not.
+
+Peak resident memory for the BF16 FastVideo T2VA run was about 129 GiB. The
+official CLI's sampling peak was 71.7 GiB; the two figures measure different
+things (process-resident memory versus the sampling phase).
+
+The example's default `--offload lazy` loads components on demand and offloads
+the encoders and VAE. It fits smaller devices but reloads the DiT on every
+request: the same T2VA request took 48.06 s (28.31 s denoise) on GB200 in that
+mode. Use `--offload resident` when the device holds every component.
+
+## Training-free acceleration
+
+The 6 forwards already use FastVideo's normal inference path. Available
+lossless switches:
+
+* `FASTVIDEO_MINIMAX_H3_EXACT_KERNELS=all` — exact fused H3 kernels;
+* `FASTVIDEO_H3_REF2VA_MEMO_ENTRIES=16` — cache reference conditioning
+  (Ref2VA);
+* `FASTVIDEO_H3_VAE_TILE_PARALLEL=1` — parallel VAE tiles;
+* `FASTVIDEO_ULYSSES_A2A=auto` — Ulysses all-to-all for multi-GPU sequence
+  parallelism;
+* `OffloadConfig(lazy_module_load=True)` and `dit_layerwise=True` for smaller
+  GPUs;
+* `CompileConfig(enabled=True)` for `torch.compile`;
+* FlashAttention-4 through `FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+  FASTVIDEO_FA4=1` (see [Attention backend](#attention-backend)).
+
+`--lossless-accel` in the example enables the first four before the generator
+starts. It does not enable FA4; set the two attention variables at launch. Kernel work beyond the existing exact kernels is intentionally not part
+of this integration.
+
+## NVFP4 on Blackwell
+
+Pika can reuse FastVideo's existing H3 NVFP4 VAE decoder and NVFP4 Qwen3-VL
+text encoder. Both are weight formats, not new sampler paths; the six-forward
+Pika schedule is unchanged.
+
+```bash
+# 36-block quality decoder: needs the full 36-block H3 VAE
+python examples/inference/basic/basic_pika_h3_flex.py --task t2va \
+    --prompt "a paper boat drifting down a rainy street" \
+    --nvfp4-decoder ./FastH3-NVFP4-VAE-Decoder/quality-36L/decoder-nvfp4.safetensors
+
+# 8-block fast decoder: depth-cuts any H3 VAE with at least 8 blocks
+python examples/inference/basic/basic_pika_h3_flex.py --task ref2va \
+    --reference image:person.png --reference video:dance.mp4 \
+    --prompt "the person dances to the music" \
+    --nvfp4-decoder ./FastH3-NVFP4-VAE-Decoder/fast-8L/decoder-nvfp4.safetensors \
+    --nvfp4-text-encoder ./fasth3-nvfp4-encoder/text_encoder
+```
+
+* `quality-36L` keeps the full decoder depth and is the closest match to Pika's
+  dense decoder. `fast-8L` is much smaller and faster but is a distilled
+  student; expect a visible quality drop. Pika's six-step re-noise latents are
+  not the student's original training distribution, so validate 8L on your
+  prompts before using it for delivery.
+* The NVFP4 decoder replaces only `vae.decoder`; keyframe and reference
+  encoding still use the FP32 VAE encoder.
+* The NVFP4 text encoder quantizes the Qwen3-VL language stack while keeping
+  the vision tower in bf16, so Ref2VA reference images/videos still work.
+* Do not select `taeh3` with an NVFP4 decoder; `h3-vae` is required.
+* Both options require Blackwell FP4 hardware. The fused decoder kernels and
+  CUDA graphs are already enabled by default; see
+  `FASTVIDEO_H3_VAE_NVFP4_FUSED` and `FASTVIDEO_H3_VAE_NVFP4_CUDA_GRAPH`.
+* There is no NVFP4 **VAE encoder** path. The H3 video VAE encoder is a causal
+  3D CNN and stays FP32; quantizing it would need a separate convolution
+  kernel/quantization effort.
+
+If a process first loads the 8L decoder it depth-cuts the VAE to eight blocks;
+loading 36L afterwards in the same process is not supported. Choose one decoder
+tier per process.
+
+## Hardware
+
+The selected H3 trunk plus Qwen3-VL does not fit comfortably on a 24 GB card
+without offload. A single 80 GB GPU can run with text-encoder offload; use
+layerwise DiT offload or multiple GPUs for smaller devices. Ref2VA encodes
+reference images at up to a 2048 short edge, so its conditioning memory is
+higher than T2VA's.

docs/inference/support_matrix.md

@@ -68,6 +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) |
+| pika_h3_flex | `Pika-Labs/Pika-H3-Flex` | T2V, I2V | [T2VA/I2VA/L2VA/FL2VA/Ref2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_pika_h3_flex.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) |

docs/training/minimax_h3_fp4_qad.md

@@ -0,0 +1,253 @@
+# FastH3 NVFP4 QAD
+
+This page describes the shared quantization-aware distillation (QAD) recipe for
+the FastH3 (MiniMax-H3) family. One recipe covers two model profiles:
+
+| Profile | Model plugin | Task | Sampling | VSA tile |
+|---|---|---|---|---|
+| OmniRef | `omniref_pdd.MiniMaxH3OmniRefPDDModel` | Ref2VA | 8-forward PDD ladder | 128 |
+| V2 T2AV QAD | `v2_t2v.MiniMaxH3V2T2VModel` | T2VA | 8-step DMD | 64 |
+
+Both profiles use the same shared recipe: NVFP4 W4A4 linears in all 50 blocks,
+the FP4 sparse-attention emulator, and the same objective and guards. QAD
+trains the student against a frozen bf16 teacher so that quantization error is
+absorbed by the weights instead of appearing as drift, ghosting or audio loss
+at inference. The published OmniRef QAD checkpoint includes FP4 attention; the
+published FastH3-QAD-v2 checkpoint predates it and keeps BF16 VSA attention (see
+[Released checkpoints](#released-checkpoints)).
+
+The runnable configs are:
+
+- `examples/train/configs/knowledge_distillation/minimax_h3/omniref_nvfp4_qad.yaml`
+- `examples/train/configs/knowledge_distillation/minimax_h3/v2_nvfp4_qad.yaml`
+
+## Released checkpoints
+
+| Role | Checkpoint |
+|---|---|
+| OmniRef QAD result | `FastVideo/FastH3-OmniRef-NVFP4-QAD` |
+| FastH3-QAD-v2 result | `FastVideo/FastH3-QAD-v2` |
+| V2 bf16 teacher / starting weights | `FastVideo/FastVideo-FastH3-8-Step-V2` |
+| NVFP4 consumer baseline | `FastVideo/FastVideo-FastH3-8-Step-V2-NVFP4-Consumer` |
+
+The OmniRef config starts from the **calibrated OmniRef checkpoint used as QAD
+step 0**, not from the released OmniRef QAD result; point `init_from` at that
+checkpoint. The V2 config starts from the bf16 V2 snapshot.
+
+The published FastH3-QAD-v2 checkpoint is the **linears-only variant** of this
+schedule: NVFP4 W4A4 linears with BF16 VSA attention, 300 steps. The shared
+FP4 recipe on this page additionally runs the FP4 attention emulator from step
+0; it is the same training code and objective, but it does not reproduce the
+published FastH3-QAD-v2 checkpoint bit for bit. Set
+`models.student.qad.fp4_attention.enabled: false` to get the published
+linears-only variant.
+
+## Target numerics
+
+The student starts from the bf16 masters, and the deployed numerics are active
+from the first forward, so every update sees the quantization error it is
+trained to absorb.
+
+- **Linears:** NVFP4 W4A4 for `attn.to_q`, `attn.to_k`, `attn.to_v`,
+  `attn.to_out`, `ff.fc_in` and `ff.fc_out` in every block.
+  - FFN inputs use static activation scales from the calibrated `amax` table.
+  - Attention and gate linears use the unit activation scale; their `amax`
+    entries are optional.
+  - `skip_blocks` is empty: no block falls back to bf16 linears.
+- **Attention:** the FP4 sparse-attention emulator for VSA tile 128 (OmniRef)
+  or tile 64 (V2 T2AV). Two-level P quantization and K smoothing are on. The
+  emulator reproduces the deployed kernel's block order and rounding on any
+  CUDA GPU and provides straight-through gradients.
+- **Kept frozen:** the VSA compression gates and every norm stay at their
+  released values. Training the gate weights drifts the sparse-attention tile
+  selection and produces ghosting; the gate linears still execute the deployed
+  NVFP4 format, they are just not updated. `proj_in`, the time embedder,
+  `proj_out` and the audio projections stay fp32 and frozen, matching
+  inference.
+
+The calibrated `amax` tables for both profiles are vendored under
+`examples/train/calibration/fasth3/`:
+
+| Table | sha256 |
+|---|---|
+| `amax-omniref-mse-bf16-top8.json` | `14b93ec3…df083d` |
+| `amax-consumer.json` | `0800516c…d39e52` |
+
+`qad.linears.amax_sha256` pins the table; the run copies it into the output
+directory and refuses to start if the bytes differ.
+
+## Training objective
+
+For each row the frozen bf16 teacher walks the profile's sampling ladder. At
+the same inputs and timesteps, the student minimizes:
+
+- **Output regression:** per-forward x0 error, normalized per forward and
+  modality by its step-0 value.
+- **Layer-wise reconstruction:** each attention and FFN sub-layer receives the
+  teacher's input (stop-gradient) and matches the teacher sub-layer's output.
+  The loss is weighted per row group (target video and audio 1.0,
+  conditioning 0.25) because audio is a small fraction of the packed tokens.
+  `recon_normalize: step0` divides it by its step-0 value so `recon_lambda`
+  balances the two terms.
+
+From `rollout_start_step`, `rollout_fraction` of the rows continue from the
+student's own trajectory instead of the teacher's, which reduces exposure bias
+in the sampled output.
+
+## Evaluation and guards
+
+- **Teacher-forced eval** every `eval_every` steps on held-out rows.
+- **Full eval** every `full_eval_every` steps: the student's own sampled output
+  against the teacher at the same seed, keyframe alignment, and saved audio and
+  video latents.
+- **Noise floor:** `noise_floor_eval` re-runs the step-0 eval at seed + 1 so a
+  metric move can be compared with teacher seed noise.
+- **Best checkpoint:** selected by the teacher-forced `score`. A checkpoint is
+  eligible only if the keys in `eligibility_keys` (audio error and endpoint)
+  stay within their step-0 guardrails.
+- **Auto-stop:** after `patience` consecutive bad evals.
+- **Step-0 cache:** `step0_cache` fingerprints the teacher, student,
+  calibration, quantization plan, eval set and loss weights, so reruns skip the
+  baseline evaluation.
+- **Resume:** checkpoints carry the step-0 baseline and guard state; do not
+  delete `method.qad_baseline` if you edit a checkpoint by hand.
+
+## Data preparation
+
+V2 T2AV is data-free: it needs prompts and their Qwen3-VL text embeddings, not
+VAE latents.
+
+1. Build the split manifests from prompt sources (parquet shards with a
+   `prompt` column, JSONL with a `prompt` field):
+
+   ```bash
+   python scripts/distill/minimax_h3_v2_qad/build_v2_prompts.py \
+     --parquet 'data/prompts/*.parquet' --out-dir data/fasth3-v2 --eval-count 30
+   ```
+
+   This writes `train.jsonl` and `eval.jsonl` plus an `embeddings/` directory.
+2. Fill the embeddings (safe to rerun and shard across GPUs):
+
+   ```bash
+   python scripts/distill/minimax_h3_v2_qad/precompute_v2_text.py \
+     --config examples/inference/basic/basic_fasth3_spark_v2_nvfp4.yaml \
+     --model-path data/models/FastH3-V2 \
+     --manifests data/fasth3-v2/train.jsonl data/fasth3-v2/eval.jsonl \
+     --shard 0 --num-shards 8
+   ```
+
+   The V2 config's `train_manifest` and `eval_manifest` point at these files.
+
+OmniRef training rows are preprocessed Ref2VA parquet rows that carry the
+reference latents, text embeddings and geometry; the encoded dataset itself is
+an input to this recipe and is not produced by the training code. Provide a
+`MANIFEST.jsonl` with `case`, `id`, `parquet` and `resolution` fields and point
+`models.student.data.manifests` at it.
+
+`scripts/distill/minimax_h3_omniref_qad/generate_omniref_latents.py` builds the
+held-out eval plan and runs the PDD ladder once per row: use `--eval-dir` to
+write the held-out `manifest.json` consumed by `models.student.data.eval_manifest`
+(and the raw latents for decoding).
+
+## Running
+
+Fill in the placeholder paths in the config (starting checkpoint, data
+manifests, output directory), then launch with `torchrun`.
+
+OmniRef, 4 nodes x 4 GPUs (sequence parallel 4):
+
+```bash
+torchrun --nnodes=4 --nproc_per_node=4 --node_rank=$NODE_RANK \
+  --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT \
+  -m fastvideo.train.entrypoint.train \
+  --config examples/train/configs/knowledge_distillation/minimax_h3/omniref_nvfp4_qad.yaml
+```
+
+V2 T2AV, 4 nodes x 4 GPUs (16-way HSDP, lr 5e-6, 300 steps):
+
+```bash
+torchrun --nnodes=4 --nproc_per_node=4 --node_rank=$NODE_RANK \
+  --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT \
+  -m fastvideo.train.entrypoint.train \
+  --config examples/train/configs/knowledge_distillation/minimax_h3/v2_nvfp4_qad.yaml
+```
+
+To reproduce the published FastH3-QAD-v2 linears-only checkpoint, disable FP4 attention
+for the same schedule:
+
+```bash
+  --models.student.qad.fp4_attention.enabled false
+```
+
+## FP4-attention adaptation
+
+The internal follow-up that this recipe tracks started from a linears-only
+300-step DCP, enabled FP4 attention, and trained 300 more steps at lr 1.5e-6
+with no warmup and a constant LR on 16 GPUs. It is an **adaptation of the
+linears-only checkpoint to FP4 attention**, not a second epoch of the FP4
+stage 1 above: the source checkpoint's attention numerics were BF16, and the
+continuation needs a DCP of that linears-only model (the packed HF release is
+not one). It is therefore not part of the public reproduction path, and the
+published FastH3-QAD-v2 checkpoint remains the linears-only result. The exact overrides
+are listed at the bottom of the V2 config.
+
+## Checkpoints
+
+Runs write regular `checkpoint-<step>/` DCP checkpoints plus a `best/`
+directory:
+
+- `best/`: the best eligible state (weights only), refreshed at every eval and
+  never pruned.
+- `checkpoint-<step>/`: weights, optimizer, RNG state and the method baseline;
+  the last step is always saved.
+
+The saved state is the bf16 master plus the quantization plan, not packed NVFP4
+weights; inference applies the same converter used for the released model. Use
+`fastvideo/train/entrypoint/dcp_to_diffusers.py` to export a checkpoint if your
+deployment path expects Diffusers weights. It detects the weights-only marker
+on a tagged checkpoint, so `best/` can be exported directly:
+
+```bash
+python -m fastvideo.train.entrypoint.dcp_to_diffusers \
+  --checkpoint outputs/v2_1_nvfp4_qad/best \
+  --output-dir outputs/v2_1_nvfp4_qad/export
+```
+
+## Code map
+
+| Path | Purpose |
+|---|---|
+| `fastvideo/train/methods/knowledge_distillation/pdd_qad.py` | Method: ladder, losses, evaluation, guards |
+| `fastvideo/train/methods/knowledge_distillation/pdd_qad_metrics.py` | Rung aggregation and score |
+| `fastvideo/train/methods/knowledge_distillation/pdd_qad_recon.py` | Layer-wise reconstruction hooks |
+| `fastvideo/train/methods/knowledge_distillation/pdd_qad_step0.py` | Step-0 baseline cache |
+| `fastvideo/train/models/minimax_h3/omniref_pdd.py` | OmniRef Ref2VA profile |
+| `fastvideo/train/models/minimax_h3/v2_t2v.py` | V2 T2AV profile |
+| `fastvideo/layers/quantization/nvfp4_qad.py` | NVFP4 linear emulation and STE |
+| `fastvideo/attention/backends/fp4_vsa_qat.py` | VSA fine-attention override |
+| `fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_train.py` | Tile 64/128 FP4 block-sparse Triton kernel |
+| `fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_reference.py` | PyTorch numeric reference |
+| `scripts/distill/minimax_h3_v2_qad/build_v2_prompts.py` | V2 prompt split manifests |
+| `scripts/distill/minimax_h3_v2_qad/precompute_v2_text.py` | V2 Qwen3-VL text embeddings |
+| `scripts/distill/minimax_h3_omniref_qad/generate_omniref_latents.py` | OmniRef held-out eval plan and latents |
+| `examples/train/calibration/fasth3/` | Pinned FFN activation `amax` tables |
+
+## Expected behavior
+
+Teacher-forced video and audio error both fall over the first tens of steps,
+and the student's own video samples improve beyond the teacher's seed noise.
+Audio samples improve more slowly and can remain within noise even when
+teacher-forced audio error clearly drops. Gains slow after roughly 60-100
+steps; the remaining error is dominated by FP4 attention. Evaluate a checkpoint
+with the full-eval samples, not the teacher-forced score alone, before
+promoting it.
+
+Released reference points (lower is better):
+
+| Run | Metric | Baseline | Selected |
+|---|---|---|---|
+| FastH3-QAD-v2 (linears-only) | mean per-forward x0 rel-L2 | 0.431 | 0.351 (video 0.382, audio 0.319) |
+| OmniRef QAD | teacher-forced score | 0.362 | 0.303 at step 120 (video 0.198, audio 0.409) |
+
+The FastH3-QAD-v2 baseline is the NVFP4 Consumer model before QAD; the OmniRef baseline
+is the calibrated checkpoint at QAD step 0.

examples/inference/basic/basic_pika_h3_flex.py

@@ -0,0 +1,297 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Generate synchronized video/audio with Pika H3 Flex.
+
+Five task shapes share one model release:
+
+``t2va``  text only
+``i2va``  first image
+``l2va``  last image
+``fl2va`` first and last images
+``ref2va`` ordered image/video/audio references
+
+The six-forward ladder is part of the adapter contract; this script deliberately
+does not expose a step-count flag.  ``--lossless-accel`` turns on the existing
+training-free FastVideo switches (bit-identical kernels, reference memo,
+parallel VAE tiles) before the generator starts.  ``--nvfp4-decoder`` and
+``--nvfp4-text-encoder`` opt into the existing Blackwell FP4 weight paths;
+``--nvfp4-dit h3_dit`` loads a packed NVFP4 DiT export (e.g.
+FastVideo/Pika-H3-Flex-NVFP4-QAD).  They change weights, not the six-forward sampler.  Kernel/VSA work is separate.
+
+``--nvfp4-attention`` turns on FlashAttention-4 NVFP4 attention (Q/K in NVFP4, P/V in BF16; Blackwell and FA4 only)
+with the named exact Q/K preconditioning, which must match the ``qk_transform`` the checkpoint was QAD-trained with
+(``smooth_k+hadamard`` for Pika-H3-Flex-NVFP4-QAD).  ``--nvfp4-attention-bf16-blocks 45,49`` keeps those
+``transformer_blocks`` on BF16 attention, as in training.  ``none`` quantizes Q/K without preconditioning.
+Without ``--nvfp4-attention`` attention stays BF16.
+
+``--resolution`` (480p/768p), ``--aspect`` (one of Pika's eleven buckets), and
+``--duration`` (5-15 s) pick the bucket and the ``17n + 5`` frame count;
+``--height``/``--width``/``--num-frames`` override them with explicit values.
+Launch with ``FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN FASTVIDEO_FA4=1`` on a
+FlashAttention-4 host (Hopper/Blackwell); without FA4 attention falls back to
+FA3/FA2 or, if neither is installed, to SDPA.
+
+``--offload resident`` keeps every component on the GPU, which is how the
+GB200 timings in docs/inference/pika_h3_flex.md were measured; the default
+``lazy`` mode loads components on demand for smaller devices and is slower.
+
+Examples:
+    FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN FASTVIDEO_FA4=1 \
+    python examples/inference/basic/basic_pika_h3_flex.py --task t2va --offload resident \
+        --prompt "a paper boat drifting down a rainy street"
+    python examples/inference/basic/basic_pika_h3_flex.py --task t2va \
+        --resolution 480p --aspect 9:16 --duration 10 --prompt "a lighthouse at dusk"
+    python examples/inference/basic/basic_pika_h3_flex.py --task fl2va \
+        --image first.png --last-image last.png --prompt "the flower opens"
+    python examples/inference/basic/basic_pika_h3_flex.py --task ref2va \
+        --reference image:person.png --reference video:dance.mp4 \
+        --reference audio:voice.wav --prompt "the person dances to the music"
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+from pathlib import Path
+from typing import Any
+
+from fastvideo import VideoGenerator, envs
+from fastvideo.api import (
+    CompileConfig,
+    ComponentConfig,
+    EngineConfig,
+    GenerationRequest,
+    GeneratorConfig,
+    InputConfig,
+    OffloadConfig,
+    OutputConfig,
+    ParallelismConfig,
+    PipelineSelection,
+    QuantizationConfig,
+    SamplingConfig,
+)
+from fastvideo.attention.utils.nvfp4_qk_precondition import QK_TRANSFORMS, parse_block_indices
+from fastvideo.models.vision_utils import load_image
+from fastvideo.pipelines.basic.minimax_h3 import MiniMaxH3Reference
+from fastvideo.pipelines.basic.pika_h3_flex.buckets import (
+    PIKA_ASPECTS,
+    PIKA_FPS,
+    PIKA_MAX_DURATION,
+    PIKA_MIN_DURATION,
+    canvas_for_aspect,
+    resolve_frames,
+)
+
+# Training-free switches FastVideo already exposes. Workers inherit the
+# environment, so they must be set before VideoGenerator starts.
+LOSSLESS_ACCEL_ENV = {
+    "FASTVIDEO_MINIMAX_H3_EXACT_KERNELS": "all",
+    "FASTVIDEO_H3_REF2VA_MEMO_ENTRIES": "16",
+    "FASTVIDEO_H3_VAE_TILE_PARALLEL": "1",
+    "FASTVIDEO_ULYSSES_A2A": "auto",
+}
+
+TASKS = ("t2va", "i2va", "l2va", "fl2va", "ref2va")
+RESOLUTIONS = {"480p": "p480", "768p": "p768"}
+
+
+def _reference(value: str) -> MiniMaxH3Reference:
+    try:
+        media_type, source = value.split(":", 1)
+    except ValueError as error:
+        raise argparse.ArgumentTypeError("A reference must be image:PATH, video:PATH, or audio:PATH.") from error
+    if media_type not in ("image", "video", "audio"):
+        raise argparse.ArgumentTypeError(f"Unknown reference type {media_type!r}.")
+    return MiniMaxH3Reference(source=source, media_type=media_type)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--model-path", default="Pika-Labs/Pika-H3-Flex")
+    parser.add_argument("--task", choices=TASKS, default="t2va")
+    parser.add_argument("--prompt", required=True)
+    parser.add_argument("--image", help="first frame (i2va/fl2va)")
+    parser.add_argument("--last-image", help="last frame (l2va/fl2va)")
+    parser.add_argument("--reference",
+                        action="append",
+                        type=_reference,
+                        default=[],
+                        help="ordered reference (repeatable; ref2va)")
+    parser.add_argument("--resolution", choices=tuple(RESOLUTIONS), default="768p", help="Pika resolution tier")
+    parser.add_argument("--aspect", choices=PIKA_ASPECTS, default="16:9", help="bucket aspect ratio (W:H)")
+    parser.add_argument("--duration",
+                        type=float,
+                        default=5.0,
+                        help=f"seconds ({PIKA_MIN_DURATION:g}-{PIKA_MAX_DURATION:g}); rounded up to the 17n+5 grid")
+    parser.add_argument("--height", type=int, help="explicit canvas height (with --width; overrides the bucket)")
+    parser.add_argument("--width", type=int, help="explicit canvas width (with --height; overrides the bucket)")
+    parser.add_argument("--num-frames", type=int, help="explicit frame count (overrides --duration)")
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--num-gpus", type=int, default=1)
+    parser.add_argument("--compile", action="store_true", help="torch.compile the DiT")
+    parser.add_argument("--lossless-accel", action="store_true", help="enable bit-identical speed-ups")
+    parser.add_argument("--offload",
+                        choices=("lazy", "resident"),
+                        default="lazy",
+                        help="lazy: load components on demand and offload the encoders/VAE (fits smaller devices); "
+                        "resident: keep everything on the GPU (fastest, about 105-130 GiB on GB200)")
+    parser.add_argument("--low-vram", action="store_true", help="stream the DiT layerwise (lazy offload only)")
+    parser.add_argument("--nvfp4-decoder",
+                        help="Blackwell-only NVFP4 H3 decoder (quality-36L or fast-8L decoder-nvfp4.safetensors)")
+    parser.add_argument("--nvfp4-text-encoder",
+                        help="Blackwell-only serialized NVFP4 Qwen3-VL text_encoder directory")
+    parser.add_argument("--nvfp4-dit",
+                        choices=("h3_dit", "h3_dit_ffn"),
+                        help="load a packed NVFP4 DiT export with this layer profile (h3_dit: attention + FFN linears)")
+    parser.add_argument("--nvfp4-attention",
+                        choices=tuple(QK_TRANSFORMS),
+                        help="FA4 NVFP4 Q/K attention with this exact Q/K preconditioning (default: BF16 attention)")
+    parser.add_argument("--nvfp4-attention-bf16-blocks",
+                        default="",
+                        help="comma-separated transformer_blocks indices that keep BF16 attention, e.g. 45,49")
+    parser.add_argument("--output", default="outputs/pika_h3_flex")
+    args = parser.parse_args()
+    if args.nvfp4_attention_bf16_blocks and not args.nvfp4_attention:
+        parser.error("--nvfp4-attention-bf16-blocks needs --nvfp4-attention")
+    try:
+        parse_block_indices(args.nvfp4_attention_bf16_blocks)
+    except ValueError as error:
+        parser.error(f"--nvfp4-attention-bf16-blocks: {error}")
+
+    if args.task in ("i2va", "fl2va") and not args.image:
+        parser.error(f"--task {args.task} needs --image")
+    if args.task in ("l2va", "fl2va") and not args.last_image:
+        parser.error(f"--task {args.task} needs --last-image")
+    if args.task == "ref2va" and not args.reference:
+        parser.error("--task ref2va needs at least one --reference")
+    if args.task != "ref2va" and args.reference:
+        parser.error("--reference is only valid with --task ref2va")
+    if (args.height is None) != (args.width is None):
+        parser.error("--height and --width must be passed together")
+    if args.low_vram and args.offload == "resident":
+        parser.error("--low-vram streams the DiT and cannot be combined with --offload resident")
+    return args
+
+
+def _offload(args: argparse.Namespace) -> OffloadConfig:
+    if args.offload == "resident":
+        return OffloadConfig(dit=False,
+                             dit_layerwise=False,
+                             text_encoder=False,
+                             vae=False,
+                             pin_cpu_memory=False,
+                             lazy_module_load=False)
+    return OffloadConfig(
+        dit=False,
+        dit_layerwise=args.low_vram,
+        text_encoder=True,
+        vae=True,
+        pin_cpu_memory=not args.low_vram,
+        lazy_module_load=True,
+    )
+
+
+def _canvas(args: argparse.Namespace) -> tuple[int, int]:
+    if args.height is not None:
+        return args.height, args.width
+    return canvas_for_aspect(args.aspect, RESOLUTIONS[args.resolution])
+
+
+def _num_frames(args: argparse.Namespace) -> int:
+    if args.num_frames is not None:
+        return args.num_frames
+    return resolve_frames(round(args.duration * PIKA_FPS))
+
+
+def _inputs(args: argparse.Namespace) -> InputConfig:
+    if args.task == "ref2va":
+        return InputConfig(references=args.reference)
+    # Keep only the keyframes the task uses, so switching --task never turns
+    # an I2VA run into FL2VA because a --last-image was also passed.
+    first = args.image if args.task in ("i2va", "fl2va") else None
+    last = args.last_image if args.task in ("l2va", "fl2va") else None
+    # The request adapter copies image fields verbatim; the stage loads
+    # ``image_path`` itself but requires ``last_image`` to already be a PIL image.
+    return InputConfig(image_path=first, last_image=load_image(last) if last else None)
+
+
+def _pipeline(args: argparse.Namespace) -> PipelineSelection:
+    components = ComponentConfig()
+    if args.task == "ref2va":
+        components.override_pipeline_cls_name = "PikaH3FlexRef2VAModularPipeline"
+    if args.nvfp4_decoder:
+        components.h3_vae_nvfp4_decoder = args.nvfp4_decoder
+    if args.nvfp4_text_encoder:
+        components.text_encoder_weights = args.nvfp4_text_encoder
+    return PipelineSelection(
+        workload_type="t2v" if args.task == "t2va" else "i2v",
+        components=components,
+        vae_tiling=True,
+    )
+
+
+def _apply_nvfp4_attention(args: argparse.Namespace) -> None:
+    """Workers inherit the environment, so set it before VideoGenerator starts."""
+    if not args.nvfp4_attention:
+        return
+    envs.FASTVIDEO_NVFP4_FA4.set(True)
+    envs.FASTVIDEO_NVFP4_FA4_QK_TRANSFORM.set(args.nvfp4_attention)
+    envs.FASTVIDEO_NVFP4_FA4_BF16_BLOCKS.set(args.nvfp4_attention_bf16_blocks)
+    envs.setdefault_external("CUTE_DSL_ENABLE_TVM_FFI", "1")
+
+
+def main() -> None:
+    args = parse_args()
+    _apply_nvfp4_attention(args)
+    if args.lossless_accel:
+        for key, value in LOSSLESS_ACCEL_ENV.items():
+            os.environ.setdefault(key, value)
+
+    height, width = _canvas(args)
+    num_frames = _num_frames(args)
+    print(f"Pika H3 Flex {args.task}: {width}x{height}, {num_frames} frames ({num_frames / PIKA_FPS:.2f} s)")
+    output_path = Path(args.output) / f"pika_h3_flex_{args.task}.mp4"
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    offload = _offload(args)
+    generator = VideoGenerator.from_config(
+        GeneratorConfig(
+            model_path=args.model_path,
+            engine=EngineConfig(
+                num_gpus=args.num_gpus,
+                use_fsdp_inference=args.num_gpus > 1,
+                parallelism=ParallelismConfig(tp_size=1, sp_size=args.num_gpus),
+                offload=offload,
+                compile=CompileConfig(enabled=args.compile),
+                quantization=(QuantizationConfig(transformer_quant="NVFP4", layer_profile=args.nvfp4_dit)
+                              if args.nvfp4_dit else None),
+            ),
+            pipeline=_pipeline(args),
+        ))
+    try:
+        result = generator.generate(
+            GenerationRequest(
+                prompt=args.prompt,
+                negative_prompt="",
+                inputs=_inputs(args),
+                sampling=SamplingConfig(
+                    height=height,
+                    width=width,
+                    num_frames=num_frames,
+                    fps=24,
+                    num_inference_steps=6,
+                    guidance_scale=1.0,
+                    batch_cfg=False,
+                    seed=args.seed,
+                ),
+                output=OutputConfig(
+                    output_path=str(output_path),
+                    save_video=True,
+                    return_frames=False,
+                ),
+            ))
+        print(f"Output written to: {result.video_path}")
+    finally:
+        generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/train/calibration/fasth3/amax-consumer.json

@@ -0,0 +1,102 @@
+{
+ "b0.ff.fc_in": 747.9899299279987,
+ "b0.ff.fc_out": 2462.017121214983,
+ "b1.ff.fc_in": 671.2121750026851,
+ "b1.ff.fc_out": 1285.509862734237,
+ "b2.ff.fc_in": 1153.55816110724,
+ "b2.ff.fc_out": 3796.948980859055,
+ "b3.ff.fc_in": 1153.55816110724,
+ "b3.ff.fc_out": 4715.270819259384,
+ "b4.ff.fc_in": 1285.509862734237,
+ "b4.ff.fc_out": 4715.270819259384,
+ "b5.ff.fc_in": 2209.302022692089,
+ "b5.ff.fc_out": 5254.635977347013,
+ "b6.ff.fc_in": 2209.302022692089,
+ "b6.ff.fc_out": 5254.635977347013,
+ "b7.ff.fc_in": 2462.017121214983,
+ "b7.ff.fc_out": 4715.270819259384,
+ "b8.ff.fc_in": 2462.017121214983,
+ "b8.ff.fc_out": 5855.696365467314,
+ "b9.ff.fc_in": 2743.6392141259485,
+ "b9.ff.fc_out": 5855.696365467314,
+ "b10.ff.fc_in": 2743.6392141259485,
+ "b10.ff.fc_out": 5855.696365467314,
+ "b11.ff.fc_in": 2743.6392141259485,
+ "b11.ff.fc_out": 6525.510437495776,
+ "b12.ff.fc_in": 2743.6392141259485,
+ "b12.ff.fc_out": 6525.510437495776,
+ "b13.ff.fc_in": 3057.4751817363253,
+ "b13.ff.fc_out": 6525.510437495776,
+ "b14.ff.fc_in": 2743.6392141259485,
+ "b14.ff.fc_out": 6525.510437495776,
+ "b15.ff.fc_in": 2743.6392141259485,
+ "b15.ff.fc_out": 6525.510437495776,
+ "b16.ff.fc_in": 3407.2097180654623,
+ "b16.ff.fc_out": 6525.510437495776,
+ "b17.ff.fc_in": 3057.4751817363253,
+ "b17.ff.fc_out": 9030.717912847476,
+ "b18.ff.fc_in": 2743.6392141259485,
+ "b18.ff.fc_out": 9030.717912847476,
+ "b19.ff.fc_in": 2743.6392141259485,
+ "b19.ff.fc_out": 9030.717912847476,
+ "b20.ff.fc_in": 3057.4751817363253,
+ "b20.ff.fc_out": 10063.712585454898,
+ "b21.ff.fc_in": 3057.4751817363253,
+ "b21.ff.fc_out": 10063.712585454898,
+ "b22.ff.fc_in": 3407.2097180654623,
+ "b22.ff.fc_out": 10063.712585454898,
+ "b23.ff.fc_in": 3407.2097180654623,
+ "b23.ff.fc_out": 12497.700799896355,
+ "b24.ff.fc_in": 3796.948980859055,
+ "b24.ff.fc_out": 13927.271620769527,
+ "b25.ff.fc_in": 3796.948980859055,
+ "b25.ff.fc_out": 19274.08888735978,
+ "b26.ff.fc_in": 4231.26976101764,
+ "b26.ff.fc_out": 21478.789600085922,
+ "b27.ff.fc_in": 4715.270819259384,
+ "b27.ff.fc_out": 29724.70802687751,
+ "b28.ff.fc_in": 4715.270819259384,
+ "b28.ff.fc_out": 45841.764386937015,
+ "b29.ff.fc_in": 4715.270819259384,
+ "b29.ff.fc_out": 56928.958163414085,
+ "b30.ff.fc_in": 5254.635977347013,
+ "b30.ff.fc_out": 78784.54787887946,
+ "b31.ff.fc_in": 5254.635977347013,
+ "b31.ff.fc_out": 109030.7109780948,
+ "b32.ff.fc_in": 5855.696365467314,
+ "b32.ff.fc_out": 135400.63235256448,
+ "b33.ff.fc_in": 5855.696365467314,
+ "b33.ff.fc_out": 168148.35127510442,
+ "b34.ff.fc_in": 6525.510437495776,
+ "b34.ff.fc_out": 168148.35127510442,
+ "b35.ff.fc_in": 6525.510437495776,
+ "b35.ff.fc_out": 232702.1493904285,
+ "b36.ff.fc_in": 6525.510437495776,
+ "b36.ff.fc_out": 232702.1493904285,
+ "b37.ff.fc_in": 6525.510437495776,
+ "b37.ff.fc_out": 208816.33399986484,
+ "b38.ff.fc_in": 6525.510437495776,
+ "b38.ff.fc_out": 288982.97321111924,
+ "b39.ff.fc_in": 7271.942168450891,
+ "b39.ff.fc_out": 358875.77201571054,
+ "b40.ff.fc_in": 7271.942168450891,
+ "b40.ff.fc_out": 445672.69186462485,
+ "b41.ff.fc_in": 8103.755565259112,
+ "b41.ff.fc_out": 496651.7297372453,
+ "b42.ff.fc_in": 9030.717912847476,
+ "b42.ff.fc_out": 616770.844395513,
+ "b43.ff.fc_in": 9030.717912847476,
+ "b43.ff.fc_out": 616770.844395513,
+ "b44.ff.fc_in": 10063.712585454898,
+ "b44.ff.fc_out": 687321.2672027495,
+ "b45.ff.fc_in": 11214.867875490954,
+ "b45.ff.fc_out": 853555.2832070139,
+ "b46.ff.fc_in": 12497.700799896355,
+ "b46.ff.fc_out": 1059994.2350776498,
+ "b47.ff.fc_in": 13927.271620769527,
+ "b47.ff.fc_out": 1316362.0994398587,
+ "b48.ff.fc_in": 17295.691014283653,
+ "b48.ff.fc_out": 2030107.8425012752,
+ "b49.ff.fc_in": 19274.08888735978,
+ "b49.ff.fc_out": 2809486.555264971
+}

examples/train/calibration/fasth3/amax-omniref-mse-bf16-top8.json

@@ -0,0 +1,102 @@
+{
+ "transformer_blocks.0.ff.fc_in": 7.34375,
+ "transformer_blocks.0.ff.fc_out": 6432.0,
+ "transformer_blocks.1.ff.fc_in": 6.3125,
+ "transformer_blocks.1.ff.fc_out": 8960.0,
+ "transformer_blocks.10.ff.fc_in": 15.25,
+ "transformer_blocks.10.ff.fc_out": 6368.0,
+ "transformer_blocks.11.ff.fc_in": 17.75,
+ "transformer_blocks.11.ff.fc_out": 16128.0,
+ "transformer_blocks.12.ff.fc_in": 18.549999999999997,
+ "transformer_blocks.12.ff.fc_out": 4064.0,
+ "transformer_blocks.13.ff.fc_in": 20.25,
+ "transformer_blocks.13.ff.fc_out": 78848.0,
+ "transformer_blocks.14.ff.fc_in": 12.0,
+ "transformer_blocks.14.ff.fc_out": 1784.0,
+ "transformer_blocks.15.ff.fc_in": 13.1875,
+ "transformer_blocks.15.ff.fc_out": 1400.0,
+ "transformer_blocks.16.ff.fc_in": 12.000625,
+ "transformer_blocks.16.ff.fc_out": 2768.16,
+ "transformer_blocks.17.ff.fc_in": 22.049999999999997,
+ "transformer_blocks.17.ff.fc_out": 4960.0,
+ "transformer_blocks.18.ff.fc_in": 8.875,
+ "transformer_blocks.18.ff.fc_out": 15616.0,
+ "transformer_blocks.19.ff.fc_in": 12.4375,
+ "transformer_blocks.19.ff.fc_out": 3203.2,
+ "transformer_blocks.2.ff.fc_in": 4.2875,
+ "transformer_blocks.2.ff.fc_out": 7648.0,
+ "transformer_blocks.20.ff.fc_in": 8.96875,
+ "transformer_blocks.20.ff.fc_out": 5696.0,
+ "transformer_blocks.21.ff.fc_in": 10.0625,
+ "transformer_blocks.21.ff.fc_out": 6464.0,
+ "transformer_blocks.22.ff.fc_in": 14.000625,
+ "transformer_blocks.22.ff.fc_out": 11712.0,
+ "transformer_blocks.23.ff.fc_in": 13.125,
+ "transformer_blocks.23.ff.fc_out": 13312.0,
+ "transformer_blocks.24.ff.fc_in": 18.125,
+ "transformer_blocks.24.ff.fc_out": 9497.599999999999,
+ "transformer_blocks.25.ff.fc_in": 10.7625,
+ "transformer_blocks.25.ff.fc_out": 12864.0,
+ "transformer_blocks.26.ff.fc_in": 10.443750000000001,
+ "transformer_blocks.26.ff.fc_out": 11776.0,
+ "transformer_blocks.27.ff.fc_in": 13.9375,
+ "transformer_blocks.27.ff.fc_out": 13504.0,
+ "transformer_blocks.28.ff.fc_in": 13.9375,
+ "transformer_blocks.28.ff.fc_out": 8448.0,
+ "transformer_blocks.29.ff.fc_in": 11.4375,
+ "transformer_blocks.29.ff.fc_out": 13058.559999999998,
+ "transformer_blocks.3.ff.fc_in": 6.28125,
+ "transformer_blocks.3.ff.fc_out": 10880.0,
+ "transformer_blocks.30.ff.fc_in": 12.000625,
+ "transformer_blocks.30.ff.fc_out": 18816.0,
+ "transformer_blocks.31.ff.fc_in": 7.75,
+ "transformer_blocks.31.ff.fc_out": 8705.279999999999,
+ "transformer_blocks.32.ff.fc_in": 7.9375,
+ "transformer_blocks.32.ff.fc_out": 13312.0,
+ "transformer_blocks.33.ff.fc_in": 10.0,
+ "transformer_blocks.33.ff.fc_out": 8035.200000000001,
+ "transformer_blocks.34.ff.fc_in": 6.95625,
+ "transformer_blocks.34.ff.fc_out": 7424.0,
+ "transformer_blocks.35.ff.fc_in": 9.25,
+ "transformer_blocks.35.ff.fc_out": 20480.0,
+ "transformer_blocks.36.ff.fc_in": 9.40625,
+ "transformer_blocks.36.ff.fc_out": 29440.0,
+ "transformer_blocks.37.ff.fc_in": 16.0,
+ "transformer_blocks.37.ff.fc_out": 14272.0,
+ "transformer_blocks.38.ff.fc_in": 15.751249999999999,
+ "transformer_blocks.38.ff.fc_out": 28160.0,
+ "transformer_blocks.39.ff.fc_in": 15.064374999999998,
+ "transformer_blocks.39.ff.fc_out": 149504.0,
+ "transformer_blocks.4.ff.fc_in": 6.75,
+ "transformer_blocks.4.ff.fc_out": 14400.0,
+ "transformer_blocks.40.ff.fc_in": 11.438125,
+ "transformer_blocks.40.ff.fc_out": 31232.0,
+ "transformer_blocks.41.ff.fc_in": 17.25,
+ "transformer_blocks.41.ff.fc_out": 37383.67999999999,
+ "transformer_blocks.42.ff.fc_in": 15.0625,
+ "transformer_blocks.42.ff.fc_out": 15296.0,
+ "transformer_blocks.43.ff.fc_in": 17.625,
+ "transformer_blocks.43.ff.fc_out": 24064.0,
+ "transformer_blocks.44.ff.fc_in": 24.375,
+ "transformer_blocks.44.ff.fc_out": 99840.0,
+ "transformer_blocks.45.ff.fc_in": 40.0025,
+ "transformer_blocks.45.ff.fc_out": 350208.0,
+ "transformer_blocks.46.ff.fc_in": 42.75749999999999,
+ "transformer_blocks.46.ff.fc_out": 20070.399999999998,
+ "transformer_blocks.47.ff.fc_in": 57.2525,
+ "transformer_blocks.47.ff.fc_out": 49920.0,
+ "transformer_blocks.48.ff.fc_in": 145.0,
+ "transformer_blocks.48.ff.fc_out": 96256.0,
+ "transformer_blocks.49.ff.fc_in": 146.0,
+ "transformer_blocks.49.ff.fc_out": 200704.0,
+ "transformer_blocks.5.ff.fc_in": 13.1875,
+ "transformer_blocks.5.ff.fc_out": 15168.0,
+ "transformer_blocks.6.ff.fc_in": 19.375,
+ "transformer_blocks.6.ff.fc_out": 11904.0,
+ "transformer_blocks.7.ff.fc_in": 33.5,
+ "transformer_blocks.7.ff.fc_out": 15872.0,
+ "transformer_blocks.8.ff.fc_in": 16.5,
+ "transformer_blocks.8.ff.fc_out": 9216.0,
+ "transformer_blocks.9.ff.fc_in": 19.375,
+ "transformer_blocks.9.ff.fc_out": 4096.0
+}

examples/train/calibration/pika_h3_flex/amax-ref2va-balanced.json

@@ -0,0 +1,102 @@
+{
+ "b0.ff.fc_in": 5.75,
+ "b0.ff.fc_out": 1792.0,
+ "b1.ff.fc_in": 8.9375,
+ "b1.ff.fc_out": 2704.0,
+ "b2.ff.fc_in": 4.5625,
+ "b2.ff.fc_out": 2096.0,
+ "b3.ff.fc_in": 6.875,
+ "b3.ff.fc_out": 2384.0,
+ "b4.ff.fc_in": 11.125,
+ "b4.ff.fc_out": 2720.0,
+ "b5.ff.fc_in": 18.549999999999997,
+ "b5.ff.fc_out": 8225.280000000006,
+ "b6.ff.fc_in": 33.0,
+ "b6.ff.fc_out": 4352.0,
+ "b7.ff.fc_in": 18.75,
+ "b7.ff.fc_out": 4608.0,
+ "b8.ff.fc_in": 13.375,
+ "b8.ff.fc_out": 2320.0,
+ "b9.ff.fc_in": 13.212499999999999,
+ "b9.ff.fc_out": 2890.880000000001,
+ "b10.ff.fc_in": 23.875,
+ "b10.ff.fc_out": 4320.0,
+ "b11.ff.fc_in": 13.737499999999999,
+ "b11.ff.fc_out": 6227.2,
+ "b12.ff.fc_in": 14.874999999999998,
+ "b12.ff.fc_out": 1632.0,
+ "b13.ff.fc_in": 15.049999999999999,
+ "b13.ff.fc_out": 22144.0,
+ "b14.ff.fc_in": 7.46875,
+ "b14.ff.fc_out": 1552.0,
+ "b15.ff.fc_in": 9.4375,
+ "b15.ff.fc_out": 1444.8,
+ "b16.ff.fc_in": 11.272500000000008,
+ "b16.ff.fc_out": 3216.0,
+ "b17.ff.fc_in": 30.5,
+ "b17.ff.fc_out": 3760.0,
+ "b18.ff.fc_in": 11.4375,
+ "b18.ff.fc_out": 14656.0,
+ "b19.ff.fc_in": 12.5,
+ "b19.ff.fc_out": 4144.0,
+ "b20.ff.fc_in": 15.5,
+ "b20.ff.fc_out": 4544.0,
+ "b21.ff.fc_in": 10.149999999999999,
+ "b21.ff.fc_out": 9920.0,
+ "b22.ff.fc_in": 14.524999999999999,
+ "b22.ff.fc_out": 6656.0,
+ "b23.ff.fc_in": 18.75,
+ "b23.ff.fc_out": 8288.0,
+ "b24.ff.fc_in": 12.987500000000004,
+ "b24.ff.fc_out": 11520.0,
+ "b25.ff.fc_in": 14.625,
+ "b25.ff.fc_out": 9963.520000000004,
+ "b26.ff.fc_in": 10.0625,
+ "b26.ff.fc_out": 9390.080000000016,
+ "b27.ff.fc_in": 11.5625,
+ "b27.ff.fc_out": 6784.0,
+ "b28.ff.fc_in": 10.587499999999999,
+ "b28.ff.fc_out": 9900.8,
+ "b29.ff.fc_in": 8.44375,
+ "b29.ff.fc_out": 11991.040000000008,
+ "b30.ff.fc_in": 10.25,
+ "b30.ff.fc_out": 15411.199999999999,
+ "b31.ff.fc_in": 11.355000000000004,
+ "b31.ff.fc_out": 8576.0,
+ "b32.ff.fc_in": 12.9375,
+ "b32.ff.fc_out": 12928.0,
+ "b33.ff.fc_in": 11.125,
+ "b33.ff.fc_out": 9907.200000000004,
+ "b34.ff.fc_in": 10.230000000000004,
+ "b34.ff.fc_out": 14656.0,
+ "b35.ff.fc_in": 13.75,
+ "b35.ff.fc_out": 15360.0,
+ "b36.ff.fc_in": 14.417500000000004,
+ "b36.ff.fc_out": 29568.0,
+ "b37.ff.fc_in": 11.4375,
+ "b37.ff.fc_out": 12759.040000000008,
+ "b38.ff.fc_in": 19.085000000000008,
+ "b38.ff.fc_out": 30848.0,
+ "b39.ff.fc_in": 8.44375,
+ "b39.ff.fc_out": 290816.0,
+ "b40.ff.fc_in": 12.272500000000008,
+ "b40.ff.fc_out": 20787.199999999997,
+ "b41.ff.fc_in": 10.8125,
+ "b41.ff.fc_out": 18887.68000000014,
+ "b42.ff.fc_in": 11.06875,
+ "b42.ff.fc_out": 11264.0,
+ "b43.ff.fc_in": 16.275,
+ "b43.ff.fc_out": 25728.0,
+ "b44.ff.fc_in": 36.75,
+ "b44.ff.fc_out": 54476.799999999996,
+ "b45.ff.fc_in": 45.51000000000005,
+ "b45.ff.fc_out": 118272.0,
+ "b46.ff.fc_in": 40.20000000000002,
+ "b46.ff.fc_out": 20070.399999999998,
+ "b47.ff.fc_in": 54.0,
+ "b47.ff.fc_out": 31232.0,
+ "b48.ff.fc_in": 113.5,
+ "b48.ff.fc_out": 40960.0,
+ "b49.ff.fc_in": 128.0,
+ "b49.ff.fc_out": 167731.19999999998
+}

examples/train/configs/knowledge_distillation/minimax_h3/omniref_nvfp4_qad.yaml

@@ -0,0 +1,135 @@
+# FastH3 OmniRef NVFP4 QAD (the shared FP4 QAD recipe, Ref2VA profile).
+#
+# Student: OmniRef Ref2VA (8-forward PDD ladder) with the deployed NVFP4
+# numerics active from step 0: NVFP4 W4A4 q/k/v/out plus FFN fc_in/fc_out in
+# all 50 blocks, and the FP4 sparse-attention emulator at VSA tile 128. The
+# VSA compression gates and every norm stay frozen at the calibrated bf16
+# values; training them drifts the tile selection and causes ghosting. The
+# teacher is the same checkpoint in bf16.
+#
+# See docs/training/minimax_h3_fp4_qad.md for the recipe shared with the V2
+# (T2AV QAD) profile, calibration, and the checkpoint contract.
+#
+# Released result: FastVideo/FastH3-OmniRef-NVFP4-QAD (step 120 of 200).
+# `init_from` is the calibrated OmniRef checkpoint used as QAD step 0,
+# not that released result. The calibrated amax table is vendored under
+# examples/train/calibration/fasth3/ and pinned by sha256 below.
+models:
+  student:
+    _target_: fastvideo.train.models.minimax_h3.omniref_pdd.MiniMaxH3OmniRefPDDModel
+    init_from: data/models/omniref-nvfp4-ptq
+    trainable: true
+    enable_gradient_checkpointing_type: full
+    qad:
+      enabled: true
+      linears:
+        enabled: true
+        amax_json: examples/train/calibration/fasth3/amax-omniref-mse-bf16-top8.json
+        # Attention and gate linears use the unit activation scale; only the
+        # FFN linears read their calibrated amax from this table.
+        attention: true
+        gate: true
+        ffn: true
+        skip_blocks: []
+        amax_sha256: 14b93ec308faf0d89bd13ded507d287c030bb3005c96b8063213eb8a09df083d
+      fp4_attention:
+        enabled: true
+        two_level_p: true
+        smooth_k: true
+        first_block_max_floor: null
+    data:
+      manifests:
+      - data/omniref/train-manifest.jsonl
+      cases:
+      - first_frame
+      - first_last_frame
+      - storyboard
+      - continue_scene
+      - continue_shot
+      - full_scene_audio_image_reference
+      per_group:
+        480p: 400
+        768p: 50
+      seed: 20261007
+      max_frames: 243
+      eval_manifest: data/omniref/eval-manifest.json
+      eval_resolutions:
+      - 480p
+      - 768p
+      eval_per_group:
+        480p: 2
+        768p: 1
+      eval_max_frames: 0
+    # Only the projections and FFN linears train: the same set the deployed
+    # NVFP4 export quantizes. Gates, norms and modulation stay frozen.
+    trainable_patterns: transformer_blocks\.\d+\.(attn\.(to_q|to_k|to_v|to_out)|ff\.(fc_in|fc_out))\.
+  teacher:
+    _target_: fastvideo.train.models.minimax_h3.omniref_pdd.MiniMaxH3OmniRefPDDModel
+    init_from: data/models/omniref-nvfp4-ptq
+    trainable: false
+    frozen_bf16: true
+
+method:
+  _target_: fastvideo.train.methods.knowledge_distillation.pdd_qad.OmniRefPDDQADMethod
+  video_weight: 1.0
+  audio_weight: 1.0
+  rollout_fraction: 0.25
+  rollout_start_step: 20
+  rungs_per_row: 4
+  eval_every: 15
+  full_eval_every: 30
+  patience: 2
+  fidelity_tolerance: 0.05
+  require_bitwise_init: true
+  save_best: true
+  select_metric: score
+  eval_save_dir: outputs/omniref_nvfp4_qad/samples
+  objective: output+recon
+  recon_lambda: 1.0
+  recon_normalize: step0
+  max_rollout_rows: 10
+  noise_floor_eval: true
+  step0_cache: outputs/omniref_nvfp4_qad/step0_cache.json
+
+training:
+  distributed:
+    num_gpus: 16
+    sp_size: 4
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 16
+    pin_cpu_memory: false
+  data:
+    train_batch_size: 1
+    seed: 1234
+  optimizer:
+    learning_rate: 5.0e-06
+    betas: [0.9, 0.999]
+    eps: 1.0e-08
+    weight_decay: 0.0
+    lr_scheduler: cosine_with_min_lr
+    lr_warmup_steps: 10
+    min_lr_ratio: 0.3
+  loop:
+    max_train_steps: 200
+    gradient_accumulation_steps: 1
+  checkpoint:
+    output_dir: outputs/omniref_nvfp4_qad
+    training_state_checkpointing_steps: 25
+    checkpoints_total_limit: 2
+    resume_from_checkpoint: ''
+  tracker:
+    trackers: [wandb]
+    project_name: fasth3-omniref-nvfp4
+    run_name: omniref-nvfp4-qad
+  model:
+    enable_gradient_checkpointing_type: full
+  dit_precision: fp32
+  vsa_sparsity: 0.0
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+
+pipeline: {}

examples/train/configs/knowledge_distillation/minimax_h3/v2_nvfp4_qad.yaml

@@ -0,0 +1,149 @@
+# FastH3 V2 NVFP4 QAD (the shared FP4 QAD recipe, T2VA profile).
+#
+# Student: V2 bf16 weights with the deployed NVFP4 numerics active from step 0:
+# NVFP4 W4A4 q/k/v/out plus FFN fc_in/fc_out in all 50 blocks, and the FP4
+# sparse-attention emulator at VSA tile 64. The VSA compression gates and every
+# norm stay frozen at the V2 bf16 values. The teacher is the same weights in
+# bf16. Middle-block token compression is disabled: pure QAD.
+#
+# This is one recipe with two profiles; see docs/training/minimax_h3_fp4_qad.md
+# for the shared objective, evaluation and checkpoint contract, and for the
+# OmniRef (Ref2VA, tile 128) config.
+#
+# Paths are placeholders. Fill in the bf16 V2 model, the calibration table
+# shipped with it, the latent manifests and the output dir.
+# Starting weights: FastVideo/FastVideo-FastH3-8-Step-V2 (bf16). The released
+# QAD result is FastVideo/FastH3-QAD-v2; the
+# NVFP4 baseline it replaces is
+# FastVideo/FastVideo-FastH3-8-Step-V2-NVFP4-Consumer.
+models:
+  student:
+    _target_: fastvideo.train.models.minimax_h3.v2_t2v.MiniMaxH3V2T2VModel
+    init_from: data/models/FastH3-V2
+    trainable: true
+    enable_gradient_checkpointing_type: full
+    qad:
+      enabled: true
+      linears:
+        enabled: true
+        amax_json: examples/train/calibration/fasth3/amax-consumer.json
+        # Attention and gate linears use the unit activation scale; only the
+        # FFN linears read their calibrated amax from this table.
+        attention: true
+        gate: true
+        ffn: true
+        skip_blocks: []
+        amax_sha256: 0800516caec04ec829d21c45ca8b18eede6b81008c040997c0b84d48e0d39e52
+      fp4_attention:
+        enabled: true
+        two_level_p: true
+        smooth_k: true
+        first_block_max_floor: null
+    compression:
+      enabled: false
+      start: 3
+      end: 41
+      mode: subsample
+    data:
+      train_manifest: data/fasth3-v2/train.jsonl
+      eval_manifest: data/fasth3-v2/eval.jsonl
+      width: 832
+      height: 480
+      num_frames: 124
+    # Only the projections and FFN linears train: the same set the deployed
+    # NVFP4 export quantizes. Gates, norms and modulation stay frozen.
+    trainable_patterns: transformer_blocks\.\d+\.(attn\.(to_q|to_k|to_v|to_out)|ff\.(fc_in|fc_out))\.
+  teacher:
+    _target_: fastvideo.train.models.minimax_h3.v2_t2v.MiniMaxH3V2T2VModel
+    init_from: data/models/FastH3-V2
+    trainable: false
+
+method:
+  _target_: fastvideo.train.methods.knowledge_distillation.pdd_qad.OmniRefPDDQADMethod
+  video_weight: 1.0
+  audio_weight: 1.0
+  rollout_fraction: 0.25
+  rollout_start_step: 20
+  rungs_per_row: 4
+  eval_every: 25
+  full_eval_every: 50
+  patience: 3
+  fidelity_tolerance: 0.05
+  require_bitwise_init: true
+  save_best: true
+  select_metric: score
+  eval_save_dir: outputs/v2_1_nvfp4_qad/samples
+  objective: output+recon
+  recon_lambda: 1.0
+  recon_normalize: step0
+  max_rollout_rows: 8
+  noise_floor_eval: true
+  # A checkpoint is eligible for "best" only if audio stays inside its
+  # step-0 guardrail; the ELBO-style score alone can trade audio for video.
+  eligibility_keys:
+  - endpoint/audio
+  - x0_rel_l2/audio/mean
+  step0_cache: outputs/v2_1_nvfp4_qad/step0_cache.json
+
+training:
+  distributed:
+    num_gpus: 16
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 16
+    pin_cpu_memory: false
+  data:
+    train_batch_size: 1
+    seed: 1234
+  optimizer:
+    learning_rate: 5.0e-06
+    betas: [0.9, 0.999]
+    eps: 1.0e-08
+    weight_decay: 0.0
+    lr_scheduler: cosine_with_min_lr
+    lr_warmup_steps: 10
+    min_lr_ratio: 0.3
+  loop:
+    max_train_steps: 300
+    gradient_accumulation_steps: 1
+  checkpoint:
+    output_dir: outputs/v2_1_nvfp4_qad
+    training_state_checkpointing_steps: 50
+    checkpoints_total_limit: 1
+    resume_from_checkpoint: ''
+  tracker:
+    trackers: [wandb]
+    project_name: fasth3-v2-nvfp4-qad
+    run_name: v2-nvfp4-qad
+  model:
+    enable_gradient_checkpointing_type: full
+  dit_precision: fp32
+  vsa_sparsity: 0.0
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+
+pipeline: {}
+
+# The published FastVideo/FastH3-QAD-v2
+# checkpoint is the linears-only variant of this schedule: NVFP4 linears with
+# BF16 VSA attention, i.e. the same config with fp4_attention.enabled=false.
+# The shared FP4 recipe above adds FP4 attention from step 0 and does not
+# reproduce that checkpoint bit for bit.
+#
+# FP4-attention adaptation (provenance note, not part of the public
+# reproduction). The internal follow-up started from a linears-only 300-step
+# DCP, enabled FP4 attention, and trained 300 more steps at lr 1.5e-6 with no
+# warmup and a constant LR. It is an adaptation of the linears-only checkpoint
+# to FP4 attention, not a second epoch of the FP4 stage 1 above, and it needs
+# a DCP of the linears-only model (the packed HF release is not one):
+#
+#   method.init_student_dcp=<linears-only 300-step DCP>
+#   training.optimizer.learning_rate=1.5e-06
+#   training.optimizer.lr_warmup_steps=0
+#   training.optimizer.min_lr_ratio=1.0
+#   training.loop.max_train_steps=300
+#   training.checkpoint.output_dir=outputs/v2_1_nvfp4_qad_fp4_adapt

examples/train/configs/knowledge_distillation/pika_h3_flex/fl2va_nvfp4_qad.yaml

@@ -0,0 +1,126 @@
+# Pika H3 Flex FL2VA NVFP4 QAD (the FastH3 QAD recipe on Flex's six-forward ladder).
+#
+# Teacher: Pika H3 Flex FL2VA in bf16, i.e. MiniMax-H3 `transformer` with the
+# released `fl2va` adapter merged by
+# scripts/distill/pika_h3_flex_qad/merge_pika_adapter.py. Student: the same
+# weights with NVFP4 W4A4 q/k/v/out plus FFN fc_in/fc_out in all 50 blocks,
+# static FFN scales from the Flex FL2VA calibration table, BF16 dense attention.
+#
+# Rows are T2VA / I2VA / L2VA / FL2VA keyframe-anchor rows built from the
+# OmniRef v4 rows by scripts/distill/pika_h3_flex_qad/build_flex_ref2va_rows.py
+# --mode fl2va (exact first/last frames, re-encoded at the canvas size).
+#
+# Paths are placeholders: the merged model, the calibration table (pin its
+# sha256 once calibrated), the prompt manifests and the output dir.
+models:
+  student:
+    _target_: fastvideo.train.models.minimax_h3.flex.MiniMaxH3FlexFL2VAModel
+    init_from: data/models/pika-flex-fl2va-merged
+    trainable: true
+    enable_gradient_checkpointing_type: full
+    qad:
+      enabled: true
+      linears:
+        enabled: true
+        amax_json: examples/train/calibration/pika_h3_flex/amax-fl2va.json
+        amax_sha256: null
+        # Attention linears use the unit activation scale; only the FFN
+        # linears read their calibrated amax from this table.
+        attention: true
+        ffn: true
+        skip_blocks: []
+      fp4_attention:
+        enabled: false
+    data:
+      manifests:
+      - data/flex-fl2va/MANIFEST.jsonl
+      cases:
+      - t2va
+      - i2va
+      - l2va
+      - fl2va
+      per_group:
+        480p: 160
+        768p: 40
+      seed: 20261007
+      max_frames: 345
+      # FL2VA rows hold out their own eval rows (top-up from the groups); the listed manifest is empty.
+      eval_manifest: data/flex-fl2va/eval-manifest.json
+      eval_resolutions:
+      - 480p
+      - 768p
+      eval_per_group:
+        480p: 2
+        768p: 1
+      eval_max_frames: 0
+  teacher:
+    _target_: fastvideo.train.models.minimax_h3.flex.MiniMaxH3FlexFL2VAModel
+    init_from: data/models/pika-flex-fl2va-merged
+    trainable: false
+    frozen_bf16: true
+
+method:
+  _target_: fastvideo.train.methods.knowledge_distillation.pdd_qad.OmniRefPDDQADMethod
+  video_weight: 1.0
+  audio_weight: 1.0
+  rollout_fraction: 0.25
+  rollout_start_step: 20
+  # Half of the six forwards train per row, as OmniRef trains 4 of 8.
+  rungs_per_row: 3
+  eval_every: 15
+  full_eval_every: 30
+  patience: 2
+  fidelity_tolerance: 0.05
+  require_bitwise_init: true
+  save_best: true
+  select_metric: score
+  eval_save_dir: outputs/pika_flex_fl2va_nvfp4_qad/samples
+  objective: output+recon
+  recon_lambda: 1.0
+  recon_normalize: step0
+  max_rollout_rows: 10
+  noise_floor_eval: true
+  step0_cache: outputs/pika_flex_fl2va_nvfp4_qad/step0_cache.json
+
+training:
+  distributed:
+    num_gpus: 16
+    sp_size: 4
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 16
+    pin_cpu_memory: false
+  data:
+    train_batch_size: 1
+    seed: 1234
+  optimizer:
+    learning_rate: 5.0e-06
+    betas: [0.9, 0.999]
+    eps: 1.0e-08
+    weight_decay: 0.0
+    lr_scheduler: cosine_with_min_lr
+    lr_warmup_steps: 10
+    min_lr_ratio: 0.3
+  loop:
+    max_train_steps: 200
+    gradient_accumulation_steps: 1
+  checkpoint:
+    output_dir: outputs/pika_flex_fl2va_nvfp4_qad
+    training_state_checkpointing_steps: 25
+    checkpoints_total_limit: 2
+    resume_from_checkpoint: ''
+  tracker:
+    trackers: [wandb]
+    project_name: pika-h3-flex-nvfp4
+    run_name: pika-flex-fl2va-nvfp4-qad
+  model:
+    enable_gradient_checkpointing_type: full
+  dit_precision: fp32
+  vsa_sparsity: 0.0
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+
+pipeline: {}

examples/train/configs/knowledge_distillation/pika_h3_flex/ref2va_nvfp4_qad.yaml

@@ -0,0 +1,136 @@
+# Pika H3 Flex Ref2VA NVFP4 QAD (the FastH3 QAD recipe on Flex's six-forward ladder).
+#
+# Teacher: Pika H3 Flex Ref2VA in bf16, i.e. MiniMax-H3 `transformer_ref` with the
+# released `ref2va` adapter merged by
+# scripts/distill/pika_h3_flex_qad/merge_pika_adapter.py. Student: the same
+# weights with NVFP4 W4A4 q/k/v/out plus FFN fc_in/fc_out in all 50 blocks
+# active from step 0, static FFN activation scales from the Flex calibration
+# table. Attention stays BF16 dense: Flex has no VSA gate, so there is no FP4
+# sparse-attention stage and no gate to freeze. Norms, AdaLN, the token refiner
+# and the FP32 boundary modules stay frozen.
+#
+# The objective, evaluation, guards and checkpoint contract are the shared
+# recipe's (docs/training/minimax_h3_fp4_qad.md). Rows are OmniRef Ref2VA
+# parquet rows; the teacher computes every target on the fly, so the OmniRef
+# held-out eval manifest is reused as is.
+#
+# Paths are placeholders: the merged model, the calibration table (pin its
+# sha256 once calibrated), the row manifests and the output dir.
+models:
+  student:
+    _target_: fastvideo.train.models.minimax_h3.flex.MiniMaxH3FlexRef2VAModel
+    init_from: data/models/pika-flex-ref2va-merged
+    trainable: true
+    enable_gradient_checkpointing_type: full
+    qad:
+      enabled: true
+      linears:
+        enabled: true
+        amax_json: examples/train/calibration/pika_h3_flex/amax-ref2va-balanced.json
+        amax_sha256: ad60b45fea00906c55e370a0dec82638eb413d6b68537d1ceb38fc52da90b373
+        # Attention linears use the unit activation scale; only the FFN
+        # linears read their calibrated amax from this table.
+        attention: true
+        ffn: true
+        skip_blocks: []
+      fp4_attention:
+        enabled: false
+    data:
+      manifests:
+      # Flex-sized rows (build_flex_ref2va_rows.py) first: they replace v4 rows with the same id.
+      - data/flex-ref2va/MANIFEST.jsonl
+      - data/omniref/train-manifest.jsonl
+      cases:
+      - first_frame
+      - first_last_frame
+      - storyboard
+      - continue_scene
+      - continue_shot
+      - audio_reference
+      - mixed_image_video
+      - many_images
+      - multi_video
+      per_group:
+        480p: 160
+        768p: 40
+      seed: 20261007
+      max_frames: 345
+      eval_manifest: data/flex-ref2va/eval-manifest.json
+      eval_resolutions:
+      - 480p
+      - 768p
+      eval_per_group:
+        480p: 2
+        768p: 1
+      eval_max_frames: 0
+  teacher:
+    _target_: fastvideo.train.models.minimax_h3.flex.MiniMaxH3FlexRef2VAModel
+    init_from: data/models/pika-flex-ref2va-merged
+    trainable: false
+    frozen_bf16: true
+
+method:
+  _target_: fastvideo.train.methods.knowledge_distillation.pdd_qad.OmniRefPDDQADMethod
+  video_weight: 1.0
+  audio_weight: 1.0
+  rollout_fraction: 0.25
+  rollout_start_step: 20
+  # Half of the six forwards train per row, as OmniRef trains 4 of 8.
+  rungs_per_row: 3
+  eval_every: 15
+  full_eval_every: 30
+  patience: 2
+  fidelity_tolerance: 0.05
+  require_bitwise_init: true
+  save_best: true
+  select_metric: score
+  eval_save_dir: outputs/pika_flex_ref2va_nvfp4_qad/samples
+  objective: output+recon
+  recon_lambda: 1.0
+  recon_normalize: step0
+  max_rollout_rows: 10
+  noise_floor_eval: true
+  step0_cache: outputs/pika_flex_ref2va_nvfp4_qad/step0_cache.json
+
+training:
+  distributed:
+    num_gpus: 16
+    sp_size: 4
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 16
+    pin_cpu_memory: false
+  data:
+    train_batch_size: 1
+    seed: 1234
+  optimizer:
+    learning_rate: 5.0e-06
+    betas: [0.9, 0.999]
+    eps: 1.0e-08
+    weight_decay: 0.0
+    lr_scheduler: cosine_with_min_lr
+    lr_warmup_steps: 10
+    min_lr_ratio: 0.3
+  loop:
+    max_train_steps: 200
+    gradient_accumulation_steps: 1
+  checkpoint:
+    output_dir: outputs/pika_flex_ref2va_nvfp4_qad
+    training_state_checkpointing_steps: 25
+    checkpoints_total_limit: 2
+    resume_from_checkpoint: ''
+  tracker:
+    trackers: [wandb]
+    project_name: pika-h3-flex-nvfp4
+    run_name: pika-flex-ref2va-nvfp4-qad
+  model:
+    enable_gradient_checkpointing_type: full
+  dit_precision: fp32
+  vsa_sparsity: 0.0
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+
+pipeline: {}

fastvideo/api/compat.py

@@ -199,6 +199,8 @@ def legacy_from_pretrained_to_config(
             components["transformer_weights"] = value
         elif key == "init_weights_from_safetensors_2":
             components["transformer_2_weights"] = value
+        elif key == "h3_vae_nvfp4_decoder":
+            components["h3_vae_nvfp4_decoder"] = value
         else:
             experimental[key] = deepcopy(value)
 
@@ -331,6 +333,8 @@ def generator_config_to_fastvideo_args(config: GeneratorConfig | Mapping[str, An
         kwargs["init_weights_from_safetensors"] = components.transformer_weights
     if components.transformer_2_weights is not None:
         kwargs["init_weights_from_safetensors_2"] = components.transformer_2_weights
+    if components.h3_vae_nvfp4_decoder is not None:
+        kwargs["h3_vae_nvfp4_decoder"] = components.h3_vae_nvfp4_decoder
     if components.config_root is not None:
         kwargs["config_model_path"] = components.config_root
     if components.upsampler_weights is not None:

fastvideo/api/pika_h3_flex.py

@@ -0,0 +1,24 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Sampling defaults for Pika H3 Flex."""
+
+from dataclasses import dataclass
+
+from fastvideo.api.sampling_param import SamplingParam
+
+
+@dataclass
+class PikaH3FlexSamplingParam(SamplingParam):
+    """Six re-noise forwards, no guidance, and the 768p bucket by default."""
+
+    num_inference_steps: int = 6
+    guidance_scale: float = 1.0
+    batch_cfg: bool = False
+    negative_prompt: str = ""
+    height: int = 768
+    width: int = 1344
+    num_frames: int = 124
+    fps: int = 24
+    seed: int = 0
+
+
+__all__ = ["PikaH3FlexSamplingParam"]

fastvideo/api/schema.py

@@ -120,6 +120,7 @@ class ComponentConfig:
     transformer_weights: str | None = None
     transformer_2_weights: str | None = None
     vae_weights: str | None = None
+    h3_vae_nvfp4_decoder: str | None = None
     upsampler_weights: str | None = None
     lora_path: str | None = None
     lora_nickname: str = "default"

fastvideo/attention/backends/flash_attn.py

@@ -27,6 +27,7 @@ from fastvideo.attention.backends.abstract import (
     PackedSequenceMetadata,
 )
 from fastvideo.attention.backends.attn_qat_infer import _resolve_fa4_pv_mode
+from fastvideo.attention.utils.nvfp4_qk_precondition import parse_qk_transform, precondition_qk
 from fastvideo.logger import init_logger
 
 logger = init_logger(__name__)
@@ -279,6 +280,12 @@ class FlashAttentionImpl(AttentionImpl):
         if nvfp4_fa4 is None:
             nvfp4_fa4 = envs.FASTVIDEO_NVFP4_FA4.get()
         self.nvfp4_fa4 = bool(nvfp4_fa4)
+        # Exact Q/K preconditioning ahead of NVFP4 quantization, shared with the QAD emulator
+        # (fastvideo.attention.utils.nvfp4_qk_precondition). Same arg-over-env precedence as nvfp4_fa4.
+        qk_transform = extra_impl_args.get("nvfp4_qk_transform")
+        if qk_transform is None:
+            qk_transform = envs.FASTVIDEO_NVFP4_FA4_QK_TRANSFORM.get()
+        self.nvfp4_smooth_k, self.nvfp4_hadamard = parse_qk_transform(qk_transform)
         # PV-mode knob for the FA4-FP4 path. Resolved + validated
         # unconditionally (shared helper) so a typo fails at construction, not
         # at the first forward on a Blackwell box, and so the ATTN_QAT_INFER
@@ -291,7 +298,8 @@ class FlashAttentionImpl(AttentionImpl):
             assert cap in [(10, 0), (10, 3)], (f"NVFP4 FA4 requires Blackwell (sm100a/sm103a), got sm{cap[0]}{cap[1]}")
             assert _FA4_FP4_AVAILABLE, ("NVFP4 FA4 requires flash-attention-fp4 (flash_attn.cute). "
                                         "Install via instructions in docs/inference/optimizations.md")
-            logger.info("NVFP4 FA4 enabled for FlashAttentionImpl (quant_qk, pv_mode=%s)", self.fa4_pv_mode)
+            logger.info("NVFP4 FA4 enabled for FlashAttentionImpl (quant_qk, pv_mode=%s, smooth_k=%s, hadamard=%s)",
+                        self.fa4_pv_mode, self.nvfp4_smooth_k, self.nvfp4_hadamard)
 
     def forward(
         self,
@@ -522,6 +530,8 @@ class FlashAttentionImpl(AttentionImpl):
         orig_seqlen_q = query.shape[1]
         orig_seqlen_k = key.shape[1]
 
+        query, key = precondition_qk(query, key, smooth_k=self.nvfp4_smooth_k, hadamard=self.nvfp4_hadamard)
+
         # Quantize Q/K to FP4 (internally pads to multiple of 128 for SF layout)
         q_fp4, q_sf = _nvfp4_quantize_for_fa4(query)
         k_fp4, k_sf = _nvfp4_quantize_for_fa4(key)

fastvideo/attention/backends/fp4_dense_qat.py

@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Dense FP4 attention emulator for QAD of models deployed with FlashAttention-4 NVFP4 Q/K (dense attention).
+
+The deployed dense FP4 path (``FLASH_ATTN`` with ``nvfp4_fa4``) quantizes Q and K to NVFP4 per 16 channels and keeps
+P and V in BF16. The FP4 block-sparse training kernel (``fp4_vsa_attn_qat`` with ``quantize_pv=False``) computes the
+same numerics with straight-through gradients; given every key block for every query block it is dense attention.
+
+``install_fp4_dense_attention(transformer, numerics)`` routes every dense attention impl of ``transformer``
+(``FlashAttentionImpl``, or ``SDPAImpl`` where the ``FLASH_ATTN`` selection fell back to SDPA because flash-attn is
+not installed) through that kernel (``None`` restores the original forward). Inputs arrive as ``[B, L, H, D]`` after the sequence-parallel
+all-to-all (full sequence, local heads); they are padded to 128-token tiles, attended over the full block lists and
+cut back.
+"""
+from __future__ import annotations
+
+import re
+import types
+from typing import Any
+
+import torch
+
+from fastvideo.attention.backends.fp4_vsa_qat import FP4AttentionNumerics, load_qat_vsa_modules
+from fastvideo.attention.utils.nvfp4_qk_precondition import precondition_qk
+
+TILE = 128
+_ORIGINAL = "_fp4_dense_original_forward"
+_BLOCK = re.compile(r"(?:^|\.)transformer_blocks\.(\d+)\.")
+
+
+def dense_block_lists(seq_len: int, batch: int, heads: int, reference: Any,
+                      device: torch.device) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Kernel block lists for full (dense) attention over ``seq_len`` tokens in 128-token tiles."""
+    n_tiles = -(-seq_len // TILE)
+    tile_valid = torch.full((n_tiles, ), TILE, dtype=torch.int32, device=device)
+    tile_valid[-1] = seq_len - (n_tiles - 1) * TILE
+    tile_mask = torch.ones((batch, heads, n_tiles, n_tiles), dtype=torch.bool, device=device)
+    return reference.tile128_mask_to_fp4_blocks(tile_mask, tile_valid)
+
+
+class FP4DenseAttention:
+    """Dense attention ``forward`` replacement: ``(q, k, v, metadata) -> [B, L, H, D]`` through the FP4 kernel."""
+
+    def __init__(self,
+                 numerics: FP4AttentionNumerics,
+                 modules: tuple[Any, Any] | None = None,
+                 qk_transform: tuple[bool, bool] = (False, False)) -> None:
+        self.numerics = numerics
+        # (smooth K, Hadamard-rotate Q/K) applied before quantization, exactly as the deployed FA4 NVFP4 path does.
+        self.smooth_k_mean, self.hadamard = qk_transform
+        self._reference, self._train = modules or load_qat_vsa_modules()
+        self._lists: dict[tuple[int, int, int, str], tuple[torch.Tensor, ...]] = {}
+
+    def _block_lists(self, seq_len: int, batch: int, heads: int, device: torch.device) -> tuple[torch.Tensor, ...]:
+        key = (seq_len, batch, heads, str(device))
+        if key not in self._lists:
+            self._lists = {key: dense_block_lists(seq_len, batch, heads, self._reference, device)}
+        return self._lists[key]
+
+    def __call__(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, metadata: Any) -> torch.Tensor:
+        if getattr(metadata, "packed_sequence", None) is not None:
+            raise NotImplementedError("dense FP4 attention emulator does not support packed sequences")
+        batch, seq_len, heads, _ = query.shape
+        query, key = precondition_qk(query, key, smooth_k=self.smooth_k_mean, hadamard=self.hadamard)
+        pad = -seq_len % TILE
+        q, k, v = (torch.nn.functional.pad(t.to(torch.bfloat16).transpose(1, 2), (0, 0, 0, pad))
+                   for t in (query, key, value))
+        q2k_idx, q2k_num, kv_valid = self._block_lists(seq_len, batch, heads, query.device)
+        out = self._train.fp4_vsa_attn_qat(q,
+                                           k,
+                                           v,
+                                           q2k_idx,
+                                           q2k_num,
+                                           kv_valid,
+                                           quantize=self.numerics.quantize,
+                                           quantize_pv=self.numerics.quantize_pv,
+                                           two_level_p=self.numerics.two_level_p,
+                                           smooth_k=self.numerics.smooth_k,
+                                           first_block_max_floor=self.numerics.first_block_max_floor)
+        return out[:, :, :seq_len].transpose(1, 2).to(query.dtype)
+
+
+def _dense_impl_types() -> tuple[type, ...]:
+    from fastvideo.attention.backends.sdpa import SDPAImpl
+
+    try:
+        from fastvideo.attention.backends.flash_attn import FlashAttentionImpl
+    except (ImportError, RuntimeError):  # no flash-attn install: FLASH_ATTN selection resolves to SDPA
+        return (SDPAImpl, )
+    return (FlashAttentionImpl, SDPAImpl)
+
+
+def named_dense_attention_impls(transformer: torch.nn.Module) -> list[tuple[str, Any]]:
+    """``(module name, impl)`` for every dense (FlashAttention or SDPA) impl behind a distributed attention layer."""
+    types_ = _dense_impl_types()
+    impls = []
+    for name, module in transformer.named_modules():
+        impl = getattr(getattr(module, "distributed_attention", None), "attn_impl", None)
+        if isinstance(impl, types_):
+            impls.append((name, impl))
+    return impls
+
+
+def dense_attention_impls(transformer: torch.nn.Module) -> list[Any]:
+    """Every dense (FlashAttention or SDPA) impl reached through the transformer's distributed attention layers."""
+    return [impl for _, impl in named_dense_attention_impls(transformer)]
+
+
+def block_index(module_name: str) -> int | None:
+    """``transformer_blocks.<i>`` index of a module name (``None`` outside the block stack)."""
+    match = _BLOCK.search(module_name + ".")
+    return int(match.group(1)) if match else None
+
+
+def install_fp4_dense_attention(transformer: torch.nn.Module,
+                                numerics: FP4AttentionNumerics | None,
+                                modules: tuple[Any, Any] | None = None,
+                                *,
+                                qk_transform: tuple[bool, bool] = (False, False),
+                                bf16_blocks: tuple[int, ...] = ()) -> int:
+    """Route dense attention layers through the FP4 emulator; returns how many now run FP4.
+
+    ``numerics=None`` restores every original forward. ``bf16_blocks`` keeps the attention of those
+    ``transformer_blocks`` indices on its original (BF16) forward, matching a deployment that opts them out of NVFP4.
+    """
+    override = FP4DenseAttention(numerics, modules, qk_transform) if numerics is not None else None
+    named = named_dense_attention_impls(transformer)
+    skip = frozenset(int(b) for b in bf16_blocks)
+    missing = skip - {block_index(name) for name, _ in named}
+    if missing:
+        raise ValueError(f"bf16_blocks {sorted(missing)} match no transformer_blocks attention layer")
+    installed = 0
+    for name, impl in named:
+        if _ORIGINAL not in impl.__dict__:
+            impl.__dict__[_ORIGINAL] = impl.forward
+        if override is None or block_index(name) in skip:
+            impl.forward = impl.__dict__[_ORIGINAL]
+        else:
+            impl.forward = types.MethodType(lambda self, q, k, v, meta, _o=override: _o(q, k, v, meta), impl)
+            installed += 1
+    return installed if override is not None else len(named)
+
+
+__all__ = [
+    "FP4DenseAttention", "block_index", "dense_block_lists", "dense_attention_impls", "install_fp4_dense_attention",
+    "named_dense_attention_impls"
+]

fastvideo/attention/backends/fp4_vsa_qat.py

@@ -105,4 +105,23 @@ class FP4VSAFineAttention:
         return out.transpose(1, 2)
 
 
-__all__ = ["FP4AttentionNumerics", "FP4VSAFineAttention", "load_qat_vsa_modules"]
+def vsa_impls(transformer: torch.nn.Module) -> list[Any]:
+    """Every VSA-H3 attention implementation in a MiniMax-H3 transformer."""
+    impls = []
+    for module in transformer.modules():
+        impl = getattr(getattr(module, "distributed_attention", None), "attn_impl", None)
+        if impl is not None and hasattr(impl, "fine_attention_override"):
+            impls.append(impl)
+    return impls
+
+
+def install_fp4_vsa_attention(transformer: torch.nn.Module, numerics: FP4AttentionNumerics | None) -> int:
+    """Route every VSA-H3 layer's fine attention through the emulator (``None`` restores the kernels)."""
+    override = FP4VSAFineAttention(numerics) if numerics is not None else None
+    impls = vsa_impls(transformer)
+    for impl in impls:
+        impl.fine_attention_override = override
+    return len(impls)
+
+
+__all__ = ["FP4AttentionNumerics", "FP4VSAFineAttention", "install_fp4_vsa_attention", "load_qat_vsa_modules"]

fastvideo/attention/utils/nvfp4_qk_precondition.py

@@ -0,0 +1,79 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Exact Q/K preconditioning ahead of NVFP4 (per-16 block) Q/K quantization.
+
+Both transforms leave the attention output unchanged in exact arithmetic and only redistribute what the FP4 blocks
+have to represent:
+
+* ``smooth_k``: subtract each head's token-mean of K. ``q . k_mean`` is the same for every key of a query row, and
+  softmax is shift-invariant per row.
+* ``hadamard``: rotate Q and K by the same orthonormal head-dim Hadamard matrix ``H`` (``Q H (K H)^T = Q K^T``),
+  which spreads outlier channels across the 16-wide scale blocks.
+
+The QAD emulator and the deployed FlashAttention-4 NVFP4 path apply the same function, so training and inference
+quantize identical tensors.
+"""
+from __future__ import annotations
+
+import math
+from functools import lru_cache
+
+import torch
+
+QK_TRANSFORMS = {
+    "none": (False, False),
+    "smooth_k": (True, False),
+    "hadamard": (False, True),
+    "smooth_k+hadamard": (True, True),
+}
+
+
+def parse_qk_transform(name: str | None) -> tuple[bool, bool]:
+    """``"smooth_k+hadamard"`` -> ``(smooth_k, hadamard)``; ``None``/``"none"`` -> no transform."""
+    key = (name or "none").strip().lower()
+    if key not in QK_TRANSFORMS:
+        raise ValueError(f"unknown NVFP4 Q/K transform {name!r}; expected one of {', '.join(QK_TRANSFORMS)}")
+    return QK_TRANSFORMS[key]
+
+
+def parse_block_indices(raw: str | None) -> frozenset[int]:
+    """``"45, 49"`` -> ``frozenset({45, 49})``; empty/``None`` -> empty set."""
+    items = [item.strip() for item in (raw or "").split(",") if item.strip()]
+    try:
+        indices = frozenset(int(item) for item in items)
+    except ValueError as error:
+        raise ValueError(f"expected a comma-separated list of block indices, got {raw!r}") from error
+    if any(index < 0 for index in indices):
+        raise ValueError(f"block indices must be non-negative, got {raw!r}")
+    return indices
+
+
+@lru_cache(maxsize=8)
+def _hadamard(n: int, device: torch.device) -> torch.Tensor:
+    if n <= 0 or n & (n - 1):
+        raise ValueError(f"Hadamard rotation needs a power-of-two head dim, got {n}")
+    h = torch.ones(1, 1)
+    while h.shape[0] < n:
+        h = torch.cat([torch.cat([h, h], 1), torch.cat([h, -h], 1)], 0)
+    return (h / math.sqrt(n)).to(device)
+
+
+def hadamard_matrix(n: int, device: torch.device | str) -> torch.Tensor:
+    """Orthonormal Sylvester Hadamard matrix of size ``n`` (fp32)."""
+    return _hadamard(n, torch.device(device))
+
+
+def precondition_qk(query: torch.Tensor, key: torch.Tensor, *, smooth_k: bool,
+                    hadamard: bool) -> tuple[torch.Tensor, torch.Tensor]:
+    """Transform ``[B, L, H, D]`` Q/K (computed in fp32, returned in the input dtype)."""
+    if not (smooth_k or hadamard):
+        return query, key
+    q, k = query.float(), key.float()
+    if smooth_k:
+        k = k - k.mean(dim=1, keepdim=True)
+    if hadamard:
+        rotation = hadamard_matrix(q.shape[-1], q.device)
+        q, k = q @ rotation, k @ rotation
+    return q.to(query.dtype), k.to(key.dtype)
+
+
+__all__ = ["QK_TRANSFORMS", "hadamard_matrix", "parse_block_indices", "parse_qk_transform", "precondition_qk"]

fastvideo/configs/models/dits/pika_h3_flex.py

@@ -0,0 +1,21 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pika H3 Flex transformer configuration.
+
+The released Pika adapters do not change MiniMax-H3's transformer: the
+checkpoint keeps the same packed video/audio trunk, Qwen3-VL conditioner, and
+VAEs, and applies a rank-128 PEFT adapter at inference.  This subclass exists so
+the Pika pipeline names its own component config without duplicating the H3
+architecture contract.
+"""
+
+from dataclasses import dataclass
+
+from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3Config
+
+
+@dataclass
+class PikaH3FlexDiTConfig(MiniMaxH3Config):
+    """The released MiniMax-H3 transformer, sampled through Pika's adapters."""
+
+
+__all__ = ["PikaH3FlexDiTConfig"]

fastvideo/configs/pipelines/pika_h3_flex.py

@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pipeline configuration for Pika H3 Flex.
+
+Pika H3 Flex is a post-trained MiniMax-H3: the transformer, Qwen3-VL
+conditioner, and video/audio VAEs are the released H3 components, and a rank-128
+adapter is merged at load time.  What changes is the sampler and the geometry
+contract:
+
+* six transformer forwards over the fixed clean-time knots
+  ``1000, 750, 500, 250, 100, 40``, with a fresh-noise re-noise step between
+  forwards (the terminal ``0`` knot is the output, never a forward);
+* video and audio clock shifts of 12 and 3;
+* 480p and 768p canvas buckets instead of a free-form canvas;
+* center-cropped keyframes and tier-dependent reference-image sizing.
+
+Everything else, including CFG being off and the 24 fps frame grid, is inherited
+from MiniMax-H3.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+from fastvideo.configs.models.dits.pika_h3_flex import PikaH3FlexDiTConfig
+from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig
+
+if TYPE_CHECKING:
+    from fastvideo.api.sampling_param import SamplingParam
+    from fastvideo.api.schema import GenerationRequest
+
+# The released re-noise ladder. The terminal 0 is a sigma-grid point, not a
+# transformer forward: six forwards, seven sigmas.
+PIKA_DENOISING_STEPS: tuple[int, ...] = (1000, 750, 500, 250, 100, 40)
+PIKA_VIDEO_SCHEDULER_SHIFT = 12.0
+PIKA_AUDIO_SCHEDULER_SHIFT = 3.0
+PIKA_RESOLUTION_TIERS: tuple[str, ...] = ("p480", "p768")
+PIKA_DEFAULT_RESOLUTION_TIER = "p768"
+
+
+@dataclass
+class PikaH3FlexPipelineConfig(MiniMaxH3PipelineConfig):
+    """MiniMax-H3 components under Pika's adapter, sampler, and bucket contract."""
+
+    dit_config: PikaH3FlexDiTConfig = field(default_factory=PikaH3FlexDiTConfig)
+    # Pika's CLI calls this ``--resolution``; the API has no dedicated tier
+    # field, so it is a pipeline config knob rather than a per-request one.
+    pika_resolution_tier: str = PIKA_DEFAULT_RESOLUTION_TIER
+    pika_denoising_steps: tuple[int, ...] = PIKA_DENOISING_STEPS
+
+    def check_pipeline_config(self) -> None:
+        super().check_pipeline_config()
+        if self.pika_resolution_tier not in PIKA_RESOLUTION_TIERS:
+            raise ValueError(f"Pika H3 Flex supports resolution tiers {list(PIKA_RESOLUTION_TIERS)}; got "
+                             f"{self.pika_resolution_tier!r}.")
+        if tuple(self.pika_denoising_steps) != PIKA_DENOISING_STEPS:
+            raise ValueError(f"Pika H3 Flex's adapter contract fixes the denoising ladder at "
+                             f"{list(PIKA_DENOISING_STEPS)}; got {list(self.pika_denoising_steps)}.")
+        if self.dmd_denoising_steps is not None:
+            raise ValueError("Pika H3 Flex uses its own six-forward re-noise schedule; dmd_denoising_steps must "
+                             "remain unset.")
+        if self.pdd_step_indices is not None:
+            raise ValueError("Pika H3 Flex is not a FastH3 PDD student; pdd_step_indices must remain unset.")
+
+    def apply_request_constraints(self, request: GenerationRequest, sampling_param: SamplingParam) -> SamplingParam:
+        """Fix the six-forward ladder and fill an omitted canvas from ``pika_resolution_tier``.
+
+        ``PikaH3FlexSamplingParam`` and the presets carry the 768p canvas, so by
+        input preparation an omitted size is indistinguishable from an explicit
+        1344x768. The request's explicit-field tracking is the last place the two
+        differ: an omitted size takes the configured tier's 16:9 bucket, while an
+        explicit size still wins and picks its own tier.
+        """
+        from fastvideo.api.compat import explicit_request_updates
+
+        explicit = explicit_request_updates(request)
+        steps = len(PIKA_DENOISING_STEPS)
+        requested = explicit.get("num_inference_steps", steps)
+        if requested != steps:
+            raise ValueError(f"Pika H3 Flex runs exactly {steps} transformer forwards; the request sets "
+                             f"num_inference_steps={requested}. The six-step ladder is part of the adapter "
+                             f"contract, not a quality/speed knob.")
+        sampling_param.num_inference_steps = steps
+
+        height, width = explicit.get("height"), explicit.get("width")
+        if (height is None) != (width is None):
+            raise ValueError("Pika H3 Flex `height` and `width` must be passed together, or neither; the request "
+                             f"sets height={height}, width={width}.")
+        if height is None:
+            from fastvideo.pipelines.basic.pika_h3_flex.buckets import DEFAULT_ASPECT, canvas_for_aspect
+
+            sampling_param.height, sampling_param.width = canvas_for_aspect(DEFAULT_ASPECT, self.pika_resolution_tier)
+        return sampling_param
+
+
+__all__ = [
+    "PIKA_AUDIO_SCHEDULER_SHIFT",
+    "PIKA_DEFAULT_RESOLUTION_TIER",
+    "PIKA_DENOISING_STEPS",
+    "PIKA_RESOLUTION_TIERS",
+    "PIKA_VIDEO_SCHEDULER_SHIFT",
+    "PikaH3FlexPipelineConfig",
+]

fastvideo/envs.py

@@ -267,6 +267,12 @@ FASTVIDEO_HOST_IP = EnvStr(
     category="distributed",
     doc="IP address of this node when the node has several network interfaces. Set it on each node for multi-node "
     "inference.")
+FASTVIDEO_PG_TIMEOUT_MIN = EnvInt(
+    0,
+    category="distributed",
+    doc="Timeout in minutes for the training process groups' collectives (0 keeps PyTorch's 10-minute default). "
+    "Raise it for large models whose rank-0 load and scatter, checkpoint saves or long evaluations outlast the "
+    "default.")
 FASTVIDEO_LOOPBACK_IP = EnvStr("",
                                category="distributed",
                                doc="Loopback IP address to use instead of the detected one.")
@@ -362,6 +368,19 @@ FASTVIDEO_NVFP4_FA4 = EnvBool(
     category="attention",
     doc="FlashAttention-4 quantizes Q and K to NVFP4. An explicit nvfp4_fa4 attention implementation argument "
     "takes precedence.")
+FASTVIDEO_NVFP4_FA4_QK_TRANSFORM = EnvChoice(
+    "none",
+    choices=("none", "smooth_k", "hadamard", "smooth_k+hadamard"),
+    category="attention",
+    doc="Exact Q/K preconditioning applied right before NVFP4 quantization on the FlashAttention-4 NVFP4 path: "
+    "smooth_k subtracts each head's token-mean of K, hadamard rotates Q and K by the head-dim Hadamard matrix. "
+    "Must match the qk_transform the model was QAD-trained with. An explicit nvfp4_qk_transform attention "
+    "implementation argument takes precedence.")
+FASTVIDEO_NVFP4_FA4_BF16_BLOCKS = EnvStr(
+    "",
+    category="attention",
+    doc="Comma-separated MiniMax-H3 transformer_blocks indices whose attention stays in BF16 while "
+    "FASTVIDEO_NVFP4_FA4 quantizes the rest, for example 45,49.")
 FASTVIDEO_FA4_PV_MODE = EnvChoice(
     "bf16",
     choices=("bf16", "fp8"),

fastvideo/layers/quantization/nvfp4_qad.py

@@ -0,0 +1,258 @@
+# SPDX-License-Identifier: Apache-2.0
+"""NVFP4 quantization-aware distillation (QAD) for the MiniMax-H3 DiT block linears.
+
+``NVFP4QADLinearMethod`` replaces a ``ReplicatedLinear``'s ``quant_method``
+in place, so parameter names, FSDP sharding and checkpoints are unchanged.
+Its forward *is* the deployed NVFP4 GEMM, computed from the layer's current
+(BF16-cast) master weight exactly as the packed export and
+``NVFP4QuantizeMethod.apply`` compute it:
+
+* weight: RTN NVFP4 with global scale ``448 * 6 / amax(W)`` (the converter's
+  ``quantize_dense_linear``; ModelOpt max calibration gives the same codes);
+* activation: per-16 E4M3 block scales under a global scale that is
+  ``static`` (``448 * 6 / calibrated amax``, the FFN linears) or ``unit``
+  (1.0, attention projections and VSA gates);
+* GEMM: FlashInfer ``mm_fp4`` with ``alpha = (1 / weight_sf) / input_sf``.
+
+With gradients enabled the forward runs inside a straight-through estimator
+whose backward is full precision. ``enabled =
+False`` delegates to the original method, so a disabled student is the
+teacher bit for bit.
+"""
+from __future__ import annotations
+
+import contextlib
+import json
+import re
+from collections.abc import Iterator
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import torch
+
+FP4_RANGE = 448.0 * 6.0
+H3_BLOCK_LINEAR = re.compile(r"transformer_blocks\.(\d+)\.(attn\.to_q|attn\.to_k|attn\.to_v|attn\.to_out"
+                             r"|attn\.to_gate_compress|ff\.fc_in|ff\.fc_out)$")
+STATIC_SUFFIXES = ("ff.fc_in", "ff.fc_out")
+
+
+def weight_global_sf(weight: torch.Tensor) -> torch.Tensor:
+    """``448 * 6 / amax(|W|)`` in fp32, as the packed-export converter computes it.
+
+    The max of a BF16 tensor is exact in BF16, so one BF16 reduction gives the converter's value
+    (``weight.float().abs().nan_to_num().max()``) for finite weights without fp32 temporaries.
+    """
+    return FP4_RANGE / torch.nan_to_num(weight.abs().amax().float()).clamp(min=1e-12)
+
+
+def nvfp4_gemm(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None,
+               input_global_sf: torch.Tensor) -> torch.Tensor:
+    """The deployed NVFP4 linear on a BF16 weight: quantize both operands, FP4 GEMM."""
+    from fastvideo.layers.quantization.nvfp4_config import _mm_fp4, _mm_fp4_backend, _nvfp4_quantize
+    from fastvideo.layers.quantization.nvfp4_config import _require_flashinfer
+
+    sf_layout, _, _ = _require_flashinfer()
+    weight = weight.to(torch.bfloat16).contiguous()
+    w_sf = weight_global_sf(weight)
+    w_fp4, w_scale = _nvfp4_quantize(weight, w_sf, sfLayout=sf_layout.layout_128x4, do_shuffle=False)
+    alpha = (1.0 / w_sf) / input_global_sf
+    shape = x.shape
+    x2d = x.to(torch.bfloat16).reshape(-1, shape[-1]).contiguous()
+    x_fp4, x_scale = _nvfp4_quantize(x2d, input_global_sf, sfLayout=sf_layout.layout_128x4, do_shuffle=False)
+    out = _mm_fp4(x_fp4, w_fp4.T, x_scale, w_scale.T, alpha, torch.bfloat16, None, backend=_mm_fp4_backend())
+    if bias is not None:
+        out = out + bias.to(out.dtype)
+    return out.view(*shape[:-1], weight.shape[0])
+
+
+class _NVFP4STE(torch.autograd.Function):
+    """Forward = ``nvfp4_gemm``; backward = full-precision linear gradients (straight-through)."""
+
+    @staticmethod
+    def forward(ctx, x, weight, bias, input_global_sf):  # type: ignore[override]
+        ctx.save_for_backward(x, weight)
+        ctx.has_bias = bias is not None
+        return nvfp4_gemm(x, weight, bias, input_global_sf)
+
+    @staticmethod
+    def backward(ctx, grad_out):  # type: ignore[override]
+        x, weight = ctx.saved_tensors
+        grad_2d = grad_out.reshape(-1, grad_out.shape[-1])
+        grad_x = (grad_2d @ weight.to(grad_2d.dtype)).reshape(x.shape).to(x.dtype)
+        grad_w = (grad_2d.t() @ x.reshape(-1, x.shape[-1]).to(grad_2d.dtype)).to(weight.dtype)
+        grad_b = grad_2d.sum(dim=0) if ctx.has_bias else None
+        return grad_x, grad_w, grad_b, None
+
+
+class NVFP4QADLinearMethod:
+    """Drop-in ``quant_method`` for one H3 block linear (see module docstring)."""
+
+    def __init__(self,
+                 original: Any,
+                 prefix: str,
+                 act_scale: str,
+                 input_amax: float | None = None,
+                 allow_uncalibrated: bool = False) -> None:
+        if act_scale not in ("static", "unit"):
+            raise ValueError(f"act_scale must be 'static' or 'unit', got {act_scale!r}")
+        if act_scale == "static" and not allow_uncalibrated and (input_amax is None or not input_amax > 0):
+            raise ValueError(f"{prefix}: a static activation scale needs a positive calibrated amax")
+        self.original = original
+        self.prefix = prefix
+        self.act_scale = act_scale
+        self.input_amax = float(input_amax) if input_amax is not None else None
+        self.enabled = True
+        # Calibration: run the original (dense) path and fold each input's abs-max into ``observed_amax``.
+        self.calibrating = False
+        self.observed_amax: torch.Tensor | None = None
+        self._sf: dict[torch.device, torch.Tensor] = {}
+
+    def set_input_amax(self, amax: float) -> None:
+        if not amax > 0:
+            raise ValueError(f"{self.prefix}: calibrated amax must be positive, got {amax}")
+        self.input_amax = float(amax)
+        self._sf.clear()
+
+    def input_global_sf(self, device: torch.device) -> torch.Tensor:
+        if self.act_scale == "static" and self.input_amax is None:
+            raise RuntimeError(f"{self.prefix}: static activation scale used before calibration")
+        if device not in self._sf:
+            value = 1.0 if self.act_scale == "unit" else FP4_RANGE / max(self.input_amax or 0.0, 1e-12)
+            self._sf[device] = torch.tensor(value, dtype=torch.float32, device=device)
+        return self._sf[device]
+
+    def apply(self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
+        if self.calibrating:
+            amax = x.detach().abs().amax().float()
+            self.observed_amax = amax if self.observed_amax is None else torch.maximum(self.observed_amax, amax)
+            return self.original.apply(layer, x, bias)
+        if not self.enabled:
+            return self.original.apply(layer, x, bias)
+        input_sf = self.input_global_sf(x.device)
+        if torch.is_grad_enabled() and (layer.weight.requires_grad or x.requires_grad):
+            return _NVFP4STE.apply(x, layer.weight, bias, input_sf)
+        return nvfp4_gemm(x, layer.weight, bias, input_sf)
+
+    def __getattr__(self, name: str) -> Any:
+        # Callers probe optional quant-method hooks (e.g. wants_prequantized_input); fall back to the original.
+        # Before __init__ (copy.deepcopy, unpickling) there is no original: AttributeError keeps hasattr working.
+        try:
+            original = self.__dict__["original"]
+        except KeyError:
+            raise AttributeError(name) from None
+        return getattr(original, name)
+
+
+@dataclass(frozen=True)
+class NVFP4QADPlan:
+    """Which block linears are quantized and how (``h3_dit_vsa``: attention, gates and FFN)."""
+    quantize_attention: bool = True
+    quantize_gate: bool = True
+    quantize_ffn: bool = True
+    skip_blocks: tuple[int, ...] = ()
+
+
+def load_amax_table(path: str | Path) -> dict[str, float]:
+    """``FASTVIDEO_NVFP4_ACT_AMAX``-format JSON (``b<block>.<sub>`` or full prefix -> amax or {"all": amax})."""
+    raw = json.loads(Path(path).read_text())
+    return {key: float(value["all"] if isinstance(value, dict) else value) for key, value in raw.items()}
+
+
+def _amax_for(table: dict[str, float], name: str, block: str, sub: str) -> float | None:
+    for key in (name, f"b{block}.{sub}", f"transformer_blocks.{block}.{sub}"):
+        if key in table:
+            return table[key]
+    return None
+
+
+def install_nvfp4_qad(transformer: torch.nn.Module,
+                      amax_table: dict[str, float] | None,
+                      plan: NVFP4QADPlan,
+                      allow_uncalibrated: bool = False) -> dict[str, NVFP4QADLinearMethod]:
+    """Install QAD methods on the selected block linears; returns them by module name.
+
+    FFN linears take ``static`` scales from ``amax_table`` (required unless
+    ``allow_uncalibrated``, which leaves them for ``calibrate``); attention
+    projections and gates use the unit scale, as the deployed export does.
+    """
+    installed: dict[str, NVFP4QADLinearMethod] = {}
+    for raw_name, module in transformer.named_modules():
+        # Activation checkpointing wraps blocks; key everything by the checkpoint (unwrapped) name.
+        name = raw_name.replace("_checkpoint_wrapped_module.", "")
+        match = H3_BLOCK_LINEAR.search(name)
+        if match is None or not hasattr(module, "quant_method") or int(match.group(1)) in plan.skip_blocks:
+            continue
+        block, sub = match.group(1), match.group(2)
+        if sub.startswith("ff.") and not plan.quantize_ffn:
+            continue
+        if sub == "attn.to_gate_compress" and not plan.quantize_gate:
+            continue
+        if sub.startswith("attn.to_") and sub != "attn.to_gate_compress" and not plan.quantize_attention:
+            continue
+        original = module.quant_method
+        if isinstance(original, NVFP4QADLinearMethod):
+            original = original.original
+        if sub in STATIC_SUFFIXES:
+            amax = _amax_for(amax_table or {}, name, block, sub)
+            if amax is None and not allow_uncalibrated:
+                raise KeyError(f"no calibrated input amax for {name} (key b{block}.{sub}); run Stage A calibration")
+            method = NVFP4QADLinearMethod(original, name, "static", amax, allow_uncalibrated=allow_uncalibrated)
+        else:
+            method = NVFP4QADLinearMethod(original, name, "unit")
+        module.quant_method = method
+        installed[name] = method
+    if not installed:
+        raise ValueError("no MiniMax-H3 block linears matched; is this a MiniMax-H3 transformer?")
+    return installed
+
+
+@contextlib.contextmanager
+def calibrating(methods: dict[str, NVFP4QADLinearMethod]) -> Iterator[None]:
+    """Max-calibrate the static linears that have no amax yet (dense forwards while active).
+
+    On exit the observed maxima are all-reduced (MAX) over the default process
+    group, so every rank, including each sequence-parallel shard, agrees.
+    """
+    pending = {name: m for name, m in methods.items() if m.act_scale == "static" and m.input_amax is None}
+    for method in pending.values():
+        method.calibrating, method.observed_amax = True, None
+    try:
+        yield
+    finally:
+        for method in pending.values():
+            method.calibrating = False
+    if not pending:
+        return
+    names = sorted(pending)
+    device = next((m.observed_amax.device for m in pending.values() if m.observed_amax is not None),
+                  torch.device("cuda", torch.cuda.current_device()) if torch.cuda.is_available() else None)
+    values = torch.stack([(pending[n].observed_amax if pending[n].observed_amax is not None else torch.zeros(
+        (), device=device)).float() for n in names])
+    if torch.distributed.is_available() and torch.distributed.is_initialized():
+        torch.distributed.all_reduce(values, op=torch.distributed.ReduceOp.MAX)
+    for name, value in zip(names, values.tolist(), strict=True):
+        if not value > 0:
+            raise RuntimeError(f"{name} saw no activations during calibration")
+        pending[name].set_input_amax(value)
+
+
+def set_nvfp4_qad_enabled(methods: dict[str, NVFP4QADLinearMethod], enabled: bool) -> None:
+    for method in methods.values():
+        method.enabled = enabled
+
+
+def nvfp4_qad_export_scales(methods: dict[str, NVFP4QADLinearMethod]) -> dict[str, float]:
+    """The static FFN amax values the export must carry (``--act-amax`` input of the packed converter)."""
+    table = {}
+    for name, method in methods.items():
+        match = H3_BLOCK_LINEAR.search(name)
+        if method.act_scale == "static" and match is not None:
+            table[f"b{match.group(1)}.{match.group(2)}"] = float(method.input_amax or 0.0)
+    return table
+
+
+__all__ = [
+    "NVFP4QADLinearMethod", "NVFP4QADPlan", "calibrating", "install_nvfp4_qad", "load_amax_table", "nvfp4_gemm",
+    "nvfp4_qad_export_scales", "set_nvfp4_qad_enabled", "weight_global_sf"
+]

fastvideo/models/dits/minimax_h3.py

@@ -19,6 +19,7 @@ from fastvideo.attention.backends.abstract import AttentionMetadata, layer_idx_f
 from fastvideo.attention.backends.video_sparse_attn_h3 import MiniMaxH3VSAMetadata
 from fastvideo.attention.layer import DistributedAttention_VSA
 from fastvideo.attention.selector import get_attn_backend
+from fastvideo.attention.utils.nvfp4_qk_precondition import parse_block_indices
 from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3Config
 from fastvideo.distributed.communication_op import (
     sequence_model_parallel_all_gather_with_unpad,
@@ -207,6 +208,7 @@ class MiniMaxH3Attention(nn.Module):
         fa4_packed_varlen: bool = False,
         exact_rope: bool = False,
         vsa_fp4: dict[str, Any] | None = None,
+        bf16_attention: bool = False,
     ) -> None:
         super().__init__()
         self.num_attention_heads = num_attention_heads
@@ -254,6 +256,8 @@ class MiniMaxH3Attention(nn.Module):
                                             supported_attention_backends=supported_attention_backends)
         use_vsa = resolved_backend.get_name() == "VIDEO_SPARSE_ATTN_H3"
         attention_cls = DistributedAttention_VSA if use_vsa else DistributedAttention
+        # Blocks listed in FASTVIDEO_NVFP4_FA4_BF16_BLOCKS opt out of the process-wide NVFP4 FA4 Q/K quantization.
+        precision_args: dict[str, Any] = {"nvfp4_fa4": False} if bf16_attention else {}
         self.distributed_attention = attention_cls(
             num_heads=num_attention_heads,
             head_size=attention_head_dim,
@@ -261,6 +265,7 @@ class MiniMaxH3Attention(nn.Module):
             supported_attention_backends=supported_attention_backends,
             prefix=prefix,
             fa4_packed_varlen=fa4_packed_varlen,
+            **precision_args,
         )
         # Opt-in inference route: VSA-H3 selection on the block-sparse FP4
         # kernel (see minimax_h3_vsa_fp4); grad and compile keep the generic path.
@@ -658,6 +663,7 @@ class MiniMaxH3TransformerBlock(nn.Module):
         fa4_packed_varlen: bool = False,
         exact_kernels: frozenset[str] = frozenset(),
         vsa_fp4: dict[str, Any] | None = None,
+        bf16_attention: bool = False,
     ) -> None:
         super().__init__()
         self.norm1 = nn.RMSNorm(hidden_size, eps=norm_eps)
@@ -673,6 +679,7 @@ class MiniMaxH3TransformerBlock(nn.Module):
             fa4_packed_varlen=fa4_packed_varlen,
             exact_rope="rope" in exact_kernels,
             vsa_fp4=vsa_fp4,
+            bf16_attention=bf16_attention,
         )
         self.norm2 = nn.RMSNorm(hidden_size, eps=norm_eps)
         self.ff = MiniMaxH3FeedForward(
@@ -916,6 +923,11 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
             config.quant_config,
             prefix=f"{config.prefix}.token_refiner",
         )
+        bf16_blocks = parse_block_indices(envs.FASTVIDEO_NVFP4_FA4_BF16_BLOCKS.get())
+        unknown_blocks = sorted(index for index in bf16_blocks if index >= arch.num_layers)
+        if unknown_blocks:
+            raise ValueError(f"FASTVIDEO_NVFP4_FA4_BF16_BLOCKS {unknown_blocks} exceed the {arch.num_layers} "
+                             "transformer_blocks")
         self.transformer_blocks = nn.ModuleList([
             MiniMaxH3TransformerBlock(
                 arch.hidden_size,
@@ -935,6 +947,7 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
                 fa4_packed_varlen=envs.FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN.get(),
                 exact_kernels=self.enabled_exact_kernels,
                 vsa_fp4=config.vsa_fp4,
+                bf16_attention=index in bf16_blocks,
             ) for index in range(arch.num_layers)
         ])
         self.norm_out = MiniMaxH3AdaLayerNormOut(

fastvideo/models/dits/minimax_h3_token_compression.py

@@ -0,0 +1,275 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Middle-block spatial token compression for FastH3 (QAD training and inference; hook based).
+
+Blocks ``[start, end)`` run on fewer video tokens. Text and audio rows are untouched; a kept video token keeps its
+original RoPE coordinates and AdaLN row; after the last compressed block every video token receives the update of
+its nearest kept token: ``H_out = H + expand(F(P H) - P H)`` where ``P`` is the compression operator. With
+compression off nothing changes (no hook acts), so the student equals the teacher bit for bit.
+
+Rather than replacing the block list, this installs forward hooks, so module names, FSDP sharding, activation
+checkpointing and QAD linears are untouched:
+
+* block ``start`` pre-hook compresses the hidden states and swaps AdaLN rows / RoPE / sequence length;
+* every block in the range gets the compressed AdaLN rows / RoPE;
+* the inner attention module of each compressed block swaps the VSA metadata for the compressed video grid (a hook
+  on the *inner* module, so activation-checkpoint recompute during backward swaps it too);
+* block ``end - 1`` post-hook expands and restores.
+
+Single sequence-parallel rank only (the compression indexes the whole packed sequence).
+"""
+from __future__ import annotations
+
+import contextlib
+from collections.abc import Callable, Iterator
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from torch import nn
+
+from fastvideo.forward_context import get_forward_context
+
+MODES = ("subsample", "pair_x", "pool2x2")
+
+
+@dataclass(frozen=True)
+class CompressionSpec:
+    start: int
+    end: int
+    mode: str = "subsample"
+    # Compress only the first ``first_steps`` of the sampling forwards (None = all); the low-noise steps that set
+    # fine detail then run on the full token grid.
+    first_steps: int | None = None
+
+    def __post_init__(self) -> None:
+        if self.mode not in MODES:
+            raise ValueError(f"compression mode must be one of {MODES}, got {self.mode!r}")
+        if not 0 <= self.start < self.end:
+            raise ValueError(f"need 0 <= start < end, got [{self.start}, {self.end})")
+
+    @classmethod
+    def parse(cls, text: str) -> CompressionSpec | None:
+        """``"off"``/empty -> None; ``"mode:start:end"`` (e.g. ``subsample:3:41``)."""
+        text = text.strip().lower()
+        if text in ("", "0", "off", "none"):
+            return None
+        parts = text.split(":")
+        if len(parts) not in (3, 4):
+            raise ValueError(f"token compression must be 'off' or 'mode:start:end[:first_steps]', got {text!r}")
+        return cls(int(parts[1]), int(parts[2]), parts[0], int(parts[3]) if len(parts) == 4 else None)
+
+    @classmethod
+    def from_config(cls, cfg: dict[str, Any] | None) -> CompressionSpec | None:
+        if not cfg or not cfg.get("enabled", True):
+            return None
+        first = cfg.get("first_steps")
+        return cls(int(cfg["start"]), int(cfg["end"]), str(cfg.get("mode", "subsample")),
+                   None if first is None else int(first))
+
+
+def compressed_grid(mode: str, grid: tuple[int, int, int]) -> tuple[int, int, int]:
+    frames, height, width = grid
+    if mode == "pair_x":
+        return frames, height, (width + 1) // 2
+    return frames, (height + 1) // 2, (width + 1) // 2
+
+
+def token_reduction(mode: str, grid: tuple[int, int, int]) -> float:
+    sub = compressed_grid(mode, grid)
+    return (grid[0] * grid[1] * grid[2]) / (sub[0] * sub[1] * sub[2])
+
+
+def _pad_to_even(x: torch.Tensor, dims: tuple[int, ...]) -> torch.Tensor:
+    for dim in dims:
+        if x.shape[dim] % 2:
+            x = torch.cat((x, x.narrow(dim, x.shape[dim] - 1, 1)), dim=dim)
+    return x
+
+
+def grid_ops(mode: str, grid: tuple[int, int, int]) -> tuple[Callable, Callable, torch.Tensor]:
+    """``(compress, expand, keep_index)`` for a video token grid ``(T, H, W)``.
+
+    ``compress``: rows ``[B, T*H*W, C]`` -> ``[B, T*h*w, C]``; ``expand``: delta ``[B, T*h*w, C]`` -> ``[B, T*H*W, C]``;
+    ``keep_index``: for every compressed token the video row whose RoPE / AdaLN row it keeps (CPU long).
+    """
+    frames, height, width = grid
+    sub = compressed_grid(mode, grid)
+    index = torch.arange(frames * height * width).reshape(frames, height, width)
+
+    def as_grid(rows: torch.Tensor) -> torch.Tensor:
+        return rows.reshape(rows.shape[0], frames, height, width, rows.shape[-1])
+
+    if mode == "subsample":
+        keep = index[:, ::2, ::2].reshape(-1)
+
+        def compress(rows: torch.Tensor) -> torch.Tensor:
+            return as_grid(rows)[:, :, ::2, ::2].reshape(rows.shape[0], -1, rows.shape[-1])
+
+    elif mode == "pair_x":
+        keep = index[:, :, ::2].reshape(-1)
+
+        def compress(rows: torch.Tensor) -> torch.Tensor:
+            g = _pad_to_even(as_grid(rows), (3, ))
+            pair = g.float().unflatten(3, (-1, 2)).mean(dim=4)
+            return pair.to(rows.dtype).reshape(rows.shape[0], -1, rows.shape[-1])
+
+    else:  # pool2x2
+        keep = index[:, ::2, ::2].reshape(-1)
+
+        def compress(rows: torch.Tensor) -> torch.Tensor:
+            g = _pad_to_even(as_grid(rows), (2, 3))
+            pooled = g.float().unflatten(3, (-1, 2)).unflatten(2, (-1, 2)).mean(dim=(3, 5))
+            return pooled.to(rows.dtype).reshape(rows.shape[0], -1, rows.shape[-1])
+
+    repeat_h = 1 if mode == "pair_x" else 2
+
+    def expand(delta: torch.Tensor) -> torch.Tensor:
+        d = delta.reshape(delta.shape[0], *sub, delta.shape[-1])
+        if repeat_h == 2:
+            d = d.repeat_interleave(2, dim=2)[:, :, :height]
+        d = d.repeat_interleave(2, dim=3)[:, :, :, :width]
+        return d.reshape(delta.shape[0], -1, delta.shape[-1])
+
+    return compress, expand, keep
+
+
+def _unwrap(module: nn.Module) -> nn.Module:
+    return getattr(module, "_checkpoint_wrapped_module", module)
+
+
+@dataclass
+class _Forward:
+    """Per-forward state (one transformer call)."""
+    full: torch.Tensor | None = None
+    compressed_in: torch.Tensor | None = None
+    adaln: torch.Tensor | None = None
+    rotary: tuple[torch.Tensor, torch.Tensor] | None = None
+    seq_len: int = 0
+    num_prefix: int = 0
+
+
+class TokenCompression:
+    """Installs the hooks; ``scope(grid, compressed_metadata)`` arms them for one forward (and its backward)."""
+
+    def __init__(self, transformer: nn.Module, spec: CompressionSpec) -> None:
+        self.spec = spec
+        self.blocks = list(transformer.transformer_blocks)
+        if spec.end > len(self.blocks):
+            raise ValueError(f"compression range [{spec.start}, {spec.end}) exceeds {len(self.blocks)} blocks")
+        self._handles: list[Any] = []
+        self._grid: tuple[int, int, int] | None = None
+        self._ops: tuple[Callable, Callable, torch.Tensor] | None = None
+        self._metadata: Any = None
+        self._saved_metadata: Any = None
+        self._state = _Forward()
+        for index in range(spec.start, spec.end):
+            block = self.blocks[index]
+            self._handles.append(block.register_forward_pre_hook(self._pre(index)))
+            if index == spec.end - 1:
+                self._handles.append(block.register_forward_hook(self._post))
+            attn = _unwrap(block).attn
+            self._handles.append(attn.register_forward_pre_hook(self._attn_pre))
+            self._handles.append(attn.register_forward_hook(self._attn_post))
+
+    @property
+    def active(self) -> bool:
+        return self._grid is not None
+
+    @contextlib.contextmanager
+    def scope(self, grid: tuple[int, int, int], compressed_metadata: Any) -> Iterator[None]:
+        """Arm the hooks for forwards (and their backward recompute) inside the block."""
+        self._grid, self._metadata = grid, compressed_metadata
+        self._ops = grid_ops(self.spec.mode, grid)
+        try:
+            yield
+        finally:
+            self._grid = self._ops = self._metadata = None
+            self._state = _Forward()
+
+    def remove(self) -> None:
+        for handle in self._handles:
+            handle.remove()
+        self._handles.clear()
+
+    # -- hooks
+    def _pre(self, index: int):
+
+        def hook(module: nn.Module, args: tuple[Any, ...]) -> tuple[Any, ...] | None:
+            if not self.active:
+                return None
+            hidden, temb, adaln, rotary, _ = args
+            if index == self.spec.start:
+                self._enter(hidden, adaln, rotary)
+                hidden = self._state.compressed_in
+            state = self._state
+            return hidden, temb, state.adaln, state.rotary, state.seq_len
+
+        return hook
+
+    def _enter(self, hidden: torch.Tensor, adaln: torch.Tensor, rotary: tuple[torch.Tensor, torch.Tensor]) -> None:
+        assert self._grid is not None and self._ops is not None
+        compress, _, keep = self._ops
+        num_video = self._grid[0] * self._grid[1] * self._grid[2]
+        num_prefix = hidden.shape[1] - num_video
+        if num_prefix < 0 or adaln.shape[0] != hidden.shape[1]:
+            raise ValueError("token compression needs the whole packed sequence on one rank "
+                             f"(hidden {tuple(hidden.shape)}, adaln {tuple(adaln.shape)}, video {num_video})")
+        rows = torch.cat((torch.arange(num_prefix), keep + num_prefix)).to(hidden.device)
+        self._state = _Forward(full=hidden,
+                               compressed_in=torch.cat((hidden[:, :num_prefix], compress(hidden[:, num_prefix:])), dim=1),
+                               adaln=adaln.index_select(0, rows),
+                               rotary=(rotary[0].index_select(0, rows), rotary[1].index_select(0, rows)),
+                               seq_len=int(rows.numel()),
+                               num_prefix=num_prefix)
+
+    def _post(self, module: nn.Module, args: tuple[Any, ...], output: torch.Tensor) -> torch.Tensor | None:
+        if not self.active:
+            return None
+        assert self._ops is not None
+        state = self._state
+        n = state.num_prefix
+        delta = self._ops[1](output[:, n:] - state.compressed_in[:, n:])
+        restored = torch.cat((output[:, :n], state.full[:, n:] + delta.to(state.full.dtype)), dim=1)
+        self._state = _Forward()
+        return restored
+
+    def _attn_pre(self, module: nn.Module, args: tuple[Any, ...], kwargs: Any = None) -> None:
+        if not self.active:
+            return
+        context = get_forward_context()
+        self._saved_metadata = context.attn_metadata
+        context.attn_metadata = self._metadata
+
+    def _attn_post(self, module: nn.Module, args: tuple[Any, ...], output: Any) -> None:
+        if not self.active:
+            return
+        get_forward_context().attn_metadata = self._saved_metadata
+
+
+def block_flops_share(spec: CompressionSpec | None, grid: tuple[int, int, int], num_blocks: int, num_other_tokens: int,
+                      attn_keep: float = 0.2) -> dict[str, float]:
+    """Relative DiT cost of the compressed model (1.0 = uncompressed), by FLOPs.
+
+    Per block, per token: linear layers ~ ``12 d^2`` MACs (qkv+out+gate ~ 5 d^2, FFN SwiGLU 3 * d * 2.67 d), attention
+    ~ ``2 * attended_keys * d``; VSA keeps ``attn_keep`` of the video keys. ``d`` cancels in the ratio when expressed in
+    ``d`` units, so only the token counts matter; ``d = 5376`` is used for the attention/linear balance.
+    """
+    d = 5376.0
+    video = grid[0] * grid[1] * grid[2]
+
+    def cost(tokens_video: float) -> float:
+        n = tokens_video + num_other_tokens
+        linear = 13.4 * d * d * n
+        attn = 2.0 * 2.0 * d * n * (attn_keep * tokens_video + num_other_tokens)
+        return linear + attn
+
+    full = cost(video)
+    if spec is None:
+        return {"relative_flops": 1.0, "speedup_flops": 1.0}
+    reduced = cost(video / token_reduction(spec.mode, grid))
+    inside = spec.end - spec.start
+    total = (num_blocks - inside) * full + inside * reduced
+    return {"relative_flops": total / (num_blocks * full), "speedup_flops": num_blocks * full / total}
+
+
+__all__ = ["CompressionSpec", "TokenCompression", "block_flops_share", "compressed_grid", "grid_ops", "token_reduction"]

fastvideo/models/schedulers/scheduling_minimax_h3.py

@@ -21,13 +21,31 @@ class MiniMaxH3Scheduler(SchedulerMixin, ConfigMixin):
     order = 1
 
     @register_to_config
-    def __init__(self, shift: float = 12.0) -> None:
+    def __init__(
+        self,
+        shift: float = 12.0,
+        stochastic_sampling: bool = False,
+        shift_provided_sigmas: bool = False,
+    ) -> None:
+        """H3's rectified-flow scheduler, with Pika's re-noise step as an option.
+
+        ``stochastic_sampling`` switches the Euler update to the re-noise
+        contract: draw fresh Gaussian noise at the next sigma and form
+        ``(1 - sigma_next) * x0 + sigma_next * noise``.  The terminal sigma
+        always returns ``x0`` without drawing noise.  ``shift_provided_sigmas``
+        records that explicit sigmas are already shifted (Pika's release uses
+        it); FastVideo always treats explicit ``set_timesteps(sigmas=...)``
+        values that way, so it is informational.
+        """
         if shift <= 0:
             raise ValueError(f"`shift` must be positive, got {shift}.")
         self.num_inference_steps: int | None = None
         self.sigmas: torch.Tensor | None = None
         self.timesteps: torch.Tensor | None = None
         self._shift = float(shift)
+        self._stochastic_sampling = bool(stochastic_sampling)
+        self._shift_provided_sigmas = bool(shift_provided_sigmas)
+        self._noise_generator: torch.Generator | None = None
         self._step_index: int | None = None
         self._begin_index: int | None = None
 
@@ -35,6 +53,18 @@ class MiniMaxH3Scheduler(SchedulerMixin, ConfigMixin):
     def shift(self) -> float:
         return self._shift
 
+    @property
+    def stochastic_sampling(self) -> bool:
+        return self._stochastic_sampling
+
+    @property
+    def shift_provided_sigmas(self) -> bool:
+        return self._shift_provided_sigmas
+
+    def set_noise_generator(self, generator: torch.Generator | None) -> None:
+        """Set the generator used by stochastic (re-noise) steps."""
+        self._noise_generator = generator
+
     @property
     def step_index(self) -> int | None:
         return self._step_index
@@ -127,11 +157,29 @@ class MiniMaxH3Scheduler(SchedulerMixin, ConfigMixin):
         denoised = sample + sigma_from_timestep * model_output
 
         compute_dtype = torch.float32 if sample.dtype in (torch.float16, torch.bfloat16) else sample.dtype
-        sigma = self.sigmas[self._step_index].to(device=sample.device, dtype=compute_dtype)
         sigma_next = self.sigmas[self._step_index + 1].to(device=sample.device, dtype=compute_dtype)
-        ratio = sigma_next / sigma
-        prev_sample = ratio * sample.to(compute_dtype) + (1.0 - ratio) * denoised.to(compute_dtype)
-        prev_sample = prev_sample.to(sample.dtype)
+        if self._stochastic_sampling:
+            if self._noise_generator is None:
+                raise ValueError("Stochastic MiniMax-H3 sampling needs a noise generator; call "
+                                 "`set_noise_generator` before `step`.")
+            if float(sigma_next) == 0.0:
+                # The terminal ladder point is the generated sample. Pika's
+                # sampler returns x0 here instead of drawing a final noise.
+                prev_sample = denoised.to(sample.dtype)
+            else:
+                noise = torch.randn(
+                    sample.shape,
+                    generator=self._noise_generator,
+                    device=sample.device,
+                    dtype=torch.float32,
+                )
+                prev_sample = ((1.0 - sigma_next) * denoised.to(compute_dtype) +
+                               sigma_next * noise).to(sample.dtype)
+        else:
+            sigma = self.sigmas[self._step_index].to(device=sample.device, dtype=compute_dtype)
+            ratio = sigma_next / sigma
+            prev_sample = ratio * sample.to(compute_dtype) + (1.0 - ratio) * denoised.to(compute_dtype)
+            prev_sample = prev_sample.to(sample.dtype)
         self._step_index += 1
 
         if not return_dict:

fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py

@@ -45,6 +45,7 @@ from fastvideo.pipelines.basic.minimax_h3.encoder_split import (
 )
 from fastvideo.pipelines.basic.minimax_h3.vsa_guard import refuse_zero_initialized_h3_vsa
 from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.lazy_module import is_lazy_module
 from fastvideo.pipelines.lora_pipeline import LoRAPipeline
 from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.platforms import AttentionBackendEnum
@@ -218,6 +219,10 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase):
 
     pipeline_config_cls: type[MiniMaxH3PipelineConfig] = MiniMaxH3PipelineConfig
     _ref2va_default = False
+    # Family subclasses override these to swap geometry and sampler stages
+    # without copying the whole stage assembly.
+    input_preparation_stage_cls: type[MiniMaxH3InputPreparationStage] = MiniMaxH3InputPreparationStage
+    denoising_stage_cls: type[MiniMaxH3DenoisingStage] = MiniMaxH3DenoisingStage
     _required_config_modules = [
         "text_encoder",
         "tokenizer",
@@ -264,8 +269,17 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase):
 
         ``_pinned_swap`` reuses a parameter's stored host copy across parks; a LoRA
         swap/merge/unmerge changes the weights, so the stale copy must not win.
+
+        A deferred component that is not resident holds no pinned copies, and
+        touching ``__dict__`` through its ``LazyModule`` would load it: the
+        constructor adapter attaches while the transformer materializes, after
+        conditioning has already released the text encoder.
         """
         for module in self.modules.values():
+            if is_lazy_module(module):
+                if not module.is_materialized:
+                    continue
+                module = module.materialize()
             getattr(module, "__dict__", {}).pop("_pinned_host_tensors", None)
 
     def set_lora_adapter(self, *args, **kwargs):  # type: ignore[override]
@@ -686,7 +700,7 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase):
     def _add_input_stage(self, fastvideo_args: FastVideoArgs, *, ref2va: bool) -> None:
         self.add_stage(
             "input_preparation_stage",
-            MiniMaxH3InputPreparationStage(
+            self.input_preparation_stage_cls(
                 vae=self._input_video_geometry(fastvideo_args),
                 audio_vae=self._input_audio_vae(fastvideo_args, ref2va=ref2va),
                 ref2va=ref2va,
@@ -734,7 +748,7 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase):
         )
         self.add_stage(
             "denoising_stage",
-            MiniMaxH3DenoisingStage(
+            self.denoising_stage_cls(
                 transformer=transformer,
                 scheduler=scheduler,
                 audio_scheduler=audio_scheduler,

fastvideo/pipelines/basic/minimax_h3/packing.py

@@ -236,20 +236,33 @@ def build_packed_batch_sequence(documents: tuple[MiniMaxH3PackedLayout, ...], )
     )
 
 
-def resolve_canvas_size(aspect_width: float, aspect_height: float) -> tuple[int, int]:
+def resolve_canvas_size(
+    aspect_width: float,
+    aspect_height: float,
+    *,
+    short_edge: int = MINIMAX_H3_SHORT_EDGE,
+    max_pixels: int = MINIMAX_H3_MAX_PIXELS,
+) -> tuple[int, int]:
+    """Resolve a canvas from an aspect ratio and a tier's short-edge/area rules.
+
+    The defaults are the released MiniMax-H3 768p rule.  Pika H3 Flex passes its
+    own 480p/768p tier values through the keyword arguments.
+    """
     if aspect_width <= 0 or aspect_height <= 0:
         raise ValueError(f"The aspect ratio must be positive, got {aspect_width}:{aspect_height}.")
+    if short_edge <= 0 or max_pixels <= 0:
+        raise ValueError(f"Canvas short_edge and max_pixels must be positive, got {short_edge}, {max_pixels}.")
     ratio = aspect_width / aspect_height
     if not MINIMAX_H3_MIN_ASPECT_RATIO <= ratio <= MINIMAX_H3_MAX_ASPECT_RATIO:
         raise ValueError(f"MiniMax-H3 supports aspect ratios from 1:4 to 4:1, got {aspect_width}:{aspect_height}.")
 
     if ratio >= 1:
-        width, height = MINIMAX_H3_SHORT_EDGE * ratio, float(MINIMAX_H3_SHORT_EDGE)
+        width, height = short_edge * ratio, float(short_edge)
     else:
-        width, height = float(MINIMAX_H3_SHORT_EDGE), MINIMAX_H3_SHORT_EDGE / ratio
+        width, height = float(short_edge), short_edge / ratio
     area = width * height
-    if area > MINIMAX_H3_MAX_PIXELS:
-        scale = (MINIMAX_H3_MAX_PIXELS / area)**0.5
+    if area > max_pixels:
+        scale = (max_pixels / area)**0.5
         width, height = width * scale, height * scale
     multiple = MINIMAX_H3_CANVAS_MULTIPLE
     return max(multiple, round(height / multiple) * multiple), max(multiple, round(width / multiple) * multiple)

fastvideo/pipelines/basic/minimax_h3/reference.py

@@ -22,6 +22,8 @@ from fastvideo.pipelines.basic.minimax_h3.packing import (
     MINIMAX_H3_FRAMES_PER_CHUNK,
     MINIMAX_H3_LATENTS_PER_CHUNK,
     MINIMAX_H3_MAX_DURATION,
+    MINIMAX_H3_MAX_PIXELS,
+    MINIMAX_H3_SHORT_EDGE,
     resolve_canvas_size,
 )
 
@@ -180,13 +182,20 @@ def decode_reference_video(source: str | os.PathLike[str]) -> tuple[np.ndarray,
     return result, float(rate), soundtrack
 
 
-def resolve_reference_image_size(width: int, height: int) -> tuple[int, int]:
-    """Resolve the released 2048-short-edge reference-image canvas."""
+def resolve_reference_image_size(
+    width: int,
+    height: int,
+    *,
+    short_edge: int = MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE,
+) -> tuple[int, int]:
+    """Resolve the released reference-image canvas at ``short_edge``."""
     if width <= 0 or height <= 0:
         raise ValueError(f"A reference image must have a positive size, got {width}x{height}.")
     if width > 4 * height or height > 4 * width:
         raise ValueError(f"A reference image must be within 1:4 and 4:1, got {width}x{height}.")
-    scale = MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE / min(width, height)
+    if short_edge <= 0:
+        raise ValueError(f"A reference-image short edge must be positive, got {short_edge}.")
+    scale = short_edge / min(width, height)
     multiple = MINIMAX_H3_CANVAS_MULTIPLE
     return (
         max(multiple,
@@ -231,12 +240,23 @@ def resample_reference_frames(frames: np.ndarray, fps: float) -> np.ndarray:
     return np.repeat(frames, repeats, axis=0)
 
 
-def prepare_reference_frames(frames: np.ndarray, num_frames: int) -> np.ndarray:
+def prepare_reference_frames(
+    frames: np.ndarray,
+    num_frames: int,
+    *,
+    short_edge: int = MINIMAX_H3_SHORT_EDGE,
+    max_pixels: int = MINIMAX_H3_MAX_PIXELS,
+) -> np.ndarray:
     """Trim a reference video to the request and resize every RGB frame."""
     if frames.ndim != 4 or frames.shape[-1] != 3 or frames.shape[0] == 0:
         raise ValueError(f"A reference video must be non-empty RGB frames, got {tuple(frames.shape)}.")
     frames = frames[:num_frames]
-    height, width = resolve_canvas_size(frames.shape[2], frames.shape[1])
+    height, width = resolve_canvas_size(
+        frames.shape[2],
+        frames.shape[1],
+        short_edge=short_edge,
+        max_pixels=max_pixels,
+    )
     if frames.shape[1:3] == (height, width):
         return frames
     return np.stack(
@@ -319,8 +339,17 @@ def prepare_reference(
     reference: MiniMaxH3Reference,
     num_frames: int,
     target_sample_rate: int,
+    *,
+    reference_short_edge: int = MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE,
+    reference_canvas_short_edge: int = MINIMAX_H3_SHORT_EDGE,
+    reference_canvas_max_pixels: int = MINIMAX_H3_MAX_PIXELS,
 ) -> MiniMaxH3PreparedReference:
-    """Decode and normalize one deferred Ref2VA medium."""
+    """Decode and normalize one deferred Ref2VA medium.
+
+    The keyword arguments let a family use its own tier's reference-image short
+    edge and reference-video canvas rule.  MiniMax-H3's released values are the
+    defaults; Pika H3 Flex passes its 480p/768p tier values.
+    """
     prepared = MiniMaxH3PreparedReference(media_type=reference.media_type)
     if reference.media_type == "image":
         source = reference.source
@@ -331,7 +360,7 @@ def prepare_reference(
                 raise ValueError(f"An image reference must be RGB, got {tuple(pixels.shape)}.")
             image = Image.fromarray(pixels)
         image = ImageOps.exif_transpose(image).convert("RGB")
-        height, width = resolve_reference_image_size(*image.size)
+        height, width = resolve_reference_image_size(*image.size, short_edge=reference_short_edge)
         prepared.image = prepare_reference_image(image, height, width)
         return prepared
 
@@ -346,7 +375,12 @@ def prepare_reference(
             frames = reference_media_to_uint8(reference.source)
             decoded_fps = float(MINIMAX_H3_FPS)
         fps = float(reference.fps if reference.fps is not None else decoded_fps)
-        prepared.frames = prepare_reference_frames(resample_reference_frames(frames, fps), num_frames)
+        prepared.frames = prepare_reference_frames(
+            resample_reference_frames(frames, fps),
+            num_frames,
+            short_edge=reference_canvas_short_edge,
+            max_pixels=reference_canvas_max_pixels,
+        )
 
         if reference.soundtrack is not None:
             waveform, source_rate = _resolve_audio_source(reference.soundtrack)

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

@@ -133,14 +133,21 @@ class MiniMaxH3DenoisingStage(PipelineStage):
         self.scheduler = scheduler
         self.audio_scheduler = audio_scheduler
 
-    def _set_dmd_schedule(self, steps: list[int], grid_points: int, device: torch.device) -> None:
-        """Run the trained rungs, shifting the shared noise clock once per modality."""
+    def _set_shifted_schedule(
+        self,
+        steps: list[int],
+        grid_points: int,
+        device: torch.device,
+        *,
+        label: str = "MiniMax-H3 DMD",
+    ) -> None:
+        """Build one modality-shifted sigma grid from explicit clean-time rungs."""
         if (not steps or any(type(step) is not int or not 0 < step <= 1000 for step in steps)
                 or any(left <= right for left, right in zip(steps, steps[1:], strict=False))):
-            raise ValueError("MiniMax-H3 DMD rungs must be strictly decreasing integers in (0, 1000].")
+            raise ValueError(f"{label} rungs must be strictly decreasing integers in (0, 1000].")
         if grid_points != len(steps) + 1:
-            raise ValueError("MiniMax-H3 num_inference_steps counts sigma-grid points: "
-                             f"{len(steps)} DMD forwards require {len(steps) + 1} grid points, got {grid_points}.")
+            raise ValueError(f"{label} num_inference_steps counts sigma-grid points: "
+                             f"{len(steps)} forwards require {len(steps) + 1} grid points, got {grid_points}.")
         base = torch.tensor([step / 1000.0 for step in steps] + [0.0], dtype=torch.float32)
         for scheduler in (self.scheduler, self.audio_scheduler):
             shift = float(scheduler.shift)
@@ -149,6 +156,38 @@ class MiniMaxH3DenoisingStage(PipelineStage):
             # clean time (1 - sigma); passing integer rungs to step() is wrong.
             scheduler.set_timesteps(sigmas=sigmas, device=device)
 
+    def _set_dmd_schedule(self, steps: list[int], grid_points: int, device: torch.device) -> None:
+        """Run the trained rungs, shifting the shared noise clock once per modality."""
+        self._set_shifted_schedule(steps, grid_points, device, label="MiniMax-H3 DMD")
+
+    def _configure_schedules(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+        device: torch.device,
+    ) -> PDDSamplingPlan | None:
+        """Install this checkpoint's sigma schedules and return a PDD plan, if any.
+
+        Families with a fixed sampler (Pika H3 Flex) override this hook; the
+        default keeps MiniMax-H3's PDD and DMD checkpoint contracts.
+        """
+        dmd_steps = fastvideo_args.pipeline_config.dmd_denoising_steps
+        pdd_plan = self._pdd_sampling_plan(fastvideo_args, device)
+        if pdd_plan is not None:
+            # One runtime step fuses one block of heads into their weighted
+            # mean velocity; the ordinary Euler step over the block's two node
+            # sigmas then applies the block's total increment per modality.
+            self.scheduler.set_timesteps(sigmas=pdd_plan.node_sigmas["video"].to(torch.float32), device=device)
+            self.audio_scheduler.set_timesteps(sigmas=pdd_plan.node_sigmas["audio"].to(torch.float32), device=device)
+            logger.info("MiniMax-H3 PDD denoising: %d fused block(s) %s over a %d-interval fine grid.",
+                        pdd_plan.num_steps, pdd_plan.block_sizes, pdd_plan.pdd_steps)
+        elif dmd_steps is None:
+            self.scheduler.set_timesteps(batch.num_inference_steps, device=device)
+            self.audio_scheduler.set_timesteps(batch.num_inference_steps, device=device)
+        else:
+            self._set_dmd_schedule(dmd_steps, batch.num_inference_steps, device)
+        return pdd_plan
+
     def _pdd_sampling_plan(self, fastvideo_args: FastVideoArgs, device: torch.device) -> PDDSamplingPlan | None:
         """The fused-block plan of a PDD checkpoint, or None for any other checkpoint.
 
@@ -207,21 +246,7 @@ class MiniMaxH3DenoisingStage(PipelineStage):
                             and not fastvideo_args.use_fsdp_inference)
         device = get_local_torch_device()
 
-        dmd_steps = fastvideo_args.pipeline_config.dmd_denoising_steps
-        pdd_plan = self._pdd_sampling_plan(fastvideo_args, device)
-        if pdd_plan is not None:
-            # One runtime step fuses one block of heads into their weighted
-            # mean velocity; the ordinary Euler step over the block's two node
-            # sigmas then applies the block's total increment per modality.
-            self.scheduler.set_timesteps(sigmas=pdd_plan.node_sigmas["video"].to(torch.float32), device=device)
-            self.audio_scheduler.set_timesteps(sigmas=pdd_plan.node_sigmas["audio"].to(torch.float32), device=device)
-            logger.info("MiniMax-H3 PDD denoising: %d fused block(s) %s over a %d-interval fine grid.",
-                        pdd_plan.num_steps, pdd_plan.block_sizes, pdd_plan.pdd_steps)
-        elif dmd_steps is None:
-            self.scheduler.set_timesteps(batch.num_inference_steps, device=device)
-            self.audio_scheduler.set_timesteps(batch.num_inference_steps, device=device)
-        else:
-            self._set_dmd_schedule(dmd_steps, batch.num_inference_steps, device)
+        pdd_plan = self._configure_schedules(batch, fastvideo_args, device)
         video_timesteps = self.scheduler.timesteps
         audio_timesteps = self.audio_scheduler.timesteps
         if video_timesteps is None or audio_timesteps is None:

fastvideo/pipelines/basic/pika_h3_flex/__init__.py

@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from fastvideo.pipelines.basic.pika_h3_flex.pipeline import (
+    EntryClass,
+    PikaH3FlexModularPipeline,
+    PikaH3FlexRef2VAModularPipeline,
+)
+from fastvideo.pipelines.basic.pika_h3_flex.adapter import PikaAdapter
+
+__all__ = [
+    "EntryClass",
+    "PikaAdapter",
+    "PikaH3FlexModularPipeline",
+    "PikaH3FlexRef2VAModularPipeline",
+]

fastvideo/pipelines/basic/pika_h3_flex/adapter.py

@@ -0,0 +1,293 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Locate and validate Pika H3 Flex's released LoRA adapters.
+
+Pika H3 Flex is MiniMax-H3 plus one of two rank-128 PEFT adapters: ``fl2va``
+for text/first-frame/last-frame generation and ``ref2va`` for ordered
+image/video/audio references.  The adapter is part of the model contract, not a
+user-tunable accelerator, so the pipeline loads it automatically and refuses a
+manifest whose rank, trunk, or sampler does not match the release.
+
+This module stays torch-free; the manifest is a plain JSON document and the
+only optional dependency is ``huggingface_hub`` for a repo-id source.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+from collections.abc import Mapping
+
+from fastvideo.configs.pipelines.pika_h3_flex import (
+    PIKA_AUDIO_SCHEDULER_SHIFT,
+    PIKA_DENOISING_STEPS,
+    PIKA_VIDEO_SCHEDULER_SHIFT,
+)
+
+PIKA_ADAPTER_FORMAT_VERSION = 1
+PIKA_ADAPTER_RANK = 128
+PIKA_ADAPTER_ALPHA = 128
+PIKA_BASE_REPO = "MiniMaxAI/MiniMax-H3"
+PIKA_ADAPTER_FILENAME = "adapter.safetensors"
+PIKA_MANIFEST_FILENAME = "manifest.json"
+
+PIKA_SAMPLING_CONTRACT = {
+    "knots": (*PIKA_DENOISING_STEPS, 0),
+    "video_shift": PIKA_VIDEO_SCHEDULER_SHIFT,
+    "audio_shift": PIKA_AUDIO_SCHEDULER_SHIFT,
+    "method": "renoise",
+}
+PIKA_TARGET_MODULES = (
+    "to_q",
+    "to_k",
+    "to_v",
+    "to_out.0",
+    "ff.net.0.proj",
+    "ff.net.2",
+)
+
+TASK_TO_TRUNK = {"fl2va": "transformer", "ref2va": "transformer_ref"}
+
+
+class PikaAdapterError(ValueError):
+    """The Pika adapter is missing or does not describe the released model."""
+
+
+@dataclass(frozen=True)
+class PikaAdapter:
+    """One validated released adapter."""
+
+    task: str
+    trunk: str
+    path: str
+    manifest_path: str
+    manifest: dict[str, Any]
+
+    @property
+    def directory(self) -> str:
+        return str(Path(self.path).parent)
+
+
+def task_for_pipeline(ref2va: bool) -> str:
+    return "ref2va" if ref2va else "fl2va"
+
+
+def expected_trunk(task: str) -> str:
+    try:
+        return TASK_TO_TRUNK[task]
+    except KeyError as error:
+        raise PikaAdapterError(
+            f"Unknown Pika H3 Flex task {task!r}; expected one of {sorted(TASK_TO_TRUNK)}.") from error
+
+
+def _as_int(value: Any, field: str) -> int:
+    if isinstance(value, bool) or not isinstance(value, int):
+        raise PikaAdapterError(f"Pika adapter manifest field {field!r} must be an integer, got {value!r}.")
+    return value
+
+
+def _as_float(value: Any, field: str) -> float:
+    if isinstance(value, bool) or not isinstance(value, int | float):
+        raise PikaAdapterError(f"Pika adapter manifest field {field!r} must be a number, got {value!r}.")
+    return float(value)
+
+
+def validate_adapter_manifest(
+    manifest: Mapping[str, Any],
+    *,
+    task: str,
+    pipeline_trunk: str | None = None,
+) -> None:
+    """Reject a manifest that is not the released Pika adapter for ``task``."""
+    trunk = expected_trunk(task)
+    if pipeline_trunk is not None and pipeline_trunk != trunk:
+        raise PikaAdapterError(f"Pika H3 Flex task {task!r} loads trunk {trunk!r}, but the pipeline requested "
+                               f"{pipeline_trunk!r}.")
+
+    format_version = _as_int(manifest.get("format_version"), "format_version")
+    if format_version != PIKA_ADAPTER_FORMAT_VERSION:
+        raise PikaAdapterError(f"Unsupported Pika adapter manifest format {format_version}; expected "
+                               f"{PIKA_ADAPTER_FORMAT_VERSION}.")
+    manifest_task = manifest.get("task")
+    if manifest_task is not None and manifest_task != task:
+        raise PikaAdapterError(f"The adapter at this path is for task {manifest_task!r}, not {task!r}.")
+    rank = _as_int(manifest.get("rank"), "rank")
+    alpha = _as_int(manifest.get("alpha"), "alpha")
+    if rank != PIKA_ADAPTER_RANK or alpha != PIKA_ADAPTER_ALPHA:
+        raise PikaAdapterError(f"Pika H3 Flex adapters are rank {PIKA_ADAPTER_RANK} with alpha "
+                               f"{PIKA_ADAPTER_ALPHA}; this manifest says rank {rank}, alpha {alpha}.")
+    manifest_trunk = manifest.get("trunk")
+    if manifest_trunk != trunk:
+        raise PikaAdapterError(f"Pika H3 Flex task {task!r} expects adapter trunk {trunk!r}, got "
+                               f"{manifest_trunk!r}.")
+
+    sampling = manifest.get("sampling")
+    if not isinstance(sampling, Mapping):
+        raise PikaAdapterError("Pika adapter manifest is missing its `sampling` contract.")
+    knots = sampling.get("knots")
+    if not isinstance(knots, list | tuple) or tuple(knots) != PIKA_SAMPLING_CONTRACT["knots"]:
+        raise PikaAdapterError(f"Pika adapter sampling knots must be "
+                               f"{list(PIKA_SAMPLING_CONTRACT['knots'])}, got {knots!r}.")
+    for field in ("video_shift", "audio_shift"):
+        value = _as_float(sampling.get(field), f"sampling.{field}")
+        expected = float(PIKA_SAMPLING_CONTRACT[field])
+        if abs(value - expected) > 1e-9:
+            raise PikaAdapterError(f"Pika adapter sampling.{field} must be {expected:g}, got {value:g}.")
+    if sampling.get("method") != PIKA_SAMPLING_CONTRACT["method"]:
+        raise PikaAdapterError(f"Pika adapter sampling.method must be "
+                               f"{PIKA_SAMPLING_CONTRACT['method']!r}, got {sampling.get('method')!r}.")
+
+    target_modules = manifest.get("target_modules")
+    if not isinstance(target_modules, list | tuple) or tuple(target_modules) != PIKA_TARGET_MODULES:
+        raise PikaAdapterError(f"Pika adapter target_modules must be {list(PIKA_TARGET_MODULES)}, got "
+                               f"{target_modules!r}.")
+
+    base = manifest.get("base")
+    if not isinstance(base, Mapping) or base.get("repo_id") != PIKA_BASE_REPO:
+        raise PikaAdapterError(f"Pika adapter base must name {PIKA_BASE_REPO!r}, got {base!r}.")
+
+
+def _load_manifest(path: Path) -> dict[str, Any]:
+    try:
+        manifest = json.loads(path.read_text(encoding="utf-8"))
+    except OSError as error:
+        raise PikaAdapterError(f"Cannot read Pika adapter manifest {str(path)!r}: {error}") from error
+    except json.JSONDecodeError as error:
+        raise PikaAdapterError(f"Pika adapter manifest {str(path)!r} is not valid JSON: {error}") from error
+    if not isinstance(manifest, dict):
+        raise PikaAdapterError(f"Pika adapter manifest {str(path)!r} must be a JSON object.")
+    return manifest
+
+
+def _adapter_in_directory(directory: Path, task: str) -> tuple[Path, Path] | None:
+    nested = directory / task / PIKA_ADAPTER_FILENAME
+    if nested.is_file():
+        manifest = directory / task / PIKA_MANIFEST_FILENAME
+        return nested, manifest
+    flat = directory / PIKA_ADAPTER_FILENAME
+    if flat.is_file():
+        manifest = directory / PIKA_MANIFEST_FILENAME
+        return flat, manifest
+    return None
+
+
+def _looks_like_repo_id(source: str) -> bool:
+    parts = source.split("/")
+    return len(parts) == 2 and all(parts) and "\\" not in source
+
+
+def _download_adapter(source: str, task: str, revision: str | None) -> tuple[Path, Path]:
+    try:
+        from huggingface_hub import hf_hub_download
+    except ImportError as error:  # pragma: no cover - huggingface_hub is a runtime dependency
+        raise PikaAdapterError("Downloading a Pika adapter by repo id requires huggingface_hub.") from error
+    try:
+        adapter = Path(hf_hub_download(
+            repo_id=source,
+            filename=f"{task}/{PIKA_ADAPTER_FILENAME}",
+            revision=revision,
+        ))
+        manifest = Path(
+            hf_hub_download(
+                repo_id=source,
+                filename=f"{task}/{PIKA_MANIFEST_FILENAME}",
+                revision=revision,
+            ))
+    except Exception as error:
+        raise PikaAdapterError(f"Could not download the Pika H3 Flex {task} adapter from {source!r}"
+                               f"{f' at revision {revision!r}' if revision else ''}: {error}") from error
+    return adapter, manifest
+
+
+def load_adapter(source: str | os.PathLike[str], *, task: str, revision: str | None = None) -> PikaAdapter:
+    """Resolve ``source`` (adapter file, release directory, task directory, or repo id)."""
+    trunk = expected_trunk(task)
+    raw = os.fspath(source)
+    path = Path(raw).expanduser()
+    if path.is_file():
+        adapter_path, manifest_path = path, path.with_name(PIKA_MANIFEST_FILENAME)
+    elif path.is_dir():
+        found = _adapter_in_directory(path, task)
+        if found is None:
+            raise PikaAdapterError(f"{raw!r} has no {task}/{PIKA_ADAPTER_FILENAME} or {PIKA_ADAPTER_FILENAME}.")
+        adapter_path, manifest_path = found
+    elif _looks_like_repo_id(raw):
+        adapter_path, manifest_path = _download_adapter(raw, task, revision)
+    else:
+        raise PikaAdapterError(f"Pika adapter path {raw!r} does not exist and is not a Hub repo id.")
+
+    if not adapter_path.is_file():
+        raise PikaAdapterError(f"Pika adapter weights are missing at {str(adapter_path)!r}.")
+    if not manifest_path.is_file():
+        raise PikaAdapterError(f"Pika adapter manifest is missing at {str(manifest_path)!r}; a bare "
+                               f"{PIKA_ADAPTER_FILENAME} cannot be checked against the release contract.")
+    manifest = _load_manifest(manifest_path)
+    validate_adapter_manifest(manifest, task=task, pipeline_trunk=trunk)
+    return PikaAdapter(
+        task=task,
+        trunk=trunk,
+        path=str(adapter_path),
+        manifest_path=str(manifest_path),
+        manifest=manifest,
+    )
+
+
+PIKA_MERGED_MARKER = "pika_h3_flex_merged.json"
+
+
+def merged_adapter(model_dir: str | os.PathLike[str], task: str) -> dict[str, Any] | None:
+    """The marker entry when ``model_dir``'s trunk for ``task`` already contains the Pika adapter, else None.
+
+    Exported checkpoints (e.g. NVFP4 QAD releases) merge the adapter into the trunk before quantizing it, so
+    applying the released adapter again would adapt the model twice. Such a checkpoint ships
+    ``pika_h3_flex_merged.json``: ``{"format_version": 1, "merged": {"<task>": {"trunk": ..., "release_id": ...,
+    "adapter_sha256": ...}}}``. A marker that does not cover ``task`` is an error: that trunk is absent or unadapted.
+    """
+    path = Path(model_dir) / PIKA_MERGED_MARKER
+    if not path.is_file():
+        return None
+    marker = _load_manifest(path)
+    if _as_int(marker.get("format_version"), "format_version") != PIKA_ADAPTER_FORMAT_VERSION:
+        raise PikaAdapterError(f"Unsupported {PIKA_MERGED_MARKER} format {marker.get('format_version')!r}.")
+    merged = marker.get("merged")
+    if not isinstance(merged, Mapping) or task not in merged:
+        raise PikaAdapterError(f"{str(path)!r} declares merged Pika adapters for {sorted(merged or {})}, not "
+                               f"{task!r}; this checkpoint does not carry the {task} model.")
+    entry = merged[task]
+    if not isinstance(entry, Mapping) or entry.get("trunk") != expected_trunk(task):
+        raise PikaAdapterError(f"{str(path)!r}: the merged {task} adapter must sit on trunk {expected_trunk(task)!r}.")
+    return dict(entry)
+
+
+def adapter_hf_patterns(task: str | None = None) -> list[str]:
+    """The adapter files the pipeline must include in a partial Hub snapshot."""
+    tasks = tuple(TASK_TO_TRUNK) if task is None else (task, )
+    patterns: list[str] = []
+    for item in tasks:
+        expected_trunk(item)
+        patterns.extend((f"{item}/{PIKA_ADAPTER_FILENAME}", f"{item}/{PIKA_MANIFEST_FILENAME}"))
+    return patterns
+
+
+__all__ = [
+    "PIKA_ADAPTER_ALPHA",
+    "PIKA_ADAPTER_FILENAME",
+    "PIKA_ADAPTER_FORMAT_VERSION",
+    "PIKA_ADAPTER_RANK",
+    "PIKA_BASE_REPO",
+    "PIKA_MANIFEST_FILENAME",
+    "PIKA_MERGED_MARKER",
+    "PIKA_SAMPLING_CONTRACT",
+    "PIKA_TARGET_MODULES",
+    "PikaAdapter",
+    "PikaAdapterError",
+    "TASK_TO_TRUNK",
+    "adapter_hf_patterns",
+    "expected_trunk",
+    "load_adapter",
+    "merged_adapter",
+    "task_for_pipeline",
+    "validate_adapter_manifest",
+]

fastvideo/pipelines/basic/pika_h3_flex/buckets.py

@@ -0,0 +1,294 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pika H3 Flex resolution buckets and frame grid.
+
+Pika H3 Flex keeps MiniMax-H3's canvas rule: sides sit on a 32-pixel lattice,
+a fixed short edge bounds the bucket, and an area cap keeps the token count
+finite.  The open-weight release samples at two of the released tiers, 480p and
+768p; the hosted API's larger tiers are intentionally not exposed here because
+the local checkpoints were not trained for them.
+
+A request selects a bucket, not a free-form canvas.  ``height`` and ``width``
+are therefore aspect hints: an exact bucket is honoured, anything else is
+snapped to the nearest aspect family of the tier implied by the requested
+short edge.  Keeping this arithmetic in a torch-free module lets the geometry
+contract be tested on a CPU-only host.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from collections.abc import Iterable
+
+PIKA_FPS = 24
+PIKA_CANVAS_MULTIPLE = 32
+PIKA_MIN_ASPECT_RATIO = 0.25
+PIKA_MAX_ASPECT_RATIO = 4.0
+
+# Landscape families first, then square, then the mirrored portraits.
+PIKA_ASPECTS: tuple[str, ...] = (
+    "21:9",
+    "16:9",
+    "3:2",
+    "4:3",
+    "5:4",
+    "1:1",
+    "4:5",
+    "3:4",
+    "2:3",
+    "9:16",
+    "9:21",
+)
+
+# 17n + 5 from 2.3 s to 15.1 s at 24 fps.  Open-weight requests are limited to
+# the documented 5-15 s range; the two smaller grid points remain part of the
+# released grid but are rejected by the request validator.
+PIKA_FRAME_GRID: tuple[int, ...] = tuple(17 * n + 5 for n in range(3, 22))
+PIKA_MIN_DURATION = 5.0
+PIKA_MAX_DURATION = 15.0
+PIKA_MIN_ALIGNED_FRAMES = 120
+PIKA_MAX_ALIGNED_FRAMES = 362
+
+
+@dataclass(frozen=True)
+class PikaTier:
+    """One released resolution tier."""
+
+    name: str
+    short_edge: int
+    max_pixels: int
+    reference_short_edge: int
+
+
+PIKA_TIERS: dict[str, PikaTier] = {
+    "p480": PikaTier("p480", short_edge=480, max_pixels=832 * 480, reference_short_edge=1024),
+    "p768": PikaTier("p768", short_edge=768, max_pixels=1344 * 768, reference_short_edge=2048),
+}
+DEFAULT_TIER = "p768"
+DEFAULT_ASPECT: tuple[int, int] = (16, 9)
+
+_TIER_ALIASES = {
+    "480p": "p480",
+    "p480": "p480",
+    "480": "p480",
+    "768p": "p768",
+    "p768": "p768",
+    "768": "p768",
+}
+
+
+def resolve_tier(value: str | None) -> str:
+    """Normalize ``480p`` / ``p480`` / ``768p`` / ``p768`` to a released tier name."""
+    if value is None:
+        return DEFAULT_TIER
+    key = str(value).strip().lower()
+    tier = _TIER_ALIASES.get(key)
+    if tier is None:
+        raise ValueError(f"Pika H3 Flex only samples the released 480p and 768p tiers; got {value!r}.")
+    return tier
+
+
+def tier_rule(tier: str | None) -> PikaTier:
+    name = resolve_tier(tier)
+    return PIKA_TIERS[name]
+
+
+def _parse_aspect(aspect: str | tuple[int, int] | tuple[float, float]) -> tuple[float, float]:
+    if isinstance(aspect, str):
+        parts = aspect.split(":")
+        if len(parts) != 2:
+            raise ValueError(f"An aspect ratio must be W:H, got {aspect!r}.")
+        width, height = float(parts[0]), float(parts[1])
+    elif isinstance(aspect, tuple) and len(aspect) == 2:
+        width, height = float(aspect[0]), float(aspect[1])
+    else:
+        raise TypeError(f"An aspect ratio must be a 'W:H' string or a (width, height) pair, got {aspect!r}.")
+    if not (math.isfinite(width) and math.isfinite(height)) or width <= 0 or height <= 0:
+        raise ValueError(f"An aspect ratio must be finite and positive, got {aspect!r}.")
+    return width, height
+
+
+def aspect_ratio(width: float, height: float) -> float:
+    width, height = _parse_aspect((width, height))
+    return width / height
+
+
+def canvas_for_aspect(aspect: str | tuple[float, float], tier: str | None = None) -> tuple[int, int]:
+    """``(height, width)`` of one tier's bucket for an aspect family."""
+    rule = tier_rule(tier)
+    aspect_width, aspect_height = _parse_aspect(aspect)
+    ratio = aspect_width / aspect_height
+    if not PIKA_MIN_ASPECT_RATIO <= ratio <= PIKA_MAX_ASPECT_RATIO:
+        raise ValueError(f"Pika H3 Flex supports aspect ratios from 1:4 to 4:1, got {aspect!r}.")
+
+    if ratio >= 1:
+        width, height = rule.short_edge * ratio, float(rule.short_edge)
+    else:
+        width, height = float(rule.short_edge), rule.short_edge / ratio
+    area = width * height
+    if area > rule.max_pixels:
+        scale = (rule.max_pixels / area)**0.5
+        width, height = width * scale, height * scale
+    multiple = PIKA_CANVAS_MULTIPLE
+    return (
+        max(multiple,
+            round(height / multiple) * multiple),
+        max(multiple,
+            round(width / multiple) * multiple),
+    )
+
+
+def aspect_families(tier: str | None = None) -> tuple[str, ...]:
+    del tier
+    return PIKA_ASPECTS
+
+
+def bucket_table(tier: str | None = None) -> tuple[tuple[int, int], ...]:
+    """All ``(width, height)`` buckets of one tier, in :data:`PIKA_ASPECTS` order."""
+    return tuple((canvas_for_aspect(aspect, tier)[1], canvas_for_aspect(aspect, tier)[0]) for aspect in PIKA_ASPECTS)
+
+
+def iter_buckets() -> Iterable[tuple[str, str, int, int]]:
+    """Yield ``(tier, aspect, width, height)`` for every released bucket."""
+    for tier in PIKA_TIERS:
+        for aspect in PIKA_ASPECTS:
+            height, width = canvas_for_aspect(aspect, tier)
+            yield tier, aspect, width, height
+
+
+def exact_bucket(height: int, width: int) -> tuple[str, str] | None:
+    """The ``(tier, aspect)`` whose bucket is exactly ``height x width``, if any."""
+    for tier, aspect, bucket_width, bucket_height in iter_buckets():
+        if (bucket_height, bucket_width) == (int(height), int(width)):
+            return tier, aspect
+    return None
+
+
+def tier_for_canvas(height: int, width: int) -> str:
+    """The smallest released tier whose short edge covers the canvas."""
+    if height <= 0 or width <= 0:
+        raise ValueError(f"A canvas must be positive, got {height}x{width}.")
+    short = min(int(height), int(width))
+    for tier in sorted(PIKA_TIERS.values(), key=lambda rule: rule.short_edge):
+        if short <= tier.short_edge:
+            return tier.name
+    return max(PIKA_TIERS.values(), key=lambda rule: rule.short_edge).name
+
+
+def reference_short_edge_for_canvas(height: int, width: int) -> int:
+    """Reference-image short edge the released tier rule uses for an output canvas."""
+    return PIKA_TIERS[tier_for_canvas(height, width)].reference_short_edge
+
+
+def _nearest_aspect(ratio: float, tier: str) -> tuple[int, int, str]:
+    best: tuple[float, str, tuple[int, int]] | None = None
+    for aspect in PIKA_ASPECTS:
+        family_ratio = aspect_ratio(*_parse_aspect(aspect))
+        distance = abs(math.log(ratio / family_ratio))
+        if best is None or distance < best[0]:
+            best = (distance, aspect, canvas_for_aspect(aspect, tier))
+        # PIKA_ASPECTS is ordered by family; equal distances keep the earlier one.
+    assert best is not None
+    height, width = best[2]
+    return height, width, best[1]
+
+
+def nearest_bucket(height: int, width: int, tier: str) -> tuple[int, int, str]:
+    """Snap ``height x width`` to the closest aspect family of ``tier``."""
+    if height <= 0 or width <= 0:
+        raise ValueError(f"A canvas must be positive, got {height}x{width}.")
+    ratio = width / height
+    if not PIKA_MIN_ASPECT_RATIO - 1e-9 <= ratio <= PIKA_MAX_ASPECT_RATIO + 1e-9:
+        raise ValueError(f"Pika H3 Flex supports aspect ratios from 1:4 to 4:1, got {width}x{height}.")
+    return _nearest_aspect(ratio, tier)
+
+
+def resolve_canvas(
+    height: int | None,
+    width: int | None,
+    *,
+    default_aspect: str | tuple[float, float] = DEFAULT_ASPECT,
+    default_tier: str | None = DEFAULT_TIER,
+) -> tuple[int, int, str, str]:
+    """Resolve a request canvas to ``(height, width, tier, aspect)``.
+
+    With neither side set, the default tier and aspect are used.  With both
+    sides set, an exact released bucket is honoured; otherwise the tier follows
+    the requested short edge (``<= 480`` is 480p, everything else 768p) and the
+    canvas snaps to the nearest aspect family.  One-sided requests are rejected.
+    """
+    if (height is None) != (width is None):
+        raise ValueError("Pika H3 Flex `height` and `width` must be passed together, or neither.")
+    if height is None:
+        tier = resolve_tier(default_tier)
+        aspect_width, aspect_height = _parse_aspect(default_aspect)
+        bucket_height, bucket_width = canvas_for_aspect((aspect_width, aspect_height), tier)
+        aspect = PIKA_ASPECTS[0]
+        for family in PIKA_ASPECTS:
+            if canvas_for_aspect(family, tier) == (bucket_height, bucket_width):
+                aspect = family
+                break
+        return bucket_height, bucket_width, tier, aspect
+
+    if not isinstance(height, int) or isinstance(height, bool) or not isinstance(width, int) or isinstance(width, bool):
+        raise TypeError("Pika H3 Flex `height` and `width` must be integers.")
+    if height <= 0 or width <= 0:
+        raise ValueError(f"Pika H3 Flex `height` and `width` must be positive, got {height}x{width}.")
+
+    exact = exact_bucket(height, width)
+    if exact is not None:
+        tier, aspect = exact
+        return height, width, tier, aspect
+
+    tier = "p480" if min(height, width) <= PIKA_TIERS["p480"].short_edge else "p768"
+    bucket_height, bucket_width, aspect = nearest_bucket(height, width, tier)
+    return bucket_height, bucket_width, tier, aspect
+
+
+def resolve_frames(num_frames: int) -> int:
+    """Round a request up to Pika's ``17n + 5`` grid inside the 5-15 s range."""
+    if not isinstance(num_frames, int) or isinstance(num_frames, bool):
+        raise TypeError(f"`num_frames` must be an integer, got {type(num_frames).__name__}.")
+    if num_frames < 1:
+        raise ValueError(f"`num_frames` must be positive, got {num_frames}.")
+    aligned = num_frames
+    while aligned % 17 != 5:
+        aligned += 1
+    if not PIKA_MIN_ALIGNED_FRAMES <= aligned <= PIKA_MAX_ALIGNED_FRAMES:
+        raise ValueError(f"Pika H3 Flex generates {PIKA_MIN_DURATION:g}-{PIKA_MAX_DURATION:g} seconds at "
+                         f"{PIKA_FPS} fps; aligned num_frames={aligned} is outside the accepted "
+                         f"{PIKA_MIN_ALIGNED_FRAMES}-{PIKA_MAX_ALIGNED_FRAMES} frame range.")
+    return aligned
+
+
+def duration_seconds(num_frames: int) -> float:
+    return int(num_frames) / PIKA_FPS
+
+
+__all__ = [
+    "DEFAULT_ASPECT",
+    "DEFAULT_TIER",
+    "PIKA_ASPECTS",
+    "PIKA_CANVAS_MULTIPLE",
+    "PIKA_FPS",
+    "PIKA_FRAME_GRID",
+    "PIKA_MAX_ALIGNED_FRAMES",
+    "PIKA_MAX_DURATION",
+    "PIKA_MIN_ALIGNED_FRAMES",
+    "PIKA_MIN_DURATION",
+    "PIKA_TIERS",
+    "PikaTier",
+    "aspect_ratio",
+    "bucket_table",
+    "canvas_for_aspect",
+    "duration_seconds",
+    "exact_bucket",
+    "iter_buckets",
+    "nearest_bucket",
+    "reference_short_edge_for_canvas",
+    "resolve_canvas",
+    "resolve_frames",
+    "resolve_tier",
+    "tier_for_canvas",
+    "tier_rule",
+]

fastvideo/pipelines/basic/pika_h3_flex/pipeline.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pika H3 Flex pipeline: MiniMax-H3 components, Pika adapter and sampler."""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+from fastvideo.configs.pipelines.pika_h3_flex import (
+    PIKA_AUDIO_SCHEDULER_SHIFT,
+    PIKA_VIDEO_SCHEDULER_SHIFT,
+    PikaH3FlexPipelineConfig,
+)
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import MiniMaxH3BasePipeline
+from fastvideo.pipelines.basic.pika_h3_flex.adapter import (PIKA_MERGED_MARKER, adapter_hf_patterns, load_adapter,
+                                                            merged_adapter, task_for_pipeline)
+from fastvideo.pipelines.basic.pika_h3_flex.stages.denoising import PikaH3FlexDenoisingStage
+from fastvideo.pipelines.basic.pika_h3_flex.stages.input_preparation import PikaH3FlexInputPreparationStage
+
+logger = init_logger(__name__)
+
+
+class PikaH3FlexBasePipeline(MiniMaxH3BasePipeline):
+    """Pika H3 Flex's adapter, bucket geometry, and six-forward re-noise sampler."""
+
+    pipeline_config_cls = PikaH3FlexPipelineConfig
+    input_preparation_stage_cls = PikaH3FlexInputPreparationStage
+    denoising_stage_cls = PikaH3FlexDenoisingStage
+    _ref2va_default = False
+    _pika_adapter: Any = None
+
+    @classmethod
+    def get_hf_download_allow_patterns(cls) -> list[str]:
+        patterns = super().get_hf_download_allow_patterns()
+        # The concrete class picks the adapter, so a T2VA download does not
+        # fetch the 2.7 GB Ref2VA adapter (and vice versa).
+        task = "ref2va" if issubclass(cls, PikaH3FlexRefPipeline) else "fl2va"
+        return [*patterns, *adapter_hf_patterns(task), PIKA_MERGED_MARKER]
+
+    def _resolve_pika_adapter(self) -> None:
+        """Point ``lora_path`` at the released adapter before LoRA layers attach.
+
+        ``_load_config`` runs before ``LoRAPipeline`` reads ``lora_path``, so the
+        automatic adapter behaves exactly like an explicitly passed one,
+        including the lazy/sequential transformer paths.
+        """
+        task = task_for_pipeline(self._ref2va)
+        explicit = self.fastvideo_args.lora_path
+        merged = merged_adapter(self.model_path, task) if explicit is None else None
+        if merged is not None:
+            # The trunk already carries the adapter (an exported, e.g. NVFP4 QAD, checkpoint): attach nothing.
+            logger.info("Pika H3 Flex: the %s trunk of %s already contains adapter %s; not applying it again", task,
+                        self.model_path, merged.get("release_id"))
+            self._pika_adapter = None
+            return
+        source = explicit if explicit is not None else self.model_path
+        revision = None if explicit is not None else self.fastvideo_args.revision
+        adapter = load_adapter(source, task=task, revision=revision)
+        source_kind = "explicit" if explicit is not None else "released"
+        logger.info("Pika H3 Flex: loaded the %s %s adapter from %s", source_kind, task, adapter.path)
+        self._pika_adapter = adapter
+        self.fastvideo_args.lora_path = adapter.path
+
+    def _load_config(self, model_path: str) -> dict[str, Any]:
+        config = super()._load_config(model_path)
+        self._resolve_pika_adapter()
+        return config
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None:
+        super().initialize_pipeline(fastvideo_args)
+        expected = {"video": PIKA_VIDEO_SCHEDULER_SHIFT, "audio": PIKA_AUDIO_SCHEDULER_SHIFT}
+        for module_name, modality in (("scheduler", "video"), ("audio_scheduler", "audio")):
+            actual = float(self.get_module(module_name).shift)
+            target = expected[modality]
+            if not math.isfinite(actual) or abs(actual - target) > 1e-9:
+                raise ValueError(f"Pika H3 Flex samples with {modality} shift {target:g}, but "
+                                 f"{module_name}/scheduler_config.json declares {actual:g}.")
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        self._add_stages(fastvideo_args, ref2va=self._ref2va_default)
+
+
+class PikaH3FlexPipeline(PikaH3FlexBasePipeline):
+    """T2VA/I2VA/L2VA/FL2VA pipeline through the ``fl2va`` adapter."""
+
+    _ref2va_default = False
+
+
+class PikaH3FlexRefPipeline(PikaH3FlexBasePipeline):
+    """Ordered-reference pipeline through the ``ref2va`` adapter and ``transformer_ref``."""
+
+    _extra_config_module_map = {"transformer": "transformer_ref"}
+    _ref2va_default = True
+
+
+class PikaH3FlexModularPipeline(PikaH3FlexPipeline):
+    """Public entry matching the released ``modular_model_index.json`` class name."""
+
+
+class PikaH3FlexRef2VAModularPipeline(PikaH3FlexRefPipeline):
+    """Public Ref2VA entry matching the released modular class name."""
+
+
+EntryClass = [PikaH3FlexModularPipeline, PikaH3FlexRef2VAModularPipeline]
+
+__all__ = [
+    "EntryClass",
+    "PikaH3FlexBasePipeline",
+    "PikaH3FlexModularPipeline",
+    "PikaH3FlexPipeline",
+    "PikaH3FlexRef2VAModularPipeline",
+    "PikaH3FlexRefPipeline",
+]

fastvideo/pipelines/basic/pika_h3_flex/presets.py

@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Inference presets for Pika H3 Flex joint video/audio generation."""
+
+from fastvideo.api.presets import InferencePreset, PresetStageSpec
+
+# The six-forward ladder is the adapter's training contract. Exposing a stage
+# override would let a request silently sample the adapter outside its
+# distribution, so the denoise stage accepts no overrides.
+_DENOISE_STAGE = PresetStageSpec(
+    name="denoise",
+    kind="denoising",
+    description="Pika's fixed six-forward re-noise ladder",
+    allowed_overrides=frozenset(),
+)
+
+_SHARED_DEFAULTS = {
+    "fps": 24,
+    "guidance_scale": 1.0,
+    "batch_cfg": False,
+    "negative_prompt": "",
+    "num_inference_steps": 6,
+    "seed": 0,
+    "height": 768,
+    "width": 1344,
+    "num_frames": 124,
+}
+
+PIKA_H3_FLEX_T2VA = InferencePreset(
+    name="pika_h3_flex_t2va",
+    version=1,
+    model_family="pika_h3_flex",
+    description="Pika H3 Flex text-to-video with synchronized stereo audio at 768p",
+    workload_type="t2v",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults=dict(_SHARED_DEFAULTS),
+)
+
+PIKA_H3_FLEX_FL2VA = InferencePreset(
+    name="pika_h3_flex_fl2va",
+    version=1,
+    model_family="pika_h3_flex",
+    description="Pika H3 Flex first/last-frame-to-video with synchronized stereo audio",
+    workload_type="i2v",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults=dict(_SHARED_DEFAULTS),
+)
+
+PIKA_H3_FLEX_REF2VA = InferencePreset(
+    name="pika_h3_flex_ref2va",
+    version=1,
+    model_family="pika_h3_flex",
+    description="Pika H3 Flex ordered image/video/audio references to joint video and stereo audio",
+    workload_type="i2v",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults=dict(_SHARED_DEFAULTS),
+)
+
+ALL_PRESETS = (PIKA_H3_FLEX_T2VA, PIKA_H3_FLEX_FL2VA, PIKA_H3_FLEX_REF2VA)
+
+__all__ = [
+    "ALL_PRESETS",
+    "PIKA_H3_FLEX_FL2VA",
+    "PIKA_H3_FLEX_REF2VA",
+    "PIKA_H3_FLEX_T2VA",
+]

fastvideo/pipelines/basic/pika_h3_flex/stages/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from fastvideo.pipelines.basic.pika_h3_flex.stages.denoising import PikaH3FlexDenoisingStage
+from fastvideo.pipelines.basic.pika_h3_flex.stages.input_preparation import PikaH3FlexInputPreparationStage
+
+__all__ = ["PikaH3FlexDenoisingStage", "PikaH3FlexInputPreparationStage"]

fastvideo/pipelines/basic/pika_h3_flex/stages/denoising.py

@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pika H3 Flex's fixed six-forward re-noise denoiser."""
+
+from __future__ import annotations
+
+import torch
+
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_denoising import MiniMaxH3DenoisingStage
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+
+
+def request_seed(batch: ForwardBatch) -> int:
+    """The request seed, falling back to the generator's initial seed."""
+    if batch.seed is not None:
+        return int(batch.seed)
+    generator = batch.generator
+    if isinstance(generator, list):
+        generator = generator[0] if generator else None
+    if generator is None:
+        return 0
+    return int(generator.initial_seed())
+
+
+class PikaH3FlexDenoisingStage(MiniMaxH3DenoisingStage):
+    """Six full-model forwards over Pika's clean-time ladder.
+
+    Pika's adapter is not a PDD student: every rung is an ordinary transformer
+    forward.  What differs from MiniMax-H3 is the ladder and the re-noise step,
+    which draws fresh Gaussian noise from a generator seeded ``seed + 1``.
+    """
+
+    def _configure_schedules(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+        device: torch.device,
+    ):
+        steps = list(fastvideo_args.pipeline_config.pika_denoising_steps)
+        if batch.num_inference_steps != len(steps):
+            raise ValueError(f"Pika H3 Flex needs {len(steps)} transformer forwards plus the terminal zero, got "
+                             f"num_inference_steps={batch.num_inference_steps}.")
+        for scheduler in (self.scheduler, self.audio_scheduler):
+            if not getattr(scheduler, "stochastic_sampling", False):
+                raise ValueError("Pika H3 Flex needs the released stochastic (re-noise) scheduler; "
+                                 f"{type(scheduler).__name__} did not load stochastic_sampling=true.")
+        self._set_shifted_schedule(steps, len(steps) + 1, device, label="Pika H3 Flex")
+        return None
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        generator = torch.Generator(device=get_local_torch_device()).manual_seed(request_seed(batch) + 1)
+        self.scheduler.set_noise_generator(generator)
+        self.audio_scheduler.set_noise_generator(generator)
+        return super().forward(batch, fastvideo_args)
+
+
+__all__ = ["PikaH3FlexDenoisingStage", "request_seed"]

fastvideo/pipelines/basic/pika_h3_flex/stages/input_preparation.py

@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pika H3 Flex input preparation: bucket canvas, cropped keyframes, tiered references."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from PIL import Image, ImageOps
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.vision_utils import load_image
+from fastvideo.pipelines.basic.minimax_h3.packing import prepare_keyframe_image
+from fastvideo.pipelines.basic.minimax_h3.reference import (
+    MiniMaxH3PreparedReference,
+    prepare_reference,
+    validate_references,
+)
+from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_input_preparation import (
+    MINIMAX_H3_KEYFRAME_ANCHORS_KEY,
+    MINIMAX_H3_KEYFRAMES_KEY,
+    MiniMaxH3InputPreparationStage,
+    prepare_common_request,
+)
+from fastvideo.pipelines.basic.pika_h3_flex.buckets import (
+    PIKA_TIERS,
+    resolve_canvas,
+    resolve_frames,
+    resolve_tier,
+)
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+
+logger = init_logger(__name__)
+
+# Pika's released reference-video floor: two seconds at 24 fps is 48 frames,
+# but the adapter's own validator requires at least 22 decoded frames.
+PIKA_MIN_REFERENCE_VIDEO_FRAMES = 22
+
+
+class PikaH3FlexInputPreparationStage(MiniMaxH3InputPreparationStage):
+    """Prepare T2VA/I2VA/L2VA/FL2VA or Ref2VA inputs on a Pika bucket."""
+
+    def __init__(self, vae: Any, audio_vae: Any | None = None, *, ref2va: bool = False) -> None:
+        super().__init__(vae, audio_vae, ref2va=ref2va)
+        self._tier = "p768"
+
+    def _request_tier(self, fastvideo_args: FastVideoArgs) -> str:
+        return resolve_tier(getattr(fastvideo_args.pipeline_config, "pika_resolution_tier", None))
+
+    def _resolve_canvas(
+        self,
+        batch: ForwardBatch,
+        default_aspect: tuple[int, int],
+    ) -> tuple[int, int, str, str]:
+        height, width, tier, aspect = resolve_canvas(
+            batch.height,
+            batch.width,
+            default_aspect=default_aspect,
+            default_tier=self._tier,
+        )
+        if (batch.height, batch.width) != (height, width):
+            logger.info("Pika H3 Flex snapped canvas %sx%s to %sx%s (%s %s).", batch.height, batch.width, height, width,
+                        tier, aspect)
+        ratio = int(self.vae.spatial_compression_ratio)
+        if height % ratio or width % ratio:
+            raise ValueError(f"Pika H3 Flex canvas {height}x{width} is not divisible by VAE ratio {ratio}.")
+        return height, width, tier, aspect
+
+    def _prepare_fl2va(self, batch: ForwardBatch) -> None:
+        if batch.references:
+            raise ValueError("Pika H3 Flex references belong to the Ref2VA pipeline.")
+
+        if batch.image_path is not None:
+            batch.pil_image = load_image(batch.image_path)
+        for name, image in (("pil_image", batch.pil_image), ("last_image", batch.last_image)):
+            if image is not None and not isinstance(image, Image.Image):
+                raise TypeError(f"Pika H3 Flex `{name}` must be a PIL image, got {type(image).__name__}.")
+
+        raw_keyframes = [(anchor, ImageOps.exif_transpose(image).convert("RGB"))
+                         for anchor, image in (("first", batch.pil_image), ("last", batch.last_image))
+                         if image is not None]
+        default_aspect = raw_keyframes[0][1].size if raw_keyframes else (16, 9)
+        height, width, _, _ = self._resolve_canvas(batch, default_aspect)
+        num_frames = resolve_frames(batch.num_frames)
+        ratio = int(self.vae.spatial_compression_ratio)
+        self._write_target_geometry(batch, height, width, ratio, num_frames)
+
+        batch.extra[MINIMAX_H3_KEYFRAME_ANCHORS_KEY] = tuple(anchor for anchor, _ in raw_keyframes)
+        # Pika crops both endpoints. MiniMax-H3 stretches the first frame; using
+        # that fit here would distort every Pika image on a bucketed canvas.
+        batch.extra[MINIMAX_H3_KEYFRAMES_KEY] = [
+            prepare_keyframe_image(image, height, width, stretch=False) for _, image in raw_keyframes
+        ]
+
+    def _prepare_ref2va(self, batch: ForwardBatch) -> None:
+        if batch.image_path is not None or batch.pil_image is not None or batch.last_image is not None:
+            raise ValueError("Pika H3 Flex Ref2VA accepts media through `references`, not FL2VA keyframe fields.")
+        references = validate_references(list(batch.references or []))
+        if self.audio_vae is None:
+            raise RuntimeError("Pika H3 Flex Ref2VA input preparation has no audio VAE.")
+
+        height, width, tier, _ = self._resolve_canvas(batch, (16, 9))
+        num_frames = resolve_frames(batch.num_frames)
+        target_sample_rate = int(self.audio_vae.sampling_rate)
+        rule = PIKA_TIERS[tier]
+        prepared: list[MiniMaxH3PreparedReference] = []
+        for reference in references:
+            item = prepare_reference(
+                reference,
+                num_frames,
+                target_sample_rate,
+                reference_short_edge=rule.reference_short_edge,
+                reference_canvas_short_edge=rule.short_edge,
+                reference_canvas_max_pixels=rule.max_pixels,
+            )
+            if (item.media_type == "video" and item.frames is not None
+                    and item.frames.shape[0] < PIKA_MIN_REFERENCE_VIDEO_FRAMES):
+                raise ValueError(f"A Pika H3 Flex reference video needs at least "
+                                 f"{PIKA_MIN_REFERENCE_VIDEO_FRAMES} decoded frames at 24 fps; got "
+                                 f"{item.frames.shape[0]}.")
+            prepared.append(item)
+        if not prepared:
+            raise ValueError("Pika H3 Flex Ref2VA requires at least one non-audio reference.")
+
+        batch.references = prepared
+        self._write_target_geometry(batch, height, width, int(self.vae.spatial_compression_ratio), num_frames)
+        batch.extra[MINIMAX_H3_KEYFRAMES_KEY] = []
+        batch.extra[MINIMAX_H3_KEYFRAME_ANCHORS_KEY] = ()
+
+    @torch.no_grad()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        prepare_common_request(batch)
+        self._tier = self._request_tier(fastvideo_args)
+        if self.ref2va:
+            self._prepare_ref2va(batch)
+        else:
+            self._prepare_fl2va(batch)
+        return batch
+
+
+__all__ = ["PIKA_MIN_REFERENCE_VIDEO_FRAMES", "PikaH3FlexInputPreparationStage"]

fastvideo/pipelines/lazy_module.py

@@ -75,9 +75,18 @@ class LazyModule:
 
         transform = object.__getattribute__(self, "_lazy_materialize_transform")
         if transform is not None:
-            module = transform(module)
-            if module is None:
-                raise ValueError(f"Materialize transform for module {name} returned None")
+            # Publish the loaded instance before the transform runs: a transform
+            # may read attributes back through this proxy (the constructor LoRA
+            # adapter reads ``param_names_mapping``), and with nothing published
+            # that access re-enters ``materialize`` and reloads the component.
+            object.__setattr__(self, "_lazy_module", module)
+            try:
+                module = transform(module)
+                if module is None:
+                    raise ValueError(f"Materialize transform for module {name} returned None")
+            except BaseException:
+                object.__setattr__(self, "_lazy_module", None)
+                raise
         object.__setattr__(self, "_lazy_module", module)
 
         allocated = _cuda_allocated_gib()

fastvideo/pipelines/lora_pipeline.py

@@ -656,6 +656,7 @@ class LoRAPipeline(ComposedPipelineBase):
         # Merge the new adapter
         adapted_count = 0
         consumed: set[str] = set()
+        untouched: list[str] = []
         for (
                 transformer_name,
                 transformer_lora_layers,
@@ -696,17 +697,21 @@ class LoRAPipeline(ComposedPipelineBase):
                             adapted_count += 1
                             consumed.update((lora_A_name, lora_B_name, lora_alpha_name))
                         else:
-                            if rank == 0:
-                                logger.warning(
-                                    "LoRA adapter %s does not contain the weights for layer %s. LoRA will not be applied to it.",
-                                    lora_path,
-                                    name,
-                                )
+                            untouched.append(name)
                             layer.disable_lora = True
                     # Repack inside the offload scope, while the block's parameters are resident.
                     repack_scope = [module] if module is not None else list(layers.values())
                     _convert_quantized_weights_after_lora_change(repack_scope,
                                                                  skip=_all_side_branch(layers.values(), repack_scope))
+        if untouched and rank == 0:
+            # One summary, not one line per wrapped layer: an adapter that trains a subset of the
+            # wrapped linears (e.g. no AdaLN projections) would otherwise log hundreds of warnings.
+            logger.warning(
+                "LoRA adapter %s has no weights for %d wrapped layers; LoRA is disabled on them (e.g. %s).",
+                lora_path,
+                len(untouched),
+                ", ".join(untouched[:5]),
+            )
         logger.info(
             "Rank %d: LoRA adapter %s applied to %d layers",
             rank,

fastvideo/registry.py

@@ -42,6 +42,7 @@ from fastvideo.configs.pipelines.flux_2 import (
 from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig
 from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig
 from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig
+from fastvideo.configs.pipelines.pika_h3_flex import PikaH3FlexPipelineConfig
 from fastvideo.configs.pipelines.mmaudio import (
     MMAudioLarge44kV2PipelineConfig,
     MMAudioLarge44kV2AConfig,
@@ -60,6 +61,7 @@ from fastvideo.api.sampling_param import SamplingParam
 from fastvideo.api.matrixgame2 import MatrixGame2SamplingParam
 from fastvideo.api.matrixgame3 import MatrixGame3SamplingParam
 from fastvideo.api.flux import FluxSamplingParam
+from fastvideo.api.pika_h3_flex import PikaH3FlexSamplingParam
 
 from fastvideo.fastvideo_args import WorkloadType
 from fastvideo.logger import init_logger
@@ -1152,6 +1154,26 @@ def _register_configs() -> None:
 
     _register_wan_configs(WAN_MODEL_DEFINITIONS)
 
+    # Pika H3 Flex. It reuses MiniMax-H3's components, but its six-forward
+    # re-noise sampler and bucket geometry are a separate family contract.
+    register_configs(
+        sampling_param_cls=PikaH3FlexSamplingParam,
+        pipeline_config_cls=PikaH3FlexPipelineConfig,
+        workload_types=(WorkloadType.T2V, WorkloadType.I2V),
+        hf_model_paths=["Pika-Labs/Pika-H3-Flex"],
+        model_detectors=[
+            lambda path: any(token in path.lower() for token in (
+                "pika-h3-flex",
+                "pika_h3_flex",
+                "pikah3flexmodularpipeline",
+                "pikah3flexref2vamodularpipeline",
+            )),
+        ],
+        model_family="pika_h3_flex",
+        default_preset="pika_h3_flex_t2va",
+        pipeline_cls_name="PikaH3FlexModularPipeline",
+    )
+
     # MiniMax H3
     register_configs(
         sampling_param_cls=None,
@@ -1350,6 +1372,8 @@ def _register_presets() -> None:
         ALL_PRESETS as MATRIXGAME3_PRESETS, )
     from fastvideo.pipelines.basic.minimax_h3.presets import (
         ALL_PRESETS as MINIMAX_H3_PRESETS, )
+    from fastvideo.pipelines.basic.pika_h3_flex.presets import (
+        ALL_PRESETS as PIKA_H3_FLEX_PRESETS, )
     from fastvideo.pipelines.basic.mmaudio.presets import (
         ALL_PRESETS as MMAUDIO_PRESETS, )
     from fastvideo.pipelines.basic.sd35.presets import (
@@ -1383,6 +1407,7 @@ def _register_presets() -> None:
         MATRIXGAME2_PRESETS,
         MATRIXGAME3_PRESETS,
         MINIMAX_H3_PRESETS,
+        PIKA_H3_FLEX_PRESETS,
         MMAUDIO_PRESETS,
         SD35_PRESETS,
         STABLE_AUDIO_PRESETS,

fastvideo/tests/api/test_compat_translation.py

@@ -11,7 +11,7 @@ from fastvideo.api.compat import (
     request_to_sampling_param,
 )
 from fastvideo.api.parser import parse_config
-from fastvideo.api.schema import CompileConfig, GenerationRequest, GeneratorConfig
+from fastvideo.api.schema import CompileConfig, ComponentConfig, GenerationRequest, GeneratorConfig, PipelineSelection
 from fastvideo.api.sampling_param import SamplingParam
 
 
@@ -78,6 +78,37 @@ class TestCompileConfigRoundTrip:
         assert args.kwargs["torch_compile_kwargs"] == {}
 
 
+class TestH3Nvfp4ComponentTranslation:
+    """Pika/H3 NVFP4 decoder and text-encoder overrides stay on ComponentConfig."""
+
+    def test_decoder_and_text_encoder_translate(self, monkeypatch) -> None:
+        _stub_fastvideo_args_from_kwargs(monkeypatch)
+        config = GeneratorConfig(
+            model_path="/models/pika",
+            pipeline=PipelineSelection(components=ComponentConfig(
+                h3_vae_nvfp4_decoder="/models/quality-36L/decoder-nvfp4.safetensors",
+                text_encoder_weights="/models/nvfp4-encoder/text_encoder",
+            )),
+        )
+        args = generator_config_to_fastvideo_args(config)
+        assert args.kwargs["h3_vae_nvfp4_decoder"] == "/models/quality-36L/decoder-nvfp4.safetensors"
+        assert args.kwargs["override_text_encoder_safetensors"] == "/models/nvfp4-encoder/text_encoder"
+
+    def test_legacy_flat_decoder_promotes_to_component(self) -> None:
+        config = legacy_from_pretrained_to_config(
+            "/models/pika",
+            {"h3_vae_nvfp4_decoder": "/models/fast-8L/decoder-nvfp4.safetensors"},
+        )
+        assert config.pipeline.components.h3_vae_nvfp4_decoder == (
+            "/models/fast-8L/decoder-nvfp4.safetensors")
+
+    def test_unset_fields_are_not_emitted(self, monkeypatch) -> None:
+        _stub_fastvideo_args_from_kwargs(monkeypatch)
+        args = generator_config_to_fastvideo_args(GeneratorConfig(model_path="/models/pika"))
+        assert "h3_vae_nvfp4_decoder" not in args.kwargs
+        assert "override_text_encoder_safetensors" not in args.kwargs
+
+
 class TestLegacyLtx2VaeTilingTranslation:
     """``ltx2_vae_tiling`` flat kwarg promotes to
     ``generator.pipeline.vae_tiling``; reverse direction emits the

fastvideo/tests/attention/test_nvfp4_qk_transform_deploy.py

@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Deployed NVFP4 FA4 Q/K preconditioning and per-block BF16 opt-out.
+
+The transform must be applied to Q/K right before NVFP4 quantization with the same function the QAD emulator uses,
+and MiniMax-H3 blocks listed in ``FASTVIDEO_NVFP4_FA4_BF16_BLOCKS`` must build BF16 attention impls.
+"""
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+import fastvideo.envs as envs  # noqa: E402
+from fastvideo.attention.utils.nvfp4_qk_precondition import (  # noqa: E402
+    parse_block_indices, parse_qk_transform, precondition_qk,
+)
+
+flash_attn_module = pytest.importorskip(
+    "fastvideo.attention.backends.flash_attn",
+    reason="no usable flash-attention package installed",
+    exc_type=ImportError,
+)
+FlashAttentionImpl = flash_attn_module.FlashAttentionImpl
+
+
+def _impl(**extra) -> "FlashAttentionImpl":
+    return FlashAttentionImpl(num_heads=2, head_size=64, causal=False, softmax_scale=0.125, num_kv_heads=2, **extra)
+
+
+@pytest.fixture
+def fa4_runtime_available():
+    with (
+        patch.object(flash_attn_module, "_FA4_FP4_AVAILABLE", True),
+        patch("torch.cuda.get_device_capability", return_value=(10, 0)),
+    ):
+        yield
+
+
+def test_block_list_parsing() -> None:
+    assert parse_block_indices("") == frozenset()
+    assert parse_block_indices(None) == frozenset()
+    assert parse_block_indices(" 45, 49,45 ") == frozenset({45, 49})
+    for bad in ("a,1", "-1", "4.5"):
+        with pytest.raises(ValueError):
+            parse_block_indices(bad)
+
+
+def test_transform_env_default_and_validation(env_overrides) -> None:
+    env_overrides.enter_context(envs.FASTVIDEO_NVFP4_FA4_QK_TRANSFORM.override(None))
+    assert envs.FASTVIDEO_NVFP4_FA4_QK_TRANSFORM.get() == "none"
+    assert _impl().nvfp4_smooth_k is False and _impl().nvfp4_hadamard is False
+    env_overrides.enter_context(envs.FASTVIDEO_NVFP4_FA4_QK_TRANSFORM.override("Smooth_K+Hadamard"))
+    impl = _impl()
+    assert (impl.nvfp4_smooth_k, impl.nvfp4_hadamard) == (True, True)
+    env_overrides.enter_context(envs.FASTVIDEO_NVFP4_FA4_QK_TRANSFORM.override("bogus"))
+    with pytest.raises(envs.EnvVarError):
+        envs.FASTVIDEO_NVFP4_FA4_QK_TRANSFORM.get()
+
+
+def test_impl_arg_wins_over_transform_env(env_overrides) -> None:
+    env_overrides.enter_context(envs.FASTVIDEO_NVFP4_FA4_QK_TRANSFORM.override("hadamard"))
+    impl = _impl(nvfp4_qk_transform="smooth_k")
+    assert (impl.nvfp4_smooth_k, impl.nvfp4_hadamard) == (True, False)
+    with pytest.raises(ValueError):
+        _impl(nvfp4_qk_transform="bogus")
+
+
+@pytest.mark.parametrize("name", ["none", "smooth_k", "hadamard", "smooth_k+hadamard"])
+def test_transform_applied_before_quantization(name, env_overrides, fa4_runtime_available) -> None:
+    env_overrides.enter_context(envs.FASTVIDEO_NVFP4_FA4.override(None))
+    impl = _impl(nvfp4_fa4=True, nvfp4_qk_transform=name)
+    torch.manual_seed(0)
+    q = torch.randn(1, 20, 2, 64, dtype=torch.bfloat16) + 3.0
+    k = torch.randn(1, 20, 2, 64, dtype=torch.bfloat16) + 3.0
+    v = torch.randn(1, 20, 2, 64, dtype=torch.bfloat16)
+    seen = []
+
+    def fake_quant(tensor):
+        seen.append(tensor.clone())
+        return tensor, tensor
+
+    fp4 = MagicMock(side_effect=lambda *a, **kw: a[2])
+    with (
+        patch.object(flash_attn_module, "_nvfp4_quantize_for_fa4", side_effect=fake_quant),
+        patch.object(flash_attn_module, "flash_attn_fp4_func", fp4),
+    ):
+        impl._forward_nvfp4(q, k, v)
+
+    smooth_k, hadamard = parse_qk_transform(name)
+    expected_q, expected_k = precondition_qk(q, k, smooth_k=smooth_k, hadamard=hadamard)
+    assert torch.equal(seen[0], expected_q) and torch.equal(seen[1], expected_k)
+    assert (name == "none") == (torch.equal(seen[0], q) and torch.equal(seen[1], k))
+
+
+def test_smooth_k_removes_mean_and_preserves_attention() -> None:
+    torch.manual_seed(1)
+    q, k, v = (torch.randn(1, 32, 2, 64) for _ in range(3))
+    k = k + 5.0
+    qt, kt = precondition_qk(q, k, smooth_k=True, hadamard=True)
+    ref = torch.nn.functional.scaled_dot_product_attention(*(t.transpose(1, 2) for t in (q, k, v)))
+    out = torch.nn.functional.scaled_dot_product_attention(*(t.transpose(1, 2) for t in (qt, kt, v)))
+    torch.testing.assert_close(out, ref, atol=1e-4, rtol=1e-4)
+
+
+def _build_h3_attention(bf16_attention: bool):
+    h3 = pytest.importorskip("fastvideo.models.dits.minimax_h3")
+    backend = MagicMock()
+    backend.get_name.return_value = "FLASH_ATTN"
+    distributed = MagicMock()
+    with (
+        patch.object(h3, "get_attn_backend", return_value=backend),
+        patch.object(h3, "get_compute_dtype", return_value=torch.bfloat16),
+        patch.object(h3, "ReplicatedLinear", side_effect=lambda *a, **kw: torch.nn.Identity()),
+        patch.object(h3, "DistributedAttention", distributed),
+    ):
+        h3.MiniMaxH3Attention(64, 2, 32, 1e-6, (), None, prefix="transformer_blocks.3.attn",
+                              bf16_attention=bf16_attention)
+    return distributed.call_args.kwargs
+
+
+def test_h3_attention_bf16_opt_out_wiring() -> None:
+    assert _build_h3_attention(True)["nvfp4_fa4"] is False
+    assert "nvfp4_fa4" not in _build_h3_attention(False)
+
+
+def test_bf16_block_impl_ignores_env_opt_in(env_overrides, fa4_runtime_available) -> None:
+    env_overrides.enter_context(envs.FASTVIDEO_NVFP4_FA4.override(True))
+    assert _impl(nvfp4_fa4=_build_h3_attention(True)["nvfp4_fa4"]).nvfp4_fa4 is False
+    assert _impl(**{k: v for k, v in _build_h3_attention(False).items() if k == "nvfp4_fa4"}).nvfp4_fa4 is True  # env wins

fastvideo/tests/loader/test_h3_community_lora.py

@@ -370,3 +370,46 @@ def test_h3_pipeline_converts_and_rejects_foreign_adapters() -> None:
     with pytest.raises(ValueError, match="no MiniMax-H3 LoRA layers"):
         MiniMaxH3BasePipeline.convert_adapter_state_dict(pipeline, {"unet.down.0.lora_A.weight": torch.zeros(1)},
                                                          "adapter")
+
+
+class _RecordingLoRALayer(torch.nn.Module):
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.disable_lora = False
+        self.applied = False
+
+    def set_lora_weights(self, *args, **kwargs) -> None:
+        self.applied = True
+
+
+def test_layers_without_adapter_weights_log_one_summary_warning(monkeypatch: pytest.MonkeyPatch) -> None:
+    """An adapter that trains a subset of the wrapped linears warns once, not once per layer."""
+    import fastvideo.pipelines.lora_pipeline as lp
+
+    names = ["blocks.0.attn.to_q"] + [f"blocks.{i}.adaln_proj.linear" for i in range(40)]
+    layers = {name: _RecordingLoRALayer() for name in names}
+    adapter = {"blocks.0.attn.to_q.lora_A": torch.zeros(2, 4), "blocks.0.attn.to_q.lora_B": torch.zeros(4, 2)}
+    warnings: list[str] = []
+    monkeypatch.setattr(lp.dist, "get_rank", lambda: 0)
+    monkeypatch.setattr(lp.logger, "warning", lambda message, *args: warnings.append(message % args))
+    pipeline = SimpleNamespace(
+        lora_adapters={"default": adapter},
+        lora_adapter_paths={"default": "pika.safetensors"},
+        lora_layers={"transformer": SimpleNamespace(lora_layers_by_block=lambda: iter([(None, layers)]))},
+        trainable_transformer_modules={},
+        cur_adapter_name="",
+        cur_adapter_path="",
+        cur_adapter_strength=1.0,
+        lora_initialized=True,
+        _setting_constructor_adapter=True,
+        _constructor_dense_lora_path=None,
+        fastvideo_args=SimpleNamespace(training_mode=False),
+    )
+
+    LoRAPipeline.set_lora_adapter(pipeline, "default", "pika.safetensors")
+
+    assert layers["blocks.0.attn.to_q"].applied and not layers["blocks.0.attn.to_q"].disable_lora
+    assert all(layers[name].disable_lora for name in names[1:])
+    assert len(warnings) == 1
+    assert "no weights for 40 wrapped layers" in warnings[0] and "blocks.0.adaln_proj.linear" in warnings[0]

fastvideo/tests/loader/test_pika_h3_flex_adapter.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pika H3 Flex adapter discovery and manifest contract."""
+
+import json
+from pathlib import Path
+
+import pytest
+
+from fastvideo.pipelines.basic.pika_h3_flex import adapter
+
+
+def _manifest(task: str = "fl2va") -> dict:
+    return {
+        "format_version": 1,
+        "model": "pika-flex",
+        "task": task,
+        "release_id": "pika-flex-fl2va-v1",
+        "rank": 128,
+        "alpha": 128,
+        "trunk": "transformer_ref" if task == "ref2va" else "transformer",
+        "sampling": {
+            "knots": [1000, 750, 500, 250, 100, 40, 0],
+            "video_shift": 12.0,
+            "audio_shift": 3.0,
+            "method": "renoise",
+        },
+        "base": {"repo_id": "MiniMaxAI/MiniMax-H3", "revision": "42ed227ee7df40d41602854ae760620d6eb651fe"},
+        "target_modules": ["to_q", "to_k", "to_v", "to_out.0", "ff.net.0.proj", "ff.net.2"],
+        "sha256": "0" * 64,
+    }
+
+
+def _write_release(root: Path, task: str = "fl2va") -> Path:
+    task_dir = root / task
+    task_dir.mkdir(parents=True)
+    (task_dir / "adapter.safetensors").write_bytes(b"")
+    (task_dir / "manifest.json").write_text(json.dumps(_manifest(task)), encoding="utf-8")
+    return task_dir
+
+
+def test_manifest_matches_released_contract() -> None:
+    adapter.validate_adapter_manifest(_manifest(), task="fl2va")
+    adapter.validate_adapter_manifest(_manifest("ref2va"), task="ref2va")
+
+
+@pytest.mark.parametrize(
+    ("field", "value", "message"),
+    [
+        ("rank", 64, "rank"),
+        ("alpha", 64, "alpha"),
+        ("trunk", "transformer_ref", "trunk"),
+        ("task", "ref2va", "task"),
+        ("format_version", 2, "format"),
+    ],
+)
+def test_manifest_rejects_wrong_release(field: str, value: object, message: str) -> None:
+    manifest = _manifest()
+    manifest[field] = value
+    with pytest.raises(adapter.PikaAdapterError, match=message):
+        adapter.validate_adapter_manifest(manifest, task="fl2va")
+
+
+def test_manifest_rejects_a_different_sampler() -> None:
+    manifest = _manifest()
+    manifest["sampling"]["knots"] = [1000, 500, 0]
+    with pytest.raises(adapter.PikaAdapterError, match="sampling knots"):
+        adapter.validate_adapter_manifest(manifest, task="fl2va")
+
+
+def test_load_adapter_finds_the_task_directory(tmp_path: Path) -> None:
+    task_dir = _write_release(tmp_path)
+    loaded = adapter.load_adapter(tmp_path, task="fl2va")
+    assert loaded.path == str(task_dir / "adapter.safetensors")
+    assert loaded.trunk == "transformer"
+    assert loaded.manifest["task"] == "fl2va"
+
+
+def test_load_adapter_finds_a_task_directory_or_adapter_file(tmp_path: Path) -> None:
+    task_dir = _write_release(tmp_path)
+    assert adapter.load_adapter(task_dir, task="fl2va").path == str(task_dir / "adapter.safetensors")
+    assert adapter.load_adapter(task_dir / "adapter.safetensors", task="fl2va").path == str(
+        task_dir / "adapter.safetensors")
+
+
+def test_load_adapter_rejects_a_missing_task(tmp_path: Path) -> None:
+    _write_release(tmp_path, "fl2va")
+    with pytest.raises(adapter.PikaAdapterError, match="ref2va"):
+        adapter.load_adapter(tmp_path, task="ref2va")
+
+
+def test_load_adapter_requires_the_manifest(tmp_path: Path) -> None:
+    task_dir = _write_release(tmp_path)
+    (task_dir / "manifest.json").unlink()
+    with pytest.raises(adapter.PikaAdapterError, match="manifest"):
+        adapter.load_adapter(tmp_path, task="fl2va")
+
+
+def _write_marker(root: Path, merged: dict) -> None:
+    (root / adapter.PIKA_MERGED_MARKER).write_text(json.dumps({"format_version": 1, "merged": merged}),
+                                                   encoding="utf-8")
+
+
+def test_merged_marker_skips_the_released_adapter(tmp_path):
+    assert adapter.merged_adapter(tmp_path, "ref2va") is None
+    _write_marker(tmp_path, {"ref2va": {"trunk": "transformer_ref", "release_id": "pika-flex-ref2va-v1"}})
+    assert adapter.merged_adapter(tmp_path, "ref2va")["release_id"] == "pika-flex-ref2va-v1"
+
+
+def test_merged_marker_must_cover_the_task_and_trunk(tmp_path):
+    _write_marker(tmp_path, {"ref2va": {"trunk": "transformer_ref"}})
+    with pytest.raises(adapter.PikaAdapterError, match="does not carry the fl2va model"):
+        adapter.merged_adapter(tmp_path, "fl2va")
+    _write_marker(tmp_path, {"ref2va": {"trunk": "transformer"}})
+    with pytest.raises(adapter.PikaAdapterError, match="must sit on trunk"):
+        adapter.merged_adapter(tmp_path, "ref2va")

fastvideo/tests/stages/test_lazy_module_load.py

@@ -932,3 +932,35 @@ def test_lora_transformer_bookkeeping_is_per_pipeline(monkeypatch):
     assert first.trainable_transformer_modules == {"transformer": first_transformer}
     assert second.trainable_transformer_modules == {"transformer": second_transformer}
     assert first.trainable_transformer_modules is not second.trainable_transformer_modules
+
+
+def test_transform_reading_through_proxy_does_not_reload():
+    """A materialize transform may read attributes back through its own proxy
+    (the constructor LoRA adapter does); that must not trigger a second load."""
+    loader, calls = _counting_loader()
+    module = LazyModule("transformer", loader)
+    seen = []
+
+    def transform(loaded):
+        seen.append(module.tag)
+        return loaded
+
+    module.set_materialize_transform(transform)
+
+    assert module.tag == "c"
+    assert calls == ["c"]
+    assert seen == ["c"]
+
+
+def test_failed_transform_leaves_module_unmaterialized():
+    loader, calls = _counting_loader()
+    module = LazyModule("transformer", loader)
+
+    def transform(loaded):
+        raise RuntimeError("boom")
+
+    module.set_materialize_transform(transform)
+
+    with pytest.raises(RuntimeError, match="boom"):
+        module.materialize()
+    assert not module.is_materialized

fastvideo/tests/stages/test_pika_h3_flex_geometry.py

@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pika H3 Flex bucket geometry contract."""
+
+import pytest
+
+from fastvideo.pipelines.basic.pika_h3_flex import buckets
+
+
+def test_released_bucket_tables_match_pika() -> None:
+    assert buckets.bucket_table("p480") == (
+        (960, 416),
+        (832, 480),
+        (704, 480),
+        (640, 480),
+        (608, 480),
+        (480, 480),
+        (480, 608),
+        (480, 640),
+        (480, 704),
+        (480, 832),
+        (416, 960),
+    )
+    assert buckets.bucket_table("p768") == (
+        (1536, 672),
+        (1344, 768),
+        (1152, 768),
+        (1024, 768),
+        (960, 768),
+        (768, 768),
+        (768, 960),
+        (768, 1024),
+        (768, 1152),
+        (768, 1344),
+        (672, 1536),
+    )
+
+
+def test_unset_canvas_uses_default_tier_and_aspect() -> None:
+    assert buckets.resolve_canvas(None, None) == (768, 1344, "p768", "16:9")
+
+
+@pytest.mark.parametrize(
+    ("height", "width", "expected"),
+    [
+        (720, 1280, (768, 1344, "p768", "16:9")),
+        (1080, 1920, (768, 1344, "p768", "16:9")),
+        (480, 832, (480, 832, "p480", "16:9")),
+        (1024, 768, (1024, 768, "p768", "3:4")),
+        (960, 416, (960, 416, "p480", "9:21")),
+    ],
+)
+def test_requested_canvas_snaps_to_bucket(height: int, width: int, expected: tuple[int, int, str, str]) -> None:
+    assert buckets.resolve_canvas(height, width) == expected
+
+
+def test_one_sided_canvas_is_rejected() -> None:
+    with pytest.raises(ValueError, match="together"):
+        buckets.resolve_canvas(720, None)
+
+
+def test_tier_reference_short_edge_follows_the_output_canvas() -> None:
+    assert buckets.reference_short_edge_for_canvas(480, 832) == 1024
+    assert buckets.reference_short_edge_for_canvas(768, 1344) == 2048
+
+
+def test_resolve_frames_aligns_to_the_17n_plus_5_grid() -> None:
+    assert buckets.resolve_frames(120) == 124
+    assert buckets.resolve_frames(124) == 124
+    assert buckets.resolve_frames(345) == 345
+    assert buckets.resolve_frames(360) == 362
+    with pytest.raises(ValueError, match="Pika H3 Flex generates"):
+        buckets.resolve_frames(107)
+
+
+def test_unknown_tier_is_rejected() -> None:
+    with pytest.raises(ValueError, match="released 480p and 768p"):
+        buckets.resolve_tier("1080p")

fastvideo/tests/stages/test_pika_h3_flex_lazy_adapter.py

@@ -0,0 +1,98 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pika H3 Flex's automatic adapter must not reload released lazy components.
+
+CPU-only: the components are tiny stand-ins and the adapter load itself is
+stubbed, so the test exercises the real LoRA constructor path, the lazy
+materialize transform, and MiniMax-H3's pinned-copy invalidation.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.lazy_module import LazyModule
+
+
+class _Component(torch.nn.Module):
+
+    def __init__(self, name: str) -> None:
+        super().__init__()
+        self.name = name
+        self.blocks = torch.nn.ModuleList([torch.nn.Linear(2, 2)])
+
+
+def _args(lora_path: str) -> SimpleNamespace:
+    return SimpleNamespace(
+        lora_target_modules=None,
+        lora_path=lora_path,
+        lora_nickname="default",
+        lora_strength=1.0,
+        training_mode=False,
+        lora_training=False,
+        dit_layerwise_offload=False,
+        pipeline_config=SimpleNamespace(dit_config=SimpleNamespace(arch_config=SimpleNamespace(
+            exclude_lora_layers=[]))),
+    )
+
+
+@pytest.fixture
+def lazy_pika(monkeypatch, tmp_path):
+    """A Pika pipeline whose four heavy components are deferred, with load tracking."""
+    from fastvideo.pipelines import lora_pipeline as lora_module
+    from fastvideo.pipelines.basic.pika_h3_flex.pipeline import PikaH3FlexPipeline
+
+    loads: list[str] = []
+    modules = {
+        name: LazyModule(name, lambda name=name: loads.append(name) or _Component(name))
+        for name in ("text_encoder", "transformer", "vae", "audio_vae")
+    }
+    args = _args(str(tmp_path / "fl2va" / "adapter.safetensors"))
+
+    def initialize_base(pipeline, *unused_args, **unused_kwargs):
+        pipeline.fastvideo_args = args
+        pipeline.modules = dict(modules)
+
+    adapter_calls: list[str] = []
+
+    def record_adapter(pipeline, lora_nickname, lora_path=None, strength=1.0, accumulate=False):
+        adapter_calls.append(lora_path)
+
+    monkeypatch.setattr(ComposedPipelineBase, "__init__", initialize_base)
+    monkeypatch.setattr(lora_module, "get_local_torch_device", lambda: torch.device("cpu"))
+    monkeypatch.setattr(lora_module.DenseLoRAPatch, "from_adapter", staticmethod(lambda path: None))
+    monkeypatch.setattr(lora_module.LoRAPipeline, "set_lora_adapter", record_adapter)
+
+    pipeline = PikaH3FlexPipeline("unused", args)
+    return SimpleNamespace(pipeline=pipeline, modules=modules, loads=loads, adapter_calls=adapter_calls)
+
+
+def test_adapter_attach_after_release_keeps_released_components_released(lazy_pika):
+    modules, loads = lazy_pika.modules, lazy_pika.loads
+    # Conditioning loads the text encoder, then the release schedule drops it.
+    modules["text_encoder"].materialize()
+    assert modules["text_encoder"].release() is True
+    assert loads == ["text_encoder"]
+
+    # Denoising loads the transformer; its materialize transform attaches the adapter.
+    modules["transformer"].materialize()
+
+    assert lazy_pika.adapter_calls == [lazy_pika.pipeline.lora_path]
+    assert loads == ["text_encoder", "transformer"], f"adapter attach reloaded {loads[2:]}"
+    assert not modules["text_encoder"].is_materialized
+    assert not modules["vae"].is_materialized
+    assert not modules["audio_vae"].is_materialized
+
+
+def test_adapter_attach_still_drops_pinned_copies_of_resident_components(lazy_pika):
+    modules = lazy_pika.modules
+    vae = modules["vae"].materialize()
+    vae.__dict__["_pinned_host_tensors"] = {"weight": torch.zeros(1)}
+
+    modules["transformer"].materialize()
+
+    assert "_pinned_host_tensors" not in vae.__dict__
+    assert lazy_pika.loads == ["vae", "transformer"]

fastvideo/tests/stages/test_pika_h3_flex_request_constraints.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pika H3 Flex request constraints: the six-step ladder and the configured tier.
+
+CPU-only; no weights are loaded.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from fastvideo.api.compat import legacy_generate_call_to_request, normalize_generation_request
+from fastvideo.api.parser import parse_config
+from fastvideo.api.pika_h3_flex import PikaH3FlexSamplingParam
+from fastvideo.api.schema import GenerationRequest, SamplingConfig
+from fastvideo.configs.pipelines.pika_h3_flex import PikaH3FlexPipelineConfig
+from fastvideo.pipelines.basic.pika_h3_flex.buckets import resolve_canvas
+
+
+def _request(sampling: dict) -> GenerationRequest:
+    return parse_config(GenerationRequest, {"prompt": "a paper boat", "sampling": sampling})
+
+
+def _constrain(tier: str, request: GenerationRequest) -> PikaH3FlexSamplingParam:
+    # The sampling-param defaults (768x1344) are what made an omitted size look explicit.
+    return PikaH3FlexPipelineConfig(pika_resolution_tier=tier).apply_request_constraints(
+        request, PikaH3FlexSamplingParam())
+
+
+@pytest.mark.parametrize(("tier", "canvas"), [("p480", (480, 832)), ("p768", (768, 1344))])
+def test_omitted_size_takes_the_configured_tier(tier, canvas):
+    sampling_param = _constrain(tier, _request({}))
+
+    assert (sampling_param.height, sampling_param.width) == canvas
+    # Input preparation then sees an exact bucket of the configured tier.
+    assert resolve_canvas(sampling_param.height, sampling_param.width)[2] == tier
+
+
+def test_generate_video_call_without_size_takes_the_configured_tier():
+    request = legacy_generate_call_to_request("a paper boat", None, legacy_kwargs={"num_frames": 124})
+
+    sampling_param = _constrain("p480", request)
+
+    assert (sampling_param.height, sampling_param.width) == (480, 832)
+
+
+def test_python_request_counts_its_schema_default_size_as_set():
+    """A GenerationRequest built in Python carries SamplingConfig's 720x1280 as a set size."""
+    request = normalize_generation_request(
+        GenerationRequest(prompt="a paper boat", sampling=SamplingConfig(num_inference_steps=6)))
+
+    sampling_param = PikaH3FlexPipelineConfig(pika_resolution_tier="p480").apply_request_constraints(
+        request, PikaH3FlexSamplingParam(height=720, width=1280))
+
+    # Explicit, so it is left for input preparation to snap (to the p768 16:9 bucket).
+    assert (sampling_param.height, sampling_param.width) == (720, 1280)
+    assert resolve_canvas(720, 1280)[:3] == (768, 1344, "p768")
+
+
+@pytest.mark.parametrize(("tier", "size"), [("p480", (768, 1344)), ("p768", (480, 832)), ("p480", (1024, 768))])
+def test_explicit_size_wins_over_the_configured_tier(tier, size):
+    height, width = size
+    sampling_param = PikaH3FlexPipelineConfig(pika_resolution_tier=tier).apply_request_constraints(
+        _request({
+            "height": height,
+            "width": width
+        }), PikaH3FlexSamplingParam(height=height, width=width))
+
+    assert (sampling_param.height, sampling_param.width) == size
+
+
+@pytest.mark.parametrize("sampling", [{"height": 480}, {"width": 832}])
+def test_one_sided_size_is_rejected(sampling):
+    with pytest.raises(ValueError, match="must be passed together"):
+        _constrain("p768", _request(sampling))
+
+
+def test_step_ladder_is_still_enforced():
+    sampling_param = _constrain("p768", _request({}))
+    assert sampling_param.num_inference_steps == 6
+    with pytest.raises(ValueError, match="runs exactly 6 transformer forwards"):
+        _constrain("p768", _request({"num_inference_steps": 8}))

fastvideo/tests/stages/test_pika_h3_flex_scheduler.py

@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pika H3 Flex's re-noise scheduler step."""
+
+import torch
+
+from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler
+
+
+def _scheduler() -> MiniMaxH3Scheduler:
+    scheduler = MiniMaxH3Scheduler(shift=12.0, stochastic_sampling=True, shift_provided_sigmas=True)
+    scheduler.set_timesteps(sigmas=[1.0, 0.5, 0.0], device="cpu")
+    scheduler.set_noise_generator(torch.Generator().manual_seed(7))
+    return scheduler
+
+
+def test_renoised_step_uses_the_next_sigma_and_fresh_noise() -> None:
+    scheduler = _scheduler()
+    sample = torch.zeros(2, 3)
+    velocity = torch.ones(2, 3)
+    timestep = scheduler.timesteps[0]
+    output = scheduler.step(velocity, timestep, sample).prev_sample
+
+    expected_noise = torch.randn(sample.shape, generator=torch.Generator().manual_seed(7), dtype=torch.float32)
+    torch.testing.assert_close(output, 0.5 * torch.ones_like(sample) + 0.5 * expected_noise)
+
+
+def test_terminal_renoised_step_returns_x0_without_drawing_noise() -> None:
+    scheduler = _scheduler()
+    _ = scheduler.step(torch.ones(2, 3), scheduler.timesteps[0], torch.zeros(2, 3))
+    state_before = scheduler._noise_generator.get_state()
+    output = scheduler.step(torch.ones(2, 3), scheduler.timesteps[1], torch.zeros(2, 3)).prev_sample
+    state_after = scheduler._noise_generator.get_state()
+    torch.testing.assert_close(output, torch.full((2, 3), 0.5))
+    assert torch.equal(state_before, state_after)
+
+
+def test_euler_step_is_unchanged_when_stochastic_sampling_is_off() -> None:
+    scheduler = MiniMaxH3Scheduler(shift=12.0)
+    scheduler.set_timesteps(sigmas=[1.0, 0.5, 0.0], device="cpu")
+    sample = torch.full((2, 3), 2.0)
+    velocity = torch.full((2, 3), 4.0)
+    output = scheduler.step(velocity, scheduler.timesteps[0], sample).prev_sample
+    # x0 = 2 + 1 * 4 = 6; ratio = 0.5, so prev = 0.5 * 2 + 0.5 * 6 = 4.
+    torch.testing.assert_close(output, torch.full((2, 3), 4.0))

fastvideo/tests/train/models/test_minimax_h3_omniref_qad.py

@@ -0,0 +1,614 @@
+# SPDX-License-Identifier: Apache-2.0
+"""OmniRef PDD QAD: student initialization, deployed NVFP4 numerics, rows and metrics.
+
+CPU tests: a disabled QAD student is bit-identical to the untouched model,
+static scales come from the calibration table, the row plan is the latent
+generator's plan, and the eval statistics are consistent. GPU tests (GB200 or
+any FlashInfer FP4 GPU): the QAD linear forward is bit-identical to the
+deployed ``NVFP4QuantizeMethod`` on a converter-packed weight, and the STE
+backward is the dense gradient.
+"""
+from __future__ import annotations
+
+import copy
+import math
+
+import pytest
+import torch
+
+from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3ArchConfig, MiniMaxH3Config
+from fastvideo.forward_context import set_forward_context
+from fastvideo.layers.quantization.nvfp4_qad import (NVFP4QADLinearMethod, NVFP4QADPlan, install_nvfp4_qad,
+                                                     nvfp4_qad_export_scales, set_nvfp4_qad_enabled)
+from fastvideo.models.dits import minimax_h3
+from fastvideo.pipelines.basic.minimax_h3.packing import MINIMAX_H3_TEXT_TAG, build_packed_sequence, build_row_timesteps
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.tests.loader.test_minimax_h3_pdd_loading import TINY_ARCH, VIDEO_PATCH_DIM, cpu_loader  # noqa: F401
+from fastvideo.tests.loader.test_minimax_h3_pdd_loading import single_process_group  # noqa: F401
+from fastvideo.train.methods.knowledge_distillation.pdd_qad_metrics import RungStats, summarize
+from fastvideo.train.models.minimax_h3.omniref_data import RowSpec, RowStream, select_rows
+
+needs_fp4 = pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10,
+                               reason="NVFP4 GEMMs need a Blackwell GPU with FlashInfer")
+
+
+def _tiny_model() -> torch.nn.Module:
+    arch = MiniMaxH3ArchConfig(**{**TINY_ARCH, "patch_size": tuple(TINY_ARCH["patch_size"])})
+    model = minimax_h3.MiniMaxH3Transformer3DModel(MiniMaxH3Config(arch_config=arch), hf_config={})
+    generator = torch.Generator().manual_seed(0)
+    with torch.no_grad():
+        for parameter in model.parameters():
+            parameter.copy_(torch.randn(parameter.shape, generator=generator) * 0.1)
+    return model.eval()
+
+
+def _amax_table(num_blocks: int = 2) -> dict[str, float]:
+    return {f"b{b}.ff.{sub}": 3.0 + b for b in range(num_blocks) for sub in ("fc_in", "fc_out")}
+
+
+def _forward(model: torch.nn.Module) -> tuple[torch.Tensor, torch.Tensor]:
+    layout = build_packed_sequence(torch.full((3, ), MINIMAX_H3_TEXT_TAG, dtype=torch.long), 2, 4, 4, 2, (1, 2, 2))
+    unique, inverse = build_row_timesteps(layout, 0.3, 0.4, 0.999, 1.0)
+    generator = torch.Generator().manual_seed(1)
+    inputs = dict(hidden_states=torch.randn(1, int(layout.video_indices.numel()), VIDEO_PATCH_DIM, generator=generator),
+                  audio_hidden_states=torch.randn(1, int(layout.audio_indices.numel()), TINY_ARCH["audio_in_channels"],
+                                                  generator=generator),
+                  encoder_hidden_states=torch.randn(1, 3, TINY_ARCH["text_dim"], generator=generator),
+                  timestep=unique, timestep_indices=inverse, token_tags=layout.token_tags,
+                  position_ids=layout.position_ids, video_indices=layout.video_indices,
+                  audio_indices=layout.audio_indices, text_indices=layout.text_indices)
+    with torch.no_grad(), set_forward_context(current_timestep=0, attn_metadata=None,
+                                              forward_batch=ForwardBatch(data_type="dummy")):
+        return model(**inputs)
+
+
+def test_disabled_qad_student_is_bit_identical_to_teacher(cpu_loader, single_process_group):  # noqa: F811
+    teacher = _tiny_model()
+    student = copy.deepcopy(teacher)
+    methods = install_nvfp4_qad(student, _amax_table(), NVFP4QADPlan())
+    # Two blocks x (q, k, v, out, fc_in, fc_out); SDPA builds no VSA gates.
+    assert len(methods) == 12
+    assert all(isinstance(m, NVFP4QADLinearMethod) for m in methods.values())
+    set_nvfp4_qad_enabled(methods, False)
+    for expected, got in zip(_forward(teacher), _forward(student), strict=True):
+        assert torch.equal(expected, got)
+    assert all(torch.equal(a, b) for a, b in zip(teacher.state_dict().values(), student.state_dict().values()))
+
+
+def test_static_scales_require_calibration_and_round_trip(cpu_loader, single_process_group):  # noqa: F811
+    model = _tiny_model()
+    table = _amax_table()
+    with pytest.raises(KeyError, match="calibrated input amax"):
+        install_nvfp4_qad(copy.deepcopy(model), {k: v for k, v in table.items() if k != "b1.ff.fc_out"},
+                          NVFP4QADPlan())
+    methods = install_nvfp4_qad(model, table, NVFP4QADPlan())
+    assert nvfp4_qad_export_scales(methods) == table
+    unit = [name for name, m in methods.items() if m.act_scale == "unit"]
+    assert len(unit) == 8 and all(".attn." in name for name in unit)
+    assert float(methods["transformer_blocks.1.ff.fc_out"].input_global_sf(torch.device("cpu"))) == pytest.approx(
+        2688.0 / 4.0)
+    # Re-installing wraps the original method, never a QAD method.
+    again = install_nvfp4_qad(model, table, NVFP4QADPlan(quantize_attention=False))
+    assert len(again) == 4 and not isinstance(again["transformer_blocks.0.ff.fc_in"].original, NVFP4QADLinearMethod)
+
+
+def _reference_plan(groups: dict[tuple[str, str], list[dict]], seed: int, per_group: int) -> list[dict]:
+    """The latent generator's selection: per-group seeded shuffle, then interleave."""
+    import random
+    picked = []
+    for key in sorted(groups):
+        entries = sorted(groups[key], key=lambda entry: entry["id"])
+        random.Random(f"{seed}:{key[0]}:{key[1]}").shuffle(entries)
+        picked.append(entries[:per_group])
+    plan = []
+    for index in range(max(len(entries) for entries in picked)):
+        for entries in picked:
+            if index < len(entries):
+                entry = entries[index]
+                source = entry["id"].split(":", 1)[-1]
+                plan.append({
+                    **entry,
+                    "id": f"omniref-{entry['case']}-{source}-{entry['resolution']}",
+                    "source": source,
+                })
+    return plan
+
+
+def test_qad_student_can_be_deep_copied_and_pickled(cpu_loader, single_process_group):  # noqa: F811
+    """EMA copies and teacher snapshots deep-copy the student after install_nvfp4_qad."""
+    import pickle
+    student = _tiny_model()
+    install_nvfp4_qad(student, _amax_table(), NVFP4QADPlan())
+    clone = copy.deepcopy(student)
+    method = clone.transformer_blocks[0].ff.fc_in.quant_method
+    assert isinstance(method, NVFP4QADLinearMethod) and method.input_amax == 3.0
+    assert method.original is not student.transformer_blocks[0].ff.fc_in.quant_method.original
+    restored = pickle.loads(pickle.dumps(student.transformer_blocks[0].ff.fc_out.quant_method))
+    assert restored.prefix.endswith("transformer_blocks.0.ff.fc_out") and restored.input_amax == 3.0
+    with pytest.raises(AttributeError):
+        object.__new__(NVFP4QADLinearMethod).anything
+
+
+def test_row_plan_is_the_latent_generators_plan():
+    groups = {(case, res): [{"id": f"{case}:src{i}", "case": case, "resolution": res, "parquet": f"/p/{case}{i}"}
+                            for i in range(12)] for case in ("first_frame", "storyboard") for res in ("480p", "768p")}
+    ours = select_rows(groups, 20261007, {"480p": 5, "768p": 5})
+    theirs = _reference_plan(groups, 20261007, 5)
+    assert [(r.id, r.parquet) for r in ours] == [(c["id"], c["parquet"]) for c in theirs]
+    excluded = select_rows(groups, 20261007, {"480p": 12}, frozenset({"src0", "src1"}))
+    assert not any(r.source in {"src0", "src1"} for r in excluded) and all(r.resolution == "480p" for r in excluded)
+
+
+def test_row_stream_feeds_each_data_parallel_group_distinct_rows():
+    plan = [RowSpec(f"r{i}", f"s{i}", "first_frame", "480p", f"/p{i}", i) for i in range(10)]
+    streams = [iter(RowStream(plan, dp_rank=r, dp_size=2, seed=3)) for r in range(2)]
+    firsts = [[next(s)["row"].id for _ in range(5)] for s in streams]
+    assert set(firsts[0]).isdisjoint(firsts[1]) and len(set(firsts[0] + firsts[1])) == 10
+    replay = iter(RowStream(plan, dp_rank=0, dp_size=2, seed=3))
+    assert [next(replay)["row"].id for _ in range(5)] == firsts[0]
+
+
+def test_rung_stats_summaries():
+    stats = RungStats(2)
+    teacher = torch.ones(4, 3)
+    for rung in range(2):
+        pair = {m: {"student": teacher * (1.1 if m == "video" else 1.0), "teacher": teacher,
+                    "student_v": teacher, "teacher_v": teacher} for m in ("video", "audio")}
+        stats.add_rung(rung, pair)
+    summary = summarize(stats.reduce(contribute=True), 2, {"video": 1.0, "audio": 1.0})
+    assert summary["x0_rel_l2/video/rung1"] == pytest.approx(0.1)
+    assert summary["x0_rel_l2/audio/rung0"] == 0.0
+    assert summary["score"] == pytest.approx(0.05)
+
+
+# ----------------------------------------------------------------------------- GPU: deployed numerics
+def _runtime_layer(weight: torch.Tensor, input_amax: float | None):
+    """A ``ReplicatedLinear`` packed the way the H3 NVFP4 converter packs a dense weight."""
+    from fastvideo.layers.linear import ReplicatedLinear
+    from fastvideo.layers.quantization.nvfp4_config import (H3_NVFP4_DIT_INPUT_SF_NAME, NVFP4QuantizeMethod,
+                                                            _nvfp4_quantize, _require_flashinfer)
+    sf_layout, _, _ = _require_flashinfer()
+    layer = ReplicatedLinear(weight.shape[1], weight.shape[0], bias=False).cuda()
+    global_sf = (448 * 6) / weight.float().abs().nan_to_num().max()
+    packed, scale = _nvfp4_quantize(weight, global_sf, sfLayout=sf_layout.layout_128x4, do_shuffle=False)
+    layer._nvfp4_weight, layer._nvfp4_weight_scale = packed, scale
+    layer._nvfp4_alpha = (1.0 / torch.as_tensor(global_sf, device="cuda", dtype=torch.float32)).float()
+    layer._weight_global_sf = torch.as_tensor(global_sf, device="cuda").to(torch.bfloat16)
+    if input_amax is not None:
+        setattr(layer, H3_NVFP4_DIT_INPUT_SF_NAME, torch.tensor(2688.0 / input_amax, device="cuda"))
+    return layer, NVFP4QuantizeMethod("transformer_blocks.0.ff.fc_out")
+
+
+@needs_fp4
+@pytest.mark.parametrize("input_amax", [None, 40.0])
+def test_qad_linear_forward_is_the_deployed_nvfp4_linear(input_amax):
+    from fastvideo.layers.quantization.nvfp4_qad import _NVFP4STE, nvfp4_gemm
+    torch.manual_seed(0)
+    weight = (torch.randn(512, 1024, device="cuda") * 0.05).to(torch.bfloat16)
+    x = (torch.randn(300, 1024, device="cuda") * 8).to(torch.bfloat16)
+    layer, runtime = _runtime_layer(weight, input_amax)
+    expected = runtime.apply(layer, x)
+    sf = torch.tensor(1.0 if input_amax is None else 2688.0 / input_amax, device="cuda")
+    assert torch.equal(nvfp4_gemm(x, weight, None, sf), expected)
+    weight_master = weight.float().requires_grad_(True)
+    x_grad = x.clone().requires_grad_(True)
+    out = _NVFP4STE.apply(x_grad, weight_master, None, sf)
+    assert torch.equal(out, expected)
+    grad = torch.randn_like(out)
+    out.backward(grad)
+    torch.testing.assert_close(weight_master.grad, (grad.float().t() @ x.float()), rtol=2e-2, atol=1e-2)
+    torch.testing.assert_close(x_grad.grad.float(), grad.float() @ weight.float(), rtol=2e-2, atol=1e-2)
+
+
+def test_block_reconstruction_is_local_and_zero_at_identity(cpu_loader, single_process_group):  # noqa: F811
+    from fastvideo.train.methods.knowledge_distillation.pdd_qad_recon import BlockReconstruction
+    teacher = _tiny_model()
+    student = copy.deepcopy(teacher).train().requires_grad_(True)
+    teacher.requires_grad_(False)
+    with BlockReconstruction(teacher, student, scale=1.0) as recon:
+        _forward(teacher)
+    assert len(recon.attn_rel) == 2 and float(recon.loss()) == 0.0
+    assert all(p.grad is None or float(p.grad.abs().max()) == 0.0 for p in student.parameters())
+    # Perturb block 1's FFN only: its FFN unit errs; block 0 and every attention unit stay exact.
+    with torch.no_grad():
+        student.transformer_blocks[1].ff.fc_out.weight.mul_(1.1)
+    student.zero_grad(set_to_none=True)
+    with BlockReconstruction(teacher, student, scale=1.0) as recon:
+        _forward(teacher)
+    assert float(recon.ff_rel[1]) > 0 and float(recon.ff_rel[0]) == 0.0
+    assert all(float(a) == 0.0 for a in recon.attn_rel)
+    grads = {n: p.grad for n, p in student.named_parameters() if p.grad is not None and p.grad.abs().max() > 0}
+    assert grads and all(n.startswith("transformer_blocks.1.ff.") for n in grads), sorted(grads)
+
+
+# ----------------------------------------------------------------------------- E0: recipe-review fixes
+def test_x0_matches_scheduler_sign():
+    """x0 = x_t + (1 - t) v, exactly what MiniMaxH3Scheduler.step denoises to (a final step lands on it)."""
+    from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler
+    from fastvideo.train.methods.knowledge_distillation.pdd_qad_metrics import denoised
+    scheduler = MiniMaxH3Scheduler(shift=12.0)
+    scheduler.set_timesteps(sigmas=torch.tensor([0.7, 0.0]))
+    x_t, v = torch.randn(5, 3), torch.randn(5, 3)
+    final = scheduler.step(v, scheduler.timesteps[0], x_t, return_dict=False)[0]
+    torch.testing.assert_close(denoised(x_t, v, scheduler.timesteps[0]), final)
+    assert not torch.allclose(final, x_t - 0.7 * v)
+
+
+def test_grouped_rel_error_weights_groups_and_drops_absent():
+    from fastvideo.train.methods.knowledge_distillation.pdd_qad_recon import GroupedRelError
+    groups = torch.tensor([0, 0, 1, 2, 2, -1])  # video, video, audio, cond, cond, padding
+    teacher = torch.tensor([[1.0, 0], [1, 0], [0, 2], [0, 100], [100, 0], [7, 7]])[None]
+    student = teacher.clone()
+    student[0, 2] += torch.tensor([0.0, 1.0])  # audio rel = 1/4
+    student[0, 5] += 9.0  # padding row never counts
+    loss, per_group = GroupedRelError(groups, {"video": 1.0, "audio": 1.0, "condition": 0.25})(student, teacher)
+    assert per_group.tolist() == [0.0, 0.25, 0.0]
+    assert float(loss) == pytest.approx(0.25 / 2.25)  # the huge conditioning rows do not dilute audio
+    no_audio = torch.tensor([0, 0, 2, 2, 2, -1])
+    _, per_group = GroupedRelError(no_audio, {"video": 1.0, "audio": 1.0})(student, teacher)
+    assert math.isnan(float(per_group[1]))
+
+
+def test_grouped_rel_error_sp_gradient_matches_unsharded(monkeypatch):
+    """FSDP averages gradients over SP ranks; the mean of the shards' gradients must be the unsharded gradient.
+
+    The x0 output loss already has this scale (its all-gather backward sums over SP ranks), so reconstruction
+    must too, or ``recon_lambda`` is silently divided by ``sp_size``.
+    """
+    import fastvideo.distributed as distributed
+    from fastvideo.train.methods.knowledge_distillation.pdd_qad_recon import GroupedRelError
+    torch.manual_seed(0)
+    weights = {"video": 1.0, "audio": 1.0, "condition": 0.25}
+    groups = torch.tensor([0, 0, 0, 1, 2, 2, 0, 1, 2, 0, 0, -1])  # 2 shards of 6 rows, last row is SP padding
+    x, teacher = torch.randn(12, 5), torch.randn(12, 3)
+    weight = torch.randn(3, 5, requires_grad=True)
+
+    def grad(loss: torch.Tensor) -> torch.Tensor:
+        weight.grad = None
+        loss.backward()
+        return weight.grad.clone()
+
+    monkeypatch.setattr(distributed, "model_parallel_is_initialized", lambda: False)
+    full_loss, full_rel = GroupedRelError(groups, weights)(x @ weight.T, teacher)
+    full_grad = grad(full_loss)
+
+    sp_size, shard = 2, 6
+    local_totals: list[torch.Tensor] = []
+
+    class _Group:  # all_reduce = the sum of every shard's totals, recorded in a first pass
+
+        def all_reduce(self, totals: torch.Tensor) -> torch.Tensor:
+            local_totals.append(totals)
+            return sum(local_totals[:sp_size]) if len(local_totals) > sp_size else totals
+
+    monkeypatch.setattr(distributed, "model_parallel_is_initialized", lambda: True)
+    monkeypatch.setattr(distributed, "get_sp_world_size", lambda: sp_size)
+    monkeypatch.setattr(distributed, "get_sp_group", lambda: _Group())
+    rows = [slice(rank * shard, (rank + 1) * shard) for rank in range(sp_size)]
+    for part in rows:  # first pass: record each shard's local totals
+        GroupedRelError(groups[part], weights)(x[part] @ weight.T, teacher[part])
+    shard_grads = []
+    for part in rows:
+        loss, rel = GroupedRelError(groups[part], weights)(x[part] @ weight.T, teacher[part])
+        torch.testing.assert_close(rel, full_rel)
+        shard_grads.append(grad(loss))
+    torch.testing.assert_close(sum(shard_grads) / sp_size, full_grad)  # FSDP's mean over the SP ranks
+
+
+def test_packed_row_groups_from_layout():
+    from types import SimpleNamespace
+    from fastvideo.train.methods.knowledge_distillation.pdd_qad_recon import local_row_groups, packed_row_groups
+    layout = SimpleNamespace(video_indices=torch.tensor([3, 4, 5, 6]), num_condition_video_rows=1,
+                             audio_indices=torch.tensor([7, 8]), num_condition_audio_rows=1)
+    groups = packed_row_groups(layout, 9)
+    assert groups.tolist() == [2, 2, 2, 2, 0, 0, 0, 2, 1]
+    assert torch.equal(local_row_groups(groups, torch.device("cpu")), groups)
+
+
+def test_grouped_block_reconstruction_measure_only_has_no_grads(cpu_loader, single_process_group):  # noqa: F811
+    from fastvideo.train.methods.knowledge_distillation.pdd_qad_recon import BlockReconstruction, GroupedRelError
+    teacher = _tiny_model()
+    student = copy.deepcopy(teacher).train().requires_grad_(True)
+    teacher.requires_grad_(False)
+    with torch.no_grad():
+        student.transformer_blocks[0].attn.to_v.weight.mul_(1.2)
+    layout = build_packed_sequence(torch.full((3, ), MINIMAX_H3_TEXT_TAG, dtype=torch.long), 2, 4, 4, 2, (1, 2, 2))
+    groups = torch.full((int(layout.position_ids.shape[0]), ), 2, dtype=torch.long)
+    groups[layout.video_indices] = 0
+    groups[layout.audio_indices] = 1
+    criterion = GroupedRelError(groups, {"video": 1.0, "audio": 1.0, "condition": 0.25})
+    with BlockReconstruction(teacher, student, 1.0, criterion, measure_only=True) as recon:
+        _forward(teacher)
+    assert float(recon.loss()) > 0 and all(p.grad is None for p in student.parameters())
+    with BlockReconstruction(teacher, student, 1.0, criterion) as recon:
+        _forward(teacher)
+    grads = {n for n, p in student.named_parameters() if p.grad is not None and p.grad.abs().max() > 0}
+    assert grads and all(n.startswith("transformer_blocks.0.attn.") for n in grads), sorted(grads)
+    assert set(recon.metrics("m")) >= {"m/video_rel", "m/audio_rel", "m/condition_rel"}
+
+
+def _bare_method(**attrs):
+    from fastvideo.train.methods.knowledge_distillation.pdd_qad import OmniRefPDDQADMethod
+    method = object.__new__(OmniRefPDDQADMethod)
+    defaults = dict(eligibility_keys=("endpoint/audio", "keyframe/vs_teacher/first"), tolerances={},
+                    fidelity_tolerance=0.05, min_eligibility_tolerance=0.02, select_metric="score",
+                    best_score=1.0, save_best=True, patience=3, bad_evals=0, bad_score_evals=0,
+                    bad_alignment_evals=0, stop_requested=False, _checkpoint_request=None)
+    method.__dict__.update({**defaults, **attrs})
+    return method
+
+
+def test_qad_data_parallel_groups_follow_the_same_collective_schedule():
+    """Different rows must not select different student forwards or reconstruction backwards under FSDP."""
+    rows = [f"omniref-case{group}-asset{group}-480p" for group in range(4)]
+    methods = [_bare_method(dp_size=4, dp_rank=group, sampling_seed=1234, rungs_per_row=4,
+                            rollout_start_step=20, rollout_fraction=0.25) for group in range(4)]
+    schedules = []
+    for iteration in range(1, 201):
+        choices = [(method._trained_rungs(row, iteration, 8), method._use_rollout(row, iteration))
+                   for method, row in zip(methods, rows, strict=True)]
+        assert all(choice == choices[0] for choice in choices)
+        assert len(choices[0][0]) == 4
+        if iteration < 20:
+            assert not choices[0][1]
+        schedules.append((tuple(sorted(choices[0][0])), choices[0][1]))
+    assert len(set(schedules)) > 8
+    assert {rollout for _, rollout in schedules[19:]} == {False, True}
+
+
+def test_qad_single_data_parallel_group_retains_row_sampling():
+    from fastvideo.train.methods.knowledge_distillation.pdd_qad import _unit_hash
+    method = _bare_method(dp_size=1, sampling_seed=1234, rungs_per_row=4,
+                          rollout_start_step=20, rollout_fraction=0.25)
+    for row in ("storyboard-a", "first_frame-b"):
+        for iteration in (1, 20, 31):
+            order = sorted(range(8), key=lambda rung: _unit_hash(row, iteration, "rung", rung))
+            assert method._trained_rungs(row, iteration, 8) == set(order[:4])
+            assert method._use_rollout(row, iteration) == (iteration >= 20 and _unit_hash(row, iteration) < 0.25)
+
+
+def test_noise_floor_sets_two_sigma_tolerances():
+    method = _bare_method()
+    step0 = {"endpoint/audio": 0.40, "keyframe/vs_teacher/first": 0.30, "score": 0.30, "endpoint/video": 0.42}
+    method._evaluate = lambda rows, rollout, save_tag="": {"endpoint/audio": 0.44, "keyframe/vs_teacher/first": 0.30,
+                                                          "score": 0.30, "endpoint/video": 0.42}
+    out = method._noise_floor([], step0)
+    sigma = 0.04 / (math.sqrt(2) * 0.42)
+    assert out["init/noise_sigma/endpoint/audio"] == pytest.approx(sigma)
+    assert method.tolerances["endpoint/audio"] == pytest.approx(2 * sigma)
+    assert method.tolerances["keyframe/vs_teacher/first"] == 0.02  # floor when two seeds agree
+
+
+def test_best_requires_eligibility_on_full_eval():
+    method = _bare_method(step0={"score": 0.30, "endpoint/audio": 0.40, "keyframe/vs_teacher/first": 0.30},
+                          best_score=0.30, tolerances={"endpoint/audio": 0.10, "keyframe/vs_teacher/first": 0.10})
+    teacher_forced_only = {"score": 0.20}
+    method._guardrails(25, teacher_forced_only)
+    assert method._checkpoint_request is None  # better score, but alignment unknown
+    audio_regressed = {"score": 0.20, "endpoint/video": 0.4, "endpoint/audio": 0.45, "keyframe/vs_teacher/first": 0.3}
+    out = method._guardrails(50, audio_regressed)
+    assert method._checkpoint_request is None and out["eval/eligible"] == 0.0 and method.bad_evals == 1
+    fine = {"score": 0.21, "endpoint/video": 0.4, "endpoint/audio": 0.43, "keyframe/vs_teacher/first": 0.31}
+    out = method._guardrails(75, fine)
+    assert method._checkpoint_request == "best" and method.best_score == 0.21 and method.bad_evals == 0
+
+
+def test_teacher_forced_evals_do_not_reset_alignment_failures():
+    """eval_every 15 / full_eval_every 30 / patience 2: persistent alignment regressions must still stop the run."""
+    method = _bare_method(step0={"score": 0.30, "endpoint/video": 0.40, "endpoint/audio": 0.40},
+                          best_score=0.30, patience=2)
+    full_regressed = {"score": 0.25, "endpoint/video": 0.40, "endpoint/audio": 0.60}
+    method._guardrails(30, full_regressed)
+    assert (method.bad_alignment_evals, method.bad_evals, method.stop_requested) == (1, 1, False)
+    method._guardrails(45, {"score": 0.24})  # teacher-forced only: better score, alignment unmeasured
+    assert (method.bad_score_evals, method.bad_alignment_evals, method.stop_requested) == (0, 1, False)
+    out = method._guardrails(60, full_regressed)
+    assert method.stop_requested and out["eval/bad_alignment_evals"] == 2.0
+    method = _bare_method(step0={"score": 0.30, "endpoint/audio": 0.40}, patience=2)
+    method._guardrails(30, {"score": 0.25, "endpoint/video": 0.4, "endpoint/audio": 0.60})
+    method._guardrails(60, {"score": 0.25, "endpoint/video": 0.4, "endpoint/audio": 0.40})  # recovered
+    assert method.bad_alignment_evals == 0 and not method.stop_requested
+    method._guardrails(75, {"score": 0.31})
+    method._guardrails(90, {"score": 0.32, "endpoint/video": 0.4, "endpoint/audio": 0.40})
+    assert method.bad_score_evals == 2 and method.stop_requested  # the score count still runs on every eval
+
+
+def test_baseline_state_round_trips_through_dcp(tmp_path, single_process_group):  # noqa: F811
+    """Resume restores the step-0 baseline, guard counters and provisional scales instead of re-measuring."""
+    import torch.distributed.checkpoint as dcp
+    from types import SimpleNamespace
+    from fastvideo.train.methods.knowledge_distillation.pdd_qad import _BaselineState
+
+    def student(amax):
+        linears = {"transformer_blocks.0.ff.fc_in": NVFP4QADLinearMethod(object(), "fc_in", "static", amax,
+                                                                         allow_uncalibrated=True)}
+        return SimpleNamespace(provisional_rows=2, qad_linears=linears)
+
+    saved = _bare_method(normalizers={"video": [0.2, 0.1], "audio": [0.4, 0.3]}, step0={"score": 0.3},
+                         recon_norm=[0.05, 0.04], tolerances={"endpoint/audio": 0.1}, best_score=0.27,
+                         bad_score_evals=1, bad_alignment_evals=2, student=student(7.5))
+    dcp.save({"method.qad_baseline": _BaselineState(saved)}, checkpoint_id=str(tmp_path))
+    resumed = _bare_method(normalizers=None, step0={}, recon_norm=None, best_score=math.inf, student=student(None))
+    dcp.load({"method.qad_baseline": _BaselineState(resumed)}, checkpoint_id=str(tmp_path))
+    for key in ("normalizers", "step0", "recon_norm", "tolerances", "best_score", "bad_score_evals",
+                "bad_alignment_evals"):
+        assert getattr(resumed, key) == getattr(saved, key), key
+    assert resumed.bad_evals == 2
+    assert resumed.student.qad_linears["transformer_blocks.0.ff.fc_in"].input_amax == 7.5
+    with pytest.raises(ValueError, match="no step-0 baseline"):
+        _bare_method(student=student(None))._load_baseline_state({**saved._baseline_state(), "normalizers": None})
+
+
+def test_provisional_calibration_runs_on_a_step0_cache_hit(single_process_group):  # noqa: F811
+    from types import SimpleNamespace
+    calls = []
+    cached = {"eval/score": 0.3, "eval/x0_mse/video/rung0": 0.2, "eval/x0_mse/audio/rung0": 0.4}
+    method = _bare_method(normalizers=None, step0={}, eval_every=15, student=SimpleNamespace(num_rungs=1),
+                          _rank_rows=lambda: [], _load_step0=lambda: cached,
+                          _restore_step0=lambda c: {k[len("eval/"):]: v for k, v in c.items()},
+                          _provisional_calibration=lambda: calls.append(1) or {"init/provisional_calibration_rows": 2})
+    metrics = method.on_validation_begin(0)
+    assert calls == [1] and metrics["init/provisional_calibration_rows"] == 2
+    assert method.normalizers == {"video": [0.2], "audio": [0.4]}
+    method.on_validation_begin(7)  # not an eval step (e.g. a resume): calibration still gets its chance
+    assert calls == [1, 1]
+
+
+def test_manifest_groups_keep_each_resolution_of_a_shared_id(tmp_path):
+    import json
+    from fastvideo.train.models.minimax_h3.omniref_data import load_manifest_groups
+
+    def entry(res, frames, case="first_frame", source="a"):
+        return {"id": f"{case}:{source}", "case": case, "resolution": res,
+                "parquet": f"/d/bucket=832x480-{frames}f/{source}-{res}.parquet"}
+
+    first, second = tmp_path / "v4.jsonl", tmp_path / "v5.jsonl"
+    first.write_text("\n".join(json.dumps(e) for e in (entry("480p", 121), entry("768p", 121))) + "\n")
+    second.write_text("\n".join(json.dumps(e) for e in (entry("480p", 121), entry("480p", 241, source="b"))) + "\n")
+    groups = load_manifest_groups([str(first), str(second)], ("first_frame", ), 243)
+    assert {key: [e["id"] for e in entries] for key, entries in groups.items()} == {
+        ("first_frame", "480p"): ["first_frame:a", "first_frame:b"],  # the v5 repeat of (a, 480p) is dropped
+        ("first_frame", "768p"): ["first_frame:a"],
+    }
+
+
+def test_eval_top_up_respects_the_eval_frame_cap():
+    from fastvideo.train.models.minimax_h3.omniref_data import limit_frames, top_up_eval_rows
+    groups = {("storyboard", "480p"): [{"id": f"storyboard:s{i}", "case": "storyboard", "resolution": "480p",
+                                        "parquet": f"/d/bucket=832x480-{121 if i % 2 else 241}f/s{i}.parquet"}
+                                       for i in range(8)]}
+    assert limit_frames(groups, 0) == groups
+    capped = limit_frames(groups, 121)
+    assert len(capped[("storyboard", "480p")]) == 4 and len(groups[("storyboard", "480p")]) == 8
+    rows = top_up_eval_rows([], capped, 7, {"480p": 3})
+    assert len(rows) == 3 and all(int(row.source[1:]) % 2 for row in rows)
+
+
+def test_eval_top_up_is_held_out_and_interleaved():
+    from fastvideo.train.models.minimax_h3.omniref_data import top_up_eval_rows
+    groups = {(case, res): [{"id": f"{case}:{res}{i}", "case": case, "resolution": res, "parquet": f"/p{i}"}
+                            for i in range(10)] for case in ("first_frame", "storyboard") for res in ("480p", "768p")}
+    existing = [RowSpec("e1", "480p0", "first_frame", "480p", "/p", 1)]
+    rows = top_up_eval_rows(existing, groups, 7, {"480p": 3, "768p": 1})
+    count = {}
+    for row in rows:
+        count[(row.case, row.resolution)] = count.get((row.case, row.resolution), 0) + 1
+    assert count == {("first_frame", "480p"): 3, ("first_frame", "768p"): 1, ("storyboard", "480p"): 3,
+                     ("storyboard", "768p"): 1}
+    assert len({row.source for row in rows if row.resolution == "480p" and row.case == "first_frame"}) == 3
+    assert [(r.case, r.resolution) for r in rows[:4]] == sorted({(r.case, r.resolution) for r in rows})
+    plan = select_rows(groups, 20261007, {"480p": 10}, frozenset(row.source for row in rows))
+    assert not {r.source for r in plan} & {r.source for r in rows}
+
+
+def test_pinned_amax_checks_sha_and_copies(tmp_path, single_process_group):  # noqa: F811
+    import hashlib
+    from fastvideo.train.models.minimax_h3.omniref_pdd import _pin_amax
+    table = tmp_path / "amax.json"
+    table.write_text('{"b0.ff.fc_in": 3.0}')
+    digest = hashlib.sha256(table.read_bytes()).hexdigest()
+    assert _pin_amax(str(table), digest, str(tmp_path / "run")) == digest
+    assert (tmp_path / "run" / f"amax-{digest[:12]}.json").is_file()
+    with pytest.raises(ValueError, match="pins"):
+        _pin_amax(str(table), "0" * 64, "")
+
+
+def test_skip_blocks_keep_whole_blocks_bf16(cpu_loader, single_process_group):  # noqa: F811
+    methods = install_nvfp4_qad(_tiny_model(), _amax_table(), NVFP4QADPlan(skip_blocks=(1, )))
+    assert methods and all(name.startswith("transformer_blocks.0.") for name in methods)
+
+
+def test_fp32_boundary_modules_match_inference_and_stay_frozen(cpu_loader, single_process_group):  # noqa: F811
+    """The QAD loader keeps H3's FP32 boundary modules FP32 (as inference) and refuses to train them."""
+    from types import SimpleNamespace
+    from fastvideo.train.models.minimax_h3.omniref_pdd import MiniMaxH3OmniRefPDDModel
+    model = _tiny_model().to(torch.float32)
+    assert not model.config.uniform_parameter_dtype
+    fp32 = {n for n, _ in model.named_parameters() if model._get_parameter_dtype(n, torch.bfloat16) == torch.float32}
+    assert fp32 and {n.split(".", 1)[0] for n in fp32} <= set(model._keep_in_fp32_modules)
+    model.requires_grad_(False)
+    fake = SimpleNamespace(transformer=model)
+    MiniMaxH3OmniRefPDDModel._check_fp32_modules_frozen(fake)
+    model.proj_in.weight.requires_grad_(True)
+    with pytest.raises(ValueError, match="must stay frozen"):
+        MiniMaxH3OmniRefPDDModel._check_fp32_modules_frozen(fake)
+
+
+def test_block_first_call_keeps_the_fsdp_root(single_process_group):  # noqa: F811
+    """Stage 2 from a checkpoint calibrates reconstruction (student blocks called directly) before any full student
+    forward; without pinning the root first, FSDP2 makes block 0 the root and the next full forward fails."""
+    from torch.distributed.device_mesh import init_device_mesh
+    from torch.distributed.fsdp import fully_shard
+
+    from fastvideo.train.methods.knowledge_distillation.pdd_qad_recon import ensure_fsdp_root_initialized
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    class Toy(torch.nn.Module):
+
+        def __init__(self) -> None:
+            super().__init__()
+            self.transformer_blocks = torch.nn.ModuleList(torch.nn.Linear(4, 4) for _ in range(2))
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            for block in self.transformer_blocks:
+                x = block(x)
+            return x
+
+    def sharded() -> Toy:
+        mesh = init_device_mesh(device, (1, ))
+        model = Toy().to(device)
+        for block in model.transformer_blocks:
+            fully_shard(block, mesh=mesh)
+        return fully_shard(model, mesh=mesh)
+
+    x = torch.randn(2, 4, device=device)
+    broken = sharded()
+    broken.transformer_blocks[0](x)
+    with pytest.raises(RuntimeError, match="already been lazily initialized"):
+        broken(x)
+    fixed = sharded()
+    ensure_fsdp_root_initialized(fixed)
+    ensure_fsdp_root_initialized(fixed)  # idempotent
+    fixed.transformer_blocks[0](x)
+    assert fixed(x).shape == (2, 4)
+    ensure_fsdp_root_initialized(torch.nn.Linear(2, 2))  # no FSDP: no-op
+
+
+def test_step0_cache_round_trip_and_fingerprint(tmp_path):
+    from fastvideo.train.methods.knowledge_distillation.pdd_qad_step0 import (load_step0, restore_step0, save_step0,
+                                                                              step0_fingerprint_inputs)
+    method = {"init_student_dcp": "/runs/e3/best", "noise_floor_eval": True, "eligibility_keys": ["endpoint/audio"],
+              "recon_group_weights": {"video": 1.0, "audio": 1.0}, "eval_every": 20}
+    data = {"eval_manifest": "/m.json", "eval_per_group": {"480p": 4, "768p": 2}, "cases": ["first_frame"]}
+    qad = {"linears": {"skip_blocks": [42, 43]}, "fp4_attention": {"two_level_p": True}}
+    models = dict(student_init_from="/ckpt/student", teacher_init_from="/ckpt/teacher")
+    inputs = step0_fingerprint_inputs(method, data, qad, "ab" * 32, 30, **models)
+    metrics = {"eval/score": 0.31, "eval/x0_mse/video/rung0": 0.2, "eval/seconds": 6000.0,
+               "init/recon_step0/rung0": 0.04, "init/recon_step0/rung1": 0.03,
+               "init/noise_sigma/endpoint/audio": 0.05, "init/noise_sigma/score": 0.001}
+    path = tmp_path / "step0.json"
+    assert load_step0(path, inputs) is None
+    save_step0(path, inputs, metrics)
+    # Settings that cannot change step 0 (eval cadence) do not invalidate the cache.
+    same = step0_fingerprint_inputs({**method, "eval_every": 40}, data, qad, "ab" * 32, 30, **models)
+    cached = load_step0(path, same)
+    assert cached is not None and "eval/seconds" not in cached and cached["eval/score"] == 0.31
+    summary, recon_norm, tolerances = restore_step0(cached, 2, recon=True, noise_floor=True, min_tolerance=0.02)
+    assert summary == {"score": 0.31, "x0_mse/video/rung0": 0.2}
+    assert recon_norm == [0.04, 0.03]
+    assert tolerances == {"endpoint/audio": pytest.approx(0.1), "score": 0.02}
+    # Training-plan keys only matter with provisional calibration (no amax_json).
+    assert load_step0(path, step0_fingerprint_inputs(method, {**data, "seed": 5}, qad, "ab" * 32, 30, **models))
+    provisional = {**qad, "linears": {"skip_blocks": [42, 43], "provisional_calibration_rows": 2}}
+    for changed in (step0_fingerprint_inputs({**method, "init_student_dcp": "/runs/e3b/best"}, data, qad, "ab" * 32, 30,
+                                             **models),
+                    step0_fingerprint_inputs(method, {**data, "eval_per_group": 1}, qad, "ab" * 32, 30, **models),
+                    step0_fingerprint_inputs(method, data, {**qad, "linears": {"skip_blocks": []}}, "ab" * 32, 30,
+                                             **models),
+                    step0_fingerprint_inputs(method, data, qad, "cd" * 32, 30, **models),
+                    step0_fingerprint_inputs(method, data, qad, "ab" * 32, 30, **{**models, "student_init_from": "/x"}),
+                    step0_fingerprint_inputs(method, data, qad, "ab" * 32, 30, **{**models, "teacher_init_from": "/x"}),
+                    step0_fingerprint_inputs({**method, "audio_weight": 2.0}, data, qad, "ab" * 32, 30, **models),
+                    step0_fingerprint_inputs(method, data, provisional, "ab" * 32, 30, **models)):
+        with pytest.raises(ValueError, match="other inputs"):
+            load_step0(path, changed)

fastvideo/tests/train/models/test_pika_h3_flex_qad.py

@@ -0,0 +1,623 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU checks for the Pika H3 Flex QAD profile: ladder, re-noise determinism, contract, adapter merge."""
+from __future__ import annotations
+
+import importlib.util
+import json
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+from safetensors.torch import save_file
+
+from fastvideo.train.methods.knowledge_distillation.pdd_qad_metrics import scheduler_step
+from fastvideo.train.models.minimax_h3.flex import (CONTRACT_FILE, CONTRACT_SCHEMA, FlexContract,
+                                                    MiniMaxH3FlexRef2VAModel, RowRenoiseScheduler, flex_sigmas,
+                                                    renoise_seed)
+
+# Pika's published per-modality ladders (knots 1000, 750, 500, 250, 100, 40, 0 shifted by 12 / 3).
+PIKA_VIDEO_SIGMAS = (1.0, 0.97297, 0.92308, 0.8, 0.57143, 0.33333, 0.0)
+PIKA_AUDIO_SIGMAS = (1.0, 0.9, 0.75, 0.5, 0.25, 0.11111, 0.0)
+STEPS = (1000, 750, 500, 250, 100, 40)
+
+
+def _merge_module():
+    path = Path(__file__).resolve().parents[4] / "scripts/distill/pika_h3_flex_qad/merge_pika_adapter.py"
+    spec = importlib.util.spec_from_file_location("merge_pika_adapter", path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+class _CPUFlexModel(MiniMaxH3FlexRef2VAModel):
+    device = torch.device("cpu")
+
+
+def _model_stub() -> MiniMaxH3FlexRef2VAModel:
+    model = object.__new__(_CPUFlexModel)
+    model.contract = FlexContract("ref2va", "transformer_ref", STEPS, 12.0, 3.0)
+    return model
+
+
+def _row(seed: int) -> SimpleNamespace:
+    return SimpleNamespace(spec=SimpleNamespace(seed=seed))
+
+
+def test_flex_sigmas_match_pika_ladders():
+    torch.testing.assert_close(flex_sigmas(STEPS, 12.0), torch.tensor(PIKA_VIDEO_SIGMAS), atol=1e-5, rtol=0)
+    torch.testing.assert_close(flex_sigmas(STEPS, 3.0), torch.tensor(PIKA_AUDIO_SIGMAS), atol=1e-5, rtol=0)
+
+
+def test_schedulers_expose_six_forwards_in_clean_time():
+    video, audio = _model_stub().schedulers(_row(7))
+    assert video.timesteps.numel() == audio.timesteps.numel() == 6
+    torch.testing.assert_close(video.timesteps, 1 - torch.tensor(PIKA_VIDEO_SIGMAS[:-1]), atol=1e-5, rtol=0)
+    assert video.stochastic_sampling and audio.stochastic_sampling
+
+
+def test_schedulers_require_the_row():
+    with pytest.raises(ValueError, match="prepared row"):
+        _model_stub().schedulers()
+
+
+def test_renoise_step_is_x0_plus_fixed_fresh_noise():
+    video, _ = _model_stub().schedulers(_row(11))
+    state, velocity = torch.randn(5, 4), torch.randn(5, 4)
+    rung = 2
+    out = scheduler_step(video, rung, velocity, state)
+    timestep = video.timesteps[rung]
+    x0 = state + (1 - timestep) * velocity
+    generator = torch.Generator().manual_seed(renoise_seed(11, rung, 0))
+    noise = torch.randn(state.shape, generator=generator, dtype=torch.float32)
+    sigma_next = video.sigmas[rung + 1]
+    torch.testing.assert_close(out, (1 - sigma_next) * x0 + sigma_next * noise)
+
+
+def test_renoise_noise_depends_only_on_row_forward_and_modality():
+    """Teacher-forced pass, rollout and re-evaluation of one row must draw identical noise at each forward."""
+    state, velocity = torch.randn(6, 3), torch.randn(6, 3)
+    first = _model_stub().schedulers(_row(3))
+    second = _model_stub().schedulers(_row(3))
+    # Out-of-order use (a rollout after a full teacher-forced pass) still reproduces forward 1.
+    for rung in range(5):
+        scheduler_step(first[0], rung, velocity, state)
+    torch.testing.assert_close(scheduler_step(first[0], 1, velocity, state),
+                               scheduler_step(second[0], 1, velocity, state))
+    assert not torch.equal(scheduler_step(first[0], 1, velocity, state), scheduler_step(first[0], 2, velocity,
+                                                                                          state))
+    other_row = _model_stub().schedulers(_row(4))
+    assert not torch.equal(scheduler_step(first[0], 1, velocity, state),
+                           scheduler_step(other_row[0], 1, velocity, state))
+    assert renoise_seed(3, 1, 0) != renoise_seed(3, 1, 1)
+
+
+def test_last_forward_returns_x0_without_noise():
+    video, audio = _model_stub().schedulers(_row(5))
+    state, velocity = torch.randn(4, 2), torch.randn(4, 2)
+    for scheduler in (video, audio):
+        out = scheduler_step(scheduler, 5, velocity, state)
+        torch.testing.assert_close(out, state + (1 - scheduler.timesteps[5]) * velocity)
+
+
+def test_row_renoise_scheduler_refuses_an_implicit_index():
+    scheduler = RowRenoiseScheduler(12.0, 0, 0)
+    scheduler.set_timesteps(sigmas=flex_sigmas(STEPS, 12.0))
+    with pytest.raises(ValueError, match="explicit step index"):
+        scheduler.step(torch.zeros(2, 2), scheduler.timesteps[0], torch.zeros(2, 2))
+
+
+def _write_contract(path: Path, **overrides) -> None:
+    data = {
+        "schema": CONTRACT_SCHEMA,
+        "task": "ref2va",
+        "trunk": "transformer_ref",
+        "denoising_steps": list(STEPS),
+        "video_scheduler_shift": 12.0,
+        "audio_scheduler_shift": 3.0,
+        "sampler": "renoise",
+    }
+    data.update(overrides)
+    (path / CONTRACT_FILE).write_text(json.dumps(data))
+
+
+def test_contract_reads_the_released_ladder(tmp_path):
+    _write_contract(tmp_path)
+    assert FlexContract.read(str(tmp_path)) == FlexContract("ref2va", "transformer_ref", STEPS, 12.0, 3.0)
+
+
+@pytest.mark.parametrize("overrides", [{
+    "denoising_steps": [1000, 500, 0]
+}, {
+    "video_scheduler_shift": 10.0
+}, {
+    "trunk": "transformer"
+}, {
+    "sampler": "euler"
+}])
+def test_contract_rejects_anything_but_the_release(tmp_path, overrides):
+    _write_contract(tmp_path, **overrides)
+    with pytest.raises(ValueError):
+        FlexContract.read(str(tmp_path))
+
+
+def test_contract_must_match_the_profile_trunk(tmp_path):
+    _write_contract(tmp_path, task="fl2va", trunk="transformer")
+    with pytest.raises(ValueError, match="needs ref2va on transformer_ref"):
+        _model_stub()._read_contract(str(tmp_path))
+
+
+def _reference_module():
+    from fastvideo.attention.backends.fp4_vsa_qat import _load
+    return _load("attn_qat_vsa_reference")
+
+
+class _ExactKernel:
+    """Stand-in for ``attn_qat_vsa_train``: exact softmax over the listed (here: all) blocks and valid keys."""
+
+    def __init__(self, reference):
+        self.reference = reference
+        self.calls = []
+
+    def fp4_vsa_attn_qat(self, q, k, v, q2k_idx, q2k_num, kv_valid, **numerics):
+        self.calls.append((q.shape, numerics))
+        n_blocks = q.shape[2] // 128
+        assert bool((q2k_num == n_blocks).all())
+        keys = self.reference.key_valid_mask(kv_valid, n_blocks, q.device).reshape(-1)
+        scores = (q.float() @ k.float().transpose(-1, -2)) / q.shape[-1]**0.5
+        scores = scores.masked_fill(~keys[None, None, None, :], float("-inf"))
+        return (scores.softmax(-1) @ v.float()).to(torch.bfloat16)
+
+
+def test_dense_block_lists_cover_every_key_tile():
+    from fastvideo.attention.backends.fp4_dense_qat import dense_block_lists
+    reference = _reference_module()
+    q2k_idx, q2k_num, kv_valid = dense_block_lists(300, 1, 2, reference, torch.device("cpu"))
+    assert q2k_num.shape == (1, 2, 3) and bool((q2k_num == 3).all())
+    assert sorted(q2k_idx[0, 0, 0].tolist()) == [0, 1, 2]
+    assert kv_valid.tolist() == [64, 64, 64, 64, 44, 0]
+
+
+def test_fp4_dense_attention_pads_to_tiles_and_matches_dense_attention():
+    from fastvideo.attention.backends.fp4_dense_qat import FP4DenseAttention
+    from fastvideo.attention.backends.fp4_vsa_qat import FP4AttentionNumerics
+    reference = _reference_module()
+    kernel = _ExactKernel(reference)
+    attention = FP4DenseAttention(FP4AttentionNumerics(quantize=True, quantize_pv=False), (reference, kernel))
+    torch.manual_seed(0)
+    q, k, v = (torch.randn(1, 200, 2, 16, dtype=torch.bfloat16) for _ in range(3))
+    out = attention(q, k, v, None)
+    expected = torch.nn.functional.scaled_dot_product_attention(*(t.float().transpose(1, 2) for t in (q, k, v)))
+    assert out.shape == q.shape and out.dtype == q.dtype
+    torch.testing.assert_close(out.float(), expected.transpose(1, 2), atol=2e-2, rtol=2e-2)
+    shape, numerics = kernel.calls[0]
+    assert shape == (1, 2, 256, 16) and numerics["quantize_pv"] is False and numerics["quantize"] is True
+
+
+def test_install_fp4_dense_attention_overrides_and_restores_flash_attention():
+    from fastvideo.attention.backends.fp4_dense_qat import install_fp4_dense_attention
+    from fastvideo.attention.backends.fp4_vsa_qat import FP4AttentionNumerics
+    try:
+        from fastvideo.attention.backends.flash_attn import FlashAttentionImpl
+    except (ImportError, RuntimeError) as error:  # FlashAttention is a CUDA-host dependency
+        pytest.skip(f"FlashAttention backend not importable here: {error}")
+
+    class Layer(torch.nn.Module):
+
+        def __init__(self):
+            super().__init__()
+            self.distributed_attention = SimpleNamespace(attn_impl=object.__new__(FlashAttentionImpl))
+
+    model = torch.nn.Sequential(Layer(), Layer())
+    impls = [layer.distributed_attention.attn_impl for layer in model]
+    originals = [impl.forward for impl in impls]
+    reference = _reference_module()
+    kernel = _ExactKernel(reference)
+    assert install_fp4_dense_attention(model, FP4AttentionNumerics(quantize_pv=False), (reference, kernel)) == 2
+    q = torch.randn(1, 130, 1, 8, dtype=torch.bfloat16)
+    impls[0].forward(q, q, q, None)
+    assert len(kernel.calls) == 1
+    install_fp4_dense_attention(model, None)
+    assert all(impl.forward == original for impl, original in zip(impls, originals, strict=True))
+
+
+def test_install_fp4_dense_attention_covers_sdpa_fallback():
+    """Without flash-attn the FLASH_ATTN selection resolves to SDPA; the emulator must still be installed."""
+    from fastvideo.attention.backends.fp4_dense_qat import install_fp4_dense_attention
+    from fastvideo.attention.backends.fp4_vsa_qat import FP4AttentionNumerics
+    from fastvideo.attention.backends.sdpa import SDPAImpl
+
+    class Layer(torch.nn.Module):
+
+        def __init__(self):
+            super().__init__()
+            self.distributed_attention = SimpleNamespace(attn_impl=SDPAImpl(1, 8, causal=False, softmax_scale=8**-0.5))
+
+    model = torch.nn.Sequential(Layer(), Layer())
+    impls = [layer.distributed_attention.attn_impl for layer in model]
+    originals = [impl.forward for impl in impls]
+    reference = _reference_module()
+    kernel = _ExactKernel(reference)
+    assert install_fp4_dense_attention(model, FP4AttentionNumerics(quantize_pv=False), (reference, kernel)) == 2
+    q = torch.randn(1, 130, 1, 8, dtype=torch.bfloat16)
+    assert impls[1].forward(q, q, q, None).shape == q.shape
+    assert len(kernel.calls) == 1
+    with pytest.raises(NotImplementedError, match="packed"):
+        impls[0].forward(q, q, q, SimpleNamespace(packed_sequence=object()))
+    install_fp4_dense_attention(model, None)
+    assert all(impl.forward == original for impl, original in zip(impls, originals, strict=True))
+
+
+@pytest.mark.parametrize("name", ["smooth_k", "hadamard", "smooth_k+hadamard"])
+def test_precondition_qk_leaves_attention_unchanged(name):
+    from fastvideo.attention.utils.nvfp4_qk_precondition import parse_qk_transform, precondition_qk
+    torch.manual_seed(0)
+    q, k, v = (torch.randn(2, 37, 3, 16, dtype=torch.float32) for _ in range(3))
+    k = k + 4.0  # shared K bias that smoothing removes
+    smooth, hadamard = parse_qk_transform(name)
+    tq, tk = precondition_qk(q, k, smooth_k=smooth, hadamard=hadamard)
+
+    def attend(a, b):
+        return torch.nn.functional.scaled_dot_product_attention(*(t.transpose(1, 2) for t in (a, b, v))).transpose(1, 2)
+
+    torch.testing.assert_close(attend(tq, tk), attend(q, k), atol=1e-5, rtol=1e-5)
+    if smooth:
+        assert float(tk.mean(dim=1).abs().max()) < 1e-4
+
+
+def test_parse_qk_transform_rejects_unknown_names():
+    from fastvideo.attention.utils.nvfp4_qk_precondition import parse_qk_transform
+    assert parse_qk_transform(None) == (False, False)
+    with pytest.raises(ValueError, match="unknown NVFP4 Q/K transform"):
+        parse_qk_transform("rotate")
+
+
+def test_fp4_dense_attention_with_qk_transform_matches_dense_attention():
+    from fastvideo.attention.backends.fp4_dense_qat import FP4DenseAttention
+    from fastvideo.attention.backends.fp4_vsa_qat import FP4AttentionNumerics
+    reference = _reference_module()
+    kernel = _ExactKernel(reference)
+    attention = FP4DenseAttention(FP4AttentionNumerics(quantize=True, quantize_pv=False), (reference, kernel),
+                                  qk_transform=(True, True))
+    torch.manual_seed(0)
+    q, k, v = (torch.randn(1, 200, 2, 16, dtype=torch.bfloat16) for _ in range(3))
+    out = attention(q, k, v, None)
+    expected = torch.nn.functional.scaled_dot_product_attention(*(t.float().transpose(1, 2) for t in (q, k, v)))
+    torch.testing.assert_close(out.float(), expected.transpose(1, 2), atol=3e-2, rtol=3e-2)
+
+
+def test_install_fp4_dense_attention_keeps_bf16_blocks_on_original_forward():
+    from fastvideo.attention.backends.fp4_dense_qat import block_index, install_fp4_dense_attention
+    from fastvideo.attention.backends.fp4_vsa_qat import FP4AttentionNumerics
+    from fastvideo.attention.backends.sdpa import SDPAImpl
+
+    class Block(torch.nn.Module):
+
+        def __init__(self):
+            super().__init__()
+            self.attn = torch.nn.Module()
+            self.attn.distributed_attention = SimpleNamespace(
+                attn_impl=SDPAImpl(1, 8, causal=False, softmax_scale=8**-0.5))
+
+    model = torch.nn.Module()
+    model.transformer_blocks = torch.nn.ModuleList([Block(), Block(), Block()])
+    impls = [block.attn.distributed_attention.attn_impl for block in model.transformer_blocks]
+    originals = [impl.forward for impl in impls]
+    reference = _reference_module()
+    kernel = _ExactKernel(reference)
+    numerics = FP4AttentionNumerics(quantize_pv=False)
+    assert block_index("transformer_blocks.12._checkpoint_wrapped_module.attn") == 12
+    assert block_index("token_refiner.blocks.0.attn") is None
+    assert install_fp4_dense_attention(model, numerics, (reference, kernel), bf16_blocks=(2, )) == 2
+    assert impls[2].forward == originals[2] and impls[0].forward != originals[0]
+    q = torch.randn(1, 130, 1, 8, dtype=torch.bfloat16)
+    impls[2].forward(q, q, q, None)
+    assert kernel.calls == []
+    impls[1].forward(q, q, q, None)
+    assert len(kernel.calls) == 1
+    with pytest.raises(ValueError, match="match no transformer_blocks"):
+        install_fp4_dense_attention(model, numerics, (reference, kernel), bf16_blocks=(7, ))
+
+
+def _write_trunk(trunk: Path, weights: dict[str, torch.Tensor]) -> None:
+    trunk.mkdir(parents=True)
+    names = sorted(weights)
+    shards = {"model-00001-of-00002.safetensors": names[:2], "model-00002-of-00002.safetensors": names[2:]}
+    weight_map = {}
+    for shard, keys in shards.items():
+        save_file({key: weights[key] for key in keys}, str(trunk / shard))
+        weight_map.update({key: shard for key in keys})
+    (trunk / "diffusion_pytorch_model.safetensors.index.json").write_text(json.dumps({"weight_map": weight_map}))
+    (trunk / "config.json").write_text("{}")
+
+
+def test_merge_trunk_adds_scaled_lora_in_fp32(tmp_path):
+    merge = _merge_module()
+    torch.manual_seed(0)
+    weights = {
+        "blocks.0.attn.to_q.weight": torch.randn(6, 4).to(torch.bfloat16),
+        "blocks.0.attn.to_q.bias": torch.randn(6).to(torch.bfloat16),
+        "blocks.0.ff.net.2.weight": torch.randn(4, 8).to(torch.bfloat16),
+        "proj_out.weight": torch.randn(3, 4),
+    }
+    factors = {
+        "blocks.0.attn.to_q": (torch.randn(2, 4), torch.randn(6, 2)),
+        "blocks.0.ff.net.2": (torch.randn(2, 8), torch.randn(4, 2)),
+    }
+    _write_trunk(tmp_path / "base", weights)
+    report = merge.merge_trunk(tmp_path / "base", tmp_path / "out", factors, scale=0.5)
+    assert report["modules"] == 2
+    from safetensors.torch import load_file
+    merged = {}
+    for shard in (tmp_path / "out").glob("*.safetensors"):
+        merged.update(load_file(str(shard)))
+    for module, (lora_a, lora_b) in factors.items():
+        key = f"{module}.weight"
+        expected = (weights[key].float() + 0.5 * lora_b @ lora_a).to(torch.bfloat16)
+        assert merged[key].dtype == torch.bfloat16
+        assert torch.equal(merged[key], expected)
+    assert torch.equal(merged["blocks.0.attn.to_q.bias"], weights["blocks.0.attn.to_q.bias"])
+    assert torch.equal(merged["proj_out.weight"], weights["proj_out.weight"])
+
+
+def test_merge_trunk_refuses_adapter_modules_without_weights(tmp_path):
+    merge = _merge_module()
+    _write_trunk(tmp_path / "base", {f"w{i}.weight": torch.zeros(2, 2) for i in range(3)})
+    with pytest.raises(ValueError, match="no weight"):
+        merge.merge_trunk(tmp_path / "base", tmp_path / "out", {"missing": (torch.zeros(1, 2), torch.zeros(2, 1))},
+                          scale=1.0)
+
+
+def test_adapter_factors_pair_peft_keys(tmp_path):
+    merge = _merge_module()
+    path = tmp_path / "adapter.safetensors"
+    save_file(
+        {
+            "a.to_q.lora_A.default.weight": torch.zeros(2, 4),
+            "a.to_q.lora_B.default.weight": torch.zeros(6, 2),
+        }, str(path))
+    factors = merge.adapter_factors(str(path))
+    assert list(factors) == ["a.to_q"] and factors["a.to_q"][0].shape == (2, 4)
+    save_file({"a.to_q.lora_A.default.weight": torch.zeros(2, 4)}, str(path))
+    with pytest.raises(ValueError, match="both LoRA factors"):
+        merge.adapter_factors(str(path))
+
+
+def _calibrate_module():
+    path = Path(__file__).resolve().parents[4] / "scripts/distill/pika_h3_flex_qad/calibrate_flex_nvfp4.py"
+    spec = importlib.util.spec_from_file_location("calibrate_flex_nvfp4", path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_calibration_keys_follow_the_qad_table():
+    calibrate = _calibrate_module()
+    assert calibrate.amax_key("transformer.transformer_blocks.7.ff.fc_out") == "b7.ff.fc_out"
+    assert calibrate.amax_key("transformer_blocks.3._checkpoint_wrapped_module.ff.fc_in") == "b3.ff.fc_in"
+    assert calibrate.amax_key("transformer_blocks.3.attn.to_q") is None
+
+
+def test_calibration_candidates_and_outliers():
+    calibrate = _calibrate_module()
+    rows = [{"id": f"r{i}", "amax": {"b0.ff.fc_in": [float(i), 0.0, 1.0], "b0.ff.fc_out": [1.0, 2.0, 3.0]}}
+            for i in range(1, 101)]
+    rows[0]["amax"]["b0.ff.fc_out"] = [50.0, 0.0, 0.0]
+    table = calibrate.candidate_table(rows)
+    assert table["b0.ff.fc_in"]["max"] == 100.0
+    assert table["b0.ff.fc_in"]["0.5x_max"] == 50.0
+    assert table["b0.ff.fc_in"]["row_p90"] == pytest.approx(90.1)
+    assert set(calibrate.outlier_row_ids(rows, 2)) == {"r1", "r100"}
+
+
+def test_calibration_choice_all_tokens_or_balanced():
+    calibrate = _calibrate_module()
+    candidates = {"b0.ff.fc_in": {"max": 10.0, "0.5x_max": 5.0}}
+    # "max" wins on all-token error (video dominates) but loses badly on audio; the balanced objective picks the clip.
+    mse = {"b0.ff.fc_in": {"ref": [100.0, 1.0, 10.0], "max": [1.0, 0.5, 0.1], "0.5x_max": [2.0, 0.05, 0.1]}}
+    table, report = calibrate.choose_table(mse, candidates, "balanced")
+    assert table == {"b0.ff.fc_in": 5.0}
+    assert report["b0.ff.fc_in"]["best"] == "0.5x_max"
+    table, _ = calibrate.choose_table(mse, candidates, "all")
+    assert table == {"b0.ff.fc_in": 10.0}
+
+
+def _builder_module():
+    path = Path(__file__).resolve().parents[4] / "scripts/distill/pika_h3_flex_qad/build_flex_ref2va_rows.py"
+    spec = importlib.util.spec_from_file_location("build_flex_ref2va_rows", path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+_PROMPT = "subject_definitions: <Picture 1> is the exact first frame of the target video.\n\nsummary: A dog runs."
+_ASSET_A = "asset-" + "a" * 32
+
+
+def _entry(case: str, clip: str, parquet: str, resolution: str = "480p") -> dict:
+    entry_id = f"{case}:{clip}" if case in ("first_frame", "first_last_frame", "storyboard") else clip
+    return {"case": case, "id": entry_id, "parquet": parquet, "resolution": resolution}
+
+
+def test_builder_adds_subject_definitions_in_order():
+    builder = _builder_module()
+    after = builder._subject(_PROMPT, "<Audio 1> is the soundtrack of the target video.")
+    assert after.startswith("subject_definitions: <Picture 1> is the exact first frame of the target video. <Audio 1>")
+    assert after.endswith("\n\nsummary: A dog runs.")
+    before = builder.apply_prompt_edits(_PROMPT, [("<Picture 1> is the first frame.", True)])
+    assert before.startswith("subject_definitions: <Picture 1> is the first frame. <Picture 1> is the exact")
+    assert builder.v4_id("storyboard", "c") == "storyboard:c" and builder.v4_id("continue_shot", "c") == "c"
+
+
+def test_builder_new_cases_from_stored_rows():
+    import random as stdlib_random
+    from collections import defaultdict
+    builder = _builder_module()
+    counts = {"ff.pq": 1, "sb.pq": 3, "ff2.pq": 1, "sb2.pq": 5, "cs.pq": 1, "cs2.pq": 1}
+    first = _entry("first_frame", f"{_ASSET_A}_f1_f124_x", "ff.pq")
+    board = _entry("storyboard", f"{_ASSET_A}_f1_f124_x", "sb.pq")
+    scene = _entry("continue_scene", f"{_ASSET_A}_omniref_continue_1", "cs.pq")
+    shot = _entry("continue_shot", f"{_ASSET_A}_omniref_continue_2", "cs2.pq")
+    by_asset = defaultdict(list)
+    by_asset[(_ASSET_A, "480p", "image")] = [first, board, _entry("first_frame", f"{_ASSET_A}_f9_x", "ff2.pq"),
+                                            _entry("storyboard", f"{_ASSET_A}_f9_x", "sb2.pq")]
+    by_asset[(_ASSET_A, "480p", "continue")] = [scene, shot]
+    rng = stdlib_random.Random(0)
+    audio = builder.new_case_spec("audio_reference", first, rng, by_asset, counts.get)
+    assert [r["kind"] for r in audio["references"]] == ["stored", "target_audio"]
+    assert audio["id"] == f"audio_reference:{_ASSET_A}_f1_f124_x" and audio["source_case"] == "first_frame"
+    mixed = builder.new_case_spec("mixed_image_video", scene, rng, by_asset, counts.get)
+    assert [r["kind"] for r in mixed["references"]] == ["target_frame", "stored"] and mixed["prompt_edits"][0][1]
+    many = builder.new_case_spec("many_images", board, rng, by_asset, counts.get)
+    assert 6 <= len(many["references"]) <= 9
+    assert {r["parquet"] for r in many["references"][3:]} <= {"ff.pq", "ff2.pq", "sb2.pq"}
+    assert len(many["prompt_edits"]) == len(many["references"]) - 3
+    multi = builder.new_case_spec("multi_video", scene, rng, by_asset, counts.get)
+    assert [r["parquet"] for r in multi["references"]] == ["cs.pq", "cs2.pq"] and multi["references"][1]["audio"] is False
+    lonely = defaultdict(list, {(_ASSET_A, "480p", "continue"): [scene]})
+    assert builder.new_case_spec("multi_video", scene, rng, lonely, counts.get) is None
+
+
+def test_builder_plan_retiers_only_480p_and_keeps_ids():
+    builder = _builder_module()
+    first = _entry("first_frame", f"{_ASSET_A}_f1_x", "a.pq")
+    hi = _entry("first_frame", f"{_ASSET_A}_f1_x", "b.pq", "768p")
+    v4 = {(first["id"], "480p"): first, (hi["id"], "768p"): hi}
+    specs = builder.plan_rows(v4, {(first["id"], "480p"), (hi["id"], "768p")}, {}, 0, set(), lambda _: 1)
+    assert len(specs) == 1 and specs[0]["id"] == first["id"] and specs[0]["case"] == "first_frame"
+    assert specs[0]["references"] == [{"kind": "stored", "parquet": "a.pq", "index": 0, "audio": True}]
+
+
+class _ToyBlock(torch.nn.Module):
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.linear = torch.nn.Linear(dim, dim)
+
+    def forward(self, hidden: torch.Tensor, shift: torch.Tensor) -> torch.Tensor:
+        return hidden + torch.tanh(self.linear(hidden)) + shift
+
+
+class _ToyStack(torch.nn.Module):
+
+    def __init__(self, blocks: int, dim: int = 4) -> None:
+        super().__init__()
+        self.transformer_blocks = torch.nn.ModuleList(_ToyBlock(dim) for _ in range(blocks))
+
+    def forward(self, hidden: torch.Tensor, shift: torch.Tensor) -> torch.Tensor:
+        for block in self.transformer_blocks:
+            hidden = block(hidden, shift)
+        return hidden
+
+
+def test_theseus_schedule_widens_short_to_long():
+    from fastvideo.train.methods.knowledge_distillation.pdd_qad_recon import theseus_schedule, theseus_widths
+    assert theseus_widths(50) == (1, 2, 4, 8, 16, 32, 50)
+    assert theseus_widths(28) == (1, 2, 4, 8, 16, 28)
+    plan = theseus_schedule(40, theseus_widths(50), 2.0)
+    assert len(plan) == 40 and plan == sorted(plan) and set(plan) == {1, 2, 4, 8, 16, 32, 50}
+    assert plan.count(50) > plan.count(32) > plan.count(4)
+    assert theseus_schedule(7, (1, 2, 4, 8, 16, 32, 50), 2.0) == [1, 2, 4, 8, 16, 32, 50]
+    with pytest.raises(ValueError):
+        theseus_schedule(3, (1, 2, 4, 8), 2.0)
+
+
+@pytest.mark.parametrize("width", [1, 3, 5])
+def test_segment_reconstruction_is_zero_at_identity_and_covers_every_block(width):
+    import copy
+    from fastvideo.train.methods.knowledge_distillation.pdd_qad_recon import SegmentReconstruction
+    torch.manual_seed(0)
+    teacher = _ToyStack(5)
+    student = copy.deepcopy(teacher)
+    hidden, shift = torch.randn(3, 4), torch.randn(4)
+    with torch.no_grad(), SegmentReconstruction(teacher, student, width, 1.0) as segments:
+        teacher(hidden, shift)
+    assert len(segments.rel) == -(-5 // width)
+    assert float(segments.loss()) == 0.0
+    assert all(p.grad is None or float(p.grad.abs().max()) == 0.0 for p in student.parameters())
+
+
+def test_segment_reconstruction_trains_from_teacher_segment_inputs():
+    import copy
+    from fastvideo.train.methods.knowledge_distillation.pdd_qad_recon import SegmentReconstruction
+    torch.manual_seed(0)
+    teacher = _ToyStack(4)
+    student = copy.deepcopy(teacher)
+    with torch.no_grad():
+        student.transformer_blocks[0].linear.weight.add_(0.5)  # only block 0 differs
+    hidden, shift = torch.randn(3, 4), torch.randn(4)
+    with torch.no_grad(), SegmentReconstruction(teacher, student, 2, 1.0) as segments:
+        teacher(hidden, shift)
+    # Segment [0, 1] carries block 0's error; segment [2, 3] starts from the teacher's block-2 input, so it is exact.
+    assert float(segments.rel[0]) > 0.0 and float(segments.rel[1]) == 0.0
+    assert float(student.transformer_blocks[0].linear.weight.grad.abs().max()) > 0.0
+    assert float(student.transformer_blocks[1].linear.weight.grad.abs().max()) > 0.0
+    assert student.transformer_blocks[2].linear.weight.grad is None or \
+        float(student.transformer_blocks[2].linear.weight.grad.abs().max()) == 0.0
+
+
+def test_fl2va_prompt_keeps_the_description_and_defines_the_keyframes():
+    builder = _builder_module()
+    caption = ("subject_definitions: <Picture 1> is the exact first frame.\n\nsummary: A dog runs.\n\n"
+               "retention_analysis: <Picture 1> is preserved.\n\noverall_soundscape: barking")
+    assert builder.fl2va_prompt(caption, "t2va") == "summary: A dog runs.\n\noverall_soundscape: barking"
+    last = builder.fl2va_prompt(caption, "l2va")
+    assert last.startswith("subject_definitions: <Picture 1> is the exact last frame of the target video.\n\n")
+    assert "retention_analysis" not in last
+    both = builder.fl2va_prompt(caption, "fl2va")
+    assert "<Picture 2> is the exact last frame" in both.split("\n\n")[0]
+
+
+def test_fl2va_plan_takes_keyframes_from_the_matching_v4_case():
+    builder = _builder_module()
+    v4 = {}
+    for n in range(10):
+        for case in ("first_frame", "first_last_frame"):
+            for res in ("480p", "768p"):
+                entry = _entry(case, f"{_ASSET_A}_f{n}_x", f"{case}-{n}-{res}.pq", res)
+                v4[(entry["id"], res)] = entry
+    specs = builder.plan_fl2va_rows(v4, {"t2va": 5, "i2va": 5, "l2va": 5, "fl2va": 5}, {"480p": 0.8, "768p": 0.2},
+                                    0, set())
+    by_case = {}
+    for row in specs:
+        by_case.setdefault(row["case"], []).append(row)
+    assert {case: len(rows) for case, rows in by_case.items()} == {"t2va": 5, "i2va": 5, "l2va": 5, "fl2va": 5}
+    assert all(not r["keyframes"] for r in by_case["t2va"])
+    assert all([k["index"] for k in r["keyframes"]] == [1] and "first_last_frame" in r["base_parquet"]
+               for r in by_case["l2va"])
+    assert all([k["index"] for k in r["keyframes"]] == [0, 1] for r in by_case["fl2va"])
+    assert all(r["id"].startswith(f"{r['case']}:") for r in specs)
+    assert sum(r["resolution"] == "768p" for r in by_case["i2va"]) == 1
+
+
+def test_fl2va_row_packs_keyframe_anchors():
+    from types import SimpleNamespace
+    import numpy as np
+    from fastvideo.train.models.minimax_h3.flex import prepare_fl2va_row
+    lat_h, lat_w, frames = 30, 52, 124
+    rows_per_frame = (lat_h // 2) * (lat_w // 2)
+
+    def row(n_keyframes):
+        def put(name, array):
+            return {f"{name}_bytes": array.astype(np.float32).tobytes(), f"{name}_shape": list(array.shape)}
+        out = {"num_frames": frames, "height": lat_h * 16, "width": lat_w * 16,
+               "reference_media_types": ["image"] * n_keyframes, "reference_has_audio": [False] * n_keyframes,
+               "reference_num_latent_frames": [1] * n_keyframes, "reference_latent_heights": [lat_h] * n_keyframes,
+               "reference_latent_widths": [lat_w] * n_keyframes, "reference_num_audio_latents": [0] * n_keyframes}
+        out.update(put("text_token_tags", np.ones(7)))
+        out.update(put("text_embedding", np.zeros((7, 8))))
+        out.update(put("reference_video_latent_rows", np.random.rand(n_keyframes * rows_per_frame, 96)))
+        out.update(put("reference_audio_latent_rows", np.zeros((0, 32))))
+        return out
+
+    spec = SimpleNamespace(id="x", case="fl2va", seed=3)
+    prepared = prepare_fl2va_row(spec, row(2), patch_size=(1, 2, 2), latent_channels=24, audio_channels=32,
+                                 device=torch.device("cpu"))
+    assert prepared.layout.num_condition_video_rows == 2 * rows_per_frame
+    assert prepared.video.shape[0] == prepared.layout.video_indices.numel()
+    t2va = prepare_fl2va_row(SimpleNamespace(id="y", case="t2va", seed=3), row(0), patch_size=(1, 2, 2),
+                             latent_channels=24, audio_channels=32, device=torch.device("cpu"))
+    assert t2va.layout.num_condition_video_rows == 0
+    with pytest.raises(ValueError, match="needs 1"):
+        prepare_fl2va_row(SimpleNamespace(id="z", case="i2va", seed=3), row(2), patch_size=(1, 2, 2),
+                          latent_channels=24, audio_channels=32, device=torch.device("cpu"))

fastvideo/tests/train/utils/test_checkpoint.py

@@ -234,6 +234,15 @@ def test_write_metadata_roundtrip_with_step(tmp_path: Path) -> None:
     assert loaded == {"step": 7}
 
 
+def test_write_metadata_marks_tagged_weights_only(tmp_path: Path) -> None:
+    mgr = _make_manager(tmp_path)
+    tagged = tmp_path / "best"
+    tagged.mkdir()
+    mgr._write_metadata(tagged, step=30, weights_only=True)
+    loaded = CheckpointManager.load_metadata(tagged)
+    assert loaded == {"step": 30, "weights_only": True}
+
+
 def test_write_metadata_includes_raw_config(tmp_path: Path) -> None:
     raw = {
         "models": {

fastvideo/tests/train/utils/test_dcp_to_diffusers_h3.py

@@ -7,6 +7,7 @@ from types import SimpleNamespace
 
 from diffusers import ConfigMixin, ModelMixin
 from diffusers.configuration_utils import register_to_config
+import pytest
 from safetensors.torch import load_file
 import torch
 from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import checkpoint_wrapper
@@ -16,6 +17,7 @@ from fastvideo.models.loader.weight_utils import safetensors_weights_iterator
 from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import MiniMaxH3Ref2VAModularPipeline
 from fastvideo.train.entrypoint.dcp_to_diffusers import (
     _native_export_state_dict,
+    _resolve_export_checkpoint,
     _save_diffusers_safetensors,
     _save_role_pretrained,
     _strict_reload_verify,
@@ -23,6 +25,20 @@ from fastvideo.train.entrypoint.dcp_to_diffusers import (
 from fastvideo.train.utils.lora import enable_lora_training, finalize_lora_training
 
 
+def test_resolve_export_checkpoint_accepts_tagged_weights_only_dir(tmp_path: Path) -> None:
+    tagged = tmp_path / "best"
+    (tagged / "dcp").mkdir(parents=True)
+    (tagged / "dcp" / ".metadata").write_text("")
+    assert _resolve_export_checkpoint(str(tagged)) == tagged.resolve()
+
+
+def test_resolve_export_checkpoint_rejects_plain_dir(tmp_path: Path) -> None:
+    plain = tmp_path / "best"
+    plain.mkdir()
+    with pytest.raises(ValueError, match="best"):
+        _resolve_export_checkpoint(str(plain))
+
+
 class _ArchConfig:
     exclude_lora_layers: list[str] = []
 

fastvideo/train/entrypoint/dcp_to_diffusers.py

@@ -40,6 +40,7 @@ import gc
 import json
 import os
 import sys
+from pathlib import Path
 from typing import Any
 
 import fastvideo.envs as envs
@@ -449,6 +450,15 @@ def _strict_reload_verify(
     logger.info("Strict reload verification passed.")
 
 
+def _resolve_export_checkpoint(checkpoint_dir: str) -> Path:
+    """A DCP image directory, including a tagged weights-only one, or a regular checkpoint path."""
+    raw_path = Path(os.path.expanduser(checkpoint_dir)).resolve()
+    if (raw_path / "dcp" / ".metadata").is_file():
+        return raw_path
+    from fastvideo.train.utils.checkpoint import _resolve_resume_checkpoint
+    return _resolve_resume_checkpoint(checkpoint_dir, output_dir=checkpoint_dir)
+
+
 def convert(
     *,
     checkpoint_dir: str,
@@ -468,10 +478,7 @@ def convert(
     from fastvideo.distributed import (
         maybe_init_distributed_environment_and_model_parallel, )
     from fastvideo.train.utils.builder import build_from_config
-    from fastvideo.train.utils.checkpoint import (
-        CheckpointManager,
-        _resolve_resume_checkpoint,
-    )
+    from fastvideo.train.utils.checkpoint import CheckpointManager
     from fastvideo.train.utils.config import (
         RunConfig,
         load_run_config,
@@ -480,20 +487,18 @@ def convert(
     import torch.distributed.checkpoint as dcp
 
     # -- Resolve checkpoint directory --
-    resolved = _resolve_resume_checkpoint(
-        checkpoint_dir,
-        output_dir=checkpoint_dir,
-    )
+    resolved = _resolve_export_checkpoint(checkpoint_dir)
     dcp_dir = resolved / "dcp"
     if not dcp_dir.is_dir():
         raise FileNotFoundError(f"Missing dcp/ under {resolved}")
+    metadata = CheckpointManager.load_metadata(resolved) if (resolved / "metadata.json").is_file() else {}
+    weights_only = bool(metadata.get("weights_only", False))
 
     # -- Obtain config --
     cfg: RunConfig
     if config_path is not None:
         cfg = load_run_config(config_path)
     else:
-        metadata = CheckpointManager.load_metadata(resolved)
         raw_config = metadata.get("config")
         if raw_config is None:
             raise ValueError("Checkpoint metadata.json does not "
@@ -526,6 +531,8 @@ def convert(
 
     # -- Load DCP weights into the model --
     states = method.checkpoint_state()
+    if weights_only:
+        states = {key: value for key, value in states.items() if key.startswith("roles.")}
     logger.info(
         "Loading DCP checkpoint from %s",
         resolved,

fastvideo/train/entrypoint/train.py

@@ -20,10 +20,12 @@ from __future__ import annotations
 import argparse
 import os
 import sys
+from datetime import timedelta
 from typing import Any
 
 import torch
 
+import fastvideo.envs as envs
 from fastvideo.logger import init_logger
 
 logger = init_logger(__name__)
@@ -66,9 +68,11 @@ def run_training_from_config(
             "VIDEO_SPARSE_ATTN",
         )
 
+    timeout_min = envs.FASTVIDEO_PG_TIMEOUT_MIN.get()
     maybe_init_distributed_environment_and_model_parallel(
         tc.distributed.tp_size,
         tc.distributed.sp_size,
+        timeout=timedelta(minutes=timeout_min) if timeout_min > 0 else None,
     )
 
     _, method, dataloader, start_step = build_from_config(cfg)

fastvideo/train/methods/knowledge_distillation/pdd_qad.py

@@ -0,0 +1,791 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Per-forward quantization-aware distillation (QAD) of a FastH3 OmniRef PDD student.
+
+Student = the teacher's own weights + NVFP4 fake-quantized block linears
+(Stage A calibrated static FFN scales, unit attention/gate scales) + the
+deployed sparse-FP4 attention numerics. Loss = per-forward regression of the
+denoiser output (x0) to the frozen bf16 teacher, at identical inputs on the
+eight-forward PDD ladder; no DMD, no adversarial term.
+
+Each optimizer micro-step takes one OmniRef row: the teacher walks its PDD
+trajectory (teacher forcing), or with probability ``rollout_fraction`` the
+student walks its own (DAgger-style exposure-bias correction); at every
+forward both models see the same state and the student's x0 error is
+backpropagated at once, inside that forward's scope.
+
+Guardrails: step-0 checks (quantization disabled -> bit-identical to the
+teacher; quantized error vs the Stage A PTQ number), per-(rung, modality)
+loss normalization by the step-0 PTQ error, held-out evaluation every
+``eval_every`` steps (teacher-forced x0 rel-L2 per forward, student-rollout
+endpoint, latent keyframe fidelity), best-checkpoint requests, and a stop
+request after ``patience`` consecutive evaluations worse than step 0 (teacher
+alignment counts full evaluations only). The step-0 baseline and the guard
+state are checkpointed (``method.qad_baseline``), so a resumed run keeps them.
+"""
+from __future__ import annotations
+
+import hashlib
+import json
+import math
+import time
+from pathlib import Path
+from typing import Any
+
+import torch
+
+from fastvideo.distributed import get_world_group
+from fastvideo.logger import init_logger
+from fastvideo.train.methods.base import LogScalar, TrainingMethod
+from fastvideo.train.methods.knowledge_distillation.pdd_qad_recon import BlockReconstruction
+from fastvideo.train.methods.knowledge_distillation.pdd_qad_step0 import (load_step0, restore_step0, save_step0,
+                                                                          step0_fingerprint_inputs)
+from fastvideo.train.methods.knowledge_distillation.pdd_qad_metrics import (MODALITIES, RungStats, keyframe_latents,
+                                                                            rel_l2, step_packed, summarize,
+                                                                            target_slices, x0_pair)
+from fastvideo.train.utils.optimizer import build_optimizer_and_scheduler
+
+logger = init_logger(__name__)
+
+# Teacher alignment: the student's own sample (same seed) and its keyframes vs the teacher's; full evaluations only.
+ALIGNMENT_KEYS = ("endpoint/video", "endpoint/audio", "keyframe/vs_teacher/first", "keyframe/vs_teacher/last")
+BASELINE_FORMAT = 1
+
+
+class _BaselineState:
+    """DCP-stateful step-0 baseline and guard state of an ``OmniRefPDDQADMethod`` (one JSON string)."""
+
+    def __init__(self, method: OmniRefPDDQADMethod) -> None:
+        self.method = method
+
+    def state_dict(self) -> dict[str, Any]:
+        return {"json": json.dumps(self.method._baseline_state(), sort_keys=True)}
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        self.method._load_baseline_state(json.loads(state_dict["json"]))
+
+
+def _unit_hash(*parts: Any) -> float:
+    digest = hashlib.sha256(":".join(str(p) for p in parts).encode()).digest()
+    return int.from_bytes(digest[:8], "little") / 2.0**64
+
+
+class OmniRefPDDQADMethod(TrainingMethod):
+    """Teacher-forced (+ optional student-rollout) per-forward x0 regression; see module docstring."""
+
+    def __init__(self, *, cfg: Any, role_models: dict[str, Any]) -> None:
+        super().__init__(cfg=cfg, role_models=role_models)
+        if "teacher" not in role_models:
+            raise ValueError("OmniRefPDDQADMethod needs a frozen 'teacher' role")
+        self.teacher = role_models["teacher"]
+        if getattr(self.teacher, "_trainable", False) or self.teacher.qad_linears:
+            raise ValueError("the teacher must be frozen and unquantized")
+        mc = self.method_config
+        self.weights = {"video": float(mc.get("video_weight", 1.0)), "audio": float(mc.get("audio_weight", 1.0))}
+        self.rollout_fraction = float(mc.get("rollout_fraction", 0.0))
+        self.rollout_start_step = int(mc.get("rollout_start_step", 0))
+        # Student forward+backward on this many of the 8 forwards per row (the teacher walks all of them).
+        self.rungs_per_row = int(mc.get("rungs_per_row", 0))
+        # "output": per-forward x0 regression. "recon": layer-wise reconstruction only.
+        # "output+recon": x0 regression + lambda * reconstruction.
+        self.objective = str(mc.get("objective", "output"))
+        if self.objective not in ("output", "recon", "output+recon"):
+            raise ValueError(f"objective must be output, recon or output+recon, got {self.objective!r}")
+        recon_lambda = float(mc.get("recon_lambda", 1.0))
+        self.output_weight = {"output": 1.0, "recon": 0.0, "output+recon": 1.0 / (1.0 + recon_lambda)}[self.objective]
+        self.recon_weight = {
+            "output": 0.0,
+            "recon": 1.0,
+            "output+recon": recon_lambda / (1.0 + recon_lambda)
+        }[self.objective]
+        if self.objective == "recon" and float(mc.get("rollout_fraction", 0.0)) > 0:
+            raise ValueError("reconstruction trains on teacher states only; set rollout_fraction: 0")
+        # Best checkpoint by this eval key (lower is better): "score" (teacher-forced x0 rel-L2) or
+        # "endpoint/video" (the student's own sample vs the teacher's: teacher alignment).
+        self.select_metric = str(mc.get("select_metric", "score"))
+        # Save each eval's final latents (student's own sample, teacher's) for side-by-side decoding.
+        self.eval_save_dir = mc.get("eval_save_dir")
+        # Full evaluation (with the student's own 8-step sample) every this many steps; teacher-forced only otherwise.
+        self.full_eval_every = int(mc.get("full_eval_every", 0)) or int(mc.get("eval_every", 50))
+        # Start the student from another run's tagged weights (e.g. Stage 1 best) instead of the teacher's.
+        self.init_student_dcp = mc.get("init_student_dcp")
+        self._eval_iteration = 0
+        self.eval_every = int(mc.get("eval_every", 50))
+        self.patience = int(mc.get("patience", 3))
+        self.fidelity_tolerance = float(mc.get("fidelity_tolerance", 0.05))
+        self.ptq_reference_json = mc.get("ptq_reference_json")
+        self.ptq_tolerance = float(mc.get("ptq_tolerance", 0.05))
+        self.require_bitwise_init = bool(mc.get("require_bitwise_init", True))
+        self.max_eval_rows = int(mc.get("max_eval_rows", 0))
+        self.save_best = bool(mc.get("save_best", True))
+        # Reconstruction per row group (target video, target audio, conditioning): over the whole packed sequence
+        # audio is ~1.4% of rows and text tokens dominate the norm, so an ungrouped ratio barely trains audio.
+        self.recon_group_weights = dict(
+            mc.get("recon_group_weights") or {
+                "video": 1.0,
+                "audio": 1.0,
+                "condition": 0.25
+            })
+        # "step0": divide each forward's reconstruction loss by its step-0 value (as the x0 terms are), so
+        # recon_lambda means what it says; raw reconstruction gradients are ~1e3-1e4x smaller than x0's.
+        self.recon_normalize = str(mc.get("recon_normalize", "none"))
+        if self.recon_normalize not in ("none", "step0"):
+            raise ValueError(f"recon_normalize must be none or step0, got {self.recon_normalize!r}")
+        self.recon_norm_rows = int(mc.get("recon_norm_rows", 4))
+        self.recon_norm: list[float] | None = None
+        # Optional Theseus warm-up (JustQuant, arXiv:2609.33601): the first ``theseus_steps`` steps train only
+        # teacher-forced segment reconstruction, widening from single blocks to the whole stack (level k gets
+        # ~theseus_gamma**k of the steps); the configured objective takes over afterwards.
+        self.theseus_steps = int(mc.get("theseus_steps", 0))
+        self.theseus_gamma = float(mc.get("theseus_gamma", 2.0))
+        self._theseus_plan: list[int] = []
+        if self.theseus_steps > 0:
+            from fastvideo.train.methods.knowledge_distillation.pdd_qad_recon import (theseus_schedule, theseus_widths)
+            widths = tuple(
+                int(w)
+                for w in (mc.get("theseus_widths") or theseus_widths(len(self.student.transformer.transformer_blocks))))
+            self._theseus_plan = theseus_schedule(self.theseus_steps, widths, self.theseus_gamma)
+        # Full evaluations roll out (8 more student forwards) on at most this many rows; others are teacher-forced.
+        self.max_rollout_rows = int(mc.get("max_rollout_rows", 0))
+        # A checkpoint is "best" only if these full-eval metrics stay <= step0 * (1 + tolerance). With
+        # noise_floor_eval the tolerance is 2 sigma of a seed+1 re-evaluation at step 0 (>= min_eligibility_tolerance);
+        # otherwise fidelity_tolerance.
+        self.eligibility_keys = tuple(
+            mc.get("eligibility_keys")
+            or ("endpoint/audio", "keyframe/vs_teacher/first", "keyframe/vs_teacher/last", "x0_rel_l2/audio/mean"))
+        self.noise_floor_eval = bool(mc.get("noise_floor_eval", False))
+        self.min_eligibility_tolerance = float(mc.get("min_eligibility_tolerance", 0.02))
+        self.tolerances: dict[str, float] = {}
+        # Stage A's per-rung reference metric for the step-0 PTQ match (v_rel is sign-independent).
+        self.ptq_reference_metric = str(mc.get("ptq_reference_metric", "v_rel_l2"))
+        self._recon_criteria: dict[str, Any] = {}
+        # Reuse (or write) this run's step-0 evaluation, keyed by a fingerprint of everything it depends on.
+        self.step0_cache = mc.get("step0_cache")
+        self._method_config_snapshot = mc
+        self.student.init_preprocessors(cfg.training)
+        tc = self.training_config
+        params = [p for p in self.student.transformer.parameters() if p.requires_grad]
+        self.optimizer, self.lr_scheduler = build_optimizer_and_scheduler(params=params,
+                                                                          optimizer_config=tc.optimizer,
+                                                                          loop_config=tc.loop,
+                                                                          learning_rate=float(
+                                                                              tc.optimizer.learning_rate),
+                                                                          betas=tuple(tc.optimizer.betas),
+                                                                          scheduler_name=str(tc.optimizer.lr_scheduler))
+        self.normalizers: dict[str, list[float]] | None = None
+        self.step0: dict[str, float] = {}
+        self.best_score = math.inf
+        # Consecutive evaluations with a worse teacher-forced score (every eval) and with worse teacher alignment
+        # (full evals only); ``bad_evals`` is the larger.
+        self.bad_score_evals = 0
+        self.bad_alignment_evals = 0
+        self.bad_evals = 0
+        self.stop_requested = False
+        self._checkpoint_request: str | None = None
+        if self.init_student_dcp:
+            import torch.distributed.checkpoint as dcp
+
+            # The same DCP wrapper the checkpoint manager saves tagged/regular student weights with.
+            states = {
+                key: value
+                for key, value in self.checkpoint_state().items() if key == "roles.student.transformer"
+            }
+            if not states:
+                raise ValueError("init_student_dcp needs a trainable 'student' role with a transformer")
+            dcp.load(states, checkpoint_id=str(Path(self.init_student_dcp) / "dcp"))
+            logger.info("QAD student initialized from %s", self.init_student_dcp)
+        world = get_world_group()
+        self.sp_size = int(tc.distributed.sp_size or 1)
+        self.dp_rank, self.dp_size = world.rank // self.sp_size, world.world_size // self.sp_size
+        self.is_sp_leader = world.rank % self.sp_size == 0
+        self.sampling_seed = int(tc.data.seed or 0)
+
+    # ------------------------------------------------------------------ optimizer plumbing
+    def get_optimizers(self, iteration: int) -> list[torch.optim.Optimizer]:
+        return [self.optimizer]
+
+    def get_lr_schedulers(self, iteration: int) -> list[Any]:
+        return [self.lr_scheduler]
+
+    @property
+    def _optimizer_dict(self) -> dict[str, Any]:
+        return {"student": self.optimizer}
+
+    @property
+    def _lr_scheduler_dict(self) -> dict[str, Any]:
+        return {"student": self.lr_scheduler}
+
+    def backward(self,
+                 loss_map: dict[str, torch.Tensor],
+                 outputs: dict[str, Any],
+                 *,
+                 grad_accum_rounds: int = 1) -> None:
+        """No-op: each forward's loss was backpropagated inside its own rung scope."""
+
+    def consume_checkpoint_request(self) -> str | None:
+        request, self._checkpoint_request = self._checkpoint_request, None
+        return request
+
+    # ------------------------------------------------------------------ resume
+    def checkpoint_state(self) -> dict[str, Any]:
+        """Roles, optimizer and scheduler, plus the step-0 baseline and guard state.
+
+        Without the baseline a resumed run would re-measure step 0 on its trained weights (wrong loss scale,
+        stop history and best-checkpoint reference); a checkpoint without ``method.qad_baseline`` cannot be
+        resumed (DCP reports the missing key).
+        """
+        return {**super().checkpoint_state(), "method.qad_baseline": _BaselineState(self)}
+
+    def _baseline_state(self) -> dict[str, Any]:
+        provisional = {}
+        if getattr(self.student, "provisional_rows", 0) > 0:
+            provisional = {
+                name: m.input_amax
+                for name, m in self.student.qad_linears.items() if m.act_scale == "static" and m.input_amax is not None
+            }
+        return {
+            "format": BASELINE_FORMAT,
+            "normalizers": self.normalizers,
+            "step0": self.step0,
+            "recon_norm": self.recon_norm,
+            "tolerances": self.tolerances,
+            "best_score": self.best_score,
+            "bad_score_evals": self.bad_score_evals,
+            "bad_alignment_evals": self.bad_alignment_evals,
+            "provisional_amax": provisional,
+        }
+
+    def _load_baseline_state(self, state: dict[str, Any]) -> None:
+        if state.get("format") != BASELINE_FORMAT:
+            raise ValueError(f"unsupported QAD baseline state format {state.get('format')!r}")
+        if state["normalizers"] is None:
+            raise ValueError("the checkpoint has no step-0 baseline (saved before the step-0 evaluation)")
+        self.normalizers = {m: [float(v) for v in values] for m, values in state["normalizers"].items()}
+        self.step0 = {k: float(v) for k, v in state["step0"].items()}
+        self.recon_norm = None if state["recon_norm"] is None else [float(v) for v in state["recon_norm"]]
+        self.tolerances = {k: float(v) for k, v in state["tolerances"].items()}
+        self.best_score = float(state["best_score"])
+        self.bad_score_evals = int(state["bad_score_evals"])
+        self.bad_alignment_evals = int(state["bad_alignment_evals"])
+        self.bad_evals = max(self.bad_score_evals, self.bad_alignment_evals)
+        for name, amax in state.get("provisional_amax", {}).items():
+            method = self.student.qad_linears.get(name)
+            if method is None:
+                raise KeyError(f"checkpoint has a provisional scale for {name}, which is not a QAD linear here")
+            if method.input_amax is None:
+                method.set_input_amax(float(amax))
+        if get_world_group().rank == 0:
+            logger.info("QAD step-0 baseline restored from the checkpoint (score %.5f, bad evals %d/%d)",
+                        self.step0.get("score", float("nan")), self.bad_score_evals, self.bad_alignment_evals)
+
+    # ------------------------------------------------------------------ one training row
+    def _row(self, spec: Any) -> Any:
+        return self.student.load_row(spec)
+
+    def single_train_step(self, batch: dict[str, Any], iteration: int) -> tuple[dict, dict, dict[str, LogScalar]]:
+        if self.normalizers is None:
+            raise RuntimeError("step-0 evaluation must run before training (it sets the loss normalizers)")
+        spec = batch["row"]
+        prepared = self._row(spec)
+        if iteration <= len(self._theseus_plan):
+            return self._theseus_step(spec, prepared, iteration)
+        rollout = self._use_rollout(spec.id, iteration)
+        schedulers = self.student.schedulers(prepared)
+        timesteps = (schedulers[0].timesteps, schedulers[1].timesteps)
+        accum = max(1, int(self.training_config.loop.gradient_accumulation_steps or 1))
+        num_rungs = self.student.num_rungs
+        trained = self._trained_rungs(spec.id, iteration, num_rungs)
+        video, audio = prepared.video, prepared.audio
+        total = torch.zeros((), device=self.student.device)
+        metrics: dict[str, LogScalar] = {
+            "train/rollout": float(rollout),
+            "train/tokens": float(prepared.layout.sequence_length)
+        }
+        started = time.perf_counter()
+        for rung in range(num_rungs):
+            recon = None
+            with torch.no_grad(), self.teacher.rung_scope(prepared, rung, timesteps) as rt:
+                if rung in trained and self.recon_weight > 0:
+                    # Student blocks train inside the teacher forward, on the teacher's sub-layer inputs.
+                    norm = self.recon_norm[rung] if self.recon_norm else 1.0
+                    scale = self.recon_weight / (norm * 2 * len(self.student.transformer.transformer_blocks) *
+                                                 len(trained) * accum)
+                    with BlockReconstruction(self.teacher.transformer, self.student.transformer, scale,
+                                             self._recon_criterion(prepared)) as recon:
+                        teacher_out = self.teacher.forward_rung(prepared, video, audio, rt)
+                else:
+                    teacher_out = self.teacher.forward_rung(prepared, video, audio, rt)
+            if recon is not None:
+                recon_loss = recon.loss()
+                if not torch.isfinite(recon_loss):
+                    raise FloatingPointError(f"non-finite reconstruction loss at step {iteration}, rung {rung}")
+                norm = self.recon_norm[rung] if self.recon_norm else 1.0
+                total += self.recon_weight * recon_loss / (norm * len(trained))
+                metrics[f"train/recon_vs_step0/rung{rung}"] = recon_loss / norm
+                for key, value in recon.metrics("train/recon").items():
+                    metrics[f"{key}/rung{rung}"] = value
+            if rung not in trained or self.output_weight == 0:
+                if rollout:  # the student's own trajectory still needs its (no-grad) forward here
+                    with torch.no_grad(), self.student.rung_scope(prepared, rung, timesteps) as rt:
+                        student_out = self.student.forward_rung(prepared, video, audio, rt)
+                video, audio = step_packed(schedulers, prepared, rung, video, audio,
+                                           *(student_out if rollout else teacher_out))
+                continue
+            with self.student.rung_scope(prepared, rung, timesteps) as rt:
+                student_out = self.student.forward_rung(prepared, video, audio, rt)
+                pair = x0_pair(prepared, schedulers, rung, video, audio, {
+                    "student": student_out,
+                    "teacher": teacher_out
+                })
+                loss = self._rung_loss(pair, rung, metrics)
+                if not torch.isfinite(loss):
+                    raise FloatingPointError(f"non-finite QAD loss at step {iteration}, row {spec.id}, rung {rung}")
+                (self.output_weight * loss / (len(trained) * accum)).backward()
+            total += self.output_weight * loss.detach() / len(trained)
+            step_out = tuple(v.detach() for v in student_out) if rollout else teacher_out
+            video, audio = step_packed(schedulers, prepared, rung, video, audio, *step_out)
+        metrics["train/row_seconds"] = time.perf_counter() - started
+        if get_world_group().rank == 0:
+            logger.info("QAD step %d row %s (%s, %d tokens, rollout=%s): loss %.5f in %.1fs", iteration, spec.id,
+                        spec.case, prepared.layout.sequence_length, rollout, float(total), metrics["train/row_seconds"])
+        return {"total_loss": total}, {}, metrics
+
+    def _theseus_step(self, spec: Any, prepared: Any, iteration: int) -> tuple[dict, dict, dict[str, LogScalar]]:
+        """One warm-up step: teacher-forced segment reconstruction only, on this row's trained rungs."""
+        from fastvideo.train.methods.knowledge_distillation.pdd_qad_recon import SegmentReconstruction
+
+        width = self._theseus_plan[iteration - 1]
+        schedulers = self.student.schedulers(prepared)
+        timesteps = (schedulers[0].timesteps, schedulers[1].timesteps)
+        accum = max(1, int(self.training_config.loop.gradient_accumulation_steps or 1))
+        num_rungs = self.student.num_rungs
+        trained = self._trained_rungs(spec.id, iteration, num_rungs)
+        num_segments = -(-len(self.student.transformer.transformer_blocks) // width)
+        video, audio = prepared.video, prepared.audio
+        total = torch.zeros((), device=self.student.device)
+        metrics: dict[str, LogScalar] = {
+            "train/theseus_width": float(width),
+            "train/tokens": float(prepared.layout.sequence_length)
+        }
+        started = time.perf_counter()
+        for rung in range(num_rungs):
+            with torch.no_grad(), self.teacher.rung_scope(prepared, rung, timesteps) as rt:
+                if rung in trained:
+                    scale = 1.0 / (num_segments * len(trained) * accum)
+                    with SegmentReconstruction(self.teacher.transformer, self.student.transformer, width, scale,
+                                               self._recon_criterion(prepared)) as segments:
+                        teacher_out = self.teacher.forward_rung(prepared, video, audio, rt)
+                    rel = segments.loss()
+                    if not torch.isfinite(rel):
+                        raise FloatingPointError(f"non-finite Theseus loss at step {iteration}, rung {rung}")
+                    total += rel / len(trained)
+                    for key, value in segments.metrics("train/theseus").items():
+                        metrics[f"{key}/rung{rung}"] = value
+                else:
+                    teacher_out = self.teacher.forward_rung(prepared, video, audio, rt)
+            video, audio = step_packed(schedulers, prepared, rung, video, audio, *teacher_out)
+        metrics["train/row_seconds"] = time.perf_counter() - started
+        if get_world_group().rank == 0:
+            logger.info("Theseus warm-up step %d (width %d) row %s (%s, %d tokens): rel %.5f in %.1fs", iteration,
+                        width, spec.id, spec.case, prepared.layout.sequence_length, float(total),
+                        metrics["train/row_seconds"])
+        return {"total_loss": total}, {}, metrics
+
+    def _recon_criterion(self, prepared: Any) -> Any:
+        from fastvideo.train.methods.knowledge_distillation.pdd_qad_recon import (GroupedRelError, local_row_groups,
+                                                                                  packed_row_groups)
+        key = prepared.spec.id
+        if key not in self._recon_criteria:
+            self._recon_criteria = {}  # one row at a time
+            groups = packed_row_groups(prepared.layout, int(prepared.layout.sequence_length))
+            self._recon_criteria[key] = GroupedRelError(local_row_groups(groups, self.student.device),
+                                                        self.recon_group_weights)
+        return self._recon_criteria[key]
+
+    @torch.no_grad()
+    def _calibrate_recon(self, rows: list[tuple[Any, bool]]) -> dict[str, LogScalar]:
+        """Step-0 reconstruction error per forward (measure-only), the Stage-2 normalizer."""
+        num_rungs = self.student.num_rungs
+        sums = torch.zeros(num_rungs, dtype=torch.float64, device=self.student.device)
+        count = torch.zeros((), dtype=torch.float64, device=self.student.device)
+        per_rank = -(-self.recon_norm_rows // self.dp_size)
+        for spec, counted in rows[:per_rank]:
+            prepared = self._row(spec)
+            schedulers = self.student.schedulers(prepared)
+            timesteps = (schedulers[0].timesteps, schedulers[1].timesteps)
+            video, audio = prepared.video, prepared.audio
+            for rung in range(num_rungs):
+                with self.teacher.rung_scope(prepared, rung,
+                                             timesteps) as rt, BlockReconstruction(self.teacher.transformer,
+                                                                                   self.student.transformer,
+                                                                                   0.0,
+                                                                                   self._recon_criterion(prepared),
+                                                                                   measure_only=True) as recon:
+                    out = self.teacher.forward_rung(prepared, video, audio, rt)
+                if counted:
+                    sums[rung] += recon.loss().double()
+                video, audio = step_packed(schedulers, prepared, rung, video, audio, *out)
+            count += float(counted)
+        if not self.is_sp_leader:
+            sums.zero_()
+            count.zero_()
+        if torch.distributed.is_initialized():
+            torch.distributed.all_reduce(sums)
+            torch.distributed.all_reduce(count)
+        self.recon_norm = [max(float(v), 1e-12) for v in (sums / count.clamp_min(1.0))]
+        return {f"init/recon_step0/rung{rung}": value for rung, value in enumerate(self.recon_norm)}
+
+    def _noise_floor(self, rows: list[tuple[Any, bool]], step0: dict[str, float]) -> dict[str, LogScalar]:
+        """Re-evaluate step 0 at seed+1; per-metric tolerance = max(2 sigma, minimum) with sigma relative."""
+        import dataclasses
+        shifted = [(dataclasses.replace(spec, seed=spec.seed + 1), counted) for spec, counted in rows]
+        other = self._evaluate(shifted, rollout=True, save_tag="seed1")
+        out: dict[str, LogScalar] = {}
+        for key in set(self.eligibility_keys) | {"endpoint/video", "score"}:
+            if key in step0 and key in other:
+                mean = 0.5 * (step0[key] + other[key])
+                sigma = abs(step0[key] - other[key]) / (math.sqrt(2.0) * max(mean, 1e-12))
+                self.tolerances[key] = max(2.0 * sigma, self.min_eligibility_tolerance)
+                out[f"init/noise_sigma/{key}"] = sigma
+                out[f"init/seed1/{key}"] = other[key]
+        return out
+
+    def _tolerance(self, key: str) -> float:
+        return self.tolerances.get(key, self.fidelity_tolerance)
+
+    def _sampling_key(self, row_id: str) -> str:
+        """Keep forward/backward choices in lockstep across the FSDP shard group.
+
+        FSDP spans data-parallel groups, whose rows differ. A row-dependent choice can make one group gather
+        teacher block weights while another gathers student block weights, corrupting both forwards. Single-group
+        runs retain their existing row-dependent sampling.
+        """
+        return row_id if self.dp_size == 1 else f"fsdp-shared:{self.sampling_seed}"
+
+    def _use_rollout(self, row_id: str, iteration: int) -> bool:
+        return (iteration >= self.rollout_start_step
+                and _unit_hash(self._sampling_key(row_id), iteration) < self.rollout_fraction)
+
+    def _trained_rungs(self, row_id: str, iteration: int, num_rungs: int) -> set[int]:
+        """All forwards, or a deterministic subset shared by every FSDP rank."""
+        if self.rungs_per_row <= 0 or self.rungs_per_row >= num_rungs:
+            return set(range(num_rungs))
+        order = sorted(range(num_rungs),
+                       key=lambda rung: _unit_hash(self._sampling_key(row_id), iteration, "rung", rung))
+        return set(order[:self.rungs_per_row])
+
+    def _rung_loss(self, pair: dict[str, dict[str, torch.Tensor]], rung: int, metrics: dict[str,
+                                                                                            LogScalar]) -> torch.Tensor:
+        assert self.normalizers is not None
+        loss = torch.zeros((), device=self.student.device)
+        for modality in MODALITIES:
+            student, teacher = pair[modality]["student"], pair[modality]["teacher"]
+            mse = (student - teacher).pow(2).mean()
+            normalized = mse / self.normalizers[modality][rung]
+            loss = loss + self.weights[modality] * normalized
+            metrics[f"train/rel_mse_vs_ptq/{modality}/rung{rung}"] = normalized.detach()
+        return loss / sum(self.weights.values())
+
+    # ------------------------------------------------------------------ evaluation and guardrails
+    def on_validation_begin(self, iteration: int = 0) -> dict[str, LogScalar]:
+        # Smoke runs without Stage A's table: calibrate the FFN scales before anything uses them, whether step 0
+        # is measured, restored from the cache, or skipped by a resume. A no-op once every scale is set.
+        calibration = self._provisional_calibration()
+        first = self.normalizers is None
+        if not first and (self.eval_every <= 0 or iteration % self.eval_every):
+            return dict(calibration)
+        rows = self._rank_rows()
+        metrics: dict[str, LogScalar] = dict(calibration)
+        started = time.perf_counter()
+        self._eval_iteration = iteration
+        cached = self._load_step0() if first else None
+        if cached is not None:
+            summary = self._restore_step0(cached)
+            metrics.update(cached)
+        else:
+            if first and iteration == 0 and not self.init_student_dcp:
+                # Fresh start only: a resumed run (e.g. Stage 2 from a Stage-1 checkpoint) is no longer the teacher.
+                metrics.update(self._initial_checks(rows))
+            if first and self._needs_recon_norm():
+                metrics.update(self._calibrate_recon(rows))
+            full = first or (iteration % self.full_eval_every == 0)
+            summary = self._evaluate(rows, rollout=full)
+            metrics.update({f"eval/{k}": v for k, v in summary.items()})
+            if first and self.noise_floor_eval:
+                metrics.update(self._noise_floor(rows, summary))
+            if first and self.step0_cache and get_world_group().rank == 0:
+                save_step0(self.step0_cache, self._step0_inputs(), {k: float(v) for k, v in metrics.items()})
+        if first:
+            self.normalizers = {
+                m: [max(summary[f"x0_mse/{m}/rung{r}"], 1e-20) for r in range(self.student.num_rungs)]
+                for m in MODALITIES
+            }
+            self.step0 = dict(summary)
+            self.best_score = summary.get(self.select_metric, summary["score"])
+        else:
+            metrics.update(self._guardrails(iteration, summary))
+        metrics["eval/seconds"] = time.perf_counter() - started
+        metrics["eval/score_vs_ptq"] = summary["score"] / max(self.step0.get("score", summary["score"]), 1e-30)
+        if get_world_group().rank == 0:
+            logger.info("QAD eval @%d: score %.5f (PTQ %.5f) %s", iteration, summary["score"],
+                        self.step0.get("score", float("nan")),
+                        json.dumps({
+                            k: round(float(v), 5)
+                            for k, v in metrics.items()
+                        }))
+        return metrics
+
+    def _needs_recon_norm(self) -> bool:
+        return self.recon_normalize == "step0" and self.recon_weight > 0
+
+    def _step0_inputs(self) -> dict[str, Any]:
+        return step0_fingerprint_inputs(self._method_config_snapshot,
+                                        self.student.data_config,
+                                        self.student.qad_config,
+                                        self.student.amax_sha256,
+                                        len(self.student.eval_rows),
+                                        student_init_from=self.student._init_from,
+                                        teacher_init_from=self.teacher._init_from)
+
+    def _load_step0(self) -> dict[str, float] | None:
+        if not self.step0_cache:
+            return None
+        cached = load_step0(self.step0_cache, self._step0_inputs())
+        if cached is not None and get_world_group().rank == 0:
+            logger.info("QAD step 0 restored from %s (no step-0 evaluation)", self.step0_cache)
+        return cached
+
+    def _restore_step0(self, cached: dict[str, float]) -> dict[str, float]:
+        summary, recon_norm, tolerances = restore_step0(cached, self.student.num_rungs, self._needs_recon_norm(),
+                                                        self.noise_floor_eval, self.min_eligibility_tolerance)
+        if recon_norm is not None:
+            self.recon_norm = recon_norm
+        self.tolerances.update(tolerances)
+        return summary
+
+    def _rank_rows(self) -> list[tuple[Any, bool]]:
+        """This data-parallel group's eval rows, padded so every group runs the same number of forwards.
+
+        FSDP all-gathers span every rank, so groups must stay in lockstep; padded
+        repeats run but are not counted (``(row, counted)`` pairs).
+        """
+        rows = self.student.eval_rows[:self.max_eval_rows] if self.max_eval_rows else self.student.eval_rows
+        if not rows:
+            return []
+        per_rank = -(-len(rows) // self.dp_size)
+        return [(rows[(i * self.dp_size + self.dp_rank) % len(rows)], i * self.dp_size + self.dp_rank < len(rows))
+                for i in range(per_rank)]
+
+    def _eligible(self, summary: dict[str, float]) -> bool:
+        """Full-eval audio/keyframe metrics within step0 * (1 + tolerance); a teacher-forced-only eval cannot tell."""
+        if "endpoint/video" not in summary:
+            return False
+        return all(summary[key] <= self.step0[key] * (1.0 + self._tolerance(key)) for key in self.eligibility_keys
+                   if key in self.step0 and key in summary)
+
+    def _guardrails(self, iteration: int, summary: dict[str, float]) -> dict[str, LogScalar]:
+        worse_score = summary["score"] > self.step0["score"]
+        self.bad_score_evals = self.bad_score_evals + 1 if worse_score else 0
+        # Teacher alignment is measured by full evaluations only; a teacher-forced-only one leaves its count as is.
+        measured = [key for key in ALIGNMENT_KEYS if key in summary and key in self.step0]
+        if measured:
+            misaligned = any(summary[key] > self.step0[key] * (1.0 + self._tolerance(key)) for key in measured)
+            self.bad_alignment_evals = self.bad_alignment_evals + 1 if misaligned else 0
+        self.bad_evals = max(self.bad_score_evals, self.bad_alignment_evals)
+        selected = summary.get(self.select_metric)
+        eligible = self._eligible(summary)
+        if selected is not None and selected < self.best_score and eligible:
+            self.best_score = selected
+            self._checkpoint_request = "best" if self.save_best else None
+        if self.bad_evals >= self.patience:
+            self.stop_requested = True
+            logger.warning("QAD stop: %d consecutive evaluations worse than the PTQ start (step %d)", self.bad_evals,
+                           iteration)
+        return {
+            "eval/eligible": float(eligible),
+            "eval/bad_evals": float(self.bad_evals),
+            "eval/bad_score_evals": float(self.bad_score_evals),
+            "eval/bad_alignment_evals": float(self.bad_alignment_evals),
+            "eval/best_score": self.best_score,
+            "eval/stop_requested": float(self.stop_requested)
+        }
+
+    @torch.no_grad()
+    def _evaluate(self,
+                  rows: list[tuple[Any, bool]],
+                  *,
+                  rollout: bool,
+                  attention: bool = True,
+                  linears: bool = True,
+                  save_tag: str = "") -> dict[str, float]:
+        num_rungs = self.student.num_rungs
+        stats = RungStats(num_rungs)
+        counted_stats = stats
+        # Same rollout count on every data-parallel group (lockstep), covering the first max_rollout_rows rows.
+        rollout_rows = -(-self.max_rollout_rows // self.dp_size) if self.max_rollout_rows else len(rows)
+        with self.student.quantization(linears=linears, attention=attention):
+            for position, (spec, counted) in enumerate(rows):
+                stats = counted_stats if counted else RungStats(num_rungs)
+                prepared = self._row(spec)
+                schedulers = self.student.schedulers(prepared)
+                timesteps = (schedulers[0].timesteps, schedulers[1].timesteps)
+                video, audio = prepared.video, prepared.audio
+                for rung in range(num_rungs):
+                    with self.teacher.rung_scope(prepared, rung, timesteps) as rt:
+                        teacher_out = self.teacher.forward_rung(prepared, video, audio, rt)
+                    with self.student.rung_scope(prepared, rung, timesteps) as rt:
+                        student_out = self.student.forward_rung(prepared, video, audio, rt)
+                    stats.add_rung(
+                        rung,
+                        x0_pair(prepared, schedulers, rung, video, audio, {
+                            "student": student_out,
+                            "teacher": teacher_out
+                        }))
+                    video, audio = step_packed(schedulers, prepared, rung, video, audio, *teacher_out)
+                if rollout and position < rollout_rows:
+                    self._rollout_metrics(stats,
+                                          prepared,
+                                          schedulers,
+                                          timesteps,
+                                          video,
+                                          audio,
+                                          save=counted,
+                                          save_tag=save_tag)
+        means = counted_stats.reduce(contribute=self.is_sp_leader)
+        summary = summarize(means, num_rungs, self.weights)
+        for name in ("first", "last"):
+            student, teacher = summary.get(f"keyframe/student/{name}"), summary.get(f"keyframe/teacher/{name}")
+            if student is not None and teacher is not None:
+                summary[f"keyframe/drop/{name}"] = student / max(teacher, 1e-12) - 1.0
+        return summary
+
+    def _rollout_metrics(self,
+                         stats: RungStats,
+                         prepared: Any,
+                         schedulers: tuple[Any, Any],
+                         timesteps: Any,
+                         teacher_video: torch.Tensor,
+                         teacher_audio: torch.Tensor,
+                         save: bool = False,
+                         save_tag: str = "") -> None:
+        """Student's own 8-forward sample vs the teacher's: endpoint rel-L2 and latent keyframe fidelity."""
+        video, audio = prepared.video, prepared.audio
+        for rung in range(self.student.num_rungs):
+            with self.student.rung_scope(prepared, rung, timesteps) as rt:
+                out = self.student.forward_rung(prepared, video, audio, rt)
+            video, audio = step_packed(schedulers, prepared, rung, video, audio, *out)
+        cut = target_slices(prepared)
+        if save and self.eval_save_dir and self.is_sp_leader:
+            step_dir = f"step{self._eval_iteration:05d}" + (f"_{save_tag}" if save_tag else "")
+            path = Path(self.eval_save_dir) / step_dir / f"{prepared.spec.id}.pt"
+            path.parent.mkdir(parents=True, exist_ok=True)
+            torch.save(
+                {
+                    "student_video": video[cut["video"]].float().cpu(),
+                    "teacher_video": teacher_video[cut["video"]].float().cpu(),
+                    # Audio for decoded stats (spectral centroid, >8 kHz share, dynamic range, log-mel distance).
+                    "student_audio": audio[cut["audio"]].float().cpu(),
+                    "teacher_audio": teacher_audio[cut["audio"]].float().cpu(),
+                    "latent_shape": prepared.latent_shape,
+                    "case": prepared.spec.case,
+                    "seed": prepared.spec.seed
+                },
+                path)
+        for modality, student, teacher in (("video", video, teacher_video), ("audio", audio, teacher_audio)):
+            s, t = student[cut[modality]].float(), teacher[cut[modality]].float()
+            stats.add_scalar(f"endpoint.{modality}", (s - t).norm() / t.norm().clamp_min(1e-12))
+        teacher_frames = keyframe_latents(prepared, teacher_video)
+        for name, (frame, reference) in keyframe_latents(prepared, video).items():
+            stats.add_scalar(f"keyframe.vs_teacher.{name}", rel_l2(frame, teacher_frames[name][0]))
+            if reference is not None:
+                stats.add_scalar(f"keyframe.student.{name}", rel_l2(frame, reference))
+                stats.add_scalar(f"keyframe.teacher.{name}", rel_l2(teacher_frames[name][0], reference))
+
+    def _initial_checks(self, rows: list[tuple[Any, bool]]) -> dict[str, LogScalar]:
+        """Bit-identity with quantization off, and the PTQ error with bf16 attention (Stage A's protocol)."""
+        out: dict[str, LogScalar] = {}
+        diff = torch.zeros((), device=self.student.device, dtype=torch.float64)
+        if rows:
+            prepared = self._row(rows[0][0])
+            schedulers = self.student.schedulers(prepared)
+            timesteps = (schedulers[0].timesteps, schedulers[1].timesteps)
+            with torch.no_grad(), self.student.quantization(linears=False, attention=False):
+                for rung in (0, self.student.num_rungs - 1):
+                    outs = []
+                    for model in (self.teacher, self.student):
+                        with model.rung_scope(prepared, rung, timesteps) as rt:
+                            outs.append(model.forward_rung(prepared, prepared.video, prepared.audio, rt))
+                    for teacher_t, student_t in zip(*outs, strict=True):
+                        diff = torch.maximum(diff, (teacher_t.double() - student_t.double()).abs().max())
+        if torch.distributed.is_initialized():
+            torch.distributed.all_reduce(diff, op=torch.distributed.ReduceOp.MAX)
+        out["init/noquant_vs_teacher_max_abs"] = float(diff)
+        sha = getattr(self.student, "amax_sha256", "")
+        if sha:  # first 32 bits of the pinned calibration table's sha256, for W&B
+            out["init/amax_sha256_u32"] = float(int(sha[:8], 16))
+        if self.require_bitwise_init and float(diff) != 0.0:
+            raise RuntimeError(f"student with quantization disabled differs from the teacher (max |diff| "
+                               f"{float(diff):.3e}); refusing to train from a student that is not the teacher")
+        weights_only = self._evaluate(rows, rollout=False, attention=False)
+        out.update({
+            f"init/ptq_bf16_attn/{k}": v
+            for k, v in weights_only.items() if k.startswith(("x0_rel", "v_rel", "score"))
+        })
+        out.update(self._compare_ptq_reference(weights_only))
+        if getattr(self.student, "compression", None) is not None:
+            # Training-free token compression alone (bf16 linears): how much of the step-0 error it contributes.
+            compress_only = self._evaluate(rows, rollout=False, linears=False)
+            out.update({
+                f"init/compress_only/{k}": v
+                for k, v in compress_only.items() if k.startswith(("x0_rel", "v_rel", "score"))
+            })
+        return out
+
+    @torch.no_grad()
+    def _provisional_calibration(self) -> dict[str, LogScalar]:
+        """Smoke runs without Stage A's table: max-calibrate the FFN inputs on dense teacher trajectories."""
+        from fastvideo.layers.quantization.nvfp4_qad import calibrating, nvfp4_qad_export_scales
+        pending = [m for m in self.student.qad_linears.values() if m.act_scale == "static" and m.input_amax is None]
+        if not pending:
+            return {}
+        count = int(getattr(self.student, "provisional_rows", 0))
+        rows = self.student.train_rows[:count]
+        per_rank = -(-len(rows) // self.dp_size)
+        mine = [rows[(i * self.dp_size + self.dp_rank) % len(rows)] for i in range(per_rank)]
+        with self.student.quantization(linears=False, attention=False), calibrating(self.student.qad_linears):
+            for spec in mine:
+                prepared = self._row(spec)
+                schedulers = self.student.schedulers(prepared)
+                timesteps = (schedulers[0].timesteps, schedulers[1].timesteps)
+                video, audio = prepared.video, prepared.audio
+                for rung in range(self.student.num_rungs):
+                    with self.student.rung_scope(prepared, rung, timesteps) as rt:
+                        out = self.student.forward_rung(prepared, video, audio, rt)
+                    video, audio = step_packed(schedulers, prepared, rung, video, audio, *out)
+        table = nvfp4_qad_export_scales(self.student.qad_linears)
+        if get_world_group().rank == 0:
+            path = Path(self.training_config.checkpoint.output_dir) / "provisional_amax.json"
+            path.parent.mkdir(parents=True, exist_ok=True)
+            path.write_text(json.dumps(table, indent=1, sort_keys=True))
+            logger.warning("PROVISIONAL FFN scales from %d rows written to %s; use Stage A's table for real runs",
+                           len(rows), path)
+        return {"init/provisional_calibration_rows": float(len(rows))}
+
+    def _compare_ptq_reference(self, weights_only: dict[str, float]) -> dict[str, LogScalar]:
+        """Compare step-0 weights-only x0 rel-L2 per rung to Stage A's (``{"video": [...], "audio": [...]}``)."""
+        if not self.ptq_reference_json:
+            return {}
+        reference = json.loads(Path(self.ptq_reference_json).read_text())
+        deviations = []
+        for modality in MODALITIES:
+            for rung, value in enumerate(reference.get(modality, [])):
+                ours = weights_only.get(f"{self.ptq_reference_metric}/{modality}/rung{rung}")
+                if ours is not None and value:
+                    deviations.append(abs(ours / float(value) - 1.0))
+        worst = max(deviations, default=float("nan"))
+        if deviations and worst > self.ptq_tolerance:
+            raise RuntimeError(f"step-0 PTQ error deviates {worst:.1%} from Stage A's ({self.ptq_reference_json}); "
+                               "the QAD student does not start at the calibrated PTQ point")
+        return {"init/ptq_vs_stage_a_max_rel_dev": worst}
+
+
+__all__ = ["OmniRefPDDQADMethod"]

fastvideo/train/methods/knowledge_distillation/pdd_qad_metrics.py

@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Per-forward (T1) error and reference-fidelity bookkeeping for PDD QAD.
+
+All quantities are on the target rows only (conditioning rows are inputs).
+``x0`` is the scheduler's own denoised estimate, ``x_t + sigma_t * v`` with
+``sigma_t = 1 - timestep``, so x0 error is ``sigma_t`` times velocity error.
+"""
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import Any
+
+import torch
+import torch.distributed as dist
+
+MODALITIES = ("video", "audio")
+
+
+def target_slices(prepared: Any) -> dict[str, slice]:
+    return {
+        "video": slice(prepared.layout.num_condition_video_rows, None),
+        "audio": slice(prepared.layout.num_condition_audio_rows, None)
+    }
+
+
+def denoised(state: torch.Tensor, velocity: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor:
+    """x0 exactly as ``MiniMaxH3Scheduler.step`` forms it (fp32)."""
+    return state.float() + (1 - timestep.to(torch.float32)) * velocity.float()
+
+
+def scheduler_step(scheduler: Any, rung: int, velocity: torch.Tensor, state: torch.Tensor) -> torch.Tensor:
+    """``scheduler.step`` at forward ``rung`` without relying on its internal counter."""
+    scheduler._step_index = rung
+    return scheduler.step(velocity.float(), scheduler.timesteps[rung], state, return_dict=False)[0]
+
+
+def step_packed(schedulers: tuple[Any, Any], prepared: Any, rung: int, video: torch.Tensor, audio: torch.Tensor,
+                video_v: torch.Tensor, audio_v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Advance the packed target rows by one PDD forward (conditioning rows untouched), as the sampler does."""
+    cut = target_slices(prepared)
+    video, audio = video.clone(), audio.clone()
+    video[cut["video"]] = scheduler_step(schedulers[0], rung, video_v[cut["video"]], video[cut["video"]])
+    audio[cut["audio"]] = scheduler_step(schedulers[1], rung, audio_v[cut["audio"]], audio[cut["audio"]])
+    return video, audio
+
+
+def x0_pair(prepared: Any, schedulers: tuple[Any, Any], rung: int, video: torch.Tensor, audio: torch.Tensor,
+            outputs: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> dict[str, dict[str, torch.Tensor]]:
+    """``{modality: {"student": x0, "teacher": x0, "student_v": v, "teacher_v": v}}`` on target rows."""
+    cut = target_slices(prepared)
+    states = {"video": video, "audio": audio}
+    result: dict[str, dict[str, torch.Tensor]] = {}
+    for index, modality in enumerate(MODALITIES):
+        timestep = schedulers[index].timesteps[rung]
+        entry = {}
+        for role, (video_v, audio_v) in outputs.items():
+            velocity = (video_v, audio_v)[index][cut[modality]]
+            entry[role] = denoised(states[modality][cut[modality]], velocity, timestep)
+            entry[f"{role}_v"] = velocity.float()
+        result[modality] = entry
+    return result
+
+
+class RungStats:
+    """Sums of per-rung error terms, reduced over data-parallel groups (one SP rank contributes)."""
+
+    KEYS = ("mse", "err_sq", "ref_sq", "verr_sq", "vref_sq", "dot", "s_sq")
+    EXTRA_KEYS = ("endpoint.video", "endpoint.audio", "keyframe.vs_teacher.first", "keyframe.vs_teacher.last",
+                  "keyframe.student.first", "keyframe.teacher.first", "keyframe.student.last", "keyframe.teacher.last")
+
+    def __init__(self, num_rungs: int) -> None:
+        self.num_rungs = num_rungs
+        self.sums: dict[str, torch.Tensor] = {}
+        self.counts: dict[str, float] = defaultdict(float)
+        # Ranks hold different rows (cases), so every rank reduces this same fixed key list.
+        self.all_keys = [f"{m}.{r}.{k}" for m in MODALITIES for r in range(num_rungs) for k in self.KEYS]
+        self.all_keys += list(self.EXTRA_KEYS)
+
+    def _add(self, key: str, value: torch.Tensor) -> None:
+        self.sums[key] = self.sums.get(key, 0.0) + value.detach().double()
+
+    def add_rung(self, rung: int, pair: dict[str, dict[str, torch.Tensor]]) -> None:
+        for modality, entry in pair.items():
+            s, t, sv, tv = entry["student"], entry["teacher"], entry["student_v"], entry["teacher_v"]
+            prefix = f"{modality}.{rung}"
+            self._add(f"{prefix}.mse", (s - t).pow(2).mean())
+            self._add(f"{prefix}.err_sq", (s - t).pow(2).sum())
+            self._add(f"{prefix}.ref_sq", t.pow(2).sum())
+            self._add(f"{prefix}.verr_sq", (sv - tv).pow(2).sum())
+            self._add(f"{prefix}.vref_sq", tv.pow(2).sum())
+            self._add(f"{prefix}.dot", (s * t).sum())
+            self._add(f"{prefix}.s_sq", s.pow(2).sum())
+            for key in self.KEYS:
+                self.counts[f"{prefix}.{key}"] += 1
+
+    def add_scalar(self, key: str, value: torch.Tensor | float) -> None:
+        if key not in self.EXTRA_KEYS:
+            raise KeyError(f"{key} is not a reduced RungStats key")
+        self._add(key, torch.as_tensor(value, dtype=torch.float64, device=self._device()))
+        self.counts[key] += 1
+
+    def _device(self) -> torch.device:
+        return torch.device("cuda", torch.cuda.current_device()) if torch.cuda.is_available() else torch.device("cpu")
+
+    def reduce(self, contribute: bool) -> dict[str, float]:
+        """All-reduce sums and counts (call on every rank in the same key order); returns means."""
+        keys = self.all_keys
+        device = self._device()
+        values = torch.tensor([float(self.sums.get(k, 0.0)) for k in keys] + [self.counts.get(k, 0.0) for k in keys],
+                              dtype=torch.float64,
+                              device=device)
+        if not contribute:
+            values.zero_()
+        if dist.is_available() and dist.is_initialized():
+            dist.all_reduce(values)
+        sums, counts = values[:len(keys)], values[len(keys):]
+        return {key: float(sums[i]) / float(counts[i]) for i, key in enumerate(keys) if float(counts[i]) > 0}
+
+
+def summarize(means: dict[str, float], num_rungs: int, weights: dict[str, float]) -> dict[str, float]:
+    """Per-rung rel-L2 / cosine / MSE plus the weighted mean x0 rel-L2 ``score`` (lower is better)."""
+    out: dict[str, float] = {}
+    score_num = score_den = 0.0
+    for modality in MODALITIES:
+        rel_all = []
+        for rung in range(num_rungs):
+            p = f"{modality}.{rung}"
+            if f"{p}.ref_sq" not in means:
+                continue
+            rel = (means[f"{p}.err_sq"] / max(means[f"{p}.ref_sq"], 1e-30))**0.5
+            vrel = (means[f"{p}.verr_sq"] / max(means[f"{p}.vref_sq"], 1e-30))**0.5
+            cos = means[f"{p}.dot"] / max((means[f"{p}.s_sq"] * means[f"{p}.ref_sq"])**0.5, 1e-30)
+            out[f"x0_rel_l2/{modality}/rung{rung}"] = rel
+            out[f"v_rel_l2/{modality}/rung{rung}"] = vrel
+            out[f"x0_cos/{modality}/rung{rung}"] = cos
+            out[f"x0_mse/{modality}/rung{rung}"] = means[f"{p}.mse"]
+            rel_all.append(rel)
+        if rel_all:
+            out[f"x0_rel_l2/{modality}/mean"] = sum(rel_all) / len(rel_all)
+            score_num += weights[modality] * out[f"x0_rel_l2/{modality}/mean"]
+            score_den += weights[modality]
+    out["score"] = score_num / max(score_den, 1e-30)
+    for key, value in means.items():
+        if key.startswith(("endpoint.", "keyframe.")):
+            out[key.replace(".", "/")] = value
+    return out
+
+
+def keyframe_latents(prepared: Any, final_video: torch.Tensor) -> dict[str, tuple[torch.Tensor, torch.Tensor | None]]:
+    """Generated keyframe latents ``[C, H, W]`` (first; last for first_last_frame) and their reference image.
+
+    OmniRef encodes keyframe references at high resolution with the target's
+    aspect ratio, so the reference latent is area-resized to the target grid: a
+    rough proxy, logged only. The guarded keyframe metric compares the student's
+    keyframe with the teacher's (``keyframe/vs_teacher``).
+    """
+    import torch.nn.functional as F
+
+    from fastvideo.pipelines.basic.minimax_h3.packing import unpatchify_video_tokens
+
+    channels, frames, height, width = prepared.latent_shape
+    patch = (1, 2, 2)
+    cut = target_slices(prepared)["video"]
+    video = unpatchify_video_tokens(final_video[cut][None].float(), frames, height, width, channels, patch)[0]
+    images: list[torch.Tensor | None] = []
+    offset = 0
+    for ref in prepared.references:
+        if ref.media_type == "audio":
+            continue
+        count = ref.num_latent_frames * (ref.latent_height // 2) * (ref.latent_width // 2)
+        image = None
+        same_aspect = abs(ref.latent_height * width - ref.latent_width * height) <= 0.02 * ref.latent_width * height
+        if ref.media_type == "image" and same_aspect:
+            rows = prepared.reference_video_rows[offset:offset + count].to(video.device)
+            latent = unpatchify_video_tokens(rows[None].float(), 1, ref.latent_height, ref.latent_width, channels,
+                                             patch)[:, :, 0]
+            if (ref.latent_height, ref.latent_width) != (height, width):
+                latent = F.interpolate(latent, size=(height, width), mode="area")
+            image = latent[0]
+        images.append(image)
+        offset += count
+    # (metric name, generated frame, keyframe index) per case; FL2VA keyframes carry the same names.
+    pairs = {
+        "first_frame": [("first", 0, 0)],
+        "first_last_frame": [("first", 0, 0), ("last", -1, 1)],
+        "i2va": [("first", 0, 0)],
+        "l2va": [("last", -1, 0)],
+        "fl2va": [("first", 0, 0), ("last", -1, 1)],
+    }
+    return {
+        name: (video[:, frame], images[ref] if ref < len(images) else None)
+        for name, frame, ref in pairs.get(prepared.spec.case, [])
+    }
+
+
+def rel_l2(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    return (a.float() - b.float()).norm() / b.float().norm().clamp_min(1e-12)
+
+
+__all__ = ["RungStats", "keyframe_latents", "rel_l2", "step_packed", "summarize", "target_slices", "x0_pair"]

fastvideo/train/methods/knowledge_distillation/pdd_qad_recon.py

@@ -0,0 +1,373 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Layer-wise reconstruction for the FastH3 OmniRef QAD student, streamed per block.
+
+Supervised units R, per transformer block: the attention sub-layer (q/k/v/gate/out NVFP4 projections, QK
+RMSNorm, RoPE and the sparse FP4 VSA-128 attention, as deployed) and the FFN sub-layer (NVFP4 fc_in/fc_out).
+Each student unit receives the frozen teacher's input to that unit (stop-gradient) and matches the teacher
+unit's output, so quantization error is corrected locally and cannot accumulate into a drift away from the
+teacher.
+
+Locality makes this cheap: inside one teacher forward, right after teacher block ``b`` runs, student block
+``b`` runs on the captured inputs and is backpropagated at once. There is no end-to-end student forward or
+backward, no activation recompute, and memory holds one block's activations.
+
+Attention is sparse (VSA-128, the teacher's trained selection), not dense as in the paper's warmup: our
+teacher was trained sparse and the student is deployed sparse, so dense attention would supervise a
+configuration that never runs.
+
+The loss per unit is the relative squared error ``||s - t||^2 / ||t||^2`` (late FFN outputs reach ~1e6, so an
+unnormalized sum would only train those), averaged over R. With a ``GroupedRelError`` it is computed per row group
+(target video, target audio, conditioning) and weighted: over the whole packed sequence audio is ~1.4% of the rows
+and massive text tokens dominate ``||t||^2``, so an ungrouped ratio barely trains audio.
+"""
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+
+def ensure_fsdp_root_initialized(module: torch.nn.Module) -> None:
+    """Make ``module`` its FSDP2 tree's root before any of its sharded blocks runs on its own.
+
+    FSDP2 picks the root lazily, at the first forward of any sharded module. If a block runs first (as in
+    reconstruction, which calls student blocks directly), that block claims the root and the model's next full
+    forward fails with "FSDP state has already been lazily initialized". Idempotent; a no-op without FSDP2.
+    """
+    get_state = getattr(module, "_get_fsdp_state", None)
+    if get_state is not None:
+        get_state()._lazy_init()
+
+
+def _unwrap(module: torch.nn.Module) -> torch.nn.Module:
+    return getattr(module, "_checkpoint_wrapped_module", module)
+
+
+GROUPS = ("video", "audio", "condition")  # target video rows, target audio rows, text + reference rows
+
+
+def packed_row_groups(layout: Any, sequence_length: int) -> torch.Tensor:
+    """Group id per packed row: 0 target video, 1 target audio, 2 conditioning (text, references)."""
+    groups = torch.full((sequence_length, ), 2, dtype=torch.long)
+    groups[layout.video_indices[layout.num_condition_video_rows:]] = 0
+    groups[layout.audio_indices[layout.num_condition_audio_rows:]] = 1
+    return groups
+
+
+def local_row_groups(groups: torch.Tensor, device: torch.device) -> torch.Tensor:
+    """This sequence-parallel rank's rows of ``groups`` as the transformer shards them; padding is -1."""
+    from fastvideo.distributed import get_sp_world_size, model_parallel_is_initialized
+    from fastvideo.distributed.communication_op import sequence_model_parallel_shard
+
+    shifted = (groups + 1).to(device)
+    if model_parallel_is_initialized() and get_sp_world_size() > 1:
+        shifted, _ = sequence_model_parallel_shard(shifted, dim=0)  # pads with 0 -> -1 after the shift
+    return shifted - 1
+
+
+class GroupedRelError:
+    """Weighted mean over row groups of ``||s - t||^2 / ||t||^2`` with sequence-parallel-global sums.
+
+    Each rank backpropagates ``sp_size * w_g * num_local / den_global``. Summed over ranks, ``w_g * num_local /
+    den_global`` is the global ratio, so a group living on one rank (e.g. target audio) is weighted exactly as
+    if unsharded. The ``sp_size`` factor matches the x0 output loss: FSDP averages gradients over every rank,
+    SP ranks included, and the output loss is identical on every SP rank, where the transformer's output
+    all-gather backward (a SUM reduce-scatter) gives each rank ``sp_size`` times its share of the gradient.
+    Reconstruction has no gather, so without the factor the averaged gradient would be ``1 / sp_size`` of the
+    global ratio's. Groups absent from the sequence are dropped from the weight normalization.
+    """
+
+    def __init__(self, groups: torch.Tensor, weights: dict[str, float]) -> None:
+        self.groups = groups
+        self.weights = torch.tensor([float(weights.get(name, 0.0)) for name in GROUPS], device=groups.device)
+
+    def __call__(self, student: torch.Tensor, teacher: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """``(local loss to backprop, global per-group relative error [3], detached)``.
+
+        The local loss is this rank's ``sp_size``-scaled share (see the class docstring); the detached global
+        per-group errors are what to log.
+        """
+        from fastvideo.distributed import get_sp_group, get_sp_world_size, model_parallel_is_initialized
+
+        teacher = teacher.float().reshape(-1, teacher.shape[-1])
+        err = (student.float().reshape(-1, teacher.shape[-1]) - teacher).pow(2).sum(-1)
+        ref = teacher.pow(2).sum(-1)
+        one_hot = (self.groups[:, None] == torch.arange(len(GROUPS), device=err.device)[None, :]).float()
+        num = err @ one_hot  # [3], differentiable
+        totals = torch.stack((num.detach(), ref @ one_hot))  # [2, 3]
+        sp_size = get_sp_world_size() if model_parallel_is_initialized() else 1
+        if sp_size > 1:
+            totals = get_sp_group().all_reduce(totals)
+        present = totals[1] > 0
+        den = totals[1].clamp_min(1e-30)
+        weights = torch.where(present, self.weights, torch.zeros_like(self.weights))
+        norm = weights.sum().clamp_min(1e-30)
+        loss = sp_size * (weights * num / den).sum() / norm
+        return loss, torch.where(present, totals[0] / den, torch.full_like(den, float("nan")))
+
+
+class _Replace:
+    """Forward pre-hook replacing a sub-layer's first input with the teacher's, plus an output capture."""
+
+    def __init__(self) -> None:
+        self.value: torch.Tensor | None = None
+        self.output: torch.Tensor | None = None
+
+    def pre(self, module: torch.nn.Module, args: tuple[Any, ...]) -> tuple[Any, ...]:
+        return (self.value, *args[1:])
+
+    def post(self, module: torch.nn.Module, args: tuple[Any, ...], output: Any) -> None:
+        self.output = output
+
+
+class _Capture:
+
+    def __init__(self) -> None:
+        self.input: torch.Tensor | None = None
+        self.output: torch.Tensor | None = None
+
+    def __call__(self, module: torch.nn.Module, args: tuple[Any, ...], output: Any) -> None:
+        self.input, self.output = args[0], output
+
+
+class BlockReconstruction:
+    """Stream reconstruction losses for every block during one teacher forward.
+
+    ``scale`` multiplies each block's loss before its backward (e.g. ``1 / (units x rungs x accum)``).
+    Use as a context manager around the teacher's ``forward_rung``, with grad disabled for the teacher.
+    """
+
+    def __init__(self,
+                 teacher: torch.nn.Module,
+                 student: torch.nn.Module,
+                 scale: float,
+                 criterion: GroupedRelError | None = None,
+                 measure_only: bool = False) -> None:
+        ensure_fsdp_root_initialized(teacher)
+        ensure_fsdp_root_initialized(student)
+        self.teacher_blocks = list(teacher.transformer_blocks)
+        self.student_blocks = list(student.transformer_blocks)
+        if len(self.teacher_blocks) != len(self.student_blocks):
+            raise ValueError("teacher and student block counts differ")
+        self.scale = scale
+        self.criterion = criterion
+        self.measure_only = measure_only
+        self.attn_rel: list[torch.Tensor] = []
+        self.ff_rel: list[torch.Tensor] = []
+        self.group_rel: list[torch.Tensor] = []  # [3] per unit, global, detached
+        self._handles: list[Any] = []
+
+    def _weighted(self, rel: torch.Tensor) -> torch.Tensor:
+        """Group-weighted mean of per-group relative errors ``[..., 3]`` (NaN = absent group, dropped)."""
+        assert self.criterion is not None
+        weights = self.criterion.weights.expand_as(rel) * (~rel.isnan())
+        return (weights * rel.nan_to_num()).sum(-1) / weights.sum(-1).clamp_min(1e-30)
+
+    def _unit_loss(self, student: torch.Tensor, teacher: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """``(loss to backprop, sequence-parallel-global relative error to log, detached)``."""
+        if self.criterion is None:
+            teacher_f = teacher.float()
+            loss = (student.float() - teacher_f).pow(2).mean() / teacher_f.pow(2).mean().clamp_min(1e-20)
+            return loss, loss.detach()
+        loss, per_group = self.criterion(student, teacher)
+        self.group_rel.append(per_group)
+        return loss, self._weighted(per_group)
+
+    def __enter__(self) -> BlockReconstruction:
+        for index, block in enumerate(self.teacher_blocks):
+            inner = _unwrap(block)
+            attn, ff = _Capture(), _Capture()
+            self._handles.append(inner.attn.register_forward_hook(attn))
+            self._handles.append(inner.ff.register_forward_hook(ff))
+            self._handles.append(block.register_forward_hook(self._after_block(index, attn, ff)))
+        return self
+
+    def __exit__(self, *exc: Any) -> None:
+        for handle in self._handles:
+            handle.remove()
+        self._handles.clear()
+
+    def _after_block(self, index: int, attn: _Capture, ff: _Capture):
+
+        def hook(module: torch.nn.Module, args: tuple[Any, ...], output: Any) -> None:
+            self._student_block(index, args, attn, ff)
+            attn.input = attn.output = ff.input = ff.output = None
+
+        return hook
+
+    def _student_block(self, index: int, block_args: tuple[Any, ...], attn: _Capture, ff: _Capture) -> None:
+        block = self.student_blocks[index]
+        inner = _unwrap(block)
+        replace_attn, replace_ff = _Replace(), _Replace()
+        replace_attn.value, replace_ff.value = attn.input, ff.input
+        handles = [
+            inner.attn.register_forward_pre_hook(replace_attn.pre),
+            inner.attn.register_forward_hook(replace_attn.post),
+            inner.ff.register_forward_pre_hook(replace_ff.pre),
+            inner.ff.register_forward_hook(replace_ff.post)
+        ]
+        try:
+            with torch.set_grad_enabled(not self.measure_only):
+                # The block's own hidden input only feeds norm1 and the residual, whose results are unused; the
+                # teacher attention input has the right shape. The call goes through the block, and the block
+                # input and output join the graph (zero weight), so FSDP's pre-/post-backward hooks re-gather the
+                # block's weights for backward and reduce-scatter its gradients.
+                hidden = attn.input.detach().requires_grad_(not self.measure_only)
+                block_out = block(hidden, *block_args[1:])
+                attn_loss, attn_rel = self._unit_loss(replace_attn.output, attn.output)
+                ff_loss, ff_rel = self._unit_loss(replace_ff.output, ff.output)
+                if not self.measure_only:
+                    ((attn_loss + ff_loss) * self.scale + 0.0 * block_out.float().sum()).backward()
+        finally:
+            for handle in handles:
+                handle.remove()
+        self.attn_rel.append(attn_rel)
+        self.ff_rel.append(ff_rel)
+
+    def loss(self) -> torch.Tensor:
+        """Mean (group-weighted) relative error over all units of this forward, sequence-parallel-global, detached."""
+        if self.criterion is not None and self.group_rel:
+            return self._weighted(torch.stack(self.group_rel)).mean()  # [units, 3], NaN for absent groups
+        return (torch.stack(self.attn_rel).sum() + torch.stack(self.ff_rel).sum()) / (2 * len(self.attn_rel))
+
+    def metrics(self, prefix: str) -> dict[str, torch.Tensor]:
+        attn, ff = torch.stack(self.attn_rel), torch.stack(self.ff_rel)
+        third = max(1, len(attn) // 3)
+        groups = {}
+        if self.group_rel:
+            per_group = torch.stack(self.group_rel).nanmean(dim=0)
+            groups = {f"{prefix}/{name}_rel": per_group[i] for i, name in enumerate(GROUPS)}
+        return {
+            **groups, f"{prefix}/attn_rel": attn.mean(),
+            f"{prefix}/ff_rel": ff.mean(),
+            f"{prefix}/attn_rel_first_third": attn[:third].mean(),
+            f"{prefix}/attn_rel_last_third": attn[-third:].mean(),
+            f"{prefix}/ff_rel_last_third": ff[-third:].mean()
+        }
+
+
+def theseus_widths(num_blocks: int) -> tuple[int, ...]:
+    """Binary-merge segment widths ``1, 2, 4, ...`` below ``num_blocks``, then the whole stack."""
+    widths, width = [], 1
+    while width < num_blocks:
+        widths.append(width)
+        width *= 2
+    return (*widths, num_blocks)
+
+
+def theseus_schedule(total_steps: int, widths: tuple[int, ...], gamma: float) -> list[int]:
+    """Segment width of each warm-up step: level ``k`` gets ``~gamma**k`` of ``total_steps`` (at least one each).
+
+    Short-to-long, as JustQuant's Theseus QAD (arXiv:2609.33601): narrow, teacher-forced targets first, then
+    progressively longer student paths, ending on the whole stack.
+    """
+    if total_steps < len(widths):
+        raise ValueError(f"{total_steps} warm-up steps cannot cover {len(widths)} segment widths {widths}")
+    weights = [gamma**k for k in range(len(widths))]
+    raw = [total_steps * w / sum(weights) for w in weights]
+    steps = [max(1, int(r)) for r in raw]
+    # Hand out what flooring lost to the levels with the largest remainders (the long ones on ties).
+    for index in sorted(range(len(widths)), key=lambda i: (raw[i] - int(raw[i]), i), reverse=True):
+        if sum(steps) >= total_steps:
+            break
+        steps[index] += 1
+    while sum(steps) > total_steps:  # the at-least-one floor overshot: shave the longest levels
+        index = max(i for i in range(len(steps)) if steps[i] > 1)
+        steps[index] -= 1
+    return [width for width, count in zip(widths, steps, strict=True) for _ in range(count)]
+
+
+class SegmentReconstruction:
+    """Theseus warm-up: segments of ``width`` consecutive blocks, each from the teacher's input, during one forward.
+
+    When teacher block ``g + width - 1`` finishes, student blocks ``g .. g + width - 1`` run chained from the
+    teacher's (stop-gradient) input to block ``g`` with the teacher's block arguments, match the teacher's output of
+    the segment's last block, and are backpropagated at once. Segments are teacher-forced at their boundaries, so a
+    width-1 segment is block-wise distillation and the full width is the whole stack. ``scale`` multiplies each
+    segment's loss before its backward. Use around the teacher's ``forward_rung`` with grad disabled.
+    """
+
+    def __init__(self,
+                 teacher: torch.nn.Module,
+                 student: torch.nn.Module,
+                 width: int,
+                 scale: float,
+                 criterion: GroupedRelError | None = None) -> None:
+        ensure_fsdp_root_initialized(teacher)
+        ensure_fsdp_root_initialized(student)
+        self.teacher_blocks = list(teacher.transformer_blocks)
+        self.student_blocks = list(student.transformer_blocks)
+        if len(self.teacher_blocks) != len(self.student_blocks):
+            raise ValueError("teacher and student block counts differ")
+        if not 1 <= width <= len(self.teacher_blocks):
+            raise ValueError(f"segment width {width} outside 1..{len(self.teacher_blocks)}")
+        self.width = width
+        self.scale = scale
+        self.criterion = criterion
+        self.rel: list[torch.Tensor] = []
+        self.group_rel: list[torch.Tensor] = []
+        self._start: torch.Tensor | None = None
+        self._args: list[tuple[Any, ...]] = []
+        self._handles: list[Any] = []
+
+    @property
+    def num_segments(self) -> int:
+        return -(-len(self.teacher_blocks) // self.width)
+
+    def __enter__(self) -> SegmentReconstruction:
+        for index, block in enumerate(self.teacher_blocks):
+            self._handles.append(block.register_forward_hook(self._after_block(index)))
+        return self
+
+    def __exit__(self, *exc: Any) -> None:
+        for handle in self._handles:
+            handle.remove()
+        self._handles.clear()
+        self._start, self._args = None, []
+
+    def _after_block(self, index: int):
+
+        def hook(module: torch.nn.Module, args: tuple[Any, ...], output: Any) -> None:
+            if index % self.width == 0:
+                self._start, self._args = args[0], []
+            self._args.append(args[1:])
+            if index % self.width == self.width - 1 or index == len(self.teacher_blocks) - 1:
+                self._student_segment(index - len(self._args) + 1, output)
+                self._start, self._args = None, []
+
+        return hook
+
+    def _student_segment(self, first: int, teacher_out: torch.Tensor) -> None:
+        assert self._start is not None
+        with torch.set_grad_enabled(True):
+            hidden = self._start.detach()
+            for offset, block_args in enumerate(self._args):
+                hidden = self.student_blocks[first + offset](hidden, *block_args)
+            if self.criterion is None:
+                target = teacher_out.float()
+                loss = (hidden.float() - target).pow(2).mean() / target.pow(2).mean().clamp_min(1e-20)
+                rel = loss.detach()
+            else:
+                loss, per_group = self.criterion(hidden, teacher_out)
+                self.group_rel.append(per_group)
+                weights = self.criterion.weights * (~per_group.isnan())
+                rel = (weights * per_group.nan_to_num()).sum() / weights.sum().clamp_min(1e-30)
+            (loss * self.scale).backward()
+        self.rel.append(rel)
+
+    def loss(self) -> torch.Tensor:
+        """Mean (group-weighted) relative error over this forward's segments, sequence-parallel-global, detached."""
+        return torch.stack(self.rel).mean()
+
+    def metrics(self, prefix: str) -> dict[str, torch.Tensor]:
+        out = {f"{prefix}/rel": self.loss(), f"{prefix}/last_segment_rel": self.rel[-1]}
+        if self.group_rel:
+            per_group = torch.stack(self.group_rel).nanmean(dim=0)
+            out.update({f"{prefix}/{name}_rel": per_group[i] for i, name in enumerate(GROUPS)})
+        return out
+
+
+__all__ = [
+    "BlockReconstruction", "GROUPS", "GroupedRelError", "SegmentReconstruction", "local_row_groups",
+    "packed_row_groups", "theseus_schedule", "theseus_widths"
+]

fastvideo/train/methods/knowledge_distillation/pdd_qad_step0.py

@@ -0,0 +1,134 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Reuse a QAD run's step-0 evaluation (PTQ baseline, seed+1 noise floor, reconstruction normalizers).
+
+A Stage-2 run from a fixed checkpoint spends most of its first ~1.75 h on step 0: the full held-out evaluation,
+the same evaluation at seed+1 for the noise floor, and the reconstruction calibration. All of it is a deterministic
+function of the teacher and student checkpoints (``init_from``, plus ``init_student_dcp``), the calibration table
+(or the provisional-calibration rows), the quantization plan, the evaluation set and the score weights, so a rerun
+with the same inputs can load it. The cache is keyed by a fingerprint of exactly those inputs; any change (another
+init checkpoint, eval rows, tolerances, modality weights, ...) is a mismatch and the run refuses to start.
+
+No torch import, so the fingerprint can be computed offline from a run's ``cfg.yaml``.
+"""
+from __future__ import annotations
+
+import hashlib
+import json
+import math
+from pathlib import Path
+from typing import Any
+
+STEP0_FORMAT = 2
+METHOD_KEYS = ("init_student_dcp", "max_eval_rows", "max_rollout_rows", "noise_floor_eval", "min_eligibility_tolerance",
+               "eligibility_keys", "fidelity_tolerance", "objective", "recon_normalize", "recon_norm_rows",
+               "recon_group_weights", "rungs_per_row")
+DATA_KEYS = ("manifests", "cases", "max_frames", "eval_manifest", "eval_resolutions", "eval_per_group",
+             "eval_max_frames", "eval_seed", "width", "height", "num_frames")
+# The training-row plan; it only reaches step 0 through provisional FFN calibration (no amax_json).
+PROVISIONAL_DATA_KEYS = ("seed", "per_group", "max_train_rows", "train_manifest")
+QAD_KEYS = ("linears", "fp4_attention")
+# Logged per evaluation but not part of the step-0 result.
+VOLATILE = ("eval/seconds", "eval/score_vs_ptq")
+
+
+def _plain(value: Any) -> Any:
+    """OmegaConf / mapping / sequence -> JSON-native containers."""
+    if hasattr(value, "items"):
+        return {str(k): _plain(v) for k, v in value.items()}
+    if isinstance(value, list | tuple) or (hasattr(value, "__iter__") and not isinstance(value, str | bytes)):
+        return [_plain(v) for v in value]
+    return value
+
+
+def step0_fingerprint_inputs(method_config: Any, data_config: Any, qad_config: Any, amax_sha256: str,
+                             num_eval_rows: int, *, student_init_from: str, teacher_init_from: str) -> dict[str, Any]:
+    """Everything step 0 depends on. The init paths stand for their weights and ``fastvideo_inference.json``."""
+    method, data, qad = _plain(method_config), _plain(data_config), _plain(qad_config)
+    linears = qad.get("linears") or {}
+    provisional = int(linears.get("provisional_calibration_rows") or 0) > 0 and not linears.get("amax_json")
+    return {
+        "format": STEP0_FORMAT,
+        "models": {
+            "student_init_from": str(student_init_from),
+            "teacher_init_from": str(teacher_init_from)
+        },
+        "method": {
+            key: method.get(key)
+            for key in METHOD_KEYS
+        },
+        # ``summarize`` weights the score with these; best_score and the guards compare against it.
+        "score_weights": {
+            modality: float(method.get(f"{modality}_weight", 1.0))
+            for modality in ("video", "audio")
+        },
+        "data": {
+            key: data.get(key)
+            for key in DATA_KEYS
+        },
+        "provisional_plan": {
+            key: data.get(key)
+            for key in PROVISIONAL_DATA_KEYS
+        } if provisional else None,
+        "qad": {
+            key: qad.get(key)
+            for key in QAD_KEYS
+        },
+        "amax_sha256": amax_sha256,
+        "num_eval_rows": int(num_eval_rows),
+    }
+
+
+def fingerprint(inputs: dict[str, Any]) -> str:
+    return hashlib.sha256(json.dumps(inputs, sort_keys=True).encode()).hexdigest()
+
+
+def save_step0(path: str | Path, inputs: dict[str, Any], metrics: dict[str, float]) -> None:
+    payload = {
+        "format": STEP0_FORMAT,
+        "fingerprint": fingerprint(inputs),
+        "inputs": inputs,
+        "metrics": {
+            key: float(value)
+            for key, value in metrics.items() if key not in VOLATILE
+        },
+    }
+    Path(path).parent.mkdir(parents=True, exist_ok=True)
+    tmp = Path(f"{path}.tmp")
+    tmp.write_text(json.dumps(payload, indent=1, sort_keys=True))
+    tmp.replace(path)
+
+
+def load_step0(path: str | Path, inputs: dict[str, Any]) -> dict[str, float] | None:
+    """The cached step-0 metrics, None if there is no cache; raises on a fingerprint mismatch."""
+    path = Path(path)
+    if not path.exists():
+        return None
+    payload = json.loads(path.read_text())
+    expected = fingerprint(inputs)
+    if payload.get("format") != STEP0_FORMAT or payload.get("fingerprint") != expected:
+        cached = payload.get("inputs", {})
+        diff = sorted(k for k in set(cached) | set(inputs) if cached.get(k) != inputs.get(k))
+        raise ValueError(f"step-0 cache {path} was computed for other inputs (differs in {diff}); "
+                         f"delete it or point step0_cache elsewhere")
+    return {key: float(value) for key, value in payload["metrics"].items()}
+
+
+def restore_step0(metrics: dict[str, float], num_rungs: int, recon: bool, noise_floor: bool,
+                  min_tolerance: float) -> tuple[dict[str, float], list[float] | None, dict[str, float]]:
+    """(eval summary, reconstruction normalizers, eligibility tolerances) from cached step-0 metrics."""
+    summary = {key[len("eval/"):]: value for key, value in metrics.items() if key.startswith("eval/")}
+    if "score" not in summary:
+        raise ValueError("step-0 cache has no eval/score")
+    recon_norm = None
+    if recon:
+        recon_norm = [max(metrics[f"init/recon_step0/rung{rung}"], 1e-12) for rung in range(num_rungs)]
+    tolerances: dict[str, float] = {}
+    if noise_floor:
+        prefix = "init/noise_sigma/"
+        tolerances = {
+            key[len(prefix):]: max(2.0 * value, min_tolerance)
+            for key, value in metrics.items() if key.startswith(prefix) and math.isfinite(value)
+        }
+        if not tolerances:
+            raise ValueError("noise_floor_eval is on but the step-0 cache has no noise-floor sigmas")
+    return summary, recon_norm, tolerances

fastvideo/train/models/minimax_h3/flex.py

@@ -0,0 +1,356 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pika H3 Flex (six-forward re-noise sampler) as a per-forward model for NVFP4 QAD.
+
+Same contract as ``omniref_pdd.MiniMaxH3OmniRefPDDModel``, so ``OmniRefPDDQADMethod`` runs unchanged: one rung per
+sampler forward, ``rung_scope`` / ``forward_rung`` / ``schedulers`` / ``prepare`` / ``quantization``. ``init_from``
+is a merged Flex model directory written by ``scripts/distill/pika_h3_flex_qad/merge_pika_adapter.py``.
+
+Differences from OmniRef:
+
+- Dense attention (``FLASH_ATTN``): Flex has no VSA compression gate, so there is no tile selection to freeze. QAD
+  covers the NVFP4 linears and, optionally (``qad.fp4_attention``), dense FP4 attention as FlashAttention-4 NVFP4
+  deploys it (Q/K NVFP4, P and V BF16; ``fp4_dense_qat``).
+- No PDD head fusion: every rung is a plain transformer forward over Pika's clean-time knots
+  (1000, 750, 500, 250, 100, 40), shifted per modality (video 12, audio 3), starting from unit noise.
+- Re-noise steps: ``x_next = (1 - sigma_next) x0 + sigma_next eps``. ``eps`` is a function of (row seed, forward,
+  modality) only, so the teacher-forced pass, the student rollout and the seed+1 noise-floor eval of a row all draw
+  the same noise. It is not Pika's sampler stream; the distribution is.
+
+``MiniMaxH3FlexRef2VAModel`` reads Ref2VA rows (``omniref_data``) on ``transformer_ref``;
+``MiniMaxH3FlexFL2VAModel`` reads T2VA / I2VA / L2VA / FL2VA keyframe-anchor rows on ``transformer``.
+"""
+from __future__ import annotations
+
+import contextlib
+import json
+from collections.abc import Iterator
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from fastvideo.configs.pipelines.pika_h3_flex import (PIKA_AUDIO_SCHEDULER_SHIFT, PIKA_DENOISING_STEPS,
+                                                      PIKA_VIDEO_SCHEDULER_SHIFT)
+from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler
+from fastvideo.train.models.minimax_h3.omniref_pdd import MiniMaxH3OmniRefPDDModel, _pin_amax
+
+if TYPE_CHECKING:
+    from fastvideo.train.models.minimax_h3.omniref_data import PreparedRow
+    from fastvideo.train.utils.training_config import TrainingConfig
+
+logger = init_logger(__name__)
+CONTRACT_FILE = "pika_flex_qad.json"
+CONTRACT_SCHEMA = "pika-flex-qad-v1"
+TASK_TO_TRUNK = {"fl2va": "transformer", "ref2va": "transformer_ref"}
+# Block projections and FFN linears: the set the NVFP4 export quantizes. Norms, AdaLN, the token refiner,
+# embeddings and the output heads stay frozen.
+FLEX_TRAINABLE = r"transformer_blocks\.\d+\.(attn\.(to_q|to_k|to_v|to_out)|ff\.(fc_in|fc_out))\."
+_VIDEO_STREAM, _AUDIO_STREAM = 0, 1
+
+
+@dataclass(frozen=True)
+class FlexContract:
+    """The merged checkpoint's sampling contract (``pika_flex_qad.json``)."""
+    task: str
+    trunk: str
+    denoising_steps: tuple[int, ...]
+    video_shift: float
+    audio_shift: float
+
+    @classmethod
+    def read(cls, model_path: str) -> FlexContract:
+        path = Path(model_path) / CONTRACT_FILE
+        if not path.is_file():
+            raise ValueError(f"{model_path} has no {CONTRACT_FILE}; build it with merge_pika_adapter.py")
+        data = json.loads(path.read_text())
+        if data.get("schema") != CONTRACT_SCHEMA or data.get("sampler") != "renoise":
+            raise ValueError(f"{path} is not a {CONTRACT_SCHEMA} re-noise contract")
+        contract = cls(str(data["task"]), str(data["trunk"]), tuple(int(s) for s in data["denoising_steps"]),
+                       float(data["video_scheduler_shift"]), float(data["audio_scheduler_shift"]))
+        released = (PIKA_DENOISING_STEPS, PIKA_VIDEO_SCHEDULER_SHIFT, PIKA_AUDIO_SCHEDULER_SHIFT)
+        if (contract.denoising_steps, contract.video_shift, contract.audio_shift) != released:
+            raise ValueError(f"{path} does not carry the released Pika H3 Flex ladder {released}")
+        if TASK_TO_TRUNK.get(contract.task) != contract.trunk:
+            raise ValueError(f"{path}: task {contract.task!r} does not run on trunk {contract.trunk!r}")
+        return contract
+
+
+def flex_sigmas(steps: tuple[int, ...], shift: float) -> torch.Tensor:
+    """Shifted sigma grid over ``steps`` plus the terminal zero, as ``MiniMaxH3DenoisingStage._set_shifted_schedule``."""
+    base = torch.tensor([step / 1000.0 for step in steps] + [0.0], dtype=torch.float32)
+    return shift * base / (1 + (shift - 1) * base)
+
+
+def renoise_seed(row_seed: int, forward: int, stream: int) -> int:
+    """Distinct generator seed per (row seed, sampler forward, modality); forwards < 8, two streams."""
+    return (int(row_seed) + 1) * 16 + 2 * int(forward) + int(stream)
+
+
+class RowRenoiseScheduler(MiniMaxH3Scheduler):
+    """H3 re-noise scheduler whose fresh noise at forward ``k`` is fixed by (row seed, ``k``, modality)."""
+
+    def __init__(self, shift: float, row_seed: int, stream: int) -> None:
+        super().__init__(shift=shift, stochastic_sampling=True, shift_provided_sigmas=True)
+        self._row_seed = int(row_seed)
+        self._stream = int(stream)
+
+    def step(self,
+             model_output: torch.FloatTensor,
+             timestep: float | torch.FloatTensor,
+             sample: torch.FloatTensor,
+             return_dict: bool = True) -> Any:
+        if self._step_index is None:
+            raise ValueError("RowRenoiseScheduler needs an explicit step index (use pdd_qad_metrics.scheduler_step)")
+        generator = torch.Generator(device=sample.device)
+        generator.manual_seed(renoise_seed(self._row_seed, self._step_index, self._stream))
+        self.set_noise_generator(generator)
+        return super().step(model_output, timestep, sample, return_dict=return_dict)
+
+
+class MiniMaxH3FlexRef2VAModel(MiniMaxH3OmniRefPDDModel):
+    """Teacher (merged bf16 Flex) or NVFP4 QAD student on the six-forward Flex ladder, Ref2VA rows."""
+
+    _transformer_module_type = "transformer_ref"
+    _attention_backend = "FLASH_ATTN"
+    _task = "ref2va"
+
+    def __init__(self,
+                 *,
+                 init_from: str,
+                 training_config: TrainingConfig,
+                 trainable: bool = True,
+                 enable_gradient_checkpointing_type: str | None = None,
+                 trainable_patterns: str = FLEX_TRAINABLE,
+                 qad: dict[str, Any] | None = None,
+                 data: dict[str, Any] | None = None,
+                 frozen_bf16: bool = False) -> None:
+        super().__init__(init_from=init_from,
+                         training_config=training_config,
+                         trainable=trainable,
+                         enable_gradient_checkpointing_type=enable_gradient_checkpointing_type,
+                         trainable_patterns=trainable_patterns,
+                         qad=qad,
+                         data=data,
+                         frozen_bf16=frozen_bf16)
+
+    # ------------------------------------------------------------------ construction hooks
+    def _read_contract(self, model_path: str) -> Any:
+        contract = FlexContract.read(model_path)
+        if contract.task != self._task or contract.trunk != self._transformer_module_type:
+            raise ValueError(f"{model_path} is a Flex {contract.task} merge on {contract.trunk}; "
+                             f"{type(self).__name__} needs {self._task} on {self._transformer_module_type}")
+        return contract
+
+    def _check_tile(self) -> None:
+        """Dense attention: no VSA tile to check."""
+
+    def _install_qad(self) -> None:
+        """NVFP4 W4A4 linears with the calibrated static FFN scales; attention stays BF16 dense."""
+        from fastvideo.layers.quantization.nvfp4_qad import NVFP4QADPlan, install_nvfp4_qad, load_amax_table
+
+        self._install_fp4_attention(dict(self.qad_config.get("fp4_attention") or {}))
+        linears = dict(self.qad_config.get("linears") or {})
+        if not linears.get("enabled", True):
+            return
+        amax_path = linears.get("amax_json")
+        table = load_amax_table(amax_path) if amax_path else None
+        if amax_path:
+            self.amax_sha256 = _pin_amax(amax_path, linears.get("amax_sha256"),
+                                         self.training_config.checkpoint.output_dir)
+        self.provisional_rows = int(linears.get("provisional_calibration_rows", 0))
+        if table is None and self.provisional_rows <= 0:
+            raise ValueError("qad.linears.amax_json (the Flex calibration table) is required")
+        # Flex checkpoints have no compression gate, so ``gate`` has nothing to quantize.
+        plan = NVFP4QADPlan(quantize_attention=bool(linears.get("attention", True)),
+                            quantize_gate=False,
+                            quantize_ffn=bool(linears.get("ffn", True)),
+                            skip_blocks=tuple(int(b) for b in linears.get("skip_blocks", ())))
+        self.qad_linears = install_nvfp4_qad(self.transformer, table, plan, allow_uncalibrated=table is None)
+        logger.info("Flex %s QAD student: %d NVFP4 QAD linears, %s attention, %.2fB trainable parameters", self._task,
+                    len(self.qad_linears), self.fp4_numerics or "BF16 dense", self.num_trainable / 1e9)
+
+    def _install_fp4_attention(self, config: dict[str, Any]) -> None:
+        """Dense FP4 attention as deployed by FlashAttention-4 NVFP4: Q/K in NVFP4, P and V in BF16 by default.
+
+        ``qk_transform`` (``none``/``smooth_k``/``hadamard``/``smooth_k+hadamard``) preconditions Q/K exactly before
+        quantization; ``bf16_blocks`` keeps those blocks' attention in BF16. Deployment must use the same settings.
+        """
+        self.fp4_numerics = None
+        self.fp4_install: dict[str, Any] = {}
+        if not config.get("enabled", False):
+            return
+        from fastvideo.attention.backends.fp4_dense_qat import install_fp4_dense_attention
+        from fastvideo.attention.backends.fp4_vsa_qat import FP4AttentionNumerics
+        from fastvideo.attention.utils.nvfp4_qk_precondition import parse_qk_transform
+
+        floor = config.get("first_block_max_floor")
+        self.fp4_numerics = FP4AttentionNumerics(quantize=True,
+                                                 quantize_pv=bool(config.get("quantize_pv", False)),
+                                                 two_level_p=bool(config.get("two_level_p", False)),
+                                                 smooth_k=bool(config.get("smooth_k", False)),
+                                                 first_block_max_floor=None if floor is None else float(floor))
+        self.fp4_install = {
+            "qk_transform": parse_qk_transform(config.get("qk_transform")),
+            "bf16_blocks": tuple(int(b) for b in config.get("bf16_blocks") or ()),
+        }
+        installed = install_fp4_dense_attention(self.transformer, self.fp4_numerics, **self.fp4_install)
+        if installed == 0:
+            raise ValueError("qad.fp4_attention needs dense (FlashAttention or SDPA) attention layers; none were found")
+        logger.info("Dense FP4 attention emulator installed on %d attention layers (qk_transform=%s, bf16_blocks=%s)",
+                    installed, config.get("qk_transform") or "none", list(self.fp4_install["bf16_blocks"]))
+
+    @contextlib.contextmanager
+    def quantization(self, *, linears: bool = True, attention: bool = True) -> Iterator[None]:
+        """Toggle the NVFP4 linears and, when installed, the dense FP4 attention (all off = the teacher)."""
+        from fastvideo.attention.backends.fp4_dense_qat import install_fp4_dense_attention
+        from fastvideo.layers.quantization.nvfp4_qad import set_nvfp4_qad_enabled
+
+        fp4 = getattr(self, "fp4_numerics", None)
+        set_nvfp4_qad_enabled(self.qad_linears, linears)
+        if fp4 is not None and not attention:
+            install_fp4_dense_attention(self.transformer, None)
+        try:
+            yield
+        finally:
+            set_nvfp4_qad_enabled(self.qad_linears, True)
+            if fp4 is not None:
+                install_fp4_dense_attention(self.transformer, fp4, **self.fp4_install)
+
+    # ------------------------------------------------------------------ Flex ladder
+    @property
+    def plan(self) -> Any:  # pragma: no cover - PDD-only
+        raise AttributeError("Pika H3 Flex has no PDD plan")
+
+    @property
+    def num_rungs(self) -> int:
+        return len(self.contract.denoising_steps)
+
+    def schedulers(self, prepared: PreparedRow | None = None) -> tuple[Any, Any]:
+        """Re-noise schedulers on the shifted Flex grids; their noise is fixed by ``prepared``'s seed."""
+        if prepared is None:
+            raise ValueError("Flex re-noise schedulers need the prepared row (its seed fixes the noise)")
+        seed = int(prepared.spec.seed)
+        out = []
+        for shift, stream in ((self.contract.video_shift, _VIDEO_STREAM), (self.contract.audio_shift, _AUDIO_STREAM)):
+            scheduler = RowRenoiseScheduler(shift, seed, stream)
+            scheduler.set_timesteps(sigmas=flex_sigmas(self.contract.denoising_steps, shift), device=self.device)
+            out.append(scheduler)
+        return out[0], out[1]
+
+    def prepare(self, spec: Any, row: dict[str, Any]) -> PreparedRow:
+        from fastvideo.train.models.minimax_h3.omniref_data import prepare_row
+
+        # Knot 1000 is sigma 1 for both modalities: target rows start as unit noise (no PDD start scaling).
+        return prepare_row(spec,
+                           row,
+                           patch_size=self.patch_size,
+                           latent_channels=24,
+                           audio_channels=32,
+                           pdd_start_sigmas=(1.0, 1.0),
+                           device=self.device)
+
+    @contextlib.contextmanager
+    def rung_scope(self, prepared: PreparedRow, rung: int, timesteps: tuple[Any, Any]) -> Iterator[Any]:
+        """Forward context of Flex forward ``rung``; keep backward inside it (activation recompute reads it)."""
+        from fastvideo.pipelines.basic.minimax_h3.packing import MINIMAX_H3_KEYFRAME_NOISE_AUG, build_row_timesteps
+        video_t, audio_t = float(timesteps[0][rung].item()), float(timesteps[1][rung].item())
+        unique, inverse = build_row_timesteps(prepared.layout,
+                                              video_timestep=video_t,
+                                              audio_timestep=audio_t,
+                                              condition_video_timestep=max(video_t, MINIMAX_H3_KEYFRAME_NOISE_AUG),
+                                              condition_audio_timestep=1.0)
+        with set_forward_context(current_timestep=rung, attn_metadata=None):
+            yield unique.to(self.device), inverse.to(self.device)
+
+
+# Keyframe anchors of each FL2VA task, in packing order (Pika's first/last-frame conditioning).
+FL2VA_ANCHORS: dict[str, tuple[str, ...]] = {
+    "t2va": (),
+    "i2va": ("first", ),
+    "l2va": ("last", ),
+    "fl2va": ("first", "last")
+}
+
+
+def prepare_fl2va_row(spec: Any, row: dict[str, Any], *, patch_size: tuple[int, int, int], latent_channels: int,
+                      audio_channels: int, device: torch.device) -> PreparedRow:
+    """Initial packed FL2VA state: keyframe anchors at the canvas size, noised to the keyframe clean time.
+
+    Mirrors ``MiniMaxH3LatentPreparationStage``: the row's image references are the keyframes (one latent frame at
+    the target's latent size each, clean), packed as ``build_packed_sequence`` condition rows in anchor order and
+    noised with ``keyframe_condition_noise`` before the target's unit noise, all from one row-seeded generator.
+    """
+    from diffusers.utils.torch_utils import randn_tensor
+
+    from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler
+    from fastvideo.pipelines.basic.minimax_h3.packing import (MINIMAX_H3_AUDIO_CHANNELS, MINIMAX_H3_KEYFRAME_NOISE_AUG,
+                                                              audio_latent_num_frames, build_packed_sequence,
+                                                              keyframe_condition_noise, patchify_video_latents,
+                                                              video_latent_num_frames)
+    from fastvideo.train.models.minimax_h3.omniref_data import PreparedRow, _tensor, prepared_references
+
+    anchors = FL2VA_ANCHORS.get(spec.case)
+    if anchors is None:
+        raise ValueError(f"{spec.id}: {spec.case!r} is not an FL2VA task ({sorted(FL2VA_ANCHORS)})")
+    num_frames, height, width = int(row["num_frames"]), int(row["height"]), int(row["width"])
+    frames, lat_h, lat_w = video_latent_num_frames(num_frames), height // 16, width // 16
+    references = prepared_references(row)
+    if len(references) != len(anchors) or any(ref.media_type != "image" or (ref.num_latent_frames, ref.latent_height,
+                                                                            ref.latent_width) != (1, lat_h, lat_w)
+                                              for ref in references):
+        raise ValueError(f"{spec.id}: {spec.case} needs {len(anchors)} canvas-size keyframe(s), got "
+                         f"{[(r.media_type, r.latent_height, r.latent_width) for r in references]}")
+    num_audio = audio_latent_num_frames(num_frames)
+    tags = _tensor(row, "text_token_tags").round().to(torch.long)
+    embeds = _tensor(row, "text_embedding").to(device=device, dtype=torch.bfloat16)[None]
+    layout = build_packed_sequence(tags, frames, lat_h, lat_w, num_audio, patch_size, keyframe_anchors=anchors)
+    generator = torch.Generator("cpu").manual_seed(spec.seed)
+    clean = _tensor(row, "reference_video_latent_rows")
+    condition = clean[:0].to(device)
+    if anchors:
+        noise = keyframe_condition_noise(((1, lat_h, lat_w), ) * len(anchors),
+                                         patch_size,
+                                         latent_channels,
+                                         generator=generator,
+                                         device=device)
+        condition = MiniMaxH3Scheduler().scale_noise(clean.to(device), MINIMAX_H3_KEYFRAME_NOISE_AUG, noise)
+    video_noise = randn_tensor((1, latent_channels, frames, lat_h, lat_w),
+                               generator=generator,
+                               device=device,
+                               dtype=torch.float32)
+    audio = randn_tensor((num_audio * MINIMAX_H3_AUDIO_CHANNELS, audio_channels),
+                         generator=generator,
+                         device=device,
+                         dtype=torch.float32)
+    video = torch.cat((condition, patchify_video_latents(video_noise, patch_size)))
+    if video.shape[0] != layout.video_indices.numel() or audio.shape[0] != layout.audio_indices.numel():
+        raise ValueError(f"{spec.id}: packed rows do not match the FL2VA layout")
+    return PreparedRow(spec, layout, embeds, video, audio, (latent_channels, frames, lat_h, lat_w), references, clean)
+
+
+class MiniMaxH3FlexFL2VAModel(MiniMaxH3FlexRef2VAModel):
+    """The FL2VA adapter on ``transformer``: T2VA / I2VA / L2VA / FL2VA rows (keyframe anchors, not references).
+
+    Rows come from ``build_flex_ref2va_rows.py --mode fl2va`` manifests (``case`` in ``FL2VA_ANCHORS``) and are
+    planned, held out and streamed exactly as the Ref2VA rows (``omniref_data``).
+    """
+
+    _transformer_module_type = "transformer"
+    _task = "fl2va"
+
+    def prepare(self, spec: Any, row: dict[str, Any]) -> PreparedRow:
+        return prepare_fl2va_row(spec,
+                                 row,
+                                 patch_size=self.patch_size,
+                                 latent_channels=24,
+                                 audio_channels=32,
+                                 device=self.device)
+
+
+__all__ = [
+    "FL2VA_ANCHORS", "FLEX_TRAINABLE", "FlexContract", "prepare_fl2va_row", "MiniMaxH3FlexFL2VAModel",
+    "MiniMaxH3FlexRef2VAModel", "RowRenoiseScheduler", "flex_sigmas", "renoise_seed"
+]

fastvideo/train/models/minimax_h3/omniref_data.py

@@ -0,0 +1,295 @@
+# SPDX-License-Identifier: Apache-2.0
+"""OmniRef (FastH3 Ref2VA) rows for per-forward distillation of the PDD student.
+
+A row is one OmniRef training parquet row (precomputed Qwen3-VL presentation,
+clean reference latent rows and their geometry); nothing is re-encoded. The
+initial packed state of a (row, seed) pair is built exactly as
+``generate_omniref_latents.py`` and ``MiniMaxH3LatentPreparationStage`` build
+it (reference rows noise-augmented to 0.999, target video/audio noise from the
+same CPU generator, PDD start scaling), so trajectories, teacher latents and
+Stage A's teacher-forced error are all reproducible from the row id and seed.
+
+Row selection reuses the generator's deterministic plan (per (case,
+resolution): seeded shuffle, then interleave), so the first rows of each group
+are the rows the OmniRef latent generator and Stage A calibration used.
+Held-out sources are excluded at both resolutions.
+"""
+from __future__ import annotations
+
+import hashlib
+import json
+import random
+import re
+from collections import defaultdict
+from collections.abc import Iterator
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+
+OMNIREF_CASES = ("first_frame", "first_last_frame", "storyboard", "continue_scene", "continue_shot",
+                 "full_scene_audio_image_reference")
+RESOLUTIONS = ("480p", "768p")
+_BUCKET_FRAMES = re.compile(r"bucket=\d+x\d+-(\d+)f")
+ROW_COLUMNS = [
+    "id", "vae_latent_shape", "audio_latent_shape", "text_embedding_bytes", "text_embedding_shape",
+    "text_token_tags_bytes", "text_token_tags_shape", "reference_video_latent_rows_bytes",
+    "reference_video_latent_rows_shape", "reference_audio_latent_rows_bytes", "reference_audio_latent_rows_shape",
+    "reference_media_types", "reference_has_audio", "reference_num_latent_frames", "reference_latent_heights",
+    "reference_latent_widths", "reference_num_audio_latents", "width", "height", "num_frames"
+]
+
+
+@dataclass(frozen=True)
+class RowSpec:
+    id: str
+    source: str
+    case: str
+    resolution: str
+    parquet: str
+    seed: int
+
+
+def clip_seed(base_seed: int, clip_id: str) -> int:
+    digest = hashlib.sha256(f"{base_seed}:{clip_id}".encode()).digest()
+    return int.from_bytes(digest[:4], "little") & 0x7FFFFFFF
+
+
+def load_manifest_groups(paths: list[str], cases: tuple[str, ...],
+                         max_frames: int) -> dict[tuple[str, str], list[dict[str, Any]]]:
+    """Manifest entries grouped by (case, resolution); several manifests (v4 + v5 additions) merge.
+
+    An entry repeated across manifests is kept once per (id, resolution): the same source at 480p and 768p is
+    two rows.
+    """
+    groups: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list)
+    seen: set[tuple[str, str]] = set()
+    for path in paths:
+        with open(path) as handle:
+            for line in handle:
+                entry = json.loads(line)
+                match = _BUCKET_FRAMES.search(entry["parquet"])
+                key = (entry["id"], entry["resolution"])
+                if (match is None or entry["case"] not in cases or entry["resolution"] not in RESOLUTIONS or key in seen
+                        or int(match.group(1)) > max_frames):
+                    continue
+                seen.add(key)
+                groups[(entry["case"], entry["resolution"])].append(entry)
+    return groups
+
+
+def limit_frames(groups: dict[tuple[str, str], list[dict[str, Any]]],
+                 max_frames: int) -> dict[tuple[str, str], list[dict[str, Any]]]:
+    """``groups`` without entries whose bucket is longer than ``max_frames`` (0 = no limit); a new mapping."""
+    if not max_frames:
+        return dict(groups)
+
+    def frames(entry: dict[str, Any]) -> int:
+        match = _BUCKET_FRAMES.search(entry["parquet"])
+        return int(match.group(1)) if match else max_frames + 1
+
+    return {key: [entry for entry in entries if frames(entry) <= max_frames] for key, entries in groups.items()}
+
+
+def select_rows(groups: dict[tuple[str, str], list[dict[str, Any]]],
+                seed: int,
+                per_group: dict[str, int],
+                exclude_sources: frozenset[str] = frozenset()) -> list[RowSpec]:
+    """The generator's plan: per group a seeded shuffle, minus excluded sources, then interleaved."""
+    picked = []
+    for key in sorted(groups):
+        entries = sorted(groups[key], key=lambda entry: entry["id"])
+        random.Random(f"{seed}:{key[0]}:{key[1]}").shuffle(entries)
+        entries = [entry for entry in entries if entry["id"].split(":", 1)[-1] not in exclude_sources]
+        picked.append(entries[:int(per_group.get(key[1], 0))])
+    plan = []
+    for index in range(max((len(entries) for entries in picked), default=0)):
+        for entries in picked:
+            if index < len(entries):
+                entry = entries[index]
+                source = entry["id"].split(":", 1)[-1]
+                clip_id = f"omniref-{entry['case']}-{source}-{entry['resolution']}"
+                plan.append(
+                    RowSpec(clip_id, source, entry["case"], entry["resolution"], entry["parquet"],
+                            clip_seed(seed, clip_id)))
+    return plan
+
+
+def load_eval_rows(manifest_json: str,
+                   cases: tuple[str, ...],
+                   resolutions: tuple[str, ...],
+                   per_group: int | dict[str, int],
+                   max_frames: int = 0) -> list[RowSpec]:
+    """Held-out rows from a ``generate_omniref_latents.py --eval-dir`` manifest (ids, parquets and seeds)."""
+    items = json.loads(Path(manifest_json).read_text())
+    counts: dict[tuple[str, str], int] = defaultdict(int)
+    rows = []
+    for item in sorted(items, key=lambda item: item["id"]):
+        key = (item["case"], item["resolution"])
+        limit = int(per_group.get(item["resolution"], 0)) if isinstance(per_group, dict) else int(per_group)
+        if item["case"] not in cases or item["resolution"] not in resolutions or counts[key] >= limit:
+            continue
+        if max_frames and int(item.get("num_frames", 0)) > max_frames:
+            continue
+        counts[key] += 1
+        rows.append(
+            RowSpec(item["id"], item["source"], item["case"], item["resolution"], item["parquet"], int(item["seed"])))
+    return rows
+
+
+def top_up_eval_rows(eval_rows: list[RowSpec], groups: dict[tuple[str, str], list[dict[str, Any]]], seed: int,
+                     per_group: dict[str, int]) -> list[RowSpec]:
+    """Fill each (case, resolution) up to ``per_group`` with fresh held-out rows (call before the training plan).
+
+    Extra rows come from the generator's plan rule at ``seed``, skipping sources already in ``eval_rows``;
+    the training plan must then exclude every eval source. Result is interleaved across groups.
+    """
+    have: dict[tuple[str, str], int] = defaultdict(int)
+    for row in eval_rows:
+        have[(row.case, row.resolution)] += 1
+    taken = frozenset(row.source for row in eval_rows)
+    extra = []
+    for key in sorted(groups):
+        need = int(per_group.get(key[1], 0)) - have[key]
+        if need > 0:
+            extra += select_rows({key: groups[key]}, seed, {key[1]: need}, taken)
+    return interleave_by_group(eval_rows + extra)
+
+
+def interleave_by_group(rows: list[RowSpec]) -> list[RowSpec]:
+    """Round-robin over (case, resolution) groups, so any prefix covers every case."""
+    buckets: dict[tuple[str, str], list[RowSpec]] = defaultdict(list)
+    for row in rows:
+        buckets[(row.case, row.resolution)].append(row)
+    ordered = [buckets[key] for key in sorted(buckets)]
+    return [
+        bucket[i] for i in range(max((len(b) for b in ordered), default=0)) for bucket in ordered if i < len(bucket)
+    ]
+
+
+class RowStream:
+    """Infinite, deterministic per-data-parallel-rank row stream (every SP rank of a group sees the same row)."""
+
+    def __init__(self, plan: list[RowSpec], *, dp_rank: int, dp_size: int, seed: int) -> None:
+        if not plan:
+            raise ValueError("the training row plan is empty")
+        self.plan = plan
+        self.dp_rank = dp_rank
+        self.dp_size = dp_size
+        self.seed = seed
+        self.position = 0
+
+    def __iter__(self) -> Iterator[dict[str, Any]]:
+        while True:
+            epoch, offset = divmod(self.position * self.dp_size + self.dp_rank, len(self.plan))
+            order = list(range(len(self.plan)))
+            random.Random(f"{self.seed}:epoch{epoch}").shuffle(order)
+            self.position += 1
+            yield {"row": self.plan[order[offset]], "epoch": epoch}
+
+    def state_dict(self) -> dict[str, int]:
+        return {"position": self.position}
+
+    def load_state_dict(self, state: dict[str, int]) -> None:
+        self.position = int(state["position"])
+
+
+# --------------------------------------------------------------------------- one row -> initial packed state
+def _tensor(row: dict[str, Any], name: str) -> torch.Tensor:
+    shape = [int(value) for value in row[f"{name}_shape"]]
+    array = np.frombuffer(row[f"{name}_bytes"], dtype=np.float32)
+    if array.size != int(np.prod(shape, dtype=np.int64)):
+        raise ValueError(f"{name}: {array.size} values do not fill shape {shape}")
+    return torch.from_numpy(array.reshape(shape).copy())
+
+
+def read_row(parquet: str) -> dict[str, Any]:
+    import pyarrow.parquet as pq
+
+    rows = pq.read_table(parquet, columns=ROW_COLUMNS).to_pylist()
+    if len(rows) != 1:
+        raise ValueError(f"{parquet}: expected one row, got {len(rows)}")
+    return rows[0]
+
+
+def prepared_references(row: dict[str, Any]) -> list[Any]:
+    from fastvideo.pipelines.basic.minimax_h3.reference import MiniMaxH3PreparedReference
+
+    fields = ("reference_media_types", "reference_has_audio", "reference_num_latent_frames", "reference_latent_heights",
+              "reference_latent_widths", "reference_num_audio_latents")
+    return [
+        MiniMaxH3PreparedReference(media_type=media,
+                                   has_audio=bool(audio),
+                                   num_latent_frames=int(frames),
+                                   latent_height=int(height),
+                                   latent_width=int(width),
+                                   num_audio_latents=int(count))
+        for media, audio, frames, height, width, count in zip(*(row[name] for name in fields), strict=True)
+    ]
+
+
+@dataclass
+class PreparedRow:
+    """Everything one PDD trajectory of a row needs; ``video``/``audio`` are the packed initial rows (fp32)."""
+    spec: RowSpec
+    layout: Any
+    prompt_embeds: torch.Tensor
+    video: torch.Tensor
+    audio: torch.Tensor
+    latent_shape: tuple[int, int, int, int]  # (channels, frames, height, width) of the target video
+    references: list[Any]
+    reference_video_rows: torch.Tensor  # clean, CPU
+
+
+def prepare_row(spec: RowSpec, row: dict[str, Any], *, patch_size: tuple[int, int, int], latent_channels: int,
+                audio_channels: int, pdd_start_sigmas: tuple[float, float], device: torch.device) -> PreparedRow:
+    """Initial packed state, bit-identical to ``OmniRefLatentGenerator._batch`` + ``scale_pdd_initial_noise``."""
+    from diffusers.utils.torch_utils import randn_tensor
+
+    from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler
+    from fastvideo.pipelines.basic.minimax_h3.packing import (MINIMAX_H3_AUDIO_CHANNELS, MINIMAX_H3_KEYFRAME_NOISE_AUG,
+                                                              audio_latent_num_frames, build_ref2va_packed_sequence,
+                                                              keyframe_condition_noise, patchify_video_latents,
+                                                              video_latent_num_frames)
+
+    num_frames, height, width = int(row["num_frames"]), int(row["height"]), int(row["width"])
+    frames, lat_h, lat_w = video_latent_num_frames(num_frames), height // 16, width // 16
+    expected = tuple(int(value) for value in row["vae_latent_shape"][1:])
+    if (frames, lat_h, lat_w) != expected or int(row["vae_latent_shape"][0]) != latent_channels:
+        raise ValueError(f"{spec.id}: target geometry {(frames, lat_h, lat_w)} disagrees with {expected}")
+    num_audio = audio_latent_num_frames(num_frames)
+    references = prepared_references(row)
+    tags = _tensor(row, "text_token_tags").round().to(torch.long)
+    embeds = _tensor(row, "text_embedding").to(device=device, dtype=torch.bfloat16)[None]
+    layout = build_ref2va_packed_sequence(tags, references, frames, lat_h, lat_w, num_audio, patch_size)
+
+    generator = torch.Generator("cpu").manual_seed(spec.seed)
+    shapes = tuple(
+        (ref.num_latent_frames, ref.latent_height, ref.latent_width) for ref in references if ref.media_type != "audio")
+    noise = keyframe_condition_noise(shapes, patch_size, latent_channels, generator=generator, device=device)
+    clean = _tensor(row, "reference_video_latent_rows")
+    condition_video = MiniMaxH3Scheduler().scale_noise(clean.to(device), MINIMAX_H3_KEYFRAME_NOISE_AUG, noise)
+    video_noise = randn_tensor((1, latent_channels, frames, lat_h, lat_w),
+                               generator=generator,
+                               device=device,
+                               dtype=torch.float32)
+    audio_rows = randn_tensor((num_audio * MINIMAX_H3_AUDIO_CHANNELS, audio_channels),
+                              generator=generator,
+                              device=device,
+                              dtype=torch.float32)
+    condition_audio = _tensor(row, "reference_audio_latent_rows").to(device)
+    video = torch.cat((condition_video, patchify_video_latents(video_noise, patch_size)))
+    audio = torch.cat((condition_audio, audio_rows))
+    if video.shape[0] != layout.video_indices.numel() or audio.shape[0] != layout.audio_indices.numel():
+        raise ValueError(f"{spec.id}: packed rows do not match the Ref2VA layout")
+    video[layout.num_condition_video_rows:] *= pdd_start_sigmas[0]
+    audio[layout.num_condition_audio_rows:] *= pdd_start_sigmas[1]
+    return PreparedRow(spec, layout, embeds, video, audio, (latent_channels, frames, lat_h, lat_w), references, clean)
+
+
+__all__ = [
+    "OMNIREF_CASES", "PreparedRow", "RowSpec", "RowStream", "interleave_by_group", "load_eval_rows",
+    "load_manifest_groups", "top_up_eval_rows", "prepare_row", "read_row", "select_rows"
+]

fastvideo/train/models/minimax_h3/omniref_pdd.py

@@ -0,0 +1,406 @@
+# SPDX-License-Identifier: Apache-2.0
+"""FastH3 OmniRef PDD transformer as a per-forward (rung) model for distillation.
+
+One instance wraps ``transformer_ref`` of a FastH3 OmniRef PDD export, loaded
+through the training FSDP loader with VSA-H3 attention at the trained tile
+size. It exposes the eight PDD forwards of the deployed sampler one at a
+time: ``rung_scope`` installs exactly what ``MiniMaxH3DenoisingStage`` sets
+for forward ``k`` (VSA-H3 metadata with the reference keep rate, the fused
+PDD output heads, per-row timesteps), ``forward_rung`` runs the transformer,
+and ``step`` applies the scheduler's Euler step. Teacher and student load
+the same way, so a student whose quantization is disabled is the teacher.
+
+The student may carry NVFP4 QAD linears (``qad.linears``) and the sparse-FP4
+attention emulator (``qad.fp4_attention``); see ``fastvideo.layers.
+quantization.nvfp4_qad`` and ``fastvideo.attention.backends.fp4_vsa_qat``.
+"""
+from __future__ import annotations
+
+import contextlib
+import json
+import re
+from collections.abc import Iterator
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from fastvideo.distributed import get_sp_group, get_world_group
+from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.train.models.base import ModelBase
+from fastvideo.train.utils.activation_checkpoint import apply_activation_checkpointing, resolve_checkpointing_type
+from fastvideo.train.utils.moduleloader import load_module_from_path
+
+if TYPE_CHECKING:
+    from fastvideo.train.models.minimax_h3.omniref_data import PreparedRow
+    from fastvideo.train.utils.training_config import TrainingConfig
+
+logger = init_logger(__name__)
+
+# Block linears and norms; AdaLN, embeddings, the token refiner and the output heads stay frozen.
+DEFAULT_TRAINABLE = (r"transformer_blocks\.\d+\.(attn\.(to_q|to_k|to_v|to_out|to_gate_compress|norm_q|norm_k)"
+                     r"|norm1|norm2|ff\.(fc_in|fc_out))\.")
+INFERENCE_FILE = "fastvideo_inference.json"
+
+
+@dataclass(frozen=True)
+class PDDContract:
+    """The checkpoint's deployed sampling contract (``fastvideo_inference.json``)."""
+    pdd_step_indices: tuple[int, ...]
+    video_shift: float
+    audio_shift: float
+    vsa_sparsity: float
+    vsa_tile_size: int
+    vsa_ref_keep_rate: float | None
+
+    @classmethod
+    def read(cls, model_path: str) -> PDDContract:
+        data = json.loads((Path(model_path) / INFERENCE_FILE).read_text())
+        if data.get("model_type") != "ref2va" or "pdd_step_indices" not in data:
+            raise ValueError(f"{model_path} is not a FastH3 Ref2VA PDD export")
+        return cls(tuple(int(i) for i in data["pdd_step_indices"]), float(data["video_scheduler_shift"]),
+                   float(data["audio_scheduler_shift"]), float(data["vsa_sparsity"]), int(data["vsa_tile_size"]),
+                   data.get("vsa_ref_keep_rate"))
+
+
+class MiniMaxH3OmniRefPDDModel(ModelBase):
+    """Teacher or QAD student for per-forward regression on the OmniRef PDD ladder."""
+
+    _transformer_cls_name = "MiniMaxH3Transformer3DModel"
+    _transformer_module_type = "transformer_ref"
+    _attention_backend = "VIDEO_SPARSE_ATTN_H3"
+
+    def __init__(self,
+                 *,
+                 init_from: str,
+                 training_config: TrainingConfig,
+                 trainable: bool = True,
+                 enable_gradient_checkpointing_type: str | None = None,
+                 trainable_patterns: str = DEFAULT_TRAINABLE,
+                 qad: dict[str, Any] | None = None,
+                 data: dict[str, Any] | None = None,
+                 frozen_bf16: bool = False) -> None:
+        super().__init__(trainable=trainable, attention_backend=self._attention_backend)
+        # Keep a frozen role's masters in bf16 (the checkpoint's dtype) instead of an fp32 upcast: compute is bf16 either
+        # way, so outputs are bit-identical, and the role needs half the memory.
+        self._frozen_bf16 = bool(frozen_bf16)
+        if training_config.pipeline_config is None:
+            raise ValueError("MiniMaxH3OmniRefPDDModel requires a resolved MiniMax-H3 pipeline config")
+        if str(training_config.dit_precision) != "fp32" and trainable:
+            raise ValueError("QAD needs fp32 master weights (training.dit_precision: fp32); BF16 masters "
+                             "round away updates at QAD learning rates")
+        self.contract = self._read_contract(init_from)
+        self._check_tile()
+        # Inference keeps proj_in/out, audio_proj_in/out and the time embedder in FP32; FSDP mixed precision would
+        # compute them in BF16 and the harness would no longer be the deployed operating point (E1: -7.6% v_rel at
+        # forward 0). They are frozen here, so they stay FP32 and replicated outside FSDP.
+        training_config.pipeline_config.dit_config.uniform_parameter_dtype = False  # type: ignore[attr-defined]
+        training_config.pipeline_config.dit_config.replicate_fp32_modules = True  # type: ignore[attr-defined]
+        self._init_from = str(init_from)
+        self.training_config = training_config
+        self.qad_config = dict(qad or {})
+        self.data_config = dict(data or {})
+        with self._load_precision(trainable):
+            self.transformer = load_module_from_path(model_path=self._init_from,
+                                                     module_type=self._transformer_module_type,
+                                                     training_config=training_config,
+                                                     override_transformer_cls_name=self._transformer_cls_name,
+                                                     attention_backend=self.attention_backend)
+        checkpointing = resolve_checkpointing_type(enable_gradient_checkpointing_type, training_config)
+        if trainable and checkpointing:
+            self.transformer = apply_activation_checkpointing(self.transformer, checkpointing_type=checkpointing)
+        self.num_trainable = self._apply_trainable(trainable_patterns if trainable else None)
+        self._check_fp32_modules_frozen()
+        self.patch_size = tuple(int(v) for v in self.transformer.patch_size)
+        self._plan: Any = None
+        self._vsa_builder: Any = None
+        self.qad_linears: dict[str, Any] = {}
+        self.fp4_numerics: Any = None
+        self.provisional_rows = 0
+        self.amax_sha256 = ""
+        self.eval_rows: list[Any] = []
+        self.train_rows: list[Any] = []
+        if self.qad_config.get("enabled", False):
+            self._install_qad()
+        self.dataloader: Any = None
+        self.start_step = 0
+
+    # ------------------------------------------------------------------ setup
+    def _load_precision(self, trainable: bool) -> contextlib.AbstractContextManager[None]:
+        """Context for the transformer load: a frozen role may keep bf16 masters (``frozen_bf16``)."""
+        if trainable or not self._frozen_bf16:
+            return contextlib.nullcontext()
+        return self._bf16_load()
+
+    @contextlib.contextmanager
+    def _bf16_load(self) -> Iterator[None]:
+        config = self.training_config
+        pipeline_config = config.pipeline_config
+        saved = (config.dit_precision, pipeline_config.dit_precision)
+        config.dit_precision = pipeline_config.dit_precision = "bf16"
+        try:
+            yield
+        finally:
+            config.dit_precision, pipeline_config.dit_precision = saved
+
+    def _read_contract(self, model_path: str) -> Any:
+        return PDDContract.read(model_path)
+
+    def _check_tile(self) -> None:
+        if self.contract.vsa_tile_size != 128:
+            raise ValueError(f"the FP4 VSA emulator models tile 128; this export uses {self.contract.vsa_tile_size}")
+
+    def load_row(self, spec: Any) -> PreparedRow:
+        from fastvideo.train.models.minimax_h3.omniref_data import read_row
+        return self.prepare(spec, read_row(spec.parquet))
+
+    def _apply_trainable(self, pattern: str | None) -> int:
+        regex = re.compile(pattern) if pattern else None
+        count = 0
+        for name, parameter in self.transformer.named_parameters():
+            clean = name.replace("_checkpoint_wrapped_module.", "")
+            wanted = regex is not None and regex.search(clean) is not None
+            parameter.requires_grad_(wanted)
+            count += parameter.numel() if wanted else 0
+        self.transformer.train(regex is not None)
+        if regex is not None and count == 0:
+            raise ValueError(f"trainable_patterns {pattern!r} matched no parameters")
+        return count
+
+    def _check_fp32_modules_frozen(self) -> None:
+        """FP32 boundary modules are replicated outside FSDP: training them would desynchronize the ranks."""
+        selector = getattr(self.transformer, "_get_parameter_dtype", None)
+        for name, parameter in self.transformer.named_parameters():
+            clean = name.replace("_checkpoint_wrapped_module.", "")
+            if callable(selector) and selector(clean, torch.bfloat16) == torch.float32:
+                if parameter.requires_grad:
+                    raise ValueError(f"{clean} is an FP32 replicated module and must stay frozen")
+                if parameter.dtype != torch.float32:
+                    raise ValueError(f"{clean} should be FP32 (as inference), got {parameter.dtype}")
+
+    def _install_qad(self) -> None:
+        from fastvideo.attention.backends.fp4_vsa_qat import FP4AttentionNumerics, install_fp4_vsa_attention
+        from fastvideo.layers.quantization.nvfp4_qad import NVFP4QADPlan, install_nvfp4_qad, load_amax_table
+
+        linears = dict(self.qad_config.get("linears") or {})
+        if linears.get("enabled", True):
+            amax_path = linears.get("amax_json")
+            table = load_amax_table(amax_path) if amax_path else None
+            if amax_path:
+                self.amax_sha256 = _pin_amax(amax_path, linears.get("amax_sha256"),
+                                             self.training_config.checkpoint.output_dir)
+            plan = NVFP4QADPlan(quantize_attention=bool(linears.get("attention", True)),
+                                quantize_gate=bool(linears.get("gate", True)),
+                                quantize_ffn=bool(linears.get("ffn", True)),
+                                skip_blocks=tuple(int(b) for b in linears.get("skip_blocks", ())))
+            # Smoke tests only: without Stage A's table, max-calibrate the FFN scales on a few rows at step 0.
+            self.provisional_rows = int(linears.get("provisional_calibration_rows", 0))
+            if table is None and self.provisional_rows <= 0:
+                raise ValueError("qad.linears.amax_json (Stage A calibration) is required")
+            self.qad_linears = install_nvfp4_qad(self.transformer, table, plan, allow_uncalibrated=table is None)
+        attention = dict(self.qad_config.get("fp4_attention") or {})
+        if attention.get("enabled", True):
+            floor = attention.get("first_block_max_floor")
+            self.fp4_numerics = FP4AttentionNumerics(quantize=True,
+                                                     two_level_p=bool(attention.get("two_level_p", False)),
+                                                     smooth_k=bool(attention.get("smooth_k", False)),
+                                                     first_block_max_floor=None if floor is None else float(floor))
+        elif self._trainable:
+            # Tile 128 has no other grad-capable kernel: train through the emulator with exact softmax.
+            self.fp4_numerics = FP4AttentionNumerics(quantize=False)
+        install_fp4_vsa_attention(self.transformer, self.fp4_numerics)
+        logger.info("OmniRef QAD student: %d NVFP4 QAD linears, attention %s, %.2fB trainable parameters",
+                    len(self.qad_linears), self.fp4_numerics, self.num_trainable / 1e9)
+
+    @contextlib.contextmanager
+    def quantization(self, *, linears: bool = True, attention: bool = True) -> Iterator[None]:
+        """Temporarily enable/disable the student's fake quantization (all off = the teacher's numerics)."""
+        from fastvideo.attention.backends.fp4_vsa_qat import install_fp4_vsa_attention
+        from fastvideo.layers.quantization.nvfp4_qad import set_nvfp4_qad_enabled
+
+        set_nvfp4_qad_enabled(self.qad_linears, linears)
+        if not attention:
+            install_fp4_vsa_attention(self.transformer, None)
+        try:
+            yield
+        finally:
+            set_nvfp4_qad_enabled(self.qad_linears, True)
+            install_fp4_vsa_attention(self.transformer, self.fp4_numerics)
+
+    def init_preprocessors(self, training_config: TrainingConfig) -> None:
+        from fastvideo.train.models.minimax_h3.omniref_data import (OMNIREF_CASES, RowStream, interleave_by_group,
+                                                                    limit_frames, load_eval_rows, load_manifest_groups,
+                                                                    select_rows, top_up_eval_rows)
+        cfg = self.data_config
+        cases = tuple(cfg.get("cases", OMNIREF_CASES))
+        resolutions = tuple(cfg.get("eval_resolutions", ("480p", )))
+        eval_per_group = cfg.get("eval_per_group", 2)
+        eval_max_frames = int(cfg.get("eval_max_frames", 0))
+        eval_rows = load_eval_rows(cfg["eval_manifest"], cases, resolutions, eval_per_group, eval_max_frames)
+        groups = load_manifest_groups(list(cfg["manifests"]), cases, int(cfg.get("max_frames", 243)))
+        missing = sorted(set(cases) - {case for case, _ in groups})
+        if missing:
+            logger.warning("OmniRef cases with no rows in the given manifests: %s", missing)
+        if isinstance(eval_per_group, dict):
+            # More held-out rows than the precomputed eval manifest holds: draw them before the training plan,
+            # under the same eval length cap as the manifest rows.
+            eval_groups = limit_frames({
+                key: entries
+                for key, entries in groups.items() if key[1] in resolutions
+            }, eval_max_frames)
+            eval_rows = top_up_eval_rows(eval_rows, eval_groups, int(cfg.get("eval_seed", 20261008)),
+                                         dict(eval_per_group))
+        else:
+            eval_rows = interleave_by_group(eval_rows)
+        plan = select_rows(groups, int(cfg.get("seed", 20261007)), dict(cfg.get("per_group", {"480p": 400})),
+                           frozenset(row.source for row in eval_rows))
+        if int(cfg.get("max_train_rows", 0)) > 0:  # smoke tests: overfit a few rows
+            plan = plan[:int(cfg["max_train_rows"])]
+        world = get_world_group()
+        sp_size = int(training_config.distributed.sp_size or 1)
+        self.eval_rows = eval_rows
+        self.train_rows = plan
+        self.dataloader = RowStream(plan,
+                                    dp_rank=world.rank // sp_size,
+                                    dp_size=world.world_size // sp_size,
+                                    seed=int(training_config.data.seed or 0))
+        logger.info("OmniRef QAD rows: %d train (%s), %d held-out eval", len(plan), dict(sorted(_count(plan).items())),
+                    len(eval_rows))
+
+    # ------------------------------------------------------------------ PDD ladder
+    @property
+    def plan(self) -> Any:
+        if self._plan is None:
+            from fastvideo.layers.pdd import PDDModalitySchedule, build_pdd_sampling_plan
+            schedules = {
+                "video": PDDModalitySchedule(shift=self.contract.video_shift),
+                "audio": PDDModalitySchedule(shift=self.contract.audio_shift)
+            }
+            self._plan = build_pdd_sampling_plan(self.contract.pdd_step_indices, schedules, device=self.device)
+        return self._plan
+
+    @property
+    def num_rungs(self) -> int:
+        return int(self.plan.num_steps)
+
+    def schedulers(self, prepared: PreparedRow | None = None) -> tuple[Any, Any]:
+        """Fresh video/audio schedulers on the PDD node sigmas (as the denoising stage sets them)."""
+        from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler
+        video = MiniMaxH3Scheduler(shift=self.contract.video_shift)
+        audio = MiniMaxH3Scheduler(shift=self.contract.audio_shift)
+        video.set_timesteps(sigmas=self.plan.node_sigmas["video"].to(torch.float32), device=self.device)
+        audio.set_timesteps(sigmas=self.plan.node_sigmas["audio"].to(torch.float32), device=self.device)
+        return video, audio
+
+    def prepare(self, spec: Any, row: dict[str, Any]) -> PreparedRow:
+        from fastvideo.train.models.minimax_h3.omniref_data import prepare_row
+        start = (float(self.plan.node_sigmas["video"][0]), float(self.plan.node_sigmas["audio"][0]))
+        return prepare_row(spec,
+                           row,
+                           patch_size=self.patch_size,
+                           latent_channels=24,
+                           audio_channels=32,
+                           pdd_start_sigmas=start,
+                           device=self.device)
+
+    def _metadata(self, prepared: PreparedRow, rung: int) -> Any:
+        from fastvideo.attention.backends.video_sparse_attn_h3 import MiniMaxH3VSAMetadataBuilder
+        from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_denoising import (_h3_vsa_ref2va_segments,
+                                                                                      _h3_vsa_single_region_segments)
+        if self._vsa_builder is None:
+            self._vsa_builder = MiniMaxH3VSAMetadataBuilder()
+        keep = self.contract.vsa_ref_keep_rate
+        segments = (_h3_vsa_single_region_segments(prepared.layout, self.patch_size)
+                    if keep is None else _h3_vsa_ref2va_segments(prepared.layout, self.patch_size))
+        return self._vsa_builder.build(current_timestep=rung,
+                                       patch_size=self.patch_size,
+                                       VSA_sparsity=self.contract.vsa_sparsity,
+                                       packed_segments=segments,
+                                       device=self.device,
+                                       exempt=True,
+                                       dense_layers=(),
+                                       tile_size=self.contract.vsa_tile_size,
+                                       ref_keep_rate=keep)
+
+    @contextlib.contextmanager
+    def rung_scope(self, prepared: PreparedRow, rung: int, timesteps: tuple[Any, Any]) -> Iterator[Any]:
+        """Forward context of PDD forward ``rung``; keep backward inside it (activation recompute reads it)."""
+        from fastvideo.pipelines.basic.minimax_h3.packing import MINIMAX_H3_KEYFRAME_NOISE_AUG, build_row_timesteps
+        video_t, audio_t = float(timesteps[0][rung].item()), float(timesteps[1][rung].item())
+        unique, inverse = build_row_timesteps(prepared.layout,
+                                              video_timestep=video_t,
+                                              audio_timestep=audio_t,
+                                              condition_video_timestep=max(video_t, MINIMAX_H3_KEYFRAME_NOISE_AUG),
+                                              condition_audio_timestep=1.0)
+        start, end = self.plan.block(rung)
+        fusion = self.transformer.fuse_pdd_block(start, end, self.plan.integration_weights, torch.float32)
+        with fusion, set_forward_context(current_timestep=rung, attn_metadata=self._metadata(prepared, rung)):
+            yield unique.to(self.device), inverse.to(self.device)
+
+    def forward_rung(self, prepared: PreparedRow, video: torch.Tensor, audio: torch.Tensor,
+                     row_timesteps: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
+        """Packed velocities ``([Lv, Cv], [La, Ca])`` for one forward; call inside ``rung_scope``."""
+        layout = prepared.layout
+        device = self.device
+        # No autocast: like inference, the transformer casts its inputs to each weight's dtype (BF16 under
+        # FSDP mixed precision), so Q/K/V reach attention in BF16 exactly as on the deployed path. Autocast
+        # would run RMSNorm/RoPE in fp32 and change what the FP4 quantizers see.
+        video_v, audio_v = self.transformer(hidden_states=video[None],
+                                            audio_hidden_states=audio[None],
+                                            encoder_hidden_states=prepared.prompt_embeds,
+                                            timestep=row_timesteps[0],
+                                            timestep_indices=row_timesteps[1],
+                                            token_tags=layout.token_tags.to(device),
+                                            position_ids=layout.position_ids.to(device),
+                                            video_indices=layout.video_indices.to(device),
+                                            audio_indices=layout.audio_indices.to(device),
+                                            text_indices=layout.text_indices.to(device))
+        return video_v[0], audio_v[0]
+
+    # ------------------------------------------------------------------ ModelBase contract (unused paths)
+    def prepare_batch(self, raw_batch: dict[str, Any], **kwargs: Any) -> Any:  # noqa: D102
+        raise NotImplementedError("use prepare() + rung_scope() + forward_rung()")
+
+    def add_noise(self, clean_latents: torch.Tensor, noise: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError("PDD rungs start from the generator's noise; see prepare()")
+
+    def predict_noise(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError("use forward_rung()")
+
+    def backward(self, loss: torch.Tensor, ctx: Any, *, grad_accum_rounds: int) -> None:
+        raise NotImplementedError("the PDD QAD method backpropagates inside rung_scope()")
+
+    def sp_group(self) -> Any:
+        return get_sp_group()
+
+
+def _pin_amax(path: str, expected_sha256: str | None, output_dir: str) -> str:
+    """Check the calibration table's sha256 (when pinned) and keep a copy next to the run's checkpoints.
+
+    The static FFN scales are fixed buffers that QAD learns around, so a run must never pick up a table that was
+    rewritten in place.
+    """
+    import hashlib
+    import shutil
+
+    digest = hashlib.sha256(Path(path).read_bytes()).hexdigest()
+    if expected_sha256 and digest != expected_sha256:
+        raise ValueError(f"{path} has sha256 {digest}, but the config pins {expected_sha256}")
+    if output_dir and get_world_group().rank == 0:
+        Path(output_dir).mkdir(parents=True, exist_ok=True)
+        shutil.copyfile(path, Path(output_dir) / f"amax-{digest[:12]}.json")
+    logger.info("NVFP4 calibration table %s sha256 %s", path, digest)
+    return digest
+
+
+def _count(rows: list[Any]) -> dict[str, int]:
+    counts: dict[str, int] = {}
+    for row in rows:
+        key = f"{row.case}/{row.resolution}"
+        counts[key] = counts.get(key, 0) + 1
+    return counts
+
+
+__all__ = ["DEFAULT_TRAINABLE", "MiniMaxH3OmniRefPDDModel", "PDDContract"]

fastvideo/train/models/minimax_h3/v2_compression.py

@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Training-side alias of the shared FastH3 token compression (the implementation serves inference too)."""
+from fastvideo.models.dits.minimax_h3_token_compression import (MODES, CompressionSpec, TokenCompression,
+                                                                block_flops_share, compressed_grid, grid_ops,
+                                                                token_reduction)
+
+__all__ = [
+    "MODES", "CompressionSpec", "TokenCompression", "block_flops_share", "compressed_grid", "grid_ops",
+    "token_reduction"
+]

fastvideo/train/models/minimax_h3/v2_data.py

@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Data-free text-to-video rows for FastH3 V2 QAD.
+
+A row is a prompt with a precomputed Qwen3-VL text embedding (``precompute_v2_text.py``) plus a seed. The initial
+packed state is pure noise (target video + audio), drawn from one CPU generator so (row, seed) reproduces the same
+trajectory everywhere; V2 is a plain 8-step DMD ladder, so there is no PDD start scaling and no conditioning rows.
+The result mimics ``omniref_data.PreparedRow`` so the QAD method and metrics run unchanged (``case == "t2v"``
+yields no keyframe metrics).
+"""
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import torch
+
+V2_CASE = "t2v"
+
+
+@dataclass(frozen=True)
+class V2RowSpec:
+    id: str
+    source: str
+    prompt: str
+    embedding: str  # path of the .pt written by precompute_v2_text.py
+    seed: int
+    width: int = 832
+    height: int = 480
+    num_frames: int = 124
+    case: str = V2_CASE
+
+    @property
+    def resolution(self) -> str:
+        return f"{self.height}p"
+
+
+def load_row_specs(manifest: str, width: int, height: int, num_frames: int) -> list[V2RowSpec]:
+    """Rows of a ``build_v2_prompts.py`` manifest (JSONL: id, source, prompt, embedding, seed)."""
+    specs = []
+    for line in Path(manifest).read_text().splitlines():
+        if not line.strip():
+            continue
+        item = json.loads(line)
+        specs.append(
+            V2RowSpec(item["id"], item.get("source", ""), item["prompt"], item["embedding"], int(item["seed"]),
+                      int(item.get("width", width)), int(item.get("height", height)),
+                      int(item.get("num_frames", num_frames))))
+    if not specs:
+        raise ValueError(f"{manifest} has no rows")
+    return specs
+
+
+def read_embedding(path: str) -> tuple[torch.Tensor, torch.Tensor]:
+    data = torch.load(path, map_location="cpu", weights_only=True)
+    embeds, tags = data["embeds"], data["tags"]
+    if embeds.ndim != 2 or tags.shape != (embeds.shape[0], ):
+        raise ValueError(f"{path}: bad embedding/tag shapes {tuple(embeds.shape)} / {tuple(tags.shape)}")
+    return embeds, tags.to(torch.long)
+
+
+def prepare_row(spec: V2RowSpec, *, patch_size: tuple[int, int, int], latent_channels: int, audio_channels: int,
+                device: torch.device) -> Any:
+    """Initial packed state of ``spec`` (fp32 noise rows), laid out exactly as ``MiniMaxH3LatentPreparationStage``."""
+    from diffusers.utils.torch_utils import randn_tensor
+
+    from fastvideo.pipelines.basic.minimax_h3.packing import (MINIMAX_H3_AUDIO_CHANNELS, audio_latent_num_frames,
+                                                              build_packed_sequence, patchify_video_latents,
+                                                              video_latent_num_frames)
+    from fastvideo.train.models.minimax_h3.omniref_data import PreparedRow
+
+    frames, lat_h, lat_w = video_latent_num_frames(spec.num_frames), spec.height // 16, spec.width // 16
+    embeds, tags = read_embedding(spec.embedding)
+    num_audio = audio_latent_num_frames(spec.num_frames)
+    layout = build_packed_sequence(tags, frames, lat_h, lat_w, num_audio, patch_size)
+    generator = torch.Generator("cpu").manual_seed(spec.seed)
+    video_noise = randn_tensor((1, latent_channels, frames, lat_h, lat_w),
+                               generator=generator,
+                               device=device,
+                               dtype=torch.float32)
+    audio = randn_tensor((num_audio * MINIMAX_H3_AUDIO_CHANNELS, audio_channels),
+                         generator=generator,
+                         device=device,
+                         dtype=torch.float32)
+    video = patchify_video_latents(video_noise, patch_size)
+    if video.shape[0] != layout.video_indices.numel() or audio.shape[0] != layout.audio_indices.numel():
+        raise ValueError(f"{spec.id}: packed rows do not match the T2V layout")
+    return PreparedRow(spec, layout,
+                       embeds.to(device=device, dtype=torch.bfloat16)[None], video, audio,
+                       (latent_channels, frames, lat_h, lat_w), [], torch.empty(0, video.shape[-1]))
+
+
+__all__ = ["V2RowSpec", "V2_CASE", "load_row_specs", "prepare_row", "read_embedding"]

fastvideo/train/models/minimax_h3/v2_t2v.py

@@ -0,0 +1,272 @@
+# SPDX-License-Identifier: Apache-2.0
+"""FastH3 V2 (8-step DMD, text-to-video+audio) as a per-forward model for NVFP4 QAD with optional token compression.
+
+Same contract as ``omniref_pdd.MiniMaxH3OmniRefPDDModel`` so ``OmniRefPDDQADMethod`` runs unchanged: eight forwards
+("rungs", one per DMD step), ``rung_scope`` / ``forward_rung`` / ``schedulers`` / ``prepare`` / ``quantization``.
+
+Differences from OmniRef: plain ``transformer`` module (no PDD head fusion), V2 sampling contract (video shift 10,
+audio shift 3, ``dmd_denoising_steps``), VSA tile 64 with sparsity 0.8 (BF16 Triton kernels by default;
+``qad.fp4_attention`` installs the same FP4 emulator as the shared recipe), text-only rows (``v2_data``), and the
+optional middle-block spatial token compression (``v2_compression``).
+
+``quantization(linears=, attention=)``: ``linears`` toggles the NVFP4 QAD linears; ``attention=False`` switches off
+the FP4 attention emulator and the token compression, so ``quantization(False, False)`` is the BF16 teacher.
+"""
+from __future__ import annotations
+
+import contextlib
+import json
+from collections.abc import Iterator
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from fastvideo.distributed import get_sp_world_size, get_world_group, model_parallel_is_initialized
+from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.train.models.minimax_h3.omniref_pdd import DEFAULT_TRAINABLE, MiniMaxH3OmniRefPDDModel, _pin_amax
+from fastvideo.train.models.minimax_h3.v2_compression import (CompressionSpec, TokenCompression, block_flops_share,
+                                                              compressed_grid)
+
+if TYPE_CHECKING:
+    from fastvideo.train.models.minimax_h3.omniref_data import PreparedRow
+    from fastvideo.train.models.minimax_h3.v2_data import V2RowSpec
+    from fastvideo.train.utils.training_config import TrainingConfig
+
+logger = init_logger(__name__)
+INFERENCE_FILE = "fastvideo_inference.json"
+
+
+@dataclass(frozen=True)
+class V2Contract:
+    """The checkpoint's deployed sampling contract (``fastvideo_inference.json``)."""
+    dmd_steps: tuple[int, ...]
+    video_shift: float
+    audio_shift: float
+    vsa_sparsity: float
+    vsa_tile_size: int
+
+    @classmethod
+    def read(cls, model_path: str) -> V2Contract:
+        data = json.loads((Path(model_path) / INFERENCE_FILE).read_text())
+        if "dmd_denoising_steps" not in data or data.get("task") != "t2av":
+            raise ValueError(f"{model_path} is not a FastH3 T2AV DMD export")
+        return cls(tuple(int(i) for i in data["dmd_denoising_steps"]), float(data["video_scheduler_shift"]),
+                   float(data["audio_scheduler_shift"]), float(data["vsa_sparsity"]), int(data["vsa_tile_size"]))
+
+
+class MiniMaxH3V2T2VModel(MiniMaxH3OmniRefPDDModel):
+    """Teacher (bf16 V2) or NVFP4 QAD student (+ token compression) on the 8-step V2 ladder."""
+
+    _transformer_module_type = "transformer"
+
+    def __init__(self,
+                 *,
+                 init_from: str,
+                 training_config: TrainingConfig,
+                 trainable: bool = True,
+                 enable_gradient_checkpointing_type: str | None = None,
+                 trainable_patterns: str = DEFAULT_TRAINABLE,
+                 qad: dict[str, Any] | None = None,
+                 compression: dict[str, Any] | None = None,
+                 frozen_bf16: bool = False,
+                 data: dict[str, Any] | None = None) -> None:
+        self._compression_cfg = dict(compression or {})
+        super().__init__(init_from=init_from,
+                         training_config=training_config,
+                         trainable=trainable,
+                         enable_gradient_checkpointing_type=enable_gradient_checkpointing_type,
+                         trainable_patterns=trainable_patterns,
+                         qad=qad,
+                         data=data,
+                         frozen_bf16=frozen_bf16)
+        self._vsa_builder: Any = self._vsa_builder
+        self.compression: TokenCompression | None = None
+        spec = CompressionSpec.from_config(self._compression_cfg)
+        if spec is not None:
+            if model_parallel_is_initialized() and get_sp_world_size() > 1:
+                raise ValueError("token compression needs sp_size 1")
+            self.compression = TokenCompression(self.transformer, spec)
+            logger.info("V2 token compression: blocks [%d, %d) mode %s", spec.start, spec.end, spec.mode)
+        self._compression_on = True
+
+    # ------------------------------------------------------------------ construction hooks
+    def _read_contract(self, model_path: str) -> Any:
+        return V2Contract.read(model_path)
+
+    def _check_tile(self) -> None:
+        if self.contract.vsa_tile_size != 64:
+            raise ValueError(f"the V2 QAD harness expects VSA tile 64, got {self.contract.vsa_tile_size}")
+
+    def _install_qad(self) -> None:
+        """NVFP4 linears plus the optional FP4 VSA attention emulator.
+
+        With ``qad.fp4_attention.enabled`` the student trains against the same FP4 attention numerics
+        used by the fast inference lane; without it V2 keeps the native bf16 VSA path.
+        """
+        from fastvideo.attention.backends.fp4_vsa_qat import FP4AttentionNumerics, install_fp4_vsa_attention
+        from fastvideo.layers.quantization.nvfp4_qad import NVFP4QADPlan, install_nvfp4_qad, load_amax_table
+
+        linears = dict(self.qad_config.get("linears") or {})
+        if not linears.get("enabled", True):
+            return
+        amax_path = linears.get("amax_json")
+        if not amax_path:
+            raise ValueError("qad.linears.amax_json (the Consumer export's static FFN scales) is required")
+        table = load_amax_table(amax_path)
+        self.amax_sha256 = _pin_amax(amax_path, linears.get("amax_sha256"), self.training_config.checkpoint.output_dir)
+        plan = NVFP4QADPlan(quantize_attention=bool(linears.get("attention", True)),
+                            quantize_gate=bool(linears.get("gate", True)),
+                            quantize_ffn=bool(linears.get("ffn", True)),
+                            skip_blocks=tuple(int(b) for b in linears.get("skip_blocks", ())))
+        self.qad_linears = install_nvfp4_qad(self.transformer, table, plan)
+        self.fp4_numerics = None
+        attention = dict(self.qad_config.get("fp4_attention") or {})
+        if attention.get("enabled", False):
+            floor = attention.get("first_block_max_floor")
+            self.fp4_numerics = FP4AttentionNumerics(quantize=True,
+                                                     two_level_p=bool(attention.get("two_level_p", False)),
+                                                     smooth_k=bool(attention.get("smooth_k", False)),
+                                                     first_block_max_floor=None if floor is None else float(floor))
+            install_fp4_vsa_attention(self.transformer, self.fp4_numerics)
+        logger.info("V2 QAD student: %d NVFP4 QAD linears, attention %s, %.2fB trainable parameters",
+                    len(self.qad_linears), self.fp4_numerics, self.num_trainable / 1e9)
+
+    @contextlib.contextmanager
+    def quantization(self, *, linears: bool = True, attention: bool = True) -> Iterator[None]:
+        from fastvideo.attention.backends.fp4_vsa_qat import install_fp4_vsa_attention
+        from fastvideo.layers.quantization.nvfp4_qad import set_nvfp4_qad_enabled
+
+        set_nvfp4_qad_enabled(self.qad_linears, linears)
+        fp4 = getattr(self, "fp4_numerics", None)
+        if fp4 is not None and not attention:
+            install_fp4_vsa_attention(self.transformer, None)
+        previous = self._compression_on
+        self._compression_on = attention
+        try:
+            yield
+        finally:
+            set_nvfp4_qad_enabled(self.qad_linears, True)
+            if fp4 is not None:
+                install_fp4_vsa_attention(self.transformer, fp4)
+            self._compression_on = previous
+
+    # ------------------------------------------------------------------ data
+    def init_preprocessors(self, training_config: TrainingConfig) -> None:
+        from fastvideo.train.models.minimax_h3.omniref_data import RowStream
+        from fastvideo.train.models.minimax_h3.v2_data import load_row_specs
+
+        cfg = self.data_config
+        size = (int(cfg.get("width", 832)), int(cfg.get("height", 480)), int(cfg.get("num_frames", 124)))
+        eval_rows = load_row_specs(cfg["eval_manifest"], *size)
+        eval_ids = {row.id for row in eval_rows}
+        train_rows = [row for row in load_row_specs(cfg["train_manifest"], *size) if row.id not in eval_ids]
+        if int(cfg.get("max_train_rows", 0)) > 0:
+            train_rows = train_rows[:int(cfg["max_train_rows"])]
+        if int(cfg.get("max_eval_rows", 0)) > 0:
+            eval_rows = eval_rows[:int(cfg["max_eval_rows"])]
+        world = get_world_group()
+        sp_size = int(training_config.distributed.sp_size or 1)
+        self.eval_rows, self.train_rows = eval_rows, train_rows
+        self.dataloader = RowStream(train_rows,
+                                    dp_rank=world.rank // sp_size,
+                                    dp_size=world.world_size // sp_size,
+                                    seed=int(training_config.data.seed or 0))
+        logger.info("V2 QAD rows: %d train, %d held-out eval", len(train_rows), len(eval_rows))
+
+    def load_row(self, spec: V2RowSpec) -> PreparedRow:
+        from fastvideo.train.models.minimax_h3.v2_data import prepare_row
+        return prepare_row(spec, patch_size=self.patch_size, latent_channels=24, audio_channels=32, device=self.device)
+
+    # ------------------------------------------------------------------ DMD ladder
+    @property
+    def plan(self) -> Any:  # pragma: no cover - PDD-only
+        raise AttributeError("V2 has no PDD plan")
+
+    @property
+    def num_rungs(self) -> int:
+        return len(self.contract.dmd_steps)
+
+    def schedulers(self, prepared: PreparedRow | None = None) -> tuple[Any, Any]:
+        """Fresh video/audio schedulers on the DMD sigma grid, as ``MiniMaxH3DenoisingStage._set_dmd_schedule``."""
+        from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler
+        base = torch.tensor([step / 1000.0 for step in self.contract.dmd_steps] + [0.0], dtype=torch.float32)
+        out = []
+        for shift in (self.contract.video_shift, self.contract.audio_shift):
+            scheduler = MiniMaxH3Scheduler(shift=shift)
+            scheduler.set_timesteps(sigmas=shift * base / (1 + (shift - 1) * base), device=self.device)
+            out.append(scheduler)
+        return out[0], out[1]
+
+    def _token_grid(self, prepared: PreparedRow) -> tuple[int, int, int]:
+        layout = self.patch_size
+        return (prepared.layout.num_video_latent_frames // layout[0], prepared.layout.latent_height // layout[1],
+                prepared.layout.latent_width // layout[2])
+
+    def _metadata(self, prepared: PreparedRow, rung: int, grid: tuple[int, int, int] | None = None) -> Any:
+        """VSA-H3 metadata for the full video grid, or for a (compressed) token ``grid``."""
+        from fastvideo.attention.backends.video_sparse_attn_h3 import MiniMaxH3VSAMetadataBuilder
+        from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_denoising import _h3_vsa_single_region_segments
+        if self._vsa_builder is None:
+            self._vsa_builder = MiniMaxH3VSAMetadataBuilder()
+        layout = prepared.layout
+        if grid is None:
+            segments = _h3_vsa_single_region_segments(layout, self.patch_size)
+        else:
+            p = self.patch_size
+            segments = (int(layout.text_indices.numel()), int(layout.audio_indices.numel()),
+                        (grid[0] * p[0], grid[1] * p[1], grid[2] * p[2]))
+        return self._vsa_builder.build(current_timestep=rung,
+                                       patch_size=self.patch_size,
+                                       VSA_sparsity=self.contract.vsa_sparsity,
+                                       packed_segments=segments,
+                                       device=self.device,
+                                       exempt=True,
+                                       dense_layers=(),
+                                       tile_size=self.contract.vsa_tile_size)
+
+    @contextlib.contextmanager
+    def rung_scope(self, prepared: PreparedRow, rung: int, timesteps: tuple[Any, Any]) -> Iterator[Any]:
+        """Forward context of DMD forward ``rung``; keep backward inside it (activation recompute reads it)."""
+        from fastvideo.pipelines.basic.minimax_h3.packing import MINIMAX_H3_KEYFRAME_NOISE_AUG, build_row_timesteps
+        video_t, audio_t = float(timesteps[0][rung].item()), float(timesteps[1][rung].item())
+        unique, inverse = build_row_timesteps(prepared.layout,
+                                              video_timestep=video_t,
+                                              audio_timestep=audio_t,
+                                              condition_video_timestep=max(video_t, MINIMAX_H3_KEYFRAME_NOISE_AUG),
+                                              condition_audio_timestep=1.0)
+        with contextlib.ExitStack() as stack:
+            stack.enter_context(set_forward_context(current_timestep=rung, attn_metadata=self._metadata(prepared,
+                                                                                                        rung)))
+            first_steps = self.compression.spec.first_steps if self.compression is not None else None
+            if self.compression is not None and self._compression_on and (first_steps is None or rung < first_steps):
+                grid = self._token_grid(prepared)
+                sub = compressed_grid(self.compression.spec.mode, grid)
+                stack.enter_context(self.compression.scope(grid, self._metadata(prepared, rung, sub)))
+            yield unique.to(self.device), inverse.to(self.device)
+
+    def set_compression(self, spec: CompressionSpec | None) -> None:
+        """Replace the token-compression spec (probes and sweeps); ``None`` removes it."""
+        if self.compression is not None:
+            self.compression.remove()
+        self.compression = None if spec is None else TokenCompression(self.transformer, spec)
+
+    def compression_summary(self, prepared: PreparedRow) -> dict[str, float]:
+        """Token and FLOP accounting of the configured compression for one row."""
+        spec = self.compression.spec if self.compression else None
+        grid = self._token_grid(prepared)
+        other = prepared.layout.sequence_length - grid[0] * grid[1] * grid[2]
+        out = block_flops_share(spec, grid, len(self.transformer.transformer_blocks), other)
+        out["video_tokens"] = float(grid[0] * grid[1] * grid[2])
+        out["sequence_tokens"] = float(prepared.layout.sequence_length)
+        return out
+
+    # ------------------------------------------------------------------ PDD-only paths of the parent
+    def forward_rung(self, prepared: PreparedRow, video: torch.Tensor, audio: torch.Tensor,
+                     row_timesteps: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
+        return super().forward_rung(prepared, video, audio, row_timesteps)
+
+
+__all__ = ["MiniMaxH3V2T2VModel", "V2Contract"]

fastvideo/train/trainer.py

@@ -139,6 +139,21 @@ class Trainer:
         if self.global_rank == 0 and validation_metrics:
             self.tracker.log(validation_metrics, iteration)
 
+    def _consume_method_requests(self, method: TrainingMethod, checkpoint_manager: Any | None, step: int) -> bool:
+        """Save a method-requested tagged checkpoint and report whether training should stop.
+
+        Runs immediately after every evaluation, including the pre-loop one, so
+        a request always refers to the weights that were just evaluated.
+        """
+        request_hook = getattr(method, "consume_checkpoint_request", None)
+        tag = request_hook() if callable(request_hook) else None
+        if tag and checkpoint_manager is not None:
+            checkpoint_manager.save_tagged(str(tag), step)
+        stop = bool(getattr(method, "stop_requested", False))
+        if stop and self.global_rank == 0:
+            self.tracker.log({"stopped_early_at": float(step)}, step)
+        return stop
+
     def run(
         self,
         method: TrainingMethod,
@@ -186,6 +201,7 @@ class Trainer:
             iteration=start_step,
         )
         self._run_method_validation(method, start_step)
+        stop_early = self._consume_method_requests(method, checkpoint_manager, start_step)
         method.optimizers_zero_grad(start_step)
 
         data_stream = self._iter_dataloader(dataloader,
@@ -210,7 +226,10 @@ class Trainer:
         )
         # Allow method-specific optimization flow (e.g. DiffusionNFT).
         method_manages_optimization = bool(method.manages_optimization())
+        last_step = start_step
         for step in progress:
+            if stop_early:
+                break
             t0 = time.perf_counter()
             if performance_enabled:
                 self.performance_monitor.reset()
@@ -332,14 +351,20 @@ class Trainer:
                 method,
                 iteration=step,
             )
+            last_step = step
+
+            if self._consume_method_requests(method, checkpoint_manager, step):
+                break
+        else:
+            last_step = max_steps
 
         self.callbacks.on_train_end(
             method,
-            iteration=max_steps,
+            iteration=last_step,
         )
 
         if checkpoint_manager is not None:
-            checkpoint_manager.save_final(max_steps)
+            checkpoint_manager.save_final(last_step)
 
         self.performance_monitor.close()
         self.tracker.finish()

fastvideo/train/utils/checkpoint.py

@@ -98,6 +98,17 @@ def _rank() -> int:
     return 0
 
 
+_CPU_GROUP: Any = None
+
+
+def _cpu_group() -> Any:
+    """A gloo group over all ranks (created once, collectively) for checkpoint coordination."""
+    global _CPU_GROUP
+    if _CPU_GROUP is None and dist.is_available() and dist.is_initialized():
+        _CPU_GROUP = dist.new_group(backend="gloo")
+    return _CPU_GROUP
+
+
 def _barrier() -> None:
     if dist.is_available() and dist.is_initialized():
         dist.barrier()
@@ -773,12 +784,43 @@ class CheckpointManager:
             shutil.rmtree(path)
             snapshots = {node: saved for node, saved in snapshots.items() if node != path and path not in node.parents}
 
+    def save_tagged(self, tag: str, step: int) -> None:
+        """Save the trainable roles' model weights only to ``<output_dir>/<tag>`` (overwritten).
+
+        Tagged directories are not ``checkpoint-<step>`` and are never pruned;
+        they hold no optimizer state (resume from regular checkpoints).
+        """
+        # A leading alphanumeric excludes "." and ".." (the run directory and its parent, which rmtree would delete).
+        if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.-]*", tag) or _CHECKPOINT_DIR_RE.match(tag):
+            raise ValueError(f"invalid checkpoint tag {tag!r}")
+        tagged_dir = Path(self.output_dir) / tag
+        if tagged_dir.resolve().parent != Path(self.output_dir).resolve():
+            raise ValueError(f"checkpoint tag {tag!r} does not name a directory directly inside {self.output_dir}")
+        dcp_dir = tagged_dir / "dcp"
+        if _rank() == 0:
+            if tagged_dir.exists():
+                shutil.rmtree(tagged_dir)
+            os.makedirs(dcp_dir, exist_ok=True)
+            self._write_metadata(tagged_dir, step, weights_only=True)
+            logger.info("Saving %s weights (step %d) to %s", tag, step, tagged_dir)
+        _barrier()
+        states = {key: value for key, value in self.method.checkpoint_state().items() if key.startswith("roles.")}
+        # Coordinate over a CPU group: DCP's metadata gather over NCCL needs device memory, and right after an
+        # evaluation the GPUs can be nearly full (seen as "NCCL Error 1: unhandled cuda error" in gather).
+        torch.cuda.empty_cache()
+        dcp.save(states, checkpoint_id=str(dcp_dir), process_group=_cpu_group())
+        _barrier()
+
     def _write_metadata(
         self,
         checkpoint_dir: Path,
         step: int,
+        *,
+        weights_only: bool = False,
     ) -> None:
         metadata: dict[str, Any] = {"step": step}
+        if weights_only:
+            metadata["weights_only"] = True
         if self._raw_config is not None:
             metadata["config"] = self._raw_config
         meta_path = checkpoint_dir / "metadata.json"

scripts/distill/minimax_h3_omniref_qad/generate_omniref_latents.py

@@ -0,0 +1,544 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Generate FastH3 OmniRef (Ref2VA) DiT latents from precomputed conditioning rows.
+
+Training and evaluation data for reference-conditioned FastH3 students. Each input
+row is one OmniRef Ref2VA training parquet row (``minimax_h3_ref2va`` record
+schema): the Qwen3-VL presentation (``text_embedding`` + ``text_token_tags``),
+the clean normalized reference latent rows (``reference_video_latent_rows`` /
+``reference_audio_latent_rows``) and their ordered geometry. Nothing is
+re-encoded: no text encoder, no reference media, no VAE encode.
+
+Per row this script does what ``MiniMaxH3LatentPreparationStage`` does after its
+encoders (0.999 noise augmentation of the reference rows with the request
+generator, then target video/audio noise, packed with
+``build_ref2va_packed_sequence``), runs the pipeline's own
+``MiniMaxH3DenoisingStage`` (PDD fused blocks, VSA-H3 with the checkpoint's
+reference keep rate), and keeps the target video rows exactly as
+``MiniMaxH3VideoDecodingStage`` would decode them. The saved latent is
+``[24, T, H, W]`` float32 in the VAE's normalized space,
+``(raw - latents_mean) / latents_std``: the DiT's own space, so this equals
+decode-stage ``denormalize_latents`` followed by normalization.
+
+Output rows match ``generate_h3_latents.py`` (``id``, ``vae_latent_bytes/shape/
+dtype``, ``generator``, ``case``, ``width``, ``height``, ``num_frames``, ``seed``)
+and are written atomically, ``FLUSH_EVERY`` rows per file. Re-running skips ids
+already written. Rows are chosen deterministically: per (case, resolution) the
+manifest entries up to ``--max-frames`` are shuffled with ``--seed`` and the first
+``--per-group`` are taken, interleaved across groups so a partial run stays
+balanced.
+
+One process per GPU (each its own single-rank external launcher), e.g.::
+
+    CUDA_VISIBLE_DEVICES=0 MASTER_PORT=29500 python generate_omniref_latents.py \\
+        --model-path /models/fasth3_omniref_composed --manifest /data/omniref/MANIFEST.jsonl \\
+        --output-dir /data/gen-latents/omniref/data \\
+        --per-group 300 --shard 0 --num-shards 4
+"""
+from __future__ import annotations
+
+import argparse
+import fcntl
+import hashlib
+import json
+import os
+import random
+import re
+import time
+from collections import defaultdict
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+
+FLUSH_EVERY = 8
+GENERATOR_NAME = "omniref_step3500"
+CASES = ("first_frame", "first_last_frame", "storyboard", "continue_scene", "continue_shot")
+RESOLUTIONS = ("480p", "768p")
+BUCKET_FRAMES = re.compile(r"bucket=\d+x\d+-(\d+)f")
+ROW_COLUMNS = [
+    "id",
+    "vae_latent_shape",
+    "audio_latent_shape",
+    "text_embedding_bytes",
+    "text_embedding_shape",
+    "text_token_tags_bytes",
+    "text_token_tags_shape",
+    "reference_video_latent_rows_bytes",
+    "reference_video_latent_rows_shape",
+    "reference_audio_latent_rows_bytes",
+    "reference_audio_latent_rows_shape",
+    "reference_media_types",
+    "reference_has_audio",
+    "reference_num_latent_frames",
+    "reference_latent_heights",
+    "reference_latent_widths",
+    "reference_num_audio_latents",
+    "width",
+    "height",
+    "num_frames",
+]
+
+
+# --------------------------------------------------------------------------- plan
+def clip_seed(base_seed: int, clip_id: str) -> int:
+    digest = hashlib.sha256(f"{base_seed}:{clip_id}".encode()).digest()
+    return int.from_bytes(digest[:4], "little") & 0x7FFFFFFF
+
+
+def load_manifest(path: str, cases: list[str], max_frames: int) -> dict[tuple[str, str], list[dict[str, Any]]]:
+    groups: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list)
+    with open(path) as handle:
+        for line in handle:
+            entry = json.loads(line)
+            match = BUCKET_FRAMES.search(entry["parquet"])
+            if match is None or entry["case"] not in cases or entry["resolution"] not in RESOLUTIONS:
+                continue
+            frames = int(match.group(1))
+            if frames <= max_frames:
+                groups[(entry["case"], entry["resolution"])].append({**entry, "frames": frames})
+    return groups
+
+
+def select_clips(groups: dict[tuple[str, str], list[dict[str, Any]]],
+                 seed: int,
+                 per_group: int,
+                 exclude_sources: frozenset[str] = frozenset()) -> list[dict[str, Any]]:
+    """Per group: seeded shuffle, drop excluded source clips, take ``per_group``; interleave the groups."""
+    picked = []
+    for key in sorted(groups):
+        entries = sorted(groups[key], key=lambda entry: entry["id"])
+        random.Random(f"{seed}:{key[0]}:{key[1]}").shuffle(entries)
+        entries = [entry for entry in entries if entry["id"].split(":", 1)[-1] not in exclude_sources]
+        picked.append(entries[:per_group])
+    plan = []
+    for index in range(max((len(entries) for entries in picked), default=0)):
+        for entries in picked:
+            if index < len(entries):
+                entry = entries[index]
+                source_id = entry["id"].split(":", 1)[-1]
+                clip_id = f"omniref-{entry['case']}-{source_id}-{entry['resolution']}"
+                plan.append({
+                    "id": clip_id,
+                    "source": source_id,
+                    "case": entry["case"],
+                    "resolution": entry["resolution"],
+                    "parquet": entry["parquet"],
+                    "seed": clip_seed(seed, clip_id)
+                })
+    return plan
+
+
+def clip_plan(args: argparse.Namespace) -> list[dict[str, Any]]:
+    """The full deterministic plan; shard ``s`` takes every ``num_shards``-th entry from ``s``."""
+    groups = load_manifest(args.manifest, args.cases, args.max_frames)
+    if args.eval_dir is None:
+        return select_clips(groups, args.seed, args.per_group)[args.shard::args.num_shards]
+    # A held-out set: no source clip (at either resolution) of the main plan.
+    main_plan = select_clips(groups, args.exclude_seed, args.exclude_per_group)
+    excluded = frozenset(clip["source"] for clip in main_plan)
+    return select_clips(groups, args.seed, args.per_group, excluded)[args.shard::args.num_shards]
+
+
+def existing_ids(data_dir: Path) -> set[str]:
+    import pyarrow.parquet as pq
+
+    done: set[str] = set()
+    for path in data_dir.glob("*.parquet"):
+        done.update(pq.read_table(path, columns=["id"]).column("id").to_pylist())
+    return done
+
+
+def write_part(data_dir: Path, shard: int, rows: list[dict[str, Any]]) -> Path:
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+
+    path = data_dir / f"omniref-shard{shard:03d}-{int(time.time() * 1000)}.parquet"
+    tmp = path.with_suffix(".tmp")
+    pq.write_table(pa.Table.from_pylist(rows), tmp)
+    tmp.replace(path)
+    return path
+
+
+# --------------------------------------------------------------------------- rows
+def _tensor(row: dict[str, Any], name: str) -> torch.Tensor:
+    shape = [int(value) for value in row[f"{name}_shape"]]
+    array = np.frombuffer(row[f"{name}_bytes"], dtype=np.float32)
+    if array.size != int(np.prod(shape, dtype=np.int64)):
+        raise ValueError(f"{name}: {array.size} values do not fill shape {shape}")
+    return torch.from_numpy(array.reshape(shape).copy())
+
+
+def read_row(parquet: str) -> dict[str, Any]:
+    import pyarrow.parquet as pq
+
+    rows = pq.read_table(parquet, columns=ROW_COLUMNS).to_pylist()
+    if len(rows) != 1:
+        raise ValueError(f"{parquet}: expected one row, got {len(rows)}")
+    return rows[0]
+
+
+def prepared_references(row: dict[str, Any]) -> list[Any]:
+    from fastvideo.pipelines.basic.minimax_h3.reference import MiniMaxH3PreparedReference
+
+    fields = ("reference_media_types", "reference_has_audio", "reference_num_latent_frames",
+              "reference_latent_heights", "reference_latent_widths", "reference_num_audio_latents")
+    return [
+        MiniMaxH3PreparedReference(media_type=media_type,
+                                   has_audio=bool(has_audio),
+                                   num_latent_frames=int(frames),
+                                   latent_height=int(height),
+                                   latent_width=int(width),
+                                   num_audio_latents=int(audio))
+        for media_type, has_audio, frames, height, width, audio in zip(*(row[name] for name in fields), strict=True)
+    ]
+
+
+# --------------------------------------------------------------------------- generation
+class OmniRefLatentGenerator:
+    """Single-rank (or SPMD) pipeline driver that injects precomputed Ref2VA conditioning."""
+
+    def __init__(self, args: argparse.Namespace) -> None:
+        from fastvideo import VideoGenerator
+        from fastvideo.api import (ComponentConfig, EngineConfig, GeneratorConfig, OffloadConfig, ParallelismConfig,
+                                   PipelineSelection)
+
+        world_size = int(os.environ.setdefault("WORLD_SIZE", "1"))
+        os.environ.setdefault("RANK", "0")
+        os.environ.setdefault("LOCAL_RANK", "0")
+        os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+        os.environ.setdefault("MASTER_PORT", str(args.master_port))
+        config = GeneratorConfig(
+            model_path=args.model_path,
+            engine=EngineConfig(
+                num_gpus=world_size,
+                execution_backend="external_launcher",
+                use_fsdp_inference=world_size > 1,
+                parallelism=ParallelismConfig(tp_size=1, sp_size=world_size),
+                offload=OffloadConfig(dit=False,
+                                      dit_layerwise=False,
+                                      text_encoder=True,
+                                      vae=True,
+                                      pin_cpu_memory=False,
+                                      lazy_module_load=True),
+            ),
+            pipeline=PipelineSelection(
+                workload_type="i2v",
+                components=ComponentConfig(override_pipeline_cls_name="MiniMaxH3Ref2VAModularPipeline"),
+            ),
+        )
+        self.generator = VideoGenerator.from_config(config)
+        worker = self.generator.executor.worker.worker
+        self.pipeline = worker.pipeline
+        self.fastvideo_args = worker.fastvideo_args
+        self.is_output_rank = bool(self.fastvideo_args.is_output_rank)
+        if not self.pipeline._denoise_stages_ready:
+            self.pipeline._add_denoise_stages(ref2va=True)
+        self.denoise = self.pipeline._stage_name_mapping["denoising_stage"]
+        self.scheduler = self.pipeline.get_module("scheduler")
+
+    def _geometry(self, row: dict[str, Any]) -> tuple[int, int, int, int]:
+        from fastvideo.pipelines.basic.minimax_h3.packing import audio_latent_num_frames, video_latent_num_frames
+
+        arch = self.fastvideo_args.pipeline_config.vae_config.arch_config
+        ratio = int(arch.spatial_compression_ratio)
+        num_frames, height, width = int(row["num_frames"]), int(row["height"]), int(row["width"])
+        geometry = (video_latent_num_frames(num_frames), height // ratio, width // ratio)
+        expected = tuple(int(value) for value in row["vae_latent_shape"][1:])
+        if geometry != expected or int(row["vae_latent_shape"][0]) != int(arch.latent_channels):
+            raise ValueError(f"target latent geometry {geometry} disagrees with the row's {row['vae_latent_shape']}")
+        num_audio = audio_latent_num_frames(num_frames)
+        if num_audio != int(row["audio_latent_shape"][-1]):
+            raise ValueError(f"{num_audio} audio latents disagree with the row's {row['audio_latent_shape']}")
+        return (*geometry, num_audio)
+
+    def _batch(self, row: dict[str, Any], seed: int) -> Any:
+        """The ForwardBatch ``MiniMaxH3LatentPreparationStage`` would hand to denoising."""
+        from diffusers.utils.torch_utils import randn_tensor
+
+        from fastvideo.distributed import get_local_torch_device
+        from fastvideo.pipelines.basic.minimax_h3.packing import (MINIMAX_H3_AUDIO_CHANNELS,
+                                                                  MINIMAX_H3_KEYFRAME_NOISE_AUG,
+                                                                  build_ref2va_packed_sequence, h3_dit_patch_size,
+                                                                  keyframe_condition_noise, patchify_video_latents)
+        from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import MINIMAX_H3_LAYOUT_KEY
+        from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+
+        device = get_local_torch_device()
+        args = self.fastvideo_args
+        patch = h3_dit_patch_size(args)
+        channels = int(args.pipeline_config.vae_config.arch_config.latent_channels)
+        audio_channels = int(args.pipeline_config.audio_vae_config.arch_config.latent_channels)
+        frames, height, width, num_audio = self._geometry(row)
+        references = prepared_references(row)
+        tags = _tensor(row, "text_token_tags").round().to(torch.long)
+        embeds = _tensor(row, "text_embedding").to(device=device, dtype=torch.bfloat16)[None]
+        layout = build_ref2va_packed_sequence(tags, references, frames, height, width, num_audio, patch)
+
+        generator = torch.Generator("cpu").manual_seed(seed)
+        shapes = tuple((ref.num_latent_frames, ref.latent_height, ref.latent_width) for ref in references
+                       if ref.media_type != "audio")
+        noise = keyframe_condition_noise(shapes, patch, channels, generator=generator, device=device)
+        clean = _tensor(row, "reference_video_latent_rows").to(device)
+        condition_video = self.scheduler.scale_noise(clean, MINIMAX_H3_KEYFRAME_NOISE_AUG, noise)
+        video_noise = randn_tensor((1, channels, frames, height, width),
+                                   generator=generator,
+                                   device=device,
+                                   dtype=torch.float32)
+        audio_rows = randn_tensor((num_audio * MINIMAX_H3_AUDIO_CHANNELS, audio_channels),
+                                  generator=generator,
+                                  device=device,
+                                  dtype=torch.float32)
+        condition_audio = _tensor(row, "reference_audio_latent_rows").to(device)
+        video_rows = torch.cat((condition_video, patchify_video_latents(video_noise, patch)))
+        audio_rows = torch.cat((condition_audio, audio_rows))
+        if video_rows.shape[0] != layout.video_indices.numel() or audio_rows.shape[0] != layout.audio_indices.numel():
+            raise ValueError("packed rows do not match the Ref2VA layout")
+
+        batch = ForwardBatch(data_type="video",
+                             prompt="",
+                             seed=seed,
+                             num_frames=int(row["num_frames"]),
+                             height=int(row["height"]),
+                             width=int(row["width"]),
+                             num_inference_steps=len(args.pipeline_config.pdd_step_indices) - 1,
+                             guidance_scale=1.0,
+                             VSA_sparsity=float(args.VSA_sparsity))
+        batch.generator = generator
+        batch.prompt_embeds = [embeds]
+        batch.latents = video_rows
+        batch.audio_latents = audio_rows
+        batch.raw_latent_shape = (1, channels, frames, height, width)
+        batch.extra[MINIMAX_H3_LAYOUT_KEY] = layout
+        return batch
+
+    @torch.no_grad()
+    def generate(self, row: dict[str, Any], seed: int) -> torch.Tensor:
+        """Normalized target-video latent ``[24, T, H, W]`` (float32, CPU)."""
+        from fastvideo.pipelines.basic.minimax_h3.packing import h3_dit_patch_size, unpatchify_video_tokens
+        from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import MINIMAX_H3_LAYOUT_KEY
+
+        batch = self.denoise.forward(self._batch(row, seed), self.fastvideo_args)
+        layout = batch.extra[MINIMAX_H3_LAYOUT_KEY]
+        _, channels, frames, height, width = batch.raw_latent_shape
+        latents = unpatchify_video_tokens(batch.latents[layout.num_condition_video_rows:], frames, height, width,
+                                          channels, h3_dit_patch_size(self.fastvideo_args))
+        return latents[0].float().cpu().contiguous()
+
+    @torch.no_grad()
+    def decode(self, normalized: torch.Tensor) -> np.ndarray:
+        """Decode one normalized ``[24, T, H, W]`` latent to uint8 ``[T, H, W, 3]`` with the original VAE."""
+        from fastvideo.distributed import get_local_torch_device
+
+        from fastvideo.models import pinned_offload
+
+        device = get_local_torch_device()
+        vae = self.pipeline.get_module("vae")
+        pinned_offload.load(vae, device, pin=False)
+        try:
+            latents = vae.denormalize_latents(normalized[None].to(device=device, dtype=torch.float32))
+            output = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32)
+            with torch.autocast(device_type=device.type, dtype=torch.float16):
+                vae.decode_to_pixels(latents, output)
+        finally:
+            pinned_offload.unload(vae)
+        return output[0].mul(255).clamp(0, 255).to(torch.uint8).permute(1, 2, 3, 0).numpy()
+
+    def shutdown(self) -> None:
+        self.generator.shutdown()
+
+
+# --------------------------------------------------------------------------- preview
+def first_reference_latent(row: dict[str, Any], patch: tuple[int, int, int]) -> torch.Tensor | None:
+    from fastvideo.pipelines.basic.minimax_h3.packing import unpatchify_video_tokens
+
+    references = [ref for ref in prepared_references(row) if ref.media_type != "audio"]
+    if not references:
+        return None
+    ref = references[0]
+    count = ref.num_latent_frames * (ref.latent_height // patch[1]) * (ref.latent_width // patch[2])
+    rows = _tensor(row, "reference_video_latent_rows")[:count]
+    return unpatchify_video_tokens(rows, ref.num_latent_frames, ref.latent_height, ref.latent_width, 24, patch)[0]
+
+
+def write_preview(driver: OmniRefLatentGenerator, clip: dict[str, Any], row: dict[str, Any], saved: bytes,
+                  shape: list[int], preview_dir: Path) -> None:
+    """Decode the serialized latent, the dataset target and the first reference for a visual check."""
+    import imageio.v2 as imageio
+    import pyarrow.parquet as pq
+
+    from fastvideo.pipelines.basic.minimax_h3.packing import h3_dit_patch_size
+
+    preview_dir.mkdir(parents=True, exist_ok=True)
+    generated = driver.decode(torch.from_numpy(np.frombuffer(saved, dtype=np.float32).reshape(shape).copy()))
+    target_row = pq.read_table(clip["parquet"], columns=["vae_latent_bytes", "vae_latent_shape"]).to_pylist()[0]
+    target_latent = _tensor(target_row, "vae_latent")
+    target = driver.decode(target_latent)
+    stem = preview_dir / clip["id"]
+    imageio.mimsave(f"{stem}.mp4", list(np.concatenate((generated, target), axis=2)), fps=24, format="mp4")
+    reference = first_reference_latent(row, h3_dit_patch_size(driver.fastvideo_args))
+    if reference is not None:
+        frames = driver.decode(reference)
+        imageio.mimsave(f"{stem}.reference.mp4", list(frames), fps=24, format="mp4")
+    stats = {
+        "id": clip["id"],
+        "generated_mean": float(np.frombuffer(saved, dtype=np.float32).mean()),
+        "generated_std": float(np.frombuffer(saved, dtype=np.float32).std()),
+        "target_mean": float(target_latent.mean()),
+        "target_std": float(target_latent.std()),
+    }
+    print(json.dumps({"preview": str(stem), **stats}), flush=True)
+
+
+# --------------------------------------------------------------------------- main
+def append_eval_manifest(path: Path, entry: dict[str, Any]) -> None:
+    """Append one row under an exclusive lock, rereading so shard processes don't drop each other's rows."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with open(path.with_suffix(".lock"), "a") as lock:
+        fcntl.flock(lock, fcntl.LOCK_EX)
+        try:
+            current = json.loads(path.read_text()) if path.is_file() else []
+            if entry["id"] not in {item["id"] for item in current}:
+                current.append(entry)
+            tmp = path.with_name(f".{path.name}.{os.getpid()}.tmp")
+            tmp.write_text(json.dumps(current, indent=1))
+            tmp.replace(path)
+        finally:
+            fcntl.flock(lock, fcntl.LOCK_UN)
+
+
+def run_eval(args: argparse.Namespace, plan: list[dict[str, Any]]) -> None:
+    """Write raw (denormalized) ``[1, 24, T, H, W]`` latents, as the decode stage decodes them, plus a manifest."""
+    eval_dir = Path(args.eval_dir)
+    eval_dir.mkdir(parents=True, exist_ok=True)
+    manifest_path = eval_dir / "manifest.json"
+    manifest = json.loads(manifest_path.read_text()) if manifest_path.is_file() else []
+    done = {item["id"] for item in manifest}
+    # The model's own vae/config.json: the pipeline config's VAE arch in this process keeps its dataclass
+    # defaults (zeros/ones) until the VAE itself loads.
+    vae_config = json.loads((Path(args.model_path) / "vae" / "config.json").read_text())
+    mean = torch.tensor(vae_config["latents_mean"], dtype=torch.float32).view(1, -1, 1, 1, 1)
+    std = torch.tensor(vae_config["latents_std"], dtype=torch.float32).view(1, -1, 1, 1, 1)
+    driver = OmniRefLatentGenerator(args)
+    for clip in plan:
+        if clip["id"] in done:
+            continue
+        start = time.perf_counter()
+        row = read_row(clip["parquet"])
+        raw = (driver.generate(row, clip["seed"])[None] * std + mean).contiguous()
+        if not bool(torch.isfinite(raw).all()):
+            print(json.dumps({"id": clip["id"], "error": "non-finite latent"}), flush=True)
+            continue
+        name = f"{clip['case']}_{clip['resolution']}_{clip['id']}.pt"
+        if driver.is_output_rank:
+            tmp = eval_dir / f"{name}.tmp"
+            torch.save(raw, tmp)
+            tmp.replace(eval_dir / name)
+            append_eval_manifest(
+                manifest_path, {
+                    "file": name,
+                    **{key: clip[key] for key in ("id", "source", "case", "resolution", "parquet", "seed")},
+                    "width": int(row["width"]),
+                    "height": int(row["height"]),
+                    "num_frames": int(row["num_frames"]),
+                    "shape": list(raw.shape),
+                    "dtype": "float32",
+                    "space": "raw (denormalized; decode-stage input)",
+                    "generator": args.generator_name,
+                })
+        print(json.dumps({"id": clip["id"], "shape": list(raw.shape), "seconds": round(time.perf_counter() - start, 1)}),
+              flush=True)
+    driver.shutdown()
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--model-path", required=True, help="composed FastH3 OmniRef PDD model directory")
+    parser.add_argument("--manifest", required=True, help="OmniRef MANIFEST.jsonl (case, id, parquet, resolution)")
+    parser.add_argument("--output-dir", default=None, help="directory receiving the parquet parts")
+    parser.add_argument("--cases", nargs="+", default=list(CASES))
+    parser.add_argument("--per-group", type=int, default=300, help="clips per (case, resolution)")
+    parser.add_argument("--max-frames", type=int, default=243)
+    parser.add_argument("--seed", type=int, default=20261007)
+    parser.add_argument("--shard", type=int, default=0)
+    parser.add_argument("--num-shards", type=int, default=1)
+    parser.add_argument("--limit", type=int, default=0, help="stop after this many new clips (0: no limit)")
+    parser.add_argument("--generator-name", default=GENERATOR_NAME)
+    parser.add_argument("--master-port", type=int, default=29500)
+    parser.add_argument("--preview-dir", default=None, help="decode the first --preview clips here as MP4")
+    parser.add_argument("--preview", type=int, default=0)
+    parser.add_argument("--eval-dir",
+                        default=None,
+                        help="held-out mode: write raw [1, 24, T, H, W] .pt latents and manifest.json here, from "
+                        "source clips outside the main plan (--exclude-seed / --exclude-per-group)")
+    parser.add_argument("--exclude-seed", type=int, default=20261007)
+    parser.add_argument("--exclude-per-group", type=int, default=300)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    plan = clip_plan(args)
+    if args.eval_dir is not None:
+        print(json.dumps({"eval_clips": len(plan)}), flush=True)
+        run_eval(args, plan)
+        return
+    if args.output_dir is None:
+        raise SystemExit("--output-dir is required unless --eval-dir is given")
+    data_dir = Path(args.output_dir)
+    data_dir.mkdir(parents=True, exist_ok=True)
+    done = existing_ids(data_dir)
+    todo = [clip for clip in plan if clip["id"] not in done]
+    if args.limit:
+        todo = todo[:args.limit]
+    print(json.dumps({"shard": args.shard, "planned": len(plan), "todo": len(todo)}), flush=True)
+    if not todo:
+        return
+    driver = OmniRefLatentGenerator(args)
+    pending: list[dict[str, Any]] = []
+    previews = 0
+    for clip in todo:
+        start = time.perf_counter()
+        try:
+            row = read_row(clip["parquet"])
+            latent = driver.generate(row, clip["seed"])
+            if not bool(torch.isfinite(latent).all()):
+                raise ValueError("non-finite latent")
+        except Exception as error:  # noqa: BLE001 - record and continue; one bad row must not stop the shard
+            print(json.dumps({"id": clip["id"], "error": repr(error)[:500]}), flush=True)
+            continue
+        record = {
+            "id": clip["id"],
+            "vae_latent_bytes": latent.numpy().tobytes(),
+            "vae_latent_shape": list(latent.shape),
+            "vae_latent_dtype": "float32",
+            "generator": args.generator_name,
+            "case": clip["case"],
+            "width": int(row["width"]),
+            "height": int(row["height"]),
+            "num_frames": int(row["num_frames"]),
+            "seed": int(clip["seed"]),
+        }
+        print(json.dumps({
+            "id": clip["id"],
+            "shape": record["vae_latent_shape"],
+            "seconds": round(time.perf_counter() - start, 1)
+        }),
+              flush=True)
+        if not driver.is_output_rank:
+            continue
+        pending.append(record)
+        if args.preview_dir and previews < args.preview:
+            write_preview(driver, clip, row, record["vae_latent_bytes"], record["vae_latent_shape"],
+                          Path(args.preview_dir))
+            previews += 1
+        if len(pending) >= FLUSH_EVERY:
+            write_part(data_dir, args.shard, pending)
+            pending = []
+    if pending and driver.is_output_rank:
+        write_part(data_dir, args.shard, pending)
+    driver.shutdown()
+
+
+if __name__ == "__main__":
+    main()

scripts/distill/minimax_h3_v2_qad/build_v2_prompts.py

@@ -0,0 +1,79 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Build the V2 QAD prompt manifests (train + held-out eval) from prompt sources.
+
+Sources: parquet shards with a ``prompt`` column (the decoder latent generator's outputs; only that column is read)
+and JSONL files with a ``prompt`` field (e.g. ``source.jsonl``, OmniRef captions). Prompts are deduplicated, then a
+seeded shuffle assigns ``--eval-count`` held-out prompts and the rest to training. Each row gets a deterministic seed.
+Embeddings are filled in by ``precompute_v2_text.py`` (rows carry the future embedding path).
+"""
+from __future__ import annotations
+
+import argparse
+import glob
+import hashlib
+import json
+import random
+from pathlib import Path
+
+
+def read_prompts(parquet_globs: list[str], jsonl_globs: list[str]) -> list[str]:
+    import pyarrow.parquet as pq
+
+    found: dict[str, str] = {}
+    for pattern in parquet_globs:
+        for path in sorted(glob.glob(pattern)):
+            if "prompt" in pq.read_schema(path).names:
+                for prompt in pq.read_table(path, columns=["prompt"]).column("prompt").to_pylist():
+                    found.setdefault(" ".join(prompt.split()), path)
+    for pattern in jsonl_globs:
+        for path in sorted(glob.glob(pattern)):
+            with open(path) as handle:
+                for line in handle:
+                    if line.strip():
+                        prompt = json.loads(line).get("prompt")
+                        if isinstance(prompt, str) and prompt.strip():
+                            found.setdefault(" ".join(prompt.split()), path)
+    return sorted(found)
+
+
+def make_rows(prompts: list[str], embed_dir: str, seed: int) -> list[dict]:
+    rows = []
+    for prompt in prompts:
+        digest = hashlib.sha256(prompt.encode()).hexdigest()[:16]
+        rows.append({
+            "id": f"v2-{digest}",
+            "source": "prompt",
+            "prompt": prompt,
+            "embedding": str(Path(embed_dir) / f"{digest}.pt"),
+            "seed": int(hashlib.sha256(f"{seed}:{digest}".encode()).hexdigest()[:8], 16) & 0x7FFFFFFF,
+        })
+    return rows
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--parquet", nargs="*", default=[], help="globs of parquet shards with a prompt column")
+    parser.add_argument("--jsonl", nargs="*", default=[], help="globs of JSONL files with a prompt field")
+    parser.add_argument("--out-dir", required=True)
+    parser.add_argument("--eval-count", type=int, default=30)
+    parser.add_argument("--max-train", type=int, default=0, help="0 = all")
+    parser.add_argument("--seed", type=int, default=20261016)
+    args = parser.parse_args()
+    prompts = read_prompts(args.parquet, args.jsonl)
+    random.Random(args.seed).shuffle(prompts)
+    if len(prompts) <= args.eval_count:
+        raise SystemExit(f"only {len(prompts)} prompts")
+    out = Path(args.out_dir)
+    embed_dir = out / "embeddings"
+    embed_dir.mkdir(parents=True, exist_ok=True)
+    eval_prompts, train_prompts = prompts[:args.eval_count], prompts[args.eval_count:]
+    if args.max_train:
+        train_prompts = train_prompts[:args.max_train]
+    for name, subset in (("eval", eval_prompts), ("train", train_prompts)):
+        rows = make_rows(subset, str(embed_dir), args.seed)
+        (out / f"{name}.jsonl").write_text("".join(json.dumps(row) + "\n" for row in rows))
+        print(json.dumps({"split": name, "rows": len(rows)}))
+
+
+if __name__ == "__main__":
+    main()

scripts/distill/minimax_h3_v2_qad/precompute_v2_text.py

@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Precompute FastH3 text embeddings (Qwen3-VL conditioner, bf16 V2 model) for the V2 QAD manifests.
+
+Builds the V2 inference pipeline once (single GPU), calls its conditioning stage on each prompt and writes
+``{"embeds": [L, D] bf16, "tags": [L] long}`` to every row's ``embedding`` path. Existing files are kept, so reruns and
+sharding (``--shard i --num-shards n``) are safe.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import types
+from pathlib import Path
+
+import torch
+
+
+def _build(config_path: str, model_path: str, shard: int):
+    """One-rank external-launcher generator, configured for latent output only."""
+    from fastvideo.api.compat import generator_config_to_fastvideo_args, normalize_generator_config
+    from fastvideo.api.parser import load_raw_config
+    from fastvideo.entrypoints.video_generator import VideoGenerator
+
+    for name, value in (("WORLD_SIZE", "1"), ("RANK", "0"), ("LOCAL_RANK", "0"), ("MASTER_ADDR", "127.0.0.1")):
+        os.environ.setdefault(name, value)
+    raw = load_raw_config(config_path)
+    generator_raw = dict(raw["generator"], model_path=model_path)
+    generator_raw["engine"] = dict(generator_raw["engine"], execution_backend="external_launcher")
+    fastvideo_args = generator_config_to_fastvideo_args(normalize_generator_config(generator_raw))
+    fastvideo_args.output_type = "latent"
+    generator = VideoGenerator.from_fastvideo_args(fastvideo_args)
+    return generator, generator.executor.worker.worker.pipeline
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", required=True, help="FastH3 run config (generator + request)")
+    parser.add_argument("--model-path", required=True)
+    parser.add_argument("--manifests", nargs="+", required=True)
+    parser.add_argument("--shard", type=int, default=0)
+    parser.add_argument("--num-shards", type=int, default=1)
+    args = parser.parse_args()
+
+    rows = []
+    for path in args.manifests:
+        rows += [json.loads(line) for line in Path(path).read_text().splitlines() if line.strip()]
+    todo = [r for r in rows[args.shard::args.num_shards] if not Path(r["embedding"]).exists()]
+    print(json.dumps({"rows": len(rows), "todo": len(todo)}), flush=True)
+    if not todo:
+        return
+    os.environ.setdefault("MASTER_PORT", str(29800 + args.shard))
+    generator, pipeline = _build(args.config, args.model_path, args.shard)
+    pipeline.post_init()  # idempotent: builds the stages
+    stage = next(s for s in pipeline._stage_name_mapping.values() if type(s).__name__ == "MiniMaxH3ConditioningStage")
+    device = torch.device("cuda")
+    with torch.no_grad():
+        for done, row in enumerate(todo, 1):
+            embeds, tags = stage._encode_fl2va(types.SimpleNamespace(prompt=row["prompt"], extra={}), device)
+            target = Path(row["embedding"])
+            tmp = target.with_suffix(".tmp")
+            torch.save({"embeds": embeds[0].to(torch.bfloat16).cpu(), "tags": tags.cpu()}, tmp)
+            tmp.replace(target)
+            if done % 50 == 0:
+                print(json.dumps({"done": done, "of": len(todo)}), flush=True)
+    generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

scripts/distill/pika_h3_flex_qad/build_flex_ref2va_rows.py

@@ -0,0 +1,662 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pika H3 Flex Ref2VA rows from the OmniRef v4 encoded rows: Flex reference sizing plus the inputs v4 lacks.
+
+The OmniRef v4 rows (the FastH3 OmniRef QAD data) carry every image reference at a 2048 short edge and every video
+reference on the 768p canvas. Pika H3 Flex sizes them by tier (1024 / 480p canvas below 768p, 2048 / 768p canvas at
+768p) and accepts input mixes v4 never encoded. This builds packed Ref2VA rows (``pyarrow_schema_ref2va``) from the
+encoded v4 rows alone; the source media is not needed. Each reference is decoded from its stored latents with the
+H3 VAE / audio VAE, then prepared and re-encoded with the released Ref2VA encode path at Flex's tier sizing:
+
+``retier``             a v4 row's own references and prompt; same case and id as the v4 row, so a manifest listed
+                       before v4 replaces it (``load_manifest_groups`` keeps the first entry per id and resolution).
+``audio_reference``    first-frame rows plus the clip's own soundtrack as a standalone audio reference.
+``mixed_image_video``  continuation rows plus the target's first frame as an image ahead of the preceding video.
+``many_images``        storyboard rows plus image references of other rows of the same asset, 6-9 images.
+``multi_video``        continuation rows plus a sibling clip's preceding video (silent) from the same asset.
+
+Target video and audio latents are copied from the v4 row (QAD regresses on the teacher; they carry geometry).
+
+    build_flex_ref2va_rows.py plan --v4-manifest <v4 MANIFEST.jsonl> --plan-config <QAD yaml> \\
+        --plan-max-frames 345 --eval-manifest <held-out manifest> --output-dir <out>
+    CUDA_VISIBLE_DEVICES=<g> build_flex_ref2va_rows.py encode --model-path <merged Flex> --output-dir <out> \\
+        --shard <g> --num-shards <n>
+    build_flex_ref2va_rows.py finalize --output-dir <out> --eval-manifest <held-out manifest>
+"""
+from __future__ import annotations
+
+import argparse
+import functools
+import hashlib
+import json
+import random
+import re
+from collections import defaultdict
+from pathlib import Path
+from typing import Any
+
+V4_CASES = ("first_frame", "first_last_frame", "storyboard", "continue_scene", "continue_shot")
+IMAGE_CASES = ("first_frame", "first_last_frame", "storyboard")
+CONTINUE_CASES = ("continue_scene", "continue_shot")
+NEW_CASES = ("audio_reference", "mixed_image_video", "many_images", "multi_video")
+MAX_IMAGES = 9
+MAX_VIDEO_SECONDS = 15.0
+FPS = 24
+_ASSET = re.compile(r"(asset-[0-9a-f]{32})")
+_PATCH = (1, 2, 2)
+
+
+# ---------------------------------------------------------------------- planning (CPU)
+def load_v4(manifest: Path) -> dict[tuple[str, str], dict[str, Any]]:
+    """Encoded v4 entries keyed by (id, resolution)."""
+    entries = {}
+    for line in manifest.read_text().splitlines():
+        entry = json.loads(line)
+        entries[(entry["id"], entry["resolution"])] = entry
+    return entries
+
+
+def asset_of(entry_id: str) -> str:
+    match = _ASSET.search(entry_id)
+    if match is None:
+        raise ValueError(f"{entry_id}: no asset id")
+    return match.group(1)
+
+
+def qad_plan_ids(config: str, overrides: list[str]) -> set[tuple[str, str]]:
+    """(id, resolution) of every training and held-out row the QAD config reads from the v4 manifest."""
+    from fastvideo.train.models.minimax_h3.omniref_data import (OMNIREF_CASES, limit_frames, load_eval_rows,
+                                                                load_manifest_groups, select_rows, top_up_eval_rows)
+    from fastvideo.train.utils.config import load_run_config
+
+    cfg = load_run_config(config, overrides=overrides).raw["models"]["student"]["data"]
+    cases = tuple(cfg.get("cases", OMNIREF_CASES))
+    resolutions = tuple(cfg.get("eval_resolutions", ("480p", )))
+    eval_per_group = cfg.get("eval_per_group", 2)
+    groups = load_manifest_groups(list(cfg["manifests"]), cases, int(cfg.get("max_frames", 243)))
+    eval_rows = load_eval_rows(cfg["eval_manifest"], cases, resolutions, eval_per_group,
+                               int(cfg.get("eval_max_frames", 0)))
+    if isinstance(eval_per_group, dict):
+        eval_groups = limit_frames({k: v for k, v in groups.items() if k[1] in resolutions},
+                                   int(cfg.get("eval_max_frames", 0)))
+        eval_rows = top_up_eval_rows(eval_rows, eval_groups, int(cfg.get("eval_seed", 20261008)), dict(eval_per_group))
+    plan = select_rows(groups, int(cfg.get("seed", 20261007)), dict(cfg.get("per_group", {"480p": 400})),
+                       frozenset(row.source for row in eval_rows))
+    return {(v4_id(row.case, row.source), row.resolution) for row in [*plan, *eval_rows]}
+
+
+def v4_id(case: str, source: str) -> str:
+    """The v4 manifest id of (case, source): image cases are ``<case>:<clip>``, continuation ids are the clip."""
+    return f"{case}:{source}" if case in IMAGE_CASES else source
+
+
+def _subject(prompt: str, added: str, *, before: bool = False) -> str:
+    """Add one definition to the prompt's ``subject_definitions:`` paragraph."""
+    head, sep, rest = prompt.partition("subject_definitions:")
+    if not sep:
+        raise ValueError("prompt has no subject_definitions section")
+    body, gap, tail = rest.partition("\n\n")
+    body = f" {added}{body}" if before else f"{body.rstrip()} {added}"
+    return f"{head}{sep}{body}{gap}{tail}"
+
+
+def spec(case: str, base: dict[str, Any], refs: list[dict[str, Any]], prompt_edits: list[tuple[str, bool]],
+         **extra: Any) -> dict[str, Any]:
+    """Row spec: ``base``'s target and prompt, ``refs`` (stored-latent sources), prompt definitions to add."""
+    source = base["id"].split(":", 1)[-1]
+    return {
+        "id": base["id"] if case == "retier" else f"{case}:{source}",
+        "case": base["case"] if case == "retier" else case,
+        "source_case": base["case"],
+        "resolution": base["resolution"],
+        "base_parquet": base["parquet"],
+        "references": refs,
+        "prompt_edits": prompt_edits,
+        **extra,
+    }
+
+
+def own_refs(entry: dict[str, Any], count: int, *, audio: bool = True) -> list[dict[str, Any]]:
+    """``count`` stored references of ``entry`` (their embedded audio kept unless ``audio`` is False)."""
+    return [{"kind": "stored", "parquet": entry["parquet"], "index": i, "audio": audio} for i in range(count)]
+
+
+def plan_rows(v4: dict[tuple[str, str], dict[str, Any]], retier_ids: set[tuple[str, str]], extra: dict[str, int],
+              seed: int, exclude_sources: set[str], ref_count: Any) -> list[dict[str, Any]]:
+    """Every spec: the 480p retier rows plus ``extra[resolution]`` rows per new case (held-out sources excluded).
+
+    ``ref_count(parquet)`` is the number of references of a stored row.
+    """
+    specs = []
+    for key in sorted(retier_ids):
+        if key in v4 and key[1] == "480p":
+            entry = v4[key]
+            specs.append(spec("retier", entry, own_refs(entry, ref_count(entry["parquet"])), []))
+    by_asset: dict[tuple[str, str, str], list[dict[str, Any]]] = defaultdict(list)
+    for (entry_id, res), entry in v4.items():
+        kind = "continue" if entry["case"] in CONTINUE_CASES else "image"
+        by_asset[(asset_of(entry_id), res, kind)].append(entry)
+    sources = {"audio_reference": ("first_frame", ), "mixed_image_video": CONTINUE_CASES,
+               "many_images": ("storyboard", ), "multi_video": CONTINUE_CASES}
+    for case in NEW_CASES:
+        for res, count in sorted(extra.items()):
+            pool = sorted((e for (i, r), e in v4.items()
+                           if r == res and e["case"] in sources[case] and i.split(":", 1)[-1] not in exclude_sources),
+                          key=lambda e: e["id"])
+            rng = random.Random(f"{seed}:{case}:{res}")
+            rng.shuffle(pool)
+            made = 0
+            for entry in pool:
+                if made >= count:
+                    break
+                row = new_case_spec(case, entry, rng, by_asset, ref_count)
+                if row is not None:
+                    specs.append(row)
+                    made += 1
+    return specs
+
+
+def new_case_spec(case: str, entry: dict[str, Any], rng: random.Random,
+                  by_asset: dict[tuple[str, str, str], list[dict[str, Any]]], ref_count: Any) -> dict[str, Any] | None:
+    count = ref_count(entry["parquet"])
+    asset, res = asset_of(entry["id"]), entry["resolution"]
+    if case == "audio_reference":
+        refs = [*own_refs(entry, count), {"kind": "target_audio", "parquet": entry["parquet"]}]
+        return spec(case, entry, refs, [("<Audio 1> is the soundtrack of the target video.", False)])
+    if case == "mixed_image_video":
+        refs = [{"kind": "target_frame", "parquet": entry["parquet"]}, *own_refs(entry, count)]
+        return spec(case, entry, refs, [("<Picture 1> is the exact first frame of the target video.", True)])
+    if case == "many_images":
+        others = [e for e in by_asset[(asset, res, "image")] if e["id"] != entry["id"]]
+        others = rng.sample(others, min(len(others), 12))
+        want = rng.randint(max(6, count + 1), MAX_IMAGES)
+        pool = [(e, i) for e in others for i in range(ref_count(e["parquet"]))]
+        if len(pool) < want - count:
+            return None
+        picks = rng.sample(pool, want - count)
+        refs = [*own_refs(entry, count), *({"kind": "stored", "parquet": e["parquet"], "index": i, "audio": False}
+                                           for e, i in picks)]
+        edits = [(f"<Picture {n}> is another frame from the same source video.", False)
+                 for n in range(count + 1, want + 1)]
+        return spec(case, entry, refs, edits)
+    if case == "multi_video":
+        siblings = [e for e in by_asset[(asset, res, "continue")] if e["id"] != entry["id"]]
+        if not siblings:
+            return None
+        other = rng.choice(siblings)
+        refs = [*own_refs(entry, count), {"kind": "stored", "parquet": other["parquet"], "index": 0, "audio": False}]
+        return spec(case, entry, refs, [("<Video 2> is an earlier shot from the same scene, without sound.", False)])
+    raise ValueError(f"unknown case {case!r}")
+
+
+def apply_prompt_edits(prompt: str, edits: list[tuple[str, bool]]) -> str:
+    for added, before in edits:
+        prompt = _subject(prompt, added, before=before)
+    return prompt
+
+
+# ---------------------------------------------------------------------- FL2VA rows (keyframe anchors)
+# FL2VA task -> (v4 case whose image references are the exact keyframes, reference index per anchor).
+FL2VA_SOURCES = {
+    "t2va": ("first_frame", ()),
+    "i2va": ("first_frame", (0, )),
+    "l2va": ("first_last_frame", (1, )),
+    "fl2va": ("first_last_frame", (0, 1)),
+}
+FL2VA_DEFINITIONS = {
+    "t2va": (),
+    "i2va": ("<Picture 1> is the exact first frame of the target video.", ),
+    "l2va": ("<Picture 1> is the exact last frame of the target video.", ),
+    "fl2va": ("<Picture 1> is the exact first frame of the target video.",
+              "<Picture 2> is the exact last frame of the target video."),
+}
+
+
+def fl2va_prompt(caption: str, case: str) -> str:
+    """The v4 caption with its reference paragraphs replaced by this task's keyframe definitions."""
+    paragraphs = [p for p in caption.split("\n\n") if not p.startswith(("subject_definitions:", "retention_analysis:"))]
+    if FL2VA_DEFINITIONS[case]:
+        paragraphs.insert(0, "subject_definitions: " + " ".join(FL2VA_DEFINITIONS[case]))
+    return "\n\n".join(paragraphs)
+
+
+def plan_fl2va_rows(v4: dict[tuple[str, str], dict[str, Any]], counts: dict[str, int], split: dict[str, float],
+                    seed: int, exclude_sources: set[str]) -> list[dict[str, Any]]:
+    """``counts[task]`` rows per FL2VA task, split over resolutions by ``split``; held-out sources excluded."""
+    specs = []
+    for case, total in sorted(counts.items()):
+        source_case, indices = FL2VA_SOURCES[case]
+        for res, share in sorted(split.items()):
+            pool = sorted((e for (i, r), e in v4.items()
+                           if r == res and e["case"] == source_case and i.split(":", 1)[-1] not in exclude_sources),
+                          key=lambda e: e["id"])
+            random.Random(f"{seed}:{case}:{res}").shuffle(pool)
+            for entry in pool[:round(total * share)]:
+                source = entry["id"].split(":", 1)[-1]
+                specs.append({
+                    "mode": "fl2va",
+                    "id": f"{case}:{source}",
+                    "case": case,
+                    "source_case": source_case,
+                    "resolution": res,
+                    "base_parquet": entry["parquet"],
+                    "keyframes": [{"parquet": entry["parquet"], "index": i} for i in indices],
+                })
+    return specs
+
+
+# ---------------------------------------------------------------------- stored-latent decoding (GPU)
+_COLUMNS = ("caption", "vae_latent_bytes", "vae_latent_shape", "audio_latent_bytes", "audio_latent_shape",
+            "reference_video_latent_rows_bytes", "reference_video_latent_rows_shape",
+            "reference_audio_latent_rows_bytes", "reference_audio_latent_rows_shape", "reference_media_types",
+            "reference_has_audio", "reference_num_latent_frames", "reference_latent_heights",
+            "reference_latent_widths", "reference_num_audio_latents")
+
+
+@functools.cache
+def reference_count(parquet: str) -> int:
+    import pyarrow.parquet as pq
+    return len(pq.ParquetFile(parquet).read(columns=["reference_media_types"]).to_pylist()[0]["reference_media_types"])
+
+
+@functools.lru_cache(maxsize=16)
+def stored_row(parquet: str) -> dict[str, Any]:
+    import pyarrow.parquet as pq
+    return pq.ParquetFile(parquet).read(columns=list(_COLUMNS)).to_pylist()[0]
+
+
+def _array(row: dict[str, Any], name: str) -> Any:
+    import numpy as np
+    import torch
+
+    shape = [int(v) for v in row[f"{name}_shape"]]
+    return torch.from_numpy(np.frombuffer(row[f"{name}_bytes"], dtype=np.float32).reshape(shape).copy())
+
+
+def reference_latents(row: dict[str, Any], index: int) -> tuple[str, Any, Any]:
+    """(media type, normalized visual latents ``[1, 24, T, H, W]`` or None, audio latents ``[2, 32, T]`` or None)."""
+    from fastvideo.pipelines.basic.minimax_h3.packing import unpack_audio_tokens, unpatchify_video_tokens
+
+    types = row["reference_media_types"]
+    visual_offset = audio_offset = 0
+    for i, media in enumerate(types):
+        tokens = 0
+        if media != "audio":
+            frames, height, width = (int(row["reference_num_latent_frames"][i]), int(row["reference_latent_heights"][i]),
+                                     int(row["reference_latent_widths"][i]))
+            tokens = frames * (height // _PATCH[1]) * (width // _PATCH[2])
+        audio_rows = 2 * int(row["reference_num_audio_latents"][i]) if row["reference_has_audio"][i] else 0
+        if i == index:
+            visual = audio = None
+            if tokens:
+                rows = _array(row, "reference_video_latent_rows")[visual_offset:visual_offset + tokens]
+                visual = unpatchify_video_tokens(rows, frames, height, width, 24, _PATCH)
+            if audio_rows:
+                rows = _array(row, "reference_audio_latent_rows")[audio_offset:audio_offset + audio_rows]
+                audio = unpack_audio_tokens(rows, audio_rows // 2)
+            return media, visual, audio
+        visual_offset += tokens
+        audio_offset += audio_rows
+    raise IndexError(f"reference {index} of {len(types)}")
+
+
+def decode_visual(latents: Any, vae: Any) -> Any:
+    """Normalized latents ``[1, 24, T, H, W]`` -> uint8 frames ``[T', H*16, W*16, 3]``.
+
+    Decodes spatially tiled (a 145-frame 768p reference or a 2048-edge image next to the resident Qwen3-VL does not
+    fit untiled); tiling is switched back off so the reference re-encode keeps the released untiled numerics.
+    """
+    import torch
+
+    tiling = vae.use_tiling
+    vae.use_tiling = True
+    try:
+        with torch.no_grad():
+            z = vae.denormalize_latents(latents.to("cuda", torch.float32))
+            sample = vae.decode(z, return_dict=False)[0]
+            frames = (vae.denormalize_pixels(sample).clamp(0, 1)[0].permute(1, 2, 3, 0) * 255).round()
+            frames = frames.to(torch.uint8).cpu().numpy()
+    finally:
+        vae.use_tiling = tiling
+    del z, sample
+    torch.cuda.empty_cache()
+    return frames
+
+
+def decode_audio(latents: Any, audio_vae: Any) -> Any:
+    """Normalized audio latents ``[2, 32, T]`` -> waveform ``[2, samples]`` at the audio VAE's rate."""
+    import torch
+
+    with torch.no_grad():
+        waveform = audio_vae.decode(audio_vae.denormalize_latents(latents.to("cuda", torch.float32)),
+                                    return_dict=False)[0]
+    return waveform.reshape(2, -1).float().cpu()
+
+
+def build_references(row_spec: dict[str, Any], vae: Any, audio_vae: Any) -> list[Any]:
+    """``MiniMaxH3Reference`` list in declared order, every medium decoded from stored latents."""
+    from PIL import Image
+
+    from fastvideo.pipelines.basic.minimax_h3.reference import MiniMaxH3Reference
+
+    rate = int(audio_vae.sampling_rate)
+    refs = []
+    for ref in row_spec["references"]:
+        row = stored_row(ref["parquet"])
+        if ref["kind"] == "target_audio":
+            audio = _array(row, "audio_latent")
+            refs.append(MiniMaxH3Reference(source=decode_audio(audio, audio_vae), media_type="audio", sample_rate=rate))
+        elif ref["kind"] == "target_frame":
+            first = _array(row, "vae_latent")[None, :, :1]
+            refs.append(MiniMaxH3Reference(source=Image.fromarray(decode_visual(first, vae)[0])))
+        else:
+            media, visual, audio = reference_latents(row, int(ref["index"]))
+            if media == "image":
+                refs.append(MiniMaxH3Reference(source=Image.fromarray(decode_visual(visual, vae)[0])))
+            elif media == "video":
+                frames = decode_visual(visual, vae)
+                soundtrack = decode_audio(audio, audio_vae) if audio is not None and ref["audio"] else None
+                refs.append(
+                    MiniMaxH3Reference(source=frames,
+                                       media_type="video",
+                                       fps=float(FPS),
+                                       soundtrack=soundtrack,
+                                       sample_rate=rate if soundtrack is not None else None))
+            else:
+                refs.append(MiniMaxH3Reference(source=decode_audio(audio, audio_vae), media_type="audio",
+                                               sample_rate=rate))
+    return refs
+
+
+def check_limits(references: list[Any]) -> None:
+    from fastvideo.pipelines.basic.minimax_h3.reference import (MINIMAX_H3_MAX_REFERENCE_AUDIOS,
+                                                                MINIMAX_H3_MAX_REFERENCE_IMAGES,
+                                                                MINIMAX_H3_MAX_REFERENCE_VIDEOS,
+                                                                MINIMAX_H3_MAX_REFERENCES)
+    counts: dict[str, int] = defaultdict(int)
+    for ref in references:
+        counts[ref.media_type] += 1
+    if (counts["image"] > MINIMAX_H3_MAX_REFERENCE_IMAGES or counts["video"] > MINIMAX_H3_MAX_REFERENCE_VIDEOS
+            or counts["audio"] > MINIMAX_H3_MAX_REFERENCE_AUDIOS or len(references) > MINIMAX_H3_MAX_REFERENCES):
+        raise ValueError(f"reference counts {dict(counts)} exceed Flex's limits")
+    video_seconds = sum(ref.frames.shape[0] / FPS for ref in references if ref.media_type == "video")
+    if video_seconds > MAX_VIDEO_SECONDS + 1e-6:
+        raise ValueError(f"combined reference video is {video_seconds:.2f} s (> {MAX_VIDEO_SECONDS} s)")
+
+
+_COMPONENTS: dict[str, Any] = {}
+
+
+@functools.cache
+def _encoders(model_path: str) -> tuple[Any, Any, Any]:
+    """``(encoding module, model_index, fastvideo_args)``; every component loads once and stays resident.
+
+    The released encode functions load their component per call; routing ``_load_component`` through a cache keeps
+    the VAE, audio VAE and Qwen3-VL resident across rows without changing their numerics.
+    """
+    from fastvideo.pipelines.preprocess import preprocess_minimax_h3_ref2va as encoding
+    from fastvideo.utils import verify_model_config_and_directory
+
+    encoding._init_single_process_distributed()
+    fastvideo_args, _ = encoding._build_fastvideo_args(Path(model_path))
+    loader = encoding._load_component
+
+    def cached(name: str, path: Path, index: dict[str, Any], args: Any) -> Any:
+        if name not in _COMPONENTS:
+            _COMPONENTS[name] = loader(name, path, index, args)
+        return _COMPONENTS[name]
+
+    encoding._load_component = cached
+    return encoding, verify_model_config_and_directory(model_path), fastvideo_args
+
+
+def encode_spec(row_spec: dict[str, Any], model_path: str) -> dict[str, Any]:
+    """The base v4 row with this spec's references and prompt, conditioning re-encoded at Flex's tier sizing."""
+    import pyarrow.parquet as pq
+
+    from fastvideo.pipelines.basic.minimax_h3.reference import prepare_reference
+    from fastvideo.pipelines.basic.pika_h3_flex.buckets import PIKA_TIERS, tier_for_canvas
+
+    encoding, model_index, fastvideo_args = _encoders(model_path)
+    path = Path(model_path)
+    vae = encoding._load_component("vae", path, model_index, fastvideo_args)
+    audio_vae = encoding._load_component("audio_vae", path, model_index, fastvideo_args)
+    record = pq.ParquetFile(row_spec["base_parquet"]).read().to_pylist()[0]
+    tier = PIKA_TIERS[tier_for_canvas(int(record["height"]), int(record["width"]))]
+    prepared = [
+        prepare_reference(ref,
+                          int(record["num_frames"]),
+                          encoding.AUDIO_SAMPLE_RATE,
+                          reference_short_edge=tier.reference_short_edge,
+                          reference_canvas_short_edge=tier.short_edge,
+                          reference_canvas_max_pixels=tier.max_pixels)
+        for ref in build_references(row_spec, vae, audio_vae)
+    ]
+    check_limits(prepared)
+    prompt = apply_prompt_edits(record["caption"], [tuple(edit) for edit in row_spec["prompt_edits"]])
+    visual = encoding.encode_ref_visual_anchor(prepared, path, model_index, fastvideo_args, _PATCH, augment=False)
+    audio = encoding.encode_ref_audio_anchor(prepared, path, model_index, fastvideo_args)
+    text, tags = encoding.encode_ref2va_conditioning(prompt, prepared, path, model_index, fastvideo_args)
+    metadata = encoding._canonical_prepared_references(prepared)
+    for name, tensor in (("text_embedding", text), ("text_token_tags", tags), ("reference_video_latent_rows", visual),
+                         ("reference_audio_latent_rows", audio)):
+        # Every packed tensor is stored as float32 (the readers decode tags as float32 and round).
+        array = tensor.detach().float().cpu().contiguous()
+        record[f"{name}_bytes"] = array.numpy().tobytes()
+        record[f"{name}_shape"] = list(array.shape)
+        record[f"{name}_dtype"] = "float32"
+    for field, key in (("reference_media_types", "media_type"), ("reference_has_audio", "has_audio"),
+                       ("reference_num_latent_frames", "num_latent_frames"),
+                       ("reference_latent_heights", "latent_height"), ("reference_latent_widths", "latent_width"),
+                       ("reference_num_audio_latents", "num_audio_latents")):
+        record[field] = [ref[key] for ref in metadata]
+    record["id"] = row_spec["id"]
+    record["caption"] = prompt
+    return record
+
+
+def encode_fl2va_conditioning(encoding: Any, prompt: str, keyframes: list[Any], model_path: Path,
+                              model_index: dict[str, Any], fastvideo_args: Any) -> tuple[Any, Any]:
+    """Qwen3-VL presentation of an FL2VA request: ``<Picture k>:`` keyframe vision tokens, then the prompt."""
+    import torch
+
+    stage = encoding.MiniMaxH3ConditioningStage(
+        conditioner=encoding._load_component("text_encoder", model_path, model_index, fastvideo_args),
+        tokenizer=encoding._load_component("tokenizer", model_path, model_index, fastvideo_args),
+        processor=encoding._load_component("processor", model_path, model_index, fastvideo_args),
+        ref2va=False,
+    )
+    batch = encoding.ForwardBatch(data_type="video", prompt=prompt)
+    batch.extra[encoding.MINIMAX_H3_KEYFRAMES_KEY] = list(keyframes)
+    batch = stage.forward(batch, fastvideo_args)
+    text = batch.prompt_embeds[0].squeeze(0).float().cpu().contiguous()
+    tags = batch.extra[encoding.MINIMAX_H3_TEXT_TOKEN_TAGS_KEY].to(dtype=torch.long, device="cpu").contiguous()
+    if text.ndim != 2 or tags.shape != text.shape[:1]:
+        raise ValueError(f"FL2VA conditioning shapes {tuple(text.shape)} / {tuple(tags.shape)} do not align")
+    return text, tags
+
+
+def encode_fl2va_spec(row_spec: dict[str, Any], model_path: str) -> dict[str, Any]:
+    """The base v4 row's target with this task's keyframes (decoded, canvas-cropped, re-encoded) and prompt."""
+    import pyarrow.parquet as pq
+    import torch
+    from PIL import Image
+
+    from fastvideo.pipelines.basic.minimax_h3.packing import prepare_keyframe_image
+    from fastvideo.pipelines.basic.minimax_h3.reference import MiniMaxH3PreparedReference
+
+    encoding, model_index, fastvideo_args = _encoders(model_path)
+    path = Path(model_path)
+    vae = encoding._load_component("vae", path, model_index, fastvideo_args)
+    record = pq.ParquetFile(row_spec["base_parquet"]).read().to_pylist()[0]
+    height, width = int(record["height"]), int(record["width"])
+    images, prepared = [], []
+    for keyframe in row_spec["keyframes"]:
+        media, visual, _ = reference_latents(stored_row(keyframe["parquet"]), int(keyframe["index"]))
+        if media != "image":
+            raise ValueError(f"{keyframe} is a {media} reference, not a keyframe image")
+        image = prepare_keyframe_image(Image.fromarray(decode_visual(visual, vae)[0]), height, width, stretch=False)
+        images.append(image)
+        reference = MiniMaxH3PreparedReference(media_type="image")
+        reference.image = image
+        prepared.append(reference)
+    visual = encoding.encode_ref_visual_anchor(prepared, path, model_index, fastvideo_args, _PATCH, augment=False)
+    prompt = fl2va_prompt(record["caption"], row_spec["case"])
+    text, tags = encode_fl2va_conditioning(encoding, prompt, images, path, model_index, fastvideo_args)
+    metadata = encoding._canonical_prepared_references(prepared)
+    for name, tensor in (("text_embedding", text), ("text_token_tags", tags), ("reference_video_latent_rows", visual),
+                         ("reference_audio_latent_rows", torch.empty(0, 32))):
+        array = tensor.detach().float().cpu().contiguous()
+        record[f"{name}_bytes"] = array.numpy().tobytes()
+        record[f"{name}_shape"] = list(array.shape)
+        record[f"{name}_dtype"] = "float32"
+    for field, key in (("reference_media_types", "media_type"), ("reference_has_audio", "has_audio"),
+                       ("reference_num_latent_frames", "num_latent_frames"),
+                       ("reference_latent_heights", "latent_height"), ("reference_latent_widths", "latent_width"),
+                       ("reference_num_audio_latents", "num_audio_latents")):
+        record[field] = [ref[key] for ref in metadata]
+    record["id"] = row_spec["id"]
+    record["caption"] = prompt
+    return record
+
+
+# ---------------------------------------------------------------------- drivers
+def run_plan(args: argparse.Namespace, overrides: list[str]) -> None:
+    from fastvideo.train.models.minimax_h3.omniref_data import load_eval_rows
+
+    v4 = load_v4(Path(args.v4_manifest))
+    if args.mode == "fl2va":
+        exclude = ({row.source
+                    for row in load_eval_rows(args.eval_manifest, V4_CASES, ("480p", "768p"), 99, 0)}
+                   if args.eval_manifest else set())
+        counts = {item.split("=")[0]: int(item.split("=")[1]) for item in args.fl2va_counts}
+        split = {item.split("=")[0]: float(item.split("=")[1]) for item in args.resolution_split}
+        specs = plan_fl2va_rows(v4, counts, split, args.seed, exclude)
+        _write_plan(Path(args.output_dir), specs)
+        return
+    retier: set[tuple[str, str]] = set()
+    for config in args.plan_config:
+        retier |= qad_plan_ids(config, overrides)
+        for frames in args.plan_max_frames:
+            retier |= qad_plan_ids(config, [*overrides, "--models.student.data.max_frames", str(frames)])
+    extra = {item.split("=")[0]: int(item.split("=")[1]) for item in args.extra_per_group}
+    exclude = {row.source for row in load_eval_rows(args.eval_manifest, V4_CASES, ("480p", "768p"), 99, 0)}
+    specs = plan_rows(v4, retier, extra, args.seed, exclude, reference_count)
+    _write_plan(Path(args.output_dir), specs)
+
+
+def _write_plan(out: Path, specs: list[dict[str, Any]]) -> None:
+    out.mkdir(parents=True, exist_ok=True)
+    (out / "plan.json").write_text(json.dumps(specs))
+    summary: dict[str, int] = defaultdict(int)
+    for row_spec in specs:
+        label = row_spec["case"] if row_spec["case"] in NEW_CASES or row_spec.get("mode") else f"retier/{row_spec['case']}"
+        summary[f"{label}/{row_spec['resolution']}"] += 1
+    print(f"{len(specs)} rows: {dict(sorted(summary.items()))}")
+
+
+def run_encode(args: argparse.Namespace) -> None:
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+
+    from fastvideo.dataset.dataloader.schema import pyarrow_schema_ref2va
+
+    import torch
+
+    out = Path(args.output_dir)
+    specs = json.loads((out / "plan.json").read_text())[args.shard::args.num_shards]
+    skipped = out / f"skipped-shard{args.shard:02d}.jsonl"
+    if args.retry_skipped:
+        if not skipped.exists():
+            return
+        retry = {(e["id"], e["resolution"]) for e in map(json.loads, skipped.read_text().splitlines())}
+        specs = [s for s in specs if (s["id"], s["resolution"]) in retry]
+        skipped.rename(skipped.with_suffix(".retried"))
+    manifest = out / f"MANIFEST-shard{args.shard:02d}.jsonl"
+    done = set()
+    if manifest.exists():
+        done = {(e["id"], e["resolution"]) for e in map(json.loads, manifest.read_text().splitlines())}
+    for index, row_spec in enumerate(specs):
+        if (row_spec["id"], row_spec["resolution"]) in done:
+            continue
+        try:
+            encode = encode_fl2va_spec if row_spec.get("mode") == "fl2va" else encode_spec
+            record = encode(row_spec, args.model_path)
+        except Exception as error:  # one bad source row must not stop the shard; it is reported and skipped
+            print(f"SKIP {row_spec['id']} {row_spec['resolution']}: {type(error).__name__}: {error}", flush=True)
+            torch.cuda.empty_cache()
+            with skipped.open("a") as handle:
+                handle.write(json.dumps({"id": row_spec["id"], "resolution": row_spec["resolution"],
+                                         "error": f"{type(error).__name__}: {error}"}) + "\n")
+            continue
+        bucket = re.search(r"bucket=[^/]+", row_spec["base_parquet"])
+        if bucket is None:
+            raise ValueError(f"{row_spec['base_parquet']} has no bucket directory")
+        digest = hashlib.sha256(f"{row_spec['id']}|{row_spec['resolution']}".encode()).hexdigest()[:24]
+        path = out / "parquet" / row_spec["case"] / bucket.group(0) / f"{digest}.parquet"
+        path.parent.mkdir(parents=True, exist_ok=True)
+        pq.write_table(pa.Table.from_pylist([record], schema=pyarrow_schema_ref2va), path)
+        pq.ParquetFile(path).read()  # round trip
+        entry = {"case": row_spec["case"], "id": row_spec["id"], "parquet": str(path),
+                 "resolution": row_spec["resolution"], "source_case": row_spec["source_case"]}
+        with manifest.open("a") as handle:
+            handle.write(json.dumps(entry) + "\n")
+        print(f"encoded {index + 1}/{len(specs)} {row_spec['case']} {row_spec['resolution']} {row_spec['id']}",
+              flush=True)
+        del record
+        torch.cuda.empty_cache()
+
+
+def run_finalize(args: argparse.Namespace) -> None:
+    """Merge the shard manifests; with ``--eval-manifest``, also write the held-out manifest with retiered rows."""
+    out = Path(args.output_dir)
+    lines = []
+    for path in sorted(out.glob("MANIFEST-shard*.jsonl")):
+        lines += [line for line in path.read_text().splitlines() if line.strip()]
+    (out / "MANIFEST.jsonl").write_text("\n".join(lines) + "\n")
+    skipped = sum(len(p.read_text().splitlines()) for p in out.glob("skipped-shard*.jsonl"))
+    print(f"MANIFEST.jsonl: {len(lines)} rows ({skipped} skipped)")
+    if not args.eval_manifest:
+        # FL2VA holds out rows from its own groups (eval_per_group top-up); the trainer still reads a manifest.
+        (out / "eval-manifest.json").write_text("[]\n")
+    else:
+        encoded = {(e["id"], e["resolution"]): e["parquet"] for e in map(json.loads, lines)}
+        items = json.loads(Path(args.eval_manifest).read_text())
+        replaced = 0
+        for item in items:
+            key = (v4_id(item["case"], item["source"]), item["resolution"])
+            if key in encoded:
+                item["parquet"] = encoded[key]
+                replaced += 1
+        (out / "eval-manifest.json").write_text(json.dumps(items, indent=1))
+        print(f"eval-manifest.json: {replaced} of {len(items)} held-out rows retiered")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("stage", choices=("plan", "encode", "finalize"))
+    parser.add_argument("--output-dir", required=True)
+    parser.add_argument("--v4-manifest", help="OmniRef v4 encoded MANIFEST.jsonl")
+    parser.add_argument("--plan-config", nargs="*", default=[], help="QAD configs whose 480p rows to retier")
+    parser.add_argument("--plan-max-frames", nargs="*", type=int, default=[], help="extra max_frames plans to cover")
+    parser.add_argument("--eval-manifest", help="held-out manifest; its sources are excluded from new cases")
+    parser.add_argument("--extra-per-group", nargs="*", default=["480p=80", "768p=20"])
+    parser.add_argument("--mode", choices=("ref2va", "fl2va"), default="ref2va")
+    parser.add_argument("--fl2va-counts", nargs="*", default=["t2va=400", "i2va=300", "l2va=100", "fl2va=300"],
+                        help="FL2VA rows per task (--mode fl2va)")
+    parser.add_argument("--resolution-split", nargs="*", default=["480p=0.8", "768p=0.2"])
+    parser.add_argument("--seed", type=int, default=20261030)
+    parser.add_argument("--model-path", help="merged Flex Ref2VA model (encode)")
+    parser.add_argument("--shard", type=int, default=0)
+    parser.add_argument("--num-shards", type=int, default=1)
+    parser.add_argument("--retry-skipped", action="store_true", help="encode only this shard's skipped rows")
+    args, overrides = parser.parse_known_args()
+    if args.stage == "plan":
+        run_plan(args, overrides)
+    elif args.stage == "encode":
+        run_encode(args)
+    else:
+        run_finalize(args)
+
+
+if __name__ == "__main__":
+    main()

scripts/distill/pika_h3_flex_qad/calibrate_flex_nvfp4.py

@@ -0,0 +1,375 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Static NVFP4 FFN activation scales for Pika H3 Flex: collect, candidates, MSE, choose.
+
+The QAD student quantizes every block's ``ff.fc_in`` / ``ff.fc_out`` input under a static global scale
+``448 * 6 / amax``. This picks that amax per layer on real Ref2VA (or T2VA) rows walked by the bf16 Flex teacher
+over its six-forward ladder, the same recipe as the FastH3 OmniRef calibration:
+
+1. ``collect`` (torchrun, one launch per tray): forward pre-hooks record each row's input amax per FFN layer,
+   split into target video, target audio and conditioning rows (max over the sequence-parallel group).
+2. ``candidates`` (CPU): per layer, the max over rows plus clipped alternatives (per-row p99 / p90 and
+   0.7 / 0.5 / 0.35 / 0.25 x max); writes the plain-max table too.
+3. ``mse`` (torchrun): on disjoint rows plus the rows that set the maxima, forward hooks compare each
+   candidate's deployed NVFP4 GEMM (``nvfp4_qad.nvfp4_gemm``) with the bf16 output, per group.
+4. ``choose`` (CPU): per layer, the candidate with the lowest all-token relative output MSE (``amax-mse.json``,
+   the OmniRef criterion), and the video/audio-balanced variant (``amax-mse-balanced.json``). The step-0 PTQ
+   score of the QAD harness then picks the table QAD pins, as the OmniRef sweep compared its arms.
+
+Launch ``collect`` / ``mse`` with the QAD config of the adapter (``--config``) and ``--shard i --num-shards n`` so
+each tray takes a disjoint slice; ``torchrun --nproc_per_node 4`` per tray (sequence parallel 4,
+``--training.distributed.num_gpus 4 --training.distributed.hsdp_shard_dim 4``). Unknown ``--dotted.key value`` pairs
+override the config. ``collect`` takes the first rows of the QAD plan (e.g. ``--models.student.data.per_group.480p
+160 --models.student.data.per_group.768p 40``); ``mse`` runs on the config's larger plan and skips collected rows.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import re
+from collections.abc import Iterator
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.distributed as dist
+
+FFN_LAYER = re.compile(r"transformer_blocks\.(\d+)\.ff\.(fc_in|fc_out)$")
+GROUP_NAMES = ("video", "audio", "condition")  # pdd_qad_recon.packed_row_groups ids 0, 1, 2
+MULTIPLIERS = (0.7, 0.5, 0.35, 0.25)
+FP4_RANGE = 448.0 * 6.0
+
+
+def amax_key(module_name: str) -> str | None:
+    """``...transformer_blocks.3.ff.fc_in`` -> ``b3.ff.fc_in`` (the QAD / converter table key)."""
+    match = FFN_LAYER.search(module_name.replace("_checkpoint_wrapped_module.", "").replace("_fsdp_wrapped_module.",
+                                                                                          ""))
+    return f"b{match.group(1)}.ff.{match.group(2)}" if match else None
+
+
+def percentile(values: list[float], q: float) -> float:
+    ordered = sorted(values)
+    position = q * (len(ordered) - 1)
+    low, high = math.floor(position), math.ceil(position)
+    return ordered[low] + (ordered[high] - ordered[low]) * (position - low)
+
+
+def candidate_table(rows: list[dict[str, Any]]) -> dict[str, dict[str, float]]:
+    """Per layer: candidate name -> amax, from the per-row amax records of ``collect``."""
+    keys = sorted({key for row in rows for key in row["amax"]}, key=lambda k: (int(k.split(".")[0][1:]), k))
+    table: dict[str, dict[str, float]] = {}
+    for key in keys:
+        per_row = [max(row["amax"][key]) for row in rows]
+        full = max(per_row)
+        if not (math.isfinite(full) and full > 0):
+            raise ValueError(f"{key}: non-positive or non-finite amax {full}")
+        candidates = {"max": full, "row_p99": percentile(per_row, 0.99), "row_p90": percentile(per_row, 0.90)}
+        candidates.update({f"{m}x_max": m * full for m in MULTIPLIERS})
+        table[key] = candidates
+    return table
+
+
+def outlier_row_ids(rows: list[dict[str, Any]], limit: int) -> list[str]:
+    """Rows that set some layer's maximum (most layers first): the MSE pass must see the outliers it clips."""
+    counts: dict[str, int] = {}
+    keys = {key for row in rows for key in row["amax"]}
+    for key in keys:
+        top = max(rows, key=lambda row: max(row["amax"][key]))
+        counts[top["id"]] = counts.get(top["id"], 0) + 1
+    return [row_id for row_id, _ in sorted(counts.items(), key=lambda item: -item[1])[:limit]]
+
+
+def choose_table(mse: dict[str, dict[str, list[float]]], candidates: dict[str, dict[str, float]],
+                 objective: str = "all") -> tuple[dict[str, float], dict[str, Any]]:
+    """Per layer, the candidate with the lowest error.
+
+    ``objective="all"`` is the FastH3 OmniRef criterion (relative output MSE over every token); ``"balanced"``
+    is the mean of the target-video and target-audio relative errors (audio is ~1% of the tokens, so the
+    all-token MSE barely sees it).
+    """
+    if objective not in ("all", "balanced"):
+        raise ValueError(f"objective must be 'all' or 'balanced', got {objective!r}")
+    table: dict[str, float] = {}
+    report: dict[str, Any] = {}
+    for key, cands in candidates.items():
+        energy = mse[key]["ref"]
+        scores: dict[str, dict[str, float]] = {}
+        for name in cands:
+            err = mse[key][name]
+            rel = {g: err[i] / energy[i] for i, g in enumerate(GROUP_NAMES) if energy[i] > 0}
+            target = [rel[g] for g in ("video", "audio") if g in rel]
+            scores[name] = {**rel, "all": sum(err) / sum(energy), "balanced": sum(target) / len(target)}
+        best = min(scores, key=lambda name: scores[name][objective])
+        table[key] = cands[best]
+        report[key] = {"best": best, "amax": cands[best], "max": cands["max"], "scores": scores}
+    return table, report
+
+
+# ---------------------------------------------------------------------- GPU passes
+def _sum_over_sp(tensor: torch.Tensor, op: Any) -> torch.Tensor:
+    from fastvideo.distributed import get_sp_group, get_sp_world_size
+    if get_sp_world_size() > 1:
+        dist.all_reduce(tensor, op=op, group=get_sp_group().device_group)
+    return tensor
+
+
+class FFNHooks:
+    """Hooks on every FFN block linear of ``transformer``; ``groups`` is this SP rank's row-group ids."""
+
+    def __init__(self, transformer: torch.nn.Module) -> None:
+        self.modules = {}
+        for name, module in transformer.named_modules():
+            key = amax_key(name)
+            if key is not None:
+                self.modules[key] = module
+        if len(self.modules) != 100:
+            raise ValueError(f"expected 100 FFN block linears (50 blocks x fc_in/fc_out), found {len(self.modules)}")
+        self.groups: torch.Tensor | None = None
+        self._handles: list[Any] = []
+
+    def _local(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        flat = x.detach().reshape(-1, x.shape[-1])
+        assert self.groups is not None
+        if flat.shape[0] != self.groups.numel():
+            raise RuntimeError(f"FFN input has {flat.shape[0]} rows, the row groups {self.groups.numel()}")
+        return flat, self.groups
+
+    @contextmanager
+    def attached(self, pre: Any = None, post: Any = None) -> Iterator[None]:
+        for key, module in self.modules.items():
+            if pre is not None:
+                self._handles.append(module.register_forward_pre_hook(pre(key)))
+            if post is not None:
+                self._handles.append(module.register_forward_hook(post(key)))
+        try:
+            yield
+        finally:
+            for handle in self._handles:
+                handle.remove()
+            self._handles = []
+
+
+def _group_reduce(values: torch.Tensor, groups: torch.Tensor, reduce: str) -> torch.Tensor:
+    out = torch.zeros(len(GROUP_NAMES), dtype=torch.float32, device=values.device)
+    valid = groups >= 0
+    return out.scatter_reduce(0, groups[valid], values[valid].float(), reduce=reduce, include_self=True)
+
+
+def _weight(module: torch.nn.Module) -> torch.Tensor:
+    weight = module.weight
+    full = getattr(weight, "full_tensor", None)
+    return full() if callable(full) else weight
+
+
+def walk_teacher(teacher: Any, prepared: Any, hooks: FFNHooks) -> None:
+    """The teacher's own six-forward trajectory of ``prepared`` (no grad), with the row groups set for hooks."""
+    from fastvideo.train.methods.knowledge_distillation.pdd_qad_metrics import step_packed
+    from fastvideo.train.methods.knowledge_distillation.pdd_qad_recon import local_row_groups, packed_row_groups
+
+    layout = prepared.layout
+    hooks.groups = local_row_groups(packed_row_groups(layout, int(layout.sequence_length)), teacher.device)
+    schedulers = teacher.schedulers(prepared)
+    timesteps = (schedulers[0].timesteps, schedulers[1].timesteps)
+    video, audio = prepared.video, prepared.audio
+    with torch.no_grad():
+        for rung in range(teacher.num_rungs):
+            with teacher.rung_scope(prepared, rung, timesteps) as rt:
+                out = teacher.forward_rung(prepared, video, audio, rt)
+            video, audio = step_packed(schedulers, prepared, rung, video, audio, *out)
+
+
+def collect_row(teacher: Any, prepared: Any, hooks: FFNHooks) -> dict[str, list[float]]:
+    found: dict[str, torch.Tensor] = {}
+
+    def pre(key: str) -> Any:
+
+        def hook(module: torch.nn.Module, args: tuple[Any, ...]) -> None:
+            flat, groups = hooks._local(args[0])
+            amax = _group_reduce(flat.abs().amax(dim=-1), groups, "amax")
+            found[key] = amax if key not in found else torch.maximum(found[key], amax)
+
+        return hook
+
+    with hooks.attached(pre=pre):
+        walk_teacher(teacher, prepared, hooks)
+    return {key: _sum_over_sp(value, dist.ReduceOp.MAX).tolist() for key, value in sorted(found.items())}
+
+
+def mse_row(teacher: Any, prepared: Any, hooks: FFNHooks, candidates: dict[str, dict[str, float]],
+            sums: dict[str, dict[str, torch.Tensor]]) -> None:
+    from fastvideo.layers.quantization.nvfp4_qad import nvfp4_gemm
+
+    def post(key: str) -> Any:
+
+        def hook(module: torch.nn.Module, args: tuple[Any, ...], output: Any) -> None:
+            x = args[0]
+            reference = (output[0] if isinstance(output, tuple) else output).detach()
+            flat_ref = reference.reshape(-1, reference.shape[-1]).float()
+            _, groups = hooks._local(x)
+            bias = None if getattr(module, "skip_bias_add", False) else module.bias
+            weight = _weight(module)
+            entry = sums.setdefault(key, {})
+            ref_energy = _group_reduce(flat_ref.pow(2).sum(dim=-1), groups, "sum")
+            entry["ref"] = entry.get("ref", torch.zeros_like(ref_energy)) + ref_energy
+            for name, amax in candidates[key].items():
+                scale = torch.tensor(FP4_RANGE / amax, dtype=torch.float32, device=x.device)
+                quant = nvfp4_gemm(x, weight, bias, scale).reshape(-1, flat_ref.shape[-1]).float()
+                err = _group_reduce((quant - flat_ref).pow(2).sum(dim=-1), groups, "sum")
+                entry[name] = entry.get(name, torch.zeros_like(err)) + err
+
+        return hook
+
+    with hooks.attached(post=post):
+        walk_teacher(teacher, prepared, hooks)
+
+
+# ---------------------------------------------------------------------- drivers
+def _build(args: argparse.Namespace) -> tuple[Any, list[Any]]:
+    """The QAD config's bf16 teacher and its training-row plan (the student is built without quantization)."""
+    from fastvideo.distributed import maybe_init_distributed_environment_and_model_parallel
+    from fastvideo.train.utils.builder import build_from_config
+    from fastvideo.train.utils.config import load_run_config
+
+    overrides = ["--models.student.qad.enabled", "false", "--training.tracker.trackers", "[]"] + list(args.override)
+    cfg = load_run_config(args.config, overrides=overrides)
+    tc = cfg.training
+    maybe_init_distributed_environment_and_model_parallel(tc.distributed.tp_size, tc.distributed.sp_size)
+    _, method, _, _ = build_from_config(cfg)
+    return method.teacher, list(method.student.train_rows)
+
+
+def _dp() -> tuple[int, int]:
+    """(data-parallel rank, size): sequence-parallel ranks of one group walk the same rows."""
+    from fastvideo.distributed import get_sp_world_size, get_world_group
+
+    world = get_world_group()
+    return world.rank // get_sp_world_size(), world.world_size // get_sp_world_size()
+
+
+def _shard(rows: list[Any], shard: int, num_shards: int) -> list[Any]:
+    """This group's rows; striding a hash order keeps resolutions mixed (the plan alternates 480p / 768p groups)."""
+    import hashlib
+
+    dp_rank, dp_size = _dp()
+    ordered = sorted(rows, key=lambda row: hashlib.sha256(row.id.encode()).hexdigest())
+    return ordered[shard * dp_size + dp_rank::num_shards * dp_size]
+
+
+def _is_writer() -> bool:
+    from fastvideo.distributed import get_sp_group
+    return get_sp_group().rank_in_group == 0
+
+
+def run_collect(args: argparse.Namespace) -> None:
+    teacher, plan = _build(args)
+    rows = _shard(plan[:args.max_rows] if args.max_rows else plan, args.shard, args.num_shards)
+    hooks = FFNHooks(teacher.transformer)
+    out = Path(args.output_dir)
+    out.mkdir(parents=True, exist_ok=True)
+    path = out / f"collect-shard{args.shard}-dp{_dp()[0]}.jsonl"
+    done = {json.loads(line)["id"] for line in path.read_text().splitlines()} if path.exists() else set()
+    for index, spec in enumerate(rows):
+        if spec.id in done:
+            continue
+        prepared = teacher.load_row(spec)
+        amax = collect_row(teacher, prepared, hooks)
+        if _is_writer():
+            record = {"id": spec.id, "case": spec.case, "resolution": spec.resolution, "amax": amax,
+                      "tokens": int(prepared.layout.sequence_length)}
+            with path.open("a") as handle:
+                handle.write(json.dumps(record) + "\n")
+            print(f"collect {index + 1}/{len(rows)} {spec.id} ({record['tokens']} tokens)", flush=True)
+
+
+def _read_rows(directory: Path, pattern: str) -> list[dict[str, Any]]:
+    rows = {}
+    for path in sorted(directory.glob(pattern)):
+        for line in path.read_text().splitlines():
+            record = json.loads(line)
+            rows[record["id"]] = record
+    if not rows:
+        raise ValueError(f"no {pattern} rows in {directory}")
+    return [rows[key] for key in sorted(rows)]
+
+
+def run_candidates(args: argparse.Namespace) -> None:
+    out = Path(args.output_dir)
+    rows = _read_rows(out, "collect-*.jsonl")
+    table = candidate_table(rows)
+    half = candidate_table(rows[:len(rows) // 2])
+    gaps = {key: 1 - half[key]["max"] / table[key]["max"] for key in table}
+    (out / "candidates.json").write_text(json.dumps(table, indent=1))
+    (out / "amax-max.json").write_text(json.dumps({k: v["max"] for k, v in table.items()}, indent=1))
+    (out / "outlier_rows.json").write_text(json.dumps(outlier_row_ids(rows, args.outlier_rows)))
+    worst = sorted(gaps.items(), key=lambda item: -item[1])[:5]
+    print(f"{len(rows)} rows, {len(table)} layers; first-half vs full max gap: worst {worst}")
+
+
+def run_mse(args: argparse.Namespace) -> None:
+    out = Path(args.output_dir)
+    candidates = json.loads((out / "candidates.json").read_text())
+    collected = {row["id"] for row in _read_rows(out, "collect-*.jsonl")}
+    outliers = set(json.loads((out / "outlier_rows.json").read_text()))
+    teacher, plan = _build(args)
+    # Rows the max was not taken from, a few per (case, resolution), plus the rows that set the maxima.
+    per_group: dict[tuple[str, str], int] = {}
+    fresh = []
+    for spec in plan:
+        key = (spec.case, spec.resolution)
+        if spec.id not in collected and per_group.get(key, 0) < args.mse_rows_per_group:
+            per_group[key] = per_group.get(key, 0) + 1
+            fresh.append(spec)
+    rows = _shard(fresh + [spec for spec in plan if spec.id in outliers], args.shard, args.num_shards)
+    hooks = FFNHooks(teacher.transformer)
+    sums: dict[str, dict[str, torch.Tensor]] = {}
+    for index, spec in enumerate(rows):
+        mse_row(teacher, teacher.load_row(spec), hooks, candidates, sums)
+        if _is_writer():
+            print(f"mse {index + 1}/{len(rows)} {spec.id}", flush=True)
+    reduced = {key: {name: _sum_over_sp(value, dist.ReduceOp.SUM).tolist() for name, value in entry.items()}
+               for key, entry in sums.items()}
+    if _is_writer():
+        path = out / f"mse-shard{args.shard}-dp{_dp()[0]}.json"
+        path.write_text(json.dumps({"rows": [spec.id for spec in rows], "sums": reduced}))
+
+
+def run_choose(args: argparse.Namespace) -> None:
+    out = Path(args.output_dir)
+    candidates = json.loads((out / "candidates.json").read_text())
+    total: dict[str, dict[str, list[float]]] = {}
+    for path in sorted(out.glob("mse-*.json")):
+        for key, entry in json.loads(path.read_text())["sums"].items():
+            slot = total.setdefault(key, {})
+            for name, values in entry.items():
+                slot[name] = [a + b for a, b in zip(slot.get(name, [0.0] * len(values)), values, strict=True)]
+    if set(total) != set(candidates):
+        raise ValueError(f"MSE sums cover {len(total)} of {len(candidates)} layers")
+    for objective, name in (("all", "amax-mse.json"), ("balanced", "amax-mse-balanced.json")):
+        table, report = choose_table(total, candidates, objective)
+        (out / name).write_text(json.dumps(table, indent=1))
+        (out / name.replace("amax-", "choose-")).write_text(json.dumps(report, indent=1))
+        picks: dict[str, int] = {}
+        for entry in report.values():
+            picks[entry["best"]] = picks.get(entry["best"], 0) + 1
+        print(f"{name} ({objective}): {dict(sorted(picks.items(), key=lambda item: -item[1]))}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("stage", choices=("collect", "candidates", "mse", "choose"))
+    parser.add_argument("--config", help="the adapter's QAD YAML (collect / mse)")
+    parser.add_argument("--output-dir", required=True)
+    parser.add_argument("--shard", type=int, default=0)
+    parser.add_argument("--num-shards", type=int, default=1)
+    parser.add_argument("--max-rows", type=int, default=0, help="cap the calibration plan (smoke tests)")
+    parser.add_argument("--mse-rows-per-group", type=int, default=3)
+    parser.add_argument("--outlier-rows", type=int, default=8)
+    # Any other ``--dotted.key value`` pair overrides the QAD config, as in fastvideo.train.entrypoint.train.
+    args, overrides = parser.parse_known_args()
+    args.override = overrides
+    {"collect": run_collect, "candidates": run_candidates, "mse": run_mse, "choose": run_choose}[args.stage](args)
+
+
+if __name__ == "__main__":
+    main()

scripts/distill/pika_h3_flex_qad/check_fp4_attention_parity.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Parity of the dense FP4 attention emulator (QAD training) with FlashAttention-4 NVFP4 (deployment).
+
+With ``--qk-transform`` both sides precondition Q/K first (the emulator through ``qk_transform=``, the deployed
+``FlashAttentionImpl._forward_nvfp4`` through the ``nvfp4_qk_transform`` impl arg). Both should compute Q/K in NVFP4
+(per-16 E4M3 block scales, unit global scale) with P and V in BF16. Reports the relative L2 distance of the emulator to the FA4 NVFP4 kernel next to the FP4-vs-BF16 distance: the emulator is a
+faithful training target when the first is a small fraction of the second. Needs a Blackwell GPU with FA4 FP4.
+
+    python scripts/distill/pika_h3_flex_qad/check_fp4_attention_parity.py --seq-len 8269 --heads 14 \\
+        --qk-transform smooth_k+hadamard
+"""
+from __future__ import annotations
+
+import argparse
+
+import torch
+
+
+def rel(a: torch.Tensor, b: torch.Tensor) -> float:
+    return float((a.float() - b.float()).norm() / b.float().norm().clamp_min(1e-12))
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--seq-len", type=int, default=8269, help="not a multiple of 128 on purpose")
+    parser.add_argument("--heads", type=int, default=14)
+    parser.add_argument("--dim", type=int, default=128)
+    parser.add_argument("--qk-transform", default="none", help="none, smooth_k, hadamard or smooth_k+hadamard")
+    parser.add_argument("--k-bias", type=float, default=4.0, help="per-channel K offset that smooth_k removes")
+    parser.add_argument("--outlier-scale", type=float, default=8.0, help="scale of a few outlier K channels")
+    args = parser.parse_args()
+
+    from fastvideo.attention.backends import flash_attn as fa
+    from fastvideo.attention.backends.fp4_dense_qat import FP4DenseAttention
+    from fastvideo.attention.backends.fp4_vsa_qat import FP4AttentionNumerics
+    from fastvideo.attention.utils.nvfp4_qk_precondition import parse_qk_transform
+
+    torch.manual_seed(0)
+    shape = (1, args.seq_len, args.heads, args.dim)
+    q, k, v = (torch.randn(shape, device="cuda", dtype=torch.bfloat16) for _ in range(3))
+    k[..., :4] *= args.outlier_scale
+    k += args.k_bias * torch.randn(1, 1, args.heads, args.dim, device="cuda", dtype=torch.bfloat16)
+    transform = parse_qk_transform(args.qk_transform)
+    scale = args.dim**-0.5
+
+    impl = fa.FlashAttentionImpl(num_heads=args.heads, head_size=args.dim, causal=False, softmax_scale=scale,
+                                 num_kv_heads=args.heads, nvfp4_fa4=True, nvfp4_qk_transform=args.qk_transform)
+    fa4_fp4 = impl._forward_nvfp4(q, k, v)
+    bf16 = torch.nn.functional.scaled_dot_product_attention(*(t.transpose(1, 2) for t in (q, k, v)),
+                                                           scale=scale).transpose(1, 2)
+    numerics = FP4AttentionNumerics(quantize=True, quantize_pv=False)
+    emulator = FP4DenseAttention(numerics, qk_transform=transform)(q, k, v, None)
+
+    print(f"qk_transform={args.qk_transform} k_bias={args.k_bias} outlier_scale={args.outlier_scale}")
+    fp4_vs_bf16 = rel(fa4_fp4, bf16)
+    emu_vs_fa4 = rel(emulator, fa4_fp4)
+    print(f"FA4 NVFP4 vs BF16 attention: rel L2 {fp4_vs_bf16:.5f}")
+    print(f"emulator  vs BF16 attention: rel L2 {rel(emulator, bf16):.5f}")
+    print(f"emulator  vs FA4 NVFP4:      rel L2 {emu_vs_fa4:.5f}  max |diff| "
+          f"{float((emulator.float() - fa4_fp4.float()).abs().max()):.4e}")
+    ok = emu_vs_fa4 <= 0.1 * fp4_vs_bf16
+    print(f"PARITY {'OK' if ok else 'MISMATCH'}: emulator-to-kernel distance is "
+          f"{emu_vs_fa4 / max(fp4_vs_bf16, 1e-12):.1%} of the FP4 quantization error")
+    raise SystemExit(0 if ok else 1)
+
+
+if __name__ == "__main__":
+    main()

scripts/distill/pika_h3_flex_qad/merge_pika_adapter.py

@@ -0,0 +1,163 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Merge one released Pika H3 Flex adapter into its MiniMax-H3 trunk for QAD.
+
+QAD needs the Flex teacher as plain weights: the student's fp32 masters start
+from them and the NVFP4 export quantizes them. This writes a model directory
+whose trunk (``transformer`` for FL2VA, ``transformer_ref`` for Ref2VA) holds
+``W + (alpha / rank) * B @ A`` for every adapted module, merged in fp32 and
+rounded once to the base dtype, so the result does not depend on the runtime
+LoRA path. Every other component is a symlink to the base snapshot.
+
+The adapter is located and validated with the inference loader
+(``fastvideo.pipelines.basic.pika_h3_flex.adapter.load_adapter``); its sha256
+must match the manifest. The output also carries ``pika_flex_qad.json``, the
+sampling contract the QAD model profile reads.
+
+    python scripts/distill/pika_h3_flex_qad/merge_pika_adapter.py \\
+        --base <MiniMax-H3 snapshot @42ed227> --adapter Pika-Labs/Pika-H3-Flex \\
+        --task ref2va --output data/models/pika-flex-ref2va-merged
+"""
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+from fastvideo.configs.pipelines.pika_h3_flex import (PIKA_AUDIO_SCHEDULER_SHIFT, PIKA_DENOISING_STEPS,
+                                                      PIKA_VIDEO_SCHEDULER_SHIFT)
+from fastvideo.pipelines.basic.pika_h3_flex.adapter import load_adapter
+
+CONTRACT_FILE = "pika_flex_qad.json"
+CONTRACT_SCHEMA = "pika-flex-qad-v1"
+_LORA_A = ".lora_A.default.weight"
+_LORA_B = ".lora_B.default.weight"
+
+
+def adapter_factors(path: str) -> dict[str, tuple[torch.Tensor, torch.Tensor]]:
+    """``{module: (A [r, in], B [out, r])}`` from a PEFT adapter file."""
+    factors: dict[str, dict[str, torch.Tensor]] = {}
+    with safe_open(path, framework="pt") as handle:
+        for key in handle.keys():  # noqa: SIM118 - safe_open is not a mapping
+            for suffix, slot in ((_LORA_A, "A"), (_LORA_B, "B")):
+                if key.endswith(suffix):
+                    factors.setdefault(key[:-len(suffix)], {})[slot] = handle.get_tensor(key)
+                    break
+            else:
+                raise ValueError(f"{path}: unexpected adapter tensor {key!r}")
+    incomplete = sorted(module for module, pair in factors.items() if set(pair) != {"A", "B"})
+    if incomplete:
+        raise ValueError(f"{path}: modules without both LoRA factors: {incomplete[:5]}")
+    return {module: (pair["A"], pair["B"]) for module, pair in factors.items()}
+
+
+def merge_weight(weight: torch.Tensor, lora_a: torch.Tensor, lora_b: torch.Tensor, scale: float) -> torch.Tensor:
+    """``weight + scale * B @ A`` in fp32, rounded once to ``weight``'s dtype."""
+    if lora_b.shape[1] != lora_a.shape[0] or (lora_b.shape[0], lora_a.shape[1]) != tuple(weight.shape):
+        raise ValueError(f"LoRA shapes A {tuple(lora_a.shape)} / B {tuple(lora_b.shape)} do not fit weight "
+                         f"{tuple(weight.shape)}")
+    delta = lora_b.to(torch.float32) @ lora_a.to(torch.float32)
+    return (weight.to(torch.float32) + scale * delta).to(weight.dtype)
+
+
+def merge_trunk(trunk_dir: Path, output_dir: Path, factors: dict[str, tuple[torch.Tensor, torch.Tensor]],
+                scale: float) -> dict[str, Any]:
+    """Write ``trunk_dir``'s shards with every adapted module merged; return a per-module report."""
+    index_path = trunk_dir / "diffusion_pytorch_model.safetensors.index.json"
+    index = json.loads(index_path.read_text())
+    shards = sorted(set(index["weight_map"].values()))
+    output_dir.mkdir(parents=True, exist_ok=True)
+    remaining = set(factors)
+    report: dict[str, float] = {}
+    for shard in shards:
+        tensors: dict[str, torch.Tensor] = {}
+        with safe_open(str(trunk_dir / shard), framework="pt") as handle:
+            metadata = handle.metadata()
+            for key in handle.keys():  # noqa: SIM118
+                tensor = handle.get_tensor(key)
+                module = key[:-len(".weight")] if key.endswith(".weight") else None
+                if module in factors:
+                    lora_a, lora_b = factors[module]
+                    merged = merge_weight(tensor, lora_a, lora_b, scale)
+                    report[module] = float((merged.float() - tensor.float()).norm() / tensor.float().norm())
+                    remaining.discard(module)
+                    tensor = merged
+                tensors[key] = tensor.contiguous()
+        save_file(tensors, str(output_dir / shard), metadata=metadata)
+    if remaining:
+        raise ValueError(f"{len(remaining)} adapter modules have no weight in {trunk_dir}: "
+                         f"{sorted(remaining)[:5]}")
+    (output_dir / index_path.name).write_text(index_path.read_text())
+    (output_dir / "config.json").write_text((trunk_dir / "config.json").read_text())
+    return {"modules": len(report), "max_rel_delta": max(report.values()), "rel_delta": report}
+
+
+def link_base(base: Path, output: Path, trunk: str) -> None:
+    """Symlink every base entry except the merged trunk."""
+    output.mkdir(parents=True, exist_ok=True)
+    for entry in sorted(base.iterdir()):
+        if entry.name == trunk:
+            continue
+        target = output / entry.name
+        if target.is_symlink() or target.exists():
+            raise FileExistsError(f"{target} already exists; use an empty output directory")
+        os.symlink(entry.resolve(), target)
+
+
+def sha256_file(path: str) -> str:
+    digest = hashlib.sha256()
+    with open(path, "rb") as handle:
+        for chunk in iter(lambda: handle.read(1 << 24), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--base", required=True, help="MiniMax-H3 snapshot at the adapter's pinned revision")
+    parser.add_argument("--adapter", required=True, help="adapter file, release directory, or Hub repo id")
+    parser.add_argument("--task", required=True, choices=("fl2va", "ref2va"))
+    parser.add_argument("--revision", default=None, help="Hub revision when --adapter is a repo id")
+    parser.add_argument("--output", required=True)
+    args = parser.parse_args()
+
+    adapter = load_adapter(args.adapter, task=args.task, revision=args.revision)
+    digest = sha256_file(adapter.path)
+    if adapter.manifest.get("sha256") and digest != adapter.manifest["sha256"]:
+        raise ValueError(f"{adapter.path} has sha256 {digest}, the manifest says {adapter.manifest['sha256']}")
+    base, output = Path(args.base).resolve(), Path(args.output)
+    if output.exists() and any(output.iterdir()):
+        raise FileExistsError(f"{output} is not empty")
+    factors = adapter_factors(adapter.path)
+    scale = float(adapter.manifest["alpha"]) / float(adapter.manifest["rank"])
+    link_base(base, output, adapter.trunk)
+    report = merge_trunk(base / adapter.trunk, output / adapter.trunk, factors, scale)
+    contract = {
+        "schema": CONTRACT_SCHEMA,
+        "task": args.task,
+        "trunk": adapter.trunk,
+        "release_id": adapter.manifest.get("release_id"),
+        "adapter_sha256": digest,
+        "base": adapter.manifest.get("base"),
+        "numerics_policy": adapter.manifest.get("numerics_policy"),
+        "denoising_steps": list(PIKA_DENOISING_STEPS),
+        "video_scheduler_shift": PIKA_VIDEO_SCHEDULER_SHIFT,
+        "audio_scheduler_shift": PIKA_AUDIO_SCHEDULER_SHIFT,
+        "sampler": "renoise",
+        "merged_modules": report["modules"],
+        "merge": "fp32 W + (alpha/rank) B@A, rounded once to the base dtype",
+    }
+    (output / CONTRACT_FILE).write_text(json.dumps(contract, indent=2) + "\n")
+    (output / "merge_report.json").write_text(json.dumps(report, indent=2) + "\n")
+    print(f"merged {report['modules']} modules into {output / adapter.trunk} "
+          f"(max relative weight change {report['max_rel_delta']:.4f})")
+
+
+if __name__ == "__main__":
+    main()

scripts/distill/pika_h3_flex_qad/probe_fp4_attention_transforms.py

@@ -0,0 +1,161 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Per-layer FP4 attention error of exact Q/K transforms on real Pika H3 Flex activations.
+
+Runs one single-GPU Pika H3 Flex generation (``basic_pika_h3_flex.py``, in-process) with every dense
+``FlashAttentionImpl`` call wrapped: the BF16 output is returned unchanged, and next to it the deployed FA4 NVFP4
+attention (Q/K per-16 NVFP4, P/V BF16) is evaluated on transformed inputs. The transforms leave Q K^T unchanged in
+exact arithmetic:
+
+* ``smooth_k``: subtract each head's token-mean of K (softmax is shift-invariant per query row);
+* ``hadamard``: rotate Q and K by the same orthonormal head-dim Hadamard matrix (spreads outlier channels).
+
+Reports the relative L2 distance to the BF16 output per variant, per sampler forward and per layer. Needs a Blackwell
+GPU with FA4 NVFP4 (``FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN FASTVIDEO_FA4=1``). Every other argument goes to the
+example unchanged:
+
+    python scripts/distill/pika_h3_flex_qad/probe_fp4_attention_transforms.py --out probe.json -- \\
+        --model-path <flex release> --task fl2va --image first.png --last-image last.png --prompt "..." \\
+        --resolution 480p --duration 5
+"""
+from __future__ import annotations
+
+import argparse
+import collections
+import json
+import math
+import runpy
+import sys
+from pathlib import Path
+from typing import Any
+
+import torch
+
+EXAMPLE = Path(__file__).resolve().parents[3] / "examples" / "inference" / "basic" / "basic_pika_h3_flex.py"
+# name -> (smooth K, Hadamard-rotate Q and K)
+VARIANTS = {"plain": (False, False), "smooth_k": (True, False), "hadamard": (False, True), "smooth_k+hadamard": (True, True)}
+
+
+def hadamard(n: int, device: torch.device) -> torch.Tensor:
+    if n & (n - 1):
+        raise ValueError(f"Hadamard rotation needs a power-of-two head dim, got {n}")
+    h = torch.ones(1, 1)
+    while h.shape[0] < n:
+        h = torch.cat([torch.cat([h, h], 1), torch.cat([h, -h], 1)], 0)
+    return (h / math.sqrt(n)).to(device)
+
+
+def transform(q: torch.Tensor, k: torch.Tensor, smooth: bool, rotation: torch.Tensor | None) -> tuple[torch.Tensor, ...]:
+    """``[B, L, H, D]`` Q/K -> transformed BF16 Q/K with the same Q K^T up to a per-query-row shift."""
+    qf, kf = q.float(), k.float()
+    if smooth:
+        kf = kf - kf.mean(dim=1, keepdim=True)
+    if rotation is not None:
+        qf, kf = qf @ rotation, kf @ rotation
+    return qf.to(q.dtype), kf.to(k.dtype)
+
+
+def fa4_nvfp4(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, scale: float) -> torch.Tensor:
+    from fastvideo.attention.backends import flash_attn as fa
+
+    q4, q_sf = fa._nvfp4_quantize_for_fa4(q)
+    k4, k_sf = fa._nvfp4_quantize_for_fa4(k)
+    out = fa.flash_attn_fp4_func(q4[:, :q.shape[1]], k4[:, :k.shape[1]], v, q_sf, k_sf, softmax_scale=scale, causal=False)
+    return out[0] if isinstance(out, tuple) else out
+
+
+def rel_l2(a: torch.Tensor, b: torch.Tensor) -> float:
+    return float((a.float() - b.float()).norm() / b.float().norm().clamp_min(1e-12))
+
+
+class Probe:
+
+    def __init__(self) -> None:
+        self.layer_of: dict[int, int] = {}
+        self.calls: collections.Counter[int] = collections.Counter()
+        self.records: list[dict[str, Any]] = []
+        self._rotations: dict[tuple[int, str], torch.Tensor] = {}
+
+    def install(self) -> None:
+        from fastvideo.attention.backends.flash_attn import FlashAttentionImpl
+
+        original = FlashAttentionImpl.forward
+        probe = self
+
+        def forward(impl: Any, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, metadata: Any) -> Any:
+            out = original(impl, query, key, value, metadata)
+            # The NVFP4 VAE decoder captures its blocks in a CUDA graph; measuring there would sync mid-capture.
+            if (getattr(metadata, "packed_sequence", None) is None and query.dim() == 4 and query.shape[-1] % 16 == 0
+                    and not torch.cuda.is_current_stream_capturing()):
+                with torch.no_grad():
+                    probe.measure(impl, query, key, value, out)
+            return out
+
+        FlashAttentionImpl.forward = forward
+
+    def measure(self, impl: Any, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, out: torch.Tensor) -> None:
+        layer = self.layer_of.setdefault(id(impl), len(self.layer_of))
+        call = self.calls[id(impl)]
+        self.calls[id(impl)] += 1
+        key = (q.shape[-1], str(q.device))
+        if key not in self._rotations:
+            self._rotations[key] = hadamard(q.shape[-1], q.device)
+        q, k, v = (t.to(torch.bfloat16) for t in (q, k, v))
+        record: dict[str, Any] = {"layer": layer, "forward": call, "tokens": int(q.shape[1]), "heads": int(q.shape[2])}
+        for name, (smooth, rotate) in VARIANTS.items():
+            tq, tk = transform(q, k, smooth, self._rotations[key] if rotate else None)
+            record[name] = rel_l2(fa4_nvfp4(tq, tk, v, impl.softmax_scale), out)
+        self.records.append(record)
+
+    def summary(self) -> dict[str, Any]:
+
+        def mean(rows: list[dict[str, Any]], name: str) -> float:
+            return sum(r[name] for r in rows) / max(len(rows), 1)
+
+        by_forward = collections.defaultdict(list)
+        by_layer = collections.defaultdict(list)
+        for r in self.records:
+            by_forward[r["forward"]].append(r)
+            by_layer[r["layer"]].append(r)
+        overall = {name: mean(self.records, name) for name in VARIANTS}
+        best = min(overall, key=overall.get)
+        layers = sorted(((layer, {n: mean(rows, n) for n in VARIANTS}) for layer, rows in by_layer.items()),
+                        key=lambda item: -item[1][best])
+        return {
+            "calls": len(self.records),
+            "layers": len(by_layer),
+            "overall": overall,
+            "best": best,
+            "per_forward": {f: {n: mean(rows, n) for n in VARIANTS} for f, rows in sorted(by_forward.items())},
+            "worst_layers_under_best": layers[:8],
+            "per_layer": dict(sorted(layers)),
+        }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--out", required=True, help="JSON report path")
+    args, example_args = parser.parse_known_args()
+    if example_args and example_args[0] == "--":
+        example_args = example_args[1:]
+    probe = Probe()
+    probe.install()
+    sys.argv = [str(EXAMPLE), *example_args]
+    try:
+        runpy.run_path(str(EXAMPLE), run_name="__main__")
+    finally:  # keep the DiT measurements even if decoding fails
+        write_report(probe, example_args, args.out)
+
+
+def write_report(probe: Probe, example_args: list[str], out: str) -> None:
+    if not probe.records:
+        print("no attention calls measured")
+        return
+    report = {"args": example_args, "summary": probe.summary(), "records": probe.records}
+    Path(out).write_text(json.dumps(report, indent=1))
+    s = report["summary"]
+    print(json.dumps({k: s[k] for k in ("calls", "layers", "overall", "best", "per_forward")}, indent=1))
+    print("worst layers under", s["best"], json.dumps(s["worst_layers_under_best"]))
+
+
+if __name__ == "__main__":
+    main()

scripts/distill/pika_h3_flex_qad/slinky/README.md

@@ -0,0 +1,15 @@
+# Pika H3 Flex Ref2VA NVFP4 QAD: Slinky (GB200) run kit
+
+Cluster scripts and resolved configs exactly as used for the first Flex Ref2VA QAD run (2026-10-30). Paths are the
+cluster's (`/mnt/nfs/vlm-aryan/pika-flex-qad`); the code there is a checkout of this branch. Jobs attach to held
+trays with `srun --overlap`.
+
+| Step | Script | Notes |
+|---|---|---|
+| Teachers | `merge.sh` | Downloads both adapters and runs `merge_pika_adapter.py` for ref2va and fl2va. |
+| Coverage rows | `build_rows.sh plan / encode / finalize` | `build_flex_ref2va_rows.py`: 480p rows at Flex reference sizing plus audio-reference, mixed image+video, 6-9 image and two-video rows, from the stored v4 latents. 1,186 rows. |
+| Calibration | `calib_pipeline.sh` (uses `calib.sh`) | collect (1,133 rows, 8 GPUs) -> candidates -> MSE -> choose -> step-0 PTQ evals. Step-0 scores: max 0.33112, MSE 0.32987, balanced 0.32595 (pinned, `examples/train/calibration/pika_h3_flex/amax-ref2va-balanced.json`). |
+| Warm-up trial | `theseus_trial.sh` | 8 steps each: baseline 0.34472 -> 0.32510, Theseus (4 warm-up + 4) -> 0.32943. Not adopted. |
+| QAD | `qad_node.sh` (4 nodes x 4 GPUs) with `ref2va-qad.yaml` | Needs `FASTVIDEO_PG_TIMEOUT_MIN=120` (the default 10-minute collective timeout fires during the model load). |
+
+`env.sh` sources the cluster's `secrets.env` by path; no credentials are stored here.

scripts/distill/pika_h3_flex_qad/slinky/build_rows.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+# usage: build_rows.sh plan | encode <num_shards> <shard_base> | finalize   (runs under env.sh; encode = one process per GPU)
+O=/mnt/nfs/vlm-aryan/pika-flex-qad/data/flex-ref2va; V4=/mnt/lustre/vlm-wlsaidhi/fastvideo/data/omniref-v4-processed/20260919T001000Z
+EVAL=/mnt/nfs/vlm-aryan/h3-nvfp4-decoder/gen-latents/omniref-eval/manifest.json
+case $1 in
+  plan) $PY /mnt/nfs/vlm-aryan/pika-flex-qad/code/scripts/distill/pika_h3_flex_qad/build_flex_ref2va_rows.py plan --output-dir $O --v4-manifest $V4/encoded/full/MANIFEST.jsonl --plan-config /mnt/nfs/vlm-aryan/pika-flex-qad/ref2va-qad.yaml --models.student.data.max_frames 345 --models.student.data.per_group.480p 160 --models.student.data.per_group.768p 40 --eval-manifest $EVAL --extra-per-group 480p=80 768p=20 ;;
+  encode) for g in 0 1 2 3; do CUDA_VISIBLE_DEVICES=$g MASTER_PORT=$((29700 + g)) $PY /mnt/nfs/vlm-aryan/pika-flex-qad/code/scripts/distill/pika_h3_flex_qad/build_flex_ref2va_rows.py encode --output-dir $O --model-path /mnt/nfs/vlm-aryan/pika-flex-qad/models/pika-flex-ref2va-merged --shard $(($3 + g)) --num-shards $2 ${@:4} >> $O/encode-$(($3 + g)).log 2>&1 & done; wait ;;
+  finalize) $PY /mnt/nfs/vlm-aryan/pika-flex-qad/code/scripts/distill/pika_h3_flex_qad/build_flex_ref2va_rows.py finalize --output-dir $O --eval-manifest $EVAL ;;
+esac
+echo BUILD_DONE $1

scripts/distill/pika_h3_flex_qad/slinky/calib.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+# usage: calib.sh <collect|mse|candidates|choose> <config> <out_dir> [shard num_shards] [--dotted.key value ...]  (runs under env.sh)
+set -uo pipefail
+STAGE=$1; CFG=$2; OUT=$3; shift 3
+S=/mnt/nfs/vlm-aryan/pika-flex-qad/code/scripts/distill/pika_h3_flex_qad/calibrate_flex_nvfp4.py
+mkdir -p $OUT
+case $STAGE in
+  collect|mse)
+    SH=$1; NS=$2; shift 2
+    FASTVIDEO_PG_TIMEOUT_MIN=120 $PY -m torch.distributed.run --nproc_per_node=${NPROC:-4} --master_port=$((29600 + SH)) \
+      --log-dir $OUT/torchlogs-$STAGE-$SH --redirects 3 --tee 3 $S $STAGE --config $CFG --output-dir $OUT \
+      --shard $SH --num-shards $NS --training.distributed.num_gpus ${NPROC:-4} --training.distributed.hsdp_shard_dim ${NPROC:-4} --training.distributed.sp_size ${NPROC:-4} "$@" ;;
+  *) $PY $S $STAGE --output-dir $OUT "$@" ;;
+esac
+echo "CALIB_STAGE_DONE $STAGE"

scripts/distill/pika_h3_flex_qad/slinky/calib_pipeline.sh

@@ -0,0 +1,31 @@
+#!/bin/bash
+# Overnight Flex Ref2VA calibration on trays 13509 + 13483 (login node, in tmux): finalize coverage rows,
+# collect, candidates, MSE, choose, then step-0 PTQ evals of each table. Log: logs/calib-pipeline.log
+R=/mnt/nfs/vlm-aryan/pika-flex-qad; C=$R/calib/ref2va; CFG=$R/ref2va-qad.yaml; L=$R/logs
+T0=13509; T1=13483
+run() {  # run <tray> <log> <cmd...>: one container step under env.sh, foreground
+  local t=$1 log=$2; shift 2
+  srun --overlap --jobid=$t --nodes=1 --ntasks=1 --kill-on-bad-exit=1 --container-image=nvcr.io/nvidia/pytorch:25.06-py3 \
+    --container-mounts=/mnt/nfs/vlm-aryan:/mnt/nfs/vlm-aryan,/mnt/lustre:/mnt/lustre:ro bash $R/env.sh "$@" > $log 2>&1
+}
+say() { echo "$(date -u +%FT%TZ) $*"; }
+say "start: coverage rows final (retry stopped)"
+run $T0 $L/build-finalize.log bash $R/jobs/build_rows.sh finalize; tail -3 $L/build-finalize.log
+A="--models.student.data.per_group.480p 120 --models.student.data.per_group.768p 30"
+say "collect"
+run $T0 $L/calib-collect-0.log bash $R/jobs/calib.sh collect $CFG $C 0 2 $A & run $T1 $L/calib-collect-1.log bash $R/jobs/calib.sh collect $CFG $C 1 2 $A & wait
+grep -h "CALIB_STAGE_DONE\|Traceback" $L/calib-collect-*.log; wc -l $C/collect-*.jsonl
+run $T0 $L/calib-candidates.log bash $R/jobs/calib.sh candidates - $C --outlier-rows 8; tail -2 $L/calib-candidates.log
+say "mse"
+run $T0 $L/calib-mse-0.log bash $R/jobs/calib.sh mse $CFG $C 0 2 --mse-rows-per-group 3 & run $T1 $L/calib-mse-1.log bash $R/jobs/calib.sh mse $CFG $C 1 2 --mse-rows-per-group 3 & wait
+run $T0 $L/calib-choose.log bash $R/jobs/calib.sh choose - $C; tail -3 $L/calib-choose.log
+say "step-0 PTQ evals"
+ptq() {  # ptq <tray> <name> <table>
+  run $1 $L/ptq-$2.log bash $R/jobs/smoke.sh $CFG $R/runs/ptq-$2 --models.student.qad.linears.amax_json $3 \
+    --training.distributed.num_gpus 4 --training.distributed.hsdp_shard_dim 4 --training.loop.max_train_steps 1 \
+    --training.optimizer.learning_rate 0.0 --method.noise_floor_eval false --training.tracker.trackers "[]"
+  grep -o "QAD eval @0: score [0-9.]* (PTQ [0-9.]*)" $L/ptq-$2.log | head -1 | sed "s/^/$2: /"
+}
+ptq $T0 max $C/amax-max.json & ptq $T1 mse $C/amax-mse.json & wait
+ptq $T0 balanced $C/amax-mse-balanced.json
+say "CALIB_PIPELINE_DONE"

scripts/distill/pika_h3_flex_qad/slinky/calib_smoke.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+# Calibration smoke on one tray: 2 collect rows, candidates, 2 MSE rows + outliers, choose.
+R=/mnt/nfs/vlm-aryan/pika-flex-qad; O=$R/runs/calib-smoke; rm -rf $O
+bash $R/jobs/calib.sh collect $R/smoke-ref2va.yaml $O 0 1 --max-rows 2 || exit 1
+bash $R/jobs/calib.sh candidates - $O --outlier-rows 1 || exit 1
+bash $R/jobs/calib.sh mse $R/smoke-ref2va.yaml $O 0 1 --mse-rows-per-group 1 || exit 1
+bash $R/jobs/calib.sh choose - $O || exit 1
+ls -la $O

scripts/distill/pika_h3_flex_qad/slinky/env.sh

@@ -0,0 +1,23 @@
+# Pika H3 Flex QAD environment (OmniRef QAD v2 env with this project code). Usage: env.sh <cmd...>
+set -a
+Q=/mnt/nfs/vlm-aryan/h3-omniref-qad
+source /mnt/nfs/vlm-aryan/fasth3-33b-20260806/secrets.env >/dev/null 2>&1
+EXP=/mnt/lustre/vlm-wlsaidhi/fastvideo/experiments/main-hf-step3500-gb200-20261005T203006Z
+KERNEL=$EXP/kernel/install
+PY=/mnt/nfs/vlm-aryan/h3-nvfp4-decoder/venv/bin/python
+CODE=/mnt/nfs/vlm-aryan/pika-flex-qad/code
+PYTHONPATH=$KERNEL:$CODE
+HF_HOME=/mnt/nfs/vlm-aryan/hf-cache HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 HF_HUB_DISABLE_TELEMETRY=1
+WANDB_MODE=${WANDB_MODE:-online} WANDB_DIR=/mnt/nfs/vlm-aryan/pika-flex-qad/wandb WANDB_PROJECT=pika-h3-flex-nvfp4 WANDB_ENTITY=aryan5v-san-jose-state-university
+PYTHONDONTWRITEBYTECODE=1 PYTHONUNBUFFERED=1 TOKENIZERS_PARALLELISM=false OMP_NUM_THREADS=8
+SCRATCH=/tmp/pika-flex-${SLURM_JOB_ID:-nojob}
+HOME=$SCRATCH/home XDG_CACHE_HOME=$SCRATCH/cache XDG_CONFIG_HOME=$SCRATCH/config TMPDIR=$SCRATCH/tmp
+TRITON_CACHE_DIR=/tmp/triton-pika-flex TORCHINDUCTOR_CACHE_DIR=/tmp/inductor-pika-flex TORCH_EXTENSIONS_DIR=/tmp/torchext-pika-flex
+NUMBA_CACHE_DIR=$SCRATCH/numba FLASHINFER_WORKSPACE_BASE=/mnt/nfs/vlm-aryan/h3-omniref-qad/cache/flashinfer
+mkdir -p $HOME $XDG_CACHE_HOME $XDG_CONFIG_HOME $TMPDIR $NUMBA_CACHE_DIR $TRITON_CACHE_DIR $TORCHINDUCTOR_CACHE_DIR $TORCH_EXTENSIONS_DIR $WANDB_DIR
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+TORCH_NCCL_ENABLE_MONITORING=0 NCCL_CUMEM_ENABLE=0 NCCL_MNNVL_ENABLE=0 NCCL_NVLS_ENABLE=0 NCCL_NET=IB NCCL_IB_HCA==mlx5_0,mlx5_1,mlx5_3,mlx5_4 NCCL_SOCKET_IFNAME=eth0 NCCL_SOCKET_FAMILY=AF_INET NCCL_DEBUG=WARN
+FASTVIDEO_MINIMAX_H3_FUSIONS=0 FASTVIDEO_H3_VSA_FP4=0 FASTVIDEO_VSA_SM100A=0
+set +a
+cd $CODE
+exec "$@"

scripts/distill/pika_h3_flex_qad/slinky/launch.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+# usage: launch_o.sh <tray jobid> <timeout seconds> <log> <cmd...>  (Pika Flex env; detached on the login node)
+J=$1; T=$2; LOG=$3; shift 3
+nohup timeout $((T+120)) srun --overlap --jobid=$J --nodes=1 --ntasks=1 --kill-on-bad-exit=1 \
+  --container-image=nvcr.io/nvidia/pytorch:25.06-py3 \
+  --container-mounts=/mnt/nfs/vlm-aryan:/mnt/nfs/vlm-aryan,/mnt/lustre:/mnt/lustre:ro \
+  timeout $T bash /mnt/nfs/vlm-aryan/pika-flex-qad/env.sh "$@" > $LOG 2>&1 &
+echo "launched pid $! on tray $J -> $LOG"

scripts/distill/pika_h3_flex_qad/slinky/merge.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+# Download the two Pika H3 Flex adapters and merge each into its MiniMax-H3 trunk (runs under env.sh).
+set -euo pipefail
+R=/mnt/nfs/vlm-aryan/pika-flex-qad
+BASE=$HF_HOME/hub/models--MiniMaxAI--MiniMax-H3/snapshots/42ed227ee7df40d41602854ae760620d6eb651fe
+ADAPTERS=$R/models/pika-h3-flex-adapters
+for task in ref2va fl2va; do
+  out=$R/models/pika-flex-$task-merged
+  if [ -f $out/pika_flex_qad.json ]; then echo "$out exists, skipping"; continue; fi
+  rm -rf $out
+  time $PY scripts/distill/pika_h3_flex_qad/merge_pika_adapter.py --base $BASE --adapter $ADAPTERS --task $task --output $out
+  cat $out/pika_flex_qad.json
+done

scripts/distill/pika_h3_flex_qad/slinky/pilot.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+# One-tray QAD pilot (SP 4): pilot.sh <tray> <name> <amax table> [--dotted.key value ...]; logs/pilot-<name>.log, runs/pilot-<name>
+R=/mnt/nfs/vlm-aryan/pika-flex-qad; T=$1; N=$2; A=$3; shift 3
+nohup srun --overlap --jobid=$T --nodes=1 --ntasks=1 --kill-on-bad-exit=1 --container-image=nvcr.io/nvidia/pytorch:25.06-py3 \
+  --container-mounts=/mnt/nfs/vlm-aryan:/mnt/nfs/vlm-aryan,/mnt/lustre:/mnt/lustre:ro \
+  bash $R/env.sh bash $R/jobs/smoke.sh $R/ref2va-qad.yaml $R/runs/pilot-$N \
+  --models.student.qad.linears.amax_json $A --training.distributed.num_gpus 4 --training.distributed.hsdp_shard_dim 4 \
+  --training.loop.max_train_steps 60 --method.eval_every 15 --method.full_eval_every 30 --method.patience 99 \
+  --method.noise_floor_eval false --method.max_rollout_rows 4 \
+  --models.student.data.eval_per_group.480p 1 --models.student.data.eval_per_group.768p 1 \
+  --training.tracker.run_name pilot-$N --training.checkpoint.training_state_checkpointing_steps 1000 "$@" \
+  > $R/logs/pilot-$N.log 2>&1 &
+echo "pilot $N on tray $T (pid $!) -> $R/logs/pilot-$N.log"

scripts/distill/pika_h3_flex_qad/slinky/qad_node.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+# usage: qad_node.sh <node_rank> <master_ip> <out_dir> <config> [overrides]  (4 nodes x 4 GPUs; runs under env.sh)
+NR=$1; MASTER=$2; OUT=$3; CFG=$4; shift 4
+mkdir -p $OUT
+FASTVIDEO_PG_TIMEOUT_MIN=120 $PY -m torch.distributed.run --nnodes=4 --nproc_per_node=4 --node_rank=$NR --master_addr=$MASTER \
+  --master_port=29650 --log-dir $OUT/torchlogs-n$NR --redirects 3 --tee 3 -m fastvideo.train.entrypoint.train --config $CFG "$@"

scripts/distill/pika_h3_flex_qad/slinky/ref2va-qad.yaml

@@ -0,0 +1,136 @@
+# Pika H3 Flex Ref2VA NVFP4 QAD (the FastH3 QAD recipe on Flex's six-forward ladder).
+#
+# Teacher: Pika H3 Flex Ref2VA in bf16, i.e. MiniMax-H3 `transformer_ref` with the
+# released `ref2va` adapter merged by
+# scripts/distill/pika_h3_flex_qad/merge_pika_adapter.py. Student: the same
+# weights with NVFP4 W4A4 q/k/v/out plus FFN fc_in/fc_out in all 50 blocks
+# active from step 0, static FFN activation scales from the Flex calibration
+# table. Attention stays BF16 dense: Flex has no VSA gate, so there is no FP4
+# sparse-attention stage and no gate to freeze. Norms, AdaLN, the token refiner
+# and the FP32 boundary modules stay frozen.
+#
+# The objective, evaluation, guards and checkpoint contract are the shared
+# recipe's (docs/training/minimax_h3_fp4_qad.md). Rows are OmniRef Ref2VA
+# parquet rows; the teacher computes every target on the fly, so the OmniRef
+# held-out eval manifest is reused as is.
+#
+# Paths are placeholders: the merged model, the calibration table (pin its
+# sha256 once calibrated), the row manifests and the output dir.
+models:
+  student:
+    _target_: fastvideo.train.models.minimax_h3.flex.MiniMaxH3FlexRef2VAModel
+    init_from: /mnt/nfs/vlm-aryan/pika-flex-qad/models/pika-flex-ref2va-merged
+    trainable: true
+    enable_gradient_checkpointing_type: full
+    qad:
+      enabled: true
+      linears:
+        enabled: true
+        amax_json: /mnt/nfs/vlm-aryan/pika-flex-qad/calib/amax-ref2va-balanced-ad60b45fea00.json
+        amax_sha256: ad60b45fea00906c55e370a0dec82638eb413d6b68537d1ceb38fc52da90b373
+        # Attention linears use the unit activation scale; only the FFN
+        # linears read their calibrated amax from this table.
+        attention: true
+        ffn: true
+        skip_blocks: []
+      fp4_attention:
+        enabled: false
+    data:
+      manifests:
+      # Flex-tier rows first: they replace the v4 rows with the same id and resolution.
+      - /mnt/nfs/vlm-aryan/pika-flex-qad/data/flex-ref2va/MANIFEST.jsonl
+      - /mnt/lustre/vlm-wlsaidhi/fastvideo/data/omniref-v4-processed/20260919T001000Z/encoded/full/MANIFEST.jsonl
+      cases:
+      - first_frame
+      - first_last_frame
+      - storyboard
+      - continue_scene
+      - continue_shot
+      - audio_reference
+      - mixed_image_video
+      - many_images
+      - multi_video
+      per_group:
+        480p: 160
+        768p: 40
+      seed: 20261007
+      max_frames: 345
+      eval_manifest: /mnt/nfs/vlm-aryan/pika-flex-qad/data/flex-ref2va/eval-manifest.json
+      eval_resolutions:
+      - 480p
+      - 768p
+      eval_per_group:
+        480p: 2
+        768p: 1
+      eval_max_frames: 0
+  teacher:
+    _target_: fastvideo.train.models.minimax_h3.flex.MiniMaxH3FlexRef2VAModel
+    init_from: /mnt/nfs/vlm-aryan/pika-flex-qad/models/pika-flex-ref2va-merged
+    trainable: false
+    frozen_bf16: true
+
+method:
+  _target_: fastvideo.train.methods.knowledge_distillation.pdd_qad.OmniRefPDDQADMethod
+  video_weight: 1.0
+  audio_weight: 1.0
+  rollout_fraction: 0.25
+  rollout_start_step: 20
+  # Half of the six forwards train per row, as OmniRef trains 4 of 8.
+  rungs_per_row: 3
+  eval_every: 15
+  full_eval_every: 30
+  patience: 2
+  fidelity_tolerance: 0.05
+  require_bitwise_init: true
+  save_best: true
+  select_metric: score
+  eval_save_dir: /mnt/nfs/vlm-aryan/pika-flex-qad/runs/ref2va-qad/samples
+  objective: output+recon
+  recon_lambda: 1.0
+  recon_normalize: step0
+  max_rollout_rows: 10
+  noise_floor_eval: true
+  step0_cache: /mnt/nfs/vlm-aryan/pika-flex-qad/runs/ref2va-qad/step0_cache.json
+
+training:
+  distributed:
+    num_gpus: 16
+    sp_size: 4
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 16
+    pin_cpu_memory: false
+  data:
+    train_batch_size: 1
+    seed: 1234
+  optimizer:
+    learning_rate: 5.0e-06
+    betas: [0.9, 0.999]
+    eps: 1.0e-08
+    weight_decay: 0.0
+    lr_scheduler: cosine_with_min_lr
+    lr_warmup_steps: 10
+    min_lr_ratio: 0.3
+  loop:
+    max_train_steps: 200
+    gradient_accumulation_steps: 1
+  checkpoint:
+    output_dir: /mnt/nfs/vlm-aryan/pika-flex-qad/runs/ref2va-qad
+    training_state_checkpointing_steps: 25
+    checkpoints_total_limit: 2
+    resume_from_checkpoint: ''
+  tracker:
+    trackers: [wandb]
+    project_name: pika-h3-flex-nvfp4
+    run_name: pika-flex-ref2va-nvfp4-qad
+  model:
+    enable_gradient_checkpointing_type: full
+  dit_precision: fp32
+  vsa_sparsity: 0.0
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+
+pipeline: {}

scripts/distill/pika_h3_flex_qad/slinky/smoke-ref2va.yaml

@@ -0,0 +1,126 @@
+# Pika H3 Flex Ref2VA NVFP4 QAD (the FastH3 QAD recipe on Flex's six-forward ladder).
+#
+# Teacher: Pika H3 Flex Ref2VA in bf16, i.e. MiniMax-H3 `transformer_ref` with the
+# released `ref2va` adapter merged by
+# scripts/distill/pika_h3_flex_qad/merge_pika_adapter.py. Student: the same
+# weights with NVFP4 W4A4 q/k/v/out plus FFN fc_in/fc_out in all 50 blocks
+# active from step 0, static FFN activation scales from the Flex calibration
+# table. Attention stays BF16 dense: Flex has no VSA gate, so there is no FP4
+# sparse-attention stage and no gate to freeze. Norms, AdaLN, the token refiner
+# and the FP32 boundary modules stay frozen.
+#
+# The objective, evaluation, guards and checkpoint contract are the shared
+# recipe's (docs/training/minimax_h3_fp4_qad.md). Rows are OmniRef Ref2VA
+# parquet rows; the teacher computes every target on the fly, so the OmniRef
+# held-out eval manifest is reused as is.
+#
+# Paths are placeholders: the merged model, the calibration table (pin its
+# sha256 once calibrated), the row manifests and the output dir.
+models:
+  student:
+    _target_: fastvideo.train.models.minimax_h3.flex.MiniMaxH3FlexRef2VAModel
+    init_from: /mnt/nfs/vlm-aryan/pika-flex-qad/models/pika-flex-ref2va-merged
+    trainable: true
+    enable_gradient_checkpointing_type: full
+    qad:
+      enabled: true
+      linears:
+        enabled: true
+        amax_json: null
+        provisional_calibration_rows: 2
+        amax_sha256: null
+        # Attention linears use the unit activation scale; only the FFN
+        # linears read their calibrated amax from this table.
+        attention: true
+        ffn: true
+        skip_blocks: []
+      fp4_attention:
+        enabled: false
+    data:
+      manifests:
+      - /mnt/lustre/vlm-wlsaidhi/fastvideo/data/omniref-v4-processed/20260919T001000Z/encoded/full/MANIFEST.jsonl
+      cases:
+      - first_frame
+      - full_scene_audio_image_reference
+      per_group:
+        480p: 400
+        768p: 50
+      seed: 20261007
+      max_frames: 124
+      max_train_rows: 4
+      eval_manifest: /mnt/nfs/vlm-aryan/h3-nvfp4-decoder/gen-latents/omniref-eval/manifest.json
+      eval_resolutions:
+      - 480p
+      eval_per_group: 1
+      eval_max_frames: 0
+  teacher:
+    _target_: fastvideo.train.models.minimax_h3.flex.MiniMaxH3FlexRef2VAModel
+    init_from: /mnt/nfs/vlm-aryan/pika-flex-qad/models/pika-flex-ref2va-merged
+    trainable: false
+    frozen_bf16: true
+
+method:
+  _target_: fastvideo.train.methods.knowledge_distillation.pdd_qad.OmniRefPDDQADMethod
+  video_weight: 1.0
+  audio_weight: 1.0
+  rollout_fraction: 0.25
+  rollout_start_step: 20
+  # Half of the six forwards train per row, as OmniRef trains 4 of 8.
+  rungs_per_row: 3
+  eval_every: 2
+  full_eval_every: 2
+  patience: 2
+  fidelity_tolerance: 0.05
+  require_bitwise_init: true
+  save_best: true
+  select_metric: score
+  eval_save_dir: /mnt/nfs/vlm-aryan/pika-flex-qad/runs/smoke-ref2va/samples
+  objective: output+recon
+  recon_lambda: 1.0
+  recon_normalize: step0
+  max_rollout_rows: 2
+  noise_floor_eval: false
+  step0_cache: /mnt/nfs/vlm-aryan/pika-flex-qad/runs/smoke-ref2va/step0_cache.json
+
+training:
+  distributed:
+    num_gpus: 4
+    sp_size: 4
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+    pin_cpu_memory: false
+  data:
+    train_batch_size: 1
+    seed: 1234
+  optimizer:
+    learning_rate: 5.0e-06
+    betas: [0.9, 0.999]
+    eps: 1.0e-08
+    weight_decay: 0.0
+    lr_scheduler: cosine_with_min_lr
+    lr_warmup_steps: 10
+    min_lr_ratio: 0.3
+  loop:
+    max_train_steps: 2
+    gradient_accumulation_steps: 1
+  checkpoint:
+    output_dir: /mnt/nfs/vlm-aryan/pika-flex-qad/runs/smoke-ref2va
+    training_state_checkpointing_steps: 25
+    checkpoints_total_limit: 2
+    resume_from_checkpoint: ''
+  tracker:
+    trackers: []
+    project_name: pika-h3-flex-nvfp4
+    run_name: pika-flex-ref2va-nvfp4-qad
+  model:
+    enable_gradient_checkpointing_type: full
+  dit_precision: fp32
+  vsa_sparsity: 0.0
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+
+pipeline: {}

scripts/distill/pika_h3_flex_qad/slinky/smoke.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+# 1-node QAD smoke: usage smoke.sh <config> <out>   (runs under env.sh)
+CFG=$1; OUT=$2; shift 2; rm -rf $OUT; mkdir -p $OUT
+FASTVIDEO_PG_TIMEOUT_MIN=120 $PY -m torch.distributed.run --nproc_per_node=4 --master_port=29571 --log-dir $OUT/torchlogs --redirects 3 --tee 3 -m fastvideo.train.entrypoint.train --config $CFG --training.checkpoint.output_dir $OUT --method.eval_save_dir $OUT/samples --method.step0_cache $OUT/step0.json "$@"