[perf]: training-free token reduction for MiniMax H3 (#2004)

.buildkite/scripts/unit_test.sh

@@ -29,6 +29,7 @@ exec pytest \
   ./fastvideo/tests/training/test_trackers.py \
   ./fastvideo/tests/training/test_ltx2_rope_fps.py \
   ./fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py \
+  ./fastvideo/tests/inference/test_basic_fasth3_token_reduction.py \
   ./fastvideo/tests/inference/test_inference_regional_compile.py \
   ./fastvideo/tests/attention/ \
   ./fastvideo/tests/layers/test_pdd_linear.py \

docs/design/inference_schema_parity_inventory.yaml

@@ -83,6 +83,7 @@ surfaces:
       vae_parallel_decode: "MiniMax-H3 sequence-parallel VAE decode opt-in; model-specific optimization not yet represented in the typed public schema."
       h3_sequential_load: "MiniMax-H3 sequential text-encoder then DiT/VAE load; model-specific optimization not yet represented in the typed public schema."
       h3_vsa_fp4: "MiniMax-H3 native FP4 attention options carried through generator.pipeline.experimental.h3_vsa_fp4 rather than a dedicated typed component setting."
+      h3_token_reduction: "MiniMax-H3 training-free token reduction options carried through generator.pipeline.experimental.h3_token_reduction rather than a dedicated typed component setting."
       h3_encoder_split: "MiniMax-H3 Ray-only split that runs the Qwen3-VL text encoder on dedicated nodes and the DiT/VAE on the rest; model-specific optimization not yet represented in the typed public schema."
       h3_encoder_nodes: "Leading node count that h3_encoder_split reserves for the MiniMax-H3 text encoder; model-specific optimization not yet represented in the typed public schema."
       video_decode_backend: "MiniMax-H3 video decoder selection (full VAE vs TAEH3 preview); model-specific optimization not yet represented in the typed public schema."

docs/inference/fasth3-distilled.md

@@ -198,6 +198,57 @@ forwards. With the DiT sharded, peak allocated memory is about 66 GiB per GPU
 (about 95 GiB in use), against 83 GiB (103 GiB) with a full DiT copy on each
 GPU; the output is identical. Model loading and decoding come on top of this.
 
+## Token reduction
+
+`--token-reduction` is an opt-in, training-free speed mode that follows the
+`--token-reduction` mode of [antirez/h3.c](https://github.com/antirez/h3.c).
+Inside DiT blocks 4 to 29, every two horizontally adjacent tokens of the
+generated video are averaged into one, so those blocks see the video at half
+width. During the first 4 of the 8 forwards, where the coarse layout forms,
+the interval extends to block 39. Before the first full-width block, each
+token gets back its own value plus the update its pair received. Text, audio,
+reference and condition tokens are never paired.
+
+FastH3 8-Step V2 and the OmniRef PDD export both support it:
+
+```bash
+python examples/inference/basic/basic_fasth3_8step.py \
+  --prompt 'Two people walk along a beach at sunset, talking and laughing' \
+  --num-gpus 4 --token-reduction
+
+python examples/inference/basic/basic_fasth3_omniref_pdd.py \
+  --model-path <OmniRef export> --image reference.png \
+  --prompt '<Picture 1> ...' --num-gpus 4 --token-reduction
+```
+
+In a recipe, set `pipeline.experimental.h3_token_reduction` to
+`{enabled: true}`. The optional keys `blocks`, `early` and `update_scale`
+match the `--token-reduction-blocks`, `--token-reduction-early` and
+`--token-reduction-scale` flags. Malformed values, and intervals past the
+checkpoint's DiT depth, fail before any weights load.
+
+Measured on one GB200 tray with 4 GPUs, 1344x768, 124 frames and BF16
+weights, after warm-up:
+
+| Checkpoint | Denoising, 8 forwards | Whole request |
+|---|---|---|
+| FastH3 8-Step V2, default profile with regional compile | 5.08 s to 3.53 s, 30.6% less | 9.13 s to 7.58 s, 17.0% less |
+| FastH3 8-Step V2, regional compile off | 6.10 s to 4.26 s, 30.2% less | 10.18 s to 8.41 s, 17.4% less |
+| OmniRef PDD with one image reference, eager | 12.07 s to 9.33 s, 22.7% less | 24.67 s to 21.62 s, 12.3% less |
+
+The saving follows how many rows the middle blocks drop. A V2 request has a
+short text and audio prefix, so the reduced blocks see about half the rows.
+An OmniRef request also carries the reference tokens, which stay unpaired, so
+it saves less. Decoding takes the same time either way.
+
+The reduced video keeps the quality of the unreduced one, but it is not the
+same video. Most scenes keep their overall layout, while details can change,
+such as one of two people in a scene or the venue of a concert. The
+checkpoints were not trained with token reduction, and fine-tuning with it
+enabled is the expected way to keep the original shot. Only the GB200 BF16
+path has been validated so far; the NVFP4 and FP8 consumer profiles have not
+been tested with it.
+
 ## Apple Silicon
 
 `mlx_fasth3.py` stays on FastH3 V1 and its uniform AdaLN cache.

examples/inference/basic/basic_fasth3.py

@@ -37,6 +37,8 @@ from fastvideo.api import (
     PipelineSelection,
     SamplingConfig,
 )
+from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3ArchConfig
+from fastvideo.configs.pipelines.minimax_h3_token_reduction import MiniMaxH3TokenReductionSchedule
 
 DEFAULT_MODEL = "FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2"
 
@@ -163,12 +165,55 @@ def build_parser(description: str | None = None) -> argparse.ArgumentParser:
                         default=None,
                         help='whole-DiT torch.compile mode, e.g. "reduce-overhead"; requires '
                         "--no-inference-torch-compile")
+    parser.add_argument("--token-reduction",
+                        action="store_true",
+                        help="training-free speed mode after antirez/h3.c: pair horizontally adjacent generated-video "
+                        "tokens inside the middle DiT blocks; faster, but can change composition")
+    parser.add_argument("--token-reduction-blocks",
+                        default=None,
+                        metavar="BEGIN:END",
+                        help="blocks that run reduced (default 4:30)")
+    parser.add_argument("--token-reduction-early",
+                        default=None,
+                        metavar="STEPS:END",
+                        help="for the first STEPS forwards, reduce up to block END instead (default 4:40; 0:0 "
+                        "disables)")
+    parser.add_argument("--token-reduction-scale",
+                        type=float,
+                        default=None,
+                        help="scale of the update each paired token gets from its pair (default 1.0); unpaired "
+                        "tokens, such as the last column of an odd-width video, take the full update as in h3.c")
     return parser
 
 
+def token_reduction_options(args: argparse.Namespace) -> dict[str, object] | None:
+    """``h3_token_reduction`` options from the command line, or None when it is off."""
+    if not getattr(args, "token_reduction", False):
+        return None
+    options: dict[str, object] = {"enabled": True}
+    if args.token_reduction_blocks is not None:
+        options["blocks"] = args.token_reduction_blocks
+    if args.token_reduction_early is not None:
+        options["early"] = args.token_reduction_early
+    if args.token_reduction_scale is not None:
+        options["update_scale"] = args.token_reduction_scale
+    return options
+
+
 def validate_args(parser: argparse.ArgumentParser, args: argparse.Namespace) -> argparse.Namespace:
     if args.repeats < 1:
         parser.error("--repeats must be at least 1")
+    tuning = (args.token_reduction_blocks, args.token_reduction_early, args.token_reduction_scale)
+    if not args.token_reduction and any(value is not None for value in tuning):
+        parser.error("--token-reduction-blocks/-early/-scale need --token-reduction")
+    if args.token_reduction:
+        # Reject malformed values and intervals past H3's 50 blocks before any download; the
+        # transformer loader checks a shallower checkpoint, such as FastH3 Trim, against its own depth.
+        try:
+            MiniMaxH3TokenReductionSchedule.from_options(token_reduction_options(args) or {},
+                                                         num_blocks=MiniMaxH3ArchConfig().num_layers)
+        except ValueError as error:
+            parser.error(str(error))
     if args.num_gpus < 1:
         parser.error("--num-gpus must be at least 1")
     if not 0.0 <= args.vsa_sparsity < 1.0:
@@ -303,6 +348,9 @@ def build_generator_config(args: argparse.Namespace) -> GeneratorConfig:
             "VSA_sparsity": args.vsa_sparsity,
             "VSA_tile_size": args.vsa_tile_size,
         })
