diff --git a/docs/contributing/env_vars.md b/docs/contributing/env_vars.md index 1f12a9f15..62b6c56d6 100644 --- a/docs/contributing/env_vars.md +++ b/docs/contributing/env_vars.md @@ -171,155 +171,159 @@ longer exists also fails the test, so the fixing pull request deletes its entry. ## Registered variables -| Variable | Type | Default | Category | Description | -| ---------------------------------------------------------- | ------------------------------ | ----------------------------------------------------- | ----------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `FASTVIDEO_CONFIG_ROOT` | path | computed | path | Root directory for FastVideo configuration files, at runtime and at installation. Defaults to ~/.config/fastvideo, or $XDG_CONFIG_HOME/fastvideo when XDG_CONFIG_HOME is set. | -| `FASTVIDEO_CACHE_ROOT` | path | computed | path | Root directory for FastVideo cache files. Defaults to ~/.cache/fastvideo, or $XDG_CACHE_HOME/fastvideo when XDG_CACHE_HOME is set. | -| `FASTVIDEO_REASON1_WEIGHTS_PATH` | str | unset | path | Local path or Hugging Face id of Reason1 weights to load instead of the checkpoint's own. | -| `FASTVIDEO_HOST_IP` | str | `""` | distributed | IP address of this node when the node has several network interfaces. Set it on each node for multi-node inference. | -| `FASTVIDEO_LOOPBACK_IP` | str | `""` | distributed | Loopback IP address to use instead of the detected one. | -| `FASTVIDEO_RAY_PER_WORKER_GPUS` | float | `1.0` | distributed | GPUs per Ray worker. A fraction lets Ray schedule several actors on one GPU, so other actors can share the GPUs with FastVideo. | -| `FASTVIDEO_NCCL_SO_PATH` | str | unset | distributed | Path to the NCCL library file. Needed because the nccl>=2.19 that PyTorch ships has a bug (https://github.com/NVIDIA/nccl/issues/1234). | -| `FASTVIDEO_HCCL_SO_PATH` | str | unset | distributed | Path to the HCCL library file on Ascend NPUs. Deprecated names: `HCCL_SO_PATH`. | -| `FASTVIDEO_WORKER_MULTIPROC_METHOD` | one of spawn, fork, forkserver | `spawn` | distributed | Multiprocessing start method for worker processes. | -| `FASTVIDEO_EXTERNAL_LAUNCHER` | bool | `0` | distributed | With the default mp backend, offline inference runs SPMD under torchrun or srun: each launched process is one worker joining the env:// rendezvous. All ranks must call generate() together; world rank 0 owns outputs. | -| `FASTVIDEO_ULYSSES_A2A` | one of off, auto | `off` | distributed | Sequence-parallel all-to-all backend. off uses the NCCL path in DistributedAutograd.AllToAll4D. auto uses the fused NVLink kernel when the group is a load-store accessible mesh of 2, 4, 6, or 8 ranks in eager execution, and the NCCL path otherwise. | -| `FASTVIDEO_ULYSSES_A2A_LONG_TRAINING` | one of auto, chunked | `auto` | distributed | Fused Ulysses all-to-all policy for grad-enabled calls whose window exceeds the long-training plane limit. auto keeps the original unchunked transport. chunked uses bounded chunks; validate the recipe's activation memory first. | -| `FASTVIDEO_CONFIGURE_LOGGING` | bool | `1` | logging | Configure logging at import. When true, FastVideo uses its default logging configuration or the file in FASTVIDEO_LOGGING_CONFIG_PATH. | -| `FASTVIDEO_LOGGING_CONFIG_PATH` | str | unset | logging | Path to a JSON logging configuration file. | -| `FASTVIDEO_LOGGING_LEVEL` | str | `INFO` | logging | Default logging level. | -| `FASTVIDEO_LOGGING_PREFIX` | str | `""` | logging | Prefix prepended to every log message. | -| `FASTVIDEO_STAGE_LOGGING` | bool | `0` | logging | Log the time that each pipeline stage takes. | -| `FASTVIDEO_LOG_ALL_PROCESSES` | bool | `0` | logging | logger.info logs from every process, ignoring the default local-main-process filter and the main_process_only and local_main_process_only arguments. Read at each call, so it can be set after importing fastvideo. Useful for debugging distributed runs. | -| `FASTVIDEO_ATTENTION_BACKEND` | str | unset | attention | Attention backend, as an AttentionBackendEnum name such as TORCH_SDPA, FLASH_ATTN, VIDEO_SPARSE_ATTN, SAGE_ATTN, or SAGE_ATTN_THREE. FastVideoArgs uses it when FastVideoArgs.attention_backend is unset. | -| `FASTVIDEO_FA4` | bool | `0` | attention | The FLASH_ATTN backend uses FlashAttention-4 (flash_attn.cute) instead of FA3 or FA2. | -| `FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN` | bool | `0` | attention | MiniMax-H3 dense DiT self-attention uses the FlashAttention-4 packed-varlen entry point. This changes the floating-point reduction order, so it is an inference-only opt-in. | -| `FASTVIDEO_VSA_SM100A` | bool | `0` | attention | VIDEO_SPARSE_ATTN_H3 sends no-grad tile-64 forwards to the data-center Blackwell (sm_100a) kernel. fastvideo-kernel reads the same variable with the same rule. | -| `FASTVIDEO_VSA_TRITON` | bool | `0` | attention | Force the Triton MiniMax-H3 sparse attention kernel. fastvideo-kernel reads the same variable. | -| `FASTVIDEO_NVFP4_FA4` | bool | `0` | attention | FlashAttention-4 quantizes Q and K to NVFP4. An explicit nvfp4_fa4 attention implementation argument takes precedence. | -| `FASTVIDEO_FA4_PV_MODE` | one of bf16, fp8 | `bf16` | attention | V dtype on the NVFP4 FlashAttention-4 path (FLASH_ATTN with nvfp4_fa4, ATTN_QAT_INFER on sm_100a/sm_103a): bf16 keeps V in BF16; fp8 casts V to float8 e4m3 without scaling. An explicit fa4_pv_mode attention implementation argument takes precedence. | -| `FASTVIDEO_DISABLE_ATTENTION_COMPILE` | bool | `1` | attention | Keep attention forward out of torch.compile graphs (torch.compiler.disable). Set it to 0 to let attention constructed under that setting be traced. Setting it explicitly to true also blocks regional compile. | -| `FASTVIDEO_MLX_WINDOW` | int | `0` | attention | MLX FastWan windowed attention size in tokens. 0 uses full attention. | -| `FASTVIDEO_MLX_WINDOW_SINK` | int | `0` | attention | Number of sink tokens that MLX windowed attention always attends to. | -| `FASTVIDEO_DISABLE_VSA64_FUSED_LAYOUT` | bool | `0` | attention | VIDEO_SPARSE_ATTN keeps the original tile scatter and BSHD->BHSD transposes instead of the fused Triton layout kernel that no-grad SM100 BF16 head_dim-128 forwards with 64-token tiles use by default. | -| `FASTVIDEO_INFERENCE_TORCH_COMPILE` | bool | `0` | performance | Compile each DiT transformer block with fullgraph torch.compile at inference. Same as FastVideoArgs.inference_torch_compile=True. | -| `FASTVIDEO_VAE_PARALLEL_DECODE` | bool | `0` | performance | MiniMax-H3 VAE decode splits its temporal chunks across the sequence-parallel ranks instead of running serially on the output rank. Same as FastVideoArgs.vae_parallel_decode=True. | -| `FASTVIDEO_VAE_PARALLEL_ENCODE` | bool | `0` | performance | MiniMax-H3 reference-video VAE encode splits its temporal chunks across the sequence-parallel ranks. Same as FastVideoArgs.vae_parallel_encode=True. | -| `FASTVIDEO_VAE_PARALLEL_DECODE_STRATEGY` | str | unset | performance | Collective that moves chunks in parallel VAE decode: gather (used when unset) or all_gather. | -| `FASTVIDEO_H3_ENCODER_SPLIT` | bool | `0` | distributed | MiniMax-H3 runs the Qwen3-VL text encoder on a dedicated group of nodes, so the denoising ranks never load it. Same as FastVideoArgs.h3_encoder_split=True. | -| `FASTVIDEO_H3_ENCODER_NODES` | int | `1` | distributed | Number of leading nodes that FASTVIDEO_H3_ENCODER_SPLIT dedicates to the MiniMax-H3 text encoder. Values below 1 count as 1. | -| `FASTVIDEO_MINIMAX_H3_FUSIONS` | str | `""` | performance | MiniMax-H3 inference-only Triton fusions: all, 1, or a comma-separated subset of modulate,qknorm_rope,swiglu. Empty, 0, or none keeps the eager implementation. | -| `FASTVIDEO_FSDP2_AUTOWRAP` | bool | `0` | performance | FSDP2 shards modules by parameter count instead of the model's shard conditions. Not supported by self-forcing distillation. | -| `FASTVIDEO_FSDP2_MIN_PARAMS` | int | `10000000` | performance | Minimum parameter count of a module that FASTVIDEO_FSDP2_AUTOWRAP shards. | -| `FASTVIDEO_MLX_COMPILE` | bool | `0` | performance | Compile the MLX DiT forward with mx.compile. | -| `FASTVIDEO_MLX_FAST_NORM` | bool | `0` | performance | Use MLX fast normalization kernels. | -| `FASTVIDEO_MLX_DQ_GEMM` | str | `1` | performance | MLX dequantized GEMM for affine-quantized weights: 0 turns it off, 1 uses the measured minimum row count, and an integer sets the minimum row count. | -| `FASTVIDEO_LTX2_VAE_CHANNELS_LAST_3D` | bool | `1` | performance | LTX-2 VAE uses the channels_last_3d memory format. | -| `FASTVIDEO_LTX2_DISABLE_AUDIO_AUTOCAST` | bool | `1` | performance | LTX-2 audio decoding runs without CUDA autocast. Deprecated names: `LTX2_DISABLE_AUDIO_AUTOCAST`. | -| `FASTVIDEO_FLUX2_DISABLE_BF16_REDUCED_PRECISION_REDUCTION` | bool | `0` | performance | Flux denoising disables reduced-precision reductions in bf16 matmuls, which tightens accumulation for the 4-step Klein model. | -| `FASTVIDEO_DISABLE_FUSED_NORM` | bool | `0` | performance | Turn off the Triton-fused residual + LayerNorm + modulate inference path that Wan blocks opt into, and use the eager path. | -| `FASTVIDEO_FFMPEG_BIN` | str | `ffmpeg` | output | ffmpeg executable used to save video with audio. | -| `FASTVIDEO_VIDEO_CODEC` | str | `libx264` | output | ffmpeg video codec for saved videos. | -| `FASTVIDEO_NVENC_PRESET` | str | `p1` | output | NVENC preset when the codec is an \*_nvenc codec. | -| `FASTVIDEO_NVENC_TUNE` | str | `ull` | output | NVENC tune option. | -| `FASTVIDEO_NVENC_RC` | str | `constqp` | output | NVENC rate-control mode. | -| `FASTVIDEO_NVENC_QP` | str | `28` | output | NVENC quantization parameter. | -| `FASTVIDEO_NVENC_BF` | str | `0` | output | NVENC number of B-frames. | -| `FASTVIDEO_X264_PRESET` | str | `ultrafast` | output | x264 preset for non-NVENC codecs. | -| `FASTVIDEO_OUTPUT_PIX_FMT` | str | `yuv420p` | output | ffmpeg pixel format for saved videos. | -| `FASTVIDEO_NVTX_PROFILE` | bool | `0` | profiling | Emit NVTX ranges for external profilers such as Nsight Systems. | -| `FASTVIDEO_TORCH_PROFILER_DIR` | path | unset | profiling | Enables the torch profiler and sets the directory for its traces. Must be an absolute path. | -| `FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES` | bool | `0` | profiling | Torch profiler records shapes. | -| `FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY` | bool | `0` | profiling | Torch profiler profiles memory. | -| `FASTVIDEO_TORCH_PROFILER_WITH_STACK` | bool | `0` | profiling | Torch profiler captures stacks. Costs about 1.5x runtime and 1.4x trace size. | -| `FASTVIDEO_TORCH_PROFILER_WITH_FLOPS` | bool | `0` | profiling | Torch profiler profiles FLOPs. | -| `FASTVIDEO_TORCH_PROFILE_REGIONS` | str | `""` | profiling | Comma-separated profiler regions to record. The torch profiler requires at least one region. | -| `FASTVIDEO_TRACE_ACTIVATIONS` | bool | `0` | debug | Enable activation trace hooks. | -| `FASTVIDEO_TRACE_LAYERS` | str | `""` | debug | Regex filter for traced module names. Empty means all. | -| `FASTVIDEO_TRACE_STATS` | str | `abs_mean,sum` | debug | Comma-separated activation statistics dumped for each output tensor. | -| `FASTVIDEO_TRACE_OUTPUT` | str | `/tmp/fv_trace_.jsonl` | debug | JSONL path for activation traces. The literal <pid> is replaced at runtime. | -| `FASTVIDEO_TRACE_STEPS` | str | `""` | debug | Comma-separated denoising step indices. Empty means all. | -| `FASTVIDEO_H3_VSA_PROBE` | str | unset | debug | Output directory for the VSA-H3 attention-mass probe, which writes one .pt file per step, layer, and rank. Keeps the model out of regional compile. | -| `FASTVIDEO_LTX2_GEMMA_LOG` | str | `""` | debug | Log file for LTX-2 Gemma text-encoder hidden states, used by parity tests. Deprecated names: `LTX2_FASTVIDEO_GEMMA_LOG`. | -| `FASTVIDEO_COSMOS25_LOG_KNOBS` | bool | `0` | debug | Log the Cosmos 2.5 latent-preparation conditioning inputs. | -| `FASTVIDEO_CFG_GATE_STEP` | float | `1.0` | sampling | CFG gating fraction in [0, 1]. Steps before len(timesteps) \* X run the conditional and unconditional forwards; later steps reuse the cached difference. 1.0 disables gating. | -| `FASTVIDEO_LTX2_USE_DISTILLED_SIGMAS` | bool | `1` | sampling | LTX-2 uses the distilled sigma schedule when FastVideoArgs.ltx2_use_distilled_sigmas is also true. Deprecated names: `LTX2_USE_DISTILLED_SIGMAS`. | -| `FASTVIDEO_EVAL_CACHE` | path | computed | eval | Cache directory for evaluation models and datasets. Defaults to $FASTVIDEO_CACHE_ROOT/eval. | -| `FASTVIDEO_PHYSICS_IQ_BUCKET_URL` | str | `https://storage.googleapis.com/physics-iq-benchmark` | eval | Base URL of the Physics-IQ benchmark bucket. | -| `FASTVIDEO_VBENCH_FULL_INFO_JSON` | str | unset | eval | Path to VBench_full_info.json, used instead of the vendored copy. Deprecated names: `VBENCH_FULL_INFO_JSON`. | -| `FASTVIDEO_FVD_REF_FEATURES` | str | unset | eval | Cached reference-feature file for the FVD metric. | -| `FASTVIDEO_FAD_REF_FEATURES` | str | unset | eval | Cached reference-feature file for the audio Frechet distance metric. | -| `FASTVIDEO_H3_VSA_FP4` | bool | `0` | attention | Run MiniMax-H3 VSA attention on the block-sparse SageAttention3 FP4 kernel (sm_120, no-grad, single sequence-parallel rank). | -| `FASTVIDEO_H3_VSA_TILE_FIRST` | bool | `0` | attention | Single-rank MiniMax-H3 VSA with one tile gather of the block input instead of separate Q/K/V/gate scatters. | -| `FASTVIDEO_H3_VSA_SM89_KERNEL` | one of original, bf16, int8 | `original` | attention | Fine-attention kernel for MiniMax-H3 VSA on sm_89: original, bf16, or int8 (INT8 QK, BF16 PV). | -| `FASTVIDEO_H3_SIM_SP_FP8` | bool | `0` | debug | Simulate the FP8 sequence-parallel exchange of the MiniMax-H3 FP4 VSA path on one rank. | -| `FASTVIDEO_H3_FFN_CHUNK_TOKENS` | int | `0` | performance | Inference-only MiniMax-H3 FFN token chunk size; 0 runs the FFN unchunked. | -| `FASTVIDEO_H3_FP8_ATTENTION` | bool | `0` | performance | With NVFP4 layer_profile h3_dit_ffn, run MiniMax-H3 attention projections in FP8. | -| `FASTVIDEO_H3_FP8_GRANULARITY` | one of tensor, channel | `tensor` | performance | FP8 scaling granularity for FASTVIDEO_H3_FP8_ATTENTION. | -| `FASTVIDEO_NVFP4_MM_BACKEND` | str | `auto` | performance | FlashInfer mm_fp4 backend for NVFP4 linears, e.g. auto or cutlass. | -| `FASTVIDEO_NVFP4_ACT_AMAX` | path | unset | performance | JSON of calibrated NVFP4 input amax per linear, keyed b<block>.<sub> or full prefix; sets a static activation scale. | -| `FASTVIDEO_NVFP4_DYNAMIC_ACT` | str | `""` | performance | NVFP4 linears that derive the activation scale per call: all, or comma-separated layer-name suffixes such as ff.fc_out. | -| `FASTVIDEO_H3_ADALN_CACHE` | bool | `0` | performance | Cache MiniMax-H3 AdaLN modulation per timestep instead of keeping the projection weights resident. | -| `FASTVIDEO_H3_ADALN_TABLE` | path | unset | performance | Precomputed MiniMax-H3 AdaLN modulation table; enables the cache and skips loading the AdaLN projection weights. | -| `FASTVIDEO_H3_ADALN_DUMP` | path | unset | debug | Write the MiniMax-H3 AdaLN modulation table to this path while sampling. | -| `FASTVIDEO_H3_SPLICE_TRANSFORMER` | path | unset | eval | Second MiniMax-H3 transformer that runs the late DMD steps (checkpoint step-splice evaluation). | -| `FASTVIDEO_H3_SPLICE_FROM_STEP` | int | `4` | eval | First denoising step run by FASTVIDEO_H3_SPLICE_TRANSFORMER. | -| `FASTVIDEO_H3_ENCODER_LAYERWISE` | bool | `0` | performance | Stream MiniMax-H3 text-encoder language layers through exact-size pinned host memory (text-only prompts). | -| `FASTVIDEO_H3_ENCODER_FUSED_DEQUANT` | bool | `0` | performance | Expand the serialized NVFP4 MiniMax-H3 text encoder with one fused Triton pass on GPUs without FP4 GEMM. | -| `FASTVIDEO_H3_VAE_TILE_BATCH` | int | `1` | performance | Spatial tiles per MiniMax-H3 video VAE decoder call; 1 decodes per tile. | -| `FASTVIDEO_H3_VAE_INT8_SHARED_QKV` | bool | `0` | performance | Share the INT8 activation rotation and quantization across the MiniMax-H3 VAE Q/K/V projections. | -| `FASTVIDEO_H3_VAE_INT8_TRANSPOSE_VIEW` | bool | `0` | performance | Use transposed weight views in the MiniMax-H3 VAE INT8 projections. | -| `FASTVIDEO_H3_VAE_INT8_FUSED_DEQUANT` | bool | `0` | performance | Fused dequantization epilogue for the MiniMax-H3 VAE INT8 projections. | -| `FASTVIDEO_H3_PINNED_SWAP` | bool | `1` | performance | Swap offloaded MiniMax-H3 modules through exact-size pinned host arenas. | -| `FASTVIDEO_H3_PARK_MODULES` | str | unset | performance | Comma-separated MiniMax-H3 denoise modules parked on the host while the text encoder runs, e.g. vae,audio_vae. | -| `FASTVIDEO_LAYERWISE_OFFLOAD_BUFFERS` | bool | `0` | performance | Layerwise offload also streams large buffers such as packed FP4/FP8 weights. | -| `FASTVIDEO_LAYERWISE_RESIDENT_BLOCKS` | int | `0` | performance | Keep the first N layerwise-offloaded blocks resident on the GPU. | -| `FASTVIDEO_H3_SP_PROFILE` | bool | `0` | profiling | CUDA-event spans per stage over one MiniMax-H3 FP4 VSA DiT forward. | -| `FASTVIDEO_H3_CAPTURE_QKV` | path | unset | debug | Directory for captured real MiniMax-H3 Q/K/V attention inputs. | -| `FASTVIDEO_CUDA_MEMORY_CAP_GIB` | float | `0.0` | debug | Cap this process's CUDA allocator at this many GiB to emulate a smaller GPU; 0 leaves it uncapped. | -| `FASTVIDEO_MEMORY_REPORT` | bool | `0` | debug | Log bytes held per pipeline component by device and dtype after loading. | -| `FASTVIDEO_TEST_WAN_S2V_MODEL_PATH` | path | unset | test | Wan2.2-S2V-14B weights for test_wan_s2v.py; unset means official_weights/Wan2.2-S2V-14B in the repository. Deprecated names: `WAN_S2V_MODEL_PATH`. | -| `FASTVIDEO_TEST_LTX2_OVERFIT_DATA_DIR` | str | `data/cats` | test | Raw data directory for preprocess_ltx2_overfit.py. Deprecated names: `LTX2_OVERFIT_DATA_DIR`. | -| `FASTVIDEO_TEST_LTX2_OVERFIT_CAPTION_JSON` | str | `videos2caption_1_sample.json` | test | Caption file, relative to the raw data directory. Deprecated names: `LTX2_OVERFIT_CAPTION_JSON`. | -| `FASTVIDEO_TEST_LTX2_OVERFIT_VIDEO_SUBDIR` | str | `video` | test | Video subdirectory, relative to the raw data directory. Deprecated names: `LTX2_OVERFIT_VIDEO_SUBDIR`. | -| `FASTVIDEO_TEST_LTX2_OVERFIT_OUTPUT_DIR` | str | `data/ltx2_overfit_preprocessed` | test | Output directory for preprocess_ltx2_overfit.py. Deprecated names: `LTX2_OVERFIT_OUTPUT_DIR`. | -| `FASTVIDEO_TEST_LTX2_OVERFIT_MODEL` | str | `FastVideo/LTX2-Distilled-Diffusers` | test | Model repository whose encoders preprocess_ltx2_overfit.py uses. Deprecated names: `LTX2_OVERFIT_MODEL`. | -| `FASTVIDEO_TEST_LTX2_OVERFIT_NUM_COPIES` | int | `4` | test | Number of copies of the overfit sample in the parquet file. Deprecated names: `LTX2_OVERFIT_NUM_COPIES`. | -| `FASTVIDEO_TEST_KANDINSKY5_OVERFIT_DATA_DIR` | str | `data/kandinsky5_overfit` | test | Raw data directory for preprocess_kandinsky5_overfit.py. Deprecated names: `KANDINSKY5_OVERFIT_DATA_DIR`. | -| `FASTVIDEO_TEST_KANDINSKY5_OVERFIT_OUTPUT_DIR` | str | `data/kandinsky5_overfit_preprocessed` | test | Output directory for preprocess_kandinsky5_overfit.py. Deprecated names: `KANDINSKY5_OVERFIT_OUTPUT_DIR`. | -| `FASTVIDEO_TEST_HUNYUAN15_OVERFIT_DATA_DIR` | str | `data/hunyuan15_overfit` | test | Raw data directory for preprocess_hunyuan15_overfit.py. | -| `FASTVIDEO_TEST_HUNYUAN15_OVERFIT_OUTPUT_DIR` | str | `data/hunyuan15_overfit_preprocessed` | test | Output directory for preprocess_hunyuan15_overfit.py. | -| `FASTVIDEO_TEST_SSIM_REFERENCE_HF_REPO` | str | `FastVideo/ssim-reference-videos` | test | Hugging Face repository that holds the SSIM reference videos. Deprecated names: `FASTVIDEO_SSIM_REFERENCE_HF_REPO`. | -| `FASTVIDEO_TEST_SSIM_REFERENCE_HF_REPO_TYPE` | str | `dataset` | test | Repository type of FASTVIDEO_TEST_SSIM_REFERENCE_HF_REPO. Deprecated names: `FASTVIDEO_SSIM_REFERENCE_HF_REPO_TYPE`. | -| `FASTVIDEO_TEST_SSIM_SKIP_REFERENCE_DOWNLOAD` | bool | `0` | test | SSIM tests use local reference videos without downloading. Deprecated names: `FASTVIDEO_SSIM_SKIP_REFERENCE_DOWNLOAD`. | -| `FASTVIDEO_TEST_SSIM_FULL_QUALITY` | bool | `0` | test | SSIM tests use the full-quality sampling configurations. Deprecated names: `FASTVIDEO_SSIM_FULL_QUALITY`. | -| `FASTVIDEO_TEST_NIGHTLY` | bool | `0` | test | Run the nightly end-to-end overfit tests. Deprecated names: `FASTVIDEO_NIGHTLY`. | -| `FASTVIDEO_TEST_ULYSSES_FAULT_RANK` | str | unset | test | Rank that fails in the Ulysses fault-injection test. The test sets it for its worker processes. Deprecated names: `FASTVIDEO_ULYSSES_FAULT_RANK`. | -| `FASTVIDEO_TEST_ULYSSES_FAULT_STAGE` | str | unset | test | Stage that fails in the Ulysses fault-injection test. The test sets it for its worker processes. Deprecated names: `FASTVIDEO_ULYSSES_FAULT_STAGE`. | -| `FASTVIDEO_TEST_GOLDEN_GATE_DIR` | str | unset | test | Local directory of golden-gate reference tensors. Deprecated names: `FASTVIDEO_GOLDEN_GATE_DIR`. | -| `FASTVIDEO_TEST_WAN22_5B_ALLOW_LOW_MEMORY` | bool | `0` | test | Run the MLX Wan2.2 5B real-weights parity test on hosts with little memory. Deprecated names: `FASTVIDEO_WAN22_5B_ALLOW_LOW_MEMORY`. | -| `FASTVIDEO_TEST_WAN22_5B_ROOT` | str | unset | test | Local Wan2.2 5B checkpoint for the MLX real-weights parity test. Deprecated names: `FASTVIDEO_WAN22_5B_ROOT`. | -| `FASTVIDEO_TEST_GRADNORM_UPDATE` | bool | `0` | test | Gradient-norm regression tests update their references. Deprecated names: `FASTVIDEO_GRADNORM_UPDATE`. | -| `FASTVIDEO_TEST_DREAMX_WORLD_SSIM_MODEL_PATH` | str | `FastVideo/DreamX-World-5B-Cam-Diffusers` | test | Model for the DreamX-World camera SSIM test. Deprecated names: `DREAMX_WORLD_SSIM_MODEL_PATH`. | -| `FASTVIDEO_TEST_DREAMX_WORLD_AR_SSIM_MODEL_PATH` | str | `FastVideo/DreamX-World-5B-Diffusers` | test | Model for the DreamX-World autoregressive SSIM test. Deprecated names: `DREAMX_WORLD_AR_SSIM_MODEL_PATH`. | -| `FASTVIDEO_TEST_FLUX_T2I_MODEL_DIR` | str | `black-forest-labs/FLUX.1-dev` | test | Model for the Flux text-to-image SSIM test. Deprecated names: `FLUX_T2I_MODEL_DIR`. | -| `FASTVIDEO_TEST_FLUX_TRANSFORMER_PATH` | str | unset | test | Local Flux transformer for the Flux transformer test. Deprecated names: `FLUX_TRANSFORMER_PATH`. | -| `FASTVIDEO_TEST_GAMECRAFT_MODEL_PATH` | str | `FastVideo/HunyuanGameCraft-Diffusers` | test | Model for the HunyuanGameCraft SSIM test. Deprecated names: `GAMECRAFT_MODEL_PATH`. | -| `FASTVIDEO_TEST_GEN3C_MODEL_PATH` | str | `FastVideo/GEN3C-Cosmos-7B-Diffusers` | test | Model for the GEN3C SSIM test. Deprecated names: `GEN3C_MODEL_PATH`. | -| `FASTVIDEO_TEST_GEN3C_IMAGE_PATH` | str | unset | test | Input image for the GEN3C SSIM test. Deprecated names: `GEN3C_TEST_IMAGE_PATH`. | -| `FASTVIDEO_TEST_GLM_IMAGE_LOCAL_WEIGHTS_DIR` | str | unset | test | Local official GLM-Image weights for the GLM-Image SSIM test. Deprecated names: `GLM_IMAGE_LOCAL_WEIGHTS_DIR`. | -| `FASTVIDEO_TEST_GLM_IMAGE_MODEL_DIR` | str | unset | test | Model for the GLM-Image SSIM test. Deprecated names: `GLM_IMAGE_MODEL_DIR`. | -| `FASTVIDEO_TEST_KANDINSKY5_E2E_NUM_GPUS` | int | `1` | test | GPUs for the Kandinsky5 nightly end-to-end overfit test. Deprecated names: `KANDINSKY5_E2E_NUM_GPUS`. | -| `FASTVIDEO_TEST_KANDINSKY5_E2E_WRITE_REFERENCE` | bool | `0` | test | The Kandinsky5 nightly end-to-end test writes a missing reference video. Deprecated names: `KANDINSKY5_E2E_WRITE_REFERENCE`. | -| `FASTVIDEO_TEST_LONGCAT_MODEL_ROOT` | str | unset | test | Local LongCat-Video checkpoint for the golden-gate test. Deprecated names: `LONGCAT_MODEL_ROOT`. | -| `FASTVIDEO_TEST_MINIMAX_H3_GATE_GOLDEN_DIR` | str | unset | test | Local directory of MiniMax-H3 golden-gate tensors. Deprecated names: `MINIMAX_H3_GATE_GOLDEN_DIR`. | -| `FASTVIDEO_TEST_MINIMAX_H3_GATE_LAYER` | int | `0` | test | Transformer layer that the MiniMax-H3 golden-gate test checks. Deprecated names: `MINIMAX_H3_GATE_LAYER`. | -| `FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT` | str | unset | test | Local MiniMax-H3 checkpoint for the golden-gate test. Deprecated names: `MINIMAX_H3_MODEL_ROOT`. | -| `FASTVIDEO_TEST_SD35_MODEL_DIR` | str | `stabilityai/stable-diffusion-3.5-medium` | test | Model for the Stable Diffusion 3.5 SSIM test. Deprecated names: `SD35_MODEL_DIR`. | -| `FASTVIDEO_TEST_TAEH3_REFERENCE_DIR` | str | unset | test | Upstream taehv checkout for the MLX TAEH3 parity test. Deprecated names: `TAEH3_REFERENCE_DIR`. | -| `FASTVIDEO_TEST_WAN_ANIMATE_MODEL_DIR` | str | unset | test | Local Wan2.2-Animate-14B checkpoint for the Wan-Animate weight tests. Deprecated names: `WAN_ANIMATE_MODEL_PATH`. | -| `FASTVIDEO_TEST_ZIMAGE_MODEL_DIR` | str | `Tongyi-MAI/Z-Image-Turbo` | test | Model for the Z-Image SSIM test. Deprecated names: `ZIMAGE_MODEL_DIR`. | -| `FASTVIDEO_TEST_ZIMAGE_MODEL_REVISION` | str | `f332072aa78be7aecdf3ee76d5c247082da564a6` | test | Hugging Face revision of the Z-Image model for its SSIM test. Deprecated names: `ZIMAGE_MODEL_REVISION`. | +| Variable | Type | Default | Category | Description | +| ---------------------------------------------------------- | ------------------------------ | ----------------------------------------------------- | ----------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `FASTVIDEO_CONFIG_ROOT` | path | computed | path | Root directory for FastVideo configuration files, at runtime and at installation. Defaults to ~/.config/fastvideo, or $XDG_CONFIG_HOME/fastvideo when XDG_CONFIG_HOME is set. | +| `FASTVIDEO_CACHE_ROOT` | path | computed | path | Root directory for FastVideo cache files. Defaults to ~/.cache/fastvideo, or $XDG_CACHE_HOME/fastvideo when XDG_CACHE_HOME is set. | +| `FASTVIDEO_REASON1_WEIGHTS_PATH` | str | unset | path | Local path or Hugging Face id of Reason1 weights to load instead of the checkpoint's own. | +| `FASTVIDEO_HOST_IP` | str | `""` | distributed | IP address of this node when the node has several network interfaces. Set it on each node for multi-node inference. | +| `FASTVIDEO_LOOPBACK_IP` | str | `""` | distributed | Loopback IP address to use instead of the detected one. | +| `FASTVIDEO_RAY_PER_WORKER_GPUS` | float | `1.0` | distributed | GPUs per Ray worker. A fraction lets Ray schedule several actors on one GPU, so other actors can share the GPUs with FastVideo. | +| `FASTVIDEO_NCCL_SO_PATH` | str | unset | distributed | Path to the NCCL library file. Needed because the nccl>=2.19 that PyTorch ships has a bug (https://github.com/NVIDIA/nccl/issues/1234). | +| `FASTVIDEO_HCCL_SO_PATH` | str | unset | distributed | Path to the HCCL library file on Ascend NPUs. Deprecated names: `HCCL_SO_PATH`. | +| `FASTVIDEO_WORKER_MULTIPROC_METHOD` | one of spawn, fork, forkserver | `spawn` | distributed | Multiprocessing start method for worker processes. | +| `FASTVIDEO_EXTERNAL_LAUNCHER` | bool | `0` | distributed | With the default mp backend, offline inference runs SPMD under torchrun or srun: each launched process is one worker joining the env:// rendezvous. All ranks must call generate() together; world rank 0 owns outputs. | +| `FASTVIDEO_ULYSSES_A2A` | one of off, auto | `off` | distributed | Sequence-parallel all-to-all backend. off uses the NCCL path in DistributedAutograd.AllToAll4D. auto uses the fused NVLink kernel when the group is a load-store accessible mesh of 2, 4, 6, or 8 ranks in eager execution, and the NCCL path otherwise. | +| `FASTVIDEO_ULYSSES_A2A_LONG_TRAINING` | one of auto, chunked | `auto` | distributed | Fused Ulysses all-to-all policy for grad-enabled calls whose window exceeds the long-training plane limit. auto keeps the original unchunked transport. chunked uses bounded chunks; validate the recipe's activation memory first. | +| `FASTVIDEO_CONFIGURE_LOGGING` | bool | `1` | logging | Configure logging at import. When true, FastVideo uses its default logging configuration or the file in FASTVIDEO_LOGGING_CONFIG_PATH. | +| `FASTVIDEO_LOGGING_CONFIG_PATH` | str | unset | logging | Path to a JSON logging configuration file. | +| `FASTVIDEO_LOGGING_LEVEL` | str | `INFO` | logging | Default logging level. | +| `FASTVIDEO_LOGGING_PREFIX` | str | `""` | logging | Prefix prepended to every log message. | +| `FASTVIDEO_STAGE_LOGGING` | bool | `0` | logging | Log the time that each pipeline stage takes. | +| `FASTVIDEO_LOG_ALL_PROCESSES` | bool | `0` | logging | logger.info logs from every process, ignoring the default local-main-process filter and the main_process_only and local_main_process_only arguments. Read at each call, so it can be set after importing fastvideo. Useful for debugging distributed runs. | +| `FASTVIDEO_ATTENTION_BACKEND` | str | unset | attention | Attention backend, as an AttentionBackendEnum name such as TORCH_SDPA, FLASH_ATTN, VIDEO_SPARSE_ATTN, SAGE_ATTN, or SAGE_ATTN_THREE. FastVideoArgs uses it when FastVideoArgs.attention_backend is unset. | +| `FASTVIDEO_FA4` | bool | `0` | attention | The FLASH_ATTN backend uses FlashAttention-4 (flash_attn.cute) instead of FA3 or FA2. | +| `FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN` | bool | `0` | attention | MiniMax-H3 dense DiT self-attention uses the FlashAttention-4 packed-varlen entry point. This changes the floating-point reduction order, so it is an inference-only opt-in. | +| `FASTVIDEO_VSA_SM100A` | bool | `0` | attention | VIDEO_SPARSE_ATTN_H3 sends no-grad tile-64 forwards to the data-center Blackwell (sm_100a) kernel. fastvideo-kernel reads the same variable with the same rule. | +| `FASTVIDEO_VSA_TRITON` | bool | `0` | attention | Force the Triton MiniMax-H3 sparse attention kernel. fastvideo-kernel reads the same variable. | +| `FASTVIDEO_NVFP4_FA4` | bool | `0` | attention | FlashAttention-4 quantizes Q and K to NVFP4. An explicit nvfp4_fa4 attention implementation argument takes precedence. | +| `FASTVIDEO_FA4_PV_MODE` | one of bf16, fp8 | `bf16` | attention | V dtype on the NVFP4 FlashAttention-4 path (FLASH_ATTN with nvfp4_fa4, ATTN_QAT_INFER on sm_100a/sm_103a): bf16 keeps V in BF16; fp8 casts V to float8 e4m3 without scaling. An explicit fa4_pv_mode attention implementation argument takes precedence. | +| `FASTVIDEO_DISABLE_ATTENTION_COMPILE` | bool | `1` | attention | Keep attention forward out of torch.compile graphs (torch.compiler.disable). Set it to 0 to let attention constructed under that setting be traced. Setting it explicitly to true also blocks regional compile. | +| `FASTVIDEO_MLX_WINDOW` | int | `0` | attention | MLX FastWan windowed attention size in tokens. 0 uses full attention. | +| `FASTVIDEO_MLX_WINDOW_SINK` | int | `0` | attention | Number of sink tokens that MLX windowed attention always attends to. | +| `FASTVIDEO_DISABLE_VSA64_FUSED_LAYOUT` | bool | `0` | attention | VIDEO_SPARSE_ATTN keeps the original tile scatter and BSHD->BHSD transposes instead of the fused Triton layout kernel that no-grad SM100 BF16 head_dim-128 forwards with 64-token tiles use by default. | +| `FASTVIDEO_INFERENCE_TORCH_COMPILE` | bool | `0` | performance | Compile each DiT transformer block with fullgraph torch.compile at inference. Same as FastVideoArgs.inference_torch_compile=True. | +| `FASTVIDEO_VAE_PARALLEL_DECODE` | bool | `0` | performance | MiniMax-H3 VAE decode splits its temporal chunks across the sequence-parallel ranks instead of running serially on the output rank. Same as FastVideoArgs.vae_parallel_decode=True. | +| `FASTVIDEO_VAE_PARALLEL_ENCODE` | bool | `0` | performance | MiniMax-H3 reference-video VAE encode splits its temporal chunks across the sequence-parallel ranks. Same as FastVideoArgs.vae_parallel_encode=True. | +| `FASTVIDEO_VAE_PARALLEL_DECODE_STRATEGY` | str | unset | performance | Collective that moves chunks in parallel VAE decode: gather (used when unset) or all_gather. | +| `FASTVIDEO_H3_ENCODER_SPLIT` | bool | `0` | distributed | MiniMax-H3 runs the Qwen3-VL text encoder on a dedicated group of nodes, so the denoising ranks never load it. Same as FastVideoArgs.h3_encoder_split=True. | +| `FASTVIDEO_H3_ENCODER_NODES` | int | `1` | distributed | Number of leading nodes that FASTVIDEO_H3_ENCODER_SPLIT dedicates to the MiniMax-H3 text encoder. Values below 1 count as 1. | +| `FASTVIDEO_MINIMAX_H3_FUSIONS` | str | `""` | performance | MiniMax-H3 inference-only Triton fusions: all, 1, or a comma-separated subset of modulate,qknorm_rope,swiglu. Empty, 0, or none keeps the eager implementation. | +| `FASTVIDEO_MINIMAX_H3_EXACT_KERNELS` | str | `""` | performance | MiniMax-H3 inference-only Triton kernels that reproduce the eager output bit for bit: all, 1, or a comma-separated subset of rope,modulate,swiglu. An op whose FASTVIDEO_MINIMAX_H3_FUSIONS fusion is on keeps the fusion. Empty, 0, or none keeps the eager implementation. | +| `FASTVIDEO_FSDP2_AUTOWRAP` | bool | `0` | performance | FSDP2 shards modules by parameter count instead of the model's shard conditions. Not supported by self-forcing distillation. | +| `FASTVIDEO_FSDP2_MIN_PARAMS` | int | `10000000` | performance | Minimum parameter count of a module that FASTVIDEO_FSDP2_AUTOWRAP shards. | +| `FASTVIDEO_MLX_COMPILE` | bool | `0` | performance | Compile the MLX DiT forward with mx.compile. | +| `FASTVIDEO_MLX_FAST_NORM` | bool | `0` | performance | Use MLX fast normalization kernels. | +| `FASTVIDEO_MLX_DQ_GEMM` | str | `1` | performance | MLX dequantized GEMM for affine-quantized weights: 0 turns it off, 1 uses the measured minimum row count, and an integer sets the minimum row count. | +| `FASTVIDEO_LTX2_VAE_CHANNELS_LAST_3D` | bool | `1` | performance | LTX-2 VAE uses the channels_last_3d memory format. | +| `FASTVIDEO_LTX2_DISABLE_AUDIO_AUTOCAST` | bool | `1` | performance | LTX-2 audio decoding runs without CUDA autocast. Deprecated names: `LTX2_DISABLE_AUDIO_AUTOCAST`. | +| `FASTVIDEO_FLUX2_DISABLE_BF16_REDUCED_PRECISION_REDUCTION` | bool | `0` | performance | Flux denoising disables reduced-precision reductions in bf16 matmuls, which tightens accumulation for the 4-step Klein model. | +| `FASTVIDEO_DISABLE_FUSED_NORM` | bool | `0` | performance | Turn off the Triton-fused residual + LayerNorm + modulate inference path that Wan blocks opt into, and use the eager path. | +| `FASTVIDEO_FFMPEG_BIN` | str | `ffmpeg` | output | ffmpeg executable used to save video with audio. | +| `FASTVIDEO_VIDEO_CODEC` | str | `libx264` | output | ffmpeg video codec for saved videos. | +| `FASTVIDEO_NVENC_PRESET` | str | `p1` | output | NVENC preset when the codec is an \*_nvenc codec. | +| `FASTVIDEO_NVENC_TUNE` | str | `ull` | output | NVENC tune option. | +| `FASTVIDEO_NVENC_RC` | str | `constqp` | output | NVENC rate-control mode. | +| `FASTVIDEO_NVENC_QP` | str | `28` | output | NVENC quantization parameter. | +| `FASTVIDEO_NVENC_BF` | str | `0` | output | NVENC number of B-frames. | +| `FASTVIDEO_X264_PRESET` | str | `ultrafast` | output | x264 preset for non-NVENC codecs. | +| `FASTVIDEO_OUTPUT_PIX_FMT` | str | `yuv420p` | output | ffmpeg pixel format for saved videos. | +| `FASTVIDEO_NVTX_PROFILE` | bool | `0` | profiling | Emit NVTX ranges for external profilers such as Nsight Systems. | +| `FASTVIDEO_TORCH_PROFILER_DIR` | path | unset | profiling | Enables the torch profiler and sets the directory for its traces. Must be an absolute path. | +| `FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES` | bool | `0` | profiling | Torch profiler records shapes. | +| `FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY` | bool | `0` | profiling | Torch profiler profiles memory. | +| `FASTVIDEO_TORCH_PROFILER_WITH_STACK` | bool | `0` | profiling | Torch profiler captures stacks. Costs about 1.5x runtime and 1.4x trace size. | +| `FASTVIDEO_TORCH_PROFILER_WITH_FLOPS` | bool | `0` | profiling | Torch profiler profiles FLOPs. | +| `FASTVIDEO_TORCH_PROFILE_REGIONS` | str | `""` | profiling | Comma-separated profiler regions to record. The torch profiler requires at least one region. | +| `FASTVIDEO_TRACE_ACTIVATIONS` | bool | `0` | debug | Enable activation trace hooks. | +| `FASTVIDEO_TRACE_LAYERS` | str | `""` | debug | Regex filter for traced module names. Empty means all. | +| `FASTVIDEO_TRACE_STATS` | str | `abs_mean,sum` | debug | Comma-separated activation statistics dumped for each output tensor. | +| `FASTVIDEO_TRACE_OUTPUT` | str | `/tmp/fv_trace_.jsonl` | debug | JSONL path for activation traces. The literal <pid> is replaced at runtime. | +| `FASTVIDEO_TRACE_STEPS` | str | `""` | debug | Comma-separated denoising step indices. Empty means all. | +| `FASTVIDEO_H3_VSA_PROBE` | str | unset | debug | Output directory for the VSA-H3 attention-mass probe, which writes one .pt file per step, layer, and rank. Keeps the model out of regional compile. | +| `FASTVIDEO_LTX2_GEMMA_LOG` | str | `""` | debug | Log file for LTX-2 Gemma text-encoder hidden states, used by parity tests. Deprecated names: `LTX2_FASTVIDEO_GEMMA_LOG`. | +| `FASTVIDEO_COSMOS25_LOG_KNOBS` | bool | `0` | debug | Log the Cosmos 2.5 latent-preparation conditioning inputs. | +| `FASTVIDEO_CFG_GATE_STEP` | float | `1.0` | sampling | CFG gating fraction in [0, 1]. Steps before len(timesteps) \* X run the conditional and unconditional forwards; later steps reuse the cached difference. 1.0 disables gating. | +| `FASTVIDEO_LTX2_USE_DISTILLED_SIGMAS` | bool | `1` | sampling | LTX-2 uses the distilled sigma schedule when FastVideoArgs.ltx2_use_distilled_sigmas is also true. Deprecated names: `LTX2_USE_DISTILLED_SIGMAS`. | +| `FASTVIDEO_EVAL_CACHE` | path | computed | eval | Cache directory for evaluation models and datasets. Defaults to $FASTVIDEO_CACHE_ROOT/eval. | +| `FASTVIDEO_PHYSICS_IQ_BUCKET_URL` | str | `https://storage.googleapis.com/physics-iq-benchmark` | eval | Base URL of the Physics-IQ benchmark bucket. | +| `FASTVIDEO_VBENCH_FULL_INFO_JSON` | str | unset | eval | Path to VBench_full_info.json, used instead of the vendored copy. Deprecated names: `VBENCH_FULL_INFO_JSON`. | +| `FASTVIDEO_FVD_REF_FEATURES` | str | unset | eval | Cached reference-feature file for the FVD metric. | +| `FASTVIDEO_FAD_REF_FEATURES` | str | unset | eval | Cached reference-feature file for the audio Frechet distance metric. | +| `FASTVIDEO_H3_VSA_FP4` | bool | `0` | attention | Run MiniMax-H3 VSA attention on the block-sparse SageAttention3 FP4 kernel (sm_120, no-grad, single sequence-parallel rank). | +| `FASTVIDEO_H3_VSA_HEADS_FIRST_TILE` | bool | `0` | performance | VSA-H3 64/128-token tiles: scatter rows straight into the heads-first layout the sparse kernels read, removing their per-call transpose copy of query, key and value. Bit-identical; no-grad, uncompiled CUDA only. | +| `FASTVIDEO_H3_VSA_TILE_FIRST` | bool | `0` | attention | Single-rank MiniMax-H3 VSA with one tile gather of the block input instead of separate Q/K/V/gate scatters. | +| `FASTVIDEO_H3_VSA_SM89_KERNEL` | one of original, bf16, int8 | `original` | attention | Fine-attention kernel for MiniMax-H3 VSA on sm_89: original, bf16, or int8 (INT8 QK, BF16 PV). | +| `FASTVIDEO_H3_SIM_SP_FP8` | bool | `0` | debug | Simulate the FP8 sequence-parallel exchange of the MiniMax-H3 FP4 VSA path on one rank. | +| `FASTVIDEO_H3_FFN_CHUNK_TOKENS` | int | `0` | performance | Inference-only MiniMax-H3 FFN token chunk size; 0 runs the FFN unchunked. | +| `FASTVIDEO_H3_FP8_ATTENTION` | bool | `0` | performance | With NVFP4 layer_profile h3_dit_ffn, run MiniMax-H3 attention projections in FP8. | +| `FASTVIDEO_H3_FP8_GRANULARITY` | one of tensor, channel | `tensor` | performance | FP8 scaling granularity for FASTVIDEO_H3_FP8_ATTENTION. | +| `FASTVIDEO_NVFP4_MM_BACKEND` | str | `auto` | performance | FlashInfer mm_fp4 backend for NVFP4 linears, e.g. auto or cutlass. | +| `FASTVIDEO_NVFP4_ACT_AMAX` | path | unset | performance | JSON of calibrated NVFP4 input amax per linear, keyed b<block>.<sub> or full prefix; sets a static activation scale. | +| `FASTVIDEO_NVFP4_DYNAMIC_ACT` | str | `""` | performance | NVFP4 linears that derive the activation scale per call: all, or comma-separated layer-name suffixes such as ff.fc_out. | +| `FASTVIDEO_H3_ADALN_CACHE` | bool | `0` | performance | Cache MiniMax-H3 AdaLN modulation per timestep instead of keeping the projection weights resident. | +| `FASTVIDEO_H3_ADALN_TABLE` | path | unset | performance | Precomputed MiniMax-H3 AdaLN modulation table; enables the cache and skips loading the AdaLN projection weights. | +| `FASTVIDEO_H3_ADALN_DUMP` | path | unset | debug | Write the MiniMax-H3 AdaLN modulation table to this path while sampling. | +| `FASTVIDEO_H3_SPLICE_TRANSFORMER` | path | unset | eval | Second MiniMax-H3 transformer that runs the late DMD steps (checkpoint step-splice evaluation). | +| `FASTVIDEO_H3_SPLICE_FROM_STEP` | int | `4` | eval | First denoising step run by FASTVIDEO_H3_SPLICE_TRANSFORMER. | +| `FASTVIDEO_H3_ENCODER_LAYERWISE` | bool | `0` | performance | Stream MiniMax-H3 text-encoder language layers through exact-size pinned host memory (text-only prompts). | +| `FASTVIDEO_H3_ENCODER_FUSED_DEQUANT` | bool | `0` | performance | Expand the serialized NVFP4 MiniMax-H3 text encoder with one fused Triton pass on GPUs without FP4 GEMM. | +| `FASTVIDEO_H3_VAE_TILE_BATCH` | int | `1` | performance | Spatial tiles per MiniMax-H3 video VAE decoder call; 1 decodes per tile. | +| `FASTVIDEO_H3_REF2VA_MEMO_ENTRIES` | int | `0` | performance | Entries in each content-keyed MiniMax-H3 memo of reference encodes (Qwen3-VL Ref2VA presentation, VAE keyframe latents); repeat references are reused exactly. 0 disables. | +| `FASTVIDEO_H3_VAE_TILE_PARALLEL` | bool | `0` | performance | Split MiniMax-H3 VAE spatial tiles (not only temporal chunks) across the sequence-parallel ranks in the parallel decode and keyframe encode; bitwise equal to the serial path. | +| `FASTVIDEO_H3_VAE_INT8_SHARED_QKV` | bool | `0` | performance | Share the INT8 activation rotation and quantization across the MiniMax-H3 VAE Q/K/V projections. | +| `FASTVIDEO_H3_VAE_INT8_TRANSPOSE_VIEW` | bool | `0` | performance | Use transposed weight views in the MiniMax-H3 VAE INT8 projections. | +| `FASTVIDEO_H3_VAE_INT8_FUSED_DEQUANT` | bool | `0` | performance | Fused dequantization epilogue for the MiniMax-H3 VAE INT8 projections. | +| `FASTVIDEO_H3_PINNED_SWAP` | bool | `1` | performance | Swap offloaded MiniMax-H3 modules through exact-size pinned host arenas. | +| `FASTVIDEO_H3_PARK_MODULES` | str | unset | performance | Comma-separated MiniMax-H3 denoise modules parked on the host while the text encoder runs, e.g. vae,audio_vae. | +| `FASTVIDEO_LAYERWISE_OFFLOAD_BUFFERS` | bool | `0` | performance | Layerwise offload also streams large buffers such as packed FP4/FP8 weights. | +| `FASTVIDEO_LAYERWISE_RESIDENT_BLOCKS` | int | `0` | performance | Keep the first N layerwise-offloaded blocks resident on the GPU. | +| `FASTVIDEO_H3_SP_PROFILE` | bool | `0` | profiling | CUDA-event spans per stage over one MiniMax-H3 FP4 VSA DiT forward. | +| `FASTVIDEO_H3_CAPTURE_QKV` | path | unset | debug | Directory for captured real MiniMax-H3 Q/K/V attention inputs. | +| `FASTVIDEO_CUDA_MEMORY_CAP_GIB` | float | `0.0` | debug | Cap this process's CUDA allocator at this many GiB to emulate a smaller GPU; 0 leaves it uncapped. | +| `FASTVIDEO_MEMORY_REPORT` | bool | `0` | debug | Log bytes held per pipeline component by device and dtype after loading. | +| `FASTVIDEO_TEST_WAN_S2V_MODEL_PATH` | path | unset | test | Wan2.2-S2V-14B weights for test_wan_s2v.py; unset means official_weights/Wan2.2-S2V-14B in the repository. Deprecated names: `WAN_S2V_MODEL_PATH`. | +| `FASTVIDEO_TEST_LTX2_OVERFIT_DATA_DIR` | str | `data/cats` | test | Raw data directory for preprocess_ltx2_overfit.py. Deprecated names: `LTX2_OVERFIT_DATA_DIR`. | +| `FASTVIDEO_TEST_LTX2_OVERFIT_CAPTION_JSON` | str | `videos2caption_1_sample.json` | test | Caption file, relative to the raw data directory. Deprecated names: `LTX2_OVERFIT_CAPTION_JSON`. | +| `FASTVIDEO_TEST_LTX2_OVERFIT_VIDEO_SUBDIR` | str | `video` | test | Video subdirectory, relative to the raw data directory. Deprecated names: `LTX2_OVERFIT_VIDEO_SUBDIR`. | +| `FASTVIDEO_TEST_LTX2_OVERFIT_OUTPUT_DIR` | str | `data/ltx2_overfit_preprocessed` | test | Output directory for preprocess_ltx2_overfit.py. Deprecated names: `LTX2_OVERFIT_OUTPUT_DIR`. | +| `FASTVIDEO_TEST_LTX2_OVERFIT_MODEL` | str | `FastVideo/LTX2-Distilled-Diffusers` | test | Model repository whose encoders preprocess_ltx2_overfit.py uses. Deprecated names: `LTX2_OVERFIT_MODEL`. | +| `FASTVIDEO_TEST_LTX2_OVERFIT_NUM_COPIES` | int | `4` | test | Number of copies of the overfit sample in the parquet file. Deprecated names: `LTX2_OVERFIT_NUM_COPIES`. | +| `FASTVIDEO_TEST_KANDINSKY5_OVERFIT_DATA_DIR` | str | `data/kandinsky5_overfit` | test | Raw data directory for preprocess_kandinsky5_overfit.py. Deprecated names: `KANDINSKY5_OVERFIT_DATA_DIR`. | +| `FASTVIDEO_TEST_KANDINSKY5_OVERFIT_OUTPUT_DIR` | str | `data/kandinsky5_overfit_preprocessed` | test | Output directory for preprocess_kandinsky5_overfit.py. Deprecated names: `KANDINSKY5_OVERFIT_OUTPUT_DIR`. | +| `FASTVIDEO_TEST_HUNYUAN15_OVERFIT_DATA_DIR` | str | `data/hunyuan15_overfit` | test | Raw data directory for preprocess_hunyuan15_overfit.py. | +| `FASTVIDEO_TEST_HUNYUAN15_OVERFIT_OUTPUT_DIR` | str | `data/hunyuan15_overfit_preprocessed` | test | Output directory for preprocess_hunyuan15_overfit.py. | +| `FASTVIDEO_TEST_SSIM_REFERENCE_HF_REPO` | str | `FastVideo/ssim-reference-videos` | test | Hugging Face repository that holds the SSIM reference videos. Deprecated names: `FASTVIDEO_SSIM_REFERENCE_HF_REPO`. | +| `FASTVIDEO_TEST_SSIM_REFERENCE_HF_REPO_TYPE` | str | `dataset` | test | Repository type of FASTVIDEO_TEST_SSIM_REFERENCE_HF_REPO. Deprecated names: `FASTVIDEO_SSIM_REFERENCE_HF_REPO_TYPE`. | +| `FASTVIDEO_TEST_SSIM_SKIP_REFERENCE_DOWNLOAD` | bool | `0` | test | SSIM tests use local reference videos without downloading. Deprecated names: `FASTVIDEO_SSIM_SKIP_REFERENCE_DOWNLOAD`. | +| `FASTVIDEO_TEST_SSIM_FULL_QUALITY` | bool | `0` | test | SSIM tests use the full-quality sampling configurations. Deprecated names: `FASTVIDEO_SSIM_FULL_QUALITY`. | +| `FASTVIDEO_TEST_NIGHTLY` | bool | `0` | test | Run the nightly end-to-end overfit tests. Deprecated names: `FASTVIDEO_NIGHTLY`. | +| `FASTVIDEO_TEST_ULYSSES_FAULT_RANK` | str | unset | test | Rank that fails in the Ulysses fault-injection test. The test sets it for its worker processes. Deprecated names: `FASTVIDEO_ULYSSES_FAULT_RANK`. | +| `FASTVIDEO_TEST_ULYSSES_FAULT_STAGE` | str | unset | test | Stage that fails in the Ulysses fault-injection test. The test sets it for its worker processes. Deprecated names: `FASTVIDEO_ULYSSES_FAULT_STAGE`. | +| `FASTVIDEO_TEST_GOLDEN_GATE_DIR` | str | unset | test | Local directory of golden-gate reference tensors. Deprecated names: `FASTVIDEO_GOLDEN_GATE_DIR`. | +| `FASTVIDEO_TEST_WAN22_5B_ALLOW_LOW_MEMORY` | bool | `0` | test | Run the MLX Wan2.2 5B real-weights parity test on hosts with little memory. Deprecated names: `FASTVIDEO_WAN22_5B_ALLOW_LOW_MEMORY`. | +| `FASTVIDEO_TEST_WAN22_5B_ROOT` | str | unset | test | Local Wan2.2 5B checkpoint for the MLX real-weights parity test. Deprecated names: `FASTVIDEO_WAN22_5B_ROOT`. | +| `FASTVIDEO_TEST_GRADNORM_UPDATE` | bool | `0` | test | Gradient-norm regression tests update their references. Deprecated names: `FASTVIDEO_GRADNORM_UPDATE`. | +| `FASTVIDEO_TEST_DREAMX_WORLD_SSIM_MODEL_PATH` | str | `FastVideo/DreamX-World-5B-Cam-Diffusers` | test | Model for the DreamX-World camera SSIM test. Deprecated names: `DREAMX_WORLD_SSIM_MODEL_PATH`. | +| `FASTVIDEO_TEST_DREAMX_WORLD_AR_SSIM_MODEL_PATH` | str | `FastVideo/DreamX-World-5B-Diffusers` | test | Model for the DreamX-World autoregressive SSIM test. Deprecated names: `DREAMX_WORLD_AR_SSIM_MODEL_PATH`. | +| `FASTVIDEO_TEST_FLUX_T2I_MODEL_DIR` | str | `black-forest-labs/FLUX.1-dev` | test | Model for the Flux text-to-image SSIM test. Deprecated names: `FLUX_T2I_MODEL_DIR`. | +| `FASTVIDEO_TEST_FLUX_TRANSFORMER_PATH` | str | unset | test | Local Flux transformer for the Flux transformer test. Deprecated names: `FLUX_TRANSFORMER_PATH`. | +| `FASTVIDEO_TEST_GAMECRAFT_MODEL_PATH` | str | `FastVideo/HunyuanGameCraft-Diffusers` | test | Model for the HunyuanGameCraft SSIM test. Deprecated names: `GAMECRAFT_MODEL_PATH`. | +| `FASTVIDEO_TEST_GEN3C_MODEL_PATH` | str | `FastVideo/GEN3C-Cosmos-7B-Diffusers` | test | Model for the GEN3C SSIM test. Deprecated names: `GEN3C_MODEL_PATH`. | +| `FASTVIDEO_TEST_GEN3C_IMAGE_PATH` | str | unset | test | Input image for the GEN3C SSIM test. Deprecated names: `GEN3C_TEST_IMAGE_PATH`. | +| `FASTVIDEO_TEST_GLM_IMAGE_LOCAL_WEIGHTS_DIR` | str | unset | test | Local official GLM-Image weights for the GLM-Image SSIM test. Deprecated names: `GLM_IMAGE_LOCAL_WEIGHTS_DIR`. | +| `FASTVIDEO_TEST_GLM_IMAGE_MODEL_DIR` | str | unset | test | Model for the GLM-Image SSIM test. Deprecated names: `GLM_IMAGE_MODEL_DIR`. | +| `FASTVIDEO_TEST_KANDINSKY5_E2E_NUM_GPUS` | int | `1` | test | GPUs for the Kandinsky5 nightly end-to-end overfit test. Deprecated names: `KANDINSKY5_E2E_NUM_GPUS`. | +| `FASTVIDEO_TEST_KANDINSKY5_E2E_WRITE_REFERENCE` | bool | `0` | test | The Kandinsky5 nightly end-to-end test writes a missing reference video. Deprecated names: `KANDINSKY5_E2E_WRITE_REFERENCE`. | +| `FASTVIDEO_TEST_LONGCAT_MODEL_ROOT` | str | unset | test | Local LongCat-Video checkpoint for the golden-gate test. Deprecated names: `LONGCAT_MODEL_ROOT`. | +| `FASTVIDEO_TEST_MINIMAX_H3_GATE_GOLDEN_DIR` | str | unset | test | Local directory of MiniMax-H3 golden-gate tensors. Deprecated names: `MINIMAX_H3_GATE_GOLDEN_DIR`. | +| `FASTVIDEO_TEST_MINIMAX_H3_GATE_LAYER` | int | `0` | test | Transformer layer that the MiniMax-H3 golden-gate test checks. Deprecated names: `MINIMAX_H3_GATE_LAYER`. | +| `FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT` | str | unset | test | Local MiniMax-H3 checkpoint for the golden-gate test. Deprecated names: `MINIMAX_H3_MODEL_ROOT`. | +| `FASTVIDEO_TEST_SD35_MODEL_DIR` | str | `stabilityai/stable-diffusion-3.5-medium` | test | Model for the Stable Diffusion 3.5 SSIM test. Deprecated names: `SD35_MODEL_DIR`. | +| `FASTVIDEO_TEST_TAEH3_REFERENCE_DIR` | str | unset | test | Upstream taehv checkout for the MLX TAEH3 parity test. Deprecated names: `TAEH3_REFERENCE_DIR`. | +| `FASTVIDEO_TEST_WAN_ANIMATE_MODEL_DIR` | str | unset | test | Local Wan2.2-Animate-14B checkpoint for the Wan-Animate weight tests. Deprecated names: `WAN_ANIMATE_MODEL_PATH`. | +| `FASTVIDEO_TEST_ZIMAGE_MODEL_DIR` | str | `Tongyi-MAI/Z-Image-Turbo` | test | Model for the Z-Image SSIM test. Deprecated names: `ZIMAGE_MODEL_DIR`. | +| `FASTVIDEO_TEST_ZIMAGE_MODEL_REVISION` | str | `f332072aa78be7aecdf3ee76d5c247082da564a6` | test | Hugging Face revision of the Z-Image model for its SSIM test. Deprecated names: `ZIMAGE_MODEL_REVISION`. | Variables that FastVideo no longer reads; setting one logs a warning: diff --git a/examples/inference/basic/basic_fasth3_omniref_pdd.py b/examples/inference/basic/basic_fasth3_omniref_pdd.py index 8f9673c82..a93d57f81 100644 --- a/examples/inference/basic/basic_fasth3_omniref_pdd.py +++ b/examples/inference/basic/basic_fasth3_omniref_pdd.py @@ -21,12 +21,19 @@ 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``. + +``--lossless-accel`` turns on every bit-identical speed-up the pipeline has +(LOSSLESS_ACCEL_ENV below, plus resident encoders, pinned host memory and the +sequence-parallel VAE); the output is bitwise the same as without it. It keeps +the 63 GB Qwen3-VL text encoder resident on every GPU, so it needs about +125 GB per GPU at four GPUs (B200/B300 class). """ from __future__ import annotations import argparse import json +import os from collections.abc import Sequence from pathlib import Path from typing import Any @@ -54,6 +61,21 @@ 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) +# Bit-identical speed-ups that FastVideo reads from the environment. Workers +# inherit it, so these must be set before the generator starts. +LOSSLESS_ACCEL_ENV = { + # Fused NVLink Ulysses all-to-all instead of NCCL plus permute copies. + "FASTVIDEO_ULYSSES_A2A": "auto", + # RoPE, AdaLN modulation / gated residual and SwiGLU kernels that round to + # bf16 exactly where the eager ops do (not the FP32 round-once fusions). + "FASTVIDEO_MINIMAX_H3_EXACT_KERNELS": "all", + # VSA tiles scattered straight into the kernels' heads-first layout. + "FASTVIDEO_H3_VSA_HEADS_FIRST_TILE": "1", + # VAE spatial tiles split across the sequence-parallel ranks. + "FASTVIDEO_H3_VAE_TILE_PARALLEL": "1", + # Content-keyed reuse of repeat reference encodes. + "FASTVIDEO_H3_REF2VA_MEMO_ENTRIES": "16", +} class _AppendReference(argparse.Action): @@ -97,6 +119,9 @@ def build_parser() -> argparse.ArgumentParser: 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)") + parser.add_argument("--lossless-accel", + action="store_true", + help="enable every bit-identical speed-up (output unchanged; about 125 GB per GPU)") return parser @@ -201,7 +226,23 @@ def resolve_model(args: argparse.Namespace) -> tuple[Path, dict[str, Any]]: return compose_model_dir(export_dir, base_dir, composed_dir), contract -def build_generator_config(model_dir: Path, num_gpus: int) -> GeneratorConfig: +def apply_lossless_accel_env() -> None: + """Set LOSSLESS_ACCEL_ENV, keeping any value already in the environment.""" + for name, value in LOSSLESS_ACCEL_ENV.items(): + os.environ.setdefault(name, value) + + +def build_generator_config(model_dir: Path, num_gpus: int, lossless_accel: bool = False) -> GeneratorConfig: + if lossless_accel: + # Resident encoders skip streaming the 63 GB text encoder (and the + # VAEs) from host memory on every clip; pinned buffers copy the + # decoded clip at full PCIe speed; the parallel VAE paths reassemble + # the serial result exactly. + offload = OffloadConfig(dit=False, dit_layerwise=False, text_encoder=False, vae=False, pin_cpu_memory=True) + experimental = {"vae_parallel_decode": num_gpus > 1, "vae_parallel_encode": num_gpus > 1} + else: + offload = OffloadConfig(dit=False, dit_layerwise=False, text_encoder=True, vae=True, pin_cpu_memory=False) + experimental = {} return GeneratorConfig( model_path=str(model_dir), engine=EngineConfig( @@ -209,13 +250,14 @@ def build_generator_config(model_dir: Path, num_gpus: int) -> GeneratorConfig: # 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), + offload=offload, ), 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"), + experimental=experimental, ), ) @@ -231,7 +273,9 @@ def main() -> None: 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)) + if args.lossless_accel: + apply_lossless_accel_env() + generator = VideoGenerator.from_config(build_generator_config(model_dir, args.num_gpus, args.lossless_accel)) try: result = generator.generate( GenerationRequest( diff --git a/fastvideo/attention/backends/video_sparse_attn_h3.py b/fastvideo/attention/backends/video_sparse_attn_h3.py index f74daf99a..10741e394 100644 --- a/fastvideo/attention/backends/video_sparse_attn_h3.py +++ b/fastvideo/attention/backends/video_sparse_attn_h3.py @@ -91,6 +91,8 @@ from fastvideo.attention.backends.video_sparse_attn import (compute_topk, constr get_non_pad_index, get_tile_partition_indices, scatter_into_tile_buf) from fastvideo.attention.backends.video_sparse_attn_h3_probe import probe_enabled, record_probe +from fastvideo.attention.backends.video_sparse_attn_h3_scatter import (scatter_rows_heads_first, + supports_heads_first_scatter) from fastvideo.logger import init_logger logger = init_logger(__name__) @@ -334,6 +336,10 @@ class _MiniMaxH3VSATileBufferHolder: def __init__(self) -> None: self.buffer: torch.Tensor | None = None self.untile_geometry: torch.Tensor | None = None + # FASTVIDEO_H3_VSA_HEADS_FIRST_TILE: a [B, H, S_pad, D] buffer whose + # transpose is the tile buffer, and the geometry it was written for. + self.heads_first_buffer: torch.Tensor | None = None + self.heads_first_geometry: torch.Tensor | None = None @dataclass @@ -669,6 +675,10 @@ class MiniMaxH3VSAImpl(AttentionImpl): holder = attn_metadata.tile_buf_holder if holder is None: raise RuntimeError("VSA-H3 metadata has no builder-owned tile buffer holder") + if (not compiling and attn_metadata.tile_elems in _SM100A_TILE_ELEMS + and envs.FASTVIDEO_H3_VSA_HEADS_FIRST_TILE.get() + and supports_heads_first_scatter(x, attn_metadata.untile_combined_index)): + return self._tile_heads_first(x, attn_metadata, holder, kernel_tiles, needs_sm100a_pair) buffer_matches = (holder.buffer is not None and holder.buffer.shape == target_shape and holder.buffer.dtype == x.dtype and holder.buffer.device == x.device) if buffer_matches and holder.untile_geometry is not attn_metadata.untile_combined_index: @@ -681,6 +691,29 @@ class MiniMaxH3VSAImpl(AttentionImpl): holder.buffer[:, n_tiles * attn_metadata.tile_elems:].zero_() return holder.buffer + @staticmethod + def _tile_heads_first(x: torch.Tensor, attn_metadata: MiniMaxH3VSAMetadata, holder: _MiniMaxH3VSATileBufferHolder, + kernel_tiles: int, needs_sm100a_pair: bool) -> torch.Tensor: + """``tile()`` into a heads-first buffer; returns its ``[B, S_pad, H, D]`` view. + + The 64/128-token kernels transpose the tiled tensor to BHSD and make it + contiguous before every call; written heads-first, that copy is free. + Same aliasing and zero-padding contract as the BSHD buffer. + """ + n_tiles = attn_metadata.variable_block_sizes.numel() + shape = (x.shape[0], x.shape[2], kernel_tiles * attn_metadata.tile_elems, x.shape[3]) + buffer = holder.heads_first_buffer + if buffer is None or buffer.shape != shape or buffer.dtype != x.dtype or buffer.device != x.device: + buffer = torch.zeros(shape, device=x.device, dtype=x.dtype) + elif holder.heads_first_geometry is not attn_metadata.untile_combined_index: + buffer.zero_() + scatter_rows_heads_first(x, attn_metadata.untile_combined_index, buffer) + holder.heads_first_buffer = buffer + holder.heads_first_geometry = attn_metadata.untile_combined_index + if needs_sm100a_pair: + buffer[:, :, n_tiles * attn_metadata.tile_elems:].zero_() + return buffer.transpose(1, 2) + def preprocess_qkv(self, qkv: torch.Tensor, attn_metadata: MiniMaxH3VSAMetadata) -> torch.Tensor: return self.tile(qkv, attn_metadata) diff --git a/fastvideo/attention/backends/video_sparse_attn_h3_scatter.py b/fastvideo/attention/backends/video_sparse_attn_h3_scatter.py new file mode 100644 index 000000000..0cb066f0d --- /dev/null +++ b/fastvideo/attention/backends/video_sparse_attn_h3_scatter.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Heads-first tile scatter for VSA-H3 (``FASTVIDEO_H3_VSA_HEADS_FIRST_TILE``). + +``MiniMaxH3VSAImpl.tile`` scatters packed ``[B, S, H, D]`` rows into a padded +``[B, S_pad, H, D]`` tile buffer, and the 64/128-token kernels then copy that +buffer to ``[B, H, S_pad, D]`` with ``transpose(1, 2).contiguous()`` before +every call, for query, key and value. This kernel scatters the rows straight +into a ``[B, H, S_pad, D]`` buffer instead, and ``tile`` returns its +``transpose(1, 2)`` view: the same logical tensor, so every consumer reads the +same values, and the kernels' ``.contiguous()`` becomes free. It only moves +bytes, so the result is bit-identical. +""" + +from __future__ import annotations + +import torch + +try: + import triton + import triton.language as tl +except ImportError: # pragma: no cover - CPU-only installs + triton = None + tl = None + +HAVE_TRITON = triton is not None +_BLOCK_S = 64 + +if HAVE_TRITON: + + @triton.jit + def _scatter_rows_heads_first_kernel(src, dst, slots, seq, s_sn, s_ss, s_sh, d_sn, d_sh, d_ss, D: tl.constexpr, + BLOCK_S: tl.constexpr): + rows = tl.program_id(0) * BLOCK_S + tl.arange(0, BLOCK_S) + h = tl.program_id(1).to(tl.int64) + n = tl.program_id(2).to(tl.int64) + keep = rows < seq + slot = tl.load(slots + rows, mask=keep, other=0).to(tl.int64) + cols = tl.arange(0, D) + values = tl.load(src + n * s_sn + rows[:, None].to(tl.int64) * s_ss + h * s_sh + cols[None, :], + mask=keep[:, None]) + tl.store(dst + n * d_sn + h * d_sh + slot[:, None] * d_ss + cols[None, :], values, mask=keep[:, None]) + + +def supports_heads_first_scatter(x: torch.Tensor, slots: torch.Tensor) -> bool: + """Whether ``scatter_rows_heads_first`` handles ``x`` ``[B, S, H, D]`` and its row slots.""" + dim = x.shape[-1] + return (HAVE_TRITON and x.is_cuda and x.dim() == 4 and x.stride(-1) == 1 and dim > 0 and dim & (dim - 1) == 0 + and slots.is_cuda and slots.dim() == 1 and slots.shape[0] == x.shape[1] and slots.is_contiguous()) + + +def scatter_rows_heads_first(x: torch.Tensor, slots: torch.Tensor, buffer: torch.Tensor) -> None: + """``buffer.transpose(1, 2)[:, slots] = x`` for ``x`` ``[B, S, H, D]`` and ``buffer`` ``[B, H, S_pad, D]``.""" + batch, seq, heads, dim = x.shape + _scatter_rows_heads_first_kernel[(triton.cdiv(seq, _BLOCK_S), heads, batch)]( + x, + buffer, + slots, + seq, + x.stride(0), + x.stride(1), + x.stride(2), + buffer.stride(0), + buffer.stride(1), + buffer.stride(2), + D=dim, + BLOCK_S=_BLOCK_S, + ) + + +__all__ = ["HAVE_TRITON", "scatter_rows_heads_first", "supports_heads_first_scatter"] diff --git a/fastvideo/envs.py b/fastvideo/envs.py index d1cfb62ac..696a32b6d 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -426,6 +426,12 @@ FASTVIDEO_MINIMAX_H3_FUSIONS = EnvStr( category="performance", doc="MiniMax-H3 inference-only Triton fusions: all, 1, or a comma-separated subset of " "modulate,qknorm_rope,swiglu. Empty, 0, or none keeps the eager implementation.") +FASTVIDEO_MINIMAX_H3_EXACT_KERNELS = EnvStr( + "", + category="performance", + doc="MiniMax-H3 inference-only Triton kernels that reproduce the eager output bit for bit: all, 1, or a " + "comma-separated subset of rope,modulate,swiglu. An op whose FASTVIDEO_MINIMAX_H3_FUSIONS fusion is on keeps " + "the fusion. Empty, 0, or none keeps the eager implementation.") FASTVIDEO_FSDP2_AUTOWRAP = EnvBool(False, category="performance", doc="FSDP2 shards modules by parameter count instead of the model's shard " @@ -576,6 +582,11 @@ FASTVIDEO_H3_VSA_FP4 = EnvBool(False, category="attention", doc="Run MiniMax-H3 VSA attention on the block-sparse SageAttention3 FP4 kernel " "(sm_120, no-grad, single sequence-parallel rank).") +FASTVIDEO_H3_VSA_HEADS_FIRST_TILE = EnvBool(False, + category="performance", + doc="VSA-H3 64/128-token tiles: scatter rows straight into the heads-first " + "layout the sparse kernels read, removing their per-call transpose copy " + "of query, key and value. Bit-identical; no-grad, uncompiled CUDA only.") FASTVIDEO_H3_VSA_TILE_FIRST = EnvBool(False, category="attention", doc="Single-rank MiniMax-H3 VSA with one tile gather of the block input " @@ -642,6 +653,16 @@ FASTVIDEO_H3_VAE_TILE_BATCH = EnvInt(1, category="performance", doc="Spatial tiles per MiniMax-H3 video VAE decoder call; 1 decodes per " "tile.") +FASTVIDEO_H3_REF2VA_MEMO_ENTRIES = EnvInt(0, + category="performance", + doc="Entries in each content-keyed MiniMax-H3 memo of reference encodes " + "(Qwen3-VL Ref2VA presentation, VAE keyframe latents); repeat " + "references are reused exactly. 0 disables.") +FASTVIDEO_H3_VAE_TILE_PARALLEL = EnvBool(False, + category="performance", + doc="Split MiniMax-H3 VAE spatial tiles (not only temporal chunks) across " + "the sequence-parallel ranks in the parallel decode and keyframe encode; " + "bitwise equal to the serial path.") FASTVIDEO_H3_VAE_INT8_SHARED_QKV = EnvBool(False, category="performance", doc="Share the INT8 activation rotation and quantization across the " diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py index dd9a645a1..070d13637 100644 --- a/fastvideo/models/dits/minimax_h3.py +++ b/fastvideo/models/dits/minimax_h3.py @@ -37,6 +37,7 @@ from fastvideo.logger import init_logger from fastvideo.models.dits.base import BaseDiT from fastvideo.models.dits.minimax_h3_vsa_fp4 import (STAGES, vsa_fp4_attention, vsa_fp4_attention_sp, vsa_fp4_requested, vsa_tile_first_attention) +from fastvideo.models.dits.minimax_h3_fusions import exact as exact_kernels from fastvideo.models.dits.minimax_h3_fusions import ( HAVE_TRITON, fused_qknorm_rope, @@ -53,6 +54,7 @@ logger = init_logger(__name__) MINIMAX_H3_MODALITY_NUM = 3 _CFG = MiniMaxH3Config() _MINIMAX_H3_FUSION_NAMES = frozenset({"modulate", "qknorm_rope", "swiglu"}) +_MINIMAX_H3_EXACT_KERNEL_NAMES = frozenset({"rope", "modulate", "swiglu"}) def _enabled_minimax_h3_fusions(value: str | None = None) -> frozenset[str]: @@ -71,6 +73,22 @@ def _enabled_minimax_h3_fusions(value: str | None = None) -> frozenset[str]: return enabled +def _enabled_minimax_h3_exact_kernels(value: str | None = None) -> frozenset[str]: + """Parse the bit-exact kernel set (``FASTVIDEO_MINIMAX_H3_EXACT_KERNELS``).""" + raw = envs.FASTVIDEO_MINIMAX_H3_EXACT_KERNELS.get() if value is None else value + normalized = raw.strip().lower() + if normalized in {"", "0", "none"}: + return frozenset() + if normalized in {"1", "all"}: + return _MINIMAX_H3_EXACT_KERNEL_NAMES + enabled = frozenset(item.strip() for item in normalized.split(",") if item.strip()) + unknown = enabled - _MINIMAX_H3_EXACT_KERNEL_NAMES + if unknown: + supported = ",".join(sorted(_MINIMAX_H3_EXACT_KERNEL_NAMES)) + raise ValueError(f"Unknown MiniMax H3 exact kernel(s) {sorted(unknown)}; expected a subset of {supported}.") + return enabled + + def _can_run_minimax_h3_fusion(tensor: torch.Tensor) -> bool: """Triton kernels are inference-only and trace as opaque custom ops. @@ -120,6 +138,7 @@ class MiniMaxH3FeedForward(nn.Module): quant_config: QuantizationConfig | None = None, prefix: str = "", fuse_swiglu: bool = False, + exact_swiglu: bool = False, ) -> None: super().__init__() self.fc_in = ReplicatedLinear( @@ -137,6 +156,7 @@ class MiniMaxH3FeedForward(nn.Module): prefix=f"{prefix}.fc_out", ) self.fuse_swiglu = fuse_swiglu + self.exact_swiglu = exact_swiglu self.use_mxfp8 = isinstance(self.fc_in.quant_method, MXFP8QuantizeMethod) and isinstance( self.fc_out.quant_method, MXFP8QuantizeMethod) # Inference-only token chunking: the 2 * ffn_dim intermediate is ~5.3x the block input @@ -162,6 +182,8 @@ class MiniMaxH3FeedForward(nn.Module): hidden_states, _ = self.fc_in(hidden_states) if self.fuse_swiglu and _can_run_minimax_h3_fusion(hidden_states): hidden_states = minimax_h3_swiglu(hidden_states) + elif self.exact_swiglu and exact_kernels.supports_rowwise(hidden_states): + hidden_states = exact_kernels.swiglu(hidden_states) else: hidden_states, gate = hidden_states.chunk(2, dim=-1) hidden_states = hidden_states * F.silu(gate) @@ -183,6 +205,7 @@ class MiniMaxH3Attention(nn.Module): prefix: str, fuse_qknorm_rope: bool = False, fa4_packed_varlen: bool = False, + exact_rope: bool = False, ) -> None: super().__init__() self.num_attention_heads = num_attention_heads @@ -219,6 +242,7 @@ class MiniMaxH3Attention(nn.Module): prefix=f"{prefix}.to_out", ) self.fuse_qknorm_rope = fuse_qknorm_rope + self.exact_rope = exact_rope # VSA carries a learned gate on its pooled-compression branch. The H3 # checkpoint has no such weight, so the loader zero-initializes it # (ALLOWED_NEW_PARAM_PATTERNS) and the branch is exactly disabled @@ -315,6 +339,11 @@ class MiniMaxH3Attention(nn.Module): hidden_states_rotary = hidden_states_rotary * cos + hidden_states_rotated * sin return torch.cat((hidden_states_rotary, hidden_states_pass), dim=-1).contiguous() + def _rotate(self, hidden_states: torch.Tensor, rotary_emb: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor: + if self.exact_rope and exact_kernels.supports_rope(hidden_states, rotary_emb): + return exact_kernels.rope_prefix(hidden_states, rotary_emb) + return self._apply_rotary_emb(hidden_states, rotary_emb) + def forward( self, hidden_states: torch.Tensor, @@ -369,8 +398,8 @@ class MiniMaxH3Attention(nn.Module): query = self.norm_q(query) key = self.norm_k(key) if rotary_emb is not None: - query = self._apply_rotary_emb(query, rotary_emb) - key = self._apply_rotary_emb(key, rotary_emb) + query = self._rotate(query, rotary_emb) + key = self._rotate(key, rotary_emb) # H3 rotates only 96/128 channels, which the generic `freqs_cis` # branch cannot express. Apply it above, then pass no RoPE here. @@ -602,6 +631,7 @@ class MiniMaxH3TransformerBlock(nn.Module): fuse_qknorm_rope: bool = False, fuse_swiglu: bool = False, fa4_packed_varlen: bool = False, + exact_kernels: frozenset[str] = frozenset(), ) -> None: super().__init__() self.norm1 = nn.RMSNorm(hidden_size, eps=norm_eps) @@ -615,6 +645,7 @@ class MiniMaxH3TransformerBlock(nn.Module): prefix=f"{prefix}.attn", fuse_qknorm_rope=fuse_qknorm_rope, fa4_packed_varlen=fa4_packed_varlen, + exact_rope="rope" in exact_kernels, ) self.norm2 = nn.RMSNorm(hidden_size, eps=norm_eps) self.ff = MiniMaxH3FeedForward( @@ -623,6 +654,7 @@ class MiniMaxH3TransformerBlock(nn.Module): quant_config=quant_config, prefix=f"{prefix}.ff", fuse_swiglu=fuse_swiglu, + exact_swiglu="swiglu" in exact_kernels, ) self.adaln_proj = MiniMaxH3AdaLayerNormModulation( time_embed_dim, @@ -632,6 +664,7 @@ class MiniMaxH3TransformerBlock(nn.Module): apply_silu=adaln_apply_silu, ) self.fuse_modulate = fuse_modulate + self.exact_modulate = "modulate" in exact_kernels def forward( self, @@ -646,6 +679,10 @@ class MiniMaxH3TransformerBlock(nn.Module): t.to(hidden_states.dtype) for t in self.adaln_proj(temb)) use_modulate_fusion = self.fuse_modulate and _can_run_minimax_h3_fusion(hidden_states) + # The bit-exact kernels replace only the eager expressions below. + use_exact_modulate = (not use_modulate_fusion and self.exact_modulate and adaln_indices.is_contiguous() + and exact_kernels.supports_rowwise(hidden_states, shift_msa, scale_msa, gate_msa, + shift_mlp, scale_mlp, gate_mlp)) if use_modulate_fusion: with nvtx_range("minimax_h3.transformer_block.modulate_fusion"): norm_hidden_states = fused_rmsnorm_modulate( @@ -656,6 +693,10 @@ class MiniMaxH3TransformerBlock(nn.Module): adaln_indices, self.norm1.eps, ) + elif use_exact_modulate: + with nvtx_range("minimax_h3.transformer_block.exact_modulate"): + norm_hidden_states = exact_kernels.modulate(self.norm1(hidden_states), scale_msa, shift_msa, + adaln_indices) else: with nvtx_range("minimax_h3.transformer_block.no_modulate_fusion"): norm_hidden_states = self.norm1(hidden_states) @@ -678,6 +719,11 @@ class MiniMaxH3TransformerBlock(nn.Module): adaln_indices, self.norm2.eps, ) + elif use_exact_modulate: + with nvtx_range("minimax_h3.transformer_block.exact_modulate"): + hidden_states = exact_kernels.gate_residual(hidden_states, gate_msa, attention_output, adaln_indices) + norm_hidden_states = exact_kernels.modulate(self.norm2(hidden_states), scale_mlp, shift_mlp, + adaln_indices) else: with nvtx_range("minimax_h3.transformer_block.no_modulate_fusion"): hidden_states = hidden_states + gate_msa.index_select(0, adaln_indices) * attention_output @@ -689,6 +735,8 @@ class MiniMaxH3TransformerBlock(nn.Module): feed_forward_output = self.ff(norm_hidden_states) if use_modulate_fusion and not torch.compiler.is_compiling(): return _gated_residual(hidden_states, gate_mlp, adaln_indices, feed_forward_output) + if use_exact_modulate: + return exact_kernels.gate_residual(hidden_states, gate_mlp, feed_forward_output, adaln_indices) return hidden_states + gate_mlp.index_select(0, adaln_indices) * feed_forward_output @@ -746,6 +794,17 @@ class MiniMaxH3Transformer3DModel(BaseDiT): logger.warning( "FASTVIDEO_MINIMAX_H3_FUSIONS requested %s but Triton is unavailable; " "every forward stays on the eager path.", ",".join(sorted(self.enabled_fusions))) + self.enabled_exact_kernels = _enabled_minimax_h3_exact_kernels() + if self.enabled_exact_kernels: + if exact_kernels.HAVE_TRITON: + logger.info( + "MiniMax H3 bit-exact kernels enabled: %s (CUDA BF16 inference-only; other inputs and any " + "op whose FASTVIDEO_MINIMAX_H3_FUSIONS fusion is on keep their existing path).", + ",".join(sorted(self.enabled_exact_kernels))) + else: + logger.warning( + "FASTVIDEO_MINIMAX_H3_EXACT_KERNELS requested %s but Triton is unavailable; " + "every forward stays on the eager path.", ",".join(sorted(self.enabled_exact_kernels))) sp_world_size = get_sp_world_size() if model_parallel_is_initialized() else 1 if arch.num_attention_heads % sp_world_size: raise ValueError(f"MiniMax H3 attention heads ({arch.num_attention_heads}) must be divisible by " @@ -847,6 +906,7 @@ class MiniMaxH3Transformer3DModel(BaseDiT): fuse_qknorm_rope="qknorm_rope" in self.enabled_fusions, fuse_swiglu="swiglu" in self.enabled_fusions, fa4_packed_varlen=envs.FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN.get(), + exact_kernels=self.enabled_exact_kernels, ) for index in range(arch.num_layers) ]) self.norm_out = MiniMaxH3AdaLayerNormOut( diff --git a/fastvideo/models/dits/minimax_h3_fusions/exact.py b/fastvideo/models/dits/minimax_h3_fusions/exact.py new file mode 100644 index 000000000..d860e7fc7 --- /dev/null +++ b/fastvideo/models/dits/minimax_h3_fusions/exact.py @@ -0,0 +1,264 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Bit-exact Triton replacements for MiniMax-H3's eager RoPE, AdaLN and SwiGLU. + +Each kernel performs exactly the BF16-rounded operation sequence the eager +PyTorch expression performs, in the same order, so its output equals the eager +output bit for bit. That is the difference from the Sol-Engine fusions in this +package (``FASTVIDEO_MINIMAX_H3_FUSIONS``), which keep the whole expression in +FP32 and round once at the store: faster still, but not the eager bits. + +The kernels gain by fusing the eager expression's separate elementwise +launches (and the ``index_select`` gathers of the AdaLN tables) into one pass +over the activations, not by changing the arithmetic. Two details keep the +rounding exact: + +- every FP32 multiply, add and divide is issued as explicit round-to-nearest + PTX (``mul.rn``/``add.rn``/``div.rn``), so the compiler cannot contract a + multiply and an add into one FMA, or fold the FP32 result and the BF16 + downcast into one BF16 instruction, either of which rounds once where eager + rounds twice; +- every intermediate eager materializes as a BF16 tensor is rounded to BF16 + (round-to-nearest-even) at the same point. + +Selected with ``FASTVIDEO_MINIMAX_H3_EXACT_KERNELS``. Every entry point has a +``supports_*`` predicate; callers fall back to the eager expression for inputs +outside it. +""" + +from __future__ import annotations + +import torch + +try: + import triton + import triton.language as tl +except ImportError: # pragma: no cover - CPU-only installs + triton = None + tl = None + +HAVE_TRITON = triton is not None + +_BLOCK_C = 1024 + +if HAVE_TRITON: + + @triton.jit + def _mul_rn(a, b): + return tl.inline_asm_elementwise("mul.rn.f32 $0, $1, $2;", + "=r,r,r", [a, b], + dtype=tl.float32, + is_pure=True, + pack=1) + + @triton.jit + def _add_rn(a, b): + return tl.inline_asm_elementwise("add.rn.f32 $0, $1, $2;", + "=r,r,r", [a, b], + dtype=tl.float32, + is_pure=True, + pack=1) + + @triton.jit + def _div_rn(a, b): + return tl.inline_asm_elementwise("div.rn.f32 $0, $1, $2;", + "=r,r,r", [a, b], + dtype=tl.float32, + is_pure=True, + pack=1) + + @triton.jit + def _bf16(v): + return v.to(tl.bfloat16, fp_downcast_rounding="rtne") + + @triton.jit + def _rope_prefix_kernel(x, cos, sin, out, heads, x_ss, x_sh, c_ss, D: tl.constexpr, R: tl.constexpr, + RP: tl.constexpr, TAIL: tl.constexpr, BLOCK_H: tl.constexpr): + # out[..., :R] = x * cos + rotate_half(x) * sin; out[..., R:] = x[..., R:]. + s = tl.program_id(0).to(tl.int64) + hs = tl.program_id(1) * BLOCK_H + tl.arange(0, BLOCK_H) + hmask = hs < heads + base = s * x_ss + hs[:, None].to(tl.int64) * x_sh + half: tl.constexpr = R // 2 + d = tl.arange(0, RP) + rmask = d < R + partner = tl.where(d < half, d + half, d - half) + m = hmask[:, None] & rmask[None, :] + xv = tl.load(x + base + d[None, :], mask=m, other=0.0).to(tl.float32) + pv = tl.load(x + base + partner[None, :], mask=m, other=0.0).to(tl.float32) + # torch.cat((-second_half, first_half)): negation is exact. + pv = _mul_rn(pv, tl.where(d < half, -1.0, 1.0)[None, :]) + cv = tl.load(cos + s * c_ss + d, mask=rmask, other=0.0).to(tl.float32) + sv = tl.load(sin + s * c_ss + d, mask=rmask, other=0.0).to(tl.float32) + a = _bf16(_mul_rn(xv, cv[None, :])).to(tl.float32) + b = _bf16(_mul_rn(pv, sv[None, :])).to(tl.float32) + tl.store(out + base + d[None, :], _bf16(_add_rn(a, b)), mask=m) + if D - R > 0: + tail = R + tl.arange(0, TAIL) + tmask = hmask[:, None] & (tail < D)[None, :] + tl.store(out + base + tail[None, :], tl.load(x + base + tail[None, :], mask=tmask), mask=tmask) + + @triton.jit + def _modulate_kernel(n, scale, shift, idx, out, cols, n_rs, sc_rs, sh_rs, o_rs, BLOCK_C: tl.constexpr): + # out = n * (1 + scale[idx]) + shift[idx], each op rounded to BF16 as eager does. + row = tl.program_id(0).to(tl.int64) + c = tl.program_id(1) * BLOCK_C + tl.arange(0, BLOCK_C) + m = c < cols + r = tl.load(idx + row).to(tl.int64) + nv = tl.load(n + row * n_rs + c, mask=m).to(tl.float32) + sc = tl.load(scale + r * sc_rs + c, mask=m).to(tl.float32) + sh = tl.load(shift + r * sh_rs + c, mask=m).to(tl.float32) + one_plus = _bf16(_add_rn(sc, tl.full([BLOCK_C], 1.0, tl.float32))).to(tl.float32) + y = _bf16(_mul_rn(nv, one_plus)).to(tl.float32) + tl.store(out + row * o_rs + c, _bf16(_add_rn(y, sh)), mask=m) + + @triton.jit + def _gate_residual_kernel(h, gate, y, idx, out, cols, h_rs, y_rs, g_rs, o_rs, BLOCK_C: tl.constexpr): + # out = h + gate[idx] * y, each op rounded to BF16 as eager does. + row = tl.program_id(0).to(tl.int64) + c = tl.program_id(1) * BLOCK_C + tl.arange(0, BLOCK_C) + m = c < cols + r = tl.load(idx + row).to(tl.int64) + hv = tl.load(h + row * h_rs + c, mask=m).to(tl.float32) + gv = tl.load(gate + r * g_rs + c, mask=m).to(tl.float32) + yv = tl.load(y + row * y_rs + c, mask=m).to(tl.float32) + prod = _bf16(_mul_rn(gv, yv)).to(tl.float32) + tl.store(out + row * o_rs + c, _bf16(_add_rn(hv, prod)), mask=m) + + @triton.jit + def _swiglu_kernel(x, out, cols, x_rs, o_rs, BLOCK_C: tl.constexpr): + # Value-first halves: out = value * silu(gate). Eager F.silu computes + # g / (1 + exp(-g)) in FP32 and stores BF16; the product rounds again. + row = tl.program_id(0).to(tl.int64) + c = tl.program_id(1) * BLOCK_C + tl.arange(0, BLOCK_C) + m = c < cols + v = tl.load(x + row * x_rs + c, mask=m).to(tl.float32) + g = tl.load(x + row * x_rs + cols + c, mask=m).to(tl.float32) + e = tl.extra.cuda.libdevice.exp(-g) + s = _bf16(_div_rn(g, _add_rn(tl.full([BLOCK_C], 1.0, tl.float32), e))).to(tl.float32) + tl.store(out + row * o_rs + c, _bf16(_mul_rn(v, s)), mask=m) + + +def _rows(x: torch.Tensor) -> torch.Tensor: + return x.reshape(-1, x.shape[-1]) + + +def _flattens_to_rows(t: torch.Tensor) -> bool: + if t.dim() <= 2: + return True + try: + t.view(-1, t.shape[-1]) + except RuntimeError: + return False + return True + + +def supports_rowwise(*tensors: torch.Tensor) -> bool: + """Whether ``modulate``/``gate_residual``/``swiglu`` reproduce eager on these tensors.""" + return (HAVE_TRITON and not torch.is_grad_enabled() and not torch.compiler.is_compiling() + and all(t.is_cuda and t.dtype == torch.bfloat16 and t.stride(-1) == 1 and _flattens_to_rows(t) + for t in tensors)) + + +def supports_rope(hidden_states: torch.Tensor, rotary_emb: tuple[torch.Tensor, torch.Tensor] | None) -> bool: + """Whether ``rope_prefix`` reproduces eager ``_apply_rotary_emb`` on these inputs.""" + if rotary_emb is None or not HAVE_TRITON or torch.is_grad_enabled() or torch.compiler.is_compiling(): + return False + cos, sin = rotary_emb + dim = hidden_states.shape[-1] + return (hidden_states.is_cuda and hidden_states.dtype == torch.bfloat16 and hidden_states.dim() == 4 + and hidden_states.shape[0] == 1 and hidden_states.is_contiguous() and dim & (dim - 1) == 0 + and cos.dim() == 2 and cos.shape == sin.shape and cos.shape[0] == hidden_states.shape[1] + and cos.shape[-1] % 2 == 0 and cos.shape[-1] <= dim) + + +def rope_prefix(hidden_states: torch.Tensor, rotary_emb: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor: + """Eager ``MiniMaxH3Attention._apply_rotary_emb`` for contiguous BF16 ``[1, S, H, D]``.""" + cos, sin = rotary_emb + cos = cos.to(hidden_states.dtype).contiguous() + sin = sin.to(hidden_states.dtype).contiguous() + _, seq, heads, dim = hidden_states.shape + rotary = cos.shape[-1] + out = torch.empty_like(hidden_states) + block_h = 4 + _rope_prefix_kernel[(seq, triton.cdiv(heads, block_h))]( + hidden_states, + cos, + sin, + out, + heads, + hidden_states.stride(1), + hidden_states.stride(2), + cos.stride(0), + D=dim, + R=rotary, + RP=triton.next_power_of_2(rotary), + TAIL=triton.next_power_of_2(max(dim - rotary, 1)), + BLOCK_H=block_h, + ) + return out + + +def modulate(normed: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: + """``normed * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)``.""" + n = _rows(normed) + indices = indices.contiguous() + out = torch.empty_like(n) + cols = n.shape[1] + _modulate_kernel[(n.shape[0], triton.cdiv(cols, _BLOCK_C))](n, + scale, + shift, + indices, + out, + cols, + n.stride(0), + scale.stride(0), + shift.stride(0), + out.stride(0), + BLOCK_C=_BLOCK_C) + return out.view(normed.shape) + + +def gate_residual(hidden: torch.Tensor, gate: torch.Tensor, update: torch.Tensor, + indices: torch.Tensor) -> torch.Tensor: + """``hidden + gate.index_select(0, indices) * update``.""" + h, y = _rows(hidden), _rows(update) + indices = indices.contiguous() + out = torch.empty_like(h) + cols = h.shape[1] + _gate_residual_kernel[(h.shape[0], triton.cdiv(cols, _BLOCK_C))](h, + gate, + y, + indices, + out, + cols, + h.stride(0), + y.stride(0), + gate.stride(0), + out.stride(0), + BLOCK_C=_BLOCK_C) + return out.view(hidden.shape) + + +def swiglu(packed: torch.Tensor) -> torch.Tensor: + """``value, gate = packed.chunk(2, -1); value * F.silu(gate)``.""" + x = _rows(packed) + cols = x.shape[1] // 2 + out = torch.empty((x.shape[0], cols), device=x.device, dtype=x.dtype) + _swiglu_kernel[(x.shape[0], triton.cdiv(cols, _BLOCK_C))](x, + out, + cols, + x.stride(0), + out.stride(0), + BLOCK_C=_BLOCK_C) + return out.view(*packed.shape[:-1], cols) + + +__all__ = [ + "HAVE_TRITON", + "gate_residual", + "modulate", + "rope_prefix", + "supports_rope", + "supports_rowwise", + "swiglu", +] diff --git a/fastvideo/models/vaes/minimax_h3_parallel.py b/fastvideo/models/vaes/minimax_h3_parallel.py index 5f738161e..415e2239c 100644 --- a/fastvideo/models/vaes/minimax_h3_parallel.py +++ b/fastvideo/models/vaes/minimax_h3_parallel.py @@ -47,10 +47,12 @@ from typing import TYPE_CHECKING import torch +import fastvideo.envs as envs from fastvideo.models.vaes.minimax_h3_video import ( AutoencoderKLMiniMaxH3, AutoencoderKLOutput, DiagonalGaussianDistribution, + _tile_batch_size, ) from fastvideo.profiler import nvtx_range @@ -233,6 +235,151 @@ def decode_to_pixels_parallel( return output +# -- spatial-tile split (FASTVIDEO_H3_VAE_TILE_PARALLEL) --------------------- +# +# The chunk split above leaves ranks idle when a clip has few temporal chunks +# (a 124-frame 768p clip has fewer chunks than eight ranks hold work for) and +# cannot split a one-frame keyframe at all. With spatial tiling on, every +# chunk (or keyframe) is itself a grid of independent tile forwards that the +# serial path stitches afterwards. The tile split hands out (chunk, tile) +# forwards round-robin instead, gathers them, and stitches each chunk on the +# assembling rank with the serial ``_stitch_tiles`` before the unchanged chunk +# join, so the result is bitwise the serial one under the same contract as the +# chunk split. + + +def _tile_split_applies(vae: AutoencoderKLMiniMaxH3, x: torch.Tensor) -> bool: + # The batched tile decode (FASTVIDEO_H3_VAE_TILE_BATCH) is a different + # serial path; the tile split replays only the per-tile one. + return (envs.FASTVIDEO_H3_VAE_TILE_PARALLEL.get() and vae.use_tiling and not vae._tile_helpers_compiled + and x.shape[0] == 1 and _tile_batch_size() == 1 and not torch.is_grad_enabled()) + + +def _tile_grid(vae: AutoencoderKLMiniMaxH3, height: int, width: int): + """The serial tile split of a ``height`` x ``width`` pixel canvas: rows, columns, and their boxes.""" + rows = vae._split_tiles(height, vae.tile_sample_min_height, vae.tile_sample_min_overlap_height) + columns = vae._split_tiles(width, vae.tile_sample_min_width, vae.tile_sample_min_overlap_width) + boxes = [(y, h, x, w) for y, h in zip(rows[0], rows[1]) for x, w in zip(columns[0], columns[1])] + return rows, columns, boxes + + +def _as_grid(tiles: list[torch.Tensor], num_columns: int) -> list[list[torch.Tensor]]: + return [tiles[start:start + num_columns] for start in range(0, len(tiles), num_columns)] + + +def _tile_grid_is_exact(rows, columns, height: int, width: int, ratio: int) -> bool: + """True when the grid tiles the canvas into equal shapes with ratio-aligned tile sizes. + + The gathers need equal tile shapes across ranks, and the decode slice must + replay the serial ``y // ratio:y // ratio + h // ratio`` exactly. Grids whose + last tile clamps (an unabsorbed ``_split_tiles`` residual) or whose tile size + is not a multiple of the spatial ratio stand down to the serial path. + """ + for starts, lengths, axis in ((rows[0], rows[1], height), (columns[0], columns[1], width)): + if starts[-1] + lengths[-1] != axis: + return False + if len(lengths) > 1 and lengths[0] % ratio != 0: + return False + return True + + +def encode_keyframe_tile_parallel(vae: AutoencoderKLMiniMaxH3, x: torch.Tensor, + group: "GroupCoordinator") -> torch.Tensor | None: + """``vae._encode_clip(x)`` for a one-frame keyframe with its tiles split across ``group``. + + All group ranks call this together with identical ``x`` and all receive + the moments. Returns ``None`` (on every rank alike) when the split does not + apply; the caller then runs the serial encode. + """ + if group.world_size == 1 or not _tile_split_applies(vae, x): + return None + rows, columns, boxes = _tile_grid(vae, *x.shape[-2:]) + if (len(boxes) < group.world_size + or not _tile_grid_is_exact(rows, columns, *x.shape[-2:], vae.spatial_compression_ratio)): + return None + world_size, rank = group.world_size, group.rank_in_group + tiles: list[torch.Tensor] = [] + placeholder = None + for round_start in range(0, len(boxes), world_size): + index = round_start + rank + if index < len(boxes): + y, h, x0, w = boxes[index] + with nvtx_range(f"minimax_h3.vae.parallel_keyframe_tile.{index}"): + mine = vae.quant_conv(vae.encoder(x[..., y:y + h, x0:x0 + w])).contiguous() + placeholder = mine + else: + # Only the last round can be short, and every rank owned a tile in + # round 0, so the placeholder shape is the tile shape of this grid. + mine = torch.zeros_like(placeholder) + gathered = group.all_gather(mine.unsqueeze(0), dim=0) + tiles.extend(gathered[slot] for slot in range(min(world_size, len(boxes) - round_start))) + ratio = vae.spatial_compression_ratio + return vae._stitch_tiles(_as_grid(tiles, len(columns[0])), [o // ratio for o in rows[2]], + [o // ratio for o in columns[2]]) + + +def _decode_single_tile_parallel( + vae: AutoencoderKLMiniMaxH3, + z: torch.Tensor, + output: torch.Tensor | None, + group: "GroupCoordinator", + strategy: str, +) -> bool: + """The tile split of ``_decode_single_parallel``; False (on every rank alike) when it does not apply.""" + if not _tile_split_applies(vae, z): + return False + pad_tokens, num_chunks, output_num_frames = vae._temporal_decode_plan(z.shape[2]) + ratio = vae.spatial_compression_ratio + rows, columns, boxes = _tile_grid(vae, z.shape[-2] * ratio, z.shape[-1] * ratio) + jobs = [(chunk, box) for chunk in range(num_chunks) for box in range(len(boxes))] + world_size, rank = group.world_size, group.rank_in_group + if (len(jobs) < world_size + or not _tile_grid_is_exact(rows, columns, z.shape[-2] * ratio, z.shape[-1] * ratio, ratio)): + return False + if pad_tokens > 0: + z = torch.cat([z, z[:, :, -1:].repeat(1, 1, pad_tokens, 1, 1)], dim=2) + assembler = None + if output is not None: + assembler = _ChunkAssembler(vae, output, output_num_frames, vae._streams_chunk_copies(z, output), z.device) + pending: dict[tuple[int, int], torch.Tensor] = {} + next_chunk = 0 + placeholder = None + try: + for round_start in range(0, len(jobs), world_size): + index = round_start + rank + if index < len(jobs): + chunk, box = jobs[index] + start = chunk * vae.tokens_chunk_size + y, h, x, w = boxes[box] + tile = z[:, :, start:start + vae.tokens_chunk_size + vae.token_overlap, y // ratio:(y + h) // ratio, + x // ratio:(x + w) // ratio] + with nvtx_range(f"minimax_h3.vae.parallel_tile.{chunk}.{box}"): + mine = vae.decoder(vae._project_decoder_tile(tile)).contiguous() + placeholder = mine + else: + mine = torch.zeros_like(placeholder) + with nvtx_range(f"minimax_h3.vae.parallel_tile_{strategy}.{round_start // world_size}"): + if strategy == "gather": + gathered = group.gather(mine.unsqueeze(0), dst=0, dim=0) + else: + gathered = group.all_gather(mine.unsqueeze(0), dim=0) + if assembler is None or gathered is None: + continue + for slot in range(min(world_size, len(jobs) - round_start)): + pending[jobs[round_start + slot]] = gathered[slot] + while next_chunk < num_chunks and all((next_chunk, box) in pending for box in range(len(boxes))): + tiles = [pending.pop((next_chunk, box)) for box in range(len(boxes))] + clip = vae._stitch_tiles(_as_grid(tiles, len(columns[0])), rows[2], columns[2]) + assembler.push(clip[:, :, vae.frame_pre_padding:].contiguous()) + next_chunk += 1 + if assembler is not None: + assembler.finalize() + finally: + if assembler is not None: + assembler.synchronize() + return True + + def _decode_single_parallel( vae: AutoencoderKLMiniMaxH3, z: torch.Tensor, @@ -240,6 +387,8 @@ def _decode_single_parallel( group: "GroupCoordinator", strategy: str, ) -> None: + if _decode_single_tile_parallel(vae, z, output, group, strategy): + return pad_tokens, num_chunks, output_num_frames = vae._temporal_decode_plan(z.shape[2]) if pad_tokens > 0: z = torch.cat([z, z[:, :, -1:].repeat(1, 1, pad_tokens, 1, 1)], dim=2) diff --git a/fastvideo/models/vaes/minimax_h3_video.py b/fastvideo/models/vaes/minimax_h3_video.py index fca78c244..fd9ea6cc6 100644 --- a/fastvideo/models/vaes/minimax_h3_video.py +++ b/fastvideo/models/vaes/minimax_h3_video.py @@ -1100,15 +1100,32 @@ class AutoencoderKLMiniMaxH3(nn.Module): """Encode one-frame conditioning inputs without video chunk padding.""" if x.ndim != 5 or x.shape[2] != 1: raise ValueError(f"`x` must contain exactly one video frame, got shape {tuple(x.shape)}.") - if self.use_slicing and x.shape[0] > 1: - moments = torch.cat([self._encode_clip(x_slice) for x_slice in x.split(1)]) - else: - moments = self._encode_clip(x) + moments = self._encode_keyframe_tile_parallel(x) + if moments is None: + if self.use_slicing and x.shape[0] > 1: + moments = torch.cat([self._encode_clip(x_slice) for x_slice in x.split(1)]) + else: + moments = self._encode_clip(x) posterior = DiagonalGaussianDistribution(moments) if not return_dict: return (posterior, ) return AutoencoderKLOutput(latent_dist=posterior) + def _encode_keyframe_tile_parallel(self, x: torch.Tensor) -> torch.Tensor | None: + """FASTVIDEO_H3_VAE_TILE_PARALLEL: split a keyframe's tiles across the sequence-parallel group. + + Every sequence-parallel rank encodes the same keyframes in the same + order (the pipeline prepares references identically on all ranks), so + the collective inside is entered uniformly. + """ + if not envs.FASTVIDEO_H3_VAE_TILE_PARALLEL.get(): + return None + from fastvideo.distributed import get_sp_group, model_parallel_is_initialized + if not model_parallel_is_initialized(): + return None + from fastvideo.models.vaes.minimax_h3_parallel import encode_keyframe_tile_parallel + return encode_keyframe_tile_parallel(self, x, get_sp_group()) + def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]: if self.use_slicing and z.shape[0] > 1: decoded = torch.cat([self._decode(z_slice) for z_slice in z.split(1)]) diff --git a/fastvideo/pipelines/basic/minimax_h3/memo.py b/fastvideo/pipelines/basic/minimax_h3/memo.py new file mode 100644 index 000000000..0c852d948 --- /dev/null +++ b/fastvideo/pipelines/basic/minimax_h3/memo.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Content-keyed memo for MiniMax-H3 reference encodes. + +Serving the same reference images again (a re-roll, a new seed, the next turn +of a session) re-runs the Qwen3-VL presentation and the VAE keyframe encode on +identical inputs. Both are pure functions of their inputs — the keyframe +posterior sample uses a fixed-seed generator — so their outputs can be reused +exactly. Keys hash the pixel bytes, never object identity, so a mutated or +re-decoded image misses. Entries are cloned on the way in and out, so callers +may mutate what they get back. + +Every sequence-parallel rank prepares the same references in the same order, +so hits and misses are uniform across ranks and an encode that enters a +collective is either skipped or run on all of them. +""" + +from __future__ import annotations + +import hashlib +from collections import OrderedDict +from typing import Any + +import numpy as np +import torch + +import fastvideo.envs as envs + + +def image_key(image: Any) -> tuple: + """Shape, dtype and a 128-bit digest of an image's pixel bytes.""" + pixels = np.ascontiguousarray(np.asarray(image)) + return pixels.shape, str(pixels.dtype), hashlib.blake2b(pixels.tobytes(), digest_size=16).hexdigest() + + +def _clone(value: Any) -> Any: + if isinstance(value, tuple): + return tuple(_clone(item) for item in value) + return value.clone() + + +class ContentMemo: + """A small LRU of cloned tensor (or tuple-of-tensor) results; capacity 0 disables it.""" + + def __init__(self, capacity: int | None = None) -> None: + self.capacity = max(0, envs.FASTVIDEO_H3_REF2VA_MEMO_ENTRIES.get() if capacity is None else capacity) + self._items: OrderedDict[Any, Any] = OrderedDict() + + @property + def enabled(self) -> bool: + # Training needs the graph through the encoders; a memo would cut it. + return self.capacity > 0 and not torch.is_grad_enabled() + + def __len__(self) -> int: + return len(self._items) + + def get_or_compute(self, key: Any, compute) -> Any: + if not self.enabled: + return compute() + cached = self._items.get(key) + if cached is None: + value = compute() + self._items[key] = _clone(value) + while len(self._items) > self.capacity: + self._items.popitem(last=False) + return value + self._items.move_to_end(key) + return _clone(cached) + + def clear(self) -> None: + self._items.clear() diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_conditioning.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_conditioning.py index 1e26ee540..02145e9b8 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_conditioning.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_conditioning.py @@ -12,6 +12,7 @@ from torch.distributed.tensor import DTensor from fastvideo.distributed import get_local_torch_device from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.models.encoders.minimax_h3_qwen3_vl import MiniMaxH3Qwen3VLConditioner +from fastvideo.pipelines.basic.minimax_h3.memo import ContentMemo, image_key from fastvideo.profiler import nvtx_range from fastvideo.pipelines.basic.minimax_h3.packing import ( MINIMAX_H3_IMAGE_PAD_TOKEN, @@ -113,6 +114,7 @@ class MiniMaxH3ConditioningStage(PipelineStage): self.tokenizer = tokenizer self.processor = processor self.ref2va = ref2va + self._ref2va_memo = ContentMemo() def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: result = VerificationResult() @@ -201,7 +203,21 @@ class MiniMaxH3ConditioningStage(PipelineStage): references = list(batch.references or []) if not references or not all(isinstance(item, MiniMaxH3PreparedReference) for item in references): raise TypeError("MiniMax-H3 Ref2VA conditioning requires prepared references.") + # Video references get their block timestamps assigned while presenting, + # so only presentations without a video reference are memoized. The key + # hashes pixels, so a disabled memo must not pay for building it. + has_video = any(reference.media_type == "video" for reference in references) + images_ready = all(reference.image is not None for reference in references if reference.media_type == "image") + if self._ref2va_memo.enabled and not has_video and images_ready: + key = (str(device), batch.prompt, + tuple((reference.media_type, bool(reference.has_audio), + image_key(reference.image) if reference.media_type == "image" else None) + for reference in references)) + return self._ref2va_memo.get_or_compute(key, lambda: self._present_ref2va(batch, references, device)) + return self._present_ref2va(batch, references, device) + def _present_ref2va(self, batch: ForwardBatch, references: list[MiniMaxH3PreparedReference], + device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: merge_area = int(self.processor.image_processor.merge_size)**2 pixel_values = None image_grid_thw = None diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_latent_preparation.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_latent_preparation.py index bd2482731..f684f01dd 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_latent_preparation.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_latent_preparation.py @@ -14,6 +14,7 @@ from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger from fastvideo.models import pinned_offload from fastvideo.models.vaes.minimax_h3_parallel import encode_pixels_parallel +from fastvideo.pipelines.basic.minimax_h3.memo import ContentMemo, image_key from fastvideo.pipelines.basic.minimax_h3.packing import ( MINIMAX_H3_AUDIO_CHANNELS, MINIMAX_H3_KEYFRAME_ENCODE_SEED, @@ -90,6 +91,7 @@ class MiniMaxH3LatentPreparationStage(PipelineStage): self.audio_vae = audio_vae self.scheduler = scheduler self.ref2va = ref2va + self._keyframe_memo = ContentMemo() def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: result = VerificationResult() @@ -112,6 +114,13 @@ class MiniMaxH3LatentPreparationStage(PipelineStage): def _encode_keyframe_latents(self, image, device: torch.device) -> torch.Tensor: """Encode one keyframe image to normalized latents (fp16 round-trip preserved for parity with the seeded references).""" + if not self._keyframe_memo.enabled: + # The key hashes pixels, so a disabled memo must not pay for building it. + return self._encode_keyframe_latents_uncached(image, device) + return self._keyframe_memo.get_or_compute((str(device), image_key(image)), + lambda: self._encode_keyframe_latents_uncached(image, device)) + + def _encode_keyframe_latents_uncached(self, image, device: torch.device) -> torch.Tensor: pixels = torch.from_numpy(np.asarray(image).copy()).permute(2, 0, 1)[None, :, None] pixels = pixels.to(device=device, dtype=torch.float32).div_(255.0) posterior = self.vae.encode_keyframe(self.vae.normalize_pixels(pixels)).latent_dist diff --git a/fastvideo/tests/attention/test_vsa_h3_heads_first_tile.py b/fastvideo/tests/attention/test_vsa_h3_heads_first_tile.py new file mode 100644 index 000000000..8b320c8a1 --- /dev/null +++ b/fastvideo/tests/attention/test_vsa_h3_heads_first_tile.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FASTVIDEO_H3_VSA_HEADS_FIRST_TILE must hand VSA-H3 the same tiled tensor, bit for bit.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +import fastvideo.envs as envs +from fastvideo.attention.backends import video_sparse_attn_h3 as vsa +from fastvideo.attention.backends.video_sparse_attn_h3_scatter import (HAVE_TRITON, scatter_rows_heads_first, + supports_heads_first_scatter) + + +def _require_cuda_triton() -> None: + if not torch.cuda.is_available(): + pytest.skip("the heads-first tile scatter requires CUDA") + if not HAVE_TRITON: + pytest.skip("the heads-first tile scatter requires Triton") + + +def _same_bits(a: torch.Tensor, b: torch.Tensor) -> bool: + return a.shape == b.shape and torch.equal(a.contiguous().view(torch.int16), b.contiguous().view(torch.int16)) + + +def _metadata(seq: int, n_tiles: int, tile_elems: int) -> SimpleNamespace: + slots = torch.randperm(n_tiles * tile_elems, device="cuda")[:seq].sort().values + sizes = torch.bincount(slots // tile_elems, minlength=n_tiles) + return SimpleNamespace(total_seq_length=seq, + variable_block_sizes=sizes, + tile_elems=tile_elems, + untile_combined_index=slots, + tile_buf_holder=vsa._MiniMaxH3VSATileBufferHolder()) + + +def _impl() -> vsa.MiniMaxH3VSAImpl: + impl = object.__new__(vsa.MiniMaxH3VSAImpl) + impl._regional_compile_sm100a_enabled = None + return impl + + +@pytest.fixture +def sm100a(): + with envs.FASTVIDEO_VSA_SM100A.override(True): + yield + + +@pytest.mark.gpu +def test_scatter_rows_heads_first_matches_indexed_assignment() -> None: + _require_cuda_triton() + torch.manual_seed(0) + seq, heads, dim, slots_total = 5000, 12, 128, 6144 + x = torch.randn(1, seq, heads, dim, device="cuda").to(torch.bfloat16) + slots = torch.randperm(slots_total, device="cuda")[:seq] + assert supports_heads_first_scatter(x, slots) + expected = torch.zeros(1, slots_total, heads, dim, device="cuda", dtype=x.dtype) + expected[:, slots] = x + buffer = torch.zeros(1, heads, slots_total, dim, device="cuda", dtype=x.dtype) + scatter_rows_heads_first(x, slots, buffer) + assert _same_bits(expected, buffer.transpose(1, 2)) + + +@pytest.mark.gpu +@pytest.mark.parametrize(("tile_elems", "n_tiles"), [(128, 37), (128, 40), (64, 41)]) +def test_heads_first_tile_equals_the_default_tile(sm100a, tile_elems: int, n_tiles: int) -> None: + _require_cuda_triton() + torch.manual_seed(1) + meta = _metadata(seq=n_tiles * tile_elems - 300, n_tiles=n_tiles, tile_elems=tile_elems) + x = torch.randn(1, meta.total_seq_length, 8, 128, device="cuda").to(torch.bfloat16) + impl = _impl() + with torch.inference_mode(): + with envs.FASTVIDEO_H3_VSA_HEADS_FIRST_TILE.override(False): + default = impl.tile(x, meta).clone() + with envs.FASTVIDEO_H3_VSA_HEADS_FIRST_TILE.override(True): + heads_first = impl.tile(x, meta) + assert meta.tile_buf_holder.heads_first_buffer is not None + assert heads_first.transpose(1, 2).is_contiguous() + assert _same_bits(default, heads_first) + # The tile-score pooling reads the logical prefix of the view directly + # (forward() drops the sm100a pairing tile); its FP32 sums must not move. + logical = n_tiles * tile_elems + assert _same_bits(vsa._pool_tiles(default[:, :logical], meta.variable_block_sizes, tile_elems), + vsa._pool_tiles(heads_first[:, :logical], meta.variable_block_sizes, tile_elems)) + # A new geometry on the reused buffer leaves no stale row behind. + other = _metadata(seq=meta.total_seq_length, n_tiles=n_tiles, tile_elems=tile_elems) + other.tile_buf_holder = meta.tile_buf_holder + with envs.FASTVIDEO_H3_VSA_HEADS_FIRST_TILE.override(False): + default = impl.tile(x, other).clone() + with envs.FASTVIDEO_H3_VSA_HEADS_FIRST_TILE.override(True): + assert _same_bits(default, impl.tile(x, other)) diff --git a/fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py b/fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py index 26db7a4c8..0c83eceda 100644 --- a/fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py +++ b/fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py @@ -2,6 +2,7 @@ """CPU checks of examples/inference/basic/basic_fasth3_omniref_pdd.py (no weights, no GPU).""" from __future__ import annotations +import contextlib import importlib.util import json from pathlib import Path @@ -120,3 +121,41 @@ def test_build_generator_config_multi_gpu_shards_dit(tmp_path, num_gpus, sharded 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" + + +def test_default_config_offloads_the_encoders(tmp_path): + offload = example.build_generator_config(tmp_path, 4).engine.offload + assert (offload.text_encoder, offload.vae, offload.pin_cpu_memory) == (True, True, False) + + +@pytest.mark.parametrize("num_gpus", [1, 4]) +def test_lossless_accel_config(tmp_path, num_gpus): + config = example.build_generator_config(tmp_path, num_gpus, lossless_accel=True) + offload = config.engine.offload + assert (offload.text_encoder, offload.vae, offload.pin_cpu_memory) == (False, False, True) + assert config.pipeline.experimental == { + "vae_parallel_decode": num_gpus > 1, + "vae_parallel_encode": num_gpus > 1, + } + + +def test_lossless_accel_env_names_registered_flags_and_keeps_user_values() -> None: + from fastvideo import envs + + with contextlib.ExitStack() as stack: + for name in example.LOSSLESS_ACCEL_ENV: + assert name in envs.environment_variables + stack.enter_context(envs.environment_variables[name].override(None)) + stack.enter_context(envs.FASTVIDEO_ULYSSES_A2A.override("off")) + example.apply_lossless_accel_env() + assert envs.FASTVIDEO_ULYSSES_A2A.get() == "off" + assert envs.FASTVIDEO_MINIMAX_H3_EXACT_KERNELS.get() == "all" + assert envs.FASTVIDEO_H3_VSA_HEADS_FIRST_TILE.get() is True + assert envs.FASTVIDEO_H3_VAE_TILE_PARALLEL.get() is True + assert envs.FASTVIDEO_H3_REF2VA_MEMO_ENTRIES.get() == 16 + + +def test_lossless_accel_flag_parses(): + argv = ["--model-path", "m", "--prompt", "p", "--image", "a.png"] + assert example.parse_args(argv).lossless_accel is False + assert example.parse_args([*argv, "--lossless-accel"]).lossless_accel is True diff --git a/fastvideo/tests/stages/test_minimax_h3_ref2va_memo.py b/fastvideo/tests/stages/test_minimax_h3_ref2va_memo.py new file mode 100644 index 000000000..b020b7aea --- /dev/null +++ b/fastvideo/tests/stages/test_minimax_h3_ref2va_memo.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU tests for the content-keyed MiniMax-H3 reference-encode memo.""" + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from PIL import Image + +import fastvideo.envs as envs +from fastvideo.pipelines.basic.minimax_h3.memo import ContentMemo, image_key +from fastvideo.pipelines.basic.minimax_h3.reference import MiniMaxH3PreparedReference +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_conditioning import MiniMaxH3ConditioningStage +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import ( + MiniMaxH3LatentPreparationStage, ) + + +def _image(seed: int, size: tuple[int, int] = (8, 6)) -> Image.Image: + rng = np.random.default_rng(seed) + return Image.fromarray(rng.integers(0, 256, (size[1], size[0], 3), dtype=np.uint8)) + + +@torch.no_grad() +def test_memo_returns_independent_clones_and_evicts_lru() -> None: + memo = ContentMemo(capacity=2) + calls = [] + + def compute(value): + calls.append(value) + return torch.full((3, ), float(value)), torch.tensor([value]) + + first = memo.get_or_compute("a", lambda: compute(1)) + first[0].add_(100) + again = memo.get_or_compute("a", lambda: compute(1)) + assert calls == [1] + assert torch.equal(again[0], torch.full((3, ), 1.0)) + again[0].add_(100) + assert torch.equal(memo.get_or_compute("a", lambda: compute(1))[0], torch.full((3, ), 1.0)) + + memo.get_or_compute("b", lambda: compute(2)) + memo.get_or_compute("a", lambda: compute(1)) + memo.get_or_compute("c", lambda: compute(3)) + assert len(memo) == 2 + memo.get_or_compute("b", lambda: compute(2)) + assert calls == [1, 2, 3, 2] + + +def test_memo_is_off_by_default_and_under_autograd() -> None: + with envs.FASTVIDEO_H3_REF2VA_MEMO_ENTRIES.override(None): + assert not ContentMemo().enabled + with envs.FASTVIDEO_H3_REF2VA_MEMO_ENTRIES.override(4): + memo = ContentMemo() + assert memo.capacity == 4 + with torch.enable_grad(): + assert not memo.enabled + memo.get_or_compute("a", lambda: torch.zeros(1)) + assert len(memo) == 0 + + +def test_image_key_is_content_keyed() -> None: + assert image_key(_image(0)) == image_key(_image(0)) + assert image_key(_image(0)) != image_key(_image(1)) + pixels = np.asarray(_image(0)).copy() + changed = pixels.copy() + changed[0, 0, 0] ^= 1 + assert image_key(pixels) != image_key(changed) + assert image_key(pixels) != image_key(pixels.astype(np.uint16)) + assert image_key(pixels) != image_key(pixels.reshape(-1, 3)) + + +@pytest.fixture +def memo_entries(): + with envs.FASTVIDEO_H3_REF2VA_MEMO_ENTRIES.override(4): + yield + + +def _conditioning_stage(monkeypatch): + stage = MiniMaxH3ConditioningStage(conditioner=None, tokenizer=None, processor=None, ref2va=True) + calls = [] + + def present(batch, references, device): + calls.append(batch.prompt) + return torch.randn(1, 5, 4), torch.zeros(5, dtype=torch.long) + + monkeypatch.setattr(stage, "_present_ref2va", present) + return stage, calls + + +@torch.no_grad() +def test_ref2va_presentation_is_reused_for_identical_content(monkeypatch, memo_entries) -> None: + stage, calls = _conditioning_stage(monkeypatch) + device = torch.device("cpu") + + def batch(prompt, *seeds, audio=False): + references = [MiniMaxH3PreparedReference("image", image=_image(seed)) for seed in seeds] + if audio: + references.append(MiniMaxH3PreparedReference("audio", has_audio=True)) + return SimpleNamespace(prompt=prompt, references=references) + + first = stage._encode_ref2va(batch("p", 0, 1), device) + repeat = stage._encode_ref2va(batch("p", 0, 1), device) + assert len(calls) == 1 + assert all(torch.equal(a, b) for a, b in zip(first, repeat, strict=True)) + + stage._encode_ref2va(batch("q", 0, 1), device) + stage._encode_ref2va(batch("p", 1, 0), device) + stage._encode_ref2va(batch("p", 0, 2), device) + stage._encode_ref2va(batch("p", 0, 1, audio=True), device) + assert len(calls) == 5 + + +@torch.no_grad() +def test_ref2va_with_a_video_reference_is_never_memoized(monkeypatch, memo_entries) -> None: + stage, calls = _conditioning_stage(monkeypatch) + references = [ + MiniMaxH3PreparedReference("image", image=_image(0)), + MiniMaxH3PreparedReference("video", frames=np.zeros((4, 6, 8, 3), dtype=np.uint8)), + ] + for _ in range(2): + stage._encode_ref2va(SimpleNamespace(prompt="p", references=references), torch.device("cpu")) + assert len(calls) == 2 + + +@torch.no_grad() +def test_keyframe_latents_are_reused_for_identical_pixels(monkeypatch, memo_entries) -> None: + stage = MiniMaxH3LatentPreparationStage(vae=None, audio_vae=None, scheduler=None, ref2va=True) + calls = [] + + def encode(image, device): + calls.append(image_key(image)) + return torch.randn(1, 4, 1, 3, 2) + + monkeypatch.setattr(stage, "_encode_keyframe_latents_uncached", encode) + device = torch.device("cpu") + first = stage._encode_keyframe_latents(_image(0), device) + assert torch.equal(stage._encode_keyframe_latents(_image(0), device), first) + stage._encode_keyframe_latents(_image(1), device) + assert len(calls) == 2 + + +@pytest.mark.parametrize("entries", (0, None)) +@torch.no_grad() +def test_stages_recompute_without_the_flag(monkeypatch, entries) -> None: + with envs.FASTVIDEO_H3_REF2VA_MEMO_ENTRIES.override(entries): + stage = MiniMaxH3LatentPreparationStage(vae=None, audio_vae=None, scheduler=None, ref2va=True) + calls = [] + monkeypatch.setattr(stage, "_encode_keyframe_latents_uncached", + lambda image, device: calls.append(1) or torch.zeros(1)) + for _ in range(3): + stage._encode_keyframe_latents(_image(0), torch.device("cpu")) + assert len(calls) == 3 diff --git a/fastvideo/tests/transformers/test_minimax_h3_exact_kernels.py b/fastvideo/tests/transformers/test_minimax_h3_exact_kernels.py new file mode 100644 index 000000000..539055d1e --- /dev/null +++ b/fastvideo/tests/transformers/test_minimax_h3_exact_kernels.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax H3's bit-exact Triton kernels must equal the eager code they replace, bit for bit. + +Unlike the Sol-Engine fusion tests (``test_minimax_h3_*_fusion.py``), which +allow a tolerance, every GPU comparison here is ``torch.equal`` on the raw +BF16 bit patterns. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from fastvideo.models.dits.minimax_h3 import MiniMaxH3Attention, _enabled_minimax_h3_exact_kernels +from fastvideo.models.dits.minimax_h3_fusions import exact + + +def _require_cuda_triton() -> None: + if not torch.cuda.is_available(): + pytest.skip("MiniMax H3 exact kernels require CUDA") + if not exact.HAVE_TRITON: + pytest.skip("MiniMax H3 exact kernels require Triton") + + +def _same_bits(a: torch.Tensor, b: torch.Tensor) -> bool: + return a.shape == b.shape and torch.equal(a.contiguous().view(torch.int16), b.contiguous().view(torch.int16)) + + +def _bf16(*shape: int, scale: float = 1.0) -> torch.Tensor: + return (torch.randn(*shape, device="cuda") * scale).to(torch.bfloat16) + + +@pytest.fixture +def inference(): + torch.manual_seed(0) + with torch.inference_mode(): + yield + + +def test_exact_kernel_set_parsing() -> None: + assert _enabled_minimax_h3_exact_kernels("") == frozenset() + assert _enabled_minimax_h3_exact_kernels("none") == frozenset() + assert _enabled_minimax_h3_exact_kernels("all") == frozenset({"rope", "modulate", "swiglu"}) + assert _enabled_minimax_h3_exact_kernels(" rope, swiglu ") == frozenset({"rope", "swiglu"}) + with pytest.raises(ValueError, match="qknorm_rope"): + _enabled_minimax_h3_exact_kernels("qknorm_rope") + + +def test_exact_kernels_refuse_grad_and_cpu_inputs() -> None: + x = torch.randn(4, 8, dtype=torch.bfloat16) + assert not exact.supports_rowwise(x) + with torch.no_grad(): + assert not exact.supports_rowwise(x) + + +@pytest.mark.gpu +@pytest.mark.parametrize("seq", [777, 9000]) +def test_modulate_and_gated_residual_are_bit_exact(inference, seq: int) -> None: + _require_cuda_triton() + hidden, rows = 6144, 24 + index = torch.randint(0, rows, (seq, ), device="cuda") + # Six chunks of one AdaLN projection: strided views, as the DiT produces them. + shift, scale, gate = _bf16(rows, 6 * hidden, scale=0.5).chunk(6, dim=-1)[:3] + normed, residual, branch = _bf16(1, seq, hidden), _bf16(1, seq, hidden, scale=4), _bf16(1, seq, hidden) + assert exact.supports_rowwise(normed, shift, scale, gate) + eager = normed * (1.0 + scale.index_select(0, index)) + shift.index_select(0, index) + assert _same_bits(eager, exact.modulate(normed, scale, shift, index)) + eager = residual + gate.index_select(0, index) * branch + assert _same_bits(eager, exact.gate_residual(residual, gate, branch, index)) + + +@pytest.mark.gpu +def test_modulate_accepts_tables_with_different_row_strides(inference) -> None: + _require_cuda_triton() + hidden, rows, seq = 1024, 8, 300 + index = torch.randint(0, rows, (seq, ), device="cuda") + scale = _bf16(rows, 3 * hidden).chunk(3, dim=-1)[1] + shift = _bf16(rows, hidden) + normed = _bf16(1, seq, hidden) + eager = normed * (1.0 + scale.index_select(0, index)) + shift.index_select(0, index) + assert _same_bits(eager, exact.modulate(normed, scale, shift, index)) + + +@pytest.mark.gpu +def test_swiglu_is_bit_exact_including_edge_values(inference) -> None: + _require_cuda_triton() + ffn = 16384 + packed = _bf16(1, 777, 2 * ffn, scale=4) + packed[0, 0, ffn:ffn + 8] = torch.tensor([0.0, -0.0, 88.0, -88.0, 100.0, -100.0, 1e-30, -1e-30]) + value, gate = packed.chunk(2, dim=-1) + assert _same_bits(value * F.silu(gate), exact.swiglu(packed)) + + +@pytest.mark.gpu +@pytest.mark.parametrize("seq", [777, 9000]) +def test_rope_prefix_matches_apply_rotary_emb_bit_for_bit(inference, seq: int) -> None: + _require_cuda_triton() + # H3 geometry: 128-channel heads with a 96-channel rotary prefix. + x = _bf16(1, seq, 48, 128, scale=3) + x[0, :5], x[0, 5:10] = 0.0, -0.0 + positions = torch.arange(seq, device="cuda", dtype=torch.float32)[:, None] + freqs = positions * torch.logspace(0, -4, 48, device="cuda")[None] + freqs = torch.cat((freqs, freqs), dim=-1) + rotary_emb = (freqs.cos(), freqs.sin()) + assert exact.supports_rope(x, rotary_emb) + assert _same_bits(MiniMaxH3Attention._apply_rotary_emb(x, rotary_emb), exact.rope_prefix(x, rotary_emb)) diff --git a/fastvideo/tests/vaes/test_minimax_h3_parallel_vae.py b/fastvideo/tests/vaes/test_minimax_h3_parallel_vae.py index a881dcb21..90e91c744 100644 --- a/fastvideo/tests/vaes/test_minimax_h3_parallel_vae.py +++ b/fastvideo/tests/vaes/test_minimax_h3_parallel_vae.py @@ -14,10 +14,12 @@ import pytest import torch from torch.testing import assert_close +import fastvideo.envs as envs from fastvideo.configs.models.vaes.minimax_h3_video import ( MiniMaxH3VideoVAEArchConfig, MiniMaxH3VideoVAEConfig, ) +from fastvideo.models.vaes import minimax_h3_parallel from fastvideo.models.vaes.minimax_h3_parallel import ( DECODE_GATHER_STRATEGIES, DEFAULT_DECODE_GATHER_STRATEGY, @@ -312,3 +314,132 @@ def test_fastvideo_args_strategy_literals_match_module() -> None: vae_parallel_decode_strategy=strategy).vae_parallel_decode_strategy == strategy with pytest.raises(ValueError, match="vae_parallel_decode_strategy"): FastVideoArgs(model_path="test/parallel-vae", vae_parallel_decode_strategy="scatter") + + +@pytest.fixture +def tile_parallel(): + with envs.FASTVIDEO_H3_VAE_TILE_PARALLEL.override(True): + yield + + +def _chunk_split_must_not_run(*args, **kwargs): + raise AssertionError("the tile split fell back to the chunk split") + + +def _tiled_vae(token_drop: int = 3) -> AutoencoderKLMiniMaxH3: + vae = _tiny_vae(token_drop=token_drop) + # 32x48 pixels split into a 3x5 grid of 16x16 tiles with 8-pixel overlaps. + vae.enable_tiling(tile_sample_min_height=16, + tile_sample_min_width=16, + tile_sample_min_overlap_height=8, + tile_sample_min_overlap_width=8) + return vae + + +# 3 latent frames is one chunk (fewer chunks than ranks: only the tile split +# keeps every rank busy); 13 is three chunks with pad trim and blending. +@pytest.mark.parametrize("world_size", (2, 4, 7)) +@pytest.mark.parametrize("latent_frames", (3, 13)) +@pytest.mark.parametrize("strategy", DECODE_GATHER_STRATEGIES) +@torch.inference_mode() +def test_tile_parallel_decode_matches_serial(monkeypatch, tile_parallel, world_size: int, latent_frames: int, + strategy: str) -> None: + torch.manual_seed(20261002 + latent_frames) + vae = _tiled_vae() + latents = torch.randn(1, 4, latent_frames, 8, 12) + expected = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32) + vae.decode_to_pixels(latents, expected) + + monkeypatch.setattr(minimax_h3_parallel, "_decode_segment", _chunk_split_must_not_run) + group = _ThreadedFakeGroup(world_size) + + def _rank_main(rank: int): + output = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32) if rank == 0 else None + return decode_to_pixels_parallel(vae, latents.clone(), output, group, strategy=strategy) + + results = group.run(_rank_main) + assert all(result is None for result in results[1:]) + assert_close(results[0], expected, atol=0.0, rtol=0.0) + + +@torch.inference_mode() +def test_tile_parallel_decode_falls_back_with_fewer_tiles_than_ranks(monkeypatch, tile_parallel) -> None: + """One chunk of a 1x1 tile grid on 3 ranks: the chunk split serves it unchanged.""" + torch.manual_seed(20261003) + vae = _tiny_vae() + vae.enable_tiling(tile_sample_min_height=16, tile_sample_min_width=16) + latents = torch.randn(1, 4, 3, 4, 4) + expected = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32) + vae.decode_to_pixels(latents, expected) + + group = _ThreadedFakeGroup(3) + + def _rank_main(rank: int): + output = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32) if rank == 0 else None + return decode_to_pixels_parallel(vae, latents.clone(), output, group) + + assert_close(group.run(_rank_main)[0], expected, atol=0.0, rtol=0.0) + + +@pytest.mark.parametrize("world_size", (2, 4, 15, 16)) +@torch.inference_mode() +def test_tile_parallel_keyframe_encode_matches_serial(monkeypatch, tile_parallel, world_size: int) -> None: + """Every rank gets the serial keyframe moments; a grid smaller than the group refuses (None).""" + from fastvideo.models.vaes.minimax_h3_parallel import encode_keyframe_tile_parallel + + torch.manual_seed(20261004 + world_size) + vae = _tiled_vae() + keyframe = torch.rand(1, 3, 1, 32, 48) * 2 - 1 + expected = vae._encode_clip(keyframe) + + group = _ThreadedFakeGroup(world_size) + results = group.run(lambda rank: encode_keyframe_tile_parallel(vae, keyframe, group)) + if world_size > 15: + assert all(result is None for result in results) + return + for moments in results: + assert_close(moments, expected, atol=0.0, rtol=0.0) + + +@torch.inference_mode() +def test_tile_parallel_is_off_by_default() -> None: + from fastvideo.models.vaes.minimax_h3_parallel import encode_keyframe_tile_parallel + + with envs.FASTVIDEO_H3_VAE_TILE_PARALLEL.override(None): + vae = _tiled_vae() + group = _ThreadedFakeGroup(2) + group._local.rank = 0 + assert encode_keyframe_tile_parallel(vae, torch.zeros(1, 3, 1, 32, 48), group) is None + assert vae._encode_keyframe_tile_parallel(torch.zeros(1, 3, 1, 32, 48)) is None + + +@torch.inference_mode() +def test_tile_parallel_stands_down_on_a_clamped_grid(tile_parallel) -> None: + """A grid whose last tile clamps cannot be gathered; the serial encode serves it.""" + from fastvideo.models.vaes.minimax_h3_parallel import encode_keyframe_tile_parallel + + vae = _tiny_vae() + vae.enable_tiling(tile_sample_min_height=16, tile_sample_min_width=16, tile_sample_min_overlap_height=8, + tile_sample_min_overlap_width=8) + group = _ThreadedFakeGroup(2) + group._local.rank = 0 + # 31 px split into 16/16/15: the clamped tile breaks the equal-shape gather. + assert encode_keyframe_tile_parallel(vae, torch.zeros(1, 3, 1, 31, 48), group) is None + + +@torch.inference_mode() +def test_encode_keyframe_routes_through_the_tile_split(monkeypatch, tile_parallel) -> None: + import fastvideo.distributed as distributed + + torch.manual_seed(20261005) + vae = _tiled_vae() + keyframe = torch.rand(1, 3, 1, 32, 48) * 2 - 1 + expected = vae.encode_keyframe(keyframe).latent_dist.parameters + + group = _ThreadedFakeGroup(4) + monkeypatch.setattr(distributed, "model_parallel_is_initialized", lambda: True) + monkeypatch.setattr(distributed, "get_sp_group", lambda: group) + monkeypatch.setattr(vae, "_encode_clip", _chunk_split_must_not_run) + results = group.run(lambda rank: vae.encode_keyframe(keyframe).latent_dist.parameters) + for moments in results: + assert_close(moments, expected, atol=0.0, rtol=0.0) diff --git a/scripts/benchmarks/fasth3_omniref_b200/README.md b/scripts/benchmarks/fasth3_omniref_b200/README.md new file mode 100644 index 000000000..1835b086d --- /dev/null +++ b/scripts/benchmarks/fasth3_omniref_b200/README.md @@ -0,0 +1,105 @@ +# FastH3 OmniRef: lossless speed-ups on 4x B200 + +`examples/inference/basic/basic_fasth3_omniref_pdd.py --lossless-accel` turns on +every bit-identical speed-up the MiniMax-H3 Ref2VA pipeline has. The output is +bitwise the same as without the flag: the same uint8 video frames and the same +float audio samples. This directory holds the case and the script that measure +it. + +## The case + +One 5 s clip (124 frames at 24 fps), 1344x768, one reference image, 4x B200, +four-way sequence parallel (`--num-gpus 4`). + +| field | value | +| --- | --- | +| checkpoint | `FastVideo/FastH3-OmniRef-v5-DMPDD8-w03-cfg2-step3500` and the base MiniMax-H3 its `fastvideo_inference.json` pins | +| prompt | `Picture 1 walks slowly toward the camera across the living room, smiling, and says, "Every detail was chosen with care." Smooth tracking shot, warm afternoon light, light footsteps on the wooden floor.` | +| reference | `speaker.png`, 720x400 RGB, sha256 `d38ba667...a6c92c` (shared with the PR; not committed) | +| seed | 413410 | +| sampling | the contract's 8 fused PDD blocks, guidance 1.0, no CFG batch | + +`case_speaker.json` holds the same values. The script checks the reference +image's sha256 before it runs. + +## Method + +```bash +python scripts/benchmarks/fasth3_omniref_b200/bench_lossless.py --mode off \ + --model-path FastVideo/FastH3-OmniRef-v5-DMPDD8-w03-cfg2-step3500 \ + --reference speaker.png --num-gpus 4 --out runs/off +python scripts/benchmarks/fasth3_omniref_b200/bench_lossless.py --mode on \ + --model-path FastVideo/FastH3-OmniRef-v5-DMPDD8-w03-cfg2-step3500 \ + --reference speaker.png --num-gpus 4 --out runs/on +python scripts/benchmarks/fasth3_omniref_b200/bench_lossless.py --compare runs/off runs/on +``` + +Run with `FASTVIDEO_STAGE_LOGGING=1` for per-stage times. On a machine whose +NVLink SHARP setup NCCL cannot use, also set `NCCL_NVLS_ENABLE=0` (both modes). + +Each mode runs in its own process, because the speed-ups are environment flags +that workers read when they start. Both modes build the generator through the +example (the flag is the only difference). Each mode then does three runs: + +1. an untimed warm-up with another prompt and a grey reference, so CUDA, NCCL + and Triton are warm but the reference cache holds nothing from the case; +2. the case, timed (`cold`: the prompt and reference have not been seen); +3. the case again, timed (`repeat`: with the flag on, the reference-encode + cache hits). + +The timing is the wall time of `generate()`: encoders, denoising, VAE decode +and the transfer of the frames to the caller, without the mp4 write. The +script hashes the decoded uint8 frames and the raw float audio with SHA-256, +and `--compare` fails unless every hash matches between the two modes. + +## What `--lossless-accel` turns on + +| setting | what it does | why the bits do not change | +| --- | --- | --- | +| `offload.text_encoder=False`, `offload.vae=False` | keeps Qwen3-VL (63 GB) and the VAEs on the GPU instead of copying them from host memory for every clip | same weights, same kernels | +| `pin_cpu_memory=True` | pinned host buffer for the decoded clip | copy only | +| `vae_parallel_decode` / `vae_parallel_encode` | the existing chunk-parallel VAE paths | already proven equal to serial | +| `FASTVIDEO_ULYSSES_A2A=auto` | the existing fused NVLink Ulysses all-to-all | data movement only | +| `FASTVIDEO_MINIMAX_H3_EXACT_KERNELS=all` | Triton RoPE, AdaLN modulate / gated residual and SwiGLU kernels | rounds to bf16 exactly where the eager ops do | +| `FASTVIDEO_H3_VSA_HEADS_FIRST_TILE=1` | VSA tiles scattered straight into the kernels' heads-first layout | data movement only | +| `FASTVIDEO_H3_VAE_TILE_PARALLEL=1` | VAE spatial tiles split across the ranks (decode and keyframe encode) | same per-tile forwards, same stitch | +| `FASTVIDEO_H3_REF2VA_MEMO_ENTRIES=16` | reuses repeat Qwen3-VL presentations and keyframe latents, keyed by content | pure functions of their inputs | + +The exact kernels are not `FASTVIDEO_MINIMAX_H3_FUSIONS`. Those fusions keep +FP32 across a chain and round once, which is more accurate but changes about +95% of the output pixels. The exact kernels do each eager op's bf16 rounding, +in the same order. They use PTX `mul.rn`/`add.rn`/`div.rn`, because Triton +would otherwise contract a multiply-add into an FMA or fold the intermediate +roundings away. If both are on, the fusions take precedence. + +## Results + +B200 (180 GB), the case above, PyTorch 2.12.0+cu130, `NCCL_NVLS_ENABLE=0`. + +4x B200. These were measured with the same techniques applied as run-time +patches over FastVideo `9edc8adf` plus the OmniRef PR, before they were ported +here: + +| run | off | on | speed-up | bits | +| --- | --- | --- | --- | --- | +| cold | 42 s | 14.7 s | 2.9x | identical | +| repeat | 42 s | 13.1 s | 3.2x | identical | + +Where the time went (seconds, cold, off -> on): conditioning 15-17 -> 1.1, +reference VAE encode 2.5 -> 0.43, denoising 13.8 -> 11.0, VAE decode 7.5 -> 1.4. + +This branch, with this script: + +| GPUs | run | off | on | speed-up | bits | +| --- | --- | --- | --- | --- | --- | +| 1 | cold | 67.1 s | 44.4 s | 1.51x | identical | +| 1 | repeat | 67.2 s | 42.1 s | 1.60x | identical | +| 2 | cold | 54.9 s | 26.6 s | 2.06x | identical | +| 2 | repeat | 54.6 s | 24.9 s | 2.19x | identical | + +Two GPUs, cold, off -> on: conditioning 17.0 -> 1.1 s (0.06 s on repeat), +reference VAE encode 2.7 -> 0.8 s (0.06 s), denoising 26.7 -> 21.5 s, VAE +decode 7.8 -> 2.7 s. All four runs and both GPU counts give the same digests +(frames `a48362b9...`, audio `0f0d42f8...`). One GPU runs none of the multi-rank +paths (all-to-all, VAE splits), and its denoising dominates (46.6 -> 36.2 s), +so the gain is smaller there. diff --git a/scripts/benchmarks/fasth3_omniref_b200/bench_lossless.py b/scripts/benchmarks/fasth3_omniref_b200/bench_lossless.py new file mode 100644 index 000000000..ab4e9966e --- /dev/null +++ b/scripts/benchmarks/fasth3_omniref_b200/bench_lossless.py @@ -0,0 +1,193 @@ +"""Apples-to-apples benchmark of the lossless FastH3 OmniRef speed-ups, with a bit-identity check. + +usage: + python bench_lossless.py --mode off --model-path --out runs/off + python bench_lossless.py --mode on --model-path --out runs/on + python bench_lossless.py --compare runs/off runs/on + +One process per mode: the speed-ups are read from the environment when the +workers start, so off and on cannot share a process. Both modes build the +generator exactly as examples/inference/basic/basic_fasth3_omniref_pdd.py does +(``--lossless-accel`` is the only difference), then: + +1. one untimed warm-up with a different prompt and a grey reference, so CUDA, + NCCL and Triton are warm but the reference memo holds nothing of the case; +2. the case, timed ("cold": prompt and reference unseen); +3. the case again, timed ("repeat": the reference-encode memo hits when on). + +Each timed run records the generate() wall time, FastVideo's per-stage +times, and SHA-256 digests of the decoded uint8 frames and of the raw float +audio. ``--compare`` requires every digest of the two runs to match (bitwise +identical video and audio) and prints the speed-up. +""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.util +import json +import sys +import time +from pathlib import Path + +import numpy as np + +HERE = Path(__file__).resolve().parent +EXAMPLE = HERE.parents[2] / "examples" / "inference" / "basic" / "basic_fasth3_omniref_pdd.py" +WARMUP_PROMPT = "A grey studio backdrop, still camera, room tone." + + +def _example(): + spec = importlib.util.spec_from_file_location("basic_fasth3_omniref_pdd", EXAMPLE) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _sha256(array) -> str: + return hashlib.sha256(np.ascontiguousarray(array).tobytes()).hexdigest() + + +def _stage_times(result) -> dict[str, float]: + stages = getattr(getattr(result, "logging_info", None), "stages", None) or {} + return { + name: round(float(metrics["execution_time"]), 3) + for name, metrics in stages.items() if metrics.get("execution_time") is not None + } + + +def _generate(generator, contract, case, image, seed, prompt, out_dir): + from fastvideo.api import GenerationRequest, InputConfig, OutputConfig, SamplingConfig + from fastvideo.pipelines.basic.minimax_h3 import MiniMaxH3Reference + + started = time.perf_counter() + result = generator.generate( + GenerationRequest( + prompt=prompt, + negative_prompt="", + inputs=InputConfig(references=[MiniMaxH3Reference(source=image, media_type="image")]), + sampling=SamplingConfig( + height=case["height"], + width=case["width"], + num_frames=case["num_frames"], + fps=24, + num_inference_steps=contract["num_inference_steps"], + guidance_scale=1.0, + batch_cfg=False, + seed=seed, + ), + output=OutputConfig(output_path=str(out_dir), save_video=False, return_frames=True), + )) + seconds = time.perf_counter() - started + frames = np.stack([np.asarray(frame, dtype=np.uint8) for frame in result.frames]) + audio = result.audio + audio = np.asarray(audio.detach().cpu() if hasattr(audio, "detach") else audio) + return { + "seconds": round(seconds, 3), + "stages": _stage_times(result), + "frames_shape": list(frames.shape), + "frames_sha256": _sha256(frames), + "audio_shape": list(audio.shape), + "audio_dtype": str(audio.dtype), + "audio_sha256": _sha256(audio), + }, frames + + +def run(args) -> None: + from PIL import Image, ImageOps + + example = _example() + case = json.loads(Path(args.case).read_text(encoding="utf-8")) + out_dir = Path(args.out) + out_dir.mkdir(parents=True, exist_ok=True) + reference = Path(args.reference or HERE / case["reference"]) + if not reference.is_file(): + raise SystemExit(f"reference image {reference} not found; pass --reference (README.md, 'The case')") + digest = hashlib.sha256(reference.read_bytes()).hexdigest() + if case.get("reference_sha256") and digest != case["reference_sha256"]: + raise SystemExit(f"{reference} is not the case's reference image (sha256 {digest}, expected " + f"{case['reference_sha256']})") + with Image.open(reference) as opened: + image = ImageOps.exif_transpose(opened).convert("RGB") + + example_args = example.build_parser().parse_args([ + "--model-path", args.model_path, "--prompt", case["prompt"], "--image", + str(reference), "--output", + str(out_dir), "--num-gpus", + str(args.num_gpus), *(["--base-model-path", args.base_model_path] if args.base_model_path else []), + *(["--revision", args.revision] if args.revision else []) + ]) + model_dir, contract = example.resolve_model(example_args) + accel = args.mode == "on" + if accel: + example.apply_lossless_accel_env() + + from fastvideo import VideoGenerator + + generator = VideoGenerator.from_config(example.build_generator_config(model_dir, args.num_gpus, accel)) + runs = [] + try: + grey = Image.new("RGB", image.size, (128, 128, 128)) + _generate(generator, contract, case, grey, 0, WARMUP_PROMPT, out_dir) + for label in ("cold", "repeat"): + record, frames = _generate(generator, contract, case, image, case["seed"], case["prompt"], out_dir) + record["run"] = label + runs.append(record) + print(f"[{args.mode}] {label}: {record['seconds']:.2f} s frames {record['frames_sha256'][:16]} " + f"audio {record['audio_sha256'][:16]} stages {record['stages']}") + if args.save_video: + import imageio + + imageio.mimsave(out_dir / f"{case['name']}_{args.mode}.mp4", list(frames), fps=24) + finally: + generator.shutdown() + + summary = { + "mode": args.mode, + "case": case, + "num_gpus": args.num_gpus, + "model_path": args.model_path, + "revision": args.revision, + "runs": runs, + } + (out_dir / "results.json").write_text(json.dumps(summary, indent=2), encoding="utf-8") + print(f"wrote {out_dir / 'results.json'}") + + +def compare(off_dir: str, on_dir: str) -> int: + off = json.loads((Path(off_dir) / "results.json").read_text(encoding="utf-8")) + on = json.loads((Path(on_dir) / "results.json").read_text(encoding="utf-8")) + identical = True + for base, fast in zip(off["runs"], on["runs"], strict=True): + same = all(base[key] == fast[key] + for key in ("frames_shape", "frames_sha256", "audio_shape", "audio_dtype", "audio_sha256")) + identical &= same + print(f"{base['run']:>6}: off {base['seconds']:6.2f} s on {fast['seconds']:6.2f} s " + f"{base['seconds'] / fast['seconds']:.2f}x {'bit-identical' if same else 'DIFFERENT'}") + print("PASS: video and audio bit-identical" if identical else "FAIL: outputs differ") + return 0 if identical else 1 + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--compare", nargs=2, metavar=("OFF_DIR", "ON_DIR")) + parser.add_argument("--mode", choices=("off", "on")) + parser.add_argument("--model-path", help="FastH3 OmniRef PDD export (local directory or HF repo id)") + parser.add_argument("--revision", default=None) + parser.add_argument("--base-model-path", default=None, help="default: the export's pinned base") + parser.add_argument("--case", default=str(HERE / "case_speaker.json")) + parser.add_argument("--reference", default=None, help="default: the case's reference next to this script") + parser.add_argument("--num-gpus", type=int, default=4) + parser.add_argument("--out", default="runs/omniref") + parser.add_argument("--save-video", action="store_true", help="also write the repeat clip as mp4 (video only)") + args = parser.parse_args() + if args.compare: + sys.exit(compare(*args.compare)) + if not args.mode or not args.model_path: + parser.error("--mode and --model-path are required unless --compare is given") + run(args) + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmarks/fasth3_omniref_b200/case_speaker.json b/scripts/benchmarks/fasth3_omniref_b200/case_speaker.json new file mode 100644 index 000000000..75b973d3f --- /dev/null +++ b/scripts/benchmarks/fasth3_omniref_b200/case_speaker.json @@ -0,0 +1,10 @@ +{ + "name": "speaker", + "prompt": "Picture 1 walks slowly toward the camera across the living room, smiling, and says, \"Every detail was chosen with care.\" Smooth tracking shot, warm afternoon light, light footsteps on the wooden floor.", + "reference": "speaker.png", + "reference_sha256": "d38ba6675b5c780be57818aea0109a8f028d09268c3f29bf00f2739167a6c92c", + "seed": 413410, + "height": 768, + "width": 1344, + "num_frames": 124 +}