[test] MiniMax-H3: use supported graph-reuse geometry

# Conflicts:
#	fastvideo/envs.py

docs/design/inference_schema_parity_inventory.yaml

@@ -76,6 +76,10 @@ surfaces:
       lora_target_modules: "Legacy LoRA configuration surface pending dedicated component API."
       output_type: "Legacy output formatting surface pending GenerationResult cleanup."
       VSA_sparsity: "Model-specific inference optimization not yet represented in the typed public schema."
+      VSA_tile_size: "VSA-H3 tile geometry request; model-specific optimization not yet represented in the typed public schema."
+      vae_parallel_decode: "MiniMax-H3 sequence-parallel VAE decode opt-in; model-specific optimization not yet represented in the typed public schema."
+      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."
       moba_config_path: "Model-specific MoBA optimization surface not yet represented in the typed public schema."
       master_port: "Executor/bootstrap compatibility field; not part of the canonical inference schema."

docs/inference/optimizations.md

@@ -337,7 +337,38 @@ Only DiT submodules that declare `_compile_conditions` are compiled
 (most shipped models). The text encoder and VAE are not compiled by this
 flag.
 
-### What to expect
+### Regional fullgraph compile (experimental)
+
+`inference_torch_compile` is a stricter, kwargs-free variant that ports the
+training-side regional compile of
+[#1718](https://github.com/hao-ai-lab/FastVideo/pull/1718) to inference: the
+loader wraps each `_compile_conditions` block in
+`torch.compile(fullgraph=True)` with inductor
+`options={"emulate_precision_casts": True}` right after the transformer
+loads. The ordinary compile path keeps its historical compiler-disabled
+attention boundary by default; regional compile opts in only the compatible
+attention instances owned by this transformer. VSA, FLASH_ATTN on flash-attn
+3, and the explicit `FASTVIDEO_DISABLE_ATTENTION_COMPILE=1` escape hatch keep
+the transformer eager with one warning instead of failing mid-denoise.
+
+```python
+generator = VideoGenerator.from_pretrained(
+    "MiniMaxAI/MiniMax-H3",
+    inference_torch_compile=True,  # or FASTVIDEO_INFERENCE_TORCH_COMPILE=1
+)
+```
+
+Do not combine it with `torch_compile_kwargs['mode']` (the loader injects
+inductor options, and torch.compile forbids mode+options); it is
+independent of `enable_torch_compile`, and when both are set the regional
+compile wins for the DiT.
+
+### What to expect from generic compile
+
+The Wan result below measures the existing generic
+`enable_torch_compile=True` path. It is useful evidence that compile can help,
+but it is **not** a benchmark or numerical gate for the stricter regional
+fullgraph path above.
 
 | Config | Effect |
 |---|---|
@@ -353,6 +384,16 @@ generations with the same input shapes. Always exclude the first
 (warmup) generation when measuring steady-state latency — measuring the
 warmup is the most common way to wrongly conclude "compile is slower".
 
+**Regional MiniMax-H3 accuracy caveat (job 2660).** On one GB200 at
+768×1344×124, the native 50-point schedule ran exactly 49 transformer
+forwards. After one warmup, three fixed-prompt/fixed-seed repeats averaged
+**185.08s → 157.01s end to end** and **174.90s → 147.12s denoising**. Each
+leg was independently pixel-deterministic, but compiled output did **not**
+match eager: mean absolute pixel error **20.247/255**, PSNR **16.67 dB**,
+mean SSIM **0.7108**, and mean MS-SSIM **0.6370** across 124 frames. Treat
+regional MiniMax-H3 compile as an opt-in performance experiment, not an
+eager-parity-safe mode.
+
 **Numerics.** Inductor's lowering is designed to preserve eager
 semantics within floating-point tolerance, but per-model equivalence is
 not asserted by any standing SSIM regression here — the SSIM tests in
@@ -371,9 +412,13 @@ MS-SSIM gate on *your* config, especially when combining
   hao-ai-lab/FastVideo#1365 — keep that fix to get a clean compiled
   region under the default offload path.
 - **`mode="reduce-overhead"` / CUDA graphs**: not yet supported
-  end-to-end. The attention dispatch is an untraceable custom op and
-  still breaks the graph, which CUDA-graph trees cannot span. Use the
-  default inductor mode (shown above) until that is resolved.
+  by regional compile because that path injects inductor `options`, and
+  PyTorch rejects `mode` together with `options`. The generic compile path
+  can accept `mode`, but CUDA-graph compatibility remains backend- and
+  shape-dependent. FA2/FA3 inference and FA4 expose traceable custom-op
+  boundaries; VSA and the FA3 grad-enabled path remain outside the regional
+  support envelope. Use the default inductor mode shown above unless your
+  exact configuration has its own gate.
 
 Extra `torch.compile` options are passed through `torch_compile_kwargs`
 (a dict), accepted by `VideoGenerator.from_pretrained(...)` and by the

examples/inference/basic/basic_fasth3.py

@@ -85,6 +85,12 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--compile-mode",
                         default=None,
                         help='torch.compile mode, e.g. "reduce-overhead" for CUDA graphs')
+    parser.add_argument("--inference-torch-compile",
+                        action="store_true",
+                        help="regional fullgraph torch.compile of each DiT block after load. NOTE: this "
+                        "script always runs the VSA-H3 backend, which is not fullgraph-traceable — the "
+                        "loader logs one warning and keeps the transformer eager. The flag is exposed "
+                        "here to exercise exactly that guard")
     parser.add_argument("--repeats",
                         type=int,
                         default=1,
@@ -119,6 +125,8 @@ def main() -> None:
     }
     if args.vsa_sparsity > 0.0:
         experimental["VSA_sparsity"] = args.vsa_sparsity