+    token_reduction = token_reduction_options(args)
+    if token_reduction is not None:
+        experimental["h3_token_reduction"] = token_reduction
     return GeneratorConfig(
         model_path=args.model_path,
         pipeline=PipelineSelection(

examples/inference/basic/basic_fasth3_omniref_pdd.py

@@ -61,7 +61,9 @@ from fastvideo.api import (
     QuantizationConfig,
     SamplingConfig,
 )
+from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3ArchConfig
 from fastvideo.configs.pipelines.minimax_h3 import parse_base_model_revision
+from fastvideo.configs.pipelines.minimax_h3_token_reduction import MiniMaxH3TokenReductionSchedule
 from fastvideo.pipelines.basic.minimax_h3 import MiniMaxH3Reference
 
 CONTRACT = "fastvideo_inference.json"
@@ -154,9 +156,41 @@ def build_parser() -> argparse.ArgumentParser:
                         help="precomputed AdaLN modulation table, so the 26 GB of adaln_proj never loads "
                         "(write one once with FASTVIDEO_H3_ADALN_DUMP). Strongly recommended: the 24/32 GB "
                         "profiles need it to keep the DiT resident")
+    parser.add_argument("--token-reduction",
+                        action="store_true",
+                        help="training-free speed mode after antirez/h3.c: pair horizontally adjacent generated-video "
+                        "tokens inside the middle DiT blocks; faster, but can change composition")
+    parser.add_argument("--token-reduction-blocks",
+                        default=None,
+                        metavar="BEGIN:END",
+                        help="blocks that run reduced (default 4:30)")
+    parser.add_argument("--token-reduction-early",
+                        default=None,
+                        metavar="STEPS:END",
+                        help="for the first STEPS forwards, reduce up to block END instead (default 4:40; 0:0 "
+                        "disables)")
+    parser.add_argument("--token-reduction-scale",
+                        type=float,
+                        default=None,
+                        help="scale of the update each paired token gets from its pair (default 1.0); unpaired "
+                        "tokens, such as the last column of an odd-width video, take the full update as in h3.c")
     return parser
 
 
+def token_reduction_options(args: argparse.Namespace) -> dict[str, Any] | None:
+    """``h3_token_reduction`` options from the command line, or None when it is off."""
+    if not args.token_reduction:
+        return None
+    options: dict[str, Any] = {"enabled": True}
+    if args.token_reduction_blocks is not None:
+        options["blocks"] = args.token_reduction_blocks
+    if args.token_reduction_early is not None:
+        options["early"] = args.token_reduction_early
+    if args.token_reduction_scale is not None:
+        options["update_scale"] = args.token_reduction_scale
+    return options
+
+
 def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
     parser = build_parser()
     args = parser.parse_args(argv)
@@ -164,6 +198,17 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
         parser.error("pass at least one ordered reference with --image, --video, or --audio")
     if not any(kind != "audio" for kind, _ in args.references):
         parser.error("Ref2VA needs at least one image or video reference")
+    tuning = (args.token_reduction_blocks, args.token_reduction_early, args.token_reduction_scale)
+    if not args.token_reduction and any(value is not None for value in tuning):
+        parser.error("--token-reduction-blocks/-early/-scale need --token-reduction")
+    if args.token_reduction:
+        # Reject malformed values and intervals past H3's 50 blocks before any download; the
+        # transformer loader checks the checkpoint's own depth.
+        try:
+            MiniMaxH3TokenReductionSchedule.from_options(token_reduction_options(args),
+                                                         num_blocks=MiniMaxH3ArchConfig().num_layers)
+        except ValueError as error:
+            parser.error(str(error))
     return args
 
 
@@ -421,7 +466,8 @@ def apply_profile(profile: DeviceProfile, decoder: str | None, adaln_table: str
 def build_generator_config(model_dir: Path,
                            num_gpus: int,
                            lossless_accel: bool = False,
-                           profile: DeviceProfile | None = None) -> GeneratorConfig:
+                           profile: DeviceProfile | None = None,
+                           token_reduction: dict[str, Any] | None = None) -> GeneratorConfig:
     quantization = None
     if profile is not None:
         # `h3_dit_vsa` matches the packed export the consumer weights ship; the
@@ -444,6 +490,8 @@ def build_generator_config(model_dir: Path,
     else:
         offload = OffloadConfig(dit=False, dit_layerwise=False, text_encoder=True, vae=True, pin_cpu_memory=False)
         experimental = {}
+    if token_reduction is not None:
+        experimental["h3_token_reduction"] = dict(token_reduction)
     return GeneratorConfig(
         model_path=str(model_dir),
         engine=EngineConfig(
@@ -496,7 +544,7 @@ def main() -> None:
     if args.height is None:
         args.height = profile.height if profile is not None else DEFAULT_HEIGHT
     generator = VideoGenerator.from_config(
-        build_generator_config(model_dir, args.num_gpus, args.lossless_accel, profile))
+        build_generator_config(model_dir, args.num_gpus, args.lossless_accel, profile, token_reduction_options(args)))
     try:
         result = generator.generate(
             GenerationRequest(

fastvideo/configs/pipelines/minimax_h3_token_reduction.py

@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Options for MiniMax-H3 training-free token reduction, parsed without importing model code.
+
+``FastVideoArgs`` validates ``h3_token_reduction`` here, so building args never
+imports a model module. The row maps and the pooling arithmetic live in
+``fastvideo/models/dits/minimax_h3_token_reduction.py``.
+"""
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Any
+
+TOKEN_REDUCTION_OPTION_KEYS = frozenset({"enabled", "blocks", "early", "update_scale"})
+
+
+@dataclass(frozen=True)
+class MiniMaxH3TokenReductionSchedule:
+    """Which DiT blocks run on paired tokens at each denoising step.
+
+    Blocks ``begin <= block < end`` run reduced. During the first
+    ``early_steps`` steps the interval extends to ``early_end``.
+
+    h3.c reduces blocks 4..29, and 4..39 during its first 10 of about 19
+    forwards, where coarse structure forms. The default keeps that share of an
+    eight-forward PDD schedule: blocks 4..39 for the first four forwards.
+    """
+
+    begin: int = 4
+    end: int = 30
+    early_steps: int = 4
+    early_end: int = 40
+    update_scale: float = 1.0
+
+    @classmethod
+    def from_options(cls,
+                     options: Mapping[str, Any],
+                     num_blocks: int | None = None) -> MiniMaxH3TokenReductionSchedule | None:
+        """Parse ``h3_token_reduction`` options; None unless ``enabled`` is true.
+
+        ``num_blocks``, when known, bounds both block intervals by the DiT depth.
+        """
+        if not isinstance(options, Mapping) or options.keys() - TOKEN_REDUCTION_OPTION_KEYS:
+            raise ValueError(
+                f"h3_token_reduction must be a mapping with keys from {sorted(TOKEN_REDUCTION_OPTION_KEYS)}")
+        enabled = options.get("enabled", False)
+        if not isinstance(enabled, bool):
+            raise ValueError("h3_token_reduction.enabled must be a boolean")
+        if not enabled:
+            return None
+        begin, end = _int_pair(options.get("blocks", (cls.begin, cls.end)), "blocks")
+        if not 0 <= begin < end:
+            raise ValueError(f"h3_token_reduction.blocks must satisfy 0 <= BEGIN < END, got {begin}:{end}")
+        # Same rule as h3.c: the deeper early interval applies only when the
+        # base interval restores before block 40.
+        default_early = (cls.early_steps, cls.early_end) if end < cls.early_end else (0, end)
+        early_steps, early_end = _int_pair(options.get("early", default_early), "early")
+        if early_steps < 0:
+            raise ValueError("h3_token_reduction.early STEPS must be non-negative")
+        if early_steps == 0:
+            early_end = end
+        elif early_end <= end:
+            raise ValueError(f"h3_token_reduction.early END must exceed blocks END {end}, got {early_end}")
+        if num_blocks is not None and max(end, early_end) > num_blocks:
+            raise ValueError(f"h3_token_reduction block intervals end at {max(end, early_end)}, past the "
+                             f"{num_blocks} DiT blocks")
+        scale = options.get("update_scale", cls.update_scale)
+        if isinstance(scale, bool) or not isinstance(scale, int | float) or not 0.0 <= float(scale) <= 2.0:
+            raise ValueError("h3_token_reduction.update_scale must be a number in [0, 2]")
+        return cls(begin=begin, end=end, early_steps=early_steps, early_end=early_end, update_scale=float(scale))
+
+    def end_for_step(self, step: int) -> int:
+        return self.early_end if step < self.early_steps else self.end
+
+
+def _int_pair(value: Any, name: str) -> tuple[int, int]:
+    if isinstance(value, str):
+        value = value.split(":")
+    if (not isinstance(value, list | tuple) or len(value) != 2
+            or any(isinstance(item, bool) or not isinstance(item, int | str) for item in value)):
+        raise ValueError(f"h3_token_reduction.{name} must be two integers, e.g. [4, 30] or '4:30'")
+    try:
+        return int(value[0]), int(value[1])
+    except ValueError as error:
+        raise ValueError(f"h3_token_reduction.{name} must be two integers, got {value!r}") from error
+
+
+__all__ = ["MiniMaxH3TokenReductionSchedule", "TOKEN_REDUCTION_OPTION_KEYS"]

fastvideo/fastvideo_args.py

@@ -12,6 +12,7 @@ from typing import Any, TYPE_CHECKING
 
 from fastvideo.configs.configs import PreprocessConfig
 from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.configs.pipelines.minimax_h3_token_reduction import MiniMaxH3TokenReductionSchedule
 from fastvideo.configs.utils import clean_cli_args
 from fastvideo.layers.quantization import QUANTIZATION_METHODS, QuantizationMethods
 from fastvideo.logger import init_logger
@@ -296,6 +297,11 @@ class FastVideoArgs:
     inference_torch_compile: bool = False
 
     h3_vsa_fp4: dict[str, Any] = field(default_factory=dict)
+    # Training-free horizontal pairing of MiniMax-H3's generated-video tokens in
+    # the middle DiT blocks (antirez/h3.c --token-reduction). Keys: enabled,
+    # blocks [BEGIN, END], early [STEPS, END], update_scale. See
+    # fastvideo/configs/pipelines/minimax_h3_token_reduction.py.
+    h3_token_reduction: dict[str, Any] = field(default_factory=dict)
 
     disable_autocast: bool = False
 
@@ -453,6 +459,8 @@ class FastVideoArgs:
             raise ValueError("h3_vsa_fp4.prefix_precision must be empty, fp8 or bf16")
         if self.h3_vsa_fp4.get("enabled") and (self.inference_torch_compile or self.enable_torch_compile):
             raise ValueError("h3_vsa_fp4 requires eager DiT inference; disable DiT torch.compile")
+        # The DiT depth is checked against the checkpoint's own config when the transformer loads.
+        MiniMaxH3TokenReductionSchedule.from_options(self.h3_token_reduction)
         self._fold_vae_parallel_env()
         self._fold_h3_encoder_split_env()
         # Runs after FASTVIDEO_ATTENTION_BACKEND is copied into attention_backend,

fastvideo/models/dits/minimax_h3.py

@@ -38,6 +38,8 @@ from fastvideo.logger import init_logger
 from fastvideo.models.dits.base import BaseDiT
 from fastvideo.models.dits.minimax_h3_vsa_fp4 import (STAGES, vsa_fp4_attention, vsa_fp4_attention_sp,
                                                        vsa_fp4_requested, vsa_tile_first_attention)
+from fastvideo.models.dits.minimax_h3_token_reduction import (MiniMaxH3TokenReduction, expand_token_pairs,
+                                                              pool_token_pairs)
 from fastvideo.models.dits.minimax_h3_fusions import exact as exact_kernels
 from fastvideo.models.dits.minimax_h3_fusions import mk as mk_kernels
 from fastvideo.models.dits.minimax_h3_fusions import (
@@ -967,6 +969,7 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
         # per-generation caches for loop-invariant work (see _rotary_for /
         # _refined_text); plain attrs, never in state_dict
         self._rope_cache: tuple | None = None
+        self._reduced_rope_cache: tuple | None = None
         self._text_cache: tuple | None = None
         self.token_refiner = MiniMaxH3TokenRefiner(
             arch.hidden_size,
@@ -1229,24 +1232,61 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
                                  (2 * arch.rope_freq_dim)))
             self.rope._buffers["inv_freq"] = inv_freq
 
-    def _rotary_for(self, position_ids: torch.Tensor, dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]:
+    def _rotary_for(self,
+                    position_ids: torch.Tensor,
+                    dtype: torch.dtype,
+                    cache_name: str = "_rope_cache") -> tuple[torch.Tensor, torch.Tensor]:
         """cos/sin depend only on positions; the denoising stage holds one
         position_ids tensor for the whole loop, so cache per layout and
         pre-cast once instead of rebuilding every step and re-casting in
         every block. Holding the source tensor keeps the identity check safe
-        against id reuse."""
+        against id reuse. Token reduction keeps its pair positions in a
+        second slot so the two layouts do not evict each other."""
         if torch.compiler.is_compiling():
             cos, sin = self.rope(position_ids)
             return cos.to(dtype), sin.to(dtype)
-        cached = self._rope_cache
+        cached = getattr(self, cache_name)
         if (cached is not None and cached[0] is position_ids and cached[1] == dtype
                 and cached[2][0].device == position_ids.device):
             return cached[2]
         cos, sin = self.rope(position_ids)
         value = (cos.to(dtype), sin.to(dtype))
-        self._rope_cache = (position_ids, dtype, value)
+        setattr(self, cache_name, (position_ids, dtype, value))
         return value
 
+    def _enter_token_reduction(self, packed_hidden_states: torch.Tensor, plan: MiniMaxH3TokenReduction,
+                               sequence_length: int,
+                               sp_world_size: int) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+        """Pool video pairs; returns the reduced block input and the full rows plus pooled baseline to restore from.
+
+        Under sequence parallelism pooling changes how many rows each rank
+        holds, so the full sequence is gathered, pooled and sharded again.
+        """
+        with nvtx_range("minimax_h3.token_reduction.pool"):
+            if sp_world_size > 1:
+                packed_hidden_states = sequence_model_parallel_all_gather_with_unpad(packed_hidden_states,
+                                                                                      sequence_length,
+                                                                                      dim=1)
+            pooled = pool_token_pairs(packed_hidden_states, plan)
+            reduced = pooled
+            if sp_world_size > 1:
+                reduced = self._block_input(sequence_model_parallel_shard(pooled, dim=1)[0])
+        return reduced, (packed_hidden_states, pooled)
+
+    def _leave_token_reduction(self, packed_hidden_states: torch.Tensor, saved: tuple[torch.Tensor, torch.Tensor],
+                               plan: MiniMaxH3TokenReduction, sp_world_size: int) -> torch.Tensor:
+        """Restore the full sequence from the reduced stack's output."""
+        original, pooled = saved
+        with nvtx_range("minimax_h3.token_reduction.expand"):
+            if sp_world_size > 1:
+                packed_hidden_states = sequence_model_parallel_all_gather_with_unpad(packed_hidden_states,
+                                                                                      plan.reduced_length,
+                                                                                      dim=1)
+            restored = expand_token_pairs(original, packed_hidden_states, pooled, plan)
+            if sp_world_size > 1:
+                restored = self._block_input(sequence_model_parallel_shard(restored, dim=1)[0])
+        return restored
+
     def attach_step_splice(self, late: nn.Module, from_step: int) -> None:
         """Hand denoising steps ``from_step`` onward to ``late`` (same architecture, other weights).
 
@@ -1339,13 +1379,18 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
         audio_indices: torch.Tensor,
         text_indices: torch.Tensor,
         text_attn_metadata: AttentionMetadata | None = None,
+        token_reduction: MiniMaxH3TokenReduction | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Predict video and audio velocities from one caller-defined packed layout."""
+        """Predict video and audio velocities from one caller-defined packed layout.
+
+        ``token_reduction`` (inference only) runs blocks ``[begin, end)`` on
+        horizontally paired generated-video rows; see minimax_h3_token_reduction.
+        """
         late = self.__dict__.get("_splice_late")
         if late is not None and get_forward_context().current_timestep >= self._splice_from_step:
             return late(hidden_states, audio_hidden_states, encoder_hidden_states, timestep, timestep_indices,
                         token_tags, position_ids, video_indices, audio_indices, text_indices,
-                        text_attn_metadata=text_attn_metadata)
+                        text_attn_metadata=text_attn_metadata, token_reduction=token_reduction)
         if position_ids.ndim != 2 or position_ids.shape[-1] != 3:
             raise ValueError(f"position_ids must have shape (seq_len, 3), got {tuple(position_ids.shape)}.")
         sequence_length = position_ids.shape[0]
@@ -1357,6 +1402,15 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
             raise ValueError("audio_hidden_states row count must match audio_indices.")
         if encoder_hidden_states.shape[1] != text_indices.numel():
             raise ValueError("encoder_hidden_states row count must match text_indices.")
+        if token_reduction is not None:
+            if torch.is_grad_enabled():
+                raise ValueError("MiniMax H3 token reduction is inference-only.")
+            if token_reduction.sequence_length != sequence_length:
+                raise ValueError(f"Token reduction was built for {token_reduction.sequence_length} rows, "
+                                 f"the packed sequence has {sequence_length}.")
+            if not 0 <= token_reduction.begin < token_reduction.end <= len(self.transformer_blocks):
+                raise ValueError(f"Token reduction blocks {token_reduction.begin}:{token_reduction.end} must lie in "
+                                 f"0:{len(self.transformer_blocks)}.")
 
         video_embeds, _ = self.proj_in(hidden_states.to(self.proj_in.weight.dtype))
         audio_embeds, _ = self.audio_proj_in(audio_hidden_states.to(self.audio_proj_in.weight.dtype))
@@ -1371,6 +1425,8 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
         pad_tokens = sequence_length - real_tokens
         if pad_tokens < 0:
             raise ValueError("Real H3 token counts exceed the packed sequence length")
+        if pad_tokens and token_reduction is not None:
+            raise ValueError("MiniMax H3 token reduction does not support capacity padding rows.")
         if pad_tokens:
             metadata = get_forward_context().attn_metadata
             packed_sequence = getattr(metadata, "packed_sequence", None)
@@ -1391,6 +1447,12 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
         adaln_indices = timestep_indices * MINIMAX_H3_MODALITY_NUM + token_tags
         local_timestep_indices = timestep_indices
         original_seq_len = sequence_length
+        reduced_rotary_emb = reduced_adaln_indices = None
+        if token_reduction is not None:
+            # A pair takes its first source's modulation row and the mean of both positions.
+            reduced_rotary_emb = self._rotary_for(token_reduction.position_ids, text_embeds.dtype,
+                                                  "_reduced_rope_cache")
+            reduced_adaln_indices = adaln_indices.index_select(0, token_reduction.first)
 
         if sp_world_size > 1:
             packed_hidden_states = self._block_input(sequence_model_parallel_shard(packed_hidden_states, dim=1)[0])
@@ -1399,6 +1461,10 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
             adaln_indices, _ = sequence_model_parallel_shard(adaln_indices, dim=0)
             local_timestep_indices, _ = sequence_model_parallel_shard(local_timestep_indices, dim=0)
             rotary_emb = (rotary_cos, rotary_sin)
+            if reduced_rotary_emb is not None:
+                reduced_rotary_emb = (sequence_model_parallel_shard(reduced_rotary_emb[0], dim=0)[0],
+                                      sequence_model_parallel_shard(reduced_rotary_emb[1], dim=0)[0])
+                reduced_adaln_indices, _ = sequence_model_parallel_shard(reduced_adaln_indices, dim=0)
 
         if getattr(self, "_adaln_host_cache", False):
             # One host read per forward keys every block's modulation cache by the timestep values.
@@ -1407,20 +1473,53 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
                 block.adaln_proj._cache_key = key
             self._move_adaln_tables(temb.device)
 
-        # The eager driver owns profiling markers while each block's compiled
-        # forward owns the graph that the marker surrounds.
-        for block_index, block in enumerate(self.transformer_blocks):
-            if STAGES.enabled:
-                logger.info("H3_MEMORY_BLOCK %d allocated=%.3f GiB reserved=%.3f GiB", block_index,
-                            torch.cuda.memory_allocated() / 2**30, torch.cuda.memory_reserved() / 2**30)
-            with nvtx_range(f"minimax_h3.transformer_block.{block_index}"), STAGES.span("block_total"):
-                packed_hidden_states = block(
-                    packed_hidden_states,
-                    temb,
-                    adaln_indices,
-                    rotary_emb,
-                    original_seq_len,
-                )
+        # Reduced blocks attend over the shorter sequence, so they read its own
+        # attention metadata from the forward context; the full one is put back
+        # on exit and on error.
+        forward_context = get_forward_context() if token_reduction is not None else None
+        full_attn_metadata = forward_context.attn_metadata if forward_context is not None else None
+        reduction_saved = None
+        try:
+            # The eager driver owns profiling markers while each block's compiled
+            # forward owns the graph that the marker surrounds.
+            for block_index, block in enumerate(self.transformer_blocks):
+                if token_reduction is not None and block_index == token_reduction.begin:
+                    assert forward_context is not None
+                    packed_hidden_states, reduction_saved = self._enter_token_reduction(
+                        packed_hidden_states, token_reduction, original_seq_len, sp_world_size)
+                    forward_context.attn_metadata = token_reduction.attn_metadata
+                if token_reduction is not None and block_index == token_reduction.end:
+                    assert forward_context is not None and reduction_saved is not None
+                    packed_hidden_states = self._leave_token_reduction(packed_hidden_states, reduction_saved,
+                                                                       token_reduction, sp_world_size)
+                    forward_context.attn_metadata = full_attn_metadata
+                    reduction_saved = None
+                if reduction_saved is None:
+                    block_adaln_indices, block_rotary_emb, block_seq_len = adaln_indices, rotary_emb, original_seq_len
+                else:
+                    assert token_reduction is not None and reduced_adaln_indices is not None
+                    assert reduced_rotary_emb is not None
+                    block_adaln_indices, block_rotary_emb = reduced_adaln_indices, reduced_rotary_emb
+                    block_seq_len = token_reduction.reduced_length
+                if STAGES.enabled:
+                    logger.info("H3_MEMORY_BLOCK %d allocated=%.3f GiB reserved=%.3f GiB", block_index,
+                                torch.cuda.memory_allocated() / 2**30, torch.cuda.memory_reserved() / 2**30)
+                with nvtx_range(f"minimax_h3.transformer_block.{block_index}"), STAGES.span("block_total"):
+                    packed_hidden_states = block(
+                        packed_hidden_states,
+                        temb,
+                        block_adaln_indices,
+                        block_rotary_emb,
+                        block_seq_len,
+                    )
+            if reduction_saved is not None:
+                assert token_reduction is not None
+                packed_hidden_states = self._leave_token_reduction(packed_hidden_states, reduction_saved,
+                                                                   token_reduction, sp_world_size)
+                reduction_saved = None
+        finally:
+            if forward_context is not None:
+                forward_context.attn_metadata = full_attn_metadata
 
         packed_hidden_states = self.norm_out(
             packed_hidden_states,

fastvideo/models/dits/minimax_h3_token_reduction.py

@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Training-free horizontal token pairing for MiniMax H3, after antirez/h3.c ``--token-reduction``.
+
+Inside an interval of transformer blocks, every pair of horizontally adjacent
+generated-video tokens is averaged into one row, so the middle blocks see the
+target video at half width. Text, audio, references and conditioning rows are
+never paired and pass through exactly. Before the first full-width block each
+original row is restored as its own saved value plus the update its pair
+received, so within-pair detail survives::
+
+    row 1, 3, 5, 9 -> pairs 2, 7 -> middle blocks 2.5, 6 -> restored 1.5, 3.5, 4, 8
+
+Pair positions are the mean of the two source positions; a pair takes the
+modulation row of its first source (every target-video row shares one).
+Unpaired rows, such as the last column of an odd-width video, take the
+reduced output directly, as in h3.c, so ``update_scale`` applies to pairs only.
+
+The option parsing and the default schedule live in
+``fastvideo/configs/pipelines/minimax_h3_token_reduction.py``.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+from fastvideo.attention.backends.abstract import AttentionMetadata
+from fastvideo.configs.pipelines.minimax_h3_token_reduction import MiniMaxH3TokenReductionSchedule
+
+
+@dataclass(frozen=True)
+class MiniMaxH3TokenReduction:
+    """Row maps for one packed sequence whose generated video is its trailing ``(t, h, w)`` token grid.
+
+    ``first``/``second`` name the full rows averaged into each reduced row
+    (equal for unpaired rows); ``parents`` names the reduced row every full
+    row restores from. ``pair_reduced_rows`` and ``pair_rows`` list the
+    reduced and full rows that belong to a true pair.
+    """
+
+    begin: int
+    end: int
+    first: torch.Tensor
+    second: torch.Tensor
+    parents: torch.Tensor
+    pair_reduced_rows: torch.Tensor
+    pair_rows: torch.Tensor
+    position_ids: torch.Tensor
+    update_scale: float = 1.0
+    # Attention metadata of the reduced sequence (VSA-H3 tiles of the half-width grid).
+    attn_metadata: AttentionMetadata | None = None
+
+    @property
+    def sequence_length(self) -> int:
+        return int(self.parents.numel())
+
+    @property
+    def reduced_length(self) -> int:
+        return int(self.first.numel())
+
+    def to(self, device: torch.device) -> MiniMaxH3TokenReduction:
+        return MiniMaxH3TokenReduction(
+            begin=self.begin,
+            end=self.end,
+            first=self.first.to(device),
+            second=self.second.to(device),
+            parents=self.parents.to(device),
+            pair_reduced_rows=self.pair_reduced_rows.to(device),
+            pair_rows=self.pair_rows.to(device),
+            position_ids=self.position_ids.to(device),
+            update_scale=self.update_scale,
+            attn_metadata=self.attn_metadata,
+        )
+
+
+def reduced_token_width(width: int) -> int:
+    return (width + 1) // 2
+
+
+def build_token_reduction(
+    position_ids: torch.Tensor,
+    video_start: int,
+    grid: tuple[int, int, int],
+    schedule: MiniMaxH3TokenReductionSchedule,
+) -> MiniMaxH3TokenReduction:
+    """Pair horizontally adjacent rows of the row-major ``(t, h, w)`` video grid ending the sequence."""
+    t, h, w = (int(v) for v in grid)
+    sequence_length = int(position_ids.shape[0])
+    if min(t, h, w) <= 0 or video_start < 0 or video_start + t * h * w != sequence_length:
+        raise ValueError(f"Token reduction needs the generated video grid {grid} to end the packed sequence "
+                         f"({video_start} + {t * h * w} != {sequence_length}).")
+    reduced_w = reduced_token_width(w)
+    prefix = torch.arange(video_start, dtype=torch.long)
+
+    reduced_local = torch.arange(t * h * reduced_w, dtype=torch.long)
+    pair_column = 2 * (reduced_local % reduced_w)
+    first_video = video_start + (reduced_local // reduced_w) * w + pair_column
+    has_pair = pair_column + 1 < w
+    second_video = first_video + has_pair.long()
+
+    full_local = torch.arange(t * h * w, dtype=torch.long)
+    column = full_local % w
+    parent_video = video_start + (full_local // w) * reduced_w + column // 2
+    in_pair = 2 * (column // 2) + 1 < w
+
+    first = torch.cat([prefix, first_video])
+    second = torch.cat([prefix, second_video])
+    # (x + x) / 2 == x, so unpaired rows keep their exact positions.
+    reduced_positions = (position_ids.index_select(0, first) + position_ids.index_select(0, second)) * 0.5
+    return MiniMaxH3TokenReduction(
+        begin=schedule.begin,
+        end=schedule.end,
+        first=first,
+        second=second,
+        parents=torch.cat([prefix, parent_video]),
+        pair_reduced_rows=video_start + torch.nonzero(has_pair).flatten(),
+        pair_rows=video_start + torch.nonzero(in_pair).flatten(),
+        position_ids=reduced_positions,
+        update_scale=schedule.update_scale,
+    )
+
+
+def pool_token_pairs(hidden_states: torch.Tensor, plan: MiniMaxH3TokenReduction) -> torch.Tensor:
+    """``[B, S, C]`` -> ``[B, R, C]``: unpaired rows copied, pairs averaged in float32."""
+    pooled = hidden_states.index_select(1, plan.first)
+    rows = plan.pair_reduced_rows
+    if rows.numel():
+        pair_mean = (hidden_states.index_select(1, plan.first.index_select(0, rows)).float() +
+                     hidden_states.index_select(1, plan.second.index_select(0, rows)).float()) * 0.5
+        pooled.index_copy_(1, rows, pair_mean.to(hidden_states.dtype))
+    return pooled
+
+
+def expand_token_pairs(original: torch.Tensor, reduced: torch.Tensor, pooled: torch.Tensor,
+                       plan: MiniMaxH3TokenReduction) -> torch.Tensor:
+    """Restore ``[B, S, C]`` from the reduced stack's output ``[B, R, C]``.
+
+    Unpaired rows take the reduced output directly. A paired row becomes its
+    value before pooling plus ``update_scale`` times the change of its pair,
+    ``reduced - pooled``.
+    """
+    restored = reduced.index_select(1, plan.parents)
+    rows = plan.pair_rows
+    if rows.numel():
+        parents = plan.parents.index_select(0, rows)
+        update = reduced.index_select(1, parents).float() - pooled.index_select(1, parents).float()
+        paired = original.index_select(1, rows).float() + plan.update_scale * update
+        restored.index_copy_(1, rows, paired.to(restored.dtype))
+    return restored
+
+
+__all__ = [
+    "MiniMaxH3TokenReduction",
+    "build_token_reduction",
+    "expand_token_pairs",
+    "pool_token_pairs",
+    "reduced_token_width",
+]

fastvideo/models/loader/component_loader.py

@@ -1241,6 +1241,14 @@ class TransformerLoader(ComponentLoader):
             if not isinstance(dit_config, MiniMaxH3Config):
                 raise ValueError("h3_vsa_fp4 is only supported by MiniMax-H3")
             dit_config.vsa_fp4 = dict(fastvideo_args.h3_vsa_fp4)
+        token_reduction = getattr(fastvideo_args, "h3_token_reduction", None)
+        if token_reduction and token_reduction.get("enabled"):
+            from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3Config
+            from fastvideo.configs.pipelines.minimax_h3_token_reduction import MiniMaxH3TokenReductionSchedule
+            if not isinstance(dit_config, MiniMaxH3Config):
+                raise ValueError("h3_token_reduction is only supported by MiniMax-H3")
+            # The checkpoint's own config fixes the DiT depth; reject intervals past it before any weight loads.
+            MiniMaxH3TokenReductionSchedule.from_options(token_reduction, num_blocks=dit_config.arch_config.num_layers)
 
         # Generator-only QAT for DMD distillation: the teacher (real_score) and
         # critic (fake_score) transformers load with this flag set and must stay

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

@@ -4,17 +4,24 @@
 from __future__ import annotations
 
 import contextlib
+import dataclasses
 from typing import Any
 
 import torch
 
 from fastvideo.attention.selector import component_attention_backend, get_attn_backend
+from fastvideo.configs.pipelines.minimax_h3_token_reduction import MiniMaxH3TokenReductionSchedule
 from fastvideo.distributed import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.forward_context import set_forward_context
 from fastvideo.hooks.activation_trace import trace_step
 from fastvideo.layers.pdd import PDDModalitySchedule, PDDSamplingPlan, build_pdd_sampling_plan
 from fastvideo.logger import init_logger
+from fastvideo.models.dits.minimax_h3_token_reduction import (
+    MiniMaxH3TokenReduction,
+    build_token_reduction,
+    reduced_token_width,
+)
 from fastvideo.profiler import nvtx_range, profiler_region
 from fastvideo.pipelines.basic.minimax_h3.packing import (
     MINIMAX_H3_KEYFRAME_NOISE_AUG,
@@ -122,6 +129,35 @@ def _h3_vsa_ref2va_segments(
     return tuple(segments)
 
 
+def _h3_token_reduction_plan(
+    layout: MiniMaxH3PackedLayout,
+    patch_size: tuple[int, int, int],
+    schedule: MiniMaxH3TokenReductionSchedule,
+) -> MiniMaxH3TokenReduction:
+    """Pair rows of the generated video, which every H3 packing places last in row-major ``(t, h, w)`` order."""
+    if layout.pad_tokens:
+        raise ValueError("MiniMax-H3 token reduction does not support capacity padding rows.")
+    grid = (layout.num_video_latent_frames // patch_size[0], layout.latent_height // patch_size[1],
+            layout.latent_width // patch_size[2])
+    num_target_rows = grid[0] * grid[1] * grid[2]
+    video_start = layout.sequence_length - num_target_rows
+    target_rows = layout.video_indices[layout.num_condition_video_rows:]
+    if target_rows.numel() != num_target_rows or not torch.equal(target_rows.cpu(),
+                                                                 torch.arange(video_start, layout.sequence_length)):
+        raise ValueError("MiniMax-H3 token reduction needs the generated video to end the packed sequence.")
+    return build_token_reduction(layout.position_ids, video_start, grid, schedule)
+
+
+def _h3_token_reduced_segments(
+    segments: tuple[int | tuple[int, int, int], ...],
+    patch_size: tuple[int, int, int],
+) -> tuple[int | tuple[int, int, int], ...]:
+    """VSA-H3 segments of the reduced sequence: the generated video (last region) at half token width."""
+    frames, height, width = segments[-1]  # type: ignore[misc]
+    reduced_width = reduced_token_width(width // patch_size[2]) * patch_size[2]
+    return (*segments[:-1], (frames, height, reduced_width))
+
+
 class MiniMaxH3DenoisingStage(PipelineStage):
     """Build both schedules and denoise both modalities in one transformer call."""
 
@@ -309,6 +345,31 @@ class MiniMaxH3DenoisingStage(PipelineStage):
                 logger.info("MiniMax-H3 VSA-H3: %d reference video region(s); %s; %d-token tiles.",
                             num_reference_regions, effect, vsa_tile_size)
 
+        token_reduction_schedule = MiniMaxH3TokenReductionSchedule.from_options(
+            getattr(fastvideo_args, "h3_token_reduction", {}))
+        token_reduction = None
+        reduced_vsa_metadata_builder = None
+        reduced_vsa_packed_segments: tuple[int | tuple[int, int, int], ...] = ()
+        if token_reduction_schedule is not None:
+            dit_patch_size = fastvideo_args.pipeline_config.dit_config.patch_size
+            token_reduction = _h3_token_reduction_plan(layout, dit_patch_size, token_reduction_schedule).to(device)
+            logger.info(
+                "MiniMax-H3 token reduction: %d -> %d rows in blocks %d:%d (%d:%d for the first %d step(s)), "
+                "update scale %g.", token_reduction.sequence_length, token_reduction.reduced_length,
+                token_reduction_schedule.begin, token_reduction_schedule.end, token_reduction_schedule.begin,
+                token_reduction_schedule.early_end, token_reduction_schedule.early_steps,
+                token_reduction_schedule.update_scale)
+            if (token_reduction_schedule.early_steps >= len(video_timesteps)
+                    and token_reduction_schedule.early_end != token_reduction_schedule.end):
+                logger.warning(
+                    "MiniMax-H3 token reduction: the early interval covers all %d forwards, so every forward "
+                    "reduces blocks %d:%d.", len(video_timesteps), token_reduction_schedule.begin,
+                    token_reduction_schedule.early_end)
+            if vsa_metadata_builder is not None:
+                # Its own builder, so each sequence length keeps its own tile buffers.
+                reduced_vsa_metadata_builder = _h3_vsa_metadata_builder(self.transformer, fastvideo_args)
+                reduced_vsa_packed_segments = _h3_token_reduced_segments(vsa_packed_segments, vsa_patch_size)
+
         try:
             if full_cpu_offload:
                 self.transformer.to(device)
@@ -327,6 +388,11 @@ class MiniMaxH3DenoisingStage(PipelineStage):
                         break
                     unique_timesteps, timestep_indices = row_timestep_plan[index]
                     attn_metadata = None
+                    step_token_reduction = None
+                    if token_reduction is not None:
+                        assert token_reduction_schedule is not None
+                        step_token_reduction = dataclasses.replace(token_reduction,
+                                                                   end=token_reduction_schedule.end_for_step(index))
                     if vsa_metadata_builder is not None:
                         # Optional schedule: run the first N steps dense (sparsity 0
                         # selects every tile — parity-proven ≡ dense ≤2e-4); early
@@ -343,6 +409,23 @@ class MiniMaxH3DenoisingStage(PipelineStage):
                             tile_size=vsa_tile_size,
                             ref_keep_rate=vsa_ref_keep_rate,
                         )
+                        if step_token_reduction is not None:
+                            assert reduced_vsa_metadata_builder is not None
+                            reduced_attn_metadata = reduced_vsa_metadata_builder.build(
+                                current_timestep=index,
+                                patch_size=vsa_patch_size,
+                                VSA_sparsity=vsa_sparsity,
+                                packed_segments=reduced_vsa_packed_segments,
+                                device=device,
+                                exempt=vsa_exempt,
+                                dense_layers=vsa_dense_layers,
+                                tile_size=vsa_tile_size,
+                                ref_keep_rate=vsa_ref_keep_rate,
+                            )
+                            step_token_reduction = dataclasses.replace(step_token_reduction,
+                                                                       attn_metadata=reduced_attn_metadata)
+                    # Only pass the argument when reducing, so transformers without it keep working.
+                    reduction_kwargs = {} if step_token_reduction is None else {"token_reduction": step_token_reduction}
                     # Under torch.compile(mode="reduce-overhead") each denoising
                     # step must be marked, or cudagraph trees flag cross-step
                     # reuse of pooled outputs as "accessing tensor output of
@@ -367,6 +450,7 @@ class MiniMaxH3DenoisingStage(PipelineStage):
                             video_indices=video_indices,
                             audio_indices=audio_indices,
                             text_indices=text_indices,
+                            **reduction_kwargs,
                         )
 
                     video_start = layout.num_condition_video_rows

fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py

@@ -349,3 +349,30 @@ def test_row_split_is_an_explicit_spark_opt_in():
     assert envs.environment_variables["FASTVIDEO_NVFP4_MM_ROW_SPLIT"].default == "off"
     assert example.DEVICE_PROFILES["spark"].env["FASTVIDEO_NVFP4_MM_ROW_SPLIT"] == "table"
     assert "FASTVIDEO_NVFP4_MM_ROW_SPLIT" not in example.DEVICE_PROFILES["rtx-pro-6000"].env
+
+
+@pytest.mark.parametrize("flags", [
+    ("--token-reduction", "--token-reduction-blocks", "4:60"),
+    ("--token-reduction", "--token-reduction-early", "4:30"),
+    ("--token-reduction", "--token-reduction-scale", "3"),
+    ("--token-reduction-blocks", "4:20"),
+])
+def test_token_reduction_flags_are_checked_before_download(flags):
+    with pytest.raises(SystemExit):
+        _args(*flags)
+
+
+def test_token_reduction_flags_reach_the_generator_config(tmp_path):
+    args = _args("--token-reduction", "--token-reduction-blocks", "4:20", "--token-reduction-early", "0:0")
+    config = example.build_generator_config(tmp_path, 4, token_reduction=example.token_reduction_options(args))
+    assert config.pipeline.experimental == {"h3_token_reduction": {"enabled": True, "blocks": "4:20", "early": "0:0"}}
+    assert example.build_generator_config(tmp_path, 4).pipeline.experimental == {}
+
+
+def test_invalid_token_reduction_options_fail_when_args_are_built():
+    from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig
+    from fastvideo.fastvideo_args import FastVideoArgs
+
+    with pytest.raises(ValueError, match="h3_token_reduction"):
+        FastVideoArgs(model_path="unused", pipeline_config=MiniMaxH3PipelineConfig(),
+                      h3_token_reduction={"enabled": True, "blocks": "30:4"})

fastvideo/tests/inference/test_basic_fasth3_token_reduction.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU checks of the token reduction flags in examples/inference/basic/basic_fasth3.py (no weights, no GPU).
+
+basic_fasth3_8step.py, the FastH3 8-Step V2 example, reuses this parser and generator config.
+"""
+from __future__ import annotations
+
+import importlib.util
+import sys
+from pathlib import Path
+
+import pytest
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+EXAMPLES = REPO_ROOT / "examples" / "inference" / "basic"
+
+
+def _load_example(name: str, path: Path):
+    # basic_fasth3_8step.py imports basic_fasth3 from its own directory.
+    sys.path.insert(0, str(EXAMPLES))
+    try:
+        spec = importlib.util.spec_from_file_location(name, path)
+        assert spec is not None and spec.loader is not None
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+        return module
+    finally:
+        sys.path.remove(str(EXAMPLES))
+
+
+fasth3 = _load_example("basic_fasth3_token_reduction_contract", EXAMPLES / "basic_fasth3.py")
+eight_step = _load_example("basic_fasth3_8step_token_reduction_contract", EXAMPLES / "basic_fasth3_8step.py")
+
+
+def _args(*overrides: str):
+    return fasth3.parse_args(["--prompt", "a test prompt", *overrides])
+
+
+def test_token_reduction_is_opt_in_experimental():
+    assert "h3_token_reduction" not in fasth3.build_generator_config(_args()).pipeline.experimental
+    config = fasth3.build_generator_config(_args("--token-reduction", "--token-reduction-blocks", "4:20"))
+    assert config.pipeline.experimental["h3_token_reduction"] == {"enabled": True, "blocks": "4:20"}
+
+
+@pytest.mark.parametrize("flags", [
+    ("--token-reduction", "--token-reduction-blocks", "4:60"),
+    ("--token-reduction", "--token-reduction-early", "4:30"),
+    ("--token-reduction", "--token-reduction-scale", "3"),
+    ("--token-reduction-scale", "0.5"),
+])
+def test_token_reduction_flags_are_checked_before_download(flags):
+    with pytest.raises(SystemExit):
+        _args(*flags)
+
+
+def test_eight_step_example_accepts_token_reduction():
+    args = eight_step.parse_args(["--prompt", "p", "--token-reduction", "--token-reduction-early", "0:0"])
+    config = fasth3.build_generator_config(args)
+    assert config.model_path == "FastVideo/FastVideo-FastH3-8-Step-V2"
+    assert config.pipeline.experimental["h3_token_reduction"] == {"enabled": True, "early": "0:0"}

fastvideo/tests/transformers/test_minimax_h3_token_reduction.py

@@ -0,0 +1,344 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MiniMax H3 training-free token reduction (antirez/h3.c ``--token-reduction``).
+
+CPU checks cover option parsing, the pair row maps, the pool/restore
+arithmetic and the denoising-stage helpers. The GPU checks run a tiny real
+transformer on the dense SDPA backend. A one-token-wide video has no pairs, so
+the reduced path must be bit-identical to the plain forward; that covers the
+reduced positions, modulation rows and the entry and exit plumbing. Separate
+checks record which attention metadata every block reads, and that the
+full-sequence metadata comes back after the forward and after an error.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from fastvideo import envs
+from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3ArchConfig, MiniMaxH3Config
+from fastvideo.configs.pipelines.minimax_h3_token_reduction import MiniMaxH3TokenReductionSchedule
+from fastvideo.forward_context import get_forward_context, set_forward_context
+from fastvideo.models.dits.minimax_h3_token_reduction import (
+    build_token_reduction,
+    expand_token_pairs,
+    pool_token_pairs,
+)
+from fastvideo.pipelines.basic.minimax_h3.packing import build_packed_sequence, build_row_timesteps
+from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_denoising import (
+    _h3_token_reduced_segments,
+    _h3_token_reduction_plan,
+)
+
+_PATCH = (1, 2, 2)
+
+
+def _schedule(**options) -> MiniMaxH3TokenReductionSchedule:
+    schedule = MiniMaxH3TokenReductionSchedule.from_options({"enabled": True, **options})
+    assert schedule is not None
+    return schedule
+
+
+# ----------------------------------------------------------------------------
+# Options
+# ----------------------------------------------------------------------------
+
+
+def test_schedule_is_off_unless_enabled() -> None:
+    assert MiniMaxH3TokenReductionSchedule.from_options({}) is None
+    assert MiniMaxH3TokenReductionSchedule.from_options({"enabled": False, "blocks": [1, 2]}) is None
+
+
+def test_schedule_defaults_follow_h3c_share_of_eight_forwards() -> None:
+    schedule = _schedule()
+    assert (schedule.begin, schedule.end, schedule.early_steps, schedule.early_end) == (4, 30, 4, 40)
+    assert [schedule.end_for_step(step) for step in range(8)] == [40, 40, 40, 40, 30, 30, 30, 30]
+
+
+def test_schedule_skips_early_interval_when_base_reaches_block_40() -> None:
+    schedule = _schedule(blocks="4:45")
+    assert (schedule.early_steps, schedule.end_for_step(0)) == (0, 45)
+
+
+def test_schedule_accepts_strings_and_disabling_early() -> None:
+    schedule = _schedule(blocks="2:20", early="0:0", update_scale=0.5)
+    assert (schedule.begin, schedule.end, schedule.end_for_step(0), schedule.update_scale) == (2, 20, 20, 0.5)
+
+
+@pytest.mark.parametrize("options,match", [
+    ({"enabled": "yes"}, "enabled must be a boolean"),
+    ({"enabled": True, "blocks": [30, 4]}, "BEGIN < END"),
+    ({"enabled": True, "blocks": "4"}, "two integers"),
+    ({"enabled": True, "early": [2, 20]}, "must exceed blocks END"),
+    ({"enabled": True, "update_scale": 3.0}, r"\[0, 2\]"),
+    ({"enabled": True, "update_scale": True}, r"\[0, 2\]"),
+    ({"enabled": True, "begin": 4}, "keys from"),
+])
+def test_schedule_rejects_bad_options(options, match) -> None:
+    with pytest.raises(ValueError, match=match):
+        MiniMaxH3TokenReductionSchedule.from_options(options)
+
+
+def test_schedule_bounds_both_intervals_by_the_dit_depth() -> None:
+    assert MiniMaxH3TokenReductionSchedule.from_options({"enabled": True, "blocks": "4:50"}, num_blocks=50) is not None
+    for options in ({"blocks": "4:51"}, {"early": "4:60"}):
+        with pytest.raises(ValueError, match="past the 50 DiT blocks"):
+            MiniMaxH3TokenReductionSchedule.from_options({"enabled": True, **options}, num_blocks=50)
+
+
+# ----------------------------------------------------------------------------
+# Row maps and arithmetic
+# ----------------------------------------------------------------------------
+
+
+def _positions(sequence_length: int) -> torch.Tensor:
+    return torch.arange(sequence_length * 3, dtype=torch.float64).reshape(sequence_length, 3)
+
+
+def test_pairs_cover_rows_and_leave_odd_column_alone() -> None:
+    # Two prefix rows, then a 1x2x3 video: columns (0, 1) pair, column 2 is alone.
+    plan = build_token_reduction(_positions(8), 2, (1, 2, 3), _schedule())
+    assert plan.first.tolist() == [0, 1, 2, 4, 5, 7]
+    assert plan.second.tolist() == [0, 1, 3, 4, 6, 7]
+    assert plan.parents.tolist() == [0, 1, 2, 2, 3, 4, 4, 5]
+    assert plan.pair_reduced_rows.tolist() == [2, 4]
+    assert plan.pair_rows.tolist() == [2, 3, 5, 6]
+    assert (plan.sequence_length, plan.reduced_length) == (8, 6)
+
+
+def test_pair_positions_are_the_mean_and_unpaired_positions_are_exact() -> None:
+    positions = _positions(8)
+    plan = build_token_reduction(positions, 2, (1, 2, 3), _schedule())
+    torch.testing.assert_close(plan.position_ids[2], (positions[2] + positions[3]) / 2, rtol=0, atol=0)
+    for reduced_row, full_row in ((0, 0), (1, 1), (3, 4), (5, 7)):
+        assert torch.equal(plan.position_ids[reduced_row], positions[full_row])
+
+
+def test_plan_requires_the_video_to_end_the_sequence() -> None:
+    with pytest.raises(ValueError, match="end the packed sequence"):
+        build_token_reduction(_positions(9), 2, (1, 2, 3), _schedule())
+
+
+def test_h3c_worked_example() -> None:
+    # One prefix row (10), then a 1x1x4 video row 1, 3, 5, 9.
+    plan = build_token_reduction(_positions(5), 1, (1, 1, 4), _schedule())
+    hidden = torch.tensor([10.0, 1.0, 3.0, 5.0, 9.0]).view(1, 5, 1)
+    pooled = pool_token_pairs(hidden, plan)
+    assert pooled.flatten().tolist() == [10.0, 2.0, 7.0]
+    # The middle blocks move the prefix to 11 and the pairs to 2.5 and 6.
+    reduced = torch.tensor([11.0, 2.5, 6.0]).view(1, 3, 1)
+    restored = expand_token_pairs(hidden, reduced, pooled, plan)
+    assert restored.flatten().tolist() == [11.0, 1.5, 3.5, 4.0, 8.0]
+
+
+def test_update_scale_scales_only_the_pair_update() -> None:
+    plan = build_token_reduction(_positions(5), 1, (1, 1, 4), _schedule(update_scale=0.5))
+    hidden = torch.tensor([10.0, 1.0, 3.0, 5.0, 9.0]).view(1, 5, 1)
+    pooled = pool_token_pairs(hidden, plan)
+    reduced = torch.tensor([11.0, 2.5, 6.0]).view(1, 3, 1)
+    assert expand_token_pairs(hidden, reduced, pooled, plan).flatten().tolist() == [11.0, 1.25, 3.25, 4.5, 8.5]
+
+
+def test_unchanged_reduced_stack_restores_bf16_input_exactly() -> None:
+    plan = build_token_reduction(_positions(2 + 2 * 3 * 5), 2, (2, 3, 5), _schedule())
+    hidden = torch.randn(1, plan.sequence_length, 8, generator=torch.Generator().manual_seed(0)).bfloat16()
+    pooled = pool_token_pairs(hidden, plan)
+    assert torch.equal(expand_token_pairs(hidden, pooled, pooled, plan), hidden)
+
+
+def test_one_token_wide_video_is_an_identity() -> None:
+    plan = build_token_reduction(_positions(2 + 3), 2, (3, 1, 1), _schedule())
+    assert plan.reduced_length == plan.sequence_length
+    assert plan.pair_rows.numel() == 0
+    hidden = torch.randn(1, plan.sequence_length, 4).bfloat16()
+    assert torch.equal(pool_token_pairs(hidden, plan), hidden)
+
+
+# ----------------------------------------------------------------------------
+# Denoising-stage helpers on a real packed layout
+# ----------------------------------------------------------------------------
+
+
+def _layout(latent_width: int, keyframes: tuple[str, ...] = ()):
+    return build_packed_sequence(
+        text_token_tags=torch.ones(4, dtype=torch.long),
+        num_latent_frames=2,
+        latent_height=4,
+        latent_width=latent_width,
+        num_audio_latents=3,
+        patch_size=_PATCH,
+        keyframe_anchors=keyframes,
+    )
+
+
+@pytest.mark.parametrize("keyframes", [(), ("first", )])
+def test_stage_plan_pairs_only_the_generated_video(keyframes) -> None:
+    layout = _layout(latent_width=6, keyframes=keyframes)
+    plan = _h3_token_reduction_plan(layout, _PATCH, _schedule())
+    video_start = layout.sequence_length - 2 * 2 * 3
+    assert plan.sequence_length == layout.sequence_length
+    assert plan.reduced_length == video_start + 2 * 2 * 2
+    assert plan.first[:video_start].tolist() == list(range(video_start))
+    assert int(plan.pair_rows.min()) >= video_start
+
+
+def test_reduced_segments_halve_the_last_region_width() -> None:
+    assert _h3_token_reduced_segments((4, 0, 6, (2, 4, 6)), _PATCH) == (4, 0, 6, (2, 4, 4))
+    assert _h3_token_reduced_segments((4, (3, 4, 4), 6, (2, 4, 8)), _PATCH) == (4, (3, 4, 4), 6, (2, 4, 4))
+
+
+# ----------------------------------------------------------------------------
+# Tiny real transformer (GPU)
+# ----------------------------------------------------------------------------
+
+
+def _tiny_transformer(device: torch.device):
+    from fastvideo.models.dits.minimax_h3 import MiniMaxH3Transformer3DModel
+
+    arch = MiniMaxH3ArchConfig(
+        num_attention_heads=2,
+        attention_head_dim=16,
+        hidden_size=32,
+        num_layers=3,
+        num_refiner_layers=1,
+        ffn_dim=64,
+        freq_dim=16,
+        time_embed_hidden_dim=32,
+        time_embed_dim=16,
+        rope_freq_dim=2,
+    )
+    transformer = MiniMaxH3Transformer3DModel(config=MiniMaxH3Config(arch_config=arch), hf_config={})
+    generator = torch.Generator(device="cpu").manual_seed(0)
+    with torch.no_grad():
+        for name, parameter in transformer.named_parameters():
+            if parameter.ndim < 2 and "norm" in name:
+                parameter.fill_(1.0)
+            else:
+                parameter.copy_(torch.randn(parameter.shape, generator=generator) * 0.05)
+    return transformer.to(device).eval()
+
+
+def _call_kwargs(layout, inputs) -> dict[str, torch.Tensor]:
+    device = inputs["hidden_states"].device
+    unique, inverse = build_row_timesteps(layout, 0.5, 0.5, 0.999, 1.0)
+    return {
+        "timestep": unique.to(device),
+        "timestep_indices": inverse.to(device),
+        "token_tags": layout.token_tags.to(device),
+        "position_ids": layout.position_ids.to(device),
+        "video_indices": layout.video_indices.to(device),
+        "audio_indices": layout.audio_indices.to(device),
+        "text_indices": layout.text_indices.to(device),
+        **inputs,
+    }
+
+
+def _forward(transformer, layout, inputs, token_reduction=None, attn_metadata=None):
+    with torch.no_grad(), set_forward_context(current_timestep=0, attn_metadata=attn_metadata):
+        outputs = transformer(token_reduction=token_reduction, **_call_kwargs(layout, inputs))
+        # The full-sequence metadata is back in place after the forward.
+        assert get_forward_context().attn_metadata is attn_metadata
+    return outputs
+
+
+def _inputs(transformer, layout, device: torch.device) -> dict[str, torch.Tensor]:
+    generator = torch.Generator(device="cpu").manual_seed(1)
+
+    def rows(count: int, width: int) -> torch.Tensor:
+        return torch.randn(1, count, width, generator=generator).to(device)
+
+    return {
+        "hidden_states": rows(layout.video_indices.numel(), transformer.proj_in.weight.shape[1]),
+        "audio_hidden_states": rows(layout.audio_indices.numel(), transformer.audio_proj_in.weight.shape[1]),
+        "encoder_hidden_states": rows(layout.text_indices.numel(), transformer.context_embedder.weight.shape[1]),
+    }
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(),
+                    reason="MiniMax H3 attention layers require a CUDA device for a real forward")
+def test_one_token_wide_reduced_forward_is_bit_identical(env_overrides, distributed_setup) -> None:
+    env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA"))
+    device = torch.device("cuda")
+    transformer = _tiny_transformer(device)
+    layout = _layout(latent_width=2)
+    inputs = _inputs(transformer, layout, device)
+    plan = _h3_token_reduction_plan(layout, _PATCH, _schedule(blocks="1:2", early="0:0")).to(device)
+    plain = _forward(transformer, layout, inputs)
+    reduced = _forward(transformer, layout, inputs, plan)
+    for expected, actual in zip(plain, reduced, strict=True):
+        assert torch.equal(expected, actual)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(),
+                    reason="MiniMax H3 attention layers require a CUDA device for a real forward")
+@pytest.mark.parametrize("blocks", ["1:2", "0:3"])
+def test_reduced_forward_keeps_shapes_and_changes_video(env_overrides, distributed_setup, blocks) -> None:
+    env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA"))
+    device = torch.device("cuda")
+    transformer = _tiny_transformer(device)
+    layout = _layout(latent_width=8)
+    inputs = _inputs(transformer, layout, device)
+    plan = _h3_token_reduction_plan(layout, _PATCH, _schedule(blocks=blocks, early="0:0")).to(device)
+    plain_video, plain_audio = _forward(transformer, layout, inputs)
+    video, audio = _forward(transformer, layout, inputs, plan)
+    assert video.shape == plain_video.shape and audio.shape == plain_audio.shape
+    assert torch.isfinite(video).all() and torch.isfinite(audio).all()
+    assert not torch.equal(video, plain_video)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(),
+                    reason="MiniMax H3 attention layers require a CUDA device for a real forward")
+@pytest.mark.parametrize("blocks", ["1:2", "1:3"])
+def test_reduced_blocks_read_their_own_attention_metadata(env_overrides, distributed_setup, blocks) -> None:
+    # "1:3" ends at the last block, so the exit and restore happen after the block loop.
+    env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA"))
+    device = torch.device("cuda")
+    transformer = _tiny_transformer(device)
+    layout = _layout(latent_width=8)
+    inputs = _inputs(transformer, layout, device)
+    full_metadata, reduced_metadata = SimpleNamespace(name="full"), SimpleNamespace(name="reduced")
+    plan = _h3_token_reduction_plan(layout, _PATCH, _schedule(blocks=blocks, early="0:0")).to(device)
+    plan = dataclasses.replace(plan, attn_metadata=reduced_metadata)
+    seen = []
+    hooks = [
+        block.register_forward_pre_hook(lambda module, args: seen.append(get_forward_context().attn_metadata))
+        for block in transformer.transformer_blocks
+    ]
+    try:
+        _forward(transformer, layout, inputs, plan, attn_metadata=full_metadata)
+    finally:
+        for hook in hooks:
+            hook.remove()
+    expected = [
+        reduced_metadata if plan.begin <= index < plan.end else full_metadata
+        for index in range(len(transformer.transformer_blocks))
+    ]
+    assert [metadata.name for metadata in seen] == [metadata.name for metadata in expected]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(),
+                    reason="MiniMax H3 attention layers require a CUDA device for a real forward")
+def test_attention_metadata_is_restored_after_an_error_in_a_reduced_block(env_overrides, distributed_setup) -> None:
+    env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA"))
+    device = torch.device("cuda")
+    transformer = _tiny_transformer(device)
+    layout = _layout(latent_width=8)
+    inputs = _inputs(transformer, layout, device)
+    full_metadata = SimpleNamespace(name="full")
+    plan = _h3_token_reduction_plan(layout, _PATCH, _schedule(blocks="1:2", early="0:0")).to(device)
+    plan = dataclasses.replace(plan, attn_metadata=SimpleNamespace(name="reduced"))
+
+    def fail(module, args):
+        raise RuntimeError("block failed")
+
+    hook = transformer.transformer_blocks[1].register_forward_pre_hook(fail)
+    try:
+        with torch.no_grad(), set_forward_context(current_timestep=0, attn_metadata=full_metadata):
+            with pytest.raises(RuntimeError, match="block failed"):
+                transformer(token_reduction=plan, **_call_kwargs(layout, inputs))
+            assert get_forward_context().attn_metadata is full_metadata
+    finally:
+        hook.remove()