diff --git a/.buildkite/scripts/unit_test.sh b/.buildkite/scripts/unit_test.sh index 7fa4b3799..aacba0fcc 100644 --- a/.buildkite/scripts/unit_test.sh +++ b/.buildkite/scripts/unit_test.sh @@ -15,8 +15,12 @@ exec pytest \ ./fastvideo/tests/ops/ \ ./fastvideo/tests/worker/ \ ./fastvideo/tests/training/test_trackers.py \ + ./fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py \ ./fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py \ ./fastvideo/tests/attention/test_vsa_h3_tile_grad_safety.py \ + ./fastvideo/tests/attention/test_vsa_h3_metadata.py \ + ./fastvideo/tests/attention/test_vsa_h3_ref2va_regions.py \ + ./fastvideo/tests/layers/test_pdd_linear.py \ ./fastvideo/tests/modal/test_kernel_build_cache.py \ ./fastvideo/tests/modal/test_pr_test.py \ ./fastvideo/tests/modal/test_ssim_test.py \ diff --git a/docs/design/inference_schema_parity_inventory.yaml b/docs/design/inference_schema_parity_inventory.yaml index 463977b4c..4309e9851 100644 --- a/docs/design/inference_schema_parity_inventory.yaml +++ b/docs/design/inference_schema_parity_inventory.yaml @@ -470,12 +470,14 @@ surfaces: coords_style: "MagiHuman internal data-proxy coordinate convention." frame_receptive_field: "MagiHuman internal data-proxy receptive-field setting." image_conditioning: "MagiHuman preset variant marker for reference-image conditioning." + pdd_step_indices: "MiniMax-H3 PDD fine-grid partition; resolve_checkpoint_settings sets it from fastvideo_inference.json." ref_audio_offset: "MagiHuman internal data-proxy audio alignment offset." scheduler_sigma_min: "Z-Image scheduler parity invariant; not part of the public typed inference API." scheduler_use_reference_discrete_timesteps: "Z-Image scheduler parity invariant; not part of the public typed inference API." sr_local_attn_layers: "MagiHuman SR internal sparse-attention layer selection." text_offset: "MagiHuman internal data-proxy text alignment offset." vae_stride: "MagiHuman internal VAE/data-proxy stride setting." + vsa_ref_keep_rate: "MiniMax-H3 PDD reference-video VSA keep rate; defaults to the checkpoint's trained value." z_dim: "MagiHuman internal VAE latent channel setting." vocoder_config: "Legacy internal component config object." vocoder_precision: "Precision override pending dedicated component precision design." diff --git a/docs/inference/fasth3-distilled.md b/docs/inference/fasth3-distilled.md index 27035b92f..d77f68bd5 100644 --- a/docs/inference/fasth3-distilled.md +++ b/docs/inference/fasth3-distilled.md @@ -74,6 +74,130 @@ This documents execution support for the published checkpoint. It is not a quality claim: compare video/audio output against base MiniMax-H3 on your own prompts before adopting it. +## Ref2VA PDD students + +A Parallel Decoding Distillation (PDD) student widens the transformer's two +output projections to `pdd_steps` heads, one per interval of a fixed fine +time grid on `[0, 0.999]`. Each transformer forward fuses a block of +consecutive heads into their integration-weighted mean, and an ordinary Euler +step over the block's two node sigmas applies it. The FastH3 OmniRef PDD-8 +student is a Ref2VA (`transformer_ref`) student with 32 heads, sampled in +eight blocks of four. + +Its export carries only what distillation changed: + +| Path | Contents | +| --- | --- | +| `transformer_ref/` | Widened `proj_out` and `audio_proj_out`, trained VSA compression gates; `config.json` records `pdd_steps` | +| `scheduler/`, `audio_scheduler/` | Video and audio shifts (12 and 3) | +| `fastvideo_inference.json` | The sampling contract below | +| `modular_model_index.json` | The Diffusers manifest | + +The text encoder, tokenizer, processor, and both VAEs are base MiniMax-H3's. +`basic_fasth3_omniref_pdd.py` composes the two into one local directory of +symlinks and runs the recipe the contract records. `--model-path` is the +export, as a local directory or a Hugging Face repo id. The base comes from +the revision the contract pins in `base_model_revision`, or from +`--base-model-path`: + +```bash +python examples/inference/basic/basic_fasth3_omniref_pdd.py \ + --model-path \ + --video reference.mp4 --image character.png \ + --prompt 'The dancer from the video performs the routine in the pictured outfit.' \ + --height 480 --width 832 --num-frames 124 \ + --output outputs/fasth3-omniref-pdd +``` + +References are ordered: pass `--image`, `--video`, and `--audio` in the order +the prompt refers to them. At least one image or video is required. + +A PDD export uses `fasth3-inference-contract-v1` with PDD fields in place of +the DMD ladder: + +```json +{ + "schema_version": "fasth3-inference-contract-v1", + "model_type": "ref2va", + "transformer_component": "transformer_ref", + "pdd_steps": 32, + "pdd_step_indices": [0, 4, 8, 12, 16, 20, 24, 28, 32], + "num_inference_steps": 8, + "transformer_forwards": 8, + "grid_max_t": 0.999, + "video_scheduler_shift": 12.0, + "audio_scheduler_shift": 3.0, + "guidance_scale": 1.0, + "attention_backend": "VIDEO_SPARSE_ATTN_H3", + "vsa_sparsity": 0.9, + "vsa_tile_size": 128, + "vsa_ref_policy": "p2_multi_region", + "vsa_ref_keep_rate": 0.1 +} +``` + +The export also records `schema`, `conditioning`, and `base_model_revision`, +the base snapshot it was distilled against, as `hf://@` +(for example `hf://MiniMaxAI/MiniMax-H3@`); any other form is an +error. FastVideo reads the file once, when the run's arguments are built and +before any weights load. The file must carry exactly these fields: a missing +or unknown field is an error. Fields that repeat a value stored elsewhere must +equal it: `pdd_steps` equals `transformer_ref/config.json`, the shifts equal +the scheduler configs, and `num_inference_steps` and `transformer_forwards` +equal the block count of `pdd_step_indices`, a strictly increasing partition +of the grid from 0 to `pdd_steps`. Only `MiniMaxH3Ref2VAModularPipeline` runs +the export. A `transformer_ref` whose `config.json` sets `pdd_steps` without +this file is an error. + +Each sampling setting has one source, the file: + +| Setting | Run leaves it unset | Run sets a different value | +| ------------------------------- | ---------------------- | ---------------------------------------------- | +| `pdd_step_indices` | Taken from the file | Error | +| `num_inference_steps` (request) | Set to the block count | Error | +| `attention_backend` | Taken from the file | Error, including `FASTVIDEO_ATTENTION_BACKEND` | +| `VSA_tile_size` | Taken from the file | Error | +| `VSA_sparsity` | Taken from the file | The run's value, with a warning | +| `vsa_ref_keep_rate` | Taken from the file | The run's value, with a warning | + +A request leaves `num_inference_steps` unset only when it is parsed from a +mapping or a config file. A `GenerationRequest` built in Python counts every +field as set, so it must pass the block count. The trained compression gates +of `transformer_ref` load only under `VIDEO_SPARSE_ATTN_H3`, so the attention +backend cannot change. `dmd_denoising_steps` must stay unset. + +A MiniMax-H3 checkpoint without PDD fields, such as a DMD export, whose +transformer carries VSA compression gates (`to_gate_compress` weights) also +runs only with `VIDEO_SPARSE_ATTN_H3`. The pipeline reads the shard index (or +the safetensors headers) and rejects any other backend, including automatic +selection, before it loads any component. + +### Reference-video sparsity + +Every PDD contract sets `"vsa_ref_policy": "p2_multi_region"`, the only value FastVideo accepts, which means: +`VIDEO_SPARSE_ATTN_H3` tiles every reference video as its own sparse region, +in place in the packed sequence. Each video query keeps `vsa_ref_keep_rate` of +every reference video's tiles and `1 - VSA_sparsity` of the target video's +tiles. Text, audio, and image references stay dense. With `VSA_sparsity` 0, +every region is dense and the reference keep rate has no effect. Checkpoints +other than PDD students keep every conditioning row dense. + +### Hardware + +The contract's 128-token tiles, `(4, 4, 8)`, run only on the sm_100a/sm_103a +CUDA block-sparse kernel (B200, B300, GB200, GB300) of a fastvideo-kernel +build with the Blackwell VSA extension. There is no Triton fallback for tile +128: the backend raises instead. Tile 128 runs eagerly. Regional compile +requires 64-token tiles. With `--num-gpus` above 1 the example shards the DiT +across the GPUs (FSDP) and splits the sequence across them. + +On GB200, a 480x832, 124-frame request fits on one GPU: the eight forwards +take about 54 s, and device memory in use peaks near 103 GiB. A 768x1344, +345-frame request with `--num-gpus 4` takes about 36 s for the eight +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. + ## Apple Silicon `mlx_fasth3.py` stays on FastH3 V1 and its uniform AdaLN cache. diff --git a/docs/inference/support_matrix.md b/docs/inference/support_matrix.md index 9b2bc90ec..5fb510219 100644 --- a/docs/inference/support_matrix.md +++ b/docs/inference/support_matrix.md @@ -101,7 +101,10 @@ to convert the official weights locally and set `MMAUDIO_MODEL_PATH`. **Note (MiniMax H3)**: T2VA, FL2VA, and Ref2VA all generate video with stereo audio. Use the Ref2VA example when passing ordered image, video, or audio -references. +references. Distilled Ref2VA PDD students (eight transformer forwards, with +reference videos as sparse VSA regions) run through +[basic_fasth3_omniref_pdd.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_fasth3_omniref_pdd.py); +see [FastH3 distilled checkpoint schedules](fasth3-distilled.md#ref2va-pdd-students). **Note (Wan-VACE)**: not currently supported — no VACE pipeline or registered model ID exists on `main` diff --git a/examples/inference/basic/README.md b/examples/inference/basic/README.md index 7c948bd9b..fcc5b0696 100644 --- a/examples/inference/basic/README.md +++ b/examples/inference/basic/README.md @@ -167,6 +167,24 @@ launchers still use sparse attention at strength `0` and require FastVideo's tile-64 VSA kernel; the dense launcher selects FA4. Each launcher writes to its own variant directory by default so comparison outputs do not collide. +### FastH3 OmniRef PDD (Ref2VA) + +[basic_fasth3_omniref_pdd.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_fasth3_omniref_pdd.py) +runs a Parallel Decoding Distillation Ref2VA student in eight transformer +forwards. The export carries only its `transformer_ref`, scheduler configs, +and `fastvideo_inference.json`; the script links them with the base +MiniMax-H3 components into one local model directory. Its 128-token VSA tiles +need the sm_100a/sm_103a kernel: + +```bash +python examples/inference/basic/basic_fasth3_omniref_pdd.py \ + --model-path \ + --video reference.mp4 --image character.png --prompt "your prompt" +``` + +See [FastH3 distilled checkpoint schedules](https://github.com/hao-ai-lab/FastVideo/blob/main/docs/inference/fasth3-distilled.md#ref2va-pdd-students) +for the contract and the reference-video sparsity policy. + ## Basic Walkthrough All you need to generate videos using multi-gpus from state-of-the-art diffusion pipelines is the following few lines! diff --git a/examples/inference/basic/basic_fasth3_omniref_pdd.py b/examples/inference/basic/basic_fasth3_omniref_pdd.py new file mode 100644 index 000000000..8f9673c82 --- /dev/null +++ b/examples/inference/basic/basic_fasth3_omniref_pdd.py @@ -0,0 +1,268 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Eight-forward Ref2VA video+audio generation with a FastH3 OmniRef PDD student. + +A FastH3 OmniRef PDD export is a Parallel Decoding Distillation (PDD) student of +MiniMax-H3's reference-conditioned ``transformer_ref`` partition. Its directory +(local, or a Hugging Face repo) carries only what distillation changed: +``transformer_ref/`` (output heads widened to the fine grid, trained VSA +compression gates), the two scheduler configs, and ``fastvideo_inference.json``. +The text encoder, tokenizer, processor, and both VAEs are base MiniMax-H3's. + +This script composes the two into one local model directory of symlinks and +runs the recipe the contract records: fused blocks of the fine grid (one +transformer forward each), the video/audio shifts, and VIDEO_SPARSE_ATTN_H3 +with its sparsity and tile size, where every reference video is its own +sparse region. The contract's 128-token tiles run only on the sm_100a/sm_103a +CUDA kernel (B200/B300/GB200/GB300) of a fastvideo-kernel build with the +Blackwell VSA extension; the VSA-H3 backend raises at the first attention call +when either is missing. +With ``--num-gpus`` above 1 the DiT is sharded across the GPUs (FSDP) and the +sequence is split across them (sequence parallelism). + +References are ordered. Pass them in order with --image / --video / --audio, +for example ``--video dance.mp4 --image outfit.png --audio voice.wav``. +""" + +from __future__ import annotations + +import argparse +import json +from collections.abc import Sequence +from pathlib import Path +from typing import Any + +from fastvideo import VideoGenerator +from fastvideo.api import ( + ComponentConfig, + EngineConfig, + GenerationRequest, + GeneratorConfig, + InputConfig, + OffloadConfig, + OutputConfig, + ParallelismConfig, + PipelineSelection, + SamplingConfig, +) +from fastvideo.configs.pipelines.minimax_h3 import parse_base_model_revision +from fastvideo.pipelines.basic.minimax_h3 import MiniMaxH3Reference + +CONTRACT = "fastvideo_inference.json" +# Components the distilled export replaces, and the ones it shares with base MiniMax-H3. +EXPORT_COMPONENTS = ("transformer_ref", "scheduler", "audio_scheduler") +BASE_COMPONENTS = ("text_encoder", "tokenizer", "processor", "vae", "audio_vae") +MANIFESTS = ("modular_model_index.json", "model_index.json") +# Components the composed directory's manifest must declare, besides its transformer_ref. +MANIFEST_COMPONENTS = ("scheduler", "audio_scheduler", *BASE_COMPONENTS) + + +class _AppendReference(argparse.Action): + """Collect --image/--video/--audio into one list that keeps command-line order.""" + + def __call__(self, parser, namespace, value, option_string=None): + references = list(getattr(namespace, self.dest, None) or []) + references.append((self.const, value)) + setattr(namespace, self.dest, references) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--model-path", + required=True, + help="FastH3 OmniRef PDD export: a local directory or a Hugging Face repo id") + parser.add_argument("--revision", default=None, help="revision of --model-path when it is a Hugging Face repo") + parser.add_argument("--base-model-path", + default=None, + help="base MiniMax-H3 snapshot (local directory or repo id). Default: the repo and revision " + "pinned by the export's base_model_revision") + parser.add_argument("--base-revision", + default=None, + help="revision of --base-model-path when it is a repo id (the export's pinned revision " + "applies only to its own base repo); without --base-model-path, overrides that pin") + parser.add_argument("--composed-dir", + default=None, + help="where to write the composed model directory of symlinks (default: " + "/fasth3_omniref_model)") + for flag in ("image", "video", "audio"): + parser.add_argument(f"--{flag}", + dest="references", + action=_AppendReference, + const=flag, + metavar="PATH", + help=f"an ordered {flag} reference (repeatable)") + parser.add_argument("--prompt", required=True) + parser.add_argument("--output", default="outputs/fasth3_omniref_pdd") + parser.add_argument("--height", type=int, default=480) + parser.add_argument("--width", type=int, default=832) + parser.add_argument("--num-frames", type=int, default=124) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--num-gpus", type=int, default=1, help="sequence-parallel GPUs (must divide 56 heads)") + return parser + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = build_parser() + args = parser.parse_args(argv) + if not args.references: + 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") + return args + + +def _snapshot(path_or_repo: str, revision: str | None, allow_patterns: list[str]) -> Path: + local = Path(path_or_repo).expanduser() + if local.is_dir(): + return local + from huggingface_hub import snapshot_download + + return Path(snapshot_download(repo_id=path_or_repo, revision=revision, allow_patterns=allow_patterns)) + + +def load_contract(export_dir: Path) -> dict[str, Any]: + path = export_dir / CONTRACT + if not path.is_file(): + raise FileNotFoundError(f"{export_dir} has no {CONTRACT}; it is not a FastH3 distilled export.") + contract = json.loads(path.read_text(encoding="utf-8")) + if "pdd_steps" not in contract or contract.get("model_type") != "ref2va": + raise ValueError(f"{path} is not a Ref2VA PDD contract; use the example matching the checkpoint.") + return contract + + +def base_model_source(args: argparse.Namespace, contract: dict[str, Any]) -> tuple[str, str | None]: + """The base snapshot and revision: --base-model-path as given, else the export's pin.""" + base_repo, base_revision = base_model_from_contract(contract) + if args.base_model_path: + return args.base_model_path, args.base_revision + return base_repo, args.base_revision or base_revision + + +def base_model_from_contract(contract: dict[str, Any]) -> tuple[str, str]: + """The base repo id and revision that the export pins as ``hf://@``.""" + return parse_base_model_revision(contract.get("base_model_revision")) + + +def _link(destination: Path, source: Path) -> None: + if not source.exists(): + raise FileNotFoundError(f"Missing checkpoint component: {source}") + source = source.resolve() + if destination.is_symlink(): + if destination.resolve() == source: + return + destination.unlink() + elif destination.exists(): + raise FileExistsError(f"{destination} exists and is not a symlink; choose another --composed-dir.") + destination.symlink_to(source, target_is_directory=source.is_dir()) + + +def _declares_components(manifest: Path) -> bool: + """Whether a Diffusers manifest declares every component of the composed directory.""" + declared = { + name + for name, spec in json.loads(manifest.read_text(encoding="utf-8")).items() + if isinstance(spec, list) and spec and spec[0] is not None + } + return set(MANIFEST_COMPONENTS) <= declared and bool({"transformer", "transformer_ref"} & declared) + + +def select_manifest(export_dir: Path, base_dir: Path) -> Path: + """The export's manifest when it declares every composed component, else the base's.""" + for root in (export_dir, base_dir): + for name in MANIFESTS: + if (root / name).is_file() and _declares_components(root / name): + return root / name + raise FileNotFoundError("Neither the export nor the base snapshot has a Diffusers model manifest that declares " + f"{', '.join(MANIFEST_COMPONENTS)} and transformer_ref.") + + +def compose_model_dir(export_dir: Path, base_dir: Path, composed_dir: Path) -> Path: + """One MiniMax-H3 model directory: distilled components from the export, the rest from the base.""" + composed_dir.mkdir(parents=True, exist_ok=True) + for name in (*EXPORT_COMPONENTS, CONTRACT): + _link(composed_dir / name, export_dir / name) + manifest_source = select_manifest(export_dir, base_dir) + for name in MANIFESTS: + stale = composed_dir / name + if name != manifest_source.name and stale.is_symlink(): + stale.unlink() + _link(composed_dir / manifest_source.name, manifest_source) + for name in BASE_COMPONENTS: + _link(composed_dir / name, base_dir / name) + return composed_dir + + +def resolve_model(args: argparse.Namespace) -> tuple[Path, dict[str, Any]]: + export_dir = _snapshot(args.model_path, args.revision, + [CONTRACT, *MANIFESTS, *(f"{name}/**" for name in EXPORT_COMPONENTS)]) + contract = load_contract(export_dir) + base_source, base_revision = base_model_source(args, contract) + base_dir = _snapshot(base_source, base_revision, [*MANIFESTS, *(f"{name}/**" for name in BASE_COMPONENTS)]) + composed_dir = Path(args.composed_dir) if args.composed_dir else Path(args.output) / "fasth3_omniref_model" + return compose_model_dir(export_dir, base_dir, composed_dir), contract + + +def build_generator_config(model_dir: Path, num_gpus: int) -> GeneratorConfig: + return GeneratorConfig( + model_path=str(model_dir), + engine=EngineConfig( + num_gpus=num_gpus, + # Shard the DiT across the GPUs rather than holding a full copy on each. + use_fsdp_inference=num_gpus > 1, + parallelism=ParallelismConfig(tp_size=1, sp_size=num_gpus), + offload=OffloadConfig(dit=False, dit_layerwise=False, text_encoder=True, vae=True, pin_cpu_memory=False), + ), + pipeline=PipelineSelection( + workload_type="i2v", + # FastVideo reads the trained fused-block partition and attention + # settings from the composed directory's fastvideo_inference.json. + components=ComponentConfig(override_pipeline_cls_name="MiniMaxH3Ref2VAModularPipeline"), + ), + ) + + +def main() -> None: + args = parse_args() + output_dir = Path(args.output) + output_dir.mkdir(parents=True, exist_ok=True) + model_dir, contract = resolve_model(args) + print(f"Composed model directory: {model_dir}") + print(f"Contract: {contract['num_inference_steps']} fused blocks of a {contract['pdd_steps']}-interval grid, " + f"{contract['attention_backend']} with sparsity {contract['vsa_sparsity']}, " + f"{contract['vsa_tile_size']}-token tiles, reference keep rate {contract['vsa_ref_keep_rate']}") + references = [MiniMaxH3Reference(source=path, media_type=kind) for kind, path in args.references] + + generator = VideoGenerator.from_config(build_generator_config(model_dir, args.num_gpus)) + try: + result = generator.generate( + GenerationRequest( + prompt=args.prompt, + negative_prompt="", + inputs=InputConfig(references=references), + sampling=SamplingConfig( + height=args.height, + width=args.width, + num_frames=args.num_frames, + fps=24, + # PDD counts fused blocks: one transformer forward each. + num_inference_steps=contract["num_inference_steps"], + guidance_scale=1.0, + batch_cfg=False, + seed=args.seed, + ), + output=OutputConfig( + output_path=str(output_dir / "fasth3_omniref_pdd.mp4"), + save_video=True, + return_frames=False, + ), + )) + print(f"Output written to: {result.video_path}") + if result.generation_time is not None: + print(f"Generation time: {result.generation_time:.1f} s") + if result.peak_memory_mb is not None: + print(f"Peak GPU memory allocated (rank 0): {result.peak_memory_mb / 1024:.1f} GiB") + finally: + generator.shutdown() + + +if __name__ == "__main__": + main() diff --git a/fastvideo/attention/backends/video_sparse_attn_h3.py b/fastvideo/attention/backends/video_sparse_attn_h3.py index a6f195ece..65401519b 100644 --- a/fastvideo/attention/backends/video_sparse_attn_h3.py +++ b/fastvideo/attention/backends/video_sparse_attn_h3.py @@ -7,8 +7,18 @@ backend differs from the Wan-tuned ``video_sparse_attn``: - Tiles are ``[segment-pure prefix chunks] + [3D video tiles]``; prefix tiles never straddle segment boundaries. The tile size is selectable at - metadata build time: 256 tokens ``(4,8,8)`` (default) or 64 tokens - ``(4,4,4)`` (see ``VSA_H3_TILE_SHAPES``). + metadata build time: 256 tokens ``(4,8,8)`` (default), 128 tokens + ``(4,4,8)``, or 64 tokens ``(4,4,4)`` (see ``VSA_H3_TILE_SHAPES``). +- The builder takes the packed sequence as ordered segments: dense row + counts (text, audio, image references) and one 3-D sparse region per + video, wherever it sits. The generated video is the last region; a Ref2VA + PDD run adds one earlier region per reference video. Rows are permuted to + ``[dense chunks][region 0 tiles][region 1 tiles]...`` and + ``untile_combined_index`` inverts the permutation, so the kernels see one + dense prefix followed by video tiles. Every video query keeps its own + top-k of EACH region (``video_tile_spans`` / ``span_sparsities``); the + reference-video regions may use a different keep rate than the generated + video. - Selection is pure Python on pooled tile scores; the block-sparse kernel consumes an explicit bool mask, so no kernel changes are needed. - The compression branch is gated by ``to_gate_compress``, which the base @@ -32,15 +42,21 @@ backward run the Triton block-sparse kernels directly (no expansion, ``FASTVIDEO_VSA_CUTEDSL`` does not apply). A third, opt-in route exists for the tile-64 FORWARD only: ``FASTVIDEO_VSA_SM100A=1`` sends no-grad forwards through the data-center Blackwell CUDA block-sparse kernel -(``fastvideo_kernel.block_sparse_attn_sm100a``, upstream PR #1719 plus -our per-q-tile ``q2k_num`` fix) when the extension is built, the device +(``fastvideo_kernel.block_sparse_attn_sm100a``, which reads a separate +``q2k_num`` key-tile count per query tile) when the extension is built, the device is sm_100 or sm_103, and the geometry qualifies. The CUDA kernel assigns adjacent pairs of query tiles to CTAs, so an odd logical tile count receives one internal, zero-valid partner tile for the no-grad call only. Score search, the trained mask, gate-compress, and the returned packed sequence remain on the original logical tiles. Grad-tracking forwards and every backward stay -on Triton unchanged. If the env is set but a precondition fails, the route +on the Triton kernels. If the env is set but a precondition fails, the route logs one warning and falls back. + +Tile 128 has exactly one implementation: the same sm_100a/sm_103a CUDA +forward, which carries a 128-token block instantiation. It needs no opt-in, +pads an odd tile count with the same zero-valid partner tile, runs no-grad +forwards only, and fails closed (no Triton fallback) when the extension or +device cannot run it. """ import functools @@ -61,7 +77,7 @@ except ImportError: try: # Optional: only present in fastvideo_kernel builds that carry the - # sm_100a/sm_103a CUDA block-sparse forward (upstream PR #1719). The module itself imports + # sm_100a/sm_103a CUDA block-sparse forward. The module itself imports # fine without the compiled symbols (`_HAS_VSA_SM100A` is then False and # `is_supported` says no), so this only guards *module* availability. from fastvideo_kernel import block_sparse_attn_sm100a as _sm100a @@ -91,7 +107,12 @@ _TILE_ELEMS = math.prod(VSA_H3_TILE_SIZE) VSA_H3_TILE_SHAPES: dict[int, tuple[int, int, int]] = { _TILE_ELEMS: VSA_H3_TILE_SIZE, 64: (4, 4, 4), + # sm_100a/sm_103a CUDA forward only (128-token block instantiation). + 128: (4, 4, 8), } +# Tile sizes the sm_100a/sm_103a CUDA kernel runs natively; its CTAs own +# adjacent pairs of query tiles, so an odd count gets one zero-valid partner. +_SM100A_TILE_ELEMS = (64, 128) @torch.library.custom_op( @@ -168,15 +189,17 @@ def token_tile_and_valid(variable_block_sizes: torch.Tensor, return token_tile, token_valid -def _validate_h3_tile_geometry( - prefix_segments: tuple[int, ...], - dit_seq_shape: tuple[int, int, int], +def _validate_h3_segment_geometry( + segments: tuple[int | tuple[int, int, int], ...], variable_block_sizes: torch.Tensor, untile_combined_index: torch.Tensor, tile_elems: int = _TILE_ELEMS, ) -> None: """Fail synchronously on out-of-bounds tile geometry. + ``segments`` is the packed-order interleave: an ``int`` is a dense + (prefix) segment's row count, a triple is one video region's token grid. + Invariants the block-sparse kernel trusts without checking: every tile's valid size is in (0, tile_elems]; the sizes sum to the packed sequence length; and ``untile_combined_index`` maps each packed @@ -185,82 +208,99 @@ def _validate_h3_tile_geometry( collective (e.g. an FSDP all-gather), which is unattributable — so raise here, once per cached geometry, with the numbers in hand. """ - total = sum(prefix_segments) + math.prod(dit_seq_shape) + total = sum(math.prod(segment) if isinstance(segment, tuple) else segment for segment in segments) n_pad = variable_block_sizes.numel() * tile_elems sizes_min = int(variable_block_sizes.min()) sizes_max = int(variable_block_sizes.max()) sizes_sum = int(variable_block_sizes.sum()) if sizes_min < 1 or sizes_max > tile_elems or sizes_sum != total: - raise ValueError(f"VSA-H3 tile sizes out of bounds for prefix={prefix_segments}, video={dit_seq_shape}, " + raise ValueError(f"VSA-H3 tile sizes out of bounds for segments={segments}, " f"tile_elems={tile_elems}: min={sizes_min}, max={sizes_max}, sum={sizes_sum}, " f"expected sum={total}.") if untile_combined_index.numel() != total: raise ValueError(f"VSA-H3 untile index has {untile_combined_index.numel()} entries for a packed " - f"sequence of {total} rows (prefix={prefix_segments}, video={dit_seq_shape}).") + f"sequence of {total} rows (segments={segments}).") idx_min = int(untile_combined_index.min()) idx_max = int(untile_combined_index.max()) if idx_min < 0 or idx_max >= n_pad: # Range first: the pad-slot gather below would itself index out of # bounds (the very async fault this guard exists to preempt). raise ValueError(f"VSA-H3 untile index is not an injective map into non-pad slots: range " - f"[{idx_min}, {idx_max}] vs padded length {n_pad} " - f"(prefix={prefix_segments}, video={dit_seq_shape}).") + f"[{idx_min}, {idx_max}] vs padded length {n_pad} (segments={segments}).") in_tile_offset = untile_combined_index % tile_elems maps_into_pad = bool((in_tile_offset >= variable_block_sizes[untile_combined_index // tile_elems]).any()) if maps_into_pad or int(torch.unique(untile_combined_index).numel()) != total: raise ValueError(f"VSA-H3 untile index is not an injective map into non-pad slots: " - f"pad-slot hit={maps_into_pad} " - f"(prefix={prefix_segments}, video={dit_seq_shape}).") + f"pad-slot hit={maps_into_pad} (segments={segments}).") @functools.lru_cache(maxsize=10) -def _h3_tile_geometry( - prefix_segments: tuple[int, ...], - dit_seq_shape: tuple[int, int, int], +def _h3_segment_tile_geometry( + segments: tuple[int | tuple[int, int, int], ...], device: torch.device, tile_shape: tuple[int, int, int] = VSA_H3_TILE_SIZE, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int]: - """Tile the packed sequence: segment-pure prefix chunks, then video tiles. +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int, tuple[tuple[int, int], ...]]: + """Tile a packed sequence of interleaved dense segments and video regions. + + ``segments`` lists, in packed order, dense segment row counts (``int``) + and video-region token grids (``(t, h, w)`` triples). Rows are permuted so + every dense chunk lands in the leading prefix tiles and each video region + contributes one contiguous run of 3-D tiles, wherever the dense segments + sit in the packed sequence. Dense chunks never straddle a segment + boundary. Returns (tile_partition_indices, variable_block_sizes, - untile_combined_index, num_prefix_tiles, num_video_tiles). + untile_combined_index, num_prefix_tiles, num_video_tiles, + video_tile_spans), where ``video_tile_spans`` holds one ``[start, end)`` + tile range per video region (prefix tiles included in the numbering). """ tile_elems = math.prod(tile_shape) - prefix_len = sum(prefix_segments) prefix_sizes: list[int] = [] - for segment in prefix_segments: + prefix_rows: list[torch.Tensor] = [] + video_entries: list[tuple[int, tuple[int, int, int]]] = [] + cursor = 0 + for segment in segments: + if isinstance(segment, tuple): + video_entries.append((cursor, segment)) + cursor += math.prod(segment) + continue full, rem = divmod(segment, tile_elems) prefix_sizes.extend([tile_elems] * full) if rem: prefix_sizes.append(rem) + prefix_rows.append(torch.arange(cursor, cursor + segment, device=device, dtype=torch.long)) + cursor += segment num_prefix_tiles = len(prefix_sizes) ts_t, ts_h, ts_w = tile_shape - t, h, w = dit_seq_shape - num_tiles = (math.ceil(t / ts_t), math.ceil(h / ts_h), math.ceil(w / ts_w)) - video_sizes = construct_variable_block_sizes(dit_seq_shape, num_tiles, device, tile_shape) - num_video_tiles = int(video_sizes.numel()) + video_index_parts: list[torch.Tensor] = [] + video_size_parts: list[torch.Tensor] = [] + video_tile_spans: list[tuple[int, int]] = [] + tile_cursor = num_prefix_tiles + for start, grid in video_entries: + t, h, w = grid + num_tiles = (math.ceil(t / ts_t), math.ceil(h / ts_h), math.ceil(w / ts_w)) + region_sizes = construct_variable_block_sizes(grid, num_tiles, device, tile_shape) + video_index_parts.append(get_tile_partition_indices(grid, tile_shape, device) + start) + video_size_parts.append(region_sizes) + video_tile_spans.append((tile_cursor, tile_cursor + int(region_sizes.numel()))) + tile_cursor += int(region_sizes.numel()) + num_video_tiles = tile_cursor - num_prefix_tiles - video_indices = get_tile_partition_indices(dit_seq_shape, tile_shape, device) + prefix_len - tile_partition_indices = torch.cat([ - torch.arange(prefix_len, device=device, dtype=torch.long), - video_indices, - ]) + tile_partition_indices = torch.cat(prefix_rows + video_index_parts) # cat promotes the int32 helper output to int64 alongside the prefix sizes - variable_block_sizes = torch.cat([ - torch.tensor(prefix_sizes, dtype=torch.long, device=device), - video_sizes, - ]) + variable_block_sizes = torch.cat([torch.tensor(prefix_sizes, dtype=torch.long, device=device)] + video_size_parts) # get_non_pad_index is lru-cached on tensor identity; variable_block_sizes # is itself cached by this function, so the identity stays stable. non_pad_index = get_non_pad_index(variable_block_sizes, tile_elems) untile_combined_index = non_pad_index[torch.argsort(tile_partition_indices)] - # One-time (lru-cached) synchronous bounds check; see _validate_h3_tile_geometry. - _validate_h3_tile_geometry(prefix_segments, dit_seq_shape, variable_block_sizes, untile_combined_index, tile_elems) - return (tile_partition_indices, variable_block_sizes, untile_combined_index, num_prefix_tiles, num_video_tiles) + # One-time (lru-cached) synchronous bounds check; see _validate_h3_segment_geometry. + _validate_h3_segment_geometry(segments, variable_block_sizes, untile_combined_index, tile_elems) + return (tile_partition_indices, variable_block_sizes, untile_combined_index, num_prefix_tiles, num_video_tiles, + tuple(video_tile_spans)) class MiniMaxH3VSABackend(AttentionBackend): @@ -308,8 +348,13 @@ class MiniMaxH3VSAMetadata(AttentionMetadata): # 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) + # One [start, end) tile range per sparse video region, in packed order + # with the generated video last, and the sparsity of that region's top-k. + video_tile_spans: tuple[tuple[int, int], ...] + span_sparsities: tuple[float, ...] + # tokens per tile (256, 128 or 64); selects the tile geometry AND the + # kernel route in forward() (256 -> VSA-256 CuTe/Triton, 128 -> sm_100a + # CUDA, 64 -> native Triton or opt-in sm_100a CUDA) tile_elems: int = _TILE_ELEMS # layers forced dense regardless of sparsity (probe-guided opt-outs) dense_layers: tuple[int, ...] = () @@ -330,26 +375,61 @@ class MiniMaxH3VSAMetadataBuilder(AttentionMetadataBuilder): def build( # type: ignore self, current_timestep: int, - raw_latent_shape: tuple[int, int, int], patch_size: tuple[int, int, int], VSA_sparsity: float, - prefix_segments: tuple[int, ...], + packed_segments: tuple[int | tuple[int, int, int], ...], device: torch.device, exempt: bool = True, dense_layers: tuple[int, ...] = (), tile_size: int = _TILE_ELEMS, + ref_keep_rate: float | None = None, **kwargs: dict[str, Any], ) -> MiniMaxH3VSAMetadata: + """Build per-step metadata for one packed H3 sequence. + + ``packed_segments`` lists the sequence in packed order: an ``int`` is + a dense segment's row count, and a ``(t, h, w)`` triple is the raw + latent shape of one sparse video region. The last region is the + generated video and follows ``VSA_sparsity``; every earlier region is + a reference video and keeps ``ref_keep_rate`` of its tiles, which a + build with reference-video regions requires. ``exempt=False`` + (compete mode) supports only the generated-video region. + """ tile_shape = VSA_H3_TILE_SHAPES.get(int(tile_size)) if tile_shape is None: raise ValueError(f"VSA-H3 tile_size must be one of {sorted(VSA_H3_TILE_SHAPES)}, got {tile_size!r}") - dit_seq_shape = (raw_latent_shape[0] // patch_size[0], raw_latent_shape[1] // patch_size[1], - raw_latent_shape[2] // patch_size[2]) - prefix_segments = tuple(int(s) for s in prefix_segments if s > 0) - total_seq_length = sum(prefix_segments) + math.prod(dit_seq_shape) + # Video regions become token grids under the patch size; empty dense segments are dropped. + token_segments: list[int | tuple[int, int, int]] = [] + for segment in packed_segments: + if isinstance(segment, tuple): + if any(int(v) <= 0 or int(v) % p for v, p in zip(segment, patch_size, strict=True)): + raise ValueError(f"VSA-H3 video region latent shape {tuple(segment)} is not a positive multiple " + f"of patch {tuple(patch_size)}.") + t, h, w = (int(v) // p for v, p in zip(segment, patch_size, strict=True)) + token_segments.append((t, h, w)) + elif segment > 0: + token_segments.append(int(segment)) + num_regions = sum(isinstance(segment, tuple) for segment in token_segments) + if num_regions == 0: + raise ValueError("VSA-H3 needs at least one video region.") + reference_sparsities: tuple[float, ...] = () + if num_regions > 1: + if ref_keep_rate is None: + raise ValueError(f"A VSA-H3 build with {num_regions - 1} reference-video region(s) needs " + "ref_keep_rate.") + if not exempt: + raise ValueError(f"vsa_mode='compete' supports only the generated-video region; this build has " + f"{num_regions - 1} reference-video region(s). Use vsa_mode='exempt'.") + reference_sparsities = (1.0 - float(ref_keep_rate), ) * (num_regions - 1) + total_seq_length = sum(math.prod(s) if isinstance(s, tuple) else s for s in token_segments) - (_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) + (_tile_partition_indices, variable_block_sizes, untile_combined_index, num_prefix_tiles, num_video_tiles, + video_tile_spans) = _h3_segment_tile_geometry(tuple(token_segments), device, tile_shape) + + # A dense build (sparsity 0: dense steps) keeps every region dense, + # whatever the reference keep rate. + VSA_sparsity = float(VSA_sparsity) + span_sparsities = (0.0, ) * num_regions if VSA_sparsity <= 0.0 else (*reference_sparsities, VSA_sparsity) dense_layers = tuple(int(layer) for layer in dense_layers) return MiniMaxH3VSAMetadata( @@ -361,9 +441,11 @@ class MiniMaxH3VSAMetadataBuilder(AttentionMetadataBuilder): exempt=exempt, variable_block_sizes=variable_block_sizes, untile_combined_index=untile_combined_index, + dense_layers_tensor=torch.tensor(dense_layers, device=device, dtype=torch.int64), + video_tile_spans=video_tile_spans, + span_sparsities=span_sparsities, tile_elems=int(tile_size), dense_layers=dense_layers, - dense_layers_tensor=torch.tensor(dense_layers, device=device, dtype=torch.int64), tile_buf_holder=self._tile_buf_holder, ) @@ -386,23 +468,46 @@ def _pool_tiles(x: torch.Tensor, variable_block_sizes: torch.Tensor, tile_elems: def _build_block_mask( scores: torch.Tensor, num_prefix_tiles: int, - num_video_tiles: int, VSA_sparsity: float, exempt: bool, + video_tile_spans: tuple[tuple[int, int], ...], + span_sparsities: tuple[float, ...], ) -> torch.Tensor: - """scores: [B, H, n_tiles, n_tiles] -> bool mask, same shape.""" + """scores: [B, H, n_tiles, n_tiles] -> bool mask, same shape. + + Prefix rows (non-video queries) stay dense. With ``exempt``, every query + keeps every prefix column, and every video query keeps its own top-k of + EACH video region, with the budget from that region's ``span_sparsities`` + entry. Without ``exempt`` (compete mode, generated-video region only), + prefix and video columns compete for one FLOP-matched top-k. + ``VSA_sparsity <= 0`` (dense steps and layers) keeps every tile. + ``video_tile_spans`` come from the cached geometry, which lays them out + contiguously after the prefix tiles, with one ``span_sparsities`` entry + each. + """ n_tiles = scores.shape[-1] - k_vid = compute_topk(VSA_sparsity, num_video_tiles) - if k_vid == num_video_tiles: + if VSA_sparsity <= 0.0: + span_sparsities = tuple(VSA_sparsity for _ in video_tile_spans) + span_topk = [ + compute_topk(sparsity, end - start) + for (start, end), sparsity in zip(video_tile_spans, span_sparsities, strict=True) + ] + if all(k == end - start for k, (start, end) in zip(span_topk, video_tile_spans, strict=True)): return torch.ones_like(scores, dtype=torch.bool) mask = torch.zeros_like(scores, dtype=torch.bool) if exempt or num_prefix_tiles == 0: - video_cols = scores[..., num_prefix_tiles:] - idx = video_cols.topk(k_vid, dim=-1).indices + num_prefix_tiles - mask.scatter_(-1, idx, True) + for (start, end), k_span in zip(video_tile_spans, span_topk, strict=True): + if k_span == end - start: + mask[..., start:end] = True + continue + idx = scores[..., start:end].topk(k_span, dim=-1).indices + start + mask.scatter_(-1, idx, True) mask[..., :num_prefix_tiles] = True else: - k_total = min(k_vid + num_prefix_tiles, n_tiles) + # compete: prefix keys enter the top-k under a FLOP-matched budget, + # which is defined for the generated-video region alone (the builder + # rejects compete mode with reference-video regions) + k_total = min(span_topk[0] + num_prefix_tiles, n_tiles) idx = scores.topk(k_total, dim=-1).indices mask.scatter_(-1, idx, True) mask[:, :, :num_prefix_tiles, :] = True @@ -452,6 +557,12 @@ class MiniMaxH3VSAImpl(AttentionImpl): # request-time env/probe/fallback behavior; only Dynamo capture reads # the prepared, static route. self._regional_compile_sm100a_enabled: bool | None = None + # Tile-128 route per (device, dtype, head size): None when the CUDA + # kernel can run it, else why not. The kernel's predicate otherwise + # depends only on the tile-128 buffer contract (contiguous BHSD, + # 128-token blocks, an even tile count with the partner tile, integer + # tile sizes), which every call meets, so it is evaluated once. + self._tile128_route: dict[tuple[torch.device, torch.dtype, int], str | None] = {} def prepare_for_compile(self, device: torch.device) -> None: """Tensorize per-layer state shared by every torch.compile route.""" @@ -511,9 +622,10 @@ class MiniMaxH3VSAImpl(AttentionImpl): buffer; callers must consume it before the next ``tile()`` (both call sites in ``forward()`` read it immediately). A grad-tracking forward instead receives a fresh buffer and leaves the holder untouched, so - the builder never retains autograd state across steps. Odd tile-64 - no-grad sm100a requests carry one additional all-zero tile internally; - metadata and all observable outputs retain the logical geometry. + the builder never retains autograd state across steps. Odd no-grad + sm100a requests (tile 128 always, tile 64 when opted in) carry one + additional all-zero tile internally; metadata and all observable + outputs retain the logical geometry. """ if x.shape[1] != attn_metadata.total_seq_length: raise ValueError(f"VSA-H3 metadata was built for sequence length {attn_metadata.total_seq_length}, " @@ -530,7 +642,9 @@ class MiniMaxH3VSAImpl(AttentionImpl): sm100a_requested = False else: sm100a_requested = envs.FASTVIDEO_VSA_SM100A.get() - needs_sm100a_pair = (attn_metadata.tile_elems == 64 and n_tiles % 2 != 0 and not grad_mode and sm100a_requested) + # Tile 128 has no route other than the sm100a CUDA kernel. + sm100a_route = attn_metadata.tile_elems == 128 or (attn_metadata.tile_elems == 64 and sm100a_requested) + needs_sm100a_pair = n_tiles % 2 != 0 and not grad_mode and sm100a_route kernel_tiles = n_tiles + int(needs_sm100a_pair) target_shape = (x.shape[0], kernel_tiles * attn_metadata.tile_elems, x.shape[-2], x.shape[-1]) @@ -584,11 +698,11 @@ class MiniMaxH3VSAImpl(AttentionImpl): tile_elems = attn_metadata.tile_elems if regional_compiling and tile_elems != 64: raise RuntimeError("VSA-H3 regional fullgraph compile requires 64-token tiles; disable " - "inference_torch_compile for tile-256/CuTe runs.") + "inference_torch_compile for tile-128/256 runs.") if tile_elems == 64: if block_sparse_attn_64_bhsd is None: raise NotImplementedError("fastvideo_kernel.block_sparse_attn is not installed") - elif block_sparse_attn_256_bshd is None: + elif tile_elems == 256 and block_sparse_attn_256_bshd is None: raise NotImplementedError("fastvideo_kernel.block_sparse_attn_256 is not installed") # Probe recording performs filesystem writes and host synchronizations, @@ -596,12 +710,12 @@ class MiniMaxH3VSAImpl(AttentionImpl): # environment switch while Dynamo captures a regional full graph. # The metadata always describes the trained logical geometry. # ``tile()`` may append exactly one transport-only partner for an odd - # tile-64 sm100a call. Keep score selection and the gate branch on the - # logical prefix, and reject every other shape before a kernel sees it. + # tile-64/128 sm100a call. Keep score selection and the gate branch on + # the logical prefix, and reject every other shape before a kernel sees it. n_tiles = attn_metadata.variable_block_sizes.numel() logical_seq_len = n_tiles * tile_elems pair_pad_seq_len = logical_seq_len + tile_elems - pair_pad_is_valid = tile_elems == 64 and n_tiles % 2 != 0 + pair_pad_is_valid = tile_elems in _SM100A_TILE_ELEMS and n_tiles % 2 != 0 allowed_seq_lengths = (logical_seq_len, pair_pad_seq_len) if pair_pad_is_valid else (logical_seq_len, ) if query.shape[1] not in allowed_seq_lengths: expected = (f"the logical length {logical_seq_len} or one sm100a partner tile " @@ -649,18 +763,19 @@ class MiniMaxH3VSAImpl(AttentionImpl): mask = _build_block_mask( scores, attn_metadata.num_prefix_tiles, - attn_metadata.num_video_tiles, layer_sparsity, attn_metadata.exempt, + attn_metadata.video_tile_spans, + attn_metadata.span_sparsities, ) 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' - # granularity. Both 64-token entries take BHSD ([B, H, S_pad, D]); + if tile_elems in _SM100A_TILE_ELEMS: + # Native 64/128-token path: the block map is already at the + # kernels' granularity. These entries take BHSD ([B, H, S_pad, D]); # mirror block_sparse_attn_256_bshd's Triton branch and transpose # around the call. q_bhsd = query.transpose(1, 2).contiguous() @@ -680,14 +795,30 @@ class MiniMaxH3VSAImpl(AttentionImpl): value=0, ) - # Opt-in sm_100a CUDA forward (upstream PR #1719 + per-q-tile - # q2k_num fix). Forward-only: grad-tracking calls stay on Triton + # sm_100a CUDA forward, opt-in at tile 64 and the only route at + # tile 128. Forward-only: tile-64 grad-tracking calls stay on Triton # 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 regional_compiling: + if tile_elems == 128: + # The CUDA kernel's 128-token instantiation is the only tile-128 + # implementation, so it needs no opt-in and never falls back. + if grad_mode: + raise NotImplementedError("VSA-H3 tile 128 supports no-grad inference forwards only.") + route_key = (q_bhsd.device, q_bhsd.dtype, int(q_bhsd.shape[-1])) + if route_key not in self._tile128_route: + reason = _sm100a_unavailable_reason(_sm100a, q_bhsd, sm100a_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" + self._tile128_route[route_key] = reason + reason = self._tile128_route[route_key] + if reason is not None: + raise RuntimeError(f"VSA-H3 tile 128 requires the sm_100a/sm_103a CUDA block-sparse kernel: " + f"{reason}. Use VSA_tile_size=64 or 256 elsewhere.") + use_sm100a = True + elif regional_compiling: if self._regional_compile_sm100a_enabled is None: raise RuntimeError( "VSA-H3 sm_100a routing was not resolved before torch.compile; " @@ -712,11 +843,12 @@ class MiniMaxH3VSAImpl(AttentionImpl): logger.warning_once(f"{VSA_SM100A_ENV}=1 but falling back to the Triton-64 kernels: {reason}") if use_sm100a: - # Regional preparation emits the compile-route receipt before + # prepare_for_regional_compile logs the selected mask route before # capture. Logging from this branch would itself break a # ``fullgraph=True`` forward. if not compiling: - logger.info_once("MiniMax-H3 VSA tile-64 forward: using the sm100a/sm103a CUDA block-sparse kernel") + logger.info_once( + f"MiniMax-H3 VSA tile-{tile_elems} forward: using the sm100a/sm103a CUDA block-sparse kernel") if regional_compiling: # The compile-safe wrapper keeps both Triton mask # compaction and the raw pybind launch behind one @@ -732,7 +864,7 @@ class MiniMaxH3VSAImpl(AttentionImpl): # 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. + # run prefix+top-k); the kernel reads each query tile's own count. q2k_idx, q2k_num = map_to_index(sm100a_mask) out_bhsd, _ = _sm100a.block_sparse_attn_sm100a( q_bhsd, diff --git a/fastvideo/configs/models/dits/minimax_h3.py b/fastvideo/configs/models/dits/minimax_h3.py index 464f75a69..b67f2993f 100644 --- a/fastvideo/configs/models/dits/minimax_h3.py +++ b/fastvideo/configs/models/dits/minimax_h3.py @@ -6,6 +6,7 @@ from __future__ import annotations from dataclasses import dataclass, field from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig +from fastvideo.layers.pdd import validate_pdd_steps from fastvideo.platforms import AttentionBackendEnum @@ -59,6 +60,11 @@ class MiniMaxH3ArchConfig(DiTArchConfig): norm_eps: float = 1e-5 qk_norm_eps: float = 1e-5 final_norm_eps: float = 1e-5 + # Parallel Decoding Distillation (PDD) students widen both output + # projections to this many fine-grid heads (``fastvideo.layers.pdd``). + # ``None`` is the released single-head transformer. A PDD export records + # the key in its transformer ``config.json``. + pdd_steps: int | None = None def __post_init__(self) -> None: super().__post_init__() @@ -67,6 +73,8 @@ class MiniMaxH3ArchConfig(DiTArchConfig): self.patch_size = (self.patch_size[0], self.patch_size[1], self.patch_size[2]) self.num_channels_latents = self.in_channels self.out_channels = self.in_channels + if self.pdd_steps is not None: + validate_pdd_steps(self.pdd_steps, "MiniMax H3 pdd_steps") if self.adaln_rank is not None and not 0 < self.adaln_rank <= self.time_embed_dim: raise ValueError(f"MiniMax H3 adaln_rank must be in (0, time_embed_dim={self.time_embed_dim}], " f"got {self.adaln_rank}.") diff --git a/fastvideo/configs/pipelines/base.py b/fastvideo/configs/pipelines/base.py index 271d06763..17c29bae3 100644 --- a/fastvideo/configs/pipelines/base.py +++ b/fastvideo/configs/pipelines/base.py @@ -2,7 +2,7 @@ import json from collections.abc import Callable from dataclasses import asdict, dataclass, field, fields -from typing import Any, cast +from typing import TYPE_CHECKING, Any, cast import torch @@ -12,6 +12,11 @@ from fastvideo.configs.utils import update_config_from_args from fastvideo.logger import init_logger from fastvideo.utils import FlexibleArgumentParser, StoreBoolean, shallow_asdict +if TYPE_CHECKING: + from fastvideo.api.sampling_param import SamplingParam + from fastvideo.api.schema import GenerationRequest + from fastvideo.fastvideo_args import FastVideoArgs + logger = init_logger(__name__) @@ -281,6 +286,25 @@ class PipelineConfig: f"Length of text postprocess functions ({len(self.postprocess_text_funcs)}) must be equal to length of text preprocessing functions ({len(self.preprocess_text_funcs)})" ) + def resolve_checkpoint_settings(self, fastvideo_args: "FastVideoArgs") -> None: + """Apply the run settings that the checkpoint at ``fastvideo_args.model_path`` fixes. + + ``FastVideoArgs.__post_init__`` calls this once in the main process, before + argument validation and before any weights load. Most checkpoints fix no + run settings, so the base implementation does nothing. + """ + + def apply_request_constraints(self, request: "GenerationRequest", + sampling_param: "SamplingParam") -> "SamplingParam": + """Apply this model's constraints on a generation request to its sampling parameters. + + ``sampling_param`` holds the preset defaults with the request's explicit + fields applied. A model sets the values it fixes when the request leaves + them unset, and raises when the request sets a value it does not accept. + The base PipelineConfig accepts every request unchanged. + """ + return sampling_param + def dump_to_json(self, file_path: str): output_dict = shallow_asdict(self) del_keys = [] diff --git a/fastvideo/configs/pipelines/minimax_h3.py b/fastvideo/configs/pipelines/minimax_h3.py index d2a0d5f14..366028600 100644 --- a/fastvideo/configs/pipelines/minimax_h3.py +++ b/fastvideo/configs/pipelines/minimax_h3.py @@ -4,6 +4,7 @@ from __future__ import annotations from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any from fastvideo.configs.models import EncoderConfig, VAEConfig from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3Config @@ -11,6 +12,137 @@ from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import MiniMaxH3Qwen3 from fastvideo.configs.models.vaes.minimax_h3_audio import MiniMaxH3AudioVAEConfig from fastvideo.configs.models.vaes.minimax_h3_video import MiniMaxH3VideoVAEConfig from fastvideo.configs.pipelines.base import PipelineConfig +from fastvideo.logger import init_logger +from fastvideo.utils import read_optional_model_json + +if TYPE_CHECKING: + from fastvideo.api.sampling_param import SamplingParam + from fastvideo.api.schema import GenerationRequest + from fastvideo.fastvideo_args import FastVideoArgs + +logger = init_logger(__name__) + +# The checkpoint file that records how a distilled FastH3 export must be sampled. +FASTH3_INFERENCE_FILE = "fastvideo_inference.json" +FASTH3_INFERENCE_SCHEMA = "fasth3-inference-contract-v1" +# The fields a FastH3 Ref2VA PDD export writes to FASTH3_INFERENCE_FILE. Missing +# and unknown fields are both errors, so no field is silently ignored. +_PDD_INFERENCE_FIELDS = frozenset({ + "schema_version", + "schema", + "model_type", + "transformer_component", + "conditioning", + "base_model_revision", + "pdd_steps", + "pdd_step_indices", + "num_inference_steps", + "transformer_forwards", + "grid_max_t", + "video_scheduler_shift", + "audio_scheduler_shift", + "guidance_scale", + "attention_backend", + "vsa_sparsity", + "vsa_tile_size", + "vsa_ref_policy", + "vsa_ref_keep_rate", +}) +# Fields that repeat a value FastVideo fixes in code: the Ref2VA pipeline, its +# conditioning, sampling without classifier-free guidance, and the per-reference +# sparse regions that only VIDEO_SPARSE_ATTN_H3 builds. +_PDD_FIXED_FIELDS: dict[str, Any] = { + "schema_version": FASTH3_INFERENCE_SCHEMA, + "schema": FASTH3_INFERENCE_SCHEMA, + "model_type": "ref2va", + "transformer_component": "transformer_ref", + "conditioning": "fixed_ordered_references_target_only_flow", + "guidance_scale": 1.0, + "attention_backend": "VIDEO_SPARSE_ATTN_H3", + "vsa_ref_policy": "p2_multi_region", +} + +# The base snapshot an export was distilled against: "hf://@". +_BASE_MODEL_REVISION_PREFIX = "hf://" + + +def parse_base_model_revision(value: Any) -> tuple[str, str]: + """Split a contract's ``base_model_revision``, ``hf://@``, into repo id and revision.""" + if isinstance(value, str) and value.startswith(_BASE_MODEL_REVISION_PREFIX): + repo, separator, revision = value[len(_BASE_MODEL_REVISION_PREFIX):].partition("@") + if separator and repo.strip() and revision.strip() and "@" not in revision: + return repo, revision + raise ValueError(f"FastH3 base_model_revision={value!r} must be hf://@.") + + +def _require_fraction(name: str, value: Any, *, allow_zero: bool) -> None: + """Raise unless *value* is a real number in [0, 1) (``allow_zero``) or in (0, 1).""" + is_real = isinstance(value, int | float) and not isinstance(value, bool) + if not (is_real and (value >= 0.0 if allow_zero else value > 0.0) and value < 1.0): + raise ValueError(f"{name} must be in {'[0, 1)' if allow_zero else '(0, 1)'}, got {value!r}.") + + +def _read_pdd_inference_file(model_path: str, revision: str | None) -> dict[str, Any] | None: + """Read and validate a FastH3 Ref2VA PDD checkpoint's fastvideo_inference.json; None for other checkpoints. + + The file must carry exactly the PDD fields. Fields that FastVideo fixes in + code must hold that value, and fields that repeat a value stored in another + checkpoint file must equal that file: ``pdd_steps`` the transformer config, + the shifts the scheduler configs, and the step counts the partition, which + must increase strictly from 0 to ``pdd_steps``. + """ + inference_file = read_optional_model_json(model_path, FASTH3_INFERENCE_FILE, revision) + if inference_file is None or "pdd_steps" not in inference_file: + return None + source = f"{model_path}/{FASTH3_INFERENCE_FILE}" + missing = sorted(_PDD_INFERENCE_FIELDS - inference_file.keys()) + unknown = sorted(inference_file.keys() - _PDD_INFERENCE_FIELDS) + if missing or unknown: + raise ValueError(f"{source} must carry exactly the FastH3 PDD fields; missing {missing}, unknown {unknown}.") + for key, expected in _PDD_FIXED_FIELDS.items(): + if inference_file[key] != expected: + raise ValueError(f"{source} {key}={inference_file[key]!r} is unsupported; expected {expected!r}.") + parse_base_model_revision(inference_file["base_model_revision"]) + from fastvideo.layers.pdd import PDD_GRID_MAX_T + if inference_file["grid_max_t"] != PDD_GRID_MAX_T: + raise ValueError(f"{source} grid_max_t={inference_file['grid_max_t']!r} is unsupported; the PDD fine grid " + f"ends at {PDD_GRID_MAX_T}.") + + owner_files = {} + for owner in ("transformer_ref/config.json", "scheduler/scheduler_config.json", + "audio_scheduler/scheduler_config.json"): + owner_files[owner] = read_optional_model_json(model_path, owner, revision) + if owner_files[owner] is None: + raise ValueError(f"{model_path} has {FASTH3_INFERENCE_FILE} but no {owner}.") + pdd_steps = owner_files["transformer_ref/config.json"].get("pdd_steps") + if inference_file["pdd_steps"] != pdd_steps: + raise ValueError(f"{source} pdd_steps={inference_file['pdd_steps']!r} disagrees with " + f"transformer_ref/config.json pdd_steps={pdd_steps!r}.") + for key, owner in (("video_scheduler_shift", "scheduler/scheduler_config.json"), + ("audio_scheduler_shift", "audio_scheduler/scheduler_config.json")): + shift = owner_files[owner].get("shift") + if inference_file[key] != shift: + raise ValueError(f"{source} {key}={inference_file[key]!r} disagrees with {owner} shift={shift!r}.") + indices = inference_file["pdd_step_indices"] + if (not isinstance(indices, list) or len(indices) < 2 or any(type(index) is not int for index in indices) + or indices[0] != 0 or indices[-1] != pdd_steps + or any(left >= right for left, right in zip(indices, indices[1:], strict=False))): + raise ValueError(f"{source} pdd_step_indices={indices!r} must increase strictly from 0 to " + f"pdd_steps={pdd_steps}.") + for key in ("num_inference_steps", "transformer_forwards"): + if inference_file[key] != len(indices) - 1: + raise ValueError(f"{source} {key}={inference_file[key]!r} disagrees with the {len(indices) - 1} blocks " + "of pdd_step_indices.") + + # Trained VSA settings. A run's override of a tunable one is checked where it is applied. + from fastvideo.attention.backends.video_sparse_attn_h3 import VSA_H3_TILE_SHAPES + tile_size = inference_file["vsa_tile_size"] + if type(tile_size) is not int or tile_size not in VSA_H3_TILE_SHAPES: + raise ValueError(f"{source} vsa_tile_size={tile_size!r} must be one of {sorted(VSA_H3_TILE_SHAPES)}.") + _require_fraction(f"{source} vsa_sparsity", inference_file["vsa_sparsity"], allow_zero=True) + # A keep rate of 1 would leave reference-video attention dense. + _require_fraction(f"{source} vsa_ref_keep_rate", inference_file["vsa_ref_keep_rate"], allow_zero=False) + return inference_file @dataclass @@ -28,6 +160,15 @@ class MiniMaxH3PipelineConfig(PipelineConfig): vae_precision: str = "fp32" text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", )) vae_sp: bool = False + # Parallel Decoding Distillation (PDD) students: fine-grid node indices of + # the fused blocks, one transformer forward each. resolve_checkpoint_settings + # sets it from the checkpoint's fastvideo_inference.json; None means the + # checkpoint is not a PDD student. + pdd_step_indices: tuple[int, ...] | None = None + # Fraction of each reference video's VSA tiles that every video query keeps, + # for PDD students. None takes the checkpoint's trained value; another value + # overrides it with a warning. + vsa_ref_keep_rate: float | None = None def check_pipeline_config(self) -> None: super().check_pipeline_config() @@ -35,5 +176,91 @@ class MiniMaxH3PipelineConfig(PipelineConfig): raise ValueError("MiniMax-H3 uses separate checkpoint-defined video/audio scheduler shifts; " "flow_shift must remain unset.") + def resolve_checkpoint_settings(self, fastvideo_args: FastVideoArgs) -> None: + """Apply a FastH3 Ref2VA PDD checkpoint's fastvideo_inference.json to this run. -__all__ = ["MiniMaxH3PipelineConfig"] + The file is the one source of the PDD partition and the trained VSA + settings, and a setting the run leaves unset takes the file's value. + The partition, attention backend and tile size are fixed by training, + so a different run value raises; VSA sparsity and the reference keep + rate may differ, with a warning. For checkpoints without PDD fields + (base MiniMax-H3, DMD exports) this method sets nothing. + """ + inference_file = _read_pdd_inference_file(fastvideo_args.model_path, fastvideo_args.revision) + if inference_file is None: + for name in ("pdd_step_indices", "vsa_ref_keep_rate"): + if getattr(self, name) is not None: + raise ValueError(f"{name} applies only to FastH3 PDD checkpoints, whose {FASTH3_INFERENCE_FILE} " + f"sets it; {fastvideo_args.model_path} has no PDD fields.") + return + + # Training-fixed settings: an unset run value takes the file's value; a different one raises. + from fastvideo.attention.selector import coerce_attn_backend + trained_backend = inference_file["attention_backend"] + if fastvideo_args.attention_backend is None: + fastvideo_args.attention_backend = trained_backend + elif coerce_attn_backend(fastvideo_args.attention_backend) != coerce_attn_backend(trained_backend): + raise ValueError(f"This FastH3 PDD checkpoint was trained with attention_backend={trained_backend}; " + f"this run requests {fastvideo_args.attention_backend} (argument or " + "FASTVIDEO_ATTENTION_BACKEND). Leave it unset or pass the trained value.") + trained_tile = inference_file["vsa_tile_size"] + if fastvideo_args.VSA_tile_size is None: + fastvideo_args.VSA_tile_size = trained_tile + elif fastvideo_args.VSA_tile_size != trained_tile: + raise ValueError(f"This FastH3 PDD checkpoint was trained with VSA_tile_size={trained_tile}; this run " + f"requests {fastvideo_args.VSA_tile_size}. Leave it unset or pass the trained value.") + indices = tuple(inference_file["pdd_step_indices"]) + if self.pdd_step_indices is not None and tuple(self.pdd_step_indices) != indices: + raise ValueError(f"pdd_step_indices comes from {FASTH3_INFERENCE_FILE} ({list(indices)}); this run sets " + f"{list(self.pdd_step_indices)}.") + if self.dmd_denoising_steps is not None: + raise ValueError(f"{fastvideo_args.model_path} is a FastH3 PDD checkpoint; dmd_denoising_steps must be " + "unset.") + self.pdd_step_indices = indices + + # Tunable settings: an unset run value takes the trained value; a different one wins, with a warning. + trained_sparsity = inference_file["vsa_sparsity"] + if fastvideo_args.VSA_sparsity is None: + fastvideo_args.VSA_sparsity = trained_sparsity + elif fastvideo_args.VSA_sparsity != trained_sparsity: + _require_fraction("VSA_sparsity", fastvideo_args.VSA_sparsity, allow_zero=True) + logger.warning("FastH3 PDD checkpoint was trained with VSA_sparsity=%s; this run uses %s.", + trained_sparsity, fastvideo_args.VSA_sparsity) + trained_keep_rate = inference_file["vsa_ref_keep_rate"] + if self.vsa_ref_keep_rate is None: + self.vsa_ref_keep_rate = trained_keep_rate + elif self.vsa_ref_keep_rate != trained_keep_rate: + _require_fraction("vsa_ref_keep_rate", self.vsa_ref_keep_rate, allow_zero=False) + logger.warning("FastH3 PDD checkpoint was trained with vsa_ref_keep_rate=%s; this run uses %s.", + trained_keep_rate, self.vsa_ref_keep_rate) + + logger.info( + "FastH3 PDD checkpoint %s: %d fused blocks %s of a %d-interval grid; attention_backend=%s, " + "VSA_tile_size=%s, VSA_sparsity=%s, vsa_ref_keep_rate=%s", fastvideo_args.model_path, + len(indices) - 1, list(indices), indices[-1], fastvideo_args.attention_backend, + fastvideo_args.VSA_tile_size, fastvideo_args.VSA_sparsity, self.vsa_ref_keep_rate) + + def apply_request_constraints(self, request: GenerationRequest, sampling_param: SamplingParam) -> SamplingParam: + """Fix ``num_inference_steps`` to the block count of a PDD checkpoint. + + A PDD checkpoint runs exactly one transformer forward per trained fused + block: a request that leaves ``num_inference_steps`` unset gets the block + count, and a request that sets another count raises. Checkpoints without + a PDD partition return ``sampling_param`` unchanged. + """ + if self.pdd_step_indices is None: + return sampling_param + from fastvideo.api.compat import explicit_request_updates + + num_blocks = len(self.pdd_step_indices) - 1 + requested_steps = explicit_request_updates(request).get("num_inference_steps", num_blocks) + if requested_steps != num_blocks: + raise ValueError(f"This FastH3 PDD checkpoint runs exactly {num_blocks} transformer forwards; the request " + f"sets num_inference_steps={requested_steps}. Pass num_inference_steps={num_blocks}. Only " + "a request parsed from a mapping or a config file can leave it unset; a " + "GenerationRequest built in Python counts every field as set.") + sampling_param.num_inference_steps = num_blocks + return sampling_param + + +__all__ = ["FASTH3_INFERENCE_FILE", "FASTH3_INFERENCE_SCHEMA", "MiniMaxH3PipelineConfig", "parse_base_model_revision"] diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 5bdc819b7..90d454aa4 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -458,9 +458,9 @@ class VideoGenerator: raise ValueError(f"Request field {key!r} is not supported by pipeline config overrides") setattr(fastvideo_args.pipeline_config, key, deepcopy(value)) - resolved_sampling_param = request_to_sampling_param( + resolved_sampling_param = fastvideo_args.pipeline_config.apply_request_constraints( request, - model_path=self.fastvideo_args.model_path, + request_to_sampling_param(request, model_path=self.fastvideo_args.model_path), ) return self._generate_video_impl( prompt=request.prompt, @@ -508,9 +508,9 @@ class VideoGenerator: raise ValueError(f"Request field {key!r} is not supported by pipeline config overrides") setattr(fastvideo_args.pipeline_config, key, deepcopy(value)) - sampling_param = request_to_sampling_param( + sampling_param = fastvideo_args.pipeline_config.apply_request_constraints( request, - model_path=self.fastvideo_args.model_path, + request_to_sampling_param(request, model_path=self.fastvideo_args.model_path), ) batch_extra = request_to_batch_extra(request) result = self._generate_video_impl( diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 8cb4ef920..d57a179a7 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -231,9 +231,10 @@ class FastVideoArgs: disable_autocast: bool = False - # VSA parameters - VSA_sparsity: float = 0.0 # inference/validation sparsity - VSA_tile_size: int = 256 # VSA-H3 tile size (256 or 64); 64 = native Triton path + # VSA parameters. None means unset: __post_init__ fills the checkpoint's + # trained value when the checkpoint fixes one, else 0.0 and 256. + VSA_sparsity: float | None = None # inference/validation sparsity + VSA_tile_size: int | None = None # VSA-H3 tokens per tile (256, 128, 64); 128 = sm_100a CUDA only # V-MoBA parameters moba_config_path: str | None = None @@ -361,6 +362,13 @@ class FastVideoArgs: 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() + # Runs after FASTVIDEO_ATTENTION_BACKEND is copied into attention_backend, + # so a backend chosen by that env var counts as the run's request. + self.pipeline_config.resolve_checkpoint_settings(self) + if self.VSA_sparsity is None: + self.VSA_sparsity = 0.0 + if self.VSA_tile_size is None: + self.VSA_tile_size = 256 import fastvideo.envs as envs envs.warn_deprecated_variables() self.check_fastvideo_args() @@ -800,13 +808,14 @@ class FastVideoArgs: "--VSA-sparsity", type=float, default=FastVideoArgs.VSA_sparsity, - help="Validation sparsity for VSA", + help="Validation sparsity for VSA (default: the checkpoint's trained value, else 0.0)", ) parser.add_argument( "--VSA-tile-size", type=int, default=FastVideoArgs.VSA_tile_size, - help="VSA-H3 tile size in tokens (256 or 64); 64 runs the native Triton block-sparse path", + help="VSA-H3 tile size in tokens (256, 128 or 64; default: the checkpoint's trained value, else 256); " + "64 runs the native Triton block-sparse path, 128 requires the sm_100a/sm_103a CUDA kernel", ) # Master port for distributed training/inference diff --git a/fastvideo/layers/pdd.py b/fastvideo/layers/pdd.py new file mode 100644 index 000000000..bcfb15fe4 --- /dev/null +++ b/fastvideo/layers/pdd.py @@ -0,0 +1,312 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Parallel Decoding Distillation (PDD) sampling primitives. + +A PDD student widens a diffusion transformer's final output projection so one +forward emits ``pdd_steps`` head predictions, one per interval of a fixed fine +time grid. Head ``j`` is the modality's rectified-flow direction frozen on +interval ``j``, so a block ``[start, end)`` of heads advances the state by +``sum_j (sigma(u_{j+1}) - sigma(u_j)) * q_j``. + +The sampler never materializes every head. For each runtime step it fuses one +block of head parameters inside the linear layer (:meth:`PDDReplicatedLinear.fuse`) +into their integration-weighted mean, and an ordinary Euler step over the +block's two node sigmas applies it: the block's total weight is exactly its +node-sigma increment. + +The fine grid lives on the base (unshifted) clock ``[0, 0.999]``; every +modality reaches its own noise level through its rational time shift. The +arithmetic matches the rectified-flow schedule the students are trained with. +""" + +from __future__ import annotations + +import contextlib +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn.functional as F + +from fastvideo.layers.linear import ReplicatedLinear, UnquantizedLinearMethod +from fastvideo.layers.quantization.base_config import QuantizationConfig + +# Rectified-flow students end their finite time domain at 0.999 and build the +# PDD fine grid on ``[0, 0.999]``. +PDD_GRID_MAX_T = 0.999 + + +def validate_finite_nonzero(value: torch.Tensor, name: str) -> None: + if not bool(torch.all(torch.isfinite(value) & (value != 0))): + raise ValueError(f"{name} must be finite and non-zero") + + +def validate_pdd_steps(value: Any, name: str = "pdd_steps") -> int: + """A PDD grid needs at least two intervals to have a block structure.""" + if isinstance(value, bool) or not isinstance(value, int) or value < 2: + raise ValueError(f"{name} must be an int >= 2, got {value!r}") + return value + + +# ---------------------------------------------------------------------------- +# Shifted rectified-flow arithmetic on the base clock +# ---------------------------------------------------------------------------- + + +def shifted_noise_amount(base: torch.Tensor, shift: float, *, max_t: float = PDD_GRID_MAX_T) -> torch.Tensor: + """Rational time shift ``f_s(u) = s * u * M / (u * (s - 1) + M)`` on ``[0, M]``. + + ``M = max_t`` is a fixed point of every shift, so the first grid node has + the same noise level in every modality. + """ + if shift == 1.0: + return base + return shift * base * max_t / (base * (shift - 1.0) + max_t) + + +def shifted_noise_delta( + base_start: torch.Tensor, + base_end: torch.Tensor, + shift: float, + *, + max_t: float = PDD_GRID_MAX_T, +) -> torch.Tensor: + """Return ``f_s(base_end) - f_s(base_start)`` without subtractive cancellation. + + The rational difference simplifies to + ``s * M^2 * (b - a) / (D(a) * D(b))`` with ``D(u) = M + (s - 1) * u``; the + identity shift keeps the direct difference bit-exact. + """ + if shift == 1.0: + return base_end - base_start + denominator_start = max_t + (shift - 1.0) * base_start + denominator_end = max_t + (shift - 1.0) * base_end + return shift * max_t**2 * (base_end - base_start) / (denominator_start * denominator_end) + + +@dataclass(frozen=True) +class PDDModalitySchedule: + """One modality's noise level as a function of the base clock: ``sigma(u) = f_shift(u)``.""" + + shift: float + max_t: float = PDD_GRID_MAX_T + + def sigma(self, base: torch.Tensor) -> torch.Tensor: + return shifted_noise_amount(base, self.shift, max_t=self.max_t) + + def integration_weights(self, fine_grid: torch.Tensor) -> torch.Tensor: + """``sigma(u_{j+1}) - sigma(u_j)`` for every fine-grid interval ``j``. + + Freezing head ``j`` on its interval advances the state by exactly this + coefficient. Negative for a descending grid. + """ + return shifted_noise_delta(fine_grid[:-1], fine_grid[1:], self.shift, max_t=self.max_t) + + +# ---------------------------------------------------------------------------- +# Fine grid and runtime partitions +# ---------------------------------------------------------------------------- + + +def pdd_fine_grid( + pdd_steps: int, + *, + max_t: float = PDD_GRID_MAX_T, + device: torch.device | str | None = None, +) -> torch.Tensor: + """The descending base-clock grid ``u_0 = max_t > ... > u_N = 0`` in float64.""" + grid = torch.linspace(max_t, 0.0, pdd_steps + 1, dtype=torch.float64, device=device) + return grid.clamp(max=max_t) + + +@dataclass(frozen=True) +class PDDSamplingPlan: + """One runtime partition of the fine grid plus every modality's coefficients.""" + + fine_grid: torch.Tensor + indices: torch.Tensor + node_sigmas: dict[str, torch.Tensor] + integration_weights: dict[str, torch.Tensor] + + @property + def pdd_steps(self) -> int: + return int(self.fine_grid.numel()) - 1 + + @property + def num_steps(self) -> int: + return int(self.indices.numel()) - 1 + + @property + def block_sizes(self) -> list[int]: + return (self.indices[1:] - self.indices[:-1]).tolist() + + def block(self, step: int) -> tuple[int, int]: + """Half-open fine-grid block ``[start, end)`` of runtime step *step*.""" + return int(self.indices[step]), int(self.indices[step + 1]) + + +def build_pdd_sampling_plan( + step_indices: Sequence[int], + schedules: Mapping[str, PDDModalitySchedule], + *, + device: torch.device | str | None = None, +) -> PDDSamplingPlan: + """Partition the fine grid into fused blocks for every modality. + + *step_indices* is the checkpoint's trained partition: fine-grid nodes that + increase strictly from 0 to the grid size, one fused block per consecutive + pair. ``MiniMaxH3PipelineConfig.resolve_checkpoint_settings`` validates it. + """ + max_ts = {schedule.max_t for schedule in schedules.values()} + if len(max_ts) != 1: + raise ValueError(f"Every PDD modality schedule must share max_t, got {sorted(max_ts)}") + fine_grid = pdd_fine_grid(int(step_indices[-1]), max_t=next(iter(max_ts)), device=device) + node_indices = torch.as_tensor(list(step_indices), dtype=torch.long, device=fine_grid.device) + weights = {} + for name, schedule in schedules.items(): + weights[name] = schedule.integration_weights(fine_grid) + validate_finite_nonzero(weights[name], f"PDD {name} integration weight") + return PDDSamplingPlan( + fine_grid=fine_grid, + indices=node_indices, + node_sigmas={ + name: schedule.sigma(fine_grid[node_indices]) + for name, schedule in schedules.items() + }, + integration_weights=weights, + ) + + +# ---------------------------------------------------------------------------- +# Widened output projection +# ---------------------------------------------------------------------------- + + +class PDDReplicatedLinear(ReplicatedLinear): + """Final linear layer widened to ``grid_size`` PDD heads. + + Output features are laid out head-major, ``(grid, output_size)``, so after + the architecture's own unpatchify the output channel axis reads + ``(grid, C)``. An ordinary ``forward(input)`` returns every head. + + Inside :meth:`fuse`, ``forward(input)`` instead returns the normalized + weighted average over the heads of one contiguous block ``[start, end)``, + computed by collapsing that block's parameter slices into one effective + linear. The architecture's inline final projection therefore needs no + seam, and the fused parameters are built inside ``forward``, where a + sharded weight is already gathered. + """ + + def __init__( + self, + input_size: int, + output_size: int, + *, + grid_size: int, + bias: bool = True, + skip_bias_add: bool = False, + params_dtype: torch.dtype | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__( + input_size, + output_size * grid_size, + bias=bias, + skip_bias_add=skip_bias_add, + params_dtype=params_dtype, + quant_config=quant_config, + prefix=prefix, + ) + if not isinstance(self.quant_method, UnquantizedLinearMethod): + raise NotImplementedError("PDDReplicatedLinear supports only unquantized weights") + self.head_output_size = int(output_size) + self.grid_size = grid_size + # Transient fusion state is a plain tuple, never a buffer, so it stays + # out of state_dict and can be restored as one value. + self._fusion_state: tuple[int, int, torch.Tensor, torch.dtype] | None = None + + @contextlib.contextmanager + def fuse( + self, + start: int, + end: int, + weights: torch.Tensor, + precision_decoding: torch.dtype, + ) -> Iterator[PDDReplicatedLinear]: + """Temporarily make ``forward(input)`` return the fused block output. + + *weights* are this modality's per-interval integration weights over the + whole grid (``[grid_size]``); the block ``[start, end)`` is fused with + their normalized form, so the layer emits the weighted mean and the + caller scales by the block's total weight. ``precision_decoding`` is + used only for the fused-parameter accumulation; the fused parameters + return to the layer's dtype before the projection runs. The previous + fusion state is restored on exit. + """ + previous = self._fusion_state + self._fusion_state = (start, end, weights, precision_decoding) + try: + yield self + finally: + self._fusion_state = previous + + def _fused_params( + self, + start: int, + end: int, + weights: torch.Tensor, + precision_decoding: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """The ``[O, I]`` fused weight (and bias) of block ``[start, end)``.""" + # Normalize in the incoming schedule precision, then cast only the + # fused-parameter accumulation. The layer emits the weighted mean and + # the caller multiplies by the block total, so together they represent + # ``sum_j w_j pred_j``. + block_weights = weights[start:end].to(device=self.weight.device) + block_scale = block_weights.sum() + normalized = (block_weights / block_scale).to(dtype=precision_decoding) + + def fuse_parameter(parameter: torch.Tensor) -> torch.Tensor: + heads = parameter.unflatten(0, (self.grid_size, -1))[start:end].to(dtype=precision_decoding) + return torch.einsum("n,n...->...", normalized, heads).to(dtype=parameter.dtype) + + return fuse_parameter(self.weight), None if self.bias is None else fuse_parameter(self.bias) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]: # type: ignore[override] + fusion_state = self._fusion_state + if fusion_state is None: + return super().forward(x) + fused_weight, fused_bias = self._fused_params(*fusion_state) + output = F.linear(x, fused_weight, None if self.skip_bias_add else fused_bias) + return output, fused_bias if self.skip_bias_add else None + + +@contextlib.contextmanager +def fuse_pdd_heads( + linears: Mapping[str, PDDReplicatedLinear], + start: int, + end: int, + integration_weights: Mapping[str, torch.Tensor], + precision_decoding: torch.dtype, +) -> Iterator[None]: + """Fuse block ``[start, end)`` on every modality head at once.""" + with contextlib.ExitStack() as stack: + for name, linear in linears.items(): + stack.enter_context(linear.fuse(start, end, integration_weights[name], precision_decoding)) + yield + + +__all__ = [ + "PDD_GRID_MAX_T", + "PDDModalitySchedule", + "PDDReplicatedLinear", + "PDDSamplingPlan", + "build_pdd_sampling_plan", + "fuse_pdd_heads", + "pdd_fine_grid", + "shifted_noise_amount", + "shifted_noise_delta", + "validate_finite_nonzero", + "validate_pdd_steps", +] diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py index 1cd1d16d0..e25650466 100644 --- a/fastvideo/models/dits/minimax_h3.py +++ b/fastvideo/models/dits/minimax_h3.py @@ -3,7 +3,9 @@ from __future__ import annotations +import contextlib import math +from collections.abc import Iterator, Mapping from typing import Any import torch @@ -23,6 +25,7 @@ from fastvideo.distributed.parallel_state import get_sp_world_size, model_parall from fastvideo.layers.linear import ReplicatedLinear from fastvideo.layers.quantization.mxfp8_config import MXFP8QuantizeMethod from fastvideo.layers.mlp import MLP +from fastvideo.layers.pdd import PDDReplicatedLinear, fuse_pdd_heads from fastvideo.layers.quantization import QuantizationConfig from fastvideo.layers.visual_embedding import Timesteps from fastvideo.logger import init_logger @@ -731,22 +734,56 @@ class MiniMaxH3Transformer3DModel(BaseDiT): prefix=f"{config.prefix}.norm_out", apply_silu=self.adaln_rank is None, ) - self.proj_out = ReplicatedLinear( - arch.hidden_size, - video_patch_dim, - bias=True, - quant_config=config.quant_config, - prefix=f"{config.prefix}.proj_out", - ) - self.audio_proj_out = ReplicatedLinear( - arch.hidden_size, - arch.audio_in_channels, - bias=True, - quant_config=config.quant_config, - prefix=f"{config.prefix}.audio_proj_out", - ) + # Parallel Decoding Distillation (PDD) students widen both output + # projections to ``pdd_steps`` heads, one per interval of a fixed fine + # time grid. The sampler fuses one block of heads per forward + # (``fuse_pdd_block``), so the fused output keeps the released width. + self.pdd_steps: int | None = arch.pdd_steps + self.proj_out = self._output_projection(arch.hidden_size, video_patch_dim, config, "proj_out") + self.audio_proj_out = self._output_projection(arch.hidden_size, arch.audio_in_channels, config, + "audio_proj_out") self.__post_init__() + def _output_projection(self, hidden_size: int, output_size: int, config: MiniMaxH3Config, + name: str) -> ReplicatedLinear: + if self.pdd_steps is None: + return ReplicatedLinear( + hidden_size, + output_size, + bias=True, + quant_config=config.quant_config, + prefix=f"{config.prefix}.{name}", + ) + return PDDReplicatedLinear( + hidden_size, + output_size, + grid_size=self.pdd_steps, + bias=True, + quant_config=config.quant_config, + prefix=f"{config.prefix}.{name}", + ) + + @property + def pdd_linears(self) -> dict[str, PDDReplicatedLinear]: + """The widened ``{"video": proj_out, "audio": audio_proj_out}`` heads.""" + if self.pdd_steps is None: + # Not AttributeError: nn.Module.__getattr__ would swallow it and + # report a misleading missing-attribute message. + raise RuntimeError("This MiniMax H3 transformer has no PDD heads (arch_config.pdd_steps is None)") + return {"video": self.proj_out, "audio": self.audio_proj_out} + + @contextlib.contextmanager + def fuse_pdd_block( + self, + start: int, + end: int, + integration_weights: Mapping[str, torch.Tensor], + precision_decoding: torch.dtype, + ) -> Iterator[None]: + """Fuse fine-grid block ``[start, end)`` on both widened heads for the enclosed forwards.""" + with fuse_pdd_heads(self.pdd_linears, start, end, integration_weights, precision_decoding): + yield + @staticmethod def _compile_setup_device(attention: MiniMaxH3Attention) -> torch.device: """Return the loaded device even when FP8 replaced the query weight.""" diff --git a/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py index e044822b5..b39f5c443 100644 --- a/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py +++ b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py @@ -12,9 +12,12 @@ from typing import Any import torch +from fastvideo.attention.selector import (_active_component_attention_backend_scope, coerce_attn_backend, + get_env_variable_attn_backend) from fastvideo.configs.models.vaes.minimax_h3_audio import MiniMaxH3AudioVAEArchConfig from fastvideo.configs.models.vaes.minimax_h3_video import MiniMaxH3VideoVAEArchConfig -from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig +from fastvideo.configs.pipelines.minimax_h3 import (FASTH3_INFERENCE_FILE, FASTH3_INFERENCE_SCHEMA, + MiniMaxH3PipelineConfig) from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger from fastvideo.models.hf_transformer_utils import get_diffusers_config @@ -29,6 +32,7 @@ from fastvideo.pipelines.basic.minimax_h3.stages import ( from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase from fastvideo.pipelines.lora_pipeline import LoRAPipeline from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.platforms import AttentionBackendEnum logger = init_logger(__name__) @@ -38,6 +42,41 @@ logger = init_logger(__name__) _DENOISE_MODULE_NAMES = ("vae", "audio_vae", "transformer") +def _requested_attention_backend(fastvideo_args: FastVideoArgs) -> AttentionBackendEnum | None: + """The backend the transformer loader builds with; None means automatic selection. + + Follows ``PipelineComponentLoader.load_module``: an active per-component + request wins, then the run's ``attention_backend``, then + ``FASTVIDEO_ATTENTION_BACKEND``. Automatic selection never picks + VIDEO_SPARSE_ATTN_H3. + """ + scope = _active_component_attention_backend_scope() + if scope is not None: + if scope.backend is not None or not scope.consult_env: + return scope.backend + return get_env_variable_attn_backend() + requested = coerce_attn_backend(getattr(fastvideo_args, "attention_backend", None)) + return requested if requested is not None else get_env_variable_attn_backend() + + +def _checkpoint_has_vsa_gates(transformer_dir: Path) -> bool: + """Whether a transformer checkpoint ships VSA-H3 compression gates (``to_gate_compress``). + + Reads only the shard index, or the safetensors headers when there is none. + """ + indexes = sorted(transformer_dir.glob("*.safetensors.index.json")) + if indexes: + return any(".to_gate_compress." in name for index in indexes + for name in json.loads(index.read_text(encoding="utf-8")).get("weight_map", {})) + from safetensors import safe_open + + for path in sorted(transformer_dir.glob("*.safetensors")): + with safe_open(str(path), framework="pt") as handle: + if any(".to_gate_compress." in name for name in handle.keys()): # noqa: SIM118 + return True + return False + + @dataclass(frozen=True) class _H3VideoGeometry: spatial_compression_ratio: int @@ -148,7 +187,7 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): assert patterns is not None # Keep the optional distilled schedule even when downloading only the # selected transformer partition. Otherwise Hub and local loads differ. - return [*patterns, "fastvideo_inference.json"] + return [*patterns, FASTH3_INFERENCE_FILE] def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None: _apply_h3_checkpoint_arch_configs(self.model_path, fastvideo_args, self._extra_config_module_map) @@ -161,14 +200,85 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): raise ValueError(f"MiniMax-H3 {modality} scheduler must expose a positive finite shift, got {shift}.") self._load_checkpoint_schedule(fastvideo_args) + def _checkpoint_facts(self) -> tuple[Any, bool]: + """The checkpoint's parsed ``fastvideo_inference.json`` (None without one) and whether its transformer + carries VSA compression gates. + + Read once per resolved local path: ``_load_config`` replaces a Hub + repo id in ``model_path`` with its downloaded snapshot. + """ + root = Path(self.model_path) + cached = getattr(self, "_checkpoint_facts_cache", None) + if cached is None or cached[0] != root: + path = root / FASTH3_INFERENCE_FILE + contract = json.loads(path.read_text(encoding="utf-8")) if path.is_file() else None + cached = (root, contract, _checkpoint_has_vsa_gates(self._transformer_dir())) + self._checkpoint_facts_cache = cached + return cached[1], cached[2] + + def _checkpoint_contract(self) -> Any: + """The checkpoint's ``fastvideo_inference.json``, or None when there is none.""" + return self._checkpoint_facts()[0] + + def _transformer_dir(self) -> Path: + return Path(self.model_path) / self._extra_config_module_map.get("transformer", "transformer") + + def _require_checkpoint_attention_backend(self, requested: AttentionBackendEnum | None) -> None: + """Reject a backend other than VIDEO_SPARSE_ATTN_H3 for a transformer that carries VSA compression gates. + + MiniMax-H3 builds its compression gates only under that backend and loads + strictly, so any other backend would otherwise fail later, while loading + weights, as a missing ``to_gate_compress`` parameter. + """ + if requested == AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3 or not self._checkpoint_facts()[1]: + return + raise ValueError( + f"This MiniMax-H3 checkpoint needs attention_backend=VIDEO_SPARSE_ATTN_H3: its " + f"{self._transformer_dir().name} carries VSA compression gates (to_gate_compress), and the gates are " + f"built only under that backend. This run requests " + f"{'automatic selection' if requested is None else requested.name}. Select VIDEO_SPARSE_ATTN_H3.") + + def _check_pdd_transformer(self, fastvideo_args: FastVideoArgs) -> None: + """Check this pipeline's transformer against the resolved PDD settings. + + A transformer whose config.json sets ``pdd_steps`` samples only with its + trained partition, which comes from the checkpoint's PDD fields; and + those fields describe the Ref2VA ``transformer_ref`` student. + """ + indices = getattr(fastvideo_args.pipeline_config, "pdd_step_indices", None) + config_path = self._transformer_dir() / "config.json" + widened = config_path.is_file() and json.loads(config_path.read_text(encoding="utf-8")).get("pdd_steps") + if widened and indices is None: + raise ValueError(f"{config_path} sets pdd_steps, but {self.model_path} has no PDD " + f"{FASTH3_INFERENCE_FILE}; a PDD student samples only with its trained partition.") + if indices is not None and not self._ref2va: + raise ValueError(f"{self.model_path} is a FastH3 Ref2VA PDD checkpoint; select " + "MiniMaxH3Ref2VAModularPipeline.") + + def _load_config(self, model_path: str) -> dict[str, Any]: + config = super()._load_config(model_path) + # super()._load_config resolved self.model_path to a local directory, + # downloading a Hub repo id. Check the transformer before any component + # loads, unless the caller supplied the transformer already built. + if getattr(self, "_check_transformer_before_loading", True): + self._check_pdd_transformer(self.fastvideo_args) + # A PDD checkpoint's backend comes from its fastvideo_inference.json + # (MiniMaxH3PipelineConfig.resolve_checkpoint_settings). + if getattr(self.fastvideo_args.pipeline_config, "pdd_step_indices", None) is None: + self._require_checkpoint_attention_backend(_requested_attention_backend(self.fastvideo_args)) + return config + def _load_checkpoint_schedule(self, fastvideo_args: FastVideoArgs) -> None: """A distilled export's schedule is explicit; never silently use a uniform grid.""" - path = Path(self.model_path) / "fastvideo_inference.json" - if not path.is_file(): + if getattr(fastvideo_args.pipeline_config, "pdd_step_indices", None) is not None: + # MiniMaxH3PipelineConfig.resolve_checkpoint_settings already applied + # this PDD checkpoint's file when the run's FastVideoArgs were built. return - contract = json.loads(path.read_text(encoding="utf-8")) - if not isinstance(contract, dict) or contract.get("schema_version") != "fasth3-inference-contract-v1": - raise ValueError("Unsupported FastH3 fastvideo_inference.json schema.") + contract = self._checkpoint_contract() + if contract is None: + return + if not isinstance(contract, dict) or contract.get("schema_version") != FASTH3_INFERENCE_SCHEMA: + raise ValueError(f"Unsupported FastH3 {FASTH3_INFERENCE_FILE} schema.") steps = contract.get("dmd_denoising_steps") if (not isinstance(steps, list) or not steps or any(type(step) is not int or not 0 < step <= 1000 for step in steps) @@ -237,6 +347,8 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): fastvideo_args: FastVideoArgs, loaded_modules: dict[str, torch.nn.Module] | None = None) -> dict[str, Any]: """Load the Qwen3-VL conditioner first; defer DiT and VAEs until after encode.""" + # _load_config checks the transformer against the run's settings before any component loads. + self._check_transformer_before_loading = loaded_modules is None or "transformer" not in loaded_modules if not self._defer_denoise_modules(fastvideo_args): if _use_taeh3_t2va(fastvideo_args, ref2va=self._ref2va): saved = list(self.required_config_modules) diff --git a/fastvideo/pipelines/basic/minimax_h3/packing.py b/fastvideo/pipelines/basic/minimax_h3/packing.py index 9d273e48a..bef2a6513 100644 --- a/fastvideo/pipelines/basic/minimax_h3/packing.py +++ b/fastvideo/pipelines/basic/minimax_h3/packing.py @@ -89,6 +89,14 @@ class MiniMaxH3PackedLayout: latent_height: int latent_width: int num_audio_latents: int + # Per-reference conditioning spans in packed order (Ref2VA only; empty + # elsewhere): (kind, row_count, latent_shape) covering exactly the rows + # between the text tokens and the target audio. kind is "audio", "image", + # or "video"; latent_shape is the reference's raw latent (frames, height, + # width) for visual kinds and (0, 0, 0) for audio. A video reference with + # sound contributes two entries, its audio rows first. VSA-H3 uses them to + # tile (and sparsify) reference-video regions in place. + reference_segments: tuple[tuple[str, int, tuple[int, int, int]], ...] = () def resolve_canvas_size(aspect_width: float, aspect_height: float) -> tuple[int, int]: @@ -395,10 +403,13 @@ def build_ref2va_packed_sequence( video_indices: list[torch.Tensor] = [] audio_indices: list[torch.Tensor] = [] + reference_segments: list[tuple[str, int, tuple[int, int, int]]] = [] cursor = num_text_tokens rotary_time = float(num_text_tokens) for reference, visual_row_count in zip(references, visual_row_counts, strict=True): + visual_latent_shape = (reference.num_latent_frames, reference.latent_height, reference.latent_width) if reference.media_type == "image": + reference_segments.append(("image", visual_row_count, visual_latent_shape)) rows = slice(cursor, cursor + visual_row_count) cursor = rows.stop video_indices.append(torch.arange(rows.start, rows.stop)) @@ -413,6 +424,7 @@ def build_ref2va_packed_sequence( rotary_time += 1.0 elif reference.media_type == "audio": count = reference.num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS + reference_segments.append(("audio", count, (0, 0, 0))) rows = slice(cursor, cursor + count) cursor = rows.stop audio_indices.append(torch.arange(rows.start, rows.stop)) @@ -420,6 +432,9 @@ def build_ref2va_packed_sequence( rotary_time += float(reference.num_audio_latents) elif reference.media_type == "video": audio_count = reference.num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS if reference.has_audio else 0 + if audio_count: + reference_segments.append(("audio", audio_count, (0, 0, 0))) + reference_segments.append(("video", visual_row_count, visual_latent_shape)) audio_rows = slice(cursor, cursor + audio_count) video_rows = slice(audio_rows.stop, audio_rows.stop + visual_row_count) cursor = video_rows.stop @@ -483,6 +498,7 @@ def build_ref2va_packed_sequence( latent_height=latent_height, latent_width=latent_width, num_audio_latents=num_audio_latents, + reference_segments=tuple(reference_segments), ) diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py index bce4ea978..c59ac6f11 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py @@ -3,6 +3,7 @@ from __future__ import annotations +import contextlib from typing import Any import torch @@ -12,6 +13,8 @@ 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.profiler import nvtx_range, profiler_region from fastvideo.pipelines.basic.minimax_h3.packing import ( MINIMAX_H3_KEYFRAME_NOISE_AUG, @@ -25,6 +28,18 @@ from fastvideo.pipelines.stages.validators import StageValidators as V from fastvideo.pipelines.stages.validators import VerificationResult from fastvideo.utils import get_compute_dtype +logger = init_logger(__name__) + + +def scale_pdd_initial_noise(video_rows: torch.Tensor, audio_rows: torch.Tensor, layout: MiniMaxH3PackedLayout, + plan: PDDSamplingPlan) -> None: + """Start PDD targets at the first node's noise level (``0.999 * eps``), as the students are trained. + + Latent preparation draws unit noise; conditioning rows are left untouched. + """ + video_rows[layout.num_condition_video_rows:] *= float(plan.node_sigmas["video"][0]) + audio_rows[layout.num_condition_audio_rows:] *= float(plan.node_sigmas["audio"][0]) + def _h3_vsa_metadata_builder(transformer: Any, fastvideo_args: FastVideoArgs) -> Any: """Builder instance when the transformer resolved to VSA-H3, else None. @@ -45,8 +60,16 @@ def _h3_vsa_metadata_builder(transformer: Any, fastvideo_args: FastVideoArgs) -> return backend.get_builder_cls()() -def _h3_vsa_prefix_segments(layout: MiniMaxH3PackedLayout, patch_size: tuple[int, int, int]) -> tuple[int, ...]: - """Segment sizes preceding the generated-video tail, validated against the layout.""" +def _h3_vsa_single_region_segments( + layout: MiniMaxH3PackedLayout, + patch_size: tuple[int, int, int], +) -> tuple[int | tuple[int, int, int], ...]: + """VSA-H3 packed segments with the generated video as the only sparse region. + + Returns ``(text rows, condition rows, audio rows, target video latent + shape)`` for ``MiniMaxH3VSAMetadataBuilder.build``, validated against the + layout; every conditioning row stays dense. + """ n_text = int(layout.text_indices.numel()) n_cond = int(layout.num_condition_video_rows) n_audio = int(layout.audio_indices.numel()) @@ -56,7 +79,45 @@ def _h3_vsa_prefix_segments(layout: MiniMaxH3PackedLayout, patch_size: tuple[int raise ValueError("VSA-H3 supports the standard [text|cond|audio|video] packing only; " f"segments ({n_text}, {n_cond}, {n_audio}) + video {n_video} do not sum to " f"sequence length {layout.sequence_length}.") - return n_text, n_cond, n_audio + return n_text, n_cond, n_audio, (layout.num_video_latent_frames, layout.latent_height, layout.latent_width) + + +def _h3_vsa_ref2va_segments( + layout: MiniMaxH3PackedLayout, + patch_size: tuple[int, int, int], +) -> tuple[int | tuple[int, int, int], ...]: + """VSA-H3 packed segments of a Ref2VA layout with every reference video as a sparse region. + + Returns the segments for ``MiniMaxH3VSAMetadataBuilder.build`` in true + packed order, ``[text | reference spans... | target audio | target + video]``: each reference VIDEO is a region (its latent shape); text, + every reference's audio rows, image references, and the target audio are + dense row counts; the target video is the last region. + """ + if not layout.reference_segments: + raise ValueError("Ref2VA VSA sparsification needs per-reference spans on the layout; " + "this layout carries none (T2VA/FL2VA packing).") + cursor = int(layout.text_indices.numel()) + segments: list[int | tuple[int, int, int]] = [cursor] + for kind, rows, latent_shape in layout.reference_segments: + rows = int(rows) + if kind == "video": + segments.append((int(latent_shape[0]), int(latent_shape[1]), int(latent_shape[2]))) + else: + # audio rows and single-frame image latents stay dense + segments.append(rows) + cursor += rows + n_target_audio = int(layout.audio_indices.numel()) - int(layout.num_condition_audio_rows) + segments.append(n_target_audio) + cursor += n_target_audio + segments.append((layout.num_video_latent_frames, layout.latent_height, layout.latent_width)) + n_target_video = ((layout.num_video_latent_frames // patch_size[0]) * (layout.latent_height // patch_size[1]) * + (layout.latent_width // patch_size[2])) + if cursor + n_target_video != layout.sequence_length: + raise ValueError(f"Ref2VA reference spans {layout.reference_segments} + target segments do not tile the " + f"packed sequence: reached row {cursor} + video {n_target_video} != " + f"{layout.sequence_length}.") + return tuple(segments) class MiniMaxH3DenoisingStage(PipelineStage): @@ -86,6 +147,30 @@ class MiniMaxH3DenoisingStage(PipelineStage): # clean time (1 - sigma); passing integer rungs to step() is wrong. scheduler.set_timesteps(sigmas=sigmas, device=device) + def _pdd_sampling_plan(self, fastvideo_args: FastVideoArgs, device: torch.device) -> PDDSamplingPlan | None: + """The fused-block plan of a PDD checkpoint, or None for any other checkpoint. + + ``pdd_step_indices`` is the checkpoint's trained partition, which + ``MiniMaxH3PipelineConfig.resolve_checkpoint_settings`` read and checked + against the transformer's ``pdd_steps``. Each modality's node sigmas + come from its own scheduler shift on the shared base clock. + """ + indices = getattr(fastvideo_args.pipeline_config, "pdd_step_indices", None) + if indices is None: + return None + schedules = { + "video": PDDModalitySchedule(shift=float(self.scheduler.shift)), + "audio": PDDModalitySchedule(shift=float(self.audio_scheduler.shift)), + } + return build_pdd_sampling_plan(indices, schedules, device=device) + + def _pdd_head_fusion(self, plan: PDDSamplingPlan | None, index: int) -> contextlib.AbstractContextManager[None]: + """Fuse runtime step *index*'s block of heads inside both widened output projections.""" + if plan is None: + return contextlib.nullcontext() + block_start, block_end = plan.block(index) + return self.transformer.fuse_pdd_block(block_start, block_end, plan.integration_weights, torch.float32) + def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: result = VerificationResult() result.add_check("layout", batch.extra.get(MINIMAX_H3_LAYOUT_KEY), V.not_none) @@ -117,7 +202,16 @@ class MiniMaxH3DenoisingStage(PipelineStage): device = get_local_torch_device() dmd_steps = fastvideo_args.pipeline_config.dmd_denoising_steps - if dmd_steps is None: + pdd_plan = self._pdd_sampling_plan(fastvideo_args, device) + if pdd_plan is not None: + # One runtime step fuses one block of heads into their weighted + # mean velocity; the ordinary Euler step over the block's two node + # sigmas then applies the block's total increment per modality. + self.scheduler.set_timesteps(sigmas=pdd_plan.node_sigmas["video"].to(torch.float32), device=device) + self.audio_scheduler.set_timesteps(sigmas=pdd_plan.node_sigmas["audio"].to(torch.float32), device=device) + logger.info("MiniMax-H3 PDD denoising: %d fused block(s) %s over a %d-interval fine grid.", + pdd_plan.num_steps, pdd_plan.block_sizes, pdd_plan.pdd_steps) + elif dmd_steps is None: self.scheduler.set_timesteps(batch.num_inference_steps, device=device) self.audio_scheduler.set_timesteps(batch.num_inference_steps, device=device) else: @@ -153,25 +247,44 @@ class MiniMaxH3DenoisingStage(PipelineStage): vsa_metadata_builder = _h3_vsa_metadata_builder(self.transformer, fastvideo_args) if vsa_metadata_builder is not None: vsa_patch_size = fastvideo_args.pipeline_config.dit_config.patch_size - vsa_prefix_segments = _h3_vsa_prefix_segments(layout, vsa_patch_size) # Per-request knobs (sweeps flip these between generate_video calls # without respawning workers); mode None defers to the env default. vsa_mode = batch.extra.get("vsa_mode", "exempt") if vsa_mode not in ("exempt", "compete"): raise ValueError(f"vsa_mode must be 'exempt' or 'compete', got {vsa_mode!r}.") vsa_exempt = vsa_mode == "exempt" + # PDD checkpoints set a reference keep rate: each reference video is + # then its own sparse region keeping that fraction of its tiles. + # Without one, every conditioning row stays dense. + vsa_ref_keep_rate = getattr(fastvideo_args.pipeline_config, "vsa_ref_keep_rate", None) + if vsa_ref_keep_rate is None: + vsa_packed_segments = _h3_vsa_single_region_segments(layout, vsa_patch_size) + else: + vsa_packed_segments = _h3_vsa_ref2va_segments(layout, vsa_patch_size) vsa_dense_layers = tuple(batch.extra.get("vsa_dense_layers", ())) vsa_dense_first_n = int(batch.extra.get("vsa_dense_first_n_steps", 0)) - # Run-level tile geometry (256 default, 64 = native Triton path), - # plumbed like the run-level sparsity; the builder validates the - # value against VSA_H3_TILE_SHAPES. + # Run-level tile geometry (256 default, 64 = native Triton path, + # 128 = sm_100a CUDA), plumbed like the run-level sparsity; the + # builder validates the value against VSA_H3_TILE_SHAPES. vsa_tile_size = int(fastvideo_args.VSA_tile_size) + if vsa_ref_keep_rate is not None: + num_reference_regions = sum(isinstance(segment, tuple) for segment in vsa_packed_segments) - 1 + run_sparsity = float(batch.VSA_sparsity) + if run_sparsity <= 0.0: + effect = "VSA_sparsity=0 keeps every region dense" + else: + effect = (f"each reference keeps {vsa_ref_keep_rate:g} of its tiles and the target keeps " + f"{1.0 - run_sparsity:g}") + logger.info("MiniMax-H3 VSA-H3: %d reference video region(s); %s; %d-token tiles.", + num_reference_regions, effect, vsa_tile_size) try: if full_cpu_offload: self.transformer.to(device) batch.latents = batch.latents.to(device) batch.audio_latents = batch.audio_latents.to(device) + if pdd_plan is not None: + scale_pdd_initial_noise(batch.latents, batch.audio_latents, layout, pdd_plan) # The stage range groups the complete denoising loop while the # indexed model ranges retain timing detail for every H3 block. @@ -187,15 +300,14 @@ class MiniMaxH3DenoisingStage(PipelineStage): vsa_sparsity = 0.0 if index < vsa_dense_first_n else float(batch.VSA_sparsity) attn_metadata = vsa_metadata_builder.build( current_timestep=index, - raw_latent_shape=(layout.num_video_latent_frames, layout.latent_height, - layout.latent_width), patch_size=vsa_patch_size, VSA_sparsity=vsa_sparsity, - prefix_segments=vsa_prefix_segments, + packed_segments=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, ) # Under torch.compile(mode="reduce-overhead") each denoising # step must be marked, or cudagraph trees flag cross-step @@ -203,7 +315,7 @@ class MiniMaxH3DenoisingStage(PipelineStage): # CUDAGraphs that has been overwritten" (surfaces at sp=1; # sp>1 is masked by collective-induced graph breaks). torch.compiler.cudagraph_mark_step_begin() - with trace_step(index), set_forward_context( + with self._pdd_head_fusion(pdd_plan, index), trace_step(index), set_forward_context( current_timestep=index, attn_metadata=attn_metadata, forward_batch=batch, @@ -247,4 +359,4 @@ class MiniMaxH3DenoisingStage(PipelineStage): return batch -__all__ = ["MiniMaxH3DenoisingStage"] +__all__ = ["MiniMaxH3DenoisingStage", "scale_pdd_initial_noise"] diff --git a/fastvideo/tests/attention/test_vsa_h3_backward.py b/fastvideo/tests/attention/test_vsa_h3_backward.py index 8d2d4395f..fed125f75 100644 --- a/fastvideo/tests/attention/test_vsa_h3_backward.py +++ b/fastvideo/tests/attention/test_vsa_h3_backward.py @@ -23,10 +23,9 @@ _DIM = 128 def _build_meta(device, sparsity=0.5): return MiniMaxH3VSAMetadataBuilder().build( current_timestep=0, - raw_latent_shape=_SPEC["raw_latent_shape"], patch_size=_SPEC["patch_size"], VSA_sparsity=sparsity, - prefix_segments=_SPEC["prefix_segments"], + packed_segments=(*_SPEC["prefix_segments"], _SPEC["raw_latent_shape"]), device=device, ) diff --git a/fastvideo/tests/attention/test_vsa_h3_metadata.py b/fastvideo/tests/attention/test_vsa_h3_metadata.py index 24292f26b..7823ccc9c 100644 --- a/fastvideo/tests/attention/test_vsa_h3_metadata.py +++ b/fastvideo/tests/attention/test_vsa_h3_metadata.py @@ -11,7 +11,7 @@ import torch.nn.functional as F from fastvideo.attention.backends.video_sparse_attn_h3 import (_TILE_ELEMS, MiniMaxH3VSAImpl, MiniMaxH3VSAMetadataBuilder, _build_block_mask, - _pool_tiles, _validate_h3_tile_geometry, + _pool_tiles, _validate_h3_segment_geometry, token_tile_and_valid) _720P = dict(raw_latent_shape=(30, 44, 80), patch_size=(1, 2, 2), prefix_segments=(512, 1760, 400)) @@ -30,10 +30,9 @@ _CPU = torch.device("cpu") def _build(spec, sparsity=0.0, device=_CPU, tile_size=_TILE_ELEMS): return MiniMaxH3VSAMetadataBuilder().build( current_timestep=0, - raw_latent_shape=spec["raw_latent_shape"], patch_size=spec["patch_size"], VSA_sparsity=sparsity, - prefix_segments=spec["prefix_segments"], + packed_segments=(*spec["prefix_segments"], spec["raw_latent_shape"]), device=device, tile_size=tile_size, ) @@ -92,16 +91,16 @@ def test_mask_policy(): k_vid = math.ceil(0.1 * V) scores = torch.randn(1, 2, n, n) - exempt = _build_block_mask(scores, P, V, 0.9, exempt=True) + exempt = _build_block_mask(scores, P, 0.9, True, meta.video_tile_spans, meta.span_sparsities) assert exempt[:, :, :P].all(), "prefix queries must be dense" assert exempt[..., :P].all(), "prefix keys must be visible to every query" assert (exempt[:, :, P:, P:].sum(-1) == k_vid).all(), "video rows select exactly k_vid video tiles" - compete = _build_block_mask(scores, P, V, 0.9, exempt=False) + compete = _build_block_mask(scores, P, 0.9, False, meta.video_tile_spans, meta.span_sparsities) assert compete[:, :, :P].all() assert (compete[:, :, P:].sum(-1) == min(k_vid + P, n)).all(), "budget-matched top-k" - dense = _build_block_mask(scores, P, V, 0.0, exempt=True) + dense = _build_block_mask(scores, P, 0.0, True, meta.video_tile_spans, meta.span_sparsities) assert dense.all(), "sparsity 0 must select everything" @@ -115,7 +114,8 @@ def test_sparsity_zero_matches_dense_sdpa(): scores = torch.matmul(_pool_tiles(tq, meta.variable_block_sizes), _pool_tiles(tk, meta.variable_block_sizes).transpose(-2, -1)) - mask = _build_block_mask(scores, meta.num_prefix_tiles, meta.num_video_tiles, 0.0, exempt=True) + mask = _build_block_mask(scores, meta.num_prefix_tiles, 0.0, True, meta.video_tile_spans, + meta.span_sparsities) sparse_out = impl.postprocess_output(reference_sparse_attention(tq, tk, tv, mask, meta), meta) dense_out = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)).transpose(1, 2) @@ -134,7 +134,8 @@ def test_prefix_queries_stay_dense_at_high_sparsity(): scores = torch.matmul(_pool_tiles(tq, meta.variable_block_sizes), _pool_tiles(tk, meta.variable_block_sizes).transpose(-2, -1)) for exempt in (True, False): - mask = _build_block_mask(scores, meta.num_prefix_tiles, meta.num_video_tiles, 0.75, exempt=exempt) + mask = _build_block_mask(scores, meta.num_prefix_tiles, 0.75, exempt, meta.video_tile_spans, + meta.span_sparsities) sparse_out = impl.postprocess_output(reference_sparse_attention(tq, tk, tv, mask, meta), meta) dense_out = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)).transpose(1, 2) @@ -220,7 +221,8 @@ def test_sparsity_zero_matches_dense_sdpa_tile64(): scores = torch.matmul(_pool_tiles(tq, meta.variable_block_sizes, meta.tile_elems), _pool_tiles(tk, meta.variable_block_sizes, meta.tile_elems).transpose(-2, -1)) - mask = _build_block_mask(scores, meta.num_prefix_tiles, meta.num_video_tiles, 0.0, exempt=True) + mask = _build_block_mask(scores, meta.num_prefix_tiles, 0.0, True, meta.video_tile_spans, + meta.span_sparsities) sparse_out = impl.postprocess_output(reference_sparse_attention(tq, tk, tv, mask, meta), meta) dense_out = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)).transpose(1, 2) @@ -230,18 +232,17 @@ def test_sparsity_zero_matches_dense_sdpa_tile64(): def test_geometry_guard_enforces_tile64_bound(): """A 65-token tile passes the 256 bound but must fail the 64 one.""" meta = _build(_TINY64, tile_size=64) - prefix = tuple(s for s in _TINY64["prefix_segments"] if s > 0) - dit_shape = (9, 10, 13) + segments = (*(s for s in _TINY64["prefix_segments"] if s > 0), (9, 10, 13)) sizes = meta.variable_block_sizes.clone() sizes[0] = 65 with pytest.raises(ValueError, match="tile sizes out of bounds"): - _validate_h3_tile_geometry(prefix, dit_shape, sizes, meta.untile_combined_index, 64) + _validate_h3_segment_geometry(segments, sizes, meta.untile_combined_index, 64) # the untampered tile-64 geometry passes its own bound - _validate_h3_tile_geometry(prefix, dit_shape, meta.variable_block_sizes, meta.untile_combined_index, 64) + _validate_h3_segment_geometry(segments, meta.variable_block_sizes, meta.untile_combined_index, 64) def test_builder_rejects_unknown_tile_size(): - for bad in (0, 128, 512): + for bad in (0, 96, 512): with pytest.raises(ValueError, match="tile_size"): _build(_TINY, tile_size=bad) diff --git a/fastvideo/tests/attention/test_vsa_h3_ref2va_regions.py b/fastvideo/tests/attention/test_vsa_h3_ref2va_regions.py new file mode 100644 index 000000000..ec48bbc84 --- /dev/null +++ b/fastvideo/tests/attention/test_vsa_h3_ref2va_regions.py @@ -0,0 +1,572 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU checks for VSA-H3 reference-video regions and 128-token tiles. + +With vsa_ref_policy="p2_multi_region" every Ref2VA reference video is its own +sparse region, tiled in place at its packed offset, while text, audio, and +image references stay dense. Each video query keeps its own top-k of every +region: the reference keep rate may differ from the target's. The kernel +route for 128-token tiles is exercised through a stand-in for the sm_100a +extension; real kernels are not needed. +""" + +import hashlib +import math + +import pytest +import torch +import torch.nn.functional as F + +import fastvideo.attention.backends.video_sparse_attn_h3 as vsa_h3 +import fastvideo.envs as envs +from fastvideo.attention.backends.video_sparse_attn import compute_topk +from fastvideo.attention.backends.video_sparse_attn_h3 import (_TILE_ELEMS, VSA_H3_TILE_SHAPES, MiniMaxH3VSAImpl, + MiniMaxH3VSAMetadataBuilder, _build_block_mask, + _h3_segment_tile_geometry, + _pool_tiles, token_tile_and_valid) +from fastvideo.pipelines.basic.minimax_h3.packing import MINIMAX_H3_TEXT_TAG, build_ref2va_packed_sequence +from fastvideo.pipelines.basic.minimax_h3.reference import MiniMaxH3PreparedReference +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_denoising import _h3_vsa_ref2va_segments + +_CPU = torch.device("cpu") +_PATCH = (1, 2, 2) + +# packed order: [text 37 | ref_audio 9 | ref_video (5,4,6)=120 | tgt_audio 11 +# | tgt_video (8,4,6)=192] -> S=369; raw latents under patch (1,2,2); the +# reference video keeps 25% of its tiles +_R2V = dict( + packed_segments=(37, 9, (5, 8, 12), 11, (8, 8, 12)), + ref_keep_rate=0.25, +) +_R2V_SEQ = 369 +_R2V_DENSE_RUNS = ((0, 37), (37, 46), (166, 177)) # packed [start, end) of text/ref_audio/tgt_audio + + +def _build_r2v(sparsity=0.0, tile_size=_TILE_ELEMS, device=_CPU, **overrides): + call = dict(_R2V, **overrides) + return MiniMaxH3VSAMetadataBuilder().build( + current_timestep=0, + patch_size=_PATCH, + VSA_sparsity=sparsity, + device=device, + tile_size=tile_size, + **call, + ) + + +def _impl(head_size=8): + return MiniMaxH3VSAImpl(num_heads=2, head_size=head_size, causal=False, softmax_scale=head_size**-0.5) + + +def _mask(meta, scores, sparsity): + return _build_block_mask(scores, meta.num_prefix_tiles, sparsity, meta.exempt, meta.video_tile_spans, + meta.span_sparsities) + + +def _assert_multi_region_metadata(metadata, *, expected_reference_video_regions, target_sparsity, ref_keep_rate): + """Test oracle: metadata holds one sparse span per reference video plus the target, with the policy's sparsities. + + The last span is the generated target and every earlier span is a + reference video; references keep ``ref_keep_rate`` of their tiles and the + target keeps ``1 - target_sparsity``. Text, reference audio and images, + and target audio are dense prefix tiles. + """ + expected_span_count = expected_reference_video_regions + 1 + if len(metadata.video_tile_spans) != expected_span_count: + raise ValueError("Multi-region metadata must contain one sparse span per reference video plus the target; " + f"expected {expected_span_count} spans, got {metadata.video_tile_spans}. Metadata whose only " + "sparse region is the generated video keeps reference-video conditioning dense.") + if len(metadata.span_sparsities) != expected_span_count: + raise ValueError(f"{len(metadata.span_sparsities)} span sparsities for {expected_span_count} spans.") + cursor = metadata.num_prefix_tiles + for index, (start, end) in enumerate(metadata.video_tile_spans): + if start != cursor or end <= start: + raise ValueError(f"span {index} is {(start, end)} after tile {cursor}; spans must be contiguous.") + cursor = end + if cursor != metadata.num_prefix_tiles + metadata.num_video_tiles: + raise ValueError(f"the spans end at tile {cursor} instead of covering every video tile.") + if target_sparsity <= 0.0: + expected = (0.0, ) * expected_span_count + else: + reference_sparsity = target_sparsity if ref_keep_rate is None else 1.0 - ref_keep_rate + if expected_reference_video_regions and reference_sparsity <= 0.0: + raise ValueError("sparse target attention would leave reference-video conditioning dense " + f"(ref_keep_rate={ref_keep_rate!r}).") + expected = (reference_sparsity, ) * expected_reference_video_regions + (target_sparsity, ) + if any(not math.isclose(actual, wanted, rel_tol=0.0, abs_tol=1e-12) + for actual, wanted in zip(metadata.span_sparsities, expected, strict=True)): + raise ValueError(f"per-region sparsities {metadata.span_sparsities} differ from the policy's {expected}.") + + +def _token_allow(meta, mask, b=0, h=0): + """Block mask mapped to packed-token coordinates: allow[i, j].""" + tile_of = meta.untile_combined_index // meta.tile_elems + return mask[b, h][tile_of][:, tile_of] + + +def _sparse_attention_oracle(query, key, value, mask, meta): + """SDPA over the padded tile buffer with the block mask expanded to tokens.""" + token_tile, token_valid = token_tile_and_valid(meta.variable_block_sizes, meta.tile_elems) + out = torch.empty_like(query) + for b in range(query.shape[0]): + for h in range(query.shape[2]): + allow = mask[b, h][token_tile][:, token_tile] & token_valid[None, :] + bias = torch.zeros(allow.shape, dtype=query.dtype, device=query.device) + bias.masked_fill_(~allow, float("-inf")) + out[b, :, h] = F.scaled_dot_product_attention(query[b, :, h][None], key[b, :, h][None], + value[b, :, h][None], attn_mask=bias[None])[0] + return out + + +# --------------------------------------------------------------------------- +# Packed-order segments for representative Ref2VA layouts +# --------------------------------------------------------------------------- + + +def _reference(kind, *, frames=1, height=4, width=6, audio=0): + return MiniMaxH3PreparedReference(media_type=kind, + has_audio=audio > 0, + num_latent_frames=frames, + latent_height=height, + latent_width=width, + num_audio_latents=audio) + + +def _layout(references, *, text=7, target=(3, 4, 6), audio=5): + tags = torch.full((text, ), MINIMAX_H3_TEXT_TAG, dtype=torch.long) + return build_ref2va_packed_sequence(tags, references, *target, audio, _PATCH) + + +# (references, expected packed segments, expected packed row of each video region). +# Rows: an image (1, 4, 6) latent is 1*2*3 = 6 rows; a (5, 4, 6) video is 30; the +# (3, 4, 6) target is 18; audio rows are 2 per latent; text is 7. +_LAYOUTS = { + "images_only": ([_reference("image"), _reference("image", height=8, width=4)], (7, 6, 8, 10, (3, 4, 6)), (31, )), + "image_and_silent_video": ([_reference("image"), _reference("video", frames=5)], (7, 6, (5, 4, 6), 10, + (3, 4, 6)), (13, 53)), + "video_with_audio_and_audio": ([_reference("video", frames=5, audio=4), + _reference("audio", audio=3)], (7, 8, (5, 4, 6), 6, 10, (3, 4, 6)), (15, 61)), + "two_videos": ([_reference("video", frames=2, height=8, width=4), + _reference("video", frames=5, audio=2)], (7, (2, 8, 4), 4, (5, 4, 6), 10, (3, 4, 6)), + (7, 27, 67)), +} + + +@pytest.mark.parametrize("name", sorted(_LAYOUTS)) +@pytest.mark.parametrize("tile_size", [64, 128, 256]) +def test_h3_vsa_ref2va_segments_tile_layouts_in_packed_order(name, tile_size): + references, segments, offsets = _LAYOUTS[name] + layout = _layout(references) + assert _h3_vsa_ref2va_segments(layout, _PATCH) == segments + videos = [segment for segment in segments if isinstance(segment, tuple)] + + meta = MiniMaxH3VSAMetadataBuilder().build(current_timestep=0, + patch_size=_PATCH, + VSA_sparsity=0.9, + packed_segments=segments, + device=_CPU, + tile_size=tile_size, + ref_keep_rate=0.25) + assert meta.total_seq_length == layout.sequence_length + # References keep 25% of their tiles, distinct from the target's 10%. + _assert_multi_region_metadata(meta, + expected_reference_video_regions=len(videos) - 1, + target_sparsity=0.9, + ref_keep_rate=0.25) + # Dense rows (text, audio, image references) map to prefix tiles only; + # each video region's rows map to its own span. + tile_of = meta.untile_combined_index // meta.tile_elems + region_rows = torch.zeros(layout.sequence_length, dtype=torch.bool) + for (t, h, w), offset, (start, end) in zip(videos, offsets, meta.video_tile_spans, strict=True): + rows = slice(offset, offset + (t // _PATCH[0]) * (h // _PATCH[1]) * (w // _PATCH[2])) + region_rows[rows] = True + assert bool(((tile_of[rows] >= start) & (tile_of[rows] < end)).all()) + assert bool((tile_of[~region_rows] < meta.num_prefix_tiles).all()) + x = torch.randn(1, layout.sequence_length, 2, 4) + assert torch.equal(_impl().tile(x, meta)[:, meta.untile_combined_index], x) + + +def test_h3_vsa_ref2va_segments_missing_reference_spans(): + layout = _layout([_reference("image")]) + stripped = type(layout)(**{**layout.__dict__, "reference_segments": ()}) + with pytest.raises(ValueError, match="per-reference spans"): + _h3_vsa_ref2va_segments(stripped, _PATCH) + + +# --------------------------------------------------------------------------- +# Multi-region geometry and per-region top-k +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("tile_size,width", [(64, 4), (128, 8)]) +def test_build_multi_region_native_tile_geometry(tile_size, width): + """Hand-computed per-region (4,4,4)/(4,4,8) tile ids and sizes at true offsets.""" + meta = _build_r2v(tile_size=tile_size) + assert meta.total_seq_length == _R2V_SEQ + assert meta.num_prefix_tiles == 3 # 37, 9, 11 each fit one tile + per_region = 4 if tile_size == 64 else 2 + assert meta.num_video_tiles == 2 * per_region + assert meta.video_tile_spans == ((3, 3 + per_region), (3 + per_region, 3 + 2 * per_region)) + assert meta.variable_block_sizes[:3].tolist() == [37, 9, 11] + + idx = meta.untile_combined_index + for (t, h, w), start, first_tile in (((5, 4, 6), 46, 3), ((8, 4, 6), 177, 3 + per_region)): + n_t, n_h, n_w = -(-t // 4), -(-h // 4), -(-w // width) + expected_sizes = torch.tensor([ + min(4, t - 4 * tt) * min(4, h - 4 * hh) * min(width, w - width * ww) for tt in range(n_t) + for hh in range(n_h) for ww in range(n_w) + ]) + n_region = n_t * n_h * n_w + assert torch.equal(meta.variable_block_sizes[first_tile:first_tile + n_region], expected_sizes) + row = torch.arange(t * h * w) + row_t, row_h, row_w = row // (h * w), (row // w) % h, row % w + expected_tile = first_tile + ((row_t // 4) * n_h + row_h // 4) * n_w + row_w // width + assert torch.equal(idx[start:start + t * h * w] // tile_size, expected_tile) + # Each prefix tile stays inside ONE dense segment even though the target + # audio sits between the two video regions. + for seg_start, seg_end in _R2V_DENSE_RUNS: + for tile in torch.unique(idx[seg_start:seg_end] // tile_size).tolist(): + assert tile < meta.num_prefix_tiles + rows = torch.nonzero(idx // tile_size == tile).flatten() + assert rows.min() >= seg_start and rows.max() < seg_end + + +# SHA-256 of the single-region tile geometry (partition indices, block sizes, +# untile index, prefix and video tile counts) recorded at commit 9edc8adf5, +# where MiniMaxH3VSAMetadataBuilder supported only one sparse video region. The +# hashes are fixed so the geometry is compared with that recorded output +# instead of with the code under test. +_SINGLE_REGION_GEOMETRY_SHA256 = { + ((37, 9, 120, 11), (8, 4, 6), 64): "e246d39c1be206be5b9ad938438d0594fea013d5e8c4fcb08bce05face5d4220", + ((37, 9, 120, 11), (8, 4, 6), 128): "bce9def60b23ef5dc71d790e367480c2c22e0e6f414678b72844b808d7f96bd9", + ((37, 9, 120, 11), (8, 4, 6), 256): "a710f91b80cb02bf8bbf2d37f06b612c2a4e805356c81f55c744dd066417a535", + ((512, 1200, 216), (16, 24, 40), 64): "22d1581bd63b2dfb77951e31e788aa93cd53b50c89290312b1340605889f21a9", + ((512, 1200, 216), (16, 24, 40), 128): "23b19d850777a4a6063164c5ceea33991fca5d0e7ded67d22be0a2d24c9ff0a0", + ((512, 1200, 216), (16, 24, 40), 256): "93d86b693252631c8f35efd2fdbdc216e2979c574d10c5d8d5b4fb8288750462", + ((5, 3, 250), (6, 10, 14), 64): "3226fc23f45bec409992d2bf0f42794c609d53bac401c5162e7c37615904cc41", + ((5, 3, 250), (6, 10, 14), 128): "2645e5a4c49b30de3f355ee0163e8703de2676a9a21e6d04b6b1a210a01242ab", + ((5, 3, 250), (6, 10, 14), 256): "220fa91d488c88832d08177e28d8bb57f2477b7be36677302ecfefb9cef7c56c", +} + + +def _geometry_sha256(outputs): + digest = hashlib.sha256() + for item in outputs: + if torch.is_tensor(item): + tensor = item.detach().cpu().contiguous() + digest.update(str(tensor.dtype).encode()) + digest.update(str(tuple(tensor.shape)).encode()) + digest.update(tensor.numpy().tobytes()) + else: + digest.update(repr(int(item)).encode()) + return digest.hexdigest() + + +@pytest.mark.parametrize("prefix_segments,dit_seq_shape,tile_size", sorted(_SINGLE_REGION_GEOMETRY_SHA256)) +def test_h3_segment_tile_geometry_single_region_matches_frozen_hashes(prefix_segments, dit_seq_shape, tile_size): + geometry = _h3_segment_tile_geometry((*prefix_segments, dit_seq_shape), _CPU, VSA_H3_TILE_SHAPES[tile_size])[:5] + assert _geometry_sha256(geometry) == _SINGLE_REGION_GEOMETRY_SHA256[prefix_segments, dit_seq_shape, tile_size] + + +def test_build_single_region_one_generated_video_span(): + single = MiniMaxH3VSAMetadataBuilder().build(current_timestep=0, + patch_size=_PATCH, + VSA_sparsity=0.5, + packed_segments=(37, 9, 120, 11, (8, 8, 12)), + device=_CPU, + tile_size=128) + # The 37, 9, 120, and 11 dense rows each fit one 128-token tile; the (8, 4, 6) video grid + # makes 2 * 1 * 1 (4, 4, 8) tiles. + assert (single.num_prefix_tiles, single.num_video_tiles) == (4, 2) + assert single.video_tile_spans == ((4, 6), ) + assert single.span_sparsities == (0.5, ) + + +def test_build_block_mask_per_region_topk_and_reference_keep_rate(): + """Every video query keeps exactly k_i tiles of EACH region; prefix stays dense.""" + torch.manual_seed(3) + meta = _build_r2v(sparsity=0.75, tile_size=64, ref_keep_rate=0.5) + P = meta.num_prefix_tiles + n = P + meta.num_video_tiles + scores = torch.randn(1, 2, n, n) + mask = _mask(meta, scores, 0.75) + assert mask[:, :, :P].all(), "prefix queries dense" + assert mask[..., :P].all(), "prefix keys visible to every query" + for (start, end), sparsity in zip(meta.video_tile_spans, meta.span_sparsities, strict=True): + assert (mask[:, :, P:, start:end].sum(-1) == compute_topk(sparsity, end - start)).all(), (start, end) + assert meta.span_sparsities == pytest.approx((0.5, 0.75)) + + # A span at sparsity 0 keeps every column; the other span keeps its own top-k. + ref_span, tgt_span = meta.video_tile_spans + mask = _build_block_mask(scores, P, 0.75, meta.exempt, meta.video_tile_spans, (0.0, 0.75)) + assert mask[..., ref_span[0]:ref_span[1]].all(), "a sparsity-0 span's columns are dense" + assert (mask[:, :, P:, tgt_span[0]:tgt_span[1]].sum(-1) == compute_topk(0.75, tgt_span[1] - tgt_span[0])).all() + + keep_quarter = _build_r2v(sparsity=0.9, tile_size=64, ref_keep_rate=0.25) + assert keep_quarter.span_sparsities == pytest.approx((0.75, 0.9)) + # A dense build (sparsity 0: dense steps) ignores the reference override. + assert _build_r2v(sparsity=0.0, ref_keep_rate=0.25).span_sparsities == (0.0, 0.0) + # Dense layers pass sparsity 0 to the mask builder, which overrides the spans. + assert _build_block_mask(scores, P, 0.0, True, meta.video_tile_spans, meta.span_sparsities).all() + + +def test_build_reference_video_region_sparsifies_reference_keys(): + """A reference-video region sparsifies that video's keys and queries; target-to-target selection is unchanged.""" + torch.manual_seed(4) + q = torch.randn(1, _R2V_SEQ, 2, 8) + k = torch.randn(1, _R2V_SEQ, 2, 8) + allows = [] + # The reference video's 120 rows as a dense segment, then as its own region. + for build in (dict(packed_segments=(37, 9, 120, 11, (8, 8, 12)), ref_keep_rate=None), {}): + meta = _build_r2v(sparsity=0.75, **build) + impl = _impl() + tq, tk = (impl.tile(t, meta).clone() for t in (q, k)) + scores = torch.matmul(_pool_tiles(tq, meta.variable_block_sizes), + _pool_tiles(tk, meta.variable_block_sizes).transpose(-2, -1)) + allows.append(torch.stack([_token_allow(meta, _mask(meta, scores, 0.75), 0, h) for h in range(2)])) + dense_reference_allow, multi_region_allow = allows + assert dense_reference_allow[..., 46:166].all(), "a dense reference video is visible to every query" + assert not torch.equal(dense_reference_allow, multi_region_allow) + for start, end in _R2V_DENSE_RUNS: + assert multi_region_allow[:, start:end].all() and multi_region_allow[..., start:end].all(), (start, end) + assert not multi_region_allow[:, 200, 46:166].all(), "a reference-video region sparsifies its keys" + assert torch.equal(dense_reference_allow[:, 177:, 177:], multi_region_allow[:, 177:, 177:]) + + +@pytest.mark.parametrize("tile_size", [_TILE_ELEMS, 64, 128]) +def test_build_multi_region_sparsity_zero_matches_dense_sdpa(tile_size): + torch.manual_seed(5) + meta = _build_r2v(tile_size=tile_size) + q, k, v = (torch.randn(1, _R2V_SEQ, 2, 8) for _ in range(3)) + impl = _impl() + # An odd 128-token tile count carries one transport-only partner tile; the + # attention problem is the logical prefix. + logical = meta.variable_block_sizes.numel() * meta.tile_elems + tq, tk, tv = (impl.tile(t, meta)[:, :logical].clone() for t in (q, k, v)) + scores = torch.matmul(_pool_tiles(tq, meta.variable_block_sizes, meta.tile_elems), + _pool_tiles(tk, meta.variable_block_sizes, meta.tile_elems).transpose(-2, -1)) + mask = _mask(meta, scores, 0.0) + assert mask.all() + sparse_out = impl.postprocess_output(_sparse_attention_oracle(tq, tk, tv, mask, meta), meta) + dense_out = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)).transpose(1, 2) + assert torch.allclose(sparse_out, dense_out, atol=1e-5), (sparse_out - dense_out).abs().max() + + +def test_build_region_shape_not_multiple_of_patch(): + with pytest.raises(ValueError, match="not a positive multiple of patch"): + _build_r2v(packed_segments=(37, 9, (5, 7, 12), 11, (8, 8, 12))) + + +def test_build_no_video_region(): + with pytest.raises(ValueError, match="at least one video region"): + _build_r2v(packed_segments=(37, 9, 11)) + + +def test_build_compete_with_reference_video_regions(): + with pytest.raises(ValueError, match="compete' supports only the generated-video region"): + _build_r2v(sparsity=0.75, exempt=False) + + +def test_build_compete_with_generated_video_only(): + meta = _build_r2v(sparsity=0.75, exempt=False, packed_segments=(37, 20, (8, 8, 12))) + assert meta.exempt is False and len(meta.video_tile_spans) == 1 + + +def test_assert_multi_region_metadata_rejects_mismatched_metadata(): + single = MiniMaxH3VSAMetadataBuilder().build(current_timestep=0, + patch_size=_PATCH, + VSA_sparsity=0.9, + packed_segments=(37, 9, 120, 11, (8, 8, 12)), + device=_CPU, + tile_size=64) + with pytest.raises(ValueError, match="only sparse region is the generated video"): + _assert_multi_region_metadata(single, expected_reference_video_regions=1, target_sparsity=0.9, + ref_keep_rate=0.1) + with pytest.raises(ValueError, match="leave reference-video conditioning dense"): + _assert_multi_region_metadata(_build_r2v(sparsity=0.9, tile_size=64, ref_keep_rate=0.1), + expected_reference_video_regions=1, + target_sparsity=0.9, + ref_keep_rate=1.0) + with pytest.raises(ValueError, match="differ from the policy"): + _assert_multi_region_metadata(_build_r2v(sparsity=0.9, tile_size=64, ref_keep_rate=0.1), + expected_reference_video_regions=1, + target_sparsity=0.9, + ref_keep_rate=0.2) + _assert_multi_region_metadata(_build_r2v(sparsity=0.0, tile_size=64, ref_keep_rate=0.1), + expected_reference_video_regions=1, + target_sparsity=0.0, + ref_keep_rate=0.1) + + +def test_build_multi_region_missing_ref_keep_rate(): + with pytest.raises(ValueError, match="needs ref_keep_rate"): + _build_r2v(sparsity=0.75, ref_keep_rate=None) + + +# --------------------------------------------------------------------------- +# 128-token tiles: sm_100a/sm_103a CUDA only +# --------------------------------------------------------------------------- + +_HEADS, _DIM = 2, 128 + + +class _FakeSm100a: + """Stands in for fastvideo_kernel.block_sparse_attn_sm100a.""" + + def __init__(self, supported=True): + self.supported = supported + self.calls = [] + + def is_supported(self, 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): + self.calls.append(dict(q=q, q2k_num=q2k_num, vbs=variable_block_sizes, need_lse=need_lse)) + return q.clone(), None + + +def _fake_map_to_index(block_map): + num = block_map.sum(dim=-1, dtype=torch.int32) + idx = torch.full(block_map.shape, -1, dtype=torch.int32) + for position in torch.nonzero(block_map.flatten(0, -2).any(-1)).flatten().tolist(): + cols = torch.nonzero(block_map.flatten(0, -2)[position]).flatten() + idx.flatten(0, -2)[position, :cols.numel()] = cols.to(torch.int32) + return idx, num + + +@pytest.fixture() +def tile128(monkeypatch, env_overrides): + fake = _FakeSm100a() + + def forbid_triton(*args, **kwargs): + raise AssertionError("tile 128 must never fall back to the Triton-64 kernel") + + monkeypatch.setattr(vsa_h3, "_sm100a", fake) + monkeypatch.setattr(vsa_h3, "map_to_index", _fake_map_to_index) + monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", forbid_triton) + # Tile 128 has no opt-in: it routes to the CUDA kernel with the switch off. + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(None)) + return fake + + +def _run_tile128(meta, requires_grad=False): + impl = _impl(_DIM) + raw = torch.randn(3, meta.total_seq_length, _HEADS, _DIM, dtype=torch.bfloat16, requires_grad=requires_grad) + tiled = impl.preprocess_qkv(raw, meta) + query, key, value = tiled.chunk(3, dim=0) + return tiled, impl.forward(query, key, value, None, meta) + + +def test_forward_tile128_odd_tile_count_adds_partner_tile(tile128): + meta = _build_r2v(sparsity=0.5, tile_size=128, ref_keep_rate=0.25) + n_tiles = meta.variable_block_sizes.numel() + assert n_tiles % 2 == 1 + tiled, out = _run_tile128(meta) + assert tiled.shape[1] == (n_tiles + 1) * 128 + assert torch.count_nonzero(tiled[:, n_tiles * 128:]) == 0 + assert len(tile128.calls) == 1 + call = tile128.calls[0] + assert call["need_lse"] is False + assert call["vbs"].dtype == torch.int32 and call["vbs"].numel() == n_tiles + 1 and call["vbs"][-1] == 0 + assert (call["q2k_num"][..., n_tiles] == 0).all() + assert out.shape == (1, n_tiles * 128, _HEADS, _DIM) + + +def test_forward_tile128_even_tile_count_no_partner_tile(tile128): + meta = _build_r2v(sparsity=0.5, tile_size=128, packed_segments=(37, 9, (5, 8, 12), 11, (8, 8, 12), 64)) + assert meta.variable_block_sizes.numel() % 2 == 0 + tiled, _ = _run_tile128(meta) + assert tiled.shape[1] == meta.variable_block_sizes.numel() * 128 + assert len(tile128.calls) == 1 + + +def test_forward_tile128_unavailable_cuda_kernel(tile128, monkeypatch): + meta = _build_r2v(sparsity=0.5, tile_size=128) + tile128.supported = False + with pytest.raises(RuntimeError, match="tile 128 requires the sm_100a/sm_103a CUDA block-sparse kernel"): + _run_tile128(meta) + monkeypatch.setattr(vsa_h3, "_sm100a", None) + with pytest.raises(RuntimeError, match="not installed"): + _run_tile128(meta) + + +def test_forward_tile128_route_resolved_once_per_device(tile128): + meta = _build_r2v(sparsity=0.5, tile_size=128) + probes = [] + tile128.is_supported = lambda q, variable_block_sizes: probes.append(q.device) or True + impl = _impl(_DIM) + for _ in range(3): + raw = torch.randn(3, meta.total_seq_length, _HEADS, _DIM, dtype=torch.bfloat16) + query, key, value = impl.preprocess_qkv(raw, meta).chunk(3, dim=0) + impl.forward(query, key, value, None, meta) + assert len(tile128.calls) == 3 and probes == [_CPU] + # A failed probe is cached too, and every later call still fails closed. + tile128.is_supported = lambda q, variable_block_sizes: probes.append(q.device) and False + impl = _impl(_DIM) + raw = torch.randn(3, meta.total_seq_length, _HEADS, _DIM, dtype=torch.bfloat16) + query, key, value = impl.preprocess_qkv(raw, meta).chunk(3, dim=0) + for _ in range(2): + with pytest.raises(RuntimeError, match="tile 128 requires the sm_100a/sm_103a CUDA block-sparse kernel"): + impl.forward(query, key, value, None, meta) + assert probes == [_CPU, _CPU] + + +def test_forward_tile128_requires_no_grad(tile128): + meta = _build_r2v(sparsity=0.5, tile_size=128) + with pytest.raises(NotImplementedError, match="no-grad inference"): + _run_tile128(meta, requires_grad=True) + assert tile128.calls == [] + + +def test_forward_tile128_rejects_regional_fullgraph_compile(tile128, monkeypatch): + meta = _build_r2v(sparsity=0.5, tile_size=128) + impl = _impl(_DIM) + impl._regional_compile_sm100a_enabled = True + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) + q = torch.zeros(1, meta.variable_block_sizes.numel() * 128, _HEADS, _DIM, dtype=torch.bfloat16) + with pytest.raises(RuntimeError, match="requires 64-token tiles"): + impl.forward(q, q, q, None, meta) + + +def test_build_tile128_uses_4x4x8_tiles(): + meta = _build_r2v(tile_size=128) + assert meta.tile_elems == 128 + assert int(meta.variable_block_sizes.max()) <= 128 + assert int(meta.variable_block_sizes.sum()) == _R2V_SEQ + # (5, 4, 6) and (8, 4, 6) token grids -> ceil(t/4) * ceil(4/4) * ceil(6/8) tiles each + assert meta.video_tile_spans == ((3, 5), (5, 7)) + assert math.prod((4, 4, 8)) == 128 + + +@pytest.mark.parametrize("sparsity,ref_keep_rate", [(0.5, 0.25), (0.9, 0.1)]) +def test_forward_tile128_real_sm100a_kernel_matches_oracle(env_overrides, sparsity, ref_keep_rate): + """The real 128-token CUDA forward on a multi-region layout against an FP32 masked-SDPA oracle.""" + if not torch.cuda.is_available() or torch.cuda.get_device_capability() not in {(10, 0), (10, 3)}: + pytest.skip("requires an sm_100a/sm_103a GPU (B200, B300, GB200, GB300)") + forwards = getattr(vsa_h3._sm100a, "_FWD_BY_BLOCK", {}) if vsa_h3._sm100a is not None else {} + if forwards.get(128) is None or vsa_h3.map_to_index is None: + pytest.skip("requires a fastvideo_kernel build with the 128-token sm_100a forward") + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(None)) + device = torch.device("cuda") + meta = _build_r2v(sparsity=sparsity, tile_size=128, ref_keep_rate=ref_keep_rate, device=device) + n_tiles = meta.variable_block_sizes.numel() + logical = n_tiles * 128 + impl = _impl(_DIM) + torch.manual_seed(19) + raw = torch.randn(3, meta.total_seq_length, _HEADS, _DIM, device=device, dtype=torch.bfloat16) + with torch.inference_mode(): + tiled = impl.preprocess_qkv(raw, meta) + assert tiled.shape[1] == (n_tiles + n_tiles % 2) * 128 + query, key, value = tiled.chunk(3, dim=0) + actual = impl.postprocess_output(impl.forward(query, key, value, None, meta), meta) + # The mask the backend selected: pooled scores of the logical tiles. + logical_query, logical_key, logical_value = (t[:, :logical] for t in (query, key, value)) + scores = torch.matmul(_pool_tiles(logical_query, meta.variable_block_sizes, 128), + _pool_tiles(logical_key, meta.variable_block_sizes, 128).transpose(-2, -1)) / _DIM**0.5 + mask = _mask(meta, scores, sparsity) + assert not mask.all(), "the oracle must exercise a sparse mask" + expected = impl.postprocess_output( + _sparse_attention_oracle(logical_query.float(), logical_key.float(), logical_value.float(), mask, meta), + meta) + torch.cuda.synchronize() + assert actual.shape == (1, meta.total_seq_length, _HEADS, _DIM) + torch.testing.assert_close(actual.float(), expected, atol=0.04, rtol=0.02) diff --git a/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py b/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py index 9b26678c8..0eaa79130 100644 --- a/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py +++ b/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py @@ -35,10 +35,9 @@ def _build_meta( builder = builder or MiniMaxH3VSAMetadataBuilder() return builder.build( current_timestep=0, - raw_latent_shape=_SPEC["raw_latent_shape"], patch_size=_SPEC["patch_size"], VSA_sparsity=sparsity, - prefix_segments=prefix_segments, + packed_segments=(*prefix_segments, _SPEC["raw_latent_shape"]), device=device, dense_layers=dense_layers, tile_size=64, @@ -224,10 +223,9 @@ def test_prepared_route_fullgraph_avoids_eager_dispatch(monkeypatch, env_overrid # 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), + packed_segments=(64, 64, _SPEC["raw_latent_shape"]), device=torch.device("cpu"), tile_size=64, ) diff --git a/fastvideo/tests/attention/test_vsa_h3_tile_grad_safety.py b/fastvideo/tests/attention/test_vsa_h3_tile_grad_safety.py index 30c9aa4b4..93ecf8dd3 100644 --- a/fastvideo/tests/attention/test_vsa_h3_tile_grad_safety.py +++ b/fastvideo/tests/attention/test_vsa_h3_tile_grad_safety.py @@ -36,10 +36,9 @@ def _build(builder=None, step=0): # A fresh metadata per step, all sharing the builder's holder. return builder.build( current_timestep=step, - raw_latent_shape=_SPEC["raw_latent_shape"], patch_size=_SPEC["patch_size"], VSA_sparsity=0.9, - prefix_segments=_SPEC["prefix_segments"], + packed_segments=(*_SPEC["prefix_segments"], _SPEC["raw_latent_shape"]), device=torch.device("cpu"), tile_size=_TILE_ELEMS, ), builder diff --git a/fastvideo/tests/entrypoints/test_video_generator.py b/fastvideo/tests/entrypoints/test_video_generator.py index e446b4389..82fbaff4b 100644 --- a/fastvideo/tests/entrypoints/test_video_generator.py +++ b/fastvideo/tests/entrypoints/test_video_generator.py @@ -1,3 +1,4 @@ +import dataclasses import json import shutil import subprocess @@ -41,6 +42,7 @@ def _new_runtime_video_generator() -> VideoGenerator: model_path="test-model", prompt_txt=None, workload_type=SimpleNamespace(value="t2v"), + pipeline_config=SimpleNamespace(apply_request_constraints=lambda request, sampling_param: sampling_param), ) generator.executor = SimpleNamespace( set_log_queue=lambda queue: None, @@ -114,7 +116,9 @@ def _single_video_args(output_type="pil"): output_type=output_type, pin_cpu_memory=False, VSA_sparsity=0.0, - pipeline_config=SimpleNamespace(flow_shift=1.0, embedded_cfg_scale=1.0), + pipeline_config=SimpleNamespace(flow_shift=1.0, + embedded_cfg_scale=1.0, + apply_request_constraints=lambda request, sampling_param: sampling_param), workload_type=SimpleNamespace(value="t2v"), ) @@ -704,6 +708,31 @@ def test_generate_uses_typed_request_path(monkeypatch): assert result.video_path == "outputs/test.mp4" +@pytest.mark.parametrize("entry", ["generate", "generate_video"]) +def test_generate_applies_pipeline_config_request_constraints(monkeypatch, entry): + """Both entry points run the sampling params that the loaded model's request constraints return.""" + generator = _new_runtime_video_generator() + _patch_sampling_param_from_pretrained(monkeypatch) + + def apply_request_constraints(request, sampling_param): + return dataclasses.replace(sampling_param, num_inference_steps=8) + + generator.fastvideo_args.pipeline_config = SimpleNamespace(apply_request_constraints=apply_request_constraints) + captured = {} + + def fake_generate_video_impl(prompt=None, sampling_param=None, **kwargs): + captured["sampling_param"] = sampling_param + return {"prompts": prompt, "video_path": "outputs/test.mp4"} + + monkeypatch.setattr(generator, "_generate_video_impl", fake_generate_video_impl) + if entry == "generate": + generator.generate(GenerationRequest(prompt="hello world")) + else: + with pytest.warns(DeprecationWarning): + generator.generate_video(prompt="hello world") + assert captured["sampling_param"].num_inference_steps == 8 + + def test_generate_rejects_stage_override_outside_registered_stage(monkeypatch) -> None: """Validate typed stage overrides at the public generation entrypoint.""" generator = _new_runtime_video_generator() @@ -923,7 +952,8 @@ def test_generate_video_legacy_call_uses_legacy_impl(monkeypatch): def test_generate_video_legacy_call_routes_compat_kwargs(monkeypatch): generator = _new_runtime_video_generator() - generator.fastvideo_args.pipeline_config = SimpleNamespace(embedded_cfg_scale=1.0) + generator.fastvideo_args.pipeline_config = SimpleNamespace( + embedded_cfg_scale=1.0, apply_request_constraints=lambda request, sampling_param: sampling_param) captured = {} def fake_generate_video_impl( diff --git a/fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py b/fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py new file mode 100644 index 000000000..26db7a4c8 --- /dev/null +++ b/fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU checks of examples/inference/basic/basic_fasth3_omniref_pdd.py (no weights, no GPU).""" +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[3] +EXAMPLE_PATH = REPO_ROOT / "examples" / "inference" / "basic" / "basic_fasth3_omniref_pdd.py" +BASE_PIN = "9bfb6693f2cf6de171db46d1aa586f67d773a1da" +CONTRACT = { + "schema_version": "fasth3-inference-contract-v1", + "base_model_revision": f"hf://MiniMaxAI/MiniMax-H3@{BASE_PIN}", + "pdd_steps": 32, + "pdd_step_indices": list(range(0, 33, 4)), + "num_inference_steps": 8, + "transformer_forwards": 8, + "video_scheduler_shift": 12.0, + "audio_scheduler_shift": 3.0, + "attention_backend": "VIDEO_SPARSE_ATTN_H3", + "vsa_sparsity": 0.9, + "vsa_tile_size": 128, + "vsa_ref_policy": "p2_multi_region", + "vsa_ref_keep_rate": 0.1, +} + + +def _load_example(): + spec = importlib.util.spec_from_file_location("basic_fasth3_omniref_pdd", EXAMPLE_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 + + +example = _load_example() + + +def _args(*overrides: str): + return example.parse_args(["--model-path", "export", "--image", "a.png", "--prompt", "p", *overrides]) + + +@pytest.mark.parametrize("overrides,expected", [ + ((), ("MiniMaxAI/MiniMax-H3", BASE_PIN)), + (("--base-revision", "abc"), ("MiniMaxAI/MiniMax-H3", "abc")), + # The export's pin names a revision of its own base repo; it never applies to another one. + (("--base-model-path", "someone/mirror"), ("someone/mirror", None)), + (("--base-model-path", "someone/mirror", "--base-revision", "abc"), ("someone/mirror", "abc")), +]) +def test_base_model_source_command_line_overrides(overrides, expected): + assert example.base_model_source(_args(*overrides), CONTRACT) == expected + + +@pytest.mark.parametrize("pin", [ + BASE_PIN, f"MiniMaxAI/MiniMax-H3@{BASE_PIN}", "hf://MiniMaxAI/MiniMax-H3", "hf://@abc", "", None, +]) +@pytest.mark.parametrize("overrides", [(), ("--base-model-path", "someone/mirror")]) +def test_base_model_source_unparsable_base_pin(pin, overrides): + """Never a silent fall back to the base repo's latest revision; the pipeline rejects the same values.""" + with pytest.raises(ValueError, match="must be hf://@"): + example.base_model_source(_args(*overrides), {**CONTRACT, "base_model_revision": pin}) + + +_COMPONENT = ["diffusers", "Component"] +_FULL_MANIFEST = { + "_class_name": "MiniMaxH3ModularPipeline", + **{name: _COMPONENT for name in ("transformer", "transformer_ref", "scheduler", "audio_scheduler", + "text_encoder", "tokenizer", "processor", "vae", "audio_vae")}, +} +_EXPORT_ONLY_MANIFEST = { + "_class_name": "MiniMaxH3ModularPipeline", + **{name: _COMPONENT for name in ("transformer_ref", "scheduler", "audio_scheduler")}, +} + + +def _snapshot_dir(root: Path, components, manifest): + root.mkdir() + for name in components: + (root / name).mkdir() + if manifest is not None: + (root / "modular_model_index.json").write_text(json.dumps(manifest)) + return root + + +@pytest.mark.parametrize("export_manifest,chosen", [ + (_FULL_MANIFEST, "export"), + (_EXPORT_ONLY_MANIFEST, "base"), + (None, "base"), +]) +def test_compose_model_dir_selects_complete_manifest(tmp_path, export_manifest, chosen): + export = _snapshot_dir(tmp_path / "export", example.EXPORT_COMPONENTS, export_manifest) + (export / example.CONTRACT).write_text(json.dumps(CONTRACT)) + base = _snapshot_dir(tmp_path / "base", example.BASE_COMPONENTS, _FULL_MANIFEST) + composed = example.compose_model_dir(export, base, tmp_path / "composed") + manifest = composed / "modular_model_index.json" + assert manifest.resolve() == ((export if chosen == "export" else base) / "modular_model_index.json").resolve() + for name in (*example.EXPORT_COMPONENTS, *example.BASE_COMPONENTS, example.CONTRACT): + assert (composed / name).exists() + + +def test_compose_model_dir_no_complete_manifest(tmp_path): + export = _snapshot_dir(tmp_path / "export", example.EXPORT_COMPONENTS, _EXPORT_ONLY_MANIFEST) + (export / example.CONTRACT).write_text(json.dumps(CONTRACT)) + base = _snapshot_dir(tmp_path / "base", example.BASE_COMPONENTS, _EXPORT_ONLY_MANIFEST) + with pytest.raises(FileNotFoundError, match="manifest that declares"): + example.compose_model_dir(export, base, tmp_path / "composed") + + +def test_build_generator_config_leaves_trained_settings_to_checkpoint(tmp_path): + # FastVideo reads the trained attention settings from the composed directory's fastvideo_inference.json. + assert example.build_generator_config(tmp_path, 1).pipeline.experimental == {} + + +@pytest.mark.parametrize("num_gpus,sharded", [(1, False), (4, True)]) +def test_build_generator_config_multi_gpu_shards_dit(tmp_path, num_gpus, sharded): + config = example.build_generator_config(tmp_path, num_gpus) + assert config.engine.use_fsdp_inference is sharded + assert config.engine.parallelism.sp_size == num_gpus + assert config.pipeline.components.override_pipeline_cls_name == "MiniMaxH3Ref2VAModularPipeline" diff --git a/fastvideo/tests/layers/test_pdd_linear.py b/fastvideo/tests/layers/test_pdd_linear.py new file mode 100644 index 000000000..1de8ce50a --- /dev/null +++ b/fastvideo/tests/layers/test_pdd_linear.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU contracts for the PDD sampling primitives in ``fastvideo.layers.pdd``. + +References are spelled out locally: the rational time shift on the +``[0, 0.999]`` base clock, its cancellation-free difference, the fused blocks +of an explicit fine-grid partition, and the fused-head forward as the +integration-weighted mean of the materialized heads. +""" + +from __future__ import annotations + +import pytest +import torch + +from fastvideo.layers.pdd import ( + PDD_GRID_MAX_T, + PDDModalitySchedule, + PDDReplicatedLinear, + build_pdd_sampling_plan, + fuse_pdd_heads, + pdd_fine_grid, + shifted_noise_amount, + shifted_noise_delta, +) +from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler + +# The OmniRef PDD-8 export: a 32-interval grid sampled in eight 4-head blocks. +GRID32_BLOCKS8 = list(range(0, 33, 4)) + + +def _time_shift(value: torch.Tensor, shift: float, max_t: float = PDD_GRID_MAX_T) -> torch.Tensor: + """``s * t * M / (t * (s - 1) + M)`` in float64.""" + value = value.to(torch.float64) + return value * shift * max_t / (value * (shift - 1.0) + max_t) + + +@pytest.mark.parametrize("shift", [1.0, 0.25, 3.0, 12.0]) +def test_shifted_noise_amount_matches_rational_time_shift(shift: float) -> None: + base = torch.linspace(0.0, PDD_GRID_MAX_T, 17, dtype=torch.float64) + torch.testing.assert_close(shifted_noise_amount(base, shift), _time_shift(base, shift), rtol=0.0, atol=1e-15) + # max_t is a fixed point of every shift, so the first node has the same + # noise level in every modality. + fixed_point = float(shifted_noise_amount(torch.tensor(PDD_GRID_MAX_T, dtype=torch.float64), shift)) + assert fixed_point == pytest.approx(PDD_GRID_MAX_T, abs=1e-15) + + +@pytest.mark.parametrize("shift", [1.0, 0.25, 12.0]) +def test_shifted_noise_delta_matches_direct_difference(shift: float) -> None: + grid = pdd_fine_grid(256) + direct = shifted_noise_amount(grid[1:], shift) - shifted_noise_amount(grid[:-1], shift) + delta = shifted_noise_delta(grid[:-1], grid[1:], shift) + torch.testing.assert_close(delta, direct, rtol=1e-12, atol=1e-15) + assert bool((delta < 0).all()) + if shift == 1.0: + assert torch.equal(delta, grid[1:] - grid[:-1]) + + +def test_pdd_fine_grid_descending_float64_from_max_t_to_zero() -> None: + grid = pdd_fine_grid(8) + assert grid.dtype == torch.float64 + assert grid.shape == (9, ) + assert float(grid[0]) == PDD_GRID_MAX_T + assert float(grid[-1]) == 0.0 + assert bool((grid[:-1] > grid[1:]).all()) + + +@pytest.mark.parametrize(("step_indices", "block_sizes"), [ + ([0, 64, 128, 192, 256], [64, 64, 64, 64]), + ([0, 2, 5, 8], [2, 3, 3]), +], ids=["uniform", "uneven"]) +def test_build_pdd_sampling_plan_blocks_and_node_sigmas(step_indices: list[int], block_sizes: list[int]) -> None: + schedules = {"video": PDDModalitySchedule(shift=12.0), "audio": PDDModalitySchedule(shift=3.0)} + plan = build_pdd_sampling_plan(step_indices, schedules) + + grid_size = step_indices[-1] + assert plan.pdd_steps == grid_size + assert plan.num_steps == len(block_sizes) + assert plan.block_sizes == block_sizes + assert plan.block(1) == (step_indices[1], step_indices[2]) + # Node i of an N-interval grid sits at base time max_t * (1 - i / N). + node_base = PDD_GRID_MAX_T * (1.0 - torch.tensor(step_indices, dtype=torch.float64) / grid_size) + for name, schedule in schedules.items(): + node_sigmas = plan.node_sigmas[name] + torch.testing.assert_close(node_sigmas, _time_shift(node_base, schedule.shift), rtol=0.0, atol=1e-12) + assert float(node_sigmas[0]) == pytest.approx(PDD_GRID_MAX_T, abs=1e-15) + assert float(node_sigmas[-1]) == 0.0 + # A fused block's total weight is exactly its node-sigma increment, + # which is what makes one Euler step over the two nodes the block advance. + for step in range(plan.num_steps): + start, end = plan.block(step) + total = plan.integration_weights[name][start:end].sum() + torch.testing.assert_close(total, node_sigmas[step + 1] - node_sigmas[step], rtol=1e-12, atol=1e-15) + + +def test_build_pdd_sampling_plan_mismatched_max_t() -> None: + schedules = {"a": PDDModalitySchedule(1.0), "b": PDDModalitySchedule(1.0, max_t=1.0)} + with pytest.raises(ValueError, match="share max_t"): + build_pdd_sampling_plan([0, 4, 8], schedules) + + +def test_build_pdd_sampling_plan_grid32_eight_blocks_feed_h3_schedulers() -> None: + """The exported contract (Grid32, blocks of 4, shifts 12/3) as the stage builds it.""" + plan = build_pdd_sampling_plan(GRID32_BLOCKS8, { + "video": PDDModalitySchedule(shift=12.0), + "audio": PDDModalitySchedule(shift=3.0) + }) + assert plan.block_sizes == [4] * 8 + for name, shift in (("video", 12.0), ("audio", 3.0)): + expected = _time_shift(pdd_fine_grid(32)[GRID32_BLOCKS8], shift) + torch.testing.assert_close(plan.node_sigmas[name], expected, rtol=0.0, atol=1e-15) + scheduler = MiniMaxH3Scheduler(shift=shift) + scheduler.set_timesteps(sigmas=plan.node_sigmas[name].to(torch.float32)) + assert scheduler.num_inference_steps == 8 + torch.testing.assert_close(scheduler.timesteps, 1.0 - plan.node_sigmas[name][:-1].to(torch.float32)) + + +def _randomized(linear: PDDReplicatedLinear) -> PDDReplicatedLinear: + # ReplicatedLinear storage is uninitialized until a checkpoint loads. + with torch.no_grad(): + linear.weight.normal_() + linear.bias.normal_() + return linear + + +def test_pdd_replicated_linear_head_major_and_state_dict_compatible() -> None: + linear = PDDReplicatedLinear(5, 3, grid_size=4, params_dtype=torch.float32) + assert linear.weight.shape == (12, 5) and linear.bias.shape == (12, ) + assert linear.head_output_size == 3 and linear.grid_size == 4 + # Fusion state is transient: the state dict is a plain linear's. + assert set(linear.state_dict()) == {"weight", "bias"} + + +def test_fused_params_use_normalized_integration_weights() -> None: + torch.manual_seed(2) + grid, channels, in_features = 5, 3, 7 + linear = _randomized(PDDReplicatedLinear(in_features, channels, grid_size=grid, params_dtype=torch.float32)) + weights = PDDModalitySchedule(shift=0.25).integration_weights(pdd_fine_grid(grid)) + start, end = 1, 5 + + actual_weight, actual_bias = linear._fused_params(start, end, weights, torch.float32) + alpha = weights[start:end] + alpha = (alpha / alpha.sum()).float() + expected_weight = torch.einsum("n,nci->ci", alpha, linear.weight.reshape(grid, channels, in_features)[start:end]) + expected_bias = torch.einsum("n,nc->c", alpha, linear.bias.reshape(grid, channels)[start:end]) + torch.testing.assert_close(actual_weight, expected_weight, rtol=1e-6, atol=1e-6) + torch.testing.assert_close(actual_bias, expected_bias, rtol=1e-6, atol=1e-6) + + +def test_forward_fused_block_weighted_mean_of_materialized_heads() -> None: + torch.manual_seed(3) + grid, out_features, in_features = 6, 4, 5 + linear = _randomized(PDDReplicatedLinear(in_features, out_features, grid_size=grid, params_dtype=torch.float32)) + x = torch.randn(2, 3, in_features) + weights = PDDModalitySchedule(shift=3.0).integration_weights(pdd_fine_grid(grid)) + + heads, extra_bias = linear(x) + assert extra_bias is None + heads = heads.unflatten(-1, (grid, out_features)) + for start, end in ((0, grid), (2, 5), (4, 5)): + alpha = (weights[start:end] / weights[start:end].sum()).float() + expected = torch.einsum("n,btnc->btc", alpha, heads[..., start:end, :]) + with linear.fuse(start, end, weights, torch.float32): + fused, fused_extra = linear(x) + assert fused_extra is None + torch.testing.assert_close(fused, expected, rtol=1e-5, atol=1e-5) + # Leaving the context restores the materialized-head forward. + torch.testing.assert_close(linear(x)[0].unflatten(-1, (grid, out_features)), heads) + + +def test_fused_params_bf16_heads_accumulate_in_decoding_precision() -> None: + torch.manual_seed(5) + linear = _randomized(PDDReplicatedLinear(8, 3, grid_size=4, params_dtype=torch.bfloat16)) + weights = PDDModalitySchedule(shift=12.0).integration_weights(pdd_fine_grid(4)) + fused_weight, fused_bias = linear._fused_params(0, 4, weights, torch.float32) + alpha = (weights / weights.sum()).float() + expected = torch.einsum("n,nci->ci", alpha, linear.weight.float().reshape(4, 3, 8)).to(torch.bfloat16) + assert fused_weight.dtype == torch.bfloat16 and fused_bias is not None and fused_bias.dtype == torch.bfloat16 + assert torch.equal(fused_weight, expected) + + +def test_fuse_pdd_heads_fuses_every_modality_head_at_once() -> None: + torch.manual_seed(4) + grid = 4 + linears = { + "video": _randomized(PDDReplicatedLinear(3, 2, grid_size=grid, params_dtype=torch.float32)), + "audio": _randomized(PDDReplicatedLinear(3, 1, grid_size=grid, params_dtype=torch.float32)), + } + weights = { + "video": PDDModalitySchedule(shift=12.0).integration_weights(pdd_fine_grid(grid)), + "audio": PDDModalitySchedule(shift=3.0).integration_weights(pdd_fine_grid(grid)), + } + x = torch.randn(2, 3) + with fuse_pdd_heads(linears, 1, 3, weights, torch.float32): + fused = {name: linear(x)[0] for name, linear in linears.items()} + for name, linear in linears.items(): + with linear.fuse(1, 3, weights[name], torch.float32): + torch.testing.assert_close(fused[name], linear(x)[0]) + assert linear._fusion_state is None + + +def test_fuse_nested_contexts_restore_outer_state() -> None: + linear = _randomized(PDDReplicatedLinear(3, 2, grid_size=4, params_dtype=torch.float32)) + weights = PDDModalitySchedule(shift=1.0).integration_weights(pdd_fine_grid(4)) + with linear.fuse(0, 4, weights, torch.float32): + outer = linear._fusion_state + with linear.fuse(1, 3, weights, torch.float32): + assert linear._fusion_state[:2] == (1, 3) + assert linear._fusion_state is outer + assert linear._fusion_state is None diff --git a/fastvideo/tests/loader/test_minimax_h3_pdd_loading.py b/fastvideo/tests/loader/test_minimax_h3_pdd_loading.py new file mode 100644 index 000000000..808a0a2ed --- /dev/null +++ b/fastvideo/tests/loader/test_minimax_h3_pdd_loading.py @@ -0,0 +1,231 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Loading a PDD-widened MiniMax-H3 ``transformer_ref`` export through the transformer loader. + +A PDD export records ``pdd_steps`` in its transformer ``config.json`` and +stores both output projections widened to that many heads, next to its trained +VSA compression gates. The loader must build the widened heads from the +config, load every tensor exactly, and fail closed when the config and the +weights disagree. Runs on CPU with a tiny synthetic checkpoint. +""" +from __future__ import annotations + +import json +import os +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import save_file + +import fastvideo.envs as envs +from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3ArchConfig, MiniMaxH3Config +from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig +from fastvideo.forward_context import set_forward_context +from fastvideo.layers.pdd import PDDModalitySchedule, PDDReplicatedLinear, pdd_fine_grid +from fastvideo.models.dits import minimax_h3 +from fastvideo.models.loader import component_loader +from fastvideo.pipelines.basic.minimax_h3.packing import MINIMAX_H3_TEXT_TAG, build_packed_sequence, build_row_timesteps +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.platforms import AttentionBackendEnum + +PDD_STEPS = 4 +TINY_ARCH = dict( + num_attention_heads=4, + attention_head_dim=16, + hidden_size=32, + num_layers=2, + num_refiner_layers=1, + ffn_dim=48, + in_channels=4, + audio_in_channels=4, + patch_size=[1, 2, 2], + text_dim=16, + freq_dim=16, + time_embed_hidden_dim=32, + time_embed_dim=24, + rope_freq_dim=2, + rope_theta=10000.0, + norm_eps=1e-5, + qk_norm_eps=1e-5, + final_norm_eps=1e-5, +) +VIDEO_PATCH_DIM = 4 * 1 * 2 * 2 + + +def _released_name(name: str) -> str: + """FastVideo parameter name -> released (diffusers) checkpoint name.""" + for fastvideo_name, released_name in (("time_embedder.fc_in.", "time_embedder.linear_1."), + ("time_embedder.fc_out.", "time_embedder.linear_2."), + (".attn.to_out.", ".attn.to_out.0."), (".ff.fc_in.", ".ff.net.0.proj."), + (".ff.fc_out.", ".ff.net.2.")): + name = name.replace(fastvideo_name, released_name) + return name + + +class _VSAH3Selection: + """What the selector returns when VIDEO_SPARSE_ATTN_H3 resolves on a CUDA host.""" + + @staticmethod + def get_name() -> str: + return "VIDEO_SPARSE_ATTN_H3" + + +@pytest.fixture() +def cpu_loader(monkeypatch, env_overrides): + """Load onto the CPU with SDPA attention, on CPU-only and CUDA hosts alike.""" + monkeypatch.setattr(component_loader, "get_local_torch_device", lambda: torch.device("cpu")) + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA")) + + +@pytest.fixture() +def vsa_gates(monkeypatch, cpu_loader): + """Construct H3 attention as the VSA-H3 backend does, so the compression gates exist on a CPU host.""" + select = minimax_h3.get_attn_backend + + def select_vsa_h3(head_size, dtype, supported_attention_backends=None, **kwargs): + if supported_attention_backends and AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3 in supported_attention_backends: + return _VSAH3Selection + return select(head_size, dtype, supported_attention_backends=supported_attention_backends, **kwargs) + + monkeypatch.setattr(minimax_h3, "get_attn_backend", select_vsa_h3) + + +@pytest.fixture() +def single_process_group(env_overrides): + """A 1x1 sequence/tensor-parallel group, which DistributedAttention needs even on one CPU process.""" + from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel + + # Keep a launcher- or CI-assigned rendezvous: packed CI lanes share the tray's network namespace, and each + # lease gets its own port range. + env_overrides.enter_context(envs.override_external("MASTER_ADDR", os.environ.get("MASTER_ADDR") or "127.0.0.1")) + env_overrides.enter_context(envs.override_external("MASTER_PORT", os.environ.get("MASTER_PORT") or "29591")) + env_overrides.enter_context(envs.override_external("RANK", "0")) + env_overrides.enter_context(envs.override_external("WORLD_SIZE", "1")) + env_overrides.enter_context(envs.override_external("LOCAL_RANK", "0")) + maybe_init_distributed_environment_and_model_parallel(1, 1) + try: + yield + finally: + cleanup_dist_env_and_memory() + + +def _write_transformer_ref(root, *, config_pdd_steps, weight_pdd_steps=PDD_STEPS, with_gates=True): + """A synthetic BF16 export: every tensor the (widened) model defines, in the released naming.""" + arch = MiniMaxH3ArchConfig(**{**TINY_ARCH, "patch_size": tuple(TINY_ARCH["patch_size"])}, + pdd_steps=weight_pdd_steps) + with torch.device("meta"): + reference = minimax_h3.MiniMaxH3Transformer3DModel(MiniMaxH3Config(arch_config=arch), hf_config={}) + generator = torch.Generator().manual_seed(0) + tensors = {} + for name, parameter in reference.named_parameters(): + if with_gates or "to_gate_compress" not in name: + tensors[_released_name(name)] = (torch.randn(parameter.shape, generator=generator) * 0.1).to(torch.bfloat16) + path = root / "transformer_ref" + path.mkdir() + save_file(tensors, str(path / "diffusion_pytorch_model.safetensors")) + config = {"_class_name": "MiniMaxH3Transformer3DModel", "_diffusers_version": "0.36.0.dev0", **TINY_ARCH} + if config_pdd_steps is not None: + config["pdd_steps"] = config_pdd_steps + (path / "config.json").write_text(json.dumps(config)) + return path, tensors + + +def _load(path): + args = SimpleNamespace( + pipeline_config=MiniMaxH3PipelineConfig(), + override_transformer_cls_name=None, + model_paths={}, + init_weights_from_safetensors=None, + hsdp_replicate_dim=1, + hsdp_shard_dim=1, + dit_cpu_offload=True, + pin_cpu_memory=False, + use_fsdp_inference=False, + training_mode=False, + enable_torch_compile=False, + torch_compile_kwargs=None, + inference_torch_compile=False, + VSA_tile_size=128, + lora_path=None, + lora_strength=1.0, + inference_mode=True, + dit_layerwise_offload=False, + ) + return component_loader.TransformerLoader().load(str(path), args) + + +def test_load_widened_heads_and_trained_gates_exactly(tmp_path, vsa_gates): + path, tensors = _write_transformer_ref(tmp_path, config_pdd_steps=PDD_STEPS) + model = _load(path) + + assert model.pdd_steps == PDD_STEPS + for head, rows in (("proj_out", VIDEO_PATCH_DIM), ("audio_proj_out", TINY_ARCH["audio_in_channels"])): + linear = getattr(model, head) + assert isinstance(linear, PDDReplicatedLinear) + assert linear.grid_size == PDD_STEPS and linear.head_output_size == rows + assert linear.weight.shape == (PDD_STEPS * rows, TINY_ARCH["hidden_size"]) + # Output projections stay FP32 at inference, like the released heads. + assert linear.weight.dtype == torch.float32 + gate_names = [name for name, _ in model.named_parameters() if "to_gate_compress" in name] + assert gate_names == [f"transformer_blocks.{index}.attn.to_gate_compress.weight" for index in range(2)] + + loaded = dict(model.named_parameters()) + assert {_released_name(name) for name in loaded} == set(tensors) + for name, parameter in loaded.items(): + assert torch.equal(parameter, tensors[_released_name(name)].to(parameter.dtype)), name + with torch.no_grad(): + assert all(block.attn._gate_active() for block in model.transformer_blocks) + + +def test_load_export_without_gates_disables_gate_branch(tmp_path, vsa_gates): + path, _ = _write_transformer_ref(tmp_path, config_pdd_steps=PDD_STEPS, with_gates=False) + model = _load(path) + with torch.no_grad(): + assert not any(block.attn._gate_active() for block in model.transformer_blocks) + + +@pytest.mark.parametrize("config_pdd_steps,weight_pdd_steps", [(None, PDD_STEPS), (PDD_STEPS, None), (8, PDD_STEPS)]) +def test_load_mismatched_config_and_weight_grids(tmp_path, vsa_gates, config_pdd_steps, weight_pdd_steps): + path, _ = _write_transformer_ref(tmp_path, config_pdd_steps=config_pdd_steps, weight_pdd_steps=weight_pdd_steps) + # ReplicatedLinear.weight_loader rejects the shape; never a silent partial load. + with pytest.raises((AssertionError, RuntimeError), match="size"): + _load(path) + + +def test_fuse_pdd_block_weighted_mean_of_widened_heads(tmp_path, cpu_loader, single_process_group): + """End to end through the DiT: a fused block equals the weighted mean of its heads' outputs.""" + path, _ = _write_transformer_ref(tmp_path, config_pdd_steps=PDD_STEPS, with_gates=False) + model = _load(path) + layout = build_packed_sequence(torch.full((3, ), MINIMAX_H3_TEXT_TAG, dtype=torch.long), 2, 4, 4, 2, (1, 2, 2)) + unique, inverse = build_row_timesteps(layout, 0.3, 0.4, 0.999, 1.0) + generator = torch.Generator().manual_seed(1) + inputs = dict( + hidden_states=torch.randn(1, int(layout.video_indices.numel()), VIDEO_PATCH_DIM, generator=generator), + audio_hidden_states=torch.randn(1, + int(layout.audio_indices.numel()), + TINY_ARCH["audio_in_channels"], + generator=generator), + encoder_hidden_states=torch.randn(1, 3, TINY_ARCH["text_dim"], generator=generator), + timestep=unique, + timestep_indices=inverse, + token_tags=layout.token_tags, + position_ids=layout.position_ids, + video_indices=layout.video_indices, + audio_indices=layout.audio_indices, + text_indices=layout.text_indices, + ) + weights = { + "video": PDDModalitySchedule(shift=12.0).integration_weights(pdd_fine_grid(PDD_STEPS)), + "audio": PDDModalitySchedule(shift=3.0).integration_weights(pdd_fine_grid(PDD_STEPS)), + } + with torch.no_grad(), set_forward_context(current_timestep=0, attn_metadata=None, + forward_batch=ForwardBatch(data_type="dummy")): + heads = model(**inputs) + with model.fuse_pdd_block(1, 3, weights, torch.float32): + fused = model(**inputs) + for name, all_heads, block in zip(("video", "audio"), heads, fused, strict=True): + all_heads = all_heads.unflatten(-1, (PDD_STEPS, -1)) + alpha = weights[name][1:3] / weights[name][1:3].sum() + expected = torch.einsum("n,brnc->brc", alpha.float(), all_heads[..., 1:3, :].float()) + assert block.shape == expected.shape + torch.testing.assert_close(block.float(), expected, rtol=1e-4, atol=1e-5) diff --git a/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa.py b/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa.py index 439da0c0e..611e04fc6 100644 --- a/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa.py +++ b/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa.py @@ -62,10 +62,9 @@ def _dit_shape(spec: dict) -> tuple[int, int, int]: def _torch_meta(spec: dict, tile_size: int = 256, sparsity: float = 0.0): return MiniMaxH3VSAMetadataBuilder().build( current_timestep=0, - raw_latent_shape=spec["raw_latent_shape"], patch_size=spec["patch_size"], VSA_sparsity=sparsity, - prefix_segments=spec["prefix_segments"], + packed_segments=(*spec["prefix_segments"], spec["raw_latent_shape"]), device=_CPU, tile_size=tile_size, ) @@ -112,7 +111,7 @@ def test_routing_mask_parity_exempt_and_compete() -> None: k_vid = compute_topk(0.9, v) for exempt in (True, False): mlx_mask = build_block_mask(scores, p, v, 0.9, exempt=exempt) - torch_mask = _build_block_mask(torch.tensor(scores), p, v, 0.9, exempt=exempt).numpy() + torch_mask = _build_block_mask(torch.tensor(scores), p, 0.9, exempt, ((p, p + v), ), (0.9, )).numpy() assert mlx_mask[:, :, :p].all() if exempt: assert mlx_mask[..., :p].all() diff --git a/fastvideo/tests/stages/test_minimax_h3_pdd_contract.py b/fastvideo/tests/stages/test_minimax_h3_pdd_contract.py new file mode 100644 index 000000000..22f892919 --- /dev/null +++ b/fastvideo/tests/stages/test_minimax_h3_pdd_contract.py @@ -0,0 +1,472 @@ +# SPDX-License-Identifier: Apache-2.0 +"""How a FastH3 Ref2VA PDD checkpoint's ``fastvideo_inference.json`` sets a run's settings. + +``MiniMaxH3PipelineConfig.resolve_checkpoint_settings`` runs inside +``FastVideoArgs.__post_init__``. It fills each unset run setting from the file, +rejects an explicit value that differs for a setting fixed by training (attention +backend, VSA tile size, PDD partition), keeps an explicit value that differs for +a tunable setting (VSA sparsity, reference keep rate) with a warning, and +rejects a file that is incomplete or disagrees with the checkpoint file that owns +a value. ``MiniMaxH3PipelineConfig.apply_request_constraints`` then fills an +unset request step count with the block count and rejects a different one. In +the worker, ``MiniMaxH3BasePipeline`` checks its transformer before loading +weights: against those settings, and, when the transformer carries VSA +compression gates, against the run's attention backend. A DMD export's schedule +comes from its ``dmd_denoising_steps`` (see test_minimax_h3_distilled_schedule.py). +""" +from __future__ import annotations + +import json +import logging +import re +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import save_file + +from fastvideo import envs +from fastvideo.api.compat import normalize_generation_request +from fastvideo.api.parser import parse_config +from fastvideo.api.sampling_param import SamplingParam +from fastvideo.api.schema import GenerationRequest +from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler +from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import ( + MiniMaxH3ModularPipeline, + MiniMaxH3Ref2VAModularPipeline, +) +from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase + +GRID32_BLOCKS8 = [0, 4, 8, 12, 16, 20, 24, 28, 32] +# The OmniRef PDD-8 export's fastvideo_inference.json, verbatim. +PDD_FILE = { + "attention_backend": "VIDEO_SPARSE_ATTN_H3", + "audio_scheduler_shift": 3.0, + "base_model_revision": "hf://MiniMaxAI/MiniMax-H3@9bfb6693f2cf6de171db46d1aa586f67d773a1da", + "conditioning": "fixed_ordered_references_target_only_flow", + "grid_max_t": 0.999, + "guidance_scale": 1.0, + "model_type": "ref2va", + "num_inference_steps": 8, + "pdd_step_indices": GRID32_BLOCKS8, + "pdd_steps": 32, + "schema": "fasth3-inference-contract-v1", + "schema_version": "fasth3-inference-contract-v1", + "transformer_component": "transformer_ref", + "transformer_forwards": 8, + "video_scheduler_shift": 12.0, + "vsa_ref_keep_rate": 0.1, + "vsa_ref_policy": "p2_multi_region", + "vsa_sparsity": 0.9, + "vsa_tile_size": 128, +} +# A DMD export's fastvideo_inference.json, which has no PDD fields. +DMD_FILE = { + "schema_version": "fasth3-inference-contract-v1", + "dmd_denoising_steps": [999, 874, 749, 624, 500, 375, 250, 125], + "num_inference_steps": 9, + "transformer_forwards": 8, + "video_scheduler_shift": 12.0, + "audio_scheduler_shift": 3.0, +} + + +@pytest.fixture(autouse=True) +def _unset_backend_env_var(env_overrides): + """Start every test without FASTVIDEO_ATTENTION_BACKEND, which FastVideoArgs reads as an explicit backend.""" + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override(None)) + + +def _checkpoint(root: Path, inference_file: dict | None, *, pdd_steps: int | None = 32) -> str: + """Write a Ref2VA checkpoint's JSON files: the inference file and the configs that own its repeated values. + + The transformer_ref config.json sets ``pdd_steps`` unless it is None; the + scheduler configs hold the video/audio shifts 12.0/3.0. No weights are written. + """ + files = { + "transformer_ref/config.json": {} if pdd_steps is None else {"pdd_steps": pdd_steps}, + "scheduler/scheduler_config.json": {"shift": 12.0}, + "audio_scheduler/scheduler_config.json": {"shift": 3.0}, + } + if inference_file is not None: + files["fastvideo_inference.json"] = inference_file + for relative_path, content in files.items(): + path = root / relative_path + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(content)) + return str(root) + + +def _run_args(model_path: str, config_settings: dict | None = None, **run_settings) -> FastVideoArgs: + """Build a run's FastVideoArgs; construction resolves the checkpoint's settings.""" + pipeline_config = MiniMaxH3PipelineConfig(**(config_settings or {})) + return FastVideoArgs(model_path=model_path, pipeline_config=pipeline_config, **run_settings) + + +def _pipeline(model_path: str, cls=MiniMaxH3Ref2VAModularPipeline): + """A pipeline without loaded weights: the checkpoint path, its schedulers, and the Ref2VA flag __init__ sets.""" + pipeline = object.__new__(cls) + pipeline.model_path = model_path + pipeline._ref2va = cls._ref2va_default + pipeline.modules = { + "scheduler": MiniMaxH3Scheduler(shift=12.0), + "audio_scheduler": MiniMaxH3Scheduler(shift=3.0), + } + return pipeline + + +# ------------------------------------------------- resolve_checkpoint_settings: run settings on a PDD checkpoint + + +def test_resolve_checkpoint_settings_unset_run_settings(tmp_path): + args = _run_args(_checkpoint(tmp_path, PDD_FILE)) + assert args.attention_backend == "VIDEO_SPARSE_ATTN_H3" + assert args.VSA_tile_size == 128 + assert args.VSA_sparsity == 0.9 + assert args.pipeline_config.pdd_step_indices == tuple(GRID32_BLOCKS8) + assert args.pipeline_config.vsa_ref_keep_rate == 0.1 + + +def test_resolve_checkpoint_settings_matching_explicit_settings(tmp_path, caplog): + with caplog.at_level(logging.WARNING): + args = _run_args(_checkpoint(tmp_path, PDD_FILE), { + "pdd_step_indices": GRID32_BLOCKS8, + "vsa_ref_keep_rate": 0.1 + }, + attention_backend="VIDEO_SPARSE_ATTN_H3", + VSA_tile_size=128, + VSA_sparsity=0.9) + assert args.pipeline_config.pdd_step_indices == tuple(GRID32_BLOCKS8) + assert "trained with" not in caplog.text + + +def test_resolve_checkpoint_settings_repeated_call(tmp_path): + args = _run_args(_checkpoint(tmp_path, PDD_FILE)) + args.pipeline_config.resolve_checkpoint_settings(args) + assert (args.attention_backend, args.VSA_tile_size, args.VSA_sparsity) == ("VIDEO_SPARSE_ATTN_H3", 128, 0.9) + assert args.pipeline_config.pdd_step_indices == tuple(GRID32_BLOCKS8) + assert args.pipeline_config.vsa_ref_keep_rate == 0.1 + + +@pytest.mark.parametrize("config_settings,run_settings,match", [ + ({}, {"attention_backend": "FLASH_ATTN"}, + "trained with attention_backend=VIDEO_SPARSE_ATTN_H3; this run requests FLASH_ATTN"), + ({}, {"VSA_tile_size": 64}, "trained with VSA_tile_size=128; this run requests 64"), + ({"pdd_step_indices": (0, 8, 16, 24, 32)}, {}, re.escape("this run sets [0, 8, 16, 24, 32]")), +], ids=["backend", "tile_size", "partition"]) +def test_resolve_checkpoint_settings_explicit_training_fixed_conflict(tmp_path, config_settings, run_settings, + match): + with pytest.raises(ValueError, match=match): + _run_args(_checkpoint(tmp_path, PDD_FILE), config_settings, **run_settings) + + +def test_resolve_checkpoint_settings_env_backend_conflict(tmp_path, env_overrides): + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("FLASH_ATTN")) + with pytest.raises(ValueError, match="this run requests FLASH_ATTN"): + _run_args(_checkpoint(tmp_path, PDD_FILE)) + + +def test_resolve_checkpoint_settings_explicit_sparsity_override(tmp_path, caplog): + with caplog.at_level(logging.WARNING): + args = _run_args(_checkpoint(tmp_path, PDD_FILE), VSA_sparsity=0.8) + assert args.VSA_sparsity == 0.8 + assert "trained with VSA_sparsity=0.9; this run uses 0.8" in caplog.text + + +def test_resolve_checkpoint_settings_explicit_keep_rate_override(tmp_path, caplog): + with caplog.at_level(logging.WARNING): + args = _run_args(_checkpoint(tmp_path, PDD_FILE), {"vsa_ref_keep_rate": 0.2}) + assert args.pipeline_config.vsa_ref_keep_rate == 0.2 + assert "trained with vsa_ref_keep_rate=0.1; this run uses 0.2" in caplog.text + + +@pytest.mark.parametrize("config_settings,run_settings,match", [ + ({}, {"VSA_sparsity": 1.0}, re.escape("VSA_sparsity must be in [0, 1), got 1.0")), + ({"vsa_ref_keep_rate": 1.0}, {}, re.escape("vsa_ref_keep_rate must be in (0, 1), got 1.0")), +], ids=["sparsity", "keep_rate"]) +def test_resolve_checkpoint_settings_tunable_out_of_range(tmp_path, config_settings, run_settings, match): + with pytest.raises(ValueError, match=match): + _run_args(_checkpoint(tmp_path, PDD_FILE), config_settings, **run_settings) + + +def test_resolve_checkpoint_settings_dmd_steps_on_pdd_checkpoint(tmp_path): + with pytest.raises(ValueError, match="dmd_denoising_steps must be unset"): + _run_args(_checkpoint(tmp_path, PDD_FILE), {"dmd_denoising_steps": [999, 749, 500, 250]}) + + +# ------------------------------------------------- resolve_checkpoint_settings: PDD file validation + + +def test_resolve_checkpoint_settings_missing_field(tmp_path): + inference_file = {key: value for key, value in PDD_FILE.items() if key != "vsa_sparsity"} + with pytest.raises(ValueError, match=re.escape("missing ['vsa_sparsity']")): + _run_args(_checkpoint(tmp_path, inference_file)) + + +@pytest.mark.parametrize("pin", ["", None, "9bfb6693", "MiniMaxAI/MiniMax-H3@9bfb6693", "hf://MiniMaxAI/MiniMax-H3"]) +def test_resolve_checkpoint_settings_unparsable_base_model_revision(tmp_path, pin): + with pytest.raises(ValueError, match="must be hf://@"): + _run_args(_checkpoint(tmp_path, {**PDD_FILE, "base_model_revision": pin})) + + +def test_resolve_checkpoint_settings_unknown_field(tmp_path): + with pytest.raises(ValueError, match=re.escape("unknown ['dmd_denoising_steps']")): + _run_args(_checkpoint(tmp_path, {**PDD_FILE, "dmd_denoising_steps": [999, 500]})) + + +@pytest.mark.parametrize("field,value", [ + ("schema_version", "fasth3-inference-contract-v2"), + ("schema", "fasth3-inference-contract-v2"), + ("model_type", "t2va"), + ("transformer_component", "transformer"), + ("conditioning", "noised_references"), + ("guidance_scale", 2.0), + ("attention_backend", "FLASH_ATTN"), + ("vsa_ref_policy", "dense"), + ("grid_max_t", 1.0), +]) +def test_resolve_checkpoint_settings_fixed_field_other_value(tmp_path, field, value): + with pytest.raises(ValueError, match=re.escape(f"{field}={value!r} is unsupported")): + _run_args(_checkpoint(tmp_path, {**PDD_FILE, field: value})) + + +@pytest.mark.parametrize("field,value,match", [ + ("pdd_steps", 16, "pdd_steps=16 disagrees with transformer_ref/config.json pdd_steps=32"), + ("video_scheduler_shift", 10.0, + "video_scheduler_shift=10.0 disagrees with scheduler/scheduler_config.json shift=12.0"), + ("audio_scheduler_shift", 4.0, + "audio_scheduler_shift=4.0 disagrees with audio_scheduler/scheduler_config.json shift=3.0"), + ("num_inference_steps", 9, "num_inference_steps=9 disagrees with the 8 blocks of pdd_step_indices"), + ("transformer_forwards", 4, "transformer_forwards=4 disagrees with the 8 blocks of pdd_step_indices"), +]) +def test_resolve_checkpoint_settings_duplicate_disagrees_with_owner(tmp_path, field, value, match): + with pytest.raises(ValueError, match=match): + _run_args(_checkpoint(tmp_path, {**PDD_FILE, field: value})) + + +@pytest.mark.parametrize("owner", ["transformer_ref/config.json", "audio_scheduler/scheduler_config.json"]) +def test_resolve_checkpoint_settings_missing_owner_file(tmp_path, owner): + model_path = _checkpoint(tmp_path, PDD_FILE) + (tmp_path / owner).unlink() + with pytest.raises(ValueError, match=f"has fastvideo_inference.json but no {owner}"): + _run_args(model_path) + + +@pytest.mark.parametrize("field,value,match", [ + ("vsa_tile_size", 96, "vsa_tile_size=96 must be one of"), + ("vsa_sparsity", "0.9", "vsa_sparsity must be in \\[0, 1\\), got '0.9'"), + ("vsa_ref_keep_rate", 1.0, "vsa_ref_keep_rate must be in \\(0, 1\\), got 1.0"), +]) +def test_resolve_checkpoint_settings_invalid_trained_vsa_value(tmp_path, field, value, match): + with pytest.raises(ValueError, match=match): + _run_args(_checkpoint(tmp_path, {**PDD_FILE, field: value})) + + +def test_resolve_checkpoint_settings_transformer_config_without_pdd_steps(tmp_path): + with pytest.raises(ValueError, match="pdd_steps=32 disagrees with transformer_ref/config.json pdd_steps=None"): + _run_args(_checkpoint(tmp_path, PDD_FILE, pdd_steps=None)) + + +@pytest.mark.parametrize("indices", [ + [0, 8, 4, 12, 16, 20, 24, 28, 32], + [1, 4, 8, 12, 16, 20, 24, 28, 32], + [0, 4, 8, 12, 16, 20, 24, 28, 31], + [0, 4.0, 8, 12, 16, 20, 24, 28, 32], + [32], +], ids=["decreasing", "starts_after_0", "ends_before_pdd_steps", "float_node", "single_node"]) +def test_resolve_checkpoint_settings_malformed_partition(tmp_path, indices): + with pytest.raises(ValueError, match=re.escape(f"pdd_step_indices={indices!r} must increase strictly from 0 to " + "pdd_steps=32")): + _run_args(_checkpoint(tmp_path, {**PDD_FILE, "pdd_step_indices": indices})) + + +# ------------------------------------------------- resolve_checkpoint_settings: checkpoints without PDD fields + + +@pytest.mark.parametrize("inference_file", [None, DMD_FILE], ids=["base", "dmd"]) +def test_resolve_checkpoint_settings_non_pdd_checkpoint(tmp_path, inference_file): + args = _run_args(_checkpoint(tmp_path, inference_file, pdd_steps=None)) + assert (args.attention_backend, args.VSA_tile_size, args.VSA_sparsity) == (None, 256, 0.0) + config = args.pipeline_config + assert (config.pdd_step_indices, config.vsa_ref_keep_rate, config.dmd_denoising_steps) == (None, None, None) + + +@pytest.mark.parametrize("inference_file", [None, DMD_FILE], ids=["base", "dmd"]) +@pytest.mark.parametrize("config_settings", [{"pdd_step_indices": (0, 16, 32)}, {"vsa_ref_keep_rate": 0.1}], + ids=["partition", "keep_rate"]) +def test_resolve_checkpoint_settings_pdd_setting_on_non_pdd_checkpoint(tmp_path, inference_file, config_settings): + with pytest.raises(ValueError, match="applies only to FastH3 PDD checkpoints"): + _run_args(_checkpoint(tmp_path, inference_file, pdd_steps=None), config_settings) + + +# ------------------------------------------------- apply_request_constraints: the step count of a request + + +def _pdd_config(): + return MiniMaxH3PipelineConfig(pdd_step_indices=tuple(GRID32_BLOCKS8)) + + +@pytest.mark.parametrize("sampling", [{}, {"num_inference_steps": 8}]) +def test_apply_request_constraints_pdd_unset_or_matching_steps(sampling): + request = parse_config(GenerationRequest, {"prompt": "fox", "sampling": sampling}) + sampling_param = _pdd_config().apply_request_constraints(request, SamplingParam(num_inference_steps=50)) + assert sampling_param.num_inference_steps == 8 + + +def test_apply_request_constraints_pdd_conflicting_steps(): + request = parse_config(GenerationRequest, {"prompt": "fox", "sampling": {"num_inference_steps": 50}}) + with pytest.raises(ValueError, match="runs exactly 8 transformer forwards.*num_inference_steps=50"): + _pdd_config().apply_request_constraints(request, SamplingParam()) + + +def test_apply_request_constraints_pdd_python_request_counts_schema_default_as_set(): + """A GenerationRequest built in Python carries the schema default 50 as a set value.""" + request = normalize_generation_request(GenerationRequest(prompt="fox")) + with pytest.raises(ValueError, match="num_inference_steps=50. Pass num_inference_steps=8"): + _pdd_config().apply_request_constraints(request, SamplingParam()) + + +def test_apply_request_constraints_non_pdd_checkpoint_keeps_request_steps(): + request = parse_config(GenerationRequest, {"prompt": "fox", "sampling": {"num_inference_steps": 50}}) + sampling_param = SamplingParam(num_inference_steps=50) + assert MiniMaxH3PipelineConfig().apply_request_constraints(request, sampling_param) is sampling_param + assert sampling_param.num_inference_steps == 50 + + +# ------------------------------------------------- worker pipeline: initialize_pipeline and the transformer check + + +def test_initialize_pipeline_pdd_checkpoint_keeps_resolved_settings(tmp_path): + model_path = _checkpoint(tmp_path, PDD_FILE) + args = _run_args(model_path) + _pipeline(model_path).initialize_pipeline(args) + config = args.pipeline_config + assert config.dit_config.arch_config.pdd_steps == 32 + assert config.pdd_step_indices == tuple(GRID32_BLOCKS8) + assert config.dmd_denoising_steps is None + + +def test_initialize_pipeline_dmd_checkpoint_applies_rungs(tmp_path): + model_path = _checkpoint(tmp_path, DMD_FILE, pdd_steps=None) + args = _run_args(model_path) + _pipeline(model_path).initialize_pipeline(args) + assert args.pipeline_config.dmd_denoising_steps == DMD_FILE["dmd_denoising_steps"] + assert args.pipeline_config.pdd_step_indices is None + + +def test_check_pdd_transformer_widened_transformer_without_pdd_file(tmp_path): + model_path = _checkpoint(tmp_path, None, pdd_steps=32) + with pytest.raises(ValueError, match="sets pdd_steps, but .* has no PDD fastvideo_inference.json"): + _pipeline(model_path)._check_pdd_transformer(_run_args(model_path)) + + +def test_check_pdd_transformer_pdd_checkpoint_on_t2va_pipeline(tmp_path): + model_path = _checkpoint(tmp_path, PDD_FILE) + with pytest.raises(ValueError, match="select MiniMaxH3Ref2VAModularPipeline"): + _pipeline(model_path, MiniMaxH3ModularPipeline)._check_pdd_transformer(_run_args(model_path)) + + +def test_check_pdd_transformer_pdd_checkpoint_on_ref2va_pipeline(tmp_path): + model_path = _checkpoint(tmp_path, PDD_FILE) + _pipeline(model_path)._check_pdd_transformer(_run_args(model_path)) + + +# ------------------------------------------------- worker pipeline: _load_config checks, before any component loads + +_GATE = "transformer_blocks.0.attn.to_gate_compress.weight" + + +def _write_gates(transformer_dir: Path, *, indexed: bool) -> None: + """Give a transformer checkpoint one VSA compression gate, listed in a shard index or in a safetensors file.""" + if indexed: + (transformer_dir / "diffusion_pytorch_model.safetensors.index.json").write_text( + json.dumps({"weight_map": {_GATE: "diffusion_pytorch_model-00001-of-00001.safetensors"}})) + else: + save_file({_GATE: torch.zeros(2, 2)}, str(transformer_dir / "diffusion_pytorch_model.safetensors")) + + +@pytest.fixture +def hub_snapshot(tmp_path, monkeypatch): + """Resolve the Hub repo id "org/fasth3" to tmp_path, as maybe_download_model does after downloading.""" + from fastvideo.pipelines import composed_pipeline_base + + downloads = [] + + def download(model_path, **kwargs): + downloads.append(model_path) + return str(tmp_path) if model_path == "org/fasth3" else model_path + + monkeypatch.setattr(composed_pipeline_base, "maybe_download_model", download) + monkeypatch.setattr(composed_pipeline_base, "verify_model_config_and_directory", + lambda model_path, **kwargs: {"_class_name": "MiniMaxH3ModularPipeline"}) + return downloads + + +@pytest.mark.parametrize("indexed", [True, False], ids=["shard_index", "safetensors_header"]) +@pytest.mark.parametrize("inference_file", [None, DMD_FILE], ids=["base", "dmd"]) +def test_load_config_gated_non_pdd_transformer_requires_vsa(tmp_path, hub_snapshot, inference_file, indexed): + """The gates are read from the resolved Hub snapshot. A PDD checkpoint's backend is checked by + resolve_checkpoint_settings instead (test_resolve_checkpoint_settings_explicit_training_fixed_conflict).""" + model_path = _checkpoint(tmp_path, inference_file, pdd_steps=None) + _write_gates(tmp_path / "transformer_ref", indexed=indexed) + pipeline = _pipeline("org/fasth3") + pipeline.fastvideo_args = _run_args(model_path, attention_backend="FLASH_ATTN") + with pytest.raises(ValueError, match="its transformer_ref carries VSA compression gates.*requests FLASH_ATTN"): + pipeline._load_config(pipeline.model_path) + assert hub_snapshot == ["org/fasth3"] and pipeline.model_path == model_path + pipeline.fastvideo_args.attention_backend = "VIDEO_SPARSE_ATTN_H3" + assert pipeline._load_config(pipeline.model_path) == {"_class_name": "MiniMaxH3ModularPipeline"} + + +@pytest.mark.parametrize("pdd_steps,gated,match", [ + (32, False, "has no PDD fastvideo_inference.json"), + (None, True, "carries VSA compression gates"), +], ids=["widened_without_pdd_file", "gated_with_other_backend"]) +def test_load_modules_transformer_check_before_loading_unless_supplied(tmp_path, monkeypatch, hub_snapshot, + pdd_steps, gated, match): + """A transformer that fails its check raises before any component loads; a transformer the caller already + built is not rebuilt, so its checkpoint is not checked.""" + loads = [] + + def base_load_modules(self, fastvideo_args, loaded_modules=None): + self._load_config(self.model_path) + loads.append(loaded_modules) + return {} + + monkeypatch.setattr(ComposedPipelineBase, "load_modules", base_load_modules) + model_path = _checkpoint(tmp_path, None, pdd_steps=pdd_steps) + if gated: + _write_gates(tmp_path / "transformer_ref", indexed=True) + pipeline = _pipeline(model_path) + args = _run_args(model_path, attention_backend="FLASH_ATTN", h3_sequential_load=False) + pipeline.fastvideo_args = args + with pytest.raises(ValueError, match=match): + pipeline.load_modules(args) + assert loads == [] + supplied = {"transformer": object()} + pipeline.load_modules(args, supplied) + assert loads == [supplied] + + +def test_checkpoint_facts_read_once_per_resolved_path(tmp_path, monkeypatch, hub_snapshot): + """_load_config (gate check) and initialize_pipeline (DMD rungs) share one parse of the inference file.""" + from fastvideo.pipelines.basic.minimax_h3 import minimax_h3_pipeline + + parsed = [] + + def recording_loads(text): + parsed.append(json.loads(text)) + return parsed[-1] + + monkeypatch.setattr(minimax_h3_pipeline, "json", SimpleNamespace(loads=recording_loads)) + model_path = _checkpoint(tmp_path, DMD_FILE, pdd_steps=None) + args = _run_args(model_path) + pipeline = _pipeline("org/fasth3") + pipeline.fastvideo_args = args + pipeline._load_config(pipeline.model_path) + pipeline.initialize_pipeline(args) + assert args.pipeline_config.dmd_denoising_steps == DMD_FILE["dmd_denoising_steps"] + assert parsed.count(DMD_FILE) == 1 diff --git a/fastvideo/tests/stages/test_minimax_h3_pdd_denoising.py b/fastvideo/tests/stages/test_minimax_h3_pdd_denoising.py new file mode 100644 index 000000000..d49706abb --- /dev/null +++ b/fastvideo/tests/stages/test_minimax_h3_pdd_denoising.py @@ -0,0 +1,327 @@ +# SPDX-License-Identifier: Apache-2.0 +"""PDD sampling and Ref2VA reference-video VSA in the MiniMax-H3 denoising stage. + +A PDD-widened transformer samples fused blocks of its fine grid: the stage +partitions the grid (the checkpoint contract's ``pdd_step_indices``), hands +each scheduler its modality's node sigmas, starts the targets at the first +node's noise level, and fuses one block of heads per forward. The ordinary +Euler step over two node sigmas then equals the block advance. With +``vsa_ref_keep_rate`` set, every reference video is its own sparse VSA region. +All checks are CPU-only; the transformer is a tiny stand-in with real widened +heads. +""" +from __future__ import annotations + +import contextlib +from types import SimpleNamespace + +import pytest +import torch + +from fastvideo.attention.backends.video_sparse_attn_h3 import MiniMaxH3VSAMetadataBuilder +from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig +from fastvideo.forward_context import get_forward_context +from fastvideo.layers.pdd import ( + PDD_GRID_MAX_T, + PDDModalitySchedule, + PDDReplicatedLinear, + fuse_pdd_heads, + pdd_fine_grid, + shifted_noise_amount, +) +from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler +from fastvideo.pipelines.basic.minimax_h3.packing import ( + MINIMAX_H3_KEYFRAME_NOISE_AUG, + MINIMAX_H3_TEXT_TAG, + build_ref2va_packed_sequence, +) +from fastvideo.pipelines.basic.minimax_h3.reference import MiniMaxH3PreparedReference +from fastvideo.pipelines.basic.minimax_h3.stages import minimax_h3_denoising as denoising +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import MINIMAX_H3_LAYOUT_KEY +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + +PATCH = (1, 2, 2) +VIDEO_CHANNELS = 96 # 24 latent channels x (1, 2, 2) patch +AUDIO_CHANNELS = 32 +GRID32_BLOCKS8 = tuple(range(0, 33, 4)) +# A 4-block partition of an 8-interval grid, for checks that need a short plan. +GRID8_BLOCKS4 = (0, 2, 4, 6, 8) + + +def _reference(kind: str, *, frames: int = 1, height: int = 4, width: int = 4, audio: int = 0): + return MiniMaxH3PreparedReference(media_type=kind, + has_audio=audio > 0, + num_latent_frames=frames, + latent_height=height, + latent_width=width, + num_audio_latents=audio) + + +def _ref2va_layout(references, *, text=5, frames=3, height=4, width=6, audio=3): + tags = torch.full((text, ), MINIMAX_H3_TEXT_TAG, dtype=torch.long) + return build_ref2va_packed_sequence(tags, references, frames, height, width, audio, PATCH) + + +def _stage() -> denoising.MiniMaxH3DenoisingStage: + return denoising.MiniMaxH3DenoisingStage( + transformer=SimpleNamespace(), + scheduler=MiniMaxH3Scheduler(shift=12.0), + audio_scheduler=MiniMaxH3Scheduler(shift=3.0), + ) + + +def _args(**pipeline_fields): + return SimpleNamespace(pipeline_config=MiniMaxH3PipelineConfig(**pipeline_fields), + dit_cpu_offload=False, + dit_layerwise_offload=False, + use_fsdp_inference=False, + VSA_tile_size=128) + + +def _plan(stage, args): + return stage._pdd_sampling_plan(args, torch.device("cpu")) + + +def test_pdd_sampling_plan_unset_step_indices(): + assert _plan(_stage(), _args()) is None + + +def test_pdd_sampling_plan_uses_contract_partition_and_modality_shifts(): + plan = _plan(_stage(), _args(pdd_step_indices=GRID32_BLOCKS8)) + assert plan.indices.tolist() == list(GRID32_BLOCKS8) + base = pdd_fine_grid(32)[list(GRID32_BLOCKS8)] + torch.testing.assert_close(plan.node_sigmas["video"], shifted_noise_amount(base, 12.0), rtol=0, atol=0) + torch.testing.assert_close(plan.node_sigmas["audio"], shifted_noise_amount(base, 3.0), rtol=0, atol=0) + + +def test_step_pdd_block_equals_fused_block_advance(): + """x_end = x_start + sum_j w_j * dx/dsigma_j with the raw H3 output m = x0 - eps = -dx/dsigma.""" + torch.manual_seed(0) + stage = _stage() + plan = _plan(stage, _args(pdd_step_indices=GRID8_BLOCKS4)) + stage.scheduler.set_timesteps(sigmas=plan.node_sigmas["video"].to(torch.float32)) + stage.audio_scheduler.set_timesteps(sigmas=plan.node_sigmas["audio"].to(torch.float32)) + video, audio = torch.randn(5, 6), torch.randn(4, 3) + for step in range(plan.num_steps): + start, end = plan.block(step) + fused_video, fused_audio = torch.randn(5, 6), torch.randn(4, 3) + stepped_video = stage.scheduler.step(fused_video, stage.scheduler.timesteps[step], video, return_dict=False)[0] + stepped_audio = stage.audio_scheduler.step(fused_audio, + stage.audio_scheduler.timesteps[step], + audio, + return_dict=False)[0] + for name, state, fused, stepped in (("video", video, fused_video, stepped_video), + ("audio", audio, fused_audio, stepped_audio)): + total = plan.integration_weights[name][start:end].sum() + expected = state.double() - total * fused.double() + torch.testing.assert_close(stepped.double(), expected, rtol=1e-5, atol=1e-5) + video, audio = stepped_video, stepped_audio + + +def test_scale_pdd_initial_noise_starts_targets_at_first_node_sigma(): + plan = _plan(_stage(), _args(pdd_step_indices=GRID8_BLOCKS4)) + layout = SimpleNamespace(num_condition_video_rows=2, num_condition_audio_rows=1) + video, audio = torch.randn(5, 6), torch.randn(4, 3) + expected_video, expected_audio = video.clone(), audio.clone() + expected_video[2:] *= PDD_GRID_MAX_T + expected_audio[1:] *= PDD_GRID_MAX_T + + denoising.scale_pdd_initial_noise(video, audio, layout, plan) + + assert float(plan.node_sigmas["video"][0]) == pytest.approx(PDD_GRID_MAX_T) + assert float(plan.node_sigmas["audio"][0]) == pytest.approx(PDD_GRID_MAX_T) + torch.testing.assert_close(video, expected_video, rtol=0, atol=0) + torch.testing.assert_close(audio, expected_audio, rtol=0, atol=0) + + +class _TinyPDDTransformer: + """Stand-in DiT with real widened heads: each row's output depends on its own value and clock.""" + + def __init__(self, pdd_steps: int, hidden: int = 8): + self.pdd_steps = pdd_steps + generator = torch.Generator().manual_seed(7) + self.video_in = torch.randn(VIDEO_CHANNELS, hidden, generator=generator) / VIDEO_CHANNELS**0.5 + self.audio_in = torch.randn(AUDIO_CHANNELS, hidden, generator=generator) / AUDIO_CHANNELS**0.5 + self.proj_out = PDDReplicatedLinear(hidden, VIDEO_CHANNELS, grid_size=pdd_steps, params_dtype=torch.float32) + self.audio_proj_out = PDDReplicatedLinear(hidden, + AUDIO_CHANNELS, + grid_size=pdd_steps, + params_dtype=torch.float32) + with torch.no_grad(): + for linear in (self.proj_out, self.audio_proj_out): + linear.weight.copy_(torch.randn(linear.weight.shape, generator=generator) * 0.1) + linear.bias.copy_(torch.randn(linear.bias.shape, generator=generator) * 0.1) + self.blocks: list[tuple[int, int]] = [] + self.calls: list[dict] = [] + + @contextlib.contextmanager + def fuse_pdd_block(self, start, end, integration_weights, precision_decoding): + self.blocks.append((start, end)) + with fuse_pdd_heads({"video": self.proj_out, "audio": self.audio_proj_out}, start, end, integration_weights, + precision_decoding): + yield + + def features(self, rows, weights, row_times): + return torch.tanh(rows.float() @ weights + row_times[:, None]) + + def __call__(self, **kwargs): + row_times = kwargs["timestep"][kwargs["timestep_indices"]] + self.calls.append({ + **{key: value.clone() for key, value in kwargs.items() if torch.is_tensor(value)}, + "attn_metadata": get_forward_context().attn_metadata, + }) + video = self.features(kwargs["hidden_states"][0], self.video_in, row_times[kwargs["video_indices"]]) + audio = self.features(kwargs["audio_hidden_states"][0], self.audio_in, row_times[kwargs["audio_indices"]]) + return self.proj_out(video)[0][None], self.audio_proj_out(audio)[0][None] + + +def _run(monkeypatch, layout, transformer, *, args, num_inference_steps=8, sparsity=0.0, extra=None, builder=None): + monkeypatch.setattr(denoising, "get_local_torch_device", lambda: torch.device("cpu")) + monkeypatch.setattr(denoising, "_h3_vsa_metadata_builder", lambda *_: builder) + stage = denoising.MiniMaxH3DenoisingStage(transformer, MiniMaxH3Scheduler(shift=12.0), + MiniMaxH3Scheduler(shift=3.0)) + generator = torch.Generator().manual_seed(3) + latents = torch.randn(int(layout.video_indices.numel()), VIDEO_CHANNELS, generator=generator) + audio_latents = torch.randn(int(layout.audio_indices.numel()), AUDIO_CHANNELS, generator=generator) + batch = ForwardBatch(data_type="video", + prompt_embeds=[torch.zeros(1, int(layout.text_indices.numel()), 8)], + latents=latents.clone(), + audio_latents=audio_latents.clone(), + num_inference_steps=num_inference_steps, + VSA_sparsity=sparsity, + extra={ + MINIMAX_H3_LAYOUT_KEY: layout, + **(extra or {}) + }) + return stage, stage.forward(batch, args), latents, audio_latents + + +def _reference_video_and_audio_references(): + return [ + _reference("image", height=4, width=6), + _reference("video", frames=5, height=4, width=6, audio=2), + _reference("audio", audio=3), + ] + + +def test_forward_pdd_runs_eight_fused_blocks_from_contract(monkeypatch): + layout = _ref2va_layout(_reference_video_and_audio_references()) + transformer = _TinyPDDTransformer(pdd_steps=32) + args = _args(pdd_step_indices=GRID32_BLOCKS8) + stage, result, latents, audio_latents = _run(monkeypatch, layout, transformer, args=args) + + assert len(transformer.calls) == 8 and result.step_index == 7 + assert transformer.blocks == list(zip(GRID32_BLOCKS8[:-1], GRID32_BLOCKS8[1:], strict=True)) + for proj in (transformer.proj_out, transformer.audio_proj_out): + assert proj._fusion_state is None + + # Independent recomputation: node sigmas on each modality's clock, targets + # start at 0.999 * noise, conditions stay fixed, and every block advances + # x by -(sigma_end - sigma_start) * (weighted mean of its raw heads). + base = pdd_fine_grid(32)[list(GRID32_BLOCKS8)] + sigmas = {"video": shifted_noise_amount(base, 12.0), "audio": shifted_noise_amount(base, 3.0)} + torch.testing.assert_close(stage.scheduler.sigmas, sigmas["video"].float(), rtol=0, atol=0) + torch.testing.assert_close(stage.audio_scheduler.sigmas, sigmas["audio"].float(), rtol=0, atol=0) + weights = {name: PDDModalitySchedule(shift).integration_weights(pdd_fine_grid(32)) + for name, shift in (("video", 12.0), ("audio", 3.0))} + video_start, audio_start = layout.num_condition_video_rows, layout.num_condition_audio_rows + video, audio = latents.double().clone(), audio_latents.double().clone() + video[video_start:] *= PDD_GRID_MAX_T + audio[audio_start:] *= PDD_GRID_MAX_T + for step, call in enumerate(transformer.calls): + start, end = GRID32_BLOCKS8[step], GRID32_BLOCKS8[step + 1] + torch.testing.assert_close(call["hidden_states"][0].double(), video, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(call["audio_hidden_states"][0].double(), audio, rtol=1e-5, atol=1e-5) + row_times = call["timestep"][call["timestep_indices"]] + video_time = 1.0 - float(sigmas["video"][step].float()) + audio_time = 1.0 - float(sigmas["audio"][step].float()) + video_rows, audio_rows = call["video_indices"], call["audio_indices"] + for rows, value in ((video_rows[video_start:], video_time), (video_rows[:video_start], + max(video_time, MINIMAX_H3_KEYFRAME_NOISE_AUG)), + (audio_rows[audio_start:], audio_time), (audio_rows[:audio_start], 1.0), + (call["text_indices"], video_time)): + torch.testing.assert_close(row_times[rows], torch.full_like(row_times[rows], value), rtol=0, atol=1e-7) + for name, state, rows, first, heads, features_in in ( + ("video", video, video_rows, video_start, transformer.proj_out, transformer.video_in), + ("audio", audio, audio_rows, audio_start, transformer.audio_proj_out, transformer.audio_in), + ): + features = transformer.features(state, features_in, row_times[rows].double().float()) + materialized = heads(features.float())[0].double().unflatten(-1, (32, -1)) + alpha = weights[name][start:end] + mean = torch.einsum("n,rnc->rc", alpha / alpha.sum(), materialized[:, start:end]) + state[first:] -= (sigmas[name][step + 1] - sigmas[name][step]) * mean[first:] + torch.testing.assert_close(result.latents.double(), video, rtol=1e-4, atol=1e-4) + torch.testing.assert_close(result.audio_latents.double(), audio, rtol=1e-4, atol=1e-4) + # Reference conditions never move. + torch.testing.assert_close(result.latents[:video_start], latents[:video_start], rtol=0, atol=0) + torch.testing.assert_close(result.audio_latents[:audio_start], audio_latents[:audio_start], rtol=0, atol=0) + + +@pytest.mark.parametrize("dense_first_n", [0, 1]) +def test_forward_ref_keep_rate_builds_one_sparse_region_per_reference_video(monkeypatch, dense_first_n): + references = [ + _reference("video", frames=5, height=4, width=6), + _reference("image", height=4, width=6), + _reference("video", frames=2, height=8, width=4, audio=2), + ] + layout = _ref2va_layout(references) + transformer = _TinyPDDTransformer(pdd_steps=32) + args = _args(pdd_step_indices=GRID32_BLOCKS8, vsa_ref_keep_rate=0.25) + _run(monkeypatch, + layout, + transformer, + args=args, + sparsity=0.9, + extra={"vsa_dense_first_n_steps": dense_first_n}, + builder=MiniMaxH3VSAMetadataBuilder()) + + assert len(transformer.calls) == 8 + for step, call in enumerate(transformer.calls): + metadata = call["attn_metadata"] + assert metadata.tile_elems == 128 + assert metadata.total_seq_length == layout.sequence_length + assert len(metadata.video_tile_spans) == 3 + expected = (0.0, 0.0, 0.0) if step < dense_first_n else (0.75, 0.75, 0.9) + assert metadata.span_sparsities == pytest.approx(expected) + + +@pytest.mark.parametrize("sparsity,effect", [ + (0.9, "each reference keeps 0.25 of its tiles and the target keeps 0.1"), + (0.0, "VSA_sparsity=0 keeps every region dense"), +]) +def test_forward_ref_keep_rate_logs_applied_sparsity(monkeypatch, sparsity, effect): + infos = [] + monkeypatch.setattr(denoising.logger, "info", lambda message, *args: infos.append(message % args)) + layout = _ref2va_layout(_reference_video_and_audio_references()) + args = _args(pdd_step_indices=GRID32_BLOCKS8, vsa_ref_keep_rate=0.25) + _run(monkeypatch, + layout, + _TinyPDDTransformer(pdd_steps=32), + args=args, + sparsity=sparsity, + builder=MiniMaxH3VSAMetadataBuilder()) + assert f"MiniMax-H3 VSA-H3: 1 reference video region(s); {effect}; 128-token tiles." in infos + + +def test_forward_unset_ref_keep_rate_keeps_references_in_dense_prefix(monkeypatch): + layout = _ref2va_layout(_reference_video_and_audio_references()) + transformer = _TinyPDDTransformer(pdd_steps=32) + args = _args(pdd_step_indices=GRID32_BLOCKS8) + _run(monkeypatch, layout, transformer, args=args, sparsity=0.9, builder=MiniMaxH3VSAMetadataBuilder()) + for call in transformer.calls: + metadata = call["attn_metadata"] + # The generated video is the only sparse region. + assert len(metadata.video_tile_spans) == 1 and metadata.span_sparsities == (0.9, ) + + +def test_forward_ref_keep_rate_requires_exempt_prefix_keys(monkeypatch): + layout = _ref2va_layout(_reference_video_and_audio_references()) + args = _args(pdd_step_indices=GRID32_BLOCKS8, vsa_ref_keep_rate=0.1) + with pytest.raises(ValueError, match="compete' supports only the generated-video region"): + _run(monkeypatch, + layout, + _TinyPDDTransformer(pdd_steps=32), + args=args, + sparsity=0.9, + extra={"vsa_mode": "compete"}, + builder=MiniMaxH3VSAMetadataBuilder()) diff --git a/fastvideo/tests/test_partial_hf_download.py b/fastvideo/tests/test_partial_hf_download.py index 6cb9df925..b83e2a503 100644 --- a/fastvideo/tests/test_partial_hf_download.py +++ b/fastvideo/tests/test_partial_hf_download.py @@ -175,7 +175,7 @@ def test_direct_pipeline_load_forwards_partial_patterns(tmp_path: Path, monkeypa monkeypatch.setattr("fastvideo.pipelines.composed_pipeline_base.maybe_download_model", fake_download) pipeline = object.__new__(MiniMaxH3Ref2VAModularPipeline) pipeline.model_path = "MiniMaxAI/MiniMax-H3" - pipeline.fastvideo_args = SimpleNamespace(revision="test-revision") + pipeline.fastvideo_args = SimpleNamespace(revision="test-revision", pipeline_config=SimpleNamespace(pdd_step_indices=None)) pipeline._load_config(pipeline.model_path) @@ -250,3 +250,40 @@ def test_model_index_supports_umbrella_repo_paths(tmp_path: Path, monkeypatch) - assert captured["filename"] == "minimax-h3/model_index.json" assert captured["revision"] == "test-revision" assert config["pipeline_name"] == "MiniMaxH3ModularPipeline" + + +def test_read_optional_model_json_local_directory(tmp_path: Path) -> None: + (tmp_path / "scheduler").mkdir() + (tmp_path / "scheduler" / "scheduler_config.json").write_text(json.dumps({"shift": 12.0})) + + assert utils.read_optional_model_json(str(tmp_path), "scheduler/scheduler_config.json") == {"shift": 12.0} + assert utils.read_optional_model_json(str(tmp_path), "fastvideo_inference.json") is None + + +def test_read_optional_model_json_umbrella_repo_file(tmp_path: Path, monkeypatch) -> None: + captured: dict = {} + + def fake_hf_hub_download(**kwargs): + captured.update(kwargs) + path = tmp_path / kwargs["filename"] + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps({"pdd_steps": 32})) + return str(path) + + monkeypatch.setattr("huggingface_hub.hf_hub_download", fake_hf_hub_download) + + config = utils.read_optional_model_json("org/repo/student", "transformer_ref/config.json", revision="rev") + + assert config == {"pdd_steps": 32} + assert captured == {"repo_id": "org/repo", "filename": "student/transformer_ref/config.json", "revision": "rev"} + + +def test_read_optional_model_json_absent_hub_file(monkeypatch) -> None: + from huggingface_hub.utils import EntryNotFoundError + + def fake_hf_hub_download(**kwargs): + raise EntryNotFoundError(f"{kwargs['filename']} is not in the repo") + + monkeypatch.setattr("huggingface_hub.hf_hub_download", fake_hf_hub_download) + + assert utils.read_optional_model_json("org/repo", "fastvideo_inference.json") is None diff --git a/fastvideo/utils.py b/fastvideo/utils.py index b42c3f606..10d142b76 100644 --- a/fastvideo/utils.py +++ b/fastvideo/utils.py @@ -747,6 +747,32 @@ def maybe_download_model_index(model_name_or_path: str, revision: str | None = N raise ValueError(f"Failed to download or parse a Diffusers manifest for {model_name_or_path}: {e}") from e +def read_optional_model_json(model_name_or_path: str, + relative_path: str, + revision: str | None = None) -> dict[str, Any] | None: + """Read one JSON file of a local model directory or Hub repo; None when the model has no such file. + + For a Hub repo only that file is downloaded, into the shared HF cache, so + a run can read small checkpoint metadata before any component weights. + """ + if os.path.exists(model_name_or_path): + path = os.path.join(model_name_or_path, relative_path) + if not os.path.isfile(path): + return None + else: + from huggingface_hub import hf_hub_download + from huggingface_hub.utils import EntryNotFoundError + + repo_id, subfolder = _split_hf_repo_subfolder(model_name_or_path) + filename = f"{subfolder}/{relative_path}" if subfolder else relative_path + try: + path = hf_hub_download(repo_id=repo_id, filename=filename, revision=revision) + except EntryNotFoundError: + return None + with open(path, encoding="utf-8") as f: + return cast(dict[str, Any], json.load(f)) + + # Token variables that `huggingface_hub` itself reads, in its priority order. HF_TOKEN_ENV_VARS = ("HF_TOKEN", "HUGGING_FACE_HUB_TOKEN") # FastVideo-specific token names, kept as deprecated aliases until the next