+    if args.inference_torch_compile:
+        experimental["inference_torch_compile"] = True
 
     generator = VideoGenerator.from_config(
         GeneratorConfig(

examples/inference/basic/basic_minimax_h3_t2v.py

@@ -15,6 +15,7 @@ from fastvideo.api import (
     OffloadConfig,
     OutputConfig,
     ParallelismConfig,
+    PipelineSelection,
     SamplingConfig,
 )
 
@@ -41,6 +42,12 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--compile-mode",
                         default=None,
                         help='torch.compile mode, e.g. "reduce-overhead" for CUDA graphs')
+    parser.add_argument("--inference-torch-compile",
+                        action="store_true",
+                        help="regional fullgraph torch.compile of each DiT block after load (the #1718 "
+                        "training-port semantics: no kwargs; fullgraph + emulate_precision_casts injected). "
+                        "First generation pays the inductor JIT (~1-2 min); use --repeats >= 2 and time "
+                        "the last repeat. FASTVIDEO_INFERENCE_TORCH_COMPILE=1 is equivalent")
     parser.add_argument("--repeats",
                         type=int,
                         default=1,
@@ -54,9 +61,16 @@ def main() -> None:
     output_dir = Path(args.output)
     output_dir.mkdir(parents=True, exist_ok=True)
 
+    # Boot-time run configuration folded into FastVideoArgs (the same
+    # experimental-dict route basic_fasth3.py uses for the VSA knobs).
+    experimental: dict[str, object] = {}
+    if args.inference_torch_compile:
+        experimental["inference_torch_compile"] = True
+
     generator = VideoGenerator.from_config(
         GeneratorConfig(
             model_path=args.model_path,
+            pipeline=PipelineSelection(experimental=experimental),
             engine=EngineConfig(
                 num_gpus=args.num_gpus,
                 use_fsdp_inference=args.num_gpus > 1,

fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py

@@ -84,6 +84,107 @@ def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool:
     return True
 
 
+@torch.library.custom_op(
+    "fastvideo_kernel::block_sparse_attn_sm100a_inference",
+    mutates_args=(),
+    device_types="cuda",
+)
+def _block_sparse_attn_sm100a_inference(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    q2k_idx: torch.Tensor,
+    q2k_num: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> torch.Tensor:
+    """Opaque no-LSE launch used by the inference-only sm_100a route.
+
+    The extension is exposed as a raw pybind function rather than a dispatcher
+    op. Calling it directly makes Dynamo descend through a Python/C++ boundary
+    that has no fake implementation, so ``torch.compile(fullgraph=True)``
+    cannot capture a sparse H3 block. Keep that boundary inside this custom op;
+    its inputs have already been normalized by the public wrapper below.
+    """
+    fwd = _FWD_BY_BLOCK[_block_size(q, variable_block_sizes)]
+    sm_scale = 1.0 / (q.shape[-1]**0.5)
+    res = fwd(q, k, v, None, q2k_idx, q2k_num, variable_block_sizes, sm_scale, False)
+    return res[0]
+
+
+@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_sm100a_inference")
+def _block_sparse_attn_sm100a_inference_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    q2k_idx: torch.Tensor,
+    q2k_num: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> torch.Tensor:
+    # The C++ binding allocates its output with torch::empty_like(q).
+    return torch.empty_like(q)
+
+
+@torch.library.custom_op(
+    "fastvideo_kernel::block_sparse_attn_sm100a_from_mask_inference",
+    mutates_args=(),
+    device_types="cuda",
+)
+def _block_sparse_attn_sm100a_from_mask_inference(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> torch.Tensor:
+    """Opaque mask compaction plus no-LSE sm_100a launch for inference.
+
+    H3 naturally produces a bool block map. Its Triton ``map_to_index`` call
+    must live behind the same opaque boundary as the raw pybind launch;
+    otherwise Dynamo sees that kernel before reaching the index-native custom
+    op and full-graph capture still fails.
+    """
+    from fastvideo_kernel.triton_kernels.index import map_to_index
+
+    q2k_idx, q2k_num = map_to_index(block_map)
+    return _block_sparse_attn_sm100a_inference(
+        q,
+        k,
+        v,
+        q2k_idx.to(torch.int32).contiguous(),
+        q2k_num.to(torch.int32).contiguous(),
+        variable_block_sizes,
+    )
+
+
+@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_sm100a_from_mask_inference")
+def _block_sparse_attn_sm100a_from_mask_inference_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> torch.Tensor:
+    return torch.empty_like(q)
+
+
+def block_sparse_attn_sm100a_from_mask(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> Tuple[torch.Tensor, None]:
+    """Inference forward from a bool block map, with compaction kept opaque."""
+    out = _block_sparse_attn_sm100a_from_mask_inference(
+        q.contiguous(),
+        k.contiguous(),
+        v.contiguous(),
+        block_map.to(torch.bool).contiguous(),
+        variable_block_sizes.to(torch.int32).contiguous(),
+    )
+    return out, None
+
+
 def block_sparse_attn_sm100a(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -92,13 +193,24 @@ def block_sparse_attn_sm100a(
     q2k_num: torch.Tensor,
     variable_block_sizes: torch.Tensor,
     need_lse: bool = True,
-) -> Tuple[torch.Tensor, torch.Tensor]:
+) -> Tuple[torch.Tensor, torch.Tensor | None]:
     """Forward pass. Returns ``(out, lse)``; ``out`` has q's layout."""
-    fwd = _FWD_BY_BLOCK[_block_size(q, variable_block_sizes)]
+    q = q.contiguous()
+    k = k.contiguous()
+    v = v.contiguous()
     idx = q2k_idx.to(torch.int32).contiguous()
     num = q2k_num.to(torch.int32).contiguous()
     vbs = variable_block_sizes.to(torch.int32).contiguous()
+
+    if not need_lse:
+        # This is the production inference path. The custom op keeps the raw
+        # pybind launch opaque to Dynamo while its fake kernel carries output
+        # metadata through full-graph capture.
+        return _block_sparse_attn_sm100a_inference(q, k, v, idx, num, vbs), None
+
+    # Preserve the established LSE-producing path for correctness tests and
+    # any future forward/backward pairing; only inference needs the opaque op.
+    fwd = _FWD_BY_BLOCK[_block_size(q, vbs)]
     sm_scale = 1.0 / (q.shape[-1]**0.5)
-    res = fwd(q.contiguous(), k.contiguous(), v.contiguous(), None,
-              idx, num, vbs, sm_scale, need_lse)
-    return (res[0], res[1]) if need_lse else (res[0], None)
+    res = fwd(q, k, v, None, idx, num, vbs, sm_scale, True)
+    return res[0], res[1]

fastvideo/attention/backends/flash_attn.py

@@ -10,6 +10,14 @@ from fastvideo.attention.utils.flash_attn_default import (
     flash_attn_func_compilable,
 )
 
+if fa_version == "4":
+    # The FA4 varlen wrapper is already a compile-safe custom op. Keep the
+    # import conditional so FA2/FA3 environments do not need flash_attn.cute.
+    from fastvideo.attention.utils.flash_attn_cute import (
+        flash_attn_varlen_func as flash_attn_varlen_func_compilable, )
+else:
+    flash_attn_varlen_func_compilable = None
+
 from fastvideo.attention.backends.abstract import (
     AttentionBackend,
     AttentionImpl,
@@ -38,6 +46,15 @@ except ImportError:
     _FA4_FP4_AVAILABLE = False
 
 _FA4_QUANT_OPS: tuple | None = None
+_FA4_PACKED_VARLEN_RECEIPT_LOGGED = False
+
+
+def _log_fa4_packed_varlen_receipt() -> None:
+    """Emit one grep-friendly route receipt per worker process."""
+    global _FA4_PACKED_VARLEN_RECEIPT_LOGGED
+    if not _FA4_PACKED_VARLEN_RECEIPT_LOGGED:
+        logger.info("MiniMax-H3 dense attention: FA4 packed-varlen inference route enabled")
+        _FA4_PACKED_VARLEN_RECEIPT_LOGGED = True
 
 
 def _import_fa4_quant_ops() -> tuple:
@@ -228,7 +245,21 @@ class FlashAttentionImpl(AttentionImpl):
     ) -> None:
         self.causal = causal
         self.softmax_scale = softmax_scale
-        self.nvfp4_fa4 = extra_impl_args.get("nvfp4_fa4", False) or os.environ.get("FASTVIDEO_NVFP4_FA4", "0") == "1"
+        # MiniMax-H3's dense DiT explicitly enables this faster FA4 entry
+        # point. It remains off for every other model and for the H3 text
+        # refiner; grad-enabled calls stay on the established fixed path.
+        self.fa4_packed_varlen = bool(extra_impl_args.get("fa4_packed_varlen", False))
+        if self.fa4_packed_varlen and fa_version == "4":
+            _log_fa4_packed_varlen_receipt()
+        # An explicit ``nvfp4_fa4`` impl arg wins over the process-wide
+        # FASTVIDEO_NVFP4_FA4 env opt-in, so precision-sensitive layers (e.g.
+        # the FP32-pinned H3 VAE attention) can force-disable FP4 Q/K
+        # quantization while the DiT keeps it. When the arg is absent the env
+        # keeps its previous semantics.
+        nvfp4_fa4 = extra_impl_args.get("nvfp4_fa4")
+        if nvfp4_fa4 is None:
+            nvfp4_fa4 = os.environ.get("FASTVIDEO_NVFP4_FA4", "0") == "1"
+        self.nvfp4_fa4 = bool(nvfp4_fa4)
         if self.nvfp4_fa4:
             cap = torch.cuda.get_device_capability()
             assert cap in [(10, 0), (10, 3)], (f"NVFP4 FA4 requires Blackwell (sm100a/sm103a), got sm{cap[0]}{cap[1]}")
@@ -323,6 +354,29 @@ class FlashAttentionImpl(AttentionImpl):
         elif self.nvfp4_fa4:
             output = self._forward_nvfp4(query, key, value)
 
+        elif (self.fa4_packed_varlen and fa_version == "4" and not torch.is_grad_enabled() and query.shape[0] == 1
+              and query.shape[1] == key.shape[1] == value.shape[1]):
+            # FA4's packed-varlen entry point is materially faster for H3's
+            # long, single-document self-attention. Flatten only the batch
+            # dimension and describe that one sequence with CUDA int32
+            # cumulative lengths; the existing custom-op wrapper keeps this
+            # route traceable under torch.compile(fullgraph=True).
+            assert flash_attn_varlen_func_compilable is not None
+            sequence_length = query.shape[1]
+            cu_seqlens = torch.arange(2, dtype=torch.int32, device=query.device) * sequence_length
+            output = flash_attn_varlen_func_compilable(
+                query.squeeze(0),
+                key.squeeze(0),
+                value.squeeze(0),
+                cu_seqlens,
+                cu_seqlens,
+                sequence_length,
+                sequence_length,
+                dropout_p=0.0,
+                softmax_scale=self.softmax_scale,
+                causal=self.causal,
+            ).unsqueeze(0)
+
         else:
             # Route through the compilable wrapper so dynamo sees a
             # registered op (no graph break) for FA2/FA3; identical

fastvideo/attention/backends/video_sparse_attn_h3.py

@@ -233,6 +233,10 @@ class MiniMaxH3VSAMetadata(AttentionMetadata):
     exempt: bool
     variable_block_sizes: torch.Tensor
     untile_combined_index: torch.Tensor
+    # Device-side copy of ``dense_layers``. Regional fullgraph capture uses
+    # this tensor with each implementation's tensor-valued layer index so the
+    # shared block code does not specialize once per Python ``layer_idx``.
+    dense_layers_tensor: torch.Tensor
     # tokens per tile (256 or 64); selects the tile geometry AND the kernel
     # route in forward() (256 -> VSA-256 CuTe/Triton, 64 -> native Triton)
     tile_elems: int = _TILE_ELEMS
@@ -278,6 +282,7 @@ class MiniMaxH3VSAMetadataBuilder(AttentionMetadataBuilder):
         (_tile_partition_indices, variable_block_sizes, untile_combined_index, num_prefix_tiles,
          num_video_tiles) = _h3_tile_geometry(prefix_segments, dit_seq_shape, device, tile_shape)
 
+        dense_layers = tuple(int(layer) for layer in dense_layers)
         return MiniMaxH3VSAMetadata(
             current_timestep=current_timestep,
             VSA_sparsity=VSA_sparsity,
@@ -288,7 +293,8 @@ class MiniMaxH3VSAMetadataBuilder(AttentionMetadataBuilder):
             variable_block_sizes=variable_block_sizes,
             untile_combined_index=untile_combined_index,
             tile_elems=int(tile_size),
-            dense_layers=tuple(int(layer) for layer in dense_layers),
+            dense_layers=dense_layers,
+            dense_layers_tensor=torch.tensor(dense_layers, device=device, dtype=torch.int64),
             tile_buf_holder=self._tile_buf_holder,
         )
 
@@ -366,6 +372,58 @@ class MiniMaxH3VSAImpl(AttentionImpl):
     ) -> None:
         self.prefix = prefix
         self.layer_idx = layer_idx_from_prefix(prefix, default=-1)
+        self.head_size = head_size
+        # None means the regional-compile preparation hook has not run.  The
+        # eager path deliberately ignores this cache and preserves its
+        # request-time env/probe/fallback behavior; only Dynamo capture reads
+        # the prepared, static route.
+        self._regional_compile_sm100a_enabled: bool | None = None
+        self._regional_compile_layer_idx: torch.Tensor | None = None
+
+    def prepare_for_regional_compile(self, device: torch.device) -> None:
+        """Resolve the inference-only sm_100a route before fullgraph capture.
+
+        The ordinary eager route probes the environment, extension, device,
+        and tensor contract at every call so it can warn and fall back.  Those
+        Python/device-capability checks are not safe inside a regional
+        ``fullgraph=True`` block.  Probe one representative tile-64 input on
+        the loaded model's device now, then let ``forward`` specialize on the
+        resulting plain bool while Dynamo is compiling.
+        """
+        # Dynamo specializes Python integer attributes of nn.Module-owned
+        # objects. All 50 H3 blocks share one forward code object, so reading
+        # ``self.layer_idx`` in every regional graph exhausts the recompile
+        # cache. Tensor contents stay dynamic and let the blocks share a graph.
+        self._regional_compile_layer_idx = torch.tensor(self.layer_idx, device=device, dtype=torch.int64)
+
+        requested = os.environ.get(VSA_SM100A_ENV, "0") == "1"
+        enabled = False
+        reason = None
+        if requested:
+            if _sm100a is None:
+                reason = "fastvideo_kernel.block_sparse_attn_sm100a is not installed"
+            elif not callable(getattr(_sm100a, "block_sparse_attn_sm100a_from_mask", None)):
+                reason = ("fastvideo_kernel.block_sparse_attn_sm100a_from_mask is not installed; "
+                          "rebuild fastvideo-kernel with the compile-safe mask entry")
+            else:
+                # Two 64-token blocks exercise the exact sm_100a inference
+                # specialization while keeping the one-time probe tiny.  The
+                # kernel predicate checks extension presence, CUDA capability,
+                # dtype/layout, head size, block size, and even block count
+                # without reading metadata tensor contents.
+                probe_query = torch.empty((1, 1, 128, self.head_size), device=device, dtype=torch.bfloat16)
+                probe_block_sizes = torch.full((2, ), 64, device=device, dtype=torch.int32)
+                reason = _sm100a_unavailable_reason(
+                    _sm100a,
+                    probe_query,
+                    probe_block_sizes,
+                    grad_mode=False,
+                )
+                enabled = reason is None
+
+        self._regional_compile_sm100a_enabled = enabled
+        if requested and reason is not None:
+            logger.warning_once(f"{VSA_SM100A_ENV}=1 but falling back to the Triton-64 kernels: {reason}")
 
     def tile(self, x: torch.Tensor, attn_metadata: MiniMaxH3VSAMetadata) -> torch.Tensor:
         """Scatter rows into the padded tile buffer (pad positions stay zero).
@@ -408,10 +466,26 @@ class MiniMaxH3VSAImpl(AttentionImpl):
         elif block_sparse_attn_256_bshd is None:
             raise NotImplementedError("fastvideo_kernel.block_sparse_attn_256 is not installed")
 
-        # probe-guided per-layer opt-out: diffuse layers run dense (all-True
-        # mask) while the rest keep the configured sparsity
-        layer_sparsity = 0.0 if self.layer_idx in attn_metadata.dense_layers else attn_metadata.VSA_sparsity
-        probe_dir = probe_enabled()
+        # Probe recording performs filesystem writes and host synchronizations,
+        # so the loader keeps probe-enabled runs eager. Avoid even reading that
+        # environment switch while Dynamo captures a regional full graph.
+        compiling = torch.compiler.is_compiling()
+
+        # Probe-guided per-layer opt-out: diffuse layers run dense (all-True
+        # mask) while the rest keep the configured sparsity. During regional
+        # capture, keep the layer decision tensor-valued so the 50 block
+        # instances reuse one graph instead of specializing on layer_idx.
+        force_dense = None
+        if compiling:
+            if self._regional_compile_layer_idx is None:
+                raise RuntimeError(
+                    "VSA-H3 layer routing was not resolved before torch.compile; "
+                    "call prepare_for_regional_compile(device) on every MiniMaxH3VSAImpl after loading weights.")
+            force_dense = (attn_metadata.dense_layers_tensor == self._regional_compile_layer_idx).any()
+            layer_sparsity = attn_metadata.VSA_sparsity
+        else:
+            layer_sparsity = 0.0 if self.layer_idx in attn_metadata.dense_layers else attn_metadata.VSA_sparsity
+        probe_dir = None if compiling else probe_enabled()
 
         scores = None
         if layer_sparsity > 0.0 or gate_compress is not None or probe_dir is not None:
@@ -432,6 +506,10 @@ class MiniMaxH3VSAImpl(AttentionImpl):
                 layer_sparsity,
                 attn_metadata.exempt,
             )
+        if force_dense is not None:
+            # A scalar bool tensor broadcasts over the block map. This exactly
+            # preserves the eager dense-layer contract without a Python branch.
+            mask = mask | force_dense
 
         if tile_elems == 64:
             # Native 64-token path: the block map is already at the kernels'
@@ -447,34 +525,55 @@ class MiniMaxH3VSAImpl(AttentionImpl):
             # so autograd keeps the Triton fwd+bwd pairing untouched. The
             # kernel does return an LSE in Triton's M format, so a future
             # fwd/bwd pairing is possible, but it is not built here.
+            grad_mode = torch.is_grad_enabled() and (query.requires_grad or key.requires_grad or value.requires_grad)
             use_sm100a = False
-            if os.environ.get(VSA_SM100A_ENV, "0") == "1":
-                grad_mode = torch.is_grad_enabled() and (query.requires_grad or key.requires_grad
-                                                         or value.requires_grad)
+            if compiling:
+                if self._regional_compile_sm100a_enabled is None:
+                    raise RuntimeError(
+                        "VSA-H3 sm_100a routing was not resolved before torch.compile; "
+                        "call prepare_for_regional_compile(device) on every MiniMaxH3VSAImpl after loading weights.")
+                # The preparation probe established module/device/kernel
+                # support.  Keep only static tensor/geometry facts here; no
+                # env access, device-capability query, or is_supported call may
+                # enter the Dynamo graph.
+                use_sm100a = (self._regional_compile_sm100a_enabled and not grad_mode and q_bhsd.dtype == torch.bfloat16
+                              and q_bhsd.shape[-1] == 128 and attn_metadata.variable_block_sizes.numel() % 2 == 0)
+            elif os.environ.get(VSA_SM100A_ENV, "0") == "1":
                 reason = _sm100a_unavailable_reason(_sm100a, q_bhsd, attn_metadata.variable_block_sizes, grad_mode)
                 if reason is None and map_to_index is None:
                     reason = "fastvideo_kernel.triton_kernels.index (map_to_index) is not importable"
                 if reason is None:
                     use_sm100a = True
-                elif not torch.compiler.is_compiling():
+                else:
                     logger.warning_once(f"{VSA_SM100A_ENV}=1 but falling back to the Triton-64 kernels: {reason}")
 
             if use_sm100a:
-                # The sm_100a entry is index-native; compact the bool map the
-                # same way the Triton bool entry does internally. Per-row
-                # counts are NON-uniform here (prefix query tiles are dense,
-                # video tiles run prefix+top-k) -- legal for the fixed kernel,
-                # silently wrong on the pre-fix upstream one.
-                q2k_idx, q2k_num = map_to_index(mask)
-                out_bhsd, _ = _sm100a.block_sparse_attn_sm100a(
-                    q_bhsd,
-                    k_bhsd,
-                    v_bhsd,
-                    q2k_idx,
-                    q2k_num,
-                    attn_metadata.variable_block_sizes.to(torch.int32),
-                    need_lse=False,
-                )
+                if compiling:
+                    # The compile-safe wrapper keeps both Triton mask
+                    # compaction and the raw pybind launch behind one
+                    # fake-backed custom-op boundary.
+                    out_bhsd, _ = _sm100a.block_sparse_attn_sm100a_from_mask(
+                        q_bhsd,
+                        k_bhsd,
+                        v_bhsd,
+                        mask,
+                        attn_metadata.variable_block_sizes,
+                    )
+                else:
+                    # Preserve the established eager/index-native route and
+                    # compatibility with older kernel wheels. Per-row counts
+                    # are non-uniform (prefix queries are dense; video queries
+                    # run prefix+top-k), which the fixed kernel supports.
+                    q2k_idx, q2k_num = map_to_index(mask)
+                    out_bhsd, _ = _sm100a.block_sparse_attn_sm100a(
+                        q_bhsd,
+                        k_bhsd,
+                        v_bhsd,
+                        q2k_idx,
+                        q2k_num,
+                        attn_metadata.variable_block_sizes.to(torch.int32),
+                        need_lse=False,
+                    )
             else:
                 out_bhsd, _ = block_sparse_attn_64_bhsd(
                     q_bhsd,

fastvideo/attention/layer.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import os
+from functools import wraps
 
 import torch
 import torch.nn as nn
@@ -20,7 +21,7 @@ def _attention_compile_disabled() -> bool:
 
     Defaults to ``True`` (the historical behavior: attention runs eager via
     ``torch.compiler.disable``). Set ``FASTVIDEO_DISABLE_ATTENTION_COMPILE=0``
-    to let attention be traced/compiled into the surrounding graph.
+    to let attention instances constructed under that environment be traced.
     """
     val = os.environ.get("FASTVIDEO_DISABLE_ATTENTION_COMPILE")
     if val is None:
@@ -28,11 +29,32 @@ def _attention_compile_disabled() -> bool:
     return val.strip().lower() not in ("0", "false", "no", "off", "")
 
 
+def _attention_compile_explicitly_disabled() -> bool:
+    """Whether the environment explicitly requests the eager boundary.
+
+    Regional compile can override the historical default for one loaded
+    transformer, but it must still honor an explicit debugging escape hatch.
+    """
+    return "FASTVIDEO_DISABLE_ATTENTION_COMPILE" in os.environ and _attention_compile_disabled()
+
+
 def _maybe_compiler_disable(fn):
-    """Apply ``torch.compiler.disable`` unless disabled via env var."""
-    if _attention_compile_disabled():
-        return torch.compiler.disable(fn)
-    return fn
+    """Defer the eager/traceable choice to each attention instance.
+
+    A class-definition-time choice makes a process-wide default the only
+    option. The deferred wrapper keeps ordinary instances on the historical
+    eager boundary while allowing the regional loader to opt in only the
+    attention modules owned by the transformer it is compiling.
+    """
+    disabled_fn = torch.compiler.disable(fn)
+
+    @wraps(fn)
+    def _dispatch(self, *args, **kwargs):
+        if self._compile_forward_enabled:
+            return fn(self, *args, **kwargs)
+        return disabled_fn(self, *args, **kwargs)
+
+    return _dispatch
 
 
 class DistributedAttention(nn.Module):
@@ -77,6 +99,13 @@ class DistributedAttention(nn.Module):
         self.num_kv_heads = num_kv_heads
         self.backend = backend_name_to_enum(attn_backend.get_name())
         self.dtype = dtype
+        # Preserve the historical compiler-disabled default. The regional
+        # inference loader may enable this one instance after validating the
+        # transformer's resolved backend; no process-global default changes.
+        self._compile_forward_enabled = not _attention_compile_disabled()
+
+    def _set_compile_forward_enabled(self, enabled: bool) -> None:
+        self._compile_forward_enabled = enabled
 
     @_maybe_compiler_disable
     def forward(

fastvideo/envs.py

@@ -21,7 +21,11 @@ if TYPE_CHECKING:
     FASTVIDEO_TRACE_FUNCTION: int = 0
     FASTVIDEO_ATTENTION_BACKEND: str | None = None
     FASTVIDEO_FA4: bool = False
+    FASTVIDEO_INFERENCE_TORCH_COMPILE: bool = False
     FASTVIDEO_MINIMAX_H3_FUSIONS: str = ""
+    FASTVIDEO_VAE_PARALLEL_DECODE: bool = False
+    FASTVIDEO_VAE_PARALLEL_ENCODE: bool = False
+    FASTVIDEO_VAE_PARALLEL_DECODE_STRATEGY: str | None = None
     FASTVIDEO_WORKER_MULTIPROC_METHOD: str = "spawn"
     FASTVIDEO_TARGET_DEVICE: str = "cuda"
     MAX_JOBS: str | None = None
@@ -219,6 +223,30 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "FASTVIDEO_FA4":
     lambda: os.getenv("FASTVIDEO_FA4", "0") != "0",
 
+    # If set (=1), enable regional (per-transformer-block) fullgraph
+    # torch.compile for the DiT at inference — the inference-side counterpart
+    # of the training regional-compile port of hao-ai-lab/FastVideo#1718.
+    # Equivalent to FastVideoArgs.inference_torch_compile=True (e.g. via
+    # PipelineSelection.experimental={"inference_torch_compile": True}). VSA
+    # and other non-fullgraph-traceable attention backends degrade to eager
+    # with one warning; see _regional_compile_unsupported_reason in
+    # fastvideo/models/loader/fsdp_load.py.
+    "FASTVIDEO_INFERENCE_TORCH_COMPILE":
+    lambda: os.getenv("FASTVIDEO_INFERENCE_TORCH_COMPILE", "0") != "0",
+
+    # If set (=1), MiniMax-H3 VAE decode (and, with the ENCODE variant,
+    # reference-video encode) round-robins its temporal chunks across the
+    # sequence-parallel ranks instead of running serially on the output rank.
+    # Folded into FastVideoArgs.vae_parallel_decode / vae_parallel_encode at
+    # construction (parse-once). The STRATEGY variant picks the chunk
+    # transport collective: "gather" (default) or "all_gather".
+    "FASTVIDEO_VAE_PARALLEL_DECODE":
+    lambda: os.getenv("FASTVIDEO_VAE_PARALLEL_DECODE", "0") != "0",
+    "FASTVIDEO_VAE_PARALLEL_ENCODE":
+    lambda: os.getenv("FASTVIDEO_VAE_PARALLEL_ENCODE", "0") != "0",
+    "FASTVIDEO_VAE_PARALLEL_DECODE_STRATEGY":
+    lambda: os.getenv("FASTVIDEO_VAE_PARALLEL_DECODE_STRATEGY", None),
+
     # Opt-in MiniMax-H3 inference-only Triton fusions adapted from the
     # NVlabs/Sana Sol-Engine implementation. Accepts `all`, `1`, or a
     # comma-separated subset of `modulate,qknorm_rope,swiglu`. An empty value

fastvideo/fastvideo_args.py

@@ -146,6 +146,19 @@ class FastVideoArgs:
     vae_cpu_offload: bool = True
     pin_cpu_memory: bool = True
 
+    # Sequence-parallel MiniMax-H3 VAE (opt-in, default off). With SP > 1 the
+    # video VAE's temporal chunks (decode) and clips (reference encode) are
+    # round-robined across the sequence-parallel ranks and reassembled
+    # bit-exactly on the group's first rank instead of running serially on
+    # one rank while the others idle. ``__post_init__`` folds the
+    # FASTVIDEO_VAE_PARALLEL_DECODE / FASTVIDEO_VAE_PARALLEL_ENCODE env vars
+    # into these fields (parse-once, like attention_backend), and
+    # FASTVIDEO_VAE_PARALLEL_DECODE_STRATEGY overrides the chunk-transport
+    # collective ("gather" or "all_gather").
+    vae_parallel_decode: bool = False
+    vae_parallel_encode: bool = False
+    vae_parallel_decode_strategy: str | None = None
+
     # Compilation
     # ``enable_torch_compile`` covers the DiT path (transformer,
     # transformer_2, and the LTX-2 stage-2 transformer_refine).
@@ -164,6 +177,18 @@ class FastVideoArgs:
     torch_compile_kwargs_text_encoder: dict[str, Any] = field(default_factory=dict)
     torch_compile_kwargs_vae: dict[str, Any] = field(default_factory=dict)
     torch_compile_kwargs_audio_vae: dict[str, Any] = field(default_factory=dict)
+    # Regional (per-transformer-block) fullgraph torch.compile of the DiT at
+    # inference — the inference-side counterpart of the training regional
+    # compile ported from hao-ai-lab/FastVideo#1718. Applied by the loader
+    # right after the transformer loads, with fullgraph=True and inductor
+    # options {emulate_precision_casts: True} injected (no user kwargs
+    # needed). Attention backends that cannot be fullgraph-traced (VSA,
+    # FLASH_ATTN on flash-attn 3) degrade the transformer to eager with one
+    # warning. Opt-in via FASTVIDEO_INFERENCE_TORCH_COMPILE=1 (folded in
+    # __post_init__) or PipelineSelection.experimental
+    # {"inference_torch_compile": true}. Distinct from ``enable_torch_compile``,
+    # which keeps the pipeline-level compile semantics.
+    inference_torch_compile: bool = False
 
     disable_autocast: bool = False
 
@@ -272,6 +297,13 @@ class FastVideoArgs:
         self._apply_ltx2_vae_overrides()
         self._resolve_refine_args()
         self._apply_transformer_quant()
+        if not self.inference_torch_compile:
+            # Parse-once adapter (same pattern as attention_backend below): the
+            # environment variable is an input read once here, so the loader
+            # only ever consults the typed field.
+            import fastvideo.envs as envs
+            if envs.FASTVIDEO_INFERENCE_TORCH_COMPILE:
+                self.inference_torch_compile = True
         if self.attention_backend is not None:
             # Fail fast on typos instead of silently auto-selecting later.
             from fastvideo.attention.selector import coerce_attn_backend
@@ -287,8 +319,27 @@ class FastVideoArgs:
             env_backend = envs.FASTVIDEO_ATTENTION_BACKEND
             if env_backend is not None and backend_name_to_enum(env_backend) is not None:
                 self.attention_backend = env_backend
+        self._fold_vae_parallel_env()
         self.check_fastvideo_args()
 
+    def _fold_vae_parallel_env(self) -> None:
+        """Parse-once adapters for the sequence-parallel VAE env vars."""
+        import fastvideo.envs as envs
+
+        # Mirrors fastvideo.models.vaes.minimax_h3_parallel.DECODE_GATHER_STRATEGIES /
+        # DEFAULT_DECODE_GATHER_STRATEGY (kept literal here so constructing args
+        # never imports model modules; a unit test pins the two in sync).
+        strategies = ("gather", "all_gather")
+        if not self.vae_parallel_decode and envs.FASTVIDEO_VAE_PARALLEL_DECODE:
+            self.vae_parallel_decode = True
+        if not self.vae_parallel_encode and envs.FASTVIDEO_VAE_PARALLEL_ENCODE:
+            self.vae_parallel_encode = True
+        if self.vae_parallel_decode_strategy is None:
+            self.vae_parallel_decode_strategy = envs.FASTVIDEO_VAE_PARALLEL_DECODE_STRATEGY or "gather"
+        if self.vae_parallel_decode_strategy not in strategies:
+            raise ValueError(f"vae_parallel_decode_strategy must be one of {strategies}, "
+                             f"got {self.vae_parallel_decode_strategy!r}.")
+
     def _apply_transformer_quant(self) -> None:
         """Pin the typed ``transformer_quant`` instance onto ``dit_config``.
 
@@ -592,6 +643,15 @@ class FastVideoArgs:
             help=
             "JSON string of kwargs to pass to torch.compile. Example: '{\"backend\":\"inductor\",\"mode\":\"reduce-overhead\"}'",
         )
+        parser.add_argument(
+            "--inference-torch-compile",
+            action=StoreBoolean,
+            default=FastVideoArgs.inference_torch_compile,
+            help="Regional fullgraph torch.compile of each DiT transformer block at inference "
+            "(port of the #1718 training-side regional compile). The loader injects fullgraph=True "
+            "and inductor options {emulate_precision_casts: true}; non-traceable attention backends "
+            "(VSA) degrade to eager with one warning. FASTVIDEO_INFERENCE_TORCH_COMPILE=1 is equivalent.",
+        )
 
         parser.add_argument(
             "--dit-cpu-offload",
@@ -632,6 +692,18 @@ class FastVideoArgs:
             "Pin memory for CPU offload. Only added as a temp workaround if it throws \"CUDA error: invalid argument\". "
             "Should be enabled in almost all cases",
         )
+        parser.add_argument(
+            "--vae-parallel-decode",
+            action=StoreBoolean,
+            help="With sequence parallelism, round-robin MiniMax-H3 VAE decode chunks across the SP ranks "
+            "and reassemble bit-exactly on the output rank (default: serial decode on the output rank)",
+        )
+        parser.add_argument(
+            "--vae-parallel-encode",
+            action=StoreBoolean,
+            help="With sequence parallelism, round-robin MiniMax-H3 reference-video VAE encode clips across "
+            "the SP ranks; every rank keeps the identical full encoding (default: serial encode on every rank)",
+        )
         parser.add_argument(
             "--disable-autocast",
             action=StoreBoolean,

fastvideo/logger.py

@@ -114,7 +114,10 @@ def _info(logger: Logger,
     is_local_main_process = local_rank == 0
 
     if (main_process_only and is_main_process) or (local_main_process_only and is_local_main_process):
-        logger.log(logging.INFO, msg, *args, stacklevel=2, **kwargs)
+        # Honor an explicit stacklevel (info_once routes through here with
+        # stacklevel already set) instead of passing the keyword twice.
+        stacklevel = kwargs.pop("stacklevel", 2)
+        logger.log(logging.INFO, msg, *args, stacklevel=stacklevel, **kwargs)
 
     global _warned_local_main_process, _warned_main_process
 

fastvideo/models/dits/minimax_h3.py

@@ -61,13 +61,13 @@ def _enabled_minimax_h3_fusions(value: str | None = None) -> frozenset[str]:
 
 
 def _can_run_minimax_h3_fusion(tensor: torch.Tensor) -> bool:
-    """Triton kernels are inference-only and stay outside Dynamo capture.
+    """Triton kernels are inference-only and trace as opaque custom ops.
 
     The ``HAVE_TRITON`` check makes the eager fallback exact: on a CUDA build
     whose Triton failed to import, an enabled fusion falls back instead of
     hitting the strict wrappers' hard RuntimeError mid-forward.
     """
-    return (HAVE_TRITON and tensor.is_cuda and not torch.is_grad_enabled() and not torch.compiler.is_compiling())
+    return HAVE_TRITON and tensor.is_cuda and not torch.is_grad_enabled()
 
 
 class MiniMaxH3RotaryPosEmbed(nn.Module):
@@ -144,6 +144,7 @@ class MiniMaxH3Attention(nn.Module):
         quant_config: QuantizationConfig | None,
         prefix: str,
         fuse_qknorm_rope: bool = False,
+        fa4_packed_varlen: bool = False,
     ) -> None:
         super().__init__()
         self.num_attention_heads = num_attention_heads
@@ -196,6 +197,7 @@ class MiniMaxH3Attention(nn.Module):
             causal=False,
             supported_attention_backends=supported_attention_backends,
             prefix=prefix,
+            fa4_packed_varlen=fa4_packed_varlen,
         )
         self.to_gate_compress: ReplicatedLinear | None = None
         # None = unchecked; the first forward tests the loaded weight once and
@@ -221,12 +223,31 @@ class MiniMaxH3Attention(nn.Module):
         """
         if torch.is_grad_enabled():
             return True
+        if self._gate_compress_active is None:
+            if torch.compiler.is_compiling():
+                raise RuntimeError(
+                    "MiniMax H3 VSA compression gate was not resolved before torch.compile; "
+                    "call prepare_for_compile() after loading weights.")
+            self._resolve_gate_compress_for_compile()
+        assert self._gate_compress_active is not None
+        return self._gate_compress_active
+
+    def _resolve_gate_compress_for_compile(self) -> None:
+        """Resolve the inference-only compression-gate branch eagerly.
+
+        Regional compilation wraps each transformer block with
+        ``fullgraph=True``. Resolving the loaded weight before capture keeps
+        the GPU-to-host bool conversion and the cache mutation out of every
+        compiled block graph. Grad-enabled forwards ignore this cached value
+        in ``_gate_active`` so training can still learn from a zero gate.
+        """
+        if self.to_gate_compress is None:
+            return
         if self._gate_compress_active is None:
             weight = self.to_gate_compress.weight
             # bool() on a DTensor reduction resolves collectively, so every
             # rank caches the same answer.
             self._gate_compress_active = bool((weight != 0).any())
-        return self._gate_compress_active
 
     @staticmethod
     def _apply_rotary_emb(
@@ -465,6 +486,10 @@ class MiniMaxH3TransformerBlock(nn.Module):
             quant_config,
             prefix=f"{prefix}.attn",
             fuse_qknorm_rope=fuse_qknorm_rope,
+            # The packed multimodal document is a single long self-attention
+            # sequence. FA4's varlen API is substantially faster for this
+            # shape; the backend keeps grad/training and non-FA4 calls fixed.
+            fa4_packed_varlen=True,
         )
         self.norm2 = nn.RMSNorm(hidden_size, eps=norm_eps)
         self.ff = MiniMaxH3FeedForward(
@@ -576,8 +601,8 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
         if self.enabled_fusions:
             if HAVE_TRITON:
                 logger.info(
-                    "MiniMax H3 inference fusions enabled: %s (CUDA inference-only; grad-enabled and "
-                    "torch.compile-captured forwards fall back to eager).",
+                    "MiniMax H3 inference fusions enabled: %s (CUDA inference-only; grad-enabled forwards "
+                    "fall back to eager; torch.compile captures opaque custom-op boundaries).",
                     ",".join(sorted(self.enabled_fusions)))
             else:
                 logger.warning(
@@ -712,16 +737,51 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
     def prepare_for_compile(self) -> None:
         """Pipeline hook, called once right before torch.compile wraps the blocks.
 
-        Dynamo capture traces the eager branch of every fusion guard, so an
-        enabled ``FASTVIDEO_MINIMAX_H3_FUSIONS`` set is silently inert inside
-        compiled block forwards (H3 compiles per-block by default). Say so
-        once instead of leaving the flag looking active.
+        Resolve each loaded VSA compression gate eagerly, then let the H3 VSA
+        implementation preselect its device-specific kernel route. This keeps
+        both choices static during fullgraph capture and prevents host-side
+        probes or cache mutation in Dynamo.
+
+        The inference-only Triton fusions expose fake-backed custom operators,
+        so Dynamo can keep them active as opaque nodes inside each fullgraph
+        block instead of tracing into their launcher implementation.
         """
+        compile_state = next(self.parameters(), None)
+        if compile_state is None:
+            compile_state = next(self.buffers(), None)
+        if compile_state is None:
+            raise RuntimeError("MiniMax H3 has no materialized parameter or buffer to identify its compile device.")
+        compile_device = compile_state.device
+
+        gate_states: list[bool] = []
+        prepared_vsa_impls = 0
+        for block in self.transformer_blocks:
+            attention = block.attn
+            if attention.to_gate_compress is not None:
+                attention._resolve_gate_compress_for_compile()
+                assert attention._gate_compress_active is not None
+                gate_states.append(attention._gate_compress_active)
+            prepare_vsa = getattr(attention.distributed_attention.attn_impl, "prepare_for_regional_compile", None)
+            if callable(prepare_vsa):
+                # Generic post-load quantization may replace to_q.weight with
+                # packed buffers before this hook runs. The transformer state
+                # remains materialized on the same local device.
+                prepare_vsa(compile_device)
+                prepared_vsa_impls += 1
+        if gate_states:
+            logger.info(
+                "Resolved MiniMax H3 VSA compression gates before torch.compile: %d active, %d inactive",
+                sum(gate_states),
+                len(gate_states) - sum(gate_states),
+            )
+        if prepared_vsa_impls:
+            logger.info("Prepared %d MiniMax H3 VSA attention implementations for regional torch.compile",
+                        prepared_vsa_impls)
         if self.enabled_fusions:
-            logger.warning(
-                "torch.compile is enabled for MiniMax H3, so the requested inference fusions (%s) are "
-                "inert inside compiled block forwards; the compiled eager path runs instead.",
-                ",".join(sorted(self.enabled_fusions)))
+            logger.info(
+                "MiniMax H3 inference fusions remain active under torch.compile through custom-op boundaries: %s",
+                ",".join(sorted(self.enabled_fusions)),
+            )
 
     def materialize_non_persistent_buffers(
         self,

fastvideo/models/dits/minimax_h3_fusions/modulation.py

@@ -202,7 +202,7 @@ def _row_addressable(table: torch.Tensor) -> torch.Tensor:
     return table if table.stride(-1) == 1 else table.contiguous()
 
 
-def fused_rmsnorm_modulate(
+def _fused_rmsnorm_modulate_impl(
     x: torch.Tensor,
     weight: torch.Tensor,
     scale: torch.Tensor,
@@ -210,13 +210,6 @@ def fused_rmsnorm_modulate(
     index: torch.Tensor,
     eps: float,
 ) -> torch.Tensor:
-    """Run RMSNorm and row-indexed modulation in one strict Triton kernel.
-
-    ``index`` values must lie in ``[0, table_rows)``. Unlike eager
-    ``index_select``, the kernel does not raise on out-of-range values (a
-    device-side bounds check would synchronize); callers are safe by
-    construction (``timestep_indices * 3 + token_tags``, SP pads with 0).
-    """
     _validate_contract(x, weight, (scale, shift), index, eps)
     _require_forward_only(x, weight, scale, shift)
     _require_triton_cuda(x)
@@ -248,7 +241,7 @@ def fused_rmsnorm_modulate(
     return output.view_as(x)
 
 
-def fused_residual_gate_rmsnorm_modulate(
+def _fused_residual_gate_rmsnorm_modulate_impl(
     residual: torch.Tensor,
     branch: torch.Tensor,
     gate: torch.Tensor,
@@ -258,11 +251,6 @@ def fused_residual_gate_rmsnorm_modulate(
     index: torch.Tensor,
     eps: float,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """Fuse residual update, row-indexed gate, RMSNorm, and modulation.
-
-    ``index`` values must lie in ``[0, table_rows)``; see
-    :func:`fused_rmsnorm_modulate` for why the wrapper does not check them.
-    """
     _validate_residual_branch(residual, branch)
     _validate_contract(residual, weight, (gate, scale, shift), index, eps)
     _require_forward_only(residual, branch, gate, weight, scale, shift)
@@ -300,3 +288,138 @@ def fused_residual_gate_rmsnorm_modulate(
         num_warps=_num_warps(block_size),
     )
     return hidden.view_as(residual), modulated.view_as(residual)
+
+
+# Dynamo must see the Triton launches as opaque nodes.  In eager mode the
+# public wrappers below keep their strict, actionable input validation; while
+# compiling, these custom-op boundaries avoid tracing into Triton's launcher
+# and let the surrounding H3 block remain one full graph.
+@torch.library.custom_op(
+    "fastvideo::_minimax_h3_rmsnorm_modulate",
+    mutates_args=(),
+    device_types="cuda",
+)
+def _fused_rmsnorm_modulate_op(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    index: torch.Tensor,
+    eps: float,
+) -> torch.Tensor:
+    return _fused_rmsnorm_modulate_impl(x, weight, scale, shift, index, eps)
+
+
+@torch.library.register_fake("fastvideo::_minimax_h3_rmsnorm_modulate")
+def _fused_rmsnorm_modulate_fake(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    index: torch.Tensor,
+    eps: float,
+) -> torch.Tensor:
+    del weight, scale, shift, index, eps
+    return x.new_empty(x.shape)
+
+
+@torch.library.custom_op(
+    "fastvideo::_minimax_h3_residual_gate_rmsnorm_modulate",
+    mutates_args=(),
+    device_types="cuda",
+)
+def _fused_residual_gate_rmsnorm_modulate_op(
+    residual: torch.Tensor,
+    branch: torch.Tensor,
+    gate: torch.Tensor,
+    weight: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    index: torch.Tensor,
+    eps: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    return _fused_residual_gate_rmsnorm_modulate_impl(
+        residual,
+        branch,
+        gate,
+        weight,
+        scale,
+        shift,
+        index,
+        eps,
+    )
+
+
+@torch.library.register_fake("fastvideo::_minimax_h3_residual_gate_rmsnorm_modulate")
+def _fused_residual_gate_rmsnorm_modulate_fake(
+    residual: torch.Tensor,
+    branch: torch.Tensor,
+    gate: torch.Tensor,
+    weight: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    index: torch.Tensor,
+    eps: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    del branch, gate, weight, scale, shift, index, eps
+    return residual.new_empty(residual.shape), residual.new_empty(residual.shape)
+
+
+def fused_rmsnorm_modulate(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    index: torch.Tensor,
+    eps: float,
+) -> torch.Tensor:
+    """Run RMSNorm and row-indexed modulation in one strict Triton kernel.
+
+    ``index`` values must lie in ``[0, table_rows)``. Unlike eager
+    ``index_select``, the kernel does not raise on out-of-range values (a
+    device-side bounds check would synchronize); callers are safe by
+    construction (``timestep_indices * 3 + token_tags``, SP pads with 0).
+    """
+    _require_forward_only(x, weight, scale, shift)
+    if torch.compiler.is_compiling():
+        return torch.ops.fastvideo._minimax_h3_rmsnorm_modulate(x, weight, scale, shift, index, eps)
+    return _fused_rmsnorm_modulate_impl(x, weight, scale, shift, index, eps)
+
+
+def fused_residual_gate_rmsnorm_modulate(
+    residual: torch.Tensor,
+    branch: torch.Tensor,
+    gate: torch.Tensor,
+    weight: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    index: torch.Tensor,
+    eps: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Fuse residual update, row-indexed gate, RMSNorm, and modulation.
+
+    ``index`` values must lie in ``[0, table_rows)``; see
+    :func:`fused_rmsnorm_modulate` for why the wrapper does not check them.
+    """
+    _require_forward_only(residual, branch, gate, weight, scale, shift)
+    if torch.compiler.is_compiling():
+        return torch.ops.fastvideo._minimax_h3_residual_gate_rmsnorm_modulate(
+            residual,
+            branch,
+            gate,
+            weight,
+            scale,
+            shift,
+            index,
+            eps,
+        )
+    return _fused_residual_gate_rmsnorm_modulate_impl(
+        residual,
+        branch,
+        gate,
+        weight,
+        scale,
+        shift,
+        index,
+        eps,
+    )

fastvideo/models/dits/minimax_h3_fusions/qknorm_rope.py

@@ -121,23 +121,13 @@ def _validate_inputs(
     return batch, seq_len, num_heads, head_dim, rotary_dim
 
 
-def fused_qknorm_rope(
+def _fused_qknorm_rope_impl(
     x: torch.Tensor,
     weight: torch.Tensor,
     cos: torch.Tensor,
     sin: torch.Tensor,
     eps: float,
 ) -> torch.Tensor:
-    """Run per-head RMSNorm and partial RoPE in one Sol-Engine-style kernel.
-
-    RMSNorm reduction and RoPE arithmetic stay in FP32 registers until the
-    final store. Triton's reduction order and the absence of eager's BF16
-    intermediate materializations can produce small, expected rounding drift.
-
-    Row offsets are computed in int64, so inputs beyond 2**31 total elements
-    (about 300k tokens per rank at H3's 56 heads x 128 head_dim) address
-    correctly.
-    """
     batch, seq_len, num_heads, head_dim, rotary_dim = _validate_inputs(x, weight, cos, sin, eps)
     if not weight.is_contiguous():
         raise ValueError("weight must be contiguous")
@@ -171,4 +161,56 @@ def fused_qknorm_rope(
     return flat_out.view(batch, seq_len, num_heads, head_dim)
 
 
+@torch.library.custom_op(
+    "fastvideo::_minimax_h3_qknorm_rope",
+    mutates_args=(),
+    device_types="cuda",
+)
+def _fused_qknorm_rope_op(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float,
+) -> torch.Tensor:
+    return _fused_qknorm_rope_impl(x, weight, cos, sin, eps)
+
+
+@torch.library.register_fake("fastvideo::_minimax_h3_qknorm_rope")
+def _fused_qknorm_rope_fake(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float,
+) -> torch.Tensor:
+    del weight, cos, sin, eps
+    return x.new_empty(x.shape)
+
+
+def fused_qknorm_rope(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float,
+) -> torch.Tensor:
+    """Run per-head RMSNorm and partial RoPE in one Sol-Engine-style kernel.
+
+    RMSNorm reduction and RoPE arithmetic stay in FP32 registers until the
+    final store. Triton's reduction order and the absence of eager's BF16
+    intermediate materializations can produce small, expected rounding drift.
+
+    Row offsets are computed in int64, so inputs beyond 2**31 total elements
+    (about 300k tokens per rank at H3's 56 heads x 128 head_dim) address
+    correctly.
+    """
+    if torch.is_grad_enabled() and any(
+            isinstance(tensor, torch.Tensor) and tensor.requires_grad for tensor in (x, weight, cos, sin)):
+        raise RuntimeError("fused_qknorm_rope is inference-only and does not implement autograd")
+    if torch.compiler.is_compiling():
+        return torch.ops.fastvideo._minimax_h3_qknorm_rope(x, weight, cos, sin, eps)
+    return _fused_qknorm_rope_impl(x, weight, cos, sin, eps)
+
+
 __all__ = ["HAVE_TRITON", "fused_qknorm_rope"]

fastvideo/models/dits/minimax_h3_fusions/swiglu.py

@@ -65,12 +65,7 @@ def _num_warps(block_size: int) -> int:
     return 4
 
 
-def minimax_h3_swiglu(x: torch.Tensor) -> torch.Tensor:
-    """Run the forward-only Triton fusion over an H3 ``(..., 2 * ffn_dim)`` input.
-
-    This is intentionally a strict kernel wrapper: callers own fallback policy and
-    must only invoke it for a supported CUDA inference path.
-    """
+def _minimax_h3_swiglu_impl(x: torch.Tensor) -> torch.Tensor:
     ffn_dim = _validate_input(x)
     if not x.is_cuda:
         raise ValueError("MiniMax H3 fused SwiGLU requires a CUDA tensor")
@@ -101,4 +96,31 @@ def minimax_h3_swiglu(x: torch.Tensor) -> torch.Tensor:
     return out.view(output_shape)
 
 
+@torch.library.custom_op(
+    "fastvideo::_minimax_h3_swiglu",
+    mutates_args=(),
+    device_types="cuda",
+)
+def _minimax_h3_swiglu_op(x: torch.Tensor) -> torch.Tensor:
+    return _minimax_h3_swiglu_impl(x)
+
+
+@torch.library.register_fake("fastvideo::_minimax_h3_swiglu")
+def _minimax_h3_swiglu_fake(x: torch.Tensor) -> torch.Tensor:
+    return x.new_empty((*x.shape[:-1], x.shape[-1] // 2))
+
+
+def minimax_h3_swiglu(x: torch.Tensor) -> torch.Tensor:
+    """Run the forward-only Triton fusion over an H3 ``(..., 2 * ffn_dim)`` input.
+
+    This is intentionally a strict kernel wrapper: callers own fallback policy and
+    must only invoke it for a supported CUDA inference path.
+    """
+    if torch.is_grad_enabled() and x.requires_grad:
+        raise RuntimeError("MiniMax H3 fused SwiGLU is forward-only and does not implement autograd")
+    if torch.compiler.is_compiling():
+        return torch.ops.fastvideo._minimax_h3_swiglu(x)
+    return _minimax_h3_swiglu_impl(x)
+
+
 __all__ = ["HAVE_TRITON", "minimax_h3_swiglu"]

fastvideo/models/encoders/minimax_h3_checkpoint_fp8.py

@@ -289,12 +289,27 @@ def _flashinfer_gemm_w8a8_block_fp8_linear_with_fallback(
     if backend == "cutlass":
         m, k = input_2d.shape
         n = weight.shape[0]
-        expected_x_scale_shape = (k // block_k, m)
-        expected_weight_scale_shape = (k // block_k, n // block_n)
         if x_scale.shape == (m, k // block_k):
             x_scale = x_scale.transpose(-1, -2).contiguous()
         if weight_scale.shape == (n // block_n, k // block_k):
             weight_scale = weight_scale.transpose(-1, -2).contiguous()
+        # FlashInfer documents that ``m`` "should be padded to a multiple of 4
+        # before calling this function", and the CUTLASS groupwise module
+        # rejects unaligned m at dispatch (``cutlass gemm.can_implement
+        # failed``). Prompt token counts are arbitrary, so pad the quantized
+        # activation rows and their per-token scales with zeros, then slice
+        # the padded rows off the GEMM output. The trtllm (sm100) route
+        # tolerates any m and stays unpadded.
+        padded_m = (m + 3) // 4 * 4
+        if padded_m != m:
+            padded_q_input = q_input.new_zeros((padded_m, q_input.shape[1]))
+            padded_q_input[:m] = q_input
+            q_input = padded_q_input
+            padded_x_scale = x_scale.new_zeros((x_scale.shape[0], padded_m))
+            padded_x_scale[:, :m] = x_scale
+            x_scale = padded_x_scale
+        expected_x_scale_shape = (k // block_k, padded_m)
+        expected_weight_scale_shape = (k // block_k, n // block_n)
         if x_scale.shape != expected_x_scale_shape or weight_scale.shape != expected_weight_scale_shape:
             raise RuntimeError("FlashInfer CUTLASS block-FP8 scale layout mismatch: "
                                f"x_scale={tuple(x_scale.shape)}, weight_scale={tuple(weight_scale.shape)}, "
@@ -310,6 +325,8 @@ def _flashinfer_gemm_w8a8_block_fp8_linear_with_fallback(
             backend="cutlass",
             scale_major_mode="MN",
         )
+        if padded_m != m:
+            output = output[:m]
     else:
         expected_x_scale_shape = (input_2d.shape[0], input_2d.shape[1] // block_k)
         expected_weight_scale_shape = (weight.shape[0] // block_n, weight.shape[1] // block_k)

fastvideo/models/loader/component_loader.py

@@ -1126,6 +1126,7 @@ class TransformerLoader(ComponentLoader):
                 training_mode=fastvideo_args.training_mode,
                 enable_torch_compile=fastvideo_args.enable_torch_compile,
                 torch_compile_kwargs=fastvideo_args.torch_compile_kwargs,
+                inference_regional_compile=fastvideo_args.inference_torch_compile,
             )
 
         total_params = sum(p.numel() for p in model.parameters())

fastvideo/models/loader/fsdp_load.py

@@ -136,6 +136,7 @@ def maybe_load_fsdp_model(
     pin_cpu_memory: bool = True,
     enable_torch_compile: bool = False,
     torch_compile_kwargs: dict[str, Any] | None = None,
+    inference_regional_compile: bool = False,
 ) -> torch.nn.Module:
     """
     Load the model with FSDP if is training, else load the model without FSDP.
@@ -235,9 +236,142 @@ def maybe_load_fsdp_model(
         logger.info("Enabling torch.compile for FSDP training module with kwargs=%s", compile_kwargs)
         model = torch.compile(model, **compile_kwargs)
         logger.info("torch.compile enabled for %s", type(model).__name__)
+    elif inference_regional_compile and not training_mode:
+        # Inference-side counterpart of the #1718 training regional compile:
+        # per-block fullgraph compile right after the transformer loads, no
+        # user kwargs needed (fullgraph + emulate_precision_casts injected).
+        unsupported = _regional_compile_unsupported_reason(init_params)
+        if unsupported is not None:
+            logger.warning(
+                "inference_torch_compile requested but disabled: %s. "
+                "Inference continues in eager mode.", unsupported)
+        else:
+            prepare_for_compile = getattr(model, "prepare_for_compile", None)
+            if callable(prepare_for_compile):
+                logger.info("Running prepare_for_compile for %s", type(model).__name__)
+                prepare_for_compile()
+            attention_count = _enable_regional_attention_compile(model)
+            logger.info("Enabled attention tracing for %d modules in %s", attention_count, type(model).__name__)
+            _compile_model_regions(model, torch_compile_kwargs or {})
     return model
 
 
+def _regional_compile_unsupported_reason(init_params: dict[str, Any]) -> str | None:
+    """Return why regional fullgraph compile cannot run, or None if it can.
+
+    FA3's grad-enabled attention path deliberately routes to the raw
+    autograd.Function at the cost of a dynamo graph break (see
+    flash_attn_default.py) — under regional ``fullgraph=True`` that break is
+    a hard RuntimeError at the first compiled forward. FA2 and FA4 route
+    through traceable custom ops and are compile-safe.
+
+    The legacy VSA backend remains outside the fullgraph support envelope.
+    MiniMax H3's VSA backend has a model-specific ``prepare_for_compile``
+    hook that resolves its loaded compression gates before block capture, so
+    VIDEO_SPARSE_ATTN_H3 is allowed through this policy gate.
+    """
+    try:
+        from fastvideo.attention.layer import _attention_compile_explicitly_disabled
+    except Exception:  # pragma: no cover - attention stack not importable
+        pass
+    else:
+        if _attention_compile_explicitly_disabled():
+            # The escape hatch wraps attention forwards in
+            # torch.compiler.disable, which is a hard dynamo error inside a
+            # fullgraph region ("Skip inlining `torch.compiler.disable()`d
+            # function"). Degrade to eager instead, matching the hatch's
+            # debugging intent.
+            return ("FASTVIDEO_DISABLE_ATTENTION_COMPILE=1 keeps attention "
+                    "forwards out of compiled graphs via torch.compiler."
+                    "disable, which fullgraph regional compile cannot trace; "
+                    "this model stays eager")
+    config = init_params.get("config")
+    resolved = getattr(config, "_resolved_attention_backend", None)
+    resolved_name = getattr(resolved, "name", "")
+    if resolved_name == "VIDEO_SPARSE_ATTN_H3" and os.environ.get("FASTVIDEO_H3_VSA_PROBE"):
+        return ("FASTVIDEO_H3_VSA_PROBE records tensors and files from the VSA-H3 attention body, which "
+                "regional fullgraph compile cannot capture; this model stays eager")
+    if resolved_name == "VIDEO_SPARSE_ATTN":
+        return (f"attention backend resolved to {resolved_name}, whose Triton "
+                "kernels, sequence-parallel collectives, and sync metadata "
+                "guard graph-break (incompatible with fullgraph regional "
+                "compile); this model stays eager")
+    if resolved is None or resolved_name != "FLASH_ATTN":
+        return None
+    try:
+        from fastvideo.attention.utils.flash_attn_default import fa_version
+    except Exception:  # pragma: no cover - flash-attn stack not importable
+        return None
+    if fa_version == "3":
+        return ("attention backend resolved to FLASH_ATTN with flash-attn 3, "
+                "whose grad-enabled path graph-breaks (incompatible with "
+                "fullgraph regional compile); use FA2, FA4 (FASTVIDEO_FA4=1), "
+                "or TORCH_SDPA for compiled runs; this model stays eager")
+    return None
+
+
+def _enable_regional_attention_compile(model: nn.Module) -> int:
+    """Opt in distributed-attention instances owned by ``model`` only."""
+    from fastvideo.attention.layer import DistributedAttention
+
+    enabled_count = 0
+    for submodule in model.modules():
+        if isinstance(submodule, DistributedAttention):
+            submodule._set_compile_forward_enabled(True)
+            enabled_count += 1
+    return enabled_count
+
+
+def _compile_model_regions(model: nn.Module, compile_kwargs: dict[str, Any]) -> int:
+    """Compile repeated mathematical regions of a loaded model.
+
+    Only the selected module ``forward`` is replaced. This keeps activation
+    checkpoint wrappers structurally transparent while any module-level hooks
+    (FSDP pre/post, layerwise offload) execute outside the compiled region.
+    """
+    compile_conditions = getattr(model, "_compile_conditions", None)
+    if not compile_conditions:
+        raise ValueError(f"{type(model).__name__} does not declare _compile_conditions")
+
+    if compile_kwargs.get("fullgraph", True) is not True:
+        raise ValueError("Regional compile requires fullgraph=True")
+    if "mode" in compile_kwargs:
+        # torch.compile forbids passing both `mode` and `options`, and
+        # regional compile always injects options (emulate_precision_casts)
+        # to match the training-side regional-compile configuration. Fail here
+        # with an actionable message instead of letting torch raise a
+        # mode/options conflict about an `options` key the user never wrote.
+        raise ValueError("Regional compile sets inductor options "
+                         "(emulate_precision_casts) and cannot be combined "
+                         "with torch_compile_kwargs['mode']. Remove 'mode' or "
+                         "express its effect via torch_compile_kwargs['options'].")
+    kwargs = {**compile_kwargs, "fullgraph": True}
+    options = {"emulate_precision_casts": True}
+    options.update(kwargs.get("options") or {})
+    kwargs["options"] = options
+    compiled_count = 0
+    for name, submodule in list(model.named_modules()):
+        if not name:
+            continue
+        if any(condition(name, submodule) for condition in compile_conditions):
+            # Activation checkpoint wrappers are control-flow boundaries, not
+            # mathematical regions. Keep their saved-tensor/recompute logic
+            # eager and compile only the repeated block they own.
+            compile_target = getattr(submodule, "_checkpoint_wrapped_module", submodule)
+            compile_target.forward = torch.compile(compile_target.forward, **kwargs)
+            compiled_count += 1
+
+    if compiled_count == 0:
+        raise ValueError(f"No submodules in {type(model).__name__} matched _compile_conditions")
+    logger.info(
+        "Enabled regional torch.compile for %d submodules in %s with kwargs=%s",
+        compiled_count,
+        type(model).__name__,
+        kwargs,
+    )
+    return compiled_count
+
+
 def shard_model(
     model,
     *,

fastvideo/models/vaes/minimax_h3_parallel.py

@@ -0,0 +1,381 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Sequence-parallel chunk scheduling for the MiniMax-H3 video VAE.
+
+The H3 video VAE decodes a video as a series of temporal-chunk decoder
+forwards whose outputs are joined by a short deterministic frame blend
+(``AutoencoderKLMiniMaxH3._decode_chunks``), and encodes videos as fully
+independent ``clip_length``-frame encoder forwards. Neither the chunk decode
+nor the clip encode has any cross-chunk data dependency — only the *joining*
+of decoded chunks (overlap blending, frame trimming) is sequential. This
+module round-robins the chunk/clip forwards across the ranks of a
+sequence-parallel group and replays the serial joining logic on the
+assembling rank, reproducing the serial result bit for bit.
+
+Bit-exactness contract:
+- every rank holds an identical copy of the inputs (the H3 DiT all-gathers
+  its outputs, and reference pixels are prepared identically on all ranks);
+- a chunk decoded on any rank is bitwise the tensor the serial loop would
+  produce (identical weights, inputs, and deterministic kernels on identical
+  GPUs), and NCCL transports it bitwise;
+- every serialization point of the serial algorithm (overlap blending, frame
+  trimming, pixel denormalization, output-buffer copies, moment
+  concatenation and token-drop trimming) runs on the assembling rank in
+  serial order via the same VAE methods the serial path uses.
+
+Collective safety: all group ranks must call these functions together with
+identically shaped inputs. Work proceeds in rounds of one collective each;
+ranks without a chunk in the final round contribute a placeholder tensor, so
+participation is uniform by construction and no rank-dependent branch guards
+a collective.
+
+Caveat — compiled decoders (``enable_torch_compile_vae``): inductor autotunes
+kernel configs per process at first call, so a compiled decoder is only
+deterministic WITHIN a process, not across processes. Chunks decoded on other
+ranks then differ from the serial rank's decode of the same chunk exactly as
+two serial runs in different processes would. Direct decoder tensors measured
+on GB200 at 124f had max absolute error 0.00268358 (0.684/255), mean absolute
+error 4.213e-05 (0.0107/255), and 24.59% nonzero values; the first chunk was
+bit-identical. A separate decoded-MP4 comparison reached 63/255 on <0.5% of
+pixels, but that includes lossy MP4 encoding and is not the decoder-tensor
+error envelope. With the eager decoder — the pipeline default — parallel
+output is bitwise equal to serial ``decode_to_pixels``.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from fastvideo.models.vaes.minimax_h3_video import (
+    AutoencoderKLMiniMaxH3,
+    AutoencoderKLOutput,
+    DiagonalGaussianDistribution,
+)
+from fastvideo.profiler import nvtx_range
+
+if TYPE_CHECKING:
+    from fastvideo.distributed.parallel_state import GroupCoordinator
+
+# Collective used to move decoded chunk segments to the assembling rank.
+# "gather" moves each segment once (destination-only); "all_gather" also
+# leaves every rank with every segment. Both are exact; the default is the
+# faster one measured on GB200 NVL72 (see the PR notes).
+DECODE_GATHER_STRATEGIES = ("gather", "all_gather")
+DEFAULT_DECODE_GATHER_STRATEGY = "gather"
+
+
+def parallel_chunk_indices(num_chunks: int, world_size: int, rank_in_group: int) -> list[int]:
+    """Round-robin chunk ownership: chunk ``i`` belongs to rank ``i % world_size``."""
+    if num_chunks < 0:
+        raise ValueError(f"num_chunks must be non-negative, got {num_chunks}.")
+    if world_size < 1:
+        raise ValueError(f"world_size must be positive, got {world_size}.")
+    if not 0 <= rank_in_group < world_size:
+        raise ValueError(f"rank_in_group {rank_in_group} out of range for world_size {world_size}.")
+    return list(range(rank_in_group, num_chunks, world_size))
+
+
+def _num_rounds(num_chunks: int, world_size: int) -> int:
+    return -(-num_chunks // world_size)
+
+
+def _decode_segment(vae: AutoencoderKLMiniMaxH3, z_padded: torch.Tensor, chunk_index: int) -> torch.Tensor:
+    """Decode one temporal chunk's clip and keep the frames the join consumes.
+
+    The serial loop uses two spans of each decoded clip: the chunk body
+    ``clip[:, :, frame_pre_padding:chunk_num_frames]`` and (when
+    ``token_drop > 0``) the blend tail
+    ``clip[:, :, chunk_num_frames + frame_pre_padding:]``. Everything from
+    ``frame_pre_padding`` on covers both, so one contiguous slice per chunk
+    travels over the wire. ``.contiguous()`` also detaches the segment from
+    any decoder-owned storage (e.g. a compiled decoder's reuse pools) before
+    the next chunk decode can overwrite it.
+    """
+    start = chunk_index * vae.tokens_chunk_size
+    with nvtx_range(f"minimax_h3.vae.parallel_chunk.{chunk_index}"):
+        clip = vae._decode_clip(z_padded[:, :, start:start + vae.tokens_chunk_size + vae.token_overlap])
+        return clip[:, :, vae.frame_pre_padding:].contiguous()
+
+
+class _ChunkAssembler:
+    """Replay the serial chunk-joining semantics of ``_decode_chunks`` +
+    ``_decode_to_pixels`` on gathered chunk segments, in chunk order.
+
+    On CUDA the joining kernels and output copies run on a dedicated side
+    stream: they depend only on already-gathered segments, so running them
+    off the main stream keeps the assembling rank's next chunk decode (and
+    therefore every other rank's next collective) off the assembly's tail.
+    Stream placement cannot change values — the ops and their order are
+    identical — so bit-exactness with the serial path is unaffected.
+    """
+
+    def __init__(self, vae: AutoencoderKLMiniMaxH3, output: torch.Tensor, output_num_frames: int,
+                 non_blocking: bool, device: torch.device) -> None:
+        self._vae = vae
+        self._output = output
+        self._output_num_frames = output_num_frames
+        self._non_blocking = non_blocking
+        self._body_frames = vae.tokens_chunk_size * vae.temporal_compression_ratio - vae.frame_pre_padding
+        self._overlap: torch.Tensor | None = None
+        self._frame_start = 0
+        self._stream = torch.cuda.Stream(device) if device.type == "cuda" else None
+
+    def push(self, segment: torch.Tensor) -> None:
+        """Consume the next chunk's segment (``clip[:, :, frame_pre_padding:]``)."""
+        if self._stream is None:
+            self._push(segment)
+            return
+        # The segment is produced on the current (collective) stream; hand it
+        # to the assembly stream and pin its storage until assembly reads it.
+        self._stream.wait_stream(torch.cuda.current_stream(segment.device))
+        segment.record_stream(self._stream)
+        with torch.cuda.stream(self._stream):
+            self._push(segment)
+
+    def _push(self, segment: torch.Tensor) -> None:
+        vae = self._vae
+        chunk = segment[:, :, :self._body_frames]
+        if self._overlap is not None:
+            chunk = vae._blend(self._overlap, chunk, vae.frame_overlap, dim=-3)
+        num_frames = min(chunk.shape[2], self._output_num_frames - self._frame_start)
+        chunk = chunk[:, :, :num_frames]
+        # The tail past the body (and its pre-padding gap) is the next
+        # chunk's blend overlap — the serial loop's ``next_overlap``.
+        self._overlap = segment[:, :, self._body_frames + vae.frame_pre_padding:] if vae.config.token_drop > 0 else None
+        if num_frames > 0:
+            self._emit(chunk)
+
+    def finalize(self) -> None:
+        """Emit the final overlap tail exactly as the serial generator does."""
+        if self._overlap is not None and self._frame_start < self._output_num_frames:
+            tail = self._overlap[:, :, :self._output_num_frames - self._frame_start]
+            if self._stream is None:
+                self._emit(tail)
+            else:
+                with torch.cuda.stream(self._stream):
+                    self._emit(tail)
+        if self._frame_start != self._output.shape[2]:
+            raise RuntimeError(
+                f"MiniMax-H3 decode wrote {self._frame_start} frames into an output buffer expecting "
+                f"{self._output.shape[2]}.")
+
+    def synchronize(self) -> None:
+        """Drain assembly kernels and output copies before the buffer is read."""
+        if self._stream is not None:
+            self._stream.synchronize()
+
+    def _emit(self, chunk: torch.Tensor) -> None:
+        pixels = self._vae.denormalize_pixels(chunk.float()).clamp_(0, 1)
+        self._vae._copy_chunk_pixels(pixels, self._output, self._frame_start, self._non_blocking)
+        self._frame_start += pixels.shape[2]
+
+
+def _broadcast_segment_meta(group: "GroupCoordinator",
+                            segment: torch.Tensor | None) -> tuple[torch.dtype, tuple[int, ...]]:
+    """Share the leader's real segment dtype/shape so placeholder tensors match.
+
+    The decoder's output dtype depends on the surrounding autocast context;
+    deriving it on the leader from an actually decoded segment (instead of
+    predicting it) keeps collective dtypes correct by construction.
+    """
+    meta = (segment.dtype, tuple(segment.shape)) if segment is not None else None
+    meta = group.broadcast_object(meta, src=0)
+    if meta is None:
+        raise RuntimeError("MiniMax-H3 parallel VAE meta broadcast returned no leader metadata.")
+    return meta
+
+
+def decode_to_pixels_parallel(
+    vae: AutoencoderKLMiniMaxH3,
+    z: torch.Tensor,
+    output: torch.Tensor | None,
+    group: "GroupCoordinator",
+    strategy: str = DEFAULT_DECODE_GATHER_STRATEGY,
+) -> torch.Tensor | None:
+    """Chunk-parallel ``decode_to_pixels`` across a sequence-parallel group.
+
+    All group ranks call this together with identical ``z``. Temporal chunks
+    are decoded round-robin across the group and their segments move to the
+    group's first rank, which assembles bitwise the serial
+    ``decode_to_pixels`` result into ``output``. Only the first rank passes
+    ``output`` (validated exactly like the serial API); other ranks pass
+    ``None`` and receive ``None``.
+    """
+    if strategy not in DECODE_GATHER_STRATEGIES:
+        raise ValueError(f"Unknown parallel-decode strategy {strategy!r}; expected one of {DECODE_GATHER_STRATEGIES}.")
+    is_leader = group.rank_in_group == 0
+    if is_leader:
+        if output is None:
+            raise ValueError("The first sequence-parallel rank must provide the CPU output buffer.")
+        expected_shape = vae.decoded_pixel_shape(z.shape)
+        if output.device.type != "cpu" or output.dtype != torch.float32 or tuple(output.shape) != expected_shape:
+            raise ValueError(
+                "`output` must be a CPU float32 tensor with shape "
+                f"{expected_shape}, got device={output.device}, dtype={output.dtype}, shape={tuple(output.shape)}.")
+    elif output is not None:
+        raise ValueError("Only the first sequence-parallel rank may provide an output buffer.")
+    if group.world_size == 1:
+        return vae.decode_to_pixels(z, output)
+
+    try:
+        if vae.use_slicing and z.shape[0] > 1:
+            for batch_index, z_slice in enumerate(z.split(1)):
+                slice_output = output[batch_index:batch_index + 1] if output is not None else None
+                _decode_single_parallel(vae, z_slice, slice_output, group, strategy)
+        else:
+            _decode_single_parallel(vae, z, output, group, strategy)
+    finally:
+        # Drain the leader's async chunk copies before the caller (or an
+        # exception handler) can read or release the pinned buffer.
+        if output is not None and vae._streams_chunk_copies(z, output):
+            torch.cuda.current_stream(z.device).synchronize()
+    return output
+
+
+def _decode_single_parallel(
+    vae: AutoencoderKLMiniMaxH3,
+    z: torch.Tensor,
+    output: torch.Tensor | None,
+    group: "GroupCoordinator",
+    strategy: str,
+) -> None:
+    pad_tokens, num_chunks, output_num_frames = vae._temporal_decode_plan(z.shape[2])
+    if pad_tokens > 0:
+        z = torch.cat([z, z[:, :, -1:].repeat(1, 1, pad_tokens, 1, 1)], dim=2)
+    world_size = group.world_size
+    rank = group.rank_in_group
+
+    # Every rank decodes its round-0 chunk BEFORE the metadata rendezvous so
+    # the first decodes run concurrently (a rank that waited on the broadcast
+    # first would idle a full chunk-decode behind the leader). The leader
+    # owns chunk 0 under round-robin assignment, so its segment supplies real
+    # dtype/shape for placeholder rounds instead of guessing autocast state.
+    first_segment = _decode_segment(vae, z, rank) if rank < num_chunks else None
+    segment_dtype, segment_shape = _broadcast_segment_meta(group, first_segment if rank == 0 else None)
+
+    assembler = None
+    if output is not None:
+        non_blocking = vae._streams_chunk_copies(z, output)
+        assembler = _ChunkAssembler(vae, output, output_num_frames, non_blocking, z.device)
+
+    try:
+        segment_frames = segment_shape[2]
+        for round_index in range(_num_rounds(num_chunks, world_size)):
+            chunk_index = round_index * world_size + rank
+            if chunk_index >= num_chunks:
+                segment = torch.zeros(segment_shape, dtype=segment_dtype, device=z.device)
+            elif round_index == 0 and first_segment is not None:
+                segment = first_segment
+            else:
+                segment = _decode_segment(vae, z, chunk_index)
+            with nvtx_range(f"minimax_h3.vae.parallel_{strategy}.{round_index}"):
+                if strategy == "gather":
+                    gathered = group.gather(segment, dst=0, dim=2)
+                else:
+                    gathered = group.all_gather(segment, dim=2)
+            if assembler is None or gathered is None:
+                continue
+            for slot in range(world_size):
+                if round_index * world_size + slot >= num_chunks:
+                    break
+                assembler.push(gathered.narrow(2, slot * segment_frames, segment_frames))
+        if assembler is not None:
+            assembler.finalize()
+    finally:
+        # Drain assembly-stream copies into ``output`` even on the error path
+        # so an exception cannot leave an in-flight DMA into a buffer the
+        # caller may release.
+        if assembler is not None:
+            assembler.synchronize()
+
+
+def _encode_clip_moments(vae: AutoencoderKLMiniMaxH3, pixels: torch.Tensor, clip_index: int) -> torch.Tensor:
+    """Encode one ``clip_length``-frame clip exactly as ``_encode_pixels`` does."""
+    clip_length = vae.config.clip_length
+    frame_start = clip_index * clip_length
+    with nvtx_range(f"minimax_h3.vae.parallel_encode_clip.{clip_index}"):
+        clip = pixels[:, :, frame_start:frame_start + clip_length].to(
+            device=vae.pixel_mean.device,
+            dtype=torch.float32,
+        )
+        if pixels.dtype == torch.uint8:
+            clip = clip / 255.0
+        if clip.shape[2] < clip_length:
+            pad_frames = clip[:, :, -1:].repeat(1, 1, clip_length - clip.shape[2], 1, 1)
+            clip = torch.cat([clip, pad_frames], dim=2)
+        clip = vae.normalize_pixels(clip)
+        return vae._encode_clip(clip).contiguous()
+
+
+def encode_pixels_parallel(
+    vae: AutoencoderKLMiniMaxH3,
+    pixels: torch.Tensor,
+    group: "GroupCoordinator",
+) -> AutoencoderKLOutput:
+    """Clip-parallel ``encode_pixels`` across a sequence-parallel group.
+
+    Encoder clips have no cross-clip dependency (no overlap, no blending), so
+    ranks encode disjoint clips and all-gather the per-clip moment tensors.
+    Every rank returns the identical full posterior — preserving the serial
+    contract that all ranks hold the same encoded latents — bitwise equal to
+    ``vae.encode_pixels(pixels)``. Moments are latent-sized (a few MB per
+    clip), so the all-gather is negligible next to the clip forwards.
+    """
+    if pixels.ndim != 5 or pixels.shape[1] != vae.config.in_channels or pixels.shape[2] <= 0:
+        raise ValueError(
+            f"`pixels` must have shape [B, {vae.config.in_channels}, T, H, W] with T > 0, "
+            f"got {tuple(pixels.shape)}.")
+    if pixels.device.type != "cpu":
+        raise ValueError(f"`pixels` must remain on CPU, got device={pixels.device}.")
+    if pixels.dtype != torch.uint8 and not pixels.is_floating_point():
+        raise TypeError(f"`pixels` must use uint8 or a floating-point dtype, got {pixels.dtype}.")
+    if group.world_size == 1:
+        return vae.encode_pixels(pixels)
+    if vae.use_slicing and pixels.shape[0] > 1:
+        moments = torch.cat([_encode_single_parallel(vae, pixel_slice, group) for pixel_slice in pixels.split(1)])
+    else:
+        moments = _encode_single_parallel(vae, pixels, group)
+    return AutoencoderKLOutput(latent_dist=DiagonalGaussianDistribution(moments))
+
+
+def _encode_single_parallel(vae: AutoencoderKLMiniMaxH3, pixels: torch.Tensor,
+                            group: "GroupCoordinator") -> torch.Tensor:
+    clip_length = vae.config.clip_length
+    num_clips = -(-pixels.shape[2] // clip_length)
+    world_size = group.world_size
+    rank = group.rank_in_group
+
+    # Same first-work-then-rendezvous ordering as the decode path: encode the
+    # round-0 clip before the metadata broadcast so first encodes overlap.
+    first_moments = _encode_clip_moments(vae, pixels, rank) if rank < num_clips else None
+    moment_dtype, moment_shape = _broadcast_segment_meta(group, first_moments if rank == 0 else None)
+
+    moment_tokens = moment_shape[2]
+    parts: list[torch.Tensor] = []
+    for round_index in range(_num_rounds(num_clips, world_size)):
+        clip_index = round_index * world_size + rank
+        if clip_index >= num_clips:
+            moments = torch.zeros(moment_shape, dtype=moment_dtype, device=vae.pixel_mean.device)
+        elif round_index == 0 and first_moments is not None:
+            moments = first_moments
+        else:
+            moments = _encode_clip_moments(vae, pixels, clip_index)
+        gathered = group.all_gather(moments, dim=2)
+        for slot in range(world_size):
+            if round_index * world_size + slot >= num_clips:
+                break
+            parts.append(gathered.narrow(2, slot * moment_tokens, moment_tokens))
+    encoded = torch.cat(parts, dim=2)
+    if vae.config.token_drop > 0:
+        encoded = encoded[:, :, :-vae.config.token_drop]
+    return encoded
+
+
+__all__ = [
+    "DECODE_GATHER_STRATEGIES",
+    "DEFAULT_DECODE_GATHER_STRATEGY",
+    "decode_to_pixels_parallel",
+    "encode_pixels_parallel",
+    "parallel_chunk_indices",
+]

fastvideo/models/vaes/minimax_h3_video.py

@@ -327,6 +327,10 @@ class MiniMaxH3VideoAttention(nn.Module):
                 softmax_scale=dim_head**-0.5,
                 num_kv_heads=heads,
                 causal=False,
+                # The FASTVIDEO_NVFP4_FA4 env opt-in targets the DiT; this VAE
+                # is FP32-pinned (_keep_in_fp32_modules), so force-disable FP4
+                # Q/K quantization for its attention regardless of the env.
+                nvfp4_fa4=False,
             )
 
     def forward(
@@ -531,6 +535,9 @@ class AutoencoderKLMiniMaxH3(nn.Module):
     _repeated_blocks = ["MiniMaxH3VideoTransformerBlock"]
     _keep_in_fp32_modules = ["encoder", "decoder", "quant_conv", "post_quant_conv"]
     _compile_conditions = [_is_minimax_h3_video_vae_decoder]
+    # ``prepare_for_compile`` flips this per instance when the pipeline opts
+    # into ``enable_torch_compile_vae``; default instances stay fully eager.
+    _tile_helpers_compiled = False
 
     def __init__(self, config: MiniMaxH3VideoVAEConfig) -> None:
         super().__init__()
@@ -696,8 +703,33 @@ class AutoencoderKLMiniMaxH3(nn.Module):
         slice_rest[dim] = slice(blend_extent, None)
         return torch.cat([blended, b[tuple(slice_rest)]], dim=dim)
 
-    # The fixed spatial tile grid reuses one compiled blend-and-concatenate graph.
-    @torch.compile(backend="inductor", mode="reduce-overhead", dynamic=False)
+    def prepare_for_compile(self) -> None:
+        """Compile the fixed-shape tile helpers for the opt-in VAE compile path.
+
+        ``ComposedPipelineBase._maybe_compile_pipeline_module`` calls this hook
+        only when ``enable_torch_compile_vae`` is set, right before the decoder
+        is compiled through ``_compile_conditions``. The spatial tile grid and
+        the per-tile decoder-input projection have fixed shapes, so
+        ``mode="reduce-overhead"`` records one CUDA graph per geometry and
+        replays it for every tile and temporal chunk. Keeping this behind the
+        opt-in means default (eager) users pay neither the inductor/triton
+        toolchain requirement and first-decode compile latency nor the
+        permanent cudagraph memory pools, and multi-resolution callers never
+        churn ``dynamic=False`` recompiles they did not ask for.
+        """
+        if self._tile_helpers_compiled:
+            return
+        # The fixed spatial tile grid reuses one compiled blend-and-concatenate graph.
+        self._stitch_tiles = torch.compile(self._stitch_tiles, backend="inductor", mode="reduce-overhead", dynamic=False)
+        # Each fixed-shape latent tile reuses one compiled decoder-input projection.
+        self._project_decoder_tile = torch.compile(
+            self._project_decoder_tile,
+            backend="inductor",
+            mode="reduce-overhead",
+            dynamic=False,
+        )
+        self._tile_helpers_compiled = True
+
     def _stitch_tiles(
         self,
         tiles: list[list[torch.Tensor]],
@@ -721,8 +753,6 @@ class AutoencoderKLMiniMaxH3(nn.Module):
             result_rows.append(torch.cat(result_row, dim=-1))
         return torch.cat(result_rows, dim=-2)
 
-    # Each fixed-shape latent tile reuses one compiled decoder-input projection.
-    @torch.compile(backend="inductor", mode="reduce-overhead", dynamic=False)
     def _project_decoder_tile(self, tile: torch.Tensor) -> torch.Tensor:
         """Project one spatial latent tile into the decoder input channels."""
         return self.post_quant_conv(tile)
@@ -750,12 +780,17 @@ class AutoencoderKLMiniMaxH3(nn.Module):
             rows.append(row)
         latent_y_overlaps = [overlap // self.spatial_compression_ratio for overlap in y_overlaps]
         latent_x_overlaps = [overlap // self.spatial_compression_ratio for overlap in x_overlaps]
-        # Under mode="reduce-overhead" the stitched canvas is a CUDA-graph
-        # static buffer that the next _stitch_tiles replay overwrites. Callers
-        # (_encode/_encode_pixels/encode_keyframe) collect per-clip results
-        # across replays before concatenating, so hand them a caller-owned
-        # tensor instead of cudagraph-pooled storage.
-        return self._stitch_tiles(rows, latent_y_overlaps, latent_x_overlaps).clone()
+        stitched = self._stitch_tiles(rows, latent_y_overlaps, latent_x_overlaps)
+        if self._tile_helpers_compiled:
+            # Under the opt-in mode="reduce-overhead" compile the stitched
+            # canvas is a CUDA-graph static buffer that the next _stitch_tiles
+            # replay overwrites. Callers (_encode/_encode_pixels/
+            # encode_keyframe) collect per-clip results across replays before
+            # concatenating, so hand them a caller-owned tensor instead of
+            # cudagraph-pooled storage. Eager instances return the fresh
+            # torch.cat result directly.
+            stitched = stitched.clone()
+        return stitched
 
     def _decode_clip(self, z: torch.Tensor) -> torch.Tensor:
         """Decode one temporal clip, with optional overlapping spatial tiles."""
@@ -801,13 +836,17 @@ class AutoencoderKLMiniMaxH3(nn.Module):
                     rows.append(row)
 
             with nvtx_range("minimax_h3.vae.decode_clip.stitch_tiles"):
-                # Same CUDA-graph output-ownership contract as _encode_clip:
-                # _decode collects chunks across _stitch_tiles replays before
-                # torch.cat, so the pooled canvas must not escape this driver.
-                # (The streaming _decode_to_pixels path copies each chunk out
-                # before the next decode and never held stale storage; the
-                # clone keeps that path correct too at one D2D copy per chunk.)
-                return self._stitch_tiles(rows, y_overlaps, x_overlaps).clone()
+                stitched = self._stitch_tiles(rows, y_overlaps, x_overlaps)
+                if self._tile_helpers_compiled:
+                    # Same CUDA-graph output-ownership contract as
+                    # _encode_clip: _decode collects chunks across
+                    # _stitch_tiles replays before torch.cat, so the pooled
+                    # canvas must not escape this driver. (The streaming
+                    # _decode_to_pixels path copies each chunk out before the
+                    # next decode; the clone keeps it correct too at one D2D
+                    # copy per chunk.)
+                    stitched = stitched.clone()
+                return stitched
 
     def _encode(self, x: torch.Tensor) -> torch.Tensor:
         clip_length = self.config.clip_length
@@ -933,31 +972,42 @@ class AutoencoderKLMiniMaxH3(nn.Module):
         """Whether finalized chunks copy to ``output`` asynchronously on the current CUDA stream."""
         return z.device.type == "cuda" and output.is_pinned()
 
-    def _decode_to_pixels(self, z: torch.Tensor, output: torch.Tensor) -> None:
-        """Decode temporal chunks and immediately copy finalized pixels to CPU.
+    @staticmethod
+    def _copy_chunk_pixels(pixels: torch.Tensor, output: torch.Tensor, frame_start: int, non_blocking: bool) -> None:
+        """Copy one finalized fp32 pixel chunk into the CPU ``output`` buffer.
 
         Device-to-host copies run per (batch, channel) plane: the temporal
         slice of ``output`` is strided across channels, but each plane is
         contiguous on both sides, so every transfer stays a direct memcpy
         instead of staging through a pageable CPU temporary. With a pinned
-        ``output`` the copies are additionally asynchronous and overlap the
-        next chunk's decode; ``decode_to_pixels`` synchronizes once before
-        returning.
+        ``output`` and ``non_blocking=True`` the copies are additionally
+        asynchronous on the current CUDA stream; callers synchronize once
+        before releasing the buffer.
+        """
+        target = output[:, :, frame_start:frame_start + pixels.shape[2]]
+        if pixels.device.type == "cuda":
+            pixels = pixels.contiguous()
+            for batch_index in range(pixels.shape[0]):
+                for channel_index in range(pixels.shape[1]):
+                    target[batch_index, channel_index].copy_(pixels[batch_index, channel_index],
+                                                             non_blocking=non_blocking)
+        else:
+            target.copy_(pixels)
+
+    def _decode_to_pixels(self, z: torch.Tensor, output: torch.Tensor) -> None:
+        """Decode temporal chunks and immediately copy finalized pixels to CPU.
+
+        Each finalized chunk streams through ``_copy_chunk_pixels`` (direct
+        per-plane memcpys; asynchronous with a pinned ``output``) so the
+        copies overlap the next chunk's decode; ``decode_to_pixels``
+        synchronizes once before returning.
         """
         non_blocking = self._streams_chunk_copies(z, output)
         output_frame_start = 0
         for chunk in self._decode_chunks(z):
             num_frames = chunk.shape[2]
             pixels = self.denormalize_pixels(chunk.float()).clamp_(0, 1)
-            target = output[:, :, output_frame_start:output_frame_start + num_frames]
-            if z.device.type == "cuda":
-                pixels = pixels.contiguous()
-                for batch_index in range(pixels.shape[0]):
-                    for channel_index in range(pixels.shape[1]):
-                        target[batch_index, channel_index].copy_(pixels[batch_index, channel_index],
-                                                                 non_blocking=non_blocking)
-            else:
-                target.copy_(pixels)
+            self._copy_chunk_pixels(pixels, output, output_frame_start, non_blocking)
             output_frame_start += num_frames
         if output_frame_start != output.shape[2]:
             raise RuntimeError(

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

@@ -7,9 +7,11 @@ from typing import Any
 
 import torch
 
-from fastvideo.distributed import get_local_torch_device, get_world_group, model_parallel_is_initialized
+from fastvideo.distributed import get_local_torch_device, get_sp_group, model_parallel_is_initialized
 from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
 from fastvideo.models.vaes.minimax_h3_audio import MiniMaxH3AudioVAE
+from fastvideo.models.vaes.minimax_h3_parallel import DEFAULT_DECODE_GATHER_STRATEGY, decode_to_pixels_parallel
 from fastvideo.models.vaes.minimax_h3_video import AutoencoderKLMiniMaxH3
 from fastvideo.profiler import nvtx_range
 from fastvideo.pipelines.basic.minimax_h3.packing import (
@@ -24,6 +26,8 @@ from fastvideo.pipelines.stages.validators import StageValidators as V
 from fastvideo.pipelines.stages.validators import VerificationResult
 from fastvideo.utils import is_pin_memory_available
 
+logger = init_logger(__name__)
+
 
 def _layout(batch: ForwardBatch) -> MiniMaxH3PackedLayout:
     layout = batch.extra.get(MINIMAX_H3_LAYOUT_KEY)
@@ -32,6 +36,23 @@ def _layout(batch: ForwardBatch) -> MiniMaxH3PackedLayout:
     return layout
 
 
+def _decode_participation(fastvideo_args: FastVideoArgs, want_parallel: bool) -> tuple[Any, bool, bool]:
+    """Resolve (sp_group, is_output_rank, parallel) for the VAE decode stages.
+
+    The executors consume rank 0's ForwardBatch and the training validation
+    callback consumes each sequence-parallel group leader's, so the output
+    rank is the SP group's first rank (identical to world rank 0 in the
+    single-group e2e case). ``parallel`` is only true when every group rank
+    will run the decode body — the collectives inside require uniform
+    participation, so no rank-dependent branch may guard them.
+    """
+    if not model_parallel_is_initialized():
+        return None, True, False
+    sp_group = get_sp_group()
+    parallel = bool(want_parallel) and sp_group.world_size > 1
+    return sp_group, sp_group.is_first_rank, parallel
+
+
 class MiniMaxH3VideoDecodingStage(PipelineStage):
     """Drop visual condition rows, unpatchify, and decode the target video."""
 
@@ -57,11 +78,13 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
     @torch.no_grad()
     def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
         """Decode H3 video latents into normalized CPU pixels."""
-        if model_parallel_is_initialized() and not get_world_group().is_first_rank:
-            # Distributed executors consume rank 0's ForwardBatch. Keep a
+        placeholder = torch.empty((0, 3, 0, 0, 0), device="cpu", dtype=torch.float32)
+        sp_group, is_output_rank, parallel = _decode_participation(fastvideo_args, fastvideo_args.vae_parallel_decode)
+        if not is_output_rank and not parallel:
+            # Consumers read the output rank's ForwardBatch. Keep a
             # verifier-compatible placeholder on other ranks and avoid
             # duplicating the full VAE decode and CPU output buffer.
-            batch.output = torch.empty((0, 3, 0, 0, 0), device="cpu", dtype=torch.float32)
+            batch.output = placeholder
             return batch
 
         layout = _layout(batch)
@@ -81,23 +104,33 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
         try:
             latents = self.vae.denormalize_latents(latents.to(device=device, dtype=torch.float32))
             if fastvideo_args.output_type == "latent":
-                batch.output = latents.detach().float().cpu()
+                # No collectives on this path, so uniform participation is
+                # trivial: every rank returns here.
+                batch.output = latents.detach().float().cpu() if is_output_rank else placeholder
                 return batch
 
-            output = torch.empty(
-                self.vae.decoded_pixel_shape(latents.shape),
-                device="cpu",
-                dtype=torch.float32,
-                pin_memory=fastvideo_args.pin_cpu_memory and is_pin_memory_available(),
-            )
+            output = None
+            if is_output_rank:
+                output = torch.empty(
+                    self.vae.decoded_pixel_shape(latents.shape),
+                    device="cpu",
+                    dtype=torch.float32,
+                    pin_memory=fastvideo_args.pin_cpu_memory and is_pin_memory_available(),
+                )
             # Attribute the streamed decoder computation while retaining
             # per-chunk device-to-host transfer and pinned-buffer reuse.
             with (
                     nvtx_range("minimax_h3.vae"),
                     torch.autocast(device_type=device.type, dtype=torch.float16, enabled=device.type == "cuda"),
             ):
-                self.vae.decode_to_pixels(latents, output)
-            batch.output = output
+                if parallel:
+                    strategy = fastvideo_args.vae_parallel_decode_strategy or DEFAULT_DECODE_GATHER_STRATEGY
+                    logger.info_once(f"MiniMax-H3 VAE decode: sequence-parallel chunks across "
+                                     f"{sp_group.world_size} ranks ({strategy})")
+                    decode_to_pixels_parallel(self.vae, latents, output, sp_group, strategy=strategy)
+                else:
+                    self.vae.decode_to_pixels(latents, output)
+            batch.output = output if is_output_rank else placeholder
             return batch
         finally:
             if fastvideo_args.vae_cpu_offload:
@@ -128,7 +161,9 @@ class MiniMaxH3AudioDecodingStage(PipelineStage):
     @torch.no_grad()
     def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
         """Decode H3 audio latents into a stereo CPU waveform."""
-        if model_parallel_is_initialized() and not get_world_group().is_first_rank:
+        # Audio decode is sub-second, so it always runs serially on the SP
+        # group's first rank (the rank whose ForwardBatch consumers read).
+        if model_parallel_is_initialized() and not get_sp_group().is_first_rank:
             batch.extra["audio"] = torch.empty((0, 2), device="cpu", dtype=torch.float32)
             batch.extra["audio_sample_rate"] = self.audio_vae.sampling_rate
             self._clear_runtime(batch)

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

@@ -9,8 +9,10 @@ import numpy as np
 import torch
 from diffusers.utils.torch_utils import randn_tensor
 
-from fastvideo.distributed import get_local_torch_device
+from fastvideo.distributed import get_local_torch_device, get_sp_group, model_parallel_is_initialized
 from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.vaes.minimax_h3_parallel import encode_pixels_parallel
 from fastvideo.pipelines.basic.minimax_h3.packing import (
     MINIMAX_H3_AUDIO_CHANNELS,
     MINIMAX_H3_KEYFRAME_ENCODE_SEED,
@@ -36,6 +38,8 @@ from fastvideo.pipelines.stages.base import PipelineStage
 from fastvideo.pipelines.stages.validators import StageValidators as V
 from fastvideo.pipelines.stages.validators import VerificationResult
 
+logger = init_logger(__name__)
+
 MINIMAX_H3_LAYOUT_KEY = "minimax_h3_layout"
 
 
@@ -105,8 +109,20 @@ class MiniMaxH3LatentPreparationStage(PipelineStage):
         self,
         references: list[MiniMaxH3PreparedReference],
         device: torch.device,
+        fastvideo_args: FastVideoArgs,
     ) -> list[torch.Tensor]:
         patch_size = self.transformer.patch_size
+        # Reference encode runs on every rank (all ranks hold identical
+        # prepared references), so clip-parallel encode keeps participation
+        # uniform by construction: each rank encodes a clip subset and the
+        # all-gather leaves the identical full posterior everywhere.
+        parallel_group = None
+        if fastvideo_args.vae_parallel_encode and model_parallel_is_initialized():
+            sp_group = get_sp_group()
+            if sp_group.world_size > 1:
+                parallel_group = sp_group
+                logger.info_once(f"MiniMax-H3 reference VAE encode: sequence-parallel clips across "
+                                 f"{sp_group.world_size} ranks")
         rows: list[torch.Tensor] = []
         for reference in references:
             if reference.media_type == "audio":
@@ -120,7 +136,10 @@ class MiniMaxH3LatentPreparationStage(PipelineStage):
                     raise ValueError("MiniMax-H3 reference video frames are missing.")
                 frames = reference.frames[:trim_reference_num_frames(reference.frames.shape[0])]
                 pixels = torch.from_numpy(np.ascontiguousarray(frames)).permute(3, 0, 1, 2)[None]
-                posterior = self.vae.encode_pixels(pixels).latent_dist
+                if parallel_group is not None:
+                    posterior = encode_pixels_parallel(self.vae, pixels, parallel_group).latent_dist
+                else:
+                    posterior = self.vae.encode_pixels(pixels).latent_dist
                 latents = self.vae.normalize_latents(_sample_visual_posterior(posterior).to(
                     torch.float16).float()).cpu()
             reference.num_latent_frames = int(latents.shape[2])
@@ -201,7 +220,7 @@ class MiniMaxH3LatentPreparationStage(PipelineStage):
         vae_device = get_local_torch_device()
         self.vae.to(vae_device)
         try:
-            video_rows = self._encode_visual_rows(references, vae_device)
+            video_rows = self._encode_visual_rows(references, vae_device, fastvideo_args)
         finally:
             if fastvideo_args.vae_cpu_offload:
                 self.vae.to("cpu")

fastvideo/pipelines/composed_pipeline_base.py

@@ -203,6 +203,14 @@ class ComposedPipelineBase(ABC):
                 vae_compile_kwargs = (self.fastvideo_args.torch_compile_kwargs_vae or global_compile_kwargs)
                 audio_vae_compile_kwargs = (self.fastvideo_args.torch_compile_kwargs_audio_vae or global_compile_kwargs)
 
+                if compile_transformer and self.fastvideo_args.inference_torch_compile:
+                    # The loader already applied the regional fullgraph
+                    # compile to the DiT blocks (inference_torch_compile);
+                    # wrapping the same forwards again here would stack
+                    # compiled callables.
+                    logger.info("inference_torch_compile already compiled the DiT regions in the "
+                                "loader; skipping the pipeline-level DiT compile")
+                    compile_transformer = False
                 if compile_transformer:
                     self._maybe_compile_pipeline_module(
                         module_name="transformer",

fastvideo/tests/attention/test_flash_attn_h3_packed_varlen.py

@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Routing guards for MiniMax-H3's packed-varlen FA4 inference path."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+
+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 _build_impl(*, fa4_packed_varlen: bool) -> FlashAttentionImpl:
+    return FlashAttentionImpl(
+        num_heads=2,
+        head_size=8,
+        causal=False,
+        softmax_scale=0.125,
+        num_kv_heads=2,
+        fa4_packed_varlen=fa4_packed_varlen,
+    )
+
+
+def _qkv(batch: int = 1) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    query = torch.randn(batch, 7, 2, 8)
+    return query, torch.randn_like(query), torch.randn_like(query)
+
+
+def test_h3_fa4_no_grad_uses_single_sequence_varlen(monkeypatch: pytest.MonkeyPatch) -> None:
+    calls = []
+
+    def varlen(query, key, value, cu_q, cu_k, max_q, max_k, **kwargs):
+        calls.append((query.shape, key.shape, value.shape, cu_q.tolist(), cu_k.tolist(), max_q, max_k, kwargs))
+        return value + 1
+
+    def unexpected_fixed(*args, **kwargs):
+        raise AssertionError("H3 FA4 inference took the fixed-length path")
+
+    monkeypatch.setattr(flash_attn_module, "fa_version", "4")
+    monkeypatch.setattr(flash_attn_module, "flash_attn_varlen_func_compilable", varlen)
+    monkeypatch.setattr(flash_attn_module, "flash_attn_func_compilable", unexpected_fixed)
+    query, key, value = _qkv()
+
+    with torch.no_grad():
+        output = _build_impl(fa4_packed_varlen=True)._forward_impl(query, key, value, None)
+
+    torch.testing.assert_close(output, value + 1)
+    assert calls == [
+        (
+            torch.Size([7, 2, 8]),
+            torch.Size([7, 2, 8]),
+            torch.Size([7, 2, 8]),
+            [0, 7],
+            [0, 7],
+            7,
+            7,
+            {"dropout_p": 0.0, "softmax_scale": 0.125, "causal": False},
+        )
+    ]
+
+
+@pytest.mark.parametrize(
+    ("fa4_packed_varlen", "version", "batch", "requires_grad"),
+    [
+        (False, "4", 1, False),
+        (True, "3", 1, False),
+        (True, "4", 2, False),
+        (True, "4", 1, True),
+    ],
+)
+def test_h3_varlen_guard_preserves_fixed_path(
+    monkeypatch: pytest.MonkeyPatch,
+    fa4_packed_varlen: bool,
+    version: str,
+    batch: int,
+    requires_grad: bool,
+) -> None:
+    def unexpected_varlen(*args, **kwargs):
+        raise AssertionError("packed-varlen route escaped its H3 FA4 inference guard")
+
+    def fixed(query, key, value, **kwargs):
+        del key, value, kwargs
+        return query + 2
+
+    monkeypatch.setattr(flash_attn_module, "fa_version", version)
+    monkeypatch.setattr(flash_attn_module, "flash_attn_varlen_func_compilable", unexpected_varlen)
+    monkeypatch.setattr(flash_attn_module, "flash_attn_func_compilable", fixed)
+    query, key, value = _qkv(batch)
+    query.requires_grad_(requires_grad)
+
+    context = torch.enable_grad() if requires_grad else torch.no_grad()
+    with context:
+        output = _build_impl(fa4_packed_varlen=fa4_packed_varlen)._forward_impl(query, key, value, None)
+
+    torch.testing.assert_close(output, query + 2)
+    if requires_grad:
+        output.sum().backward()
+        torch.testing.assert_close(query.grad, torch.ones_like(query))

fastvideo/tests/attention/test_flash_attn_nvfp4_opt_out.py

@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+"""NVFP4 FA4 opt-in/opt-out precedence for ``FlashAttentionImpl``.
+
+The process-wide ``FASTVIDEO_NVFP4_FA4=1`` env opt-in targets the DiT. An
+explicit ``nvfp4_fa4`` impl arg must win over the env in both directions so
+precision-sensitive layers (e.g. the FP32-pinned MiniMax-H3 VAE attention)
+can force-disable FP4 Q/K quantization while the DiT keeps it.
+"""
+
+from unittest.mock import patch
+
+import pytest
+
+# The backend module needs a flash-attention package (FA2/FA3, or FA4 with
+# FASTVIDEO_FA4=1) at import. exc_type=ImportError also covers partial
+# installs (e.g. an FA4-only environment without the opt-in env set raises
+# ImportError rather than ModuleNotFoundError).
+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 _build_impl(**extra_impl_args) -> FlashAttentionImpl:
+    return FlashAttentionImpl(
+        num_heads=2,
+        head_size=64,
+        causal=False,
+        softmax_scale=0.125,
+        num_kv_heads=2,
+        **extra_impl_args,
+    )
+
+
+@pytest.fixture
+def fa4_runtime_available():
+    """Satisfy the Blackwell + flash-attention-fp4 asserts without hardware."""
+    with (
+        patch.object(flash_attn_module, "_FA4_FP4_AVAILABLE", True),
+        patch("torch.cuda.get_device_capability", return_value=(10, 0)),
+    ):
+        yield
+
+
+def test_nvfp4_disabled_by_default(monkeypatch) -> None:
+    monkeypatch.delenv("FASTVIDEO_NVFP4_FA4", raising=False)
+    assert _build_impl().nvfp4_fa4 is False
+
+
+def test_nvfp4_env_opt_in_enables_when_arg_absent(monkeypatch, fa4_runtime_available) -> None:
+    monkeypatch.setenv("FASTVIDEO_NVFP4_FA4", "1")
+    assert _build_impl().nvfp4_fa4 is True
+
+
+def test_explicit_disable_wins_over_env_opt_in(monkeypatch) -> None:
+    """The H3 VAE constructs its impl with ``nvfp4_fa4=False``; the DiT env
+    opt-in must not FP4-quantize Q/K inside the FP32-pinned VAE."""
+    monkeypatch.setenv("FASTVIDEO_NVFP4_FA4", "1")
+    assert _build_impl(nvfp4_fa4=False).nvfp4_fa4 is False
+
+
+def test_explicit_enable_works_without_env(monkeypatch, fa4_runtime_available) -> None:
+    monkeypatch.delenv("FASTVIDEO_NVFP4_FA4", raising=False)
+    assert _build_impl(nvfp4_fa4=True).nvfp4_fa4 is True

fastvideo/tests/attention/test_vsa_h3_sm100a_route.py

@@ -20,14 +20,15 @@ _SPEC = dict(raw_latent_shape=(4, 8, 16), patch_size=(1, 2, 2), prefix_segments=
 _HEADS, _DIM = 2, 128
 
 
-def _build_meta():
+def _build_meta(*, sparsity=0.0, dense_layers=(), prefix_segments=_SPEC["prefix_segments"]):
     return MiniMaxH3VSAMetadataBuilder().build(
         current_timestep=0,
         raw_latent_shape=_SPEC["raw_latent_shape"],
         patch_size=_SPEC["patch_size"],
-        VSA_sparsity=0.0,
-        prefix_segments=_SPEC["prefix_segments"],
+        VSA_sparsity=sparsity,
+        prefix_segments=prefix_segments,
         device=torch.device("cpu"),
+        dense_layers=dense_layers,
         tile_size=64,
     )
 
@@ -46,8 +47,11 @@ class _FakeSm100a:
     def __init__(self, supported=True):
         self.supported = supported
         self.calls = []
+        self.support_calls = []
+        self.mask_calls = []
 
     def is_supported(self, q, variable_block_sizes):
+        self.support_calls.append((q, variable_block_sizes))
         return self.supported
 
     def block_sparse_attn_sm100a(self, q, k, v, q2k_idx, q2k_num, variable_block_sizes, need_lse=True):
@@ -55,6 +59,10 @@ class _FakeSm100a:
                                need_lse=need_lse))
         return q.clone(), None
 
+    def block_sparse_attn_sm100a_from_mask(self, q, k, v, block_map, variable_block_sizes):
+        self.mask_calls.append(dict(q=q, block_map=block_map, vbs=variable_block_sizes))
+        return q.clone(), None
+
 
 def _fake_map_to_index(block_map):
     """Pure-torch stand-in for the Triton map_to_index (same contract)."""
@@ -107,6 +115,155 @@ def test_reason_covers_every_precondition():
     assert _sm100a_unavailable_reason(ok, q, vbs, grad_mode=False) is None
 
 
+def test_prepare_for_regional_compile_resolves_supported_route(monkeypatch):
+    fake_sm = _FakeSm100a(supported=True)
+    monkeypatch.setattr(vsa_h3, "_sm100a", fake_sm)
+    monkeypatch.setenv(VSA_SM100A_ENV, "1")
+    impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5)
+
+    impl.prepare_for_regional_compile(torch.device("cpu"))
+
+    assert impl._regional_compile_sm100a_enabled is True
+    assert len(fake_sm.support_calls) == 1
+    probe_q, probe_vbs = fake_sm.support_calls[0]
+    assert probe_q.shape == (1, 1, 128, _DIM)
+    assert probe_q.dtype == torch.bfloat16
+    assert probe_vbs.dtype == torch.int32
+    assert probe_vbs.tolist() == [64, 64]
+    assert impl._regional_compile_layer_idx is not None
+    assert impl._regional_compile_layer_idx.item() == -1
+
+
+def test_prepare_for_regional_compile_env_off_skips_probe(monkeypatch):
+    fake_sm = _FakeSm100a(supported=True)
+    monkeypatch.setattr(vsa_h3, "_sm100a", fake_sm)
+    monkeypatch.delenv(VSA_SM100A_ENV, raising=False)
+    impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5)
+
+    impl.prepare_for_regional_compile(torch.device("cpu"))
+
+    assert impl._regional_compile_sm100a_enabled is False
+    assert fake_sm.support_calls == []
+
+
+def test_prepare_for_regional_compile_requires_mask_entry(monkeypatch):
+
+    class _IndexOnlySm100a:
+
+        def is_supported(self, q, variable_block_sizes):
+            raise AssertionError("missing mask entry must be rejected before the support probe")
+
+    monkeypatch.setattr(vsa_h3, "_sm100a", _IndexOnlySm100a())
+    monkeypatch.setenv(VSA_SM100A_ENV, "1")
+    warnings = []
+    monkeypatch.setattr(vsa_h3.logger, "warning_once", warnings.append)
+    impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5)
+
+    impl.prepare_for_regional_compile(torch.device("cpu"))
+
+    assert impl._regional_compile_sm100a_enabled is False
+    assert len(warnings) == 1
+    assert "block_sparse_attn_sm100a_from_mask" in warnings[0]
+
+
+def test_prepared_route_fullgraph_avoids_eager_dispatch(monkeypatch):
+    """Capture must not revisit env/capability checks or raw map_to_index."""
+    fake_sm = _FakeSm100a(supported=True)
+    monkeypatch.setattr(vsa_h3, "_sm100a", fake_sm)
+    monkeypatch.setenv(VSA_SM100A_ENV, "1")
+    monkeypatch.setattr(vsa_h3, "probe_enabled", lambda: None)
+    impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5)
+    impl.prepare_for_regional_compile(torch.device("cpu"))
+
+    # Four total tiles (two prefix + two video) satisfy the real kernel's
+    # adjacent-query-block contract.
+    meta = MiniMaxH3VSAMetadataBuilder().build(
+        current_timestep=0,
+        raw_latent_shape=_SPEC["raw_latent_shape"],
+        patch_size=_SPEC["patch_size"],
+        VSA_sparsity=0.0,
+        prefix_segments=(64, 64),
+        device=torch.device("cpu"),
+        tile_size=64,
+    )
+    q, k, v = _tiled_qkv(meta)
+
+    def fail_eager_dispatch(*args, **kwargs):
+        raise AssertionError("compiled route re-entered eager dispatch")
+
+    def compile_safe_from_mask(q, k, v, block_map, variable_block_sizes):
+        assert block_map.dtype == torch.bool
+        return q.clone(), None
+
+    monkeypatch.setattr(fake_sm, "is_supported", fail_eager_dispatch)
+    monkeypatch.setattr(fake_sm, "block_sparse_attn_sm100a", fail_eager_dispatch)
+    monkeypatch.setattr(fake_sm, "block_sparse_attn_sm100a_from_mask", compile_safe_from_mask)
+    monkeypatch.setattr(vsa_h3, "map_to_index", fail_eager_dispatch)
+    monkeypatch.setattr(vsa_h3, "_sm100a_unavailable_reason", fail_eager_dispatch)
+    monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", fail_eager_dispatch)
+
+    def run(q, k, v):
+        return impl.forward(q, k, v, None, meta)
+
+    try:
+        compiled = torch.compile(run, backend="eager", fullgraph=True)
+        with torch.inference_mode():
+            actual = compiled(q, k, v)
+        torch.testing.assert_close(actual, q, atol=0, rtol=0)
+    finally:
+        torch._dynamo.reset()
+
+
+def test_prepared_route_reuses_graph_across_layer_indices_and_preserves_dense_overrides(monkeypatch):
+    """Fifty H3 blocks must not specialize the shared graph on layer_idx."""
+    fake_sm = _FakeSm100a(supported=True)
+    monkeypatch.setattr(vsa_h3, "_sm100a", fake_sm)
+    monkeypatch.setenv(VSA_SM100A_ENV, "1")
+    monkeypatch.setattr(vsa_h3, "probe_enabled", lambda: None)
+    meta = _build_meta(sparsity=0.5, dense_layers=(0, 17), prefix_segments=(64, 64))
+    assert meta.dense_layers_tensor.tolist() == [0, 17]
+    q, k, v = _tiled_qkv(meta)
+
+    def fail_eager_dispatch(*args, **kwargs):
+        raise AssertionError("compiled route re-entered eager dispatch")
+
+    def compile_safe_from_mask(q, k, v, block_map, variable_block_sizes):
+        del k, v, variable_block_sizes
+        # Encode the dense-layer decision in the result so this is a semantic
+        # test as well as a recompile-cache stress test.
+        return q + block_map.all().to(q.dtype), None
+
+    implementations = []
+    for layer_idx in range(20):
+        impl = MiniMaxH3VSAImpl(
+            num_heads=_HEADS,
+            head_size=_DIM,
+            causal=False,
+            softmax_scale=_DIM**-0.5,
+            prefix=f"transformer_blocks.{layer_idx}.attn",
+        )
+        impl.prepare_for_regional_compile(torch.device("cpu"))
+        implementations.append(impl)
+
+    monkeypatch.setattr(fake_sm, "is_supported", fail_eager_dispatch)
+    monkeypatch.setattr(fake_sm, "block_sparse_attn_sm100a", fail_eager_dispatch)
+    monkeypatch.setattr(fake_sm, "block_sparse_attn_sm100a_from_mask", compile_safe_from_mask)
+    monkeypatch.setattr(vsa_h3, "map_to_index", fail_eager_dispatch)
+    monkeypatch.setattr(vsa_h3, "_sm100a_unavailable_reason", fail_eager_dispatch)
+    monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", fail_eager_dispatch)
+
+    torch._dynamo.reset()
+    try:
+        compiled = [torch.compile(impl.forward, backend="eager", fullgraph=True) for impl in implementations]
+        with torch.inference_mode():
+            for layer_idx, run in enumerate(compiled):
+                actual = run(q, k, v, None, meta)
+                expected_delta = 1.0 if layer_idx in meta.dense_layers else 0.0
+                torch.testing.assert_close(actual, q + expected_delta, atol=0, rtol=0)
+    finally:
+        torch._dynamo.reset()
+
+
 def test_default_off_routes_triton(routed, monkeypatch):
     fake_sm, fake_triton, run, _ = routed
     monkeypatch.delenv(VSA_SM100A_ENV, raising=False)

fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_fp8.py

@@ -279,3 +279,119 @@ def test_flashinfer_groupwise_path_pins_output_dtype_and_trtllm_scale_layout(mon
     assert receipt["checkpoint_weight"] is weight
     assert receipt["checkpoint_scale"] is weight_scale
     assert receipt["activation_scale"] is input_scale
+
+
+def test_flashinfer_groupwise_cutlass_path_pads_m_to_a_multiple_of_four(monkeypatch) -> None:
+    """FlashInfer documents that ``m`` must be padded to a multiple of 4, and
+    the CUTLASS groupwise module (the sm12x route) rejects unaligned ``m`` at
+    dispatch with ``cutlass gemm.can_implement failed``. m=559 is this PR's
+    own benchmark prompt length. The route must pad the quantized activation
+    rows and per-token scales, pin the MN scale layout, and slice the padded
+    rows off the output."""
+    m, k, n = 559, 256, 128
+    padded_m = 560
+    input_tensor = torch.zeros(m, k, dtype=torch.bfloat16)
+    weight = torch.zeros(n, k, dtype=torch.float8_e4m3fn)
+    weight_scale = torch.ones(n // 128, k // 128, dtype=torch.float32)
+    quantized_input = torch.zeros(m, k, dtype=torch.float8_e4m3fn)
+    input_scale = torch.ones(m, k // 128, dtype=torch.float32)
+    receipt: dict[str, object] = {}
+
+    def fake_quantize(
+        value: torch.Tensor,
+        group_size: int,
+        *,
+        column_major_scales: bool,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        assert value.data_ptr() == input_tensor.data_ptr()
+        assert group_size == 128
+        assert column_major_scales is False
+        return quantized_input, input_scale
+
+    def fake_gemm(
+        activation: torch.Tensor,
+        checkpoint_weight: torch.Tensor,
+        activation_scale: torch.Tensor,
+        checkpoint_scale: torch.Tensor,
+        *,
+        out_dtype: torch.dtype,
+        backend: str,
+        scale_major_mode: str,
+    ) -> torch.Tensor:
+        assert activation.shape[0] % 4 == 0, "cutlass GEMM requires m padded to a multiple of 4"
+        receipt.update(
+            activation=activation,
+            activation_scale=activation_scale,
+            checkpoint_scale=checkpoint_scale,
+            backend=backend,
+            scale_major_mode=scale_major_mode,
+        )
+        # Tag each row with a bf16-exact value (integers above 256 are not
+        # exactly representable in bf16) so the caller-side slice of the
+        # first m rows is observable.
+        rows = torch.arange(activation.shape[0], dtype=torch.float32) % 256
+        return rows.unsqueeze(1).expand(activation.shape[0], checkpoint_weight.shape[0]).contiguous().to(out_dtype)
+
+    monkeypatch.setattr(h3_fp8, "_get_flashinfer_groupwise_backend", lambda device: "cutlass")
+    monkeypatch.setattr(h3_fp8, "_sglang_per_token_group_quant_fp8", fake_quantize)
+    monkeypatch.setattr(h3_fp8, "_get_flashinfer_groupwise_fp8_gemm", lambda: fake_gemm)
+
+    output = h3_fp8._flashinfer_gemm_w8a8_block_fp8_linear_with_fallback(
+        input_tensor,
+        weight,
+        (128, 128),
+        weight_scale,
+    )
+
+    activation = receipt["activation"]
+    assert activation.shape == (padded_m, k)
+    assert torch.equal(activation[:m].view(torch.uint8), quantized_input.view(torch.uint8))
+    assert not activation[m:].view(torch.uint8).any(), "padded activation rows must be zero"
+
+    activation_scale = receipt["activation_scale"]
+    assert activation_scale.shape == (k // 128, padded_m)
+    assert torch.equal(activation_scale[:, :m], input_scale.transpose(-1, -2))
+    assert not activation_scale[:, m:].any(), "padded scale columns must be zero"
+
+    assert receipt["backend"] == "cutlass"
+    assert receipt["scale_major_mode"] == "MN"
+    assert receipt["checkpoint_scale"].shape == (k // 128, n // 128)
+
+    # The padded rows never reach the caller.
+    assert output.shape == (m, n)
+    assert output.dtype == torch.bfloat16
+    assert torch.equal(output.float()[:, 0], torch.arange(m, dtype=torch.float32) % 256)
+
+
+def test_flashinfer_groupwise_cutlass_path_leaves_aligned_m_unpadded(monkeypatch) -> None:
+    """Aligned token counts must not pay a padding copy on the cutlass route."""
+    m, k, n = 560, 256, 128
+    input_tensor = torch.zeros(m, k, dtype=torch.bfloat16)
+    weight = torch.zeros(n, k, dtype=torch.float8_e4m3fn)
+    weight_scale = torch.ones(k // 128, n // 128, dtype=torch.float32)
+    quantized_input = torch.zeros(m, k, dtype=torch.float8_e4m3fn)
+    input_scale = torch.ones(k // 128, m, dtype=torch.float32)
+    receipt: dict[str, object] = {}
+
+    def fake_quantize(value, group_size, *, column_major_scales):
+        return quantized_input, input_scale
+
+    def fake_gemm(activation, checkpoint_weight, activation_scale, checkpoint_scale, *,
+                  out_dtype, backend, scale_major_mode):
+        receipt.update(activation=activation, activation_scale=activation_scale)
+        return torch.zeros(activation.shape[0], checkpoint_weight.shape[0], dtype=out_dtype)
+
+    monkeypatch.setattr(h3_fp8, "_get_flashinfer_groupwise_backend", lambda device: "cutlass")
+    monkeypatch.setattr(h3_fp8, "_sglang_per_token_group_quant_fp8", fake_quantize)
+    monkeypatch.setattr(h3_fp8, "_get_flashinfer_groupwise_fp8_gemm", lambda: fake_gemm)
+
+    output = h3_fp8._flashinfer_gemm_w8a8_block_fp8_linear_with_fallback(
+        input_tensor,
+        weight,
+        (128, 128),
+        weight_scale,
+    )
+
+    assert receipt["activation"] is quantized_input
+    assert receipt["activation_scale"] is input_scale
+    assert output.shape == (m, n)

fastvideo/tests/inference/test_inference_regional_compile.py

@@ -0,0 +1,331 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Contract tests for the inference-side regional torch.compile port.
+
+The loader applies a per-transformer-block fullgraph compile after the
+transformer loads (``FastVideoArgs.inference_torch_compile``, env
+``FASTVIDEO_INFERENCE_TORCH_COMPILE=1``). These tests pin the two pieces that
+must not drift from the #1718 training-port semantics:
+
+- ``_regional_compile_unsupported_reason``: legacy VSA (and the attention
+  eager escape hatch) degrades to eager, while prepared MiniMax-H3 VSA is
+  admitted to regional fullgraph capture.
+- attention forward dispatch: ordinary instances retain the historical
+  compiler-disabled boundary; regional compile opts in only the selected
+  model's instances.
+- ``_compile_model_regions``: exactly the ``_compile_conditions`` blocks are
+  compiled, fullgraph=True plus inductor ``emulate_precision_casts`` are
+  injected, and ``mode`` kwargs are rejected (torch.compile forbids
+  mode+options).
+
+CPU-safe: the real capture check uses torch's eager recording backend; no CUDA
+or inductor kernel build is needed.
+"""
+
+import sys
+from types import ModuleType, SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from fastvideo.attention import layer as attention_layer
+from fastvideo.attention.layer import DistributedAttention
+from fastvideo.models.dits.minimax_h3 import MiniMaxH3Attention, MiniMaxH3Transformer3DModel
+from fastvideo.models.loader import fsdp_load
+from fastvideo.models.loader.fsdp_load import (
+    _compile_model_regions,
+    _enable_regional_attention_compile,
+    _regional_compile_unsupported_reason,
+)
+
+
+def _init_params_for(backend_name: str | None) -> dict:
+    resolved = None if backend_name is None else SimpleNamespace(name=backend_name)
+    return {"config": SimpleNamespace(_resolved_attention_backend=resolved)}
+
+
+def test_legacy_vsa_backend_degrades_to_eager(monkeypatch) -> None:
+    monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False)
+    backend_name = "VIDEO_SPARSE_ATTN"
+    reason = _regional_compile_unsupported_reason(_init_params_for(backend_name))
+    assert reason is not None
+    assert backend_name in reason
+    assert "eager" in reason
+
+
+@pytest.mark.parametrize("backend_name", [None, "TORCH_SDPA", "VIDEO_SPARSE_ATTN_H3"])
+def test_supported_backends_allow_compile(backend_name, monkeypatch) -> None:
+    monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False)
+    monkeypatch.delenv("FASTVIDEO_H3_VSA_PROBE", raising=False)
+    assert _regional_compile_unsupported_reason(_init_params_for(backend_name)) is None
+
+
+def test_h3_vsa_probe_degrades_regional_compile_to_eager(monkeypatch) -> None:
+    monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False)
+    monkeypatch.setenv("FASTVIDEO_H3_VSA_PROBE", "/tmp/h3-vsa-probe")
+
+    reason = _regional_compile_unsupported_reason(_init_params_for("VIDEO_SPARSE_ATTN_H3"))
+
+    assert reason is not None
+    assert "FASTVIDEO_H3_VSA_PROBE" in reason
+    assert "eager" in reason
+
+
+class _RegionalPrepareProbe:
+
+    def __init__(self) -> None:
+        self.devices: list[torch.device] = []
+
+    def prepare_for_regional_compile(self, device: torch.device) -> None:
+        self.devices.append(device)
+
+
+class _GateBlock(nn.Module):
+
+    def __init__(self, gate_weight: float) -> None:
+        super().__init__()
+        attention = MiniMaxH3Attention.__new__(MiniMaxH3Attention)
+        nn.Module.__init__(attention)
+        attention.to_q = nn.Linear(1, 1, bias=False)
+        attention.to_gate_compress = nn.Linear(1, 1, bias=False)
+        with torch.no_grad():
+            attention.to_gate_compress.weight.fill_(gate_weight)
+        attention._gate_compress_active = None
+        attention.distributed_attention = SimpleNamespace(attn_impl=_RegionalPrepareProbe())
+        self.attn = attention
+
+
+class _GateProbe(nn.Module):
+
+    def __init__(self, attention: MiniMaxH3Attention) -> None:
+        super().__init__()
+        self.attention = attention
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.attention._gate_active():
+            return x + 1
+        return x - 1
+
+
+def _gate_model(*weights: float) -> MiniMaxH3Transformer3DModel:
+    model = MiniMaxH3Transformer3DModel.__new__(MiniMaxH3Transformer3DModel)
+    nn.Module.__init__(model)
+    model.transformer_blocks = nn.ModuleList([_GateBlock(weight) for weight in weights])
+    model.enabled_fusions = frozenset()
+    return model
+
+
+def test_minimax_h3_prepare_for_compile_resolves_loaded_vsa_gates() -> None:
+    model = _gate_model(0.0, 2.0)
+    expected_device = next(model.parameters()).device
+
+    model.prepare_for_compile()
+
+    assert [block.attn._gate_compress_active for block in model.transformer_blocks] == [False, True]
+    for block in model.transformer_blocks:
+        impl = block.attn.distributed_attention.attn_impl
+        assert impl.devices == [expected_device]
+
+
+def test_minimax_h3_prepare_for_compile_does_not_require_quantized_q_weight() -> None:
+    model = _gate_model(2.0)
+    attention = model.transformer_blocks[0].attn
+    removed_weight = attention.to_q._parameters.pop("weight")
+    attention.to_q.register_buffer("_fp8_weight", removed_weight.detach(), persistent=False)
+    expected_device = next(model.parameters()).device
+
+    model.prepare_for_compile()
+
+    impl = attention.distributed_attention.attn_impl
+    assert impl.devices == [expected_device]
+
+
+def test_unprepared_vsa_gate_cannot_mutate_cache_during_compile(monkeypatch) -> None:
+    attention = _gate_model(2.0).transformer_blocks[0].attn
+    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
+
+    with torch.no_grad(), pytest.raises(RuntimeError, match="not resolved before torch.compile"):
+        attention._gate_active()
+    assert attention._gate_compress_active is None
+
+
+@pytest.mark.parametrize(("gate_weight", "expected"), [(0.0, -1.0), (2.0, 1.0)])
+def test_prepared_vsa_gate_is_static_under_fullgraph_compile(gate_weight, expected) -> None:
+    model = _gate_model(gate_weight)
+    model.prepare_for_compile()
+    attention = model.transformer_blocks[0].attn
+    resolved = attention._gate_compress_active
+
+    try:
+        compiled = torch.compile(_GateProbe(attention), backend="eager", fullgraph=True)
+        with torch.no_grad():
+            actual = compiled(torch.zeros(1))
+        torch.testing.assert_close(actual, torch.full((1, ), expected))
+        assert attention._gate_compress_active is resolved
+    finally:
+        torch._dynamo.reset()
+
+
+def test_prepared_zero_vsa_gate_still_runs_in_grad_enabled_training() -> None:
+    model = _gate_model(0.0)
+    model.prepare_for_compile()
+    attention = model.transformer_blocks[0].attn
+
+    with torch.enable_grad():
+        assert attention._gate_active() is True
+    assert attention._gate_compress_active is False
+
+
+def test_attention_compile_escape_hatch_degrades_to_eager(monkeypatch) -> None:
+    monkeypatch.setenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", "1")
+    reason = _regional_compile_unsupported_reason(_init_params_for("TORCH_SDPA"))
+    assert reason is not None
+    assert "FASTVIDEO_DISABLE_ATTENTION_COMPILE" in reason
+
+
+@pytest.mark.parametrize(("fa_version", "allowed"), [("2", True), ("3", False), ("4", True)])
+def test_flash_attention_3_degrades_to_eager(fa_version, allowed, monkeypatch) -> None:
+    monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False)
+    fake_module = ModuleType("fastvideo.attention.utils.flash_attn_default")
+    fake_module.fa_version = fa_version
+    monkeypatch.setitem(sys.modules, fake_module.__name__, fake_module)
+
+    reason = _regional_compile_unsupported_reason(_init_params_for("FLASH_ATTN"))
+    assert (reason is None) is allowed
+    if not allowed:
+        assert "flash-attn 3" in reason
+        assert "eager" in reason
+
+
+def test_default_attention_dispatch_stays_compiler_disabled(monkeypatch) -> None:
+    monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False)
+    disabled_calls: list[str] = []
+
+    def _fake_disable(fn):
+
+        def _disabled(self):
+            disabled_calls.append("disabled")
+            return fn(self)
+
+        return _disabled
+
+    monkeypatch.setattr(attention_layer.torch.compiler, "disable", _fake_disable)
+
+    class _ToyAttention:
+
+        def __init__(self) -> None:
+            self._compile_forward_enabled = not attention_layer._attention_compile_disabled()
+
+        @attention_layer._maybe_compiler_disable
+        def forward(self) -> str:
+            return "forward"
+
+    ordinary = _ToyAttention()
+    assert ordinary.forward() == "forward"
+    assert disabled_calls == ["disabled"]
+
+    # Preserve the existing process-wide escape hatch for callers that opt in
+    # before constructing their attention modules.
+    monkeypatch.setenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", "0")
+    explicit = _ToyAttention()
+    assert explicit.forward() == "forward"
+    assert disabled_calls == ["disabled"]
+
+
+class _BareDistributedAttention(DistributedAttention):
+
+    def __init__(self) -> None:
+        nn.Module.__init__(self)
+        self._compile_forward_enabled = False
+
+
+class _CompileProbe(nn.Module):
+
+    def __init__(self, enabled: bool) -> None:
+        super().__init__()
+        self._compile_forward_enabled = enabled
+
+    @attention_layer._maybe_compiler_disable
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x + 1
+
+
+def test_attention_instance_flag_controls_real_fullgraph_capture() -> None:
+    try:
+        ordinary = torch.compile(_CompileProbe(False), backend="eager", fullgraph=True)
+        with pytest.raises(torch._dynamo.exc.Unsupported, match="Skip calling"):
+            ordinary(torch.ones(2))
+
+        regional = torch.compile(_CompileProbe(True), backend="eager", fullgraph=True)
+        torch.testing.assert_close(regional(torch.ones(2)), torch.full((2,), 2.0))
+    finally:
+        torch._dynamo.reset()
+
+
+def test_regional_compile_enables_only_loaded_model_attention() -> None:
+    selected_model = nn.Sequential(_BareDistributedAttention(), _BareDistributedAttention())
+    unrelated_model = nn.Sequential(_BareDistributedAttention())
+
+    assert _enable_regional_attention_compile(selected_model) == 2
+    assert all(module._compile_forward_enabled for module in selected_model)
+    assert not unrelated_model[0]._compile_forward_enabled
+
+
+class _Block(nn.Module):
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.linear = nn.Linear(4, 4)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.linear(x)
+
+
+class _Toy(nn.Module):
+    _compile_conditions = [lambda name, module: name.startswith("blocks.") and name.count(".") == 1]
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.blocks = nn.ModuleList([_Block() for _ in range(3)])
+        self.proj_out = nn.Linear(4, 4)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for block in self.blocks:
+            x = block(x)
+        return self.proj_out(x)
+
+
+def test_compile_model_regions_injects_fullgraph_and_precision_casts(monkeypatch) -> None:
+    captured: list[dict] = []
+
+    def _fake_compile(fn, **kwargs):
+        captured.append(kwargs)
+        return fn
+
+    monkeypatch.setattr(fsdp_load.torch, "compile", _fake_compile)
+    model = _Toy()
+    count = _compile_model_regions(model, {})
+    # The three repeated blocks compile; proj_out and the root stay eager.
+    assert count == 3
+    assert len(captured) == 3
+    for kwargs in captured:
+        assert kwargs["fullgraph"] is True
+        assert kwargs["options"] == {"emulate_precision_casts": True}
+
+
+def test_compile_model_regions_rejects_mode_kwargs() -> None:
+    with pytest.raises(ValueError, match="mode"):
+        _compile_model_regions(_Toy(), {"mode": "reduce-overhead"})
+
+
+def test_compile_model_regions_requires_conditions_and_matches(monkeypatch) -> None:
+    monkeypatch.setattr(fsdp_load.torch, "compile", lambda fn, **kwargs: fn)
+    plain = nn.Linear(4, 4)
+    with pytest.raises(ValueError, match="_compile_conditions"):
+        _compile_model_regions(plain, {})
+
+    class _NoMatch(_Toy):
+        _compile_conditions = [lambda name, module: False]
+
+    with pytest.raises(ValueError, match="matched"):
+        _compile_model_regions(_NoMatch(), {})

fastvideo/tests/stages/test_minimax_h3_vae_streaming.py

@@ -61,7 +61,8 @@ def test_reference_video_encode_keeps_pixels_on_cpu() -> None:
         media_type="video",
         frames=np.zeros((22, 16, 16, 3), dtype=np.uint8),
     )
-    rows = stage._encode_visual_rows([reference], torch.device("cpu"))
+    args = SimpleNamespace(vae_parallel_encode=False)
+    rows = stage._encode_visual_rows([reference], torch.device("cpu"), args)
 
     assert observed["pixels"].dtype == torch.uint8
     assert observed["pixels"].device.type == "cpu"
@@ -96,7 +97,7 @@ def test_decode_stage_uses_cpu_output_buffer(monkeypatch) -> None:
     monkeypatch.setattr(minimax_h3_decoding, "get_local_torch_device", lambda: torch.device("cpu"))
     result = MiniMaxH3VideoDecodingStage(VAE(), SimpleNamespace(patch_size=(1, 1, 1))).forward(
         batch,
-        SimpleNamespace(output_type="pil", pin_cpu_memory=False, vae_cpu_offload=False),
+        SimpleNamespace(output_type="pil", pin_cpu_memory=False, vae_cpu_offload=False, vae_parallel_decode=False),
     )
 
     torch.testing.assert_close(observed["latents"], latents)
@@ -114,8 +115,9 @@ def test_decode_stages_skip_vae_on_non_output_rank(monkeypatch) -> None:
             raise AssertionError("non-output ranks must not execute a VAE")
 
     monkeypatch.setattr(minimax_h3_decoding, "model_parallel_is_initialized", lambda: True)
-    monkeypatch.setattr(minimax_h3_decoding, "get_world_group", lambda: SimpleNamespace(is_first_rank=False))
-    args = SimpleNamespace(output_type="pil", pin_cpu_memory=False, vae_cpu_offload=True)
+    monkeypatch.setattr(minimax_h3_decoding, "get_sp_group",
+                        lambda: SimpleNamespace(is_first_rank=False, world_size=4, rank_in_group=1))
+    args = SimpleNamespace(output_type="pil", pin_cpu_memory=False, vae_cpu_offload=True, vae_parallel_decode=False)
 
     video = MiniMaxH3VideoDecodingStage(VAE(), SimpleNamespace()).forward(ForwardBatch(data_type="video"), args)
     assert video.output.shape == (0, 3, 0, 0, 0)
@@ -128,3 +130,56 @@ def test_decode_stages_skip_vae_on_non_output_rank(monkeypatch) -> None:
     assert audio.latents is None
     assert audio.audio_latents is None
     assert MINIMAX_H3_LAYOUT_KEY not in audio.extra
+
+
+def test_parallel_decode_runs_on_every_rank(monkeypatch) -> None:
+    """With vae_parallel_decode, non-leader ranks must enter the decode body
+    (the collectives inside require uniform participation) and only the
+    leader owns the CPU output buffer."""
+    latent_shape = (1, 4, 2, 4, 4)
+    rows = patchify_video_latents(torch.randn(latent_shape), (1, 1, 1))
+    calls = []
+
+    class VAE:
+
+        def to(self, device):
+            return self
+
+        def denormalize_latents(self, decoded_latents):
+            return decoded_latents
+
+        def decoded_pixel_shape(self, shape):
+            return (1, 3, 5, 16, 16)
+
+    def fake_parallel(vae, latents, output, group, strategy):
+        calls.append((group.rank_in_group, output, strategy))
+        if output is not None:
+            output.fill_(0.5)
+        return output
+
+    monkeypatch.setattr(minimax_h3_decoding, "get_local_torch_device", lambda: torch.device("cpu"))
+    monkeypatch.setattr(minimax_h3_decoding, "model_parallel_is_initialized", lambda: True)
+    monkeypatch.setattr(minimax_h3_decoding, "decode_to_pixels_parallel", fake_parallel)
+    args = SimpleNamespace(output_type="pil",
+                           pin_cpu_memory=False,
+                           vae_cpu_offload=False,
+                           vae_parallel_decode=True,
+                           vae_parallel_decode_strategy="gather")
+
+    for rank, is_first in ((0, True), (2, False)):
+        monkeypatch.setattr(
+            minimax_h3_decoding, "get_sp_group",
+            lambda rank=rank, is_first=is_first: SimpleNamespace(is_first_rank=is_first,
+                                                                 world_size=4,
+                                                                 rank_in_group=rank))
+        batch = ForwardBatch(data_type="video", latents=rows.clone(), raw_latent_shape=latent_shape)
+        batch.extra[MINIMAX_H3_LAYOUT_KEY] = _layout(rows.shape[0], latent_shape)
+        result = MiniMaxH3VideoDecodingStage(VAE(), SimpleNamespace(patch_size=(1, 1, 1))).forward(batch, args)
+        if is_first:
+            assert result.output.shape == (1, 3, 5, 16, 16)
+            assert torch.all(result.output == 0.5)
+        else:
+            assert result.output.shape == (0, 3, 0, 0, 0)
+
+    assert [(rank, output is not None) for rank, output, _ in calls] == [(0, True), (2, False)]
+    assert all(strategy == "gather" for _, _, strategy in calls)

fastvideo/tests/transformers/test_minimax_h3_fusion_compile.py

@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fullgraph trace contract for the MiniMax-H3 inference fusions."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+from torch._subclasses.fake_tensor import FakeTensorMode
+
+from fastvideo.models.dits.minimax_h3_fusions import (
+    fused_qknorm_rope,
+    fused_residual_gate_rmsnorm_modulate,
+    fused_rmsnorm_modulate,
+    minimax_h3_swiglu,
+)
+
+
+def _all_fusions(
+    x: torch.Tensor,
+    branch: torch.Tensor,
+    weight: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    gate: torch.Tensor,
+    index: torch.Tensor,
+    q: torch.Tensor,
+    q_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    packed_mlp: torch.Tensor,
+) -> tuple[torch.Tensor, ...]:
+    norm1 = fused_rmsnorm_modulate(x, weight, scale, shift, index, 1e-5)
+    hidden, norm2 = fused_residual_gate_rmsnorm_modulate(
+        x,
+        branch,
+        gate,
+        weight,
+        scale,
+        shift,
+        index,
+        1e-5,
+    )
+    q_out = fused_qknorm_rope(q, q_weight, cos, sin, 1e-5)
+    mlp_out = minimax_h3_swiglu(packed_mlp)
+    return norm1, hidden, norm2, q_out, mlp_out
+
+
+def test_minimax_h3_fusions_are_dynamic_fullgraph_traceable() -> None:
+    """Dynamo sees four opaque ops without requiring a physical CUDA GPU."""
+    graphs: list[torch.fx.GraphModule] = []
+
+    def capture_graph(graph: torch.fx.GraphModule, example_inputs: list[torch.Tensor]):
+        del example_inputs
+        graphs.append(graph)
+        return graph.forward
+
+    compiled = torch.compile(_all_fusions, backend=capture_graph, fullgraph=True, dynamic=True)
+
+    with torch.no_grad(), FakeTensorMode():
+        x = torch.empty((2, 3, 8), device="cuda", dtype=torch.bfloat16)
+        weight = torch.empty((8,), device="cuda", dtype=torch.bfloat16)
+        table = torch.empty((6, 8), device="cuda", dtype=torch.bfloat16)
+        index = torch.empty((3,), device="cuda", dtype=torch.int64)
+        q = torch.empty((2, 3, 4, 8), device="cuda", dtype=torch.bfloat16)
+        rotary = torch.empty((3, 8), device="cuda", dtype=torch.bfloat16)
+        packed_mlp = torch.empty((2, 3, 16), device="cuda", dtype=torch.bfloat16)
+
+        outputs = compiled(
+            x,
+            x,
+            weight,
+            table,
+            table,
+            table,
+            index,
+            q,
+            weight,
+            rotary,
+            rotary,
+            packed_mlp,
+        )
+
+    assert tuple(output.shape for output in outputs) == (
+        (2, 3, 8),
+        (2, 3, 8),
+        (2, 3, 8),
+        (2, 3, 4, 8),
+        (2, 3, 8),
+    )
+    assert all(output.is_contiguous() for output in outputs)
+    assert len(graphs) == 1
+    call_targets = {node.target for node in graphs[0].graph.nodes if node.op == "call_function"}
+    assert {
+        torch.ops.fastvideo._minimax_h3_rmsnorm_modulate,
+        torch.ops.fastvideo._minimax_h3_residual_gate_rmsnorm_modulate,
+        torch.ops.fastvideo._minimax_h3_qknorm_rope,
+        torch.ops.fastvideo._minimax_h3_swiglu,
+    } <= call_targets
+
+
+@pytest.mark.gpu
+def test_minimax_h3_fusion_custom_ops_match_real_fullgraph_cuda() -> None:
+    """Cross-check fake metadata and execute the real CUDA dispatcher path."""
+    if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported():
+        pytest.skip("BF16 CUDA is required")
+    pytest.importorskip("triton")
+
+    device = torch.device("cuda")
+    generator = torch.Generator(device=device).manual_seed(2026)
+    x = torch.randn((2, 3, 128), generator=generator, device=device, dtype=torch.bfloat16)
+    branch = torch.randn_like(x)
+    weight = torch.randn((128,), generator=generator, device=device, dtype=torch.bfloat16)
+    scale = torch.randn((6, 128), generator=generator, device=device, dtype=torch.bfloat16)
+    shift = torch.randn((6, 128), generator=generator, device=device, dtype=torch.bfloat16)
+    gate = torch.randn((6, 128), generator=generator, device=device, dtype=torch.bfloat16)
+    index = torch.tensor((0, 3, 5), device=device, dtype=torch.int64)
+    q = torch.randn((2, 3, 4, 128), generator=generator, device=device, dtype=torch.bfloat16)
+    q_weight = torch.randn((128,), generator=generator, device=device, dtype=torch.bfloat16).contiguous()
+    rotary = torch.randn((3, 96), generator=generator, device=device, dtype=torch.bfloat16).contiguous()
+    packed_mlp = torch.randn((2, 3, 256), generator=generator, device=device, dtype=torch.bfloat16)
+
+    opchecks = (
+        (
+            torch.ops.fastvideo._minimax_h3_rmsnorm_modulate,
+            (x, weight, scale, shift, index, 1e-5),
+        ),
+        (
+            torch.ops.fastvideo._minimax_h3_residual_gate_rmsnorm_modulate,
+            (x, branch, gate, weight, scale, shift, index, 1e-5),
+        ),
+        (
+            torch.ops.fastvideo._minimax_h3_qknorm_rope,
+            (q, q_weight, rotary, rotary.clone(), 1e-5),
+        ),
+        (torch.ops.fastvideo._minimax_h3_swiglu, (packed_mlp,)),
+    )
+    for op, args in opchecks:
+        torch.library.opcheck(op, args)
+
+    compiled = torch.compile(_all_fusions, fullgraph=True, dynamic=True)
+    inputs = (
+        x,
+        branch,
+        weight,
+        scale,
+        shift,
+        gate,
+        index,
+        q,
+        q_weight,
+        rotary,
+        rotary.clone(),
+        packed_mlp,
+    )
+    with torch.no_grad():
+        expected = _all_fusions(*inputs)
+        actual = compiled(*inputs)
+    for actual_output, expected_output in zip(actual, expected, strict=True):
+        torch.testing.assert_close(actual_output, expected_output, atol=0.0, rtol=0.0)

fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py

@@ -9,6 +9,42 @@ import torch
 from fastvideo.platforms import AttentionBackendEnum
 
 
+def test_only_packed_minimax_h3_blocks_enable_fa4_varlen(monkeypatch: pytest.MonkeyPatch) -> None:
+    import fastvideo.models.dits.minimax_h3 as h3
+
+    routed: list[bool] = []
+
+    class RecordingAttention(torch.nn.Module):
+
+        def __init__(self, *args, fa4_packed_varlen: bool = False, **kwargs) -> None:
+            super().__init__()
+            del args, kwargs
+            routed.append(fa4_packed_varlen)
+
+    monkeypatch.setattr(h3, "MiniMaxH3Attention", RecordingAttention)
+    common = dict(
+        hidden_size=8,
+        num_attention_heads=1,
+        attention_head_dim=8,
+        ffn_dim=16,
+        norm_eps=1e-5,
+        qk_norm_eps=1e-5,
+        supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, ),
+        quant_config=None,
+    )
+    h3.MiniMaxH3TransformerBlock(
+        **common,
+        time_embed_dim=4,
+        prefix="minimax_h3.transformer_blocks.0",
+    )
+    h3.MiniMaxH3TokenRefinerBlock(
+        **common,
+        prefix="minimax_h3.token_refiner.refiner_blocks.0",
+    )
+
+    assert routed == [True, False]
+
+
 @pytest.mark.parametrize(
     ("raw", "expected"),
     [

fastvideo/tests/vaes/test_minimax_h3_parallel_vae.py

@@ -0,0 +1,314 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU tests for sequence-parallel MiniMax-H3 VAE chunk decode / clip encode.
+
+The collective transport is simulated with a threaded fake group (one thread
+per simulated rank, barrier-synchronized slots), so the REAL drivers in
+``fastvideo.models.vaes.minimax_h3_parallel`` — chunk assignment, placeholder
+rounds, metadata broadcast, gathered-segment assembly, halo/blend math — run
+end to end on CPU and are checked bit-exactly against the serial APIs.
+"""
+
+import threading
+
+import pytest
+import torch
+from torch.testing import assert_close
+
+from fastvideo.configs.models.vaes.minimax_h3_video import (
+    MiniMaxH3VideoVAEArchConfig,
+    MiniMaxH3VideoVAEConfig,
+)
+from fastvideo.models.vaes.minimax_h3_parallel import (
+    DECODE_GATHER_STRATEGIES,
+    DEFAULT_DECODE_GATHER_STRATEGY,
+    decode_to_pixels_parallel,
+    encode_pixels_parallel,
+    parallel_chunk_indices,
+)
+from fastvideo.models.vaes.minimax_h3_video import AutoencoderKLMiniMaxH3
+
+
+def _tiny_vae(token_drop: int = 3) -> AutoencoderKLMiniMaxH3:
+    arch = MiniMaxH3VideoVAEArchConfig(
+        latent_channels=4,
+        block_out_channels=(32, 32),
+        layers_per_block=1,
+        spatial_downsample_factors=(2, 2),
+        temporal_downsample_factors=(2, 2),
+        decoder_num_layers=1,
+        decoder_num_attention_heads=1,
+        decoder_attention_head_dim=8,
+        decoder_num_register_tokens=2,
+        decoder_ffn_mult=1,
+        token_drop=token_drop,
+        latents_mean=(0.0, ) * 4,
+        latents_std=(1.0, ) * 4,
+    )
+    return AutoencoderKLMiniMaxH3(
+        MiniMaxH3VideoVAEConfig(
+            arch_config=arch,
+            use_tiling=False,
+            use_temporal_tiling=False,
+            use_parallel_tiling=False,
+        )).eval()
+
+
+class _ThreadedFakeGroup:
+    """Barrier-synchronized in-process stand-in for a GroupCoordinator.
+
+    One thread per simulated rank runs the SPMD driver; ``gather`` /
+    ``all_gather`` / ``broadcast_object`` rendezvous through shared slots
+    with a double barrier (all writes land, everyone reads, then slots are
+    reusable). Matches the GroupCoordinator call signatures the drivers use.
+    """
+
+    def __init__(self, world_size: int) -> None:
+        self.world_size = world_size
+        self._local = threading.local()
+        self._barrier = threading.Barrier(world_size)
+        self._slots: list = [None] * world_size
+        self._object = None
+
+    @property
+    def rank_in_group(self) -> int:
+        return self._local.rank
+
+    def broadcast_object(self, obj=None, src: int = 0):
+        if self.world_size == 1:
+            return obj
+        if self.rank_in_group == src:
+            self._object = obj
+        self._barrier.wait()
+        received = self._object
+        self._barrier.wait()
+        return received
+
+    def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
+        if self.world_size == 1:
+            return input_
+        self._slots[self.rank_in_group] = input_
+        self._barrier.wait()
+        gathered = torch.cat([slot for slot in self._slots], dim=dim)
+        self._barrier.wait()
+        return gathered
+
+    def gather(self, input_: torch.Tensor, dst: int = 0, dim: int = -1):
+        if self.world_size == 1:
+            return input_
+        self._slots[self.rank_in_group] = input_
+        self._barrier.wait()
+        gathered = torch.cat([slot for slot in self._slots], dim=dim) if self.rank_in_group == dst else None
+        self._barrier.wait()
+        return gathered
+
+    def run(self, fn) -> list:
+        """Run ``fn(rank)`` on one thread per rank; re-raise the first error."""
+        results: list = [None] * self.world_size
+        errors: list = [None] * self.world_size
+
+        def _target(rank: int) -> None:
+            self._local.rank = rank
+            try:
+                # inference_mode is thread-local; the drivers run inference-only.
+                with torch.inference_mode():
+                    results[rank] = fn(rank)
+            except BaseException as error:  # noqa: BLE001 - propagate to the test
+                errors[rank] = error
+                self._barrier.abort()
+
+        threads = [threading.Thread(target=_target, args=(rank, )) for rank in range(self.world_size)]
+        for thread in threads:
+            thread.start()
+        for thread in threads:
+            thread.join()
+        for error in errors:
+            if error is not None and not isinstance(error, threading.BrokenBarrierError):
+                raise error
+        for error in errors:
+            if error is not None:
+                raise error
+        return results
+
+
+@pytest.mark.parametrize("num_chunks,world_size", ((0, 4), (1, 4), (7, 4), (8, 4), (20, 4), (5, 3), (2, 5)))
+def test_parallel_chunk_indices_partition(num_chunks: int, world_size: int) -> None:
+    """Round-robin ownership covers every chunk exactly once, in order."""
+    owned = [parallel_chunk_indices(num_chunks, world_size, rank) for rank in range(world_size)]
+    flattened = sorted(index for indices in owned for index in indices)
+    assert flattened == list(range(num_chunks))
+    for rank, indices in enumerate(owned):
+        assert indices == sorted(indices)
+        assert all(index % world_size == rank for index in indices)
+        # Round-robin balance: no rank holds more than one extra chunk.
+        assert len(indices) in (num_chunks // world_size, -(-num_chunks // world_size))
+
+
+def test_parallel_chunk_indices_validates() -> None:
+    with pytest.raises(ValueError, match="world_size"):
+        parallel_chunk_indices(4, 0, 0)
+    with pytest.raises(ValueError, match="rank_in_group"):
+        parallel_chunk_indices(4, 2, 2)
+    with pytest.raises(ValueError, match="num_chunks"):
+        parallel_chunk_indices(-1, 2, 0)
+
+
+# Latent frames cover: one padded chunk (3), pad on the intra-clip tail (6),
+# two blended chunks (12), three chunks plus pad trim (13). World sizes cover
+# fewer chunks than ranks, uneven rounds, and the exact-multiple case.
+@pytest.mark.parametrize("world_size", (2, 3, 4, 5))
+@pytest.mark.parametrize("latent_frames", (3, 6, 12, 13))
+@pytest.mark.parametrize("strategy", DECODE_GATHER_STRATEGIES)
+@torch.inference_mode()
+def test_parallel_decode_matches_serial(world_size: int, latent_frames: int, strategy: str) -> None:
+    torch.manual_seed(20260821 + latent_frames)
+    vae = _tiny_vae()
+    latents = torch.randn(1, 4, latent_frames, 4, 4)
+    expected = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32)
+    vae.decode_to_pixels(latents, expected)
+
+    group = _ThreadedFakeGroup(world_size)
+
+    def _rank_main(rank: int):
+        output = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32) if rank == 0 else None
+        return decode_to_pixels_parallel(vae, latents.clone(), output, group, strategy=strategy)
+
+    results = group.run(_rank_main)
+    assert all(result is None for result in results[1:])
+    assert_close(results[0], expected, atol=0.0, rtol=0.0)
+
+
+@torch.inference_mode()
+def test_parallel_decode_without_token_drop() -> None:
+    """token_drop == 0 has no overlap halo; the assembler must skip blending."""
+    torch.manual_seed(20260822)
+    vae = _tiny_vae(token_drop=0)
+    assert vae.frame_overlap == 0
+    latents = torch.randn(1, 4, 10, 4, 4)
+    expected = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32)
+    vae.decode_to_pixels(latents, expected)
+
+    group = _ThreadedFakeGroup(3)
+
+    def _rank_main(rank: int):
+        output = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32) if rank == 0 else None
+        return decode_to_pixels_parallel(vae, latents.clone(), output, group)
+
+    results = group.run(_rank_main)
+    assert_close(results[0], expected, atol=0.0, rtol=0.0)
+
+
+@torch.inference_mode()
+def test_parallel_decode_batched_slicing_matches_serial() -> None:
+    torch.manual_seed(20260823)
+    vae = _tiny_vae()
+    vae.enable_slicing()
+    latents = torch.randn(2, 4, 7, 4, 4)
+    expected = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32)
+    vae.decode_to_pixels(latents, expected)
+
+    group = _ThreadedFakeGroup(2)
+
+    def _rank_main(rank: int):
+        output = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32) if rank == 0 else None
+        return decode_to_pixels_parallel(vae, latents.clone(), output, group)
+
+    results = group.run(_rank_main)
+    assert_close(results[0], expected, atol=0.0, rtol=0.0)
+
+
+@torch.inference_mode()
+def test_parallel_decode_world_size_one_is_serial() -> None:
+    torch.manual_seed(20260824)
+    vae = _tiny_vae()
+    latents = torch.randn(1, 4, 7, 4, 4)
+    expected = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32)
+    vae.decode_to_pixels(latents, expected)
+
+    group = _ThreadedFakeGroup(1)
+
+    def _rank_main(rank: int):
+        output = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32)
+        return decode_to_pixels_parallel(vae, latents, output, group)
+
+    results = group.run(_rank_main)
+    assert_close(results[0], expected, atol=0.0, rtol=0.0)
+
+
+def test_parallel_decode_validates_buffers_and_strategy() -> None:
+    vae = _tiny_vae()
+    latents = torch.randn(1, 4, 7, 4, 4)
+    output = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32)
+    group = _ThreadedFakeGroup(1)
+    group._local.rank = 0
+
+    with pytest.raises(ValueError, match="strategy"):
+        decode_to_pixels_parallel(vae, latents, output, group, strategy="scatter")
+    with pytest.raises(ValueError, match="must provide the CPU output buffer"):
+        decode_to_pixels_parallel(vae, latents, None, group)
+    with pytest.raises(ValueError, match="CPU float32 tensor"):
+        decode_to_pixels_parallel(vae, latents, output[:, :, :-1], group)
+
+    group._local.rank = 1  # simulate a non-leader passing a buffer
+    group.world_size = 2
+    with pytest.raises(ValueError, match="Only the first sequence-parallel rank"):
+        decode_to_pixels_parallel(vae, latents, output, group)
+
+
+@pytest.mark.parametrize("world_size", (2, 4))
+@pytest.mark.parametrize("num_frames", (16, 22, 40))
+@torch.inference_mode()
+def test_parallel_encode_matches_serial(world_size: int, num_frames: int) -> None:
+    """Every rank must hold the full serial moments, bit for bit."""
+    torch.manual_seed(20260825 + num_frames)
+    vae = _tiny_vae()
+    pixels = torch.randint(0, 256, (1, 3, num_frames, 16, 16), dtype=torch.uint8)
+    expected = vae.encode_pixels(pixels).latent_dist.parameters
+
+    group = _ThreadedFakeGroup(world_size)
+    results = group.run(lambda rank: encode_pixels_parallel(vae, pixels, group).latent_dist.parameters)
+    for moments in results:
+        assert_close(moments, expected, atol=0.0, rtol=0.0)
+
+
+@torch.inference_mode()
+def test_parallel_encode_float_and_batched_slicing() -> None:
+    torch.manual_seed(20260826)
+    vae = _tiny_vae()
+    vae.enable_slicing()
+    pixels = torch.rand(2, 3, 22, 16, 16)
+    expected = vae.encode_pixels(pixels).latent_dist.parameters
+
+    group = _ThreadedFakeGroup(3)
+    results = group.run(lambda rank: encode_pixels_parallel(vae, pixels, group).latent_dist.parameters)
+    for moments in results:
+        assert_close(moments, expected, atol=0.0, rtol=0.0)
+
+
+def test_parallel_encode_validates_input() -> None:
+    vae = _tiny_vae()
+    group = _ThreadedFakeGroup(1)
+    group._local.rank = 0
+    with pytest.raises(ValueError, match="must remain on CPU"):
+        encode_pixels_parallel(vae, torch.empty(1, 3, 4, 16, 16, device="meta"), group)
+    with pytest.raises(TypeError, match="uint8 or a floating-point"):
+        encode_pixels_parallel(vae, torch.zeros(1, 3, 4, 16, 16, dtype=torch.int32), group)
+    with pytest.raises(ValueError, match="must have shape"):
+        encode_pixels_parallel(vae, torch.zeros(1, 4, 4, 16, 16), group)
+
+
+def test_fastvideo_args_strategy_literals_match_module() -> None:
+    """fastvideo_args mirrors the strategy literals to avoid importing model
+    modules at args construction; keep the two in sync."""
+    from fastvideo.fastvideo_args import FastVideoArgs
+
+    args = FastVideoArgs(model_path="test/parallel-vae")
+    assert args.vae_parallel_decode is False
+    assert args.vae_parallel_encode is False
+    assert args.vae_parallel_decode_strategy == DEFAULT_DECODE_GATHER_STRATEGY
+    assert args.vae_parallel_decode_strategy in DECODE_GATHER_STRATEGIES
+
+    for strategy in DECODE_GATHER_STRATEGIES:
+        assert FastVideoArgs(model_path="test/parallel-vae",
+                             vae_parallel_decode_strategy=strategy).vae_parallel_decode_strategy == strategy
+    with pytest.raises(ValueError, match="vae_parallel_decode_strategy"):
+        FastVideoArgs(model_path="test/parallel-vae", vae_parallel_decode_strategy="scatter")

fastvideo/tests/vaes/test_minimax_h3_parallel_vae_gpu.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+"""GPU regression test for sequence-parallel MiniMax-H3 VAE decode/encode.
+
+Requires a multi-GPU torchrun launch (real NCCL collectives across an SP
+group); skipped otherwise:
+
+    torchrun --nproc-per-node=4 -m pytest \
+        fastvideo/tests/vaes/test_minimax_h3_parallel_vae_gpu.py -q
+
+Asserts the parallel drivers are bitwise equal to the serial rank-local
+decode/encode under the pipeline's fp16 autocast, for both transport
+strategies, and that repeated parallel runs are deterministic.
+"""
+
+import os
+
+import pytest
+import torch
+from torch.testing import assert_close
+
+from fastvideo.configs.models.vaes.minimax_h3_video import (
+    MiniMaxH3VideoVAEArchConfig,
+    MiniMaxH3VideoVAEConfig,
+)
+from fastvideo.models.vaes.minimax_h3_video import AutoencoderKLMiniMaxH3
+
+_WORLD_SIZE = int(os.environ.get("WORLD_SIZE", "1"))
+
+
+def _tiny_vae() -> AutoencoderKLMiniMaxH3:
+    """Same tiny geometry as test_minimax_h3_parallel_vae (test dirs are not packages)."""
+    arch = MiniMaxH3VideoVAEArchConfig(
+        latent_channels=4,
+        block_out_channels=(32, 32),
+        layers_per_block=1,
+        spatial_downsample_factors=(2, 2),
+        temporal_downsample_factors=(2, 2),
+        decoder_num_layers=1,
+        decoder_num_attention_heads=1,
+        decoder_attention_head_dim=8,
+        decoder_num_register_tokens=2,
+        decoder_ffn_mult=1,
+        latents_mean=(0.0, ) * 4,
+        latents_std=(1.0, ) * 4,
+    )
+    return AutoencoderKLMiniMaxH3(
+        MiniMaxH3VideoVAEConfig(
+            arch_config=arch,
+            use_tiling=False,
+            use_temporal_tiling=False,
+            use_parallel_tiling=False,
+        )).eval()
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA"),
+    pytest.mark.skipif(_WORLD_SIZE < 2, reason="requires a torchrun launch with WORLD_SIZE > 1"),
+]
+
+
+@pytest.fixture(scope="module")
+def sp_group():
+    from fastvideo.distributed import get_sp_group, maybe_init_distributed_environment_and_model_parallel
+    maybe_init_distributed_environment_and_model_parallel(1, _WORLD_SIZE)
+    return get_sp_group()
+
+
+@pytest.mark.parametrize("strategy", ("gather", "all_gather"))
+@pytest.mark.parametrize("latent_frames", (3, 13))
+@torch.no_grad()
+def test_parallel_decode_bitwise_matches_serial_on_gpu(sp_group, strategy: str, latent_frames: int) -> None:
+    from fastvideo.models.vaes.minimax_h3_parallel import decode_to_pixels_parallel
+
+    device = torch.device("cuda", torch.cuda.current_device())
+    torch.manual_seed(20260821)  # identical weights on every rank
+    vae = _tiny_vae().to(device)
+    latents = torch.randn(1, 4, latent_frames, 4, 4, generator=torch.Generator().manual_seed(7)).to(device)
+
+    with torch.autocast(device_type="cuda", dtype=torch.float16):
+        expected = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32, pin_memory=True)
+        vae.decode_to_pixels(latents, expected)
+
+        outputs = []
+        for _ in range(3):  # repeat-determinism
+            output = None
+            if sp_group.is_first_rank:
+                output = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32, pin_memory=True)
+            result = decode_to_pixels_parallel(vae, latents, output, sp_group, strategy=strategy)
+            outputs.append(result.clone() if result is not None else None)
+
+    if sp_group.is_first_rank:
+        for output in outputs:
+            assert_close(output, expected, atol=0.0, rtol=0.0)
+    else:
+        assert all(output is None for output in outputs)
+
+
+@torch.no_grad()
+def test_parallel_encode_bitwise_matches_serial_on_gpu(sp_group) -> None:
+    from fastvideo.models.vaes.minimax_h3_parallel import encode_pixels_parallel
+
+    device = torch.device("cuda", torch.cuda.current_device())
+    torch.manual_seed(20260821)
+    vae = _tiny_vae().to(device)
+    pixels = torch.randint(0, 256, (1, 3, 40, 16, 16), dtype=torch.uint8,
+                           generator=torch.Generator().manual_seed(9))
+
+    expected = vae.encode_pixels(pixels).latent_dist.parameters
+    for _ in range(3):
+        moments = encode_pixels_parallel(vae, pixels, sp_group).latent_dist.parameters
+        assert_close(moments, expected, atol=0.0, rtol=0.0)

fastvideo/tests/vaes/test_minimax_h3_vae_compile.py

@@ -6,7 +6,7 @@ The final test is a CUDA regression gate for the reduce-overhead tile path
 """
 
 from contextlib import contextmanager
-from types import MethodType, SimpleNamespace
+from types import SimpleNamespace
 from typing import Any
 from unittest.mock import Mock, patch
 
@@ -141,6 +141,8 @@ def test_video_attention_uses_selected_fastvideo_backend() -> None:
         "softmax_scale": 0.5,
         "num_kv_heads": 2,
         "causal": False,
+        # The FP32-pinned VAE opts out of the FASTVIDEO_NVFP4_FA4 env opt-in.
+        "nvfp4_fa4": False,
     }
     assert backend_call["shapes"] == ((1, 3, 2, 4), ) * 3
     assert backend_call["metadata"] is None
@@ -178,14 +180,91 @@ def test_compile_with_conditions_selects_minimax_h3_video_decoder() -> None:
     _assert_dynamic_compile_selects_decoder(AutoencoderKLMiniMaxH3, MiniMaxH3VideoViTDecoder3d)
 
 
-def test_project_decoder_tile_uses_reduce_overhead_compile() -> None:
-    """Compile the per-tile decoder-input projection with CUDA Graph replay."""
-    _assert_reduce_overhead_compile(AutoencoderKLMiniMaxH3._project_decoder_tile)
+def test_tile_helpers_stay_eager_by_default() -> None:
+    """Leave the tile helpers uncompiled unless the VAE compile opt-in runs.
+
+    Default users must keep pre-PR eager behavior: no inductor/triton
+    requirement, no first-decode compile latency, and no permanent cudagraph
+    memory pools unless ``enable_torch_compile_vae`` was requested.
+    """
+    assert AutoencoderKLMiniMaxH3._tile_helpers_compiled is False
+    for helper in (AutoencoderKLMiniMaxH3._stitch_tiles, AutoencoderKLMiniMaxH3._project_decoder_tile):
+        assert not hasattr(helper, "get_compiler_config")
+        assert not hasattr(helper, "_torchdynamo_orig_callable")
 
 
-def test_stitch_tiles_uses_reduce_overhead_compile() -> None:
-    """Compile spatial tile blending and concatenation with CUDA Graph replay."""
-    _assert_reduce_overhead_compile(AutoencoderKLMiniMaxH3._stitch_tiles)
+def test_prepare_for_compile_installs_reduce_overhead_tile_helpers() -> None:
+    """Compile both tile helpers with CUDA Graph replay under the opt-in hook.
+
+    ``ComposedPipelineBase._maybe_compile_pipeline_module`` invokes
+    ``prepare_for_compile`` only when ``enable_torch_compile_vae`` is set,
+    right before the decoder is compiled through ``_compile_conditions``.
+    """
+    vae = _empty_typed_module(AutoencoderKLMiniMaxH3)
+    assert vae._tile_helpers_compiled is False
+
+    vae.prepare_for_compile()
+
+    assert vae._tile_helpers_compiled is True
+    _assert_reduce_overhead_compile(vae._stitch_tiles)
+    _assert_reduce_overhead_compile(vae._project_decoder_tile)
+
+    # Idempotent: the pipeline hook may run more than once per instance.
+    compiled_stitch = vae._stitch_tiles
+    compiled_project = vae._project_decoder_tile
+    vae.prepare_for_compile()
+    assert vae._stitch_tiles is compiled_stitch
+    assert vae._project_decoder_tile is compiled_project
+
+    # The opt-in mutates only this instance; the class (and therefore every
+    # default instance) stays eager.
+    assert AutoencoderKLMiniMaxH3._tile_helpers_compiled is False
+    assert not hasattr(AutoencoderKLMiniMaxH3._stitch_tiles, "get_compiler_config")
+    assert not hasattr(AutoencoderKLMiniMaxH3._project_decoder_tile, "get_compiler_config")
+
+
+def test_tile_drivers_return_caller_owned_tensors_when_compiled() -> None:
+    """Clone the stitched canvas out of cudagraph-pooled storage under the opt-in.
+
+    With ``mode="reduce-overhead"`` the ``_stitch_tiles`` output is a
+    CUDA-graph static buffer that the next replay overwrites, while the
+    collect-then-``torch.cat`` consumers (``_decode``/``_encode``/
+    ``_encode_pixels``/``encode_keyframe``) hold each chunk or clip result
+    across replays. The tile drivers therefore hand back a caller-owned copy
+    when the helpers are compiled — and return the stitched tensor by identity
+    when eager, keeping default behavior and peak memory unchanged.
+    """
+
+    def _mock_tiled_vae(stitched: torch.Tensor) -> nn.Module:
+        vae = _empty_typed_module(AutoencoderKLMiniMaxH3)
+        vae.use_tiling = True
+        vae.spatial_compression_ratio = 1
+        vae.tile_sample_min_height = 1
+        vae.tile_sample_min_width = 1
+        vae.tile_sample_min_overlap_height = 0
+        vae.tile_sample_min_overlap_width = 0
+        vae._split_tiles = Mock(return_value=([0, 1], [1, 1], [0]))
+        vae.post_quant_conv = nn.Identity()
+        vae._project_decoder_tile = vae.post_quant_conv
+        vae.quant_conv = nn.Identity()
+        vae.encoder = nn.Identity()
+        vae.decoder = nn.Identity()
+        vae._stitch_tiles = Mock(return_value=stitched)
+        return vae
+
+    latent_clip = torch.zeros((1, 1, 1, 2, 2))
+    for driver in ("_decode_clip", "_encode_clip"):
+        stitched = torch.zeros((1, 1, 1, 2, 2))
+
+        # Default instances return the stitched tensor without copying.
+        eager_vae = _mock_tiled_vae(stitched)
+        assert getattr(eager_vae, driver)(latent_clip) is stitched
+
+        compiled_vae = _mock_tiled_vae(stitched)
+        compiled_vae._tile_helpers_compiled = True
+        owned = getattr(compiled_vae, driver)(latent_clip)
+        assert owned is not stitched
+        assert torch.equal(owned, stitched)
 
 
 def test_compile_with_conditions_selects_minimax_h3_audio_decoder() -> None:
@@ -302,11 +381,10 @@ def test_decode_clip_emits_tiled_stage_ranges() -> None:
     with patch("fastvideo.models.vaes.minimax_h3_video.nvtx_range", record_range):
         decoded_clip = vae._decode_clip(torch.zeros((1, 1, 1, 2, 2)))
 
-    # The tile driver must hand back a caller-owned copy: under
-    # mode="reduce-overhead" the stitched canvas is CUDA-graph pooled storage
-    # that the next replay overwrites, so returning it by identity is a bug.
-    assert decoded_clip is not stitched_clip
-    assert torch.equal(decoded_clip, stitched_clip)
+    # Default (eager) instances return the stitched canvas by identity; the
+    # caller-owned copy under the compile opt-in is pinned by
+    # test_tile_drivers_return_caller_owned_tensors_when_compiled.
+    assert decoded_clip is stitched_clip
     assert vae._split_tiles.call_count == 2
     assert vae._project_decoder_tile.call_count == 4
     assert vae._stitch_tiles.call_count == 1
@@ -368,32 +446,29 @@ def _tiny_real_vae() -> AutoencoderKLMiniMaxH3:
         )).eval()
 
 
-def _dynamo_original(compiled_function: Any) -> Any:
-    """Return the eager callable behind a ``torch.compile``-decorated function."""
-    original = getattr(compiled_function, "_torchdynamo_orig_callable", None)
-    if original is None:
-        original = getattr(compiled_function, "__wrapped__", None)
-    assert original is not None, "cannot recover the eager tile helpers"
-    return original
-
-
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="reduce-overhead tile compile requires CUDA graphs")
 @torch.inference_mode()
 def test_tiled_decode_and_encode_survive_cudagraph_buffer_reuse_on_cuda() -> None:
     """Real tiled decode()/encode() with an unmocked reduce-overhead ``_stitch_tiles``.
 
-    Regression gate for the CUDA-graph output-clobbering bug: the stitched
-    canvas is a cudagraph static buffer, and the collect-then-``torch.cat``
-    consumers (``_decode``/``_encode``/``_encode_pixels``) hold chunk/clip
-    results across subsequent ``_stitch_tiles`` replays. Without the eager
-    ``.clone()`` at the tile-driver returns, the first tiled ``decode()`` with
-    >=2 temporal chunks raises ``accessing tensor output of CUDAGraphs that
-    has been overwritten by a subsequent run``. This test needs >=2 chunks
-    (decode), >=2 clips (encode), and a >=2x2 spatial tile grid.
+    Regression gate for the CUDA-graph output-clobbering bug under the
+    ``enable_torch_compile_vae`` opt-in (``prepare_for_compile``): the
+    stitched canvas is a cudagraph static buffer, and the
+    collect-then-``torch.cat`` consumers (``_decode``/``_encode``/
+    ``_encode_pixels``) hold chunk/clip results across subsequent
+    ``_stitch_tiles`` replays. Without the eager ``.clone()`` at the
+    tile-driver returns, the first tiled ``decode()`` with >=2 temporal
+    chunks raises ``accessing tensor output of CUDAGraphs that has been
+    overwritten by a subsequent run``. This test needs >=2 chunks (decode),
+    >=2 clips (encode), and a >=2x2 spatial tile grid.
     """
     torch.manual_seed(20260821)
     vae = _tiny_real_vae().to("cuda")
     vae.enable_tiling(16, 16, 4, 4)
+    # The hook ComposedPipelineBase._maybe_compile_pipeline_module runs for
+    # the opt-in; it installs the reduce-overhead compiled tile helpers.
+    vae.prepare_for_compile()
+    assert vae._tile_helpers_compiled
 
     # 8 latent tokens = 2 temporal chunks (tokens_chunk_size 5); 8x8 latents =
     # 32x32 pixels = a 2x2 grid of 16px tiles.
@@ -416,13 +491,12 @@ def test_tiled_decode_and_encode_survive_cudagraph_buffer_reuse_on_cuda() -> Non
     streamed_second = vae.encode_pixels(uint8_pixels).latent_dist.parameters
     assert torch.equal(streamed_first, streamed_second)
 
-    # Output parity vs the fully eager tile helpers (same weights, same math;
-    # the tolerance absorbs inductor fusion reassociation only).
+    # Output parity vs a default (fully eager, never-prepared) instance with
+    # the same weights and tiling — the tolerance absorbs inductor fusion
+    # reassociation only.
     eager_vae = _tiny_real_vae().to("cuda")
     eager_vae.load_state_dict(vae.state_dict())
     eager_vae.enable_tiling(16, 16, 4, 4)
-    eager_vae._stitch_tiles = MethodType(_dynamo_original(AutoencoderKLMiniMaxH3._stitch_tiles), eager_vae)
-    eager_vae._project_decoder_tile = MethodType(_dynamo_original(AutoencoderKLMiniMaxH3._project_decoder_tile),
-                                                 eager_vae)
+    assert not eager_vae._tile_helpers_compiled
     torch.testing.assert_close(decoded_first, eager_vae.decode(z).sample, atol=2e-4, rtol=2e-4)
     torch.testing.assert_close(encoded_first, eager_vae.encode(pixels).latent_dist.parameters, atol=2e-4, rtol=2e-4)

tests/test_block_sparse_sm100a.py

@@ -49,6 +49,15 @@ def make_case(block, num_blocks=8, topk=4, heads=4, batch=1, ragged=False, seed=
     return q, k, v, idx, num, vbs
 
 
+def make_block_map(idx, num, batch, heads, num_blocks):
+    """Expand the test's compact rows to the bool map H3 produces."""
+    block_map = torch.zeros((batch, heads, num_blocks, num_blocks), dtype=torch.bool, device=idx.device)
+    rows = block_map.view(-1, num_blocks)
+    for row in range(rows.shape[0]):
+        rows[row, idx[row, :int(num[row])].long()] = True
+    return block_map
+
+
 def reference(q, k, v, idx, num, vbs, block):
     """Dense attention restricted to the selected blocks, with padded keys masked."""
     if not vsa.BHSD:
@@ -123,6 +132,66 @@ def test_lse_is_not_vacuous(block):
     assert (got_lse.float() - (ref_lse + 1.0)).abs().max().item() > 0.5
 
 
+@pytest.mark.parametrize("block", BLOCK_SIZES)
+def test_inference_no_lse_matches_eager_under_fullgraph_compile(block):
+    """The production no-LSE route is one opaque, fake-backed graph node."""
+    q, k, v, idx, num, vbs = make_case(block, ragged=True)
+
+    def inference_path(q, k, v, idx, num, vbs):
+        out, lse = vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs, need_lse=False)
+        assert lse is None
+        return out
+
+    with torch.inference_mode():
+        expected = inference_path(q, k, v, idx, num, vbs)
+        compiled = torch.compile(inference_path, backend="eager", fullgraph=True)
+        actual = compiled(q, k, v, idx, num, vbs)
+    torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+
+
+@pytest.mark.parametrize("block", BLOCK_SIZES)
+def test_inference_custom_op_fake_matches_real(block):
+    """Pin the custom op schema and fake output metadata used by Dynamo."""
+    q, k, v, idx, num, vbs = make_case(block, ragged=True)
+    torch.library.opcheck(
+        torch.ops.fastvideo_kernel.block_sparse_attn_sm100a_inference.default,
+        (q, k, v, idx, num, vbs),
+        test_utils=("test_schema", "test_faketensor", "test_aot_dispatch_dynamic"),
+    )
+
+
+@pytest.mark.parametrize("block", BLOCK_SIZES)
+def test_mask_inference_matches_index_route_under_fullgraph_compile(block):
+    """H3's bool mask compaction and sm_100a launch stay in one opaque op."""
+    batch, heads, num_blocks = 1, 4, 8
+    q, k, v, idx, num, vbs = make_case(block, batch=batch, heads=heads, num_blocks=num_blocks, ragged=True)
+    block_map = make_block_map(idx, num, batch, heads, num_blocks)
+
+    def inference_path(q, k, v, block_map, vbs):
+        out, lse = vsa.block_sparse_attn_sm100a_from_mask(q, k, v, block_map, vbs)
+        assert lse is None
+        return out
+
+    with torch.inference_mode():
+        expected, _ = vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs, need_lse=False)
+        compiled = torch.compile(inference_path, backend="eager", fullgraph=True)
+        actual = compiled(q, k, v, block_map, vbs)
+    torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+
+
+@pytest.mark.parametrize("block", BLOCK_SIZES)
+def test_mask_inference_custom_op_fake_matches_real(block):
+    """Pin schema/fake metadata for H3's mask-taking opaque boundary."""
+    batch, heads, num_blocks = 1, 4, 8
+    q, k, v, idx, num, vbs = make_case(block, batch=batch, heads=heads, num_blocks=num_blocks, ragged=True)
+    block_map = make_block_map(idx, num, batch, heads, num_blocks)
+    torch.library.opcheck(
+        torch.ops.fastvideo_kernel.block_sparse_attn_sm100a_from_mask_inference.default,
+        (q, k, v, block_map, vbs),
+        test_utils=("test_schema", "test_faketensor", "test_aot_dispatch_dynamic"),
+    )
+
+
 def test_unsupported_is_rejected():
     q, _, _, _, _, vbs = make_case(64)
     assert not vsa.is_supported(q.float(), vbs)                 # wrong dtype