From 376e445f0e8d192c2d9a5f28010a6b3a7012667e Mon Sep 17 00:00:00 2001 From: Aryan Kumar <156166681+aryan5v@users.noreply.github.com> Date: Mon, 19 Oct 2026 14:57:36 -0700 Subject: [PATCH] [feat]: FastH3 OmniRef inference on consumer devices (RTX 5090 / PRO 6000 / DGX Spark / 4090) (#1980) Co-authored-by: Aryan Kumar Co-authored-by: Claude Opus 5.5 --- docs-site/navigation.json | 4 + docs/contributing/env_vars.md | 305 +++---- .../inference_schema_parity_inventory.yaml | 1 + .../installation/spark_pair.md | 8 +- .../installation/spark_performance.md | 125 ++- docs/inference/fasth3_omniref.md | 209 +++++ docs/inference/support_matrix.md | 11 + examples/inference/basic/README.md | 5 +- .../basic/basic_fasth3_omniref_pdd.py | 239 +++++- .../basic/basic_fasth3_omniref_rtx4090.yaml | 78 ++ .../basic/basic_fasth3_omniref_rtx5090.yaml | 82 ++ .../basic_fasth3_omniref_rtx_pro6000.yaml | 78 ++ .../basic/basic_fasth3_omniref_spark.yaml | 83 ++ .../basic_fasth3_omniref_spark_pair.yaml | 91 ++ .../basic_fasth3_spark_pair_qad_v2_nvfp4.yaml | 72 ++ .../basic_fasth3_spark_pair_v2_nvfp4.yaml | 7 +- .../basic_fasth3_spark_pruned_nvfp4.yaml | 13 +- .../basic_fasth3_spark_qad_v2_nvfp4.yaml | 69 ++ .../basic/basic_fasth3_spark_v2_nvfp4.yaml | 11 +- fastvideo-kernel/attn_qat_infer/api.py | 15 +- .../triton_kernels/attn_qat_vsa_reference.py | 264 ++++++ .../triton_kernels/attn_qat_vsa_train.py | 515 ++++++++++++ .../tests/test_attn_qat_infer_sparse.py | 15 + .../tests/test_attn_qat_vsa_train.py | 303 +++++++ fastvideo/attention/backends/fp4_vsa_qat.py | 108 +++ .../backends/video_sparse_attn_h3.py | 85 +- fastvideo/configs/models/dits/minimax_h3.py | 2 + fastvideo/envs.py | 66 +- fastvideo/fastvideo_args.py | 12 + fastvideo/layers/fp4linear.py | 14 +- fastvideo/layers/quantization/nvfp4_config.py | 190 ++++- fastvideo/models/dits/minimax_h3.py | 13 +- fastvideo/models/dits/minimax_h3_vsa_fp4.py | 182 +++- fastvideo/models/loader/component_loader.py | 49 +- fastvideo/models/loader/fsdp_load.py | 65 +- fastvideo/models/loader/utils.py | 7 +- fastvideo/models/pinned_offload.py | 17 + .../models/vaes/minimax_h3_int8_convrot.py | 10 + .../models/vaes/minimax_h3_int8_kernels.py | 8 +- .../vaes/minimax_h3_nvfp4_checkpoint.py | 157 ++++ .../models/vaes/minimax_h3_nvfp4_decoder.py | 497 +++++++++++ .../models/vaes/minimax_h3_nvfp4_fused.py | 683 +++++++++++++++ fastvideo/models/vaes/minimax_h3_video.py | 13 +- .../basic/minimax_h3/minimax_h3_pipeline.py | 12 +- .../pipelines/basic/minimax_h3/reference.py | 22 +- .../stages/minimax_h3_conditioning.py | 6 +- .../minimax_h3/stages/minimax_h3_decoding.py | 42 +- .../minimax_h3/stages/minimax_h3_denoising.py | 12 +- .../test_minimax_h3_fine_override.py | 99 +++ .../test_qat_vsa_inference_memory.py | 83 ++ .../attention/test_vsa_fp4_prefix_split.py | 297 +++++++ .../attention/test_vsa_h3_ref2va_regions.py | 50 +- .../test_basic_fasth3_omniref_pdd.py | 192 ++++- fastvideo/tests/loader/test_lora_patch.py | 16 +- .../loader/test_minimax_h3_pdd_loading.py | 15 +- .../tests/loader/test_prequantized_fp8.py | 118 +++ .../ops/quantization/test_nvfp4_row_split.py | 57 ++ .../quantization/test_nvfp4_skip_blocks.py | 124 +++ .../quantization/test_nvfp4_spark_fence.py | 99 +++ .../tests/stages/test_lazy_module_load.py | 9 +- .../test_minimax_h3_consumer_regressions.py | 92 ++ .../test_minimax_h3_sequential_start.py | 46 + .../stages/test_minimax_h3_vae_streaming.py | 61 ++ .../vaes/test_minimax_h3_int8_convrot.py | 18 + .../tests/vaes/test_minimax_h3_int8_shared.py | 22 + .../vaes/test_minimax_h3_nvfp4_decoder.py | 783 ++++++++++++++++++ pyproject.toml | 9 +- .../convert_minimax_h3_modelopt_nvfp4_dit.py | 65 +- .../minimax_h3_nvfp4_decoder/export_deploy.py | 59 ++ .../minimax_h3_nvfp4_decoder/export_packed.py | 69 ++ 70 files changed, 6956 insertions(+), 302 deletions(-) create mode 100644 docs/inference/fasth3_omniref.md create mode 100644 examples/inference/basic/basic_fasth3_omniref_rtx4090.yaml create mode 100644 examples/inference/basic/basic_fasth3_omniref_rtx5090.yaml create mode 100644 examples/inference/basic/basic_fasth3_omniref_rtx_pro6000.yaml create mode 100644 examples/inference/basic/basic_fasth3_omniref_spark.yaml create mode 100644 examples/inference/basic/basic_fasth3_omniref_spark_pair.yaml create mode 100644 examples/inference/basic/basic_fasth3_spark_pair_qad_v2_nvfp4.yaml create mode 100644 examples/inference/basic/basic_fasth3_spark_qad_v2_nvfp4.yaml create mode 100644 fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_reference.py create mode 100644 fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_train.py create mode 100644 fastvideo-kernel/tests/test_attn_qat_vsa_train.py create mode 100644 fastvideo/attention/backends/fp4_vsa_qat.py create mode 100644 fastvideo/models/vaes/minimax_h3_nvfp4_checkpoint.py create mode 100644 fastvideo/models/vaes/minimax_h3_nvfp4_decoder.py create mode 100644 fastvideo/models/vaes/minimax_h3_nvfp4_fused.py create mode 100644 fastvideo/tests/attention/test_minimax_h3_fine_override.py create mode 100644 fastvideo/tests/attention/test_qat_vsa_inference_memory.py create mode 100644 fastvideo/tests/attention/test_vsa_fp4_prefix_split.py create mode 100644 fastvideo/tests/loader/test_prequantized_fp8.py create mode 100644 fastvideo/tests/ops/quantization/test_nvfp4_row_split.py create mode 100644 fastvideo/tests/ops/quantization/test_nvfp4_skip_blocks.py create mode 100644 fastvideo/tests/stages/test_minimax_h3_consumer_regressions.py create mode 100644 fastvideo/tests/vaes/test_minimax_h3_nvfp4_decoder.py create mode 100644 scripts/distill/minimax_h3_nvfp4_decoder/export_deploy.py create mode 100644 scripts/distill/minimax_h3_nvfp4_decoder/export_packed.py diff --git a/docs-site/navigation.json b/docs-site/navigation.json index 11a045ddc..f265cec93 100644 --- a/docs-site/navigation.json +++ b/docs-site/navigation.json @@ -93,6 +93,10 @@ "label": "RTX PRO 6000 / FastH3", "href": "/inference/fasth3_rtx_pro_6000" }, + { + "label": "FastH3 OmniRef on consumer GPUs", + "href": "/inference/fasth3_omniref" + }, { "label": "DGX Spark tuning", "href": "/getting_started/installation/spark_performance" diff --git a/docs/contributing/env_vars.md b/docs/contributing/env_vars.md index ec6c464e4..7939d2aa7 100644 --- a/docs/contributing/env_vars.md +++ b/docs/contributing/env_vars.md @@ -171,153 +171,164 @@ 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_PTX_BLACKWELL` | bool | `0` | attention | VIDEO_SPARSE_ATTN_H3 sends no-grad tile-64 forwards to the data-center Blackwell (sm_100a / sm_103a) kernel. fastvideo-kernel reads the same variable with the same rule. Deprecated names: `FASTVIDEO_VSA_SM100A`. | -| `FASTVIDEO_VSA_TRITON` | bool | `0` | attention | Force the Triton MiniMax-H3 sparse attention kernel. fastvideo-kernel reads the same variable. | -| `FASTVIDEO_NVFP4_FA4` | bool | `0` | attention | FlashAttention-4 quantizes Q and K to NVFP4. An explicit nvfp4_fa4 attention implementation argument takes precedence. | -| `FASTVIDEO_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_ENABLE_FUSED_NORM` | bool | `0` | performance | Opt into Wan's inference-only Triton residual + LayerNorm + modulation fusion. Its reduction order changes numerical outputs; the default keeps eager reference compatibility. | -| `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_LORA_SIDE_BRANCH` | bool | `0` | performance | Apply inference LoRA adapters as a bf16 side branch next to every target linear instead of merging them, so the adapter strength can change at runtime. Packed NVFP4 or FP8 linears without their original dense weight always use the side branch. | -| `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_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_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_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_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`. | +| 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_PTX_BLACKWELL` | bool | `0` | attention | VIDEO_SPARSE_ATTN_H3 sends no-grad tile-64 forwards to the data-center Blackwell (sm_100a / sm_103a) kernel. fastvideo-kernel reads the same variable with the same rule. Deprecated names: `FASTVIDEO_VSA_SM100A`. | +| `FASTVIDEO_VSA_TRITON` | bool | `0` | attention | Force the Triton MiniMax-H3 sparse attention kernel. fastvideo-kernel reads the same variable. | +| `FASTVIDEO_NVFP4_FA4` | bool | `0` | attention | FlashAttention-4 quantizes Q and K to NVFP4. An explicit nvfp4_fa4 attention implementation argument takes precedence. | +| `FASTVIDEO_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_ENABLE_FUSED_NORM` | bool | `0` | performance | Opt into Wan's inference-only Triton residual + LayerNorm + modulation fusion. Its reduction order changes numerical outputs; the default keeps eager reference compatibility. | +| `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 or GB10 sm_121, no-grad, eager inference). | +| `FASTVIDEO_H3_VSA_FP4_PREFIX` | one of , fp8, bf16 | `""` | attention | Prefix split for the FP4 sparse attention kernel: keep the dense prefix query tiles (text, reference and audio rows) off the FP4 kernel and run those rows densely with fp8 or bf16 Q/K and bf16 P.V, while video query tiles stay on the FP4 kernel. The fp8 variant is the validated training-time numerics and recovers most of the audio error that FP4 quantizes away. Use 64/128-token tiles. Empty keeps each row on the FP4 kernel. | +| `FASTVIDEO_H3_VSA_FP4_PREFIX_STEPS` | int | `-1` | attention | Number of leading denoising steps that use the FP4 prefix split (see FASTVIDEO_H3_VSA_FP4_PREFIX). -1 applies it to every step; 0 disables it. Early steps set the global structure, so a small value keeps most of the split's quality benefit for a fraction of its cost. | +| `FASTVIDEO_H3_VSA_FP4_TWO_LEVEL_P` | bool | `0` | attention | Quantize the FP4 attention probabilities with the two-level scheme the QAD student trained with (rowmax-derived s_P1, then per-block s_P2) instead of the kernel's single-level default. | +| `FASTVIDEO_H3_VSA_FP4_K_SMOOTH` | bool | `0` | attention | Subtract the per-head key mean before FP4 key quantization, matching the emulator and QAD training numerics. Exact for the softmax; it only changes quantization error. | +| `FASTVIDEO_H3_GATE_CHUNK_ROWS` | int | `0` | performance | Rows per chunk when the VSA FP4 gate branch expands the pooled compression output back to token rows. 0 keeps the single full-size index_select; a positive value caps the transient at chunk_rows\*heads\*head_dim and is bit-identical. | +| `FASTVIDEO_H3_VSA_FINE_ATTENTION` | one of , emulator-mix8, emulator-mixmx, emulator-qk, emulator-bf16 | `""` | attention | Replace tile-128 fine attention, preserving VSA selection and compression, with the Triton emulator of the sparse FP4 kernel: emulator-mix8 = NVFP4 Q/K video, per-token FP8 Q/K prefix; emulator-mixmx = MXFP8 prefix; emulator-qk = NVFP4 Q/K everywhere; emulator-bf16 = unquantized BF16. emulator-bf16 runs on any CUDA GPU. Nonempty overrides take priority over FASTVIDEO_H3_VSA_FP4. Empty keeps the native route. | +| `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_LORA_SIDE_BRANCH` | bool | `0` | performance | Apply inference LoRA adapters as a bf16 side branch next to every target linear instead of merging them, so the adapter strength can change at runtime. Packed NVFP4 or FP8 linears without their original dense weight always use the side branch. | +| `FASTVIDEO_NVFP4_MM_BACKEND` | str | `auto` | performance | FlashInfer mm_fp4 backend for NVFP4 linears, e.g. auto or cutlass. | +| `FASTVIDEO_NVFP4_MM_ROW_SPLIT` | one of off, table, auto | `off` | performance | Run large cuDNN NVFP4 GEMMs over row chunks where its plan heuristic is slow: off, table (fixed per-shape chunk counts) or auto (time 1/2/4 chunks once per shape; not deterministic across runs). | +| `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_NVFP4_FUSED` | bool | `1` | performance | Run NVFP4 MiniMax-H3 VAE decoder blocks through the bit-exact fused Triton kernels at inference. | +| `FASTVIDEO_H3_VAE_NVFP4_CUDA_GRAPH` | bool | `1` | performance | Replay small fused NVFP4 MiniMax-H3 VAE decoder calls (at most four 1797-token tiles) as per-shape CUDA graphs. | +| `FASTVIDEO_H3_VAE_NVFP4_DECODER` | path | unset | performance | Exported NVFP4 MiniMax-H3 video decoder (.pt, or .safetensors with bf16 masters or pre-packed NVFP4 weights) that replaces the VAE decoder at load; skips the INT8 ConvRot overlay. | +| `FASTVIDEO_H3_VAE_INT8_OVERLAY` | bool | `1` | performance | Apply the INT8 ConvRot decoder overlay when the MiniMax-H3 VAE folder ships one; 0 decodes with the dense weights. | +| `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_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_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_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_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`. | Variables that FastVideo no longer reads; setting one logs a warning: diff --git a/docs/design/inference_schema_parity_inventory.yaml b/docs/design/inference_schema_parity_inventory.yaml index 258cbaad1..2c2eada72 100644 --- a/docs/design/inference_schema_parity_inventory.yaml +++ b/docs/design/inference_schema_parity_inventory.yaml @@ -82,6 +82,7 @@ surfaces: inference_torch_compile: "Regional inference compile opt-in currently carried through PipelineSelection.experimental rather than CompileConfig." vae_parallel_decode: "MiniMax-H3 sequence-parallel VAE decode opt-in; model-specific optimization not yet represented in the typed public schema." h3_sequential_load: "MiniMax-H3 sequential text-encoder then DiT/VAE load; model-specific optimization not yet represented in the typed public schema." + h3_vsa_fp4: "MiniMax-H3 native FP4 attention options carried through generator.pipeline.experimental.h3_vsa_fp4 rather than a dedicated typed component setting." h3_encoder_split: "MiniMax-H3 Ray-only split that runs the Qwen3-VL text encoder on dedicated nodes and the DiT/VAE on the rest; model-specific optimization not yet represented in the typed public schema." h3_encoder_nodes: "Leading node count that h3_encoder_split reserves for the MiniMax-H3 text encoder; model-specific optimization not yet represented in the typed public schema." video_decode_backend: "MiniMax-H3 video decoder selection (full VAE vs TAEH3 preview); model-specific optimization not yet represented in the typed public schema." diff --git a/docs/getting_started/installation/spark_pair.md b/docs/getting_started/installation/spark_pair.md index e6dab20ca..10853eaa9 100644 --- a/docs/getting_started/installation/spark_pair.md +++ b/docs/getting_started/installation/spark_pair.md @@ -149,7 +149,7 @@ fastvideo generate --config examples/inference/basic/basic_fasth3_spark_pair.yam Stop the cluster when you are done: `ray stop` on both nodes. -## Released V2 and Trim NVFP4 stacks +## Released V2, FastH3-QAD-v2 and Trim NVFP4 stacks The public eight-forward V2 and Trim stacks include the NVFP4 encoder and lightweight video VAE. Use the environment above plus the release kernel @@ -161,10 +161,14 @@ export FASTVIDEO_H3_VAE_TILE_BATCH=1 FASTVIDEO_VSA_TRITON=1 export FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0 export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 FASTVIDEO_STAGE_LOGGING=1 fastvideo generate --config examples/inference/basic/basic_fasth3_spark_pair_v2_nvfp4.yaml +# Or basic_fasth3_spark_pair_qad_v2_nvfp4.yaml for FastH3-QAD-v2. # Or basic_fasth3_spark_pair_pruned_nvfp4.yaml for FastH3 Trim. ``` -Both configs use their released Hugging Face model paths, `h3_dit_vsa`, +The configs select separate Hugging Face checkpoints for V2, FastH3-QAD-v2 and Trim. +V2 and FastH3-QAD-v2 enable native FP4 attention in YAML; the extension must be built +for GB10. FastH3-QAD-v2 also selects the FP8 prefix split for text and audio queries. +The examples use `h3_dit_vsa`, 832x480, 124 frames, seed 1234, VSA 0.8 and 64-token tiles. All components stay resident; lazy/sequential loading and compilation are disabled for these compact stacks. The older BF16/Preview memory guidance below applies diff --git a/docs/getting_started/installation/spark_performance.md b/docs/getting_started/installation/spark_performance.md index 2a4740bf6..f2b97746f 100644 --- a/docs/getting_started/installation/spark_performance.md +++ b/docs/getting_started/installation/spark_performance.md @@ -113,7 +113,7 @@ this hardware, for reasons specific to it: | bf16 VAE decode | ~1.14×, lossless; ~5–7% e2e on few-step | ✅ default for Wan | | VSA (video sparse attention) | works out of the box (Triton kernel auto-selects on `sm_121`) | ✅ automatic | | Building FlashAttention | **no speedup** — Torch SDPA already hits an efficient flash kernel on `sm_121`, and FA2 ties it | ❌ not worth building | -| `torch.compile` of the VAE decode | recompile storm (per-frame varying shapes) → ~1.1× | ❌ dead end | +| `torch.compile` of the Wan VAE decode | recompile storm (per-frame varying shapes) → ~1.1× | ❌ dead end (the H3 ViT decoder compiles well, see [FastH3](#fasth3-decode-gemms-and-attention-on-gb10)) | | Linear (fp8 / nvfp4) quantization on long-sequence models (e.g. Cosmos) | ~nothing — see below | ❌ wrong lever here | | FP4 attention (`ATTN_QAT_INFER`) | works on `sm_121` (runtime allowlist landed in #1647; kernel build is #1598); helps, but needs a QAT-trained checkpoint | ⚠️ opt-in — see below | | FP4 linear on short-sequence models (LTX2) | up to −24% denoise at 1080p (#1594) | ⚠️ model/resolution-dependent | @@ -248,10 +248,12 @@ in a local directory. It declares the trained eight-forward ladder and video/audio shifts of 10/3. The recipes use `num_inference_steps: 9` for nine sigma points and eight DiT forwards. -On GB10 with FlashInfer 0.6.18, FastVideo fences activation quantization before -releasing its padded input. Without this completion fence, identical H3 -requests produced different DiT latents and occasionally corrupt video. -The fence applies to `sm_121`; other architectures retain asynchronous execution. +On GB10 (FlashInfer 0.6.18 and 0.7.0), FastVideo completes each NVFP4 +activation quantization on the host before its GEMM. Without this completion +fence, identical H3 requests produced different DiT latents and occasionally +corrupt video. The fence costs no measurable denoise time, because the GPU stays +saturated. It applies to `sm_121`; other architectures retain asynchronous +execution. Run `examples/inference/basic/basic_fasth3_spark_v2_nvfp4.yaml` from the repository root. It uses 832x480, 124 frames and seed 1234, VSA sparsity 0.8 with @@ -259,9 +261,10 @@ repository root. It uses 832x480, 124 frames and seed 1234, VSA sparsity 0.8 wit It does not use frame dropping or spatial upscaling. ```bash +FASTVIDEO_H3_VSA_FP4=1 \ FASTVIDEO_MINIMAX_H3_FUSIONS=all \ -FASTVIDEO_NVFP4_MM_BACKEND=cutlass \ -FASTVIDEO_H3_VAE_TILE_BATCH=1 \ +FASTVIDEO_NVFP4_MM_BACKEND=cudnn FASTVIDEO_NVFP4_MM_ROW_SPLIT=table \ +FASTVIDEO_H3_VAE_TILE_BATCH=12 FASTVIDEO_H3_VAE_INT8_OVERLAY=0 \ FASTVIDEO_VSA_TRITON=1 FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0 \ FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 \ FASTVIDEO_STAGE_LOGGING=1 \ @@ -276,7 +279,7 @@ H3 permits frame counts of `17n+5`; 124 is the nearest legal count above five seconds and 243 is the nearest above ten seconds. For Trim, use `basic_fasth3_spark_pruned_nvfp4.yaml` with the same environment. -Both recipes keep the encoder, DiT and VAEs resident, disable compilation, +Both recipes keep the encoder, DiT and VAEs resident, compile only the VAE decoder, and use `h3_dit_vsa`. On two Sparks, use the corresponding `basic_fasth3_spark_pair_{pruned,v2}_nvfp4.yaml` after following the [pair setup guide](spark_pair.md). The pair uses SP2/TP1 and parallel VAE @@ -309,6 +312,112 @@ The prompt JSON must contain `latency-ceramics-005` and `latency-harbor-005`. Pass `--width 1344 --height 768` for the native 768p protocol. Use the Trim repository and config for its corresponding run. +### FastH3 decode, GEMMs and attention on GB10 + +The recipes decode with the light VAE's dense weights +(`FASTVIDEO_H3_VAE_INT8_OVERLAY=0`), 12 spatial tiles per decoder call and a +compiled decoder (`compile.vae_enabled: true`). The INT8 ConvRot overlay that +ships next to the VAE is slower on GB10 and stops tile batching from helping. +Decoder-only timings on identical 832x480, 124-frame latents (one GB10): + +| Light VAE decoder | 1 tile per call | 12 tiles per call | PSNR vs FP32 decode | +|---|---:|---:|---:| +| INT8 ConvRot overlay, fp16 autocast | 42.2 s | 43.5 s | 63.4 dB | +| Dense, fp16 autocast | 24.4 s | 14.4 s | 77.5 dB | +| Dense, fp16 autocast, compiled | 21.6 s | 9.8 s | 77.7 dB | + +The first one or two calls in a process include decoder compiles; later calls +in the same process reuse them. + +NVFP4 linears use FlashInfer's cuDNN backend (`FASTVIDEO_NVFP4_MM_BACKEND=cudnn`). +The CUTLASS backend used on RTX Blackwell is much slower on GB10 for the +FFN projections, and both backends return identical outputs. FlashInfer +`mm_fp4` with 19,392 activation rows (FastH3 at 832x480, 124 frames): + +| Projection (K→N) | `cutlass` | `cudnn` | `auto` | +|---|---:|---:|---:| +| attention q/k/v, gate (5376→7168) | 5.5 ms | 10.0 ms | 10.0 ms | +| attention out (7168→5376) | 5.5 ms | 4.6 ms | 8.3 ms | +| FFN in (5376→28672) | 52.9 ms | 19.6 ms | 20.0 ms | +| FFN out (14336→5376) | 26.6 ms | 8.9 ms | 9.0 ms | + +`FASTVIDEO_H3_VSA_FP4=1` runs VSA-H3's block-sparse attention on the FP4 +SageAttention3 kernel, as on RTX Blackwell. It needs `fastvideo-kernel` built for +`12.1a` on GB10 or `12.0a` on sm_120. A missing extension raises an error; +select BF16 VSA explicitly by setting `h3_vsa_fp4.enabled: false` in the YAML. On a captured FastH3 Trim +attention input (19,392 tokens, 56 heads, 64-token tiles) it runs in 29 ms +against 59 ms for Triton VSA, with a relative L2 error of 0.13 versus an FP32 +reference (0.002 for Triton BF16). + +The FastH3-QAD-v2 recipes, `basic_fasth3_spark_qad_v2_nvfp4.yaml` and +`basic_fasth3_spark_pair_qad_v2_nvfp4.yaml`, select +`FastVideo/FastH3-QAD-v2` explicitly. The older +`v2` files continue to select V2. Both recipes enable FP4 in YAML through +`pipeline.experimental.h3_vsa_fp4.enabled`; shell flags are not required for +that selection. FastH3-QAD-v2 keeps its 64-token tiles and sparsity 0.8, and protects +text and audio query rows with FP8 Q/K and BF16 dense attention. These are +inference settings. The published FastH3-QAD-v2 checkpoint used linears-only QAD with +BF16 attention, so this does not reproduce its training attention numerics. + +The earlier unprotected FP4 path is faster (on a 5090, 480p5 8L +measured about 12.6 s versus 19 s on bf16 VSA) but changes the output. On V2 at +832x480 with four prompts, denoising was 24-33% faster (60-68 s to 43-51 s) and +the clips scored LPIPS 0.39-0.55 against the BF16-attention clips. Validate clips +before a quality comparison. On a 32 GB RTX 5090 the measured recipe also sets +`FASTVIDEO_H3_ENCODER_LAYERWISE=1` and `FASTVIDEO_H3_PARK_MODULES=vae,audio_vae`. + +Measured on one Spark (832x480, 124 frames, seed 1234, the ceramics release +prompt), timing warm `generate()` calls in one process through finished MP4 +output: + +| Model and configuration | End to end | Denoise | Video decode | +|---|---:|---:|---:| +| Trim, previous recipe (CUTLASS, INT8 overlay, 1 tile per call) | 121-122 s | 77-78 s | 43 s | +| Trim, current recipe | 68 s | 58 s | 9.5 s | +| Trim, current recipe + `FASTVIDEO_H3_VSA_FP4=1` | 54 s | 43 s | 9.6 s | +| V2, previous recipe | 139-141 s | 95-96 s | 42-43 s | +| V2, current recipe | 79 s | 68 s | 9.6 s | + +Within each configuration, repeated calls produced bit-identical frames +except the previous V2 recipe, whose repeated calls diverged. + +Row splitting is off globally. The single-Spark cuDNN recipe opts into +`FASTVIDEO_NVFP4_MM_ROW_SPLIT=table` for the shapes measured on GB10. Leave it +off on other GPUs unless a matched benchmark validates the split there. + +### Optional NVFP4 video decoder + +`FASTVIDEO_H3_VAE_NVFP4_DECODER=/path/to/decoder.pt` swaps in an NVFP4 +decoder distilled with quantization-aware training and exported by +`scripts/distill/minimax_h3_nvfp4_decoder/export_deploy.py`. At load the +INT8 overlay is skipped, the VAE is depth-cut to the decoder's block count +(the 26-layer light VAE hosts an 8-block decoder; a 36-block decoder needs +the full `MiniMaxAI/MiniMax-H3` `vae/` in the stack's `vae/` folder), and +only the packed FP4 weights stay resident (0.31 GiB for 8 blocks instead of +2.3 GiB). The decode stage runs bf16 autocast for these decoders, the +precision they were trained in. The decoder GEMMs use the cuDNN backend on +GB10; outputs match CUTLASS bit for bit. + +The variable also accepts `.safetensors` files written by +`fastvideo.models.vaes.minimax_h3_nvfp4_checkpoint`: either the same bf16 +master weights, or pre-packed NVFP4 weights from +`scripts/distill/minimax_h3_nvfp4_decoder/export_packed.py`. A packed file is +read straight onto the GPU and never materializes the full-precision weights +of the quantized linears; it decodes bit-identically to packing the bf16 file +at load. + +V2 at 832x480, 124 frames, seed 1234, one Spark, warm calls: + +| Video decoder | End to end | Denoise | Video decode | Peak allocated | +|---|---:|---:|---:|---:| +| Light VAE, INT8 overlay, 1 tile per call, eager | 112.4 s | 68.4 s | 42.7 s | 59.4 GiB | +| Light VAE, dense, 12 tiles per call, compiled (recipe) | 78.9 s | 68.1 s | 9.6 s | 64.3 GiB | +| 8-block NVFP4 decoder | 70.9 s | 68.2 s | 1.6 s | 58.0 GiB | + +The 8-block decoder decodes the same with or without `compile.vae_enabled` +and does not need tile batching. Its output is lossy against the light VAE +(31 dB PSNR on this clip); review clips before using it. + ### Released model measurements The released Trim stack at revision `cae9ceb6feefe77d34a56640782cda3909363f19` diff --git a/docs/inference/fasth3_omniref.md b/docs/inference/fasth3_omniref.md new file mode 100644 index 000000000..c2a669383 --- /dev/null +++ b/docs/inference/fasth3_omniref.md @@ -0,0 +1,209 @@ +# FastH3 OmniRef inference + +FastH3 OmniRef (Ref2VA) is a reference-conditioned PDD student: a prompt plus +ordered image, video and audio references become one video and audio clip in +eight fused transformer forwards. Its extra conditioning prefix - the text, the +Qwen3-VL vision tokens of every reference and any reference audio - sits on top +of the target video and audio rows. See +[distilled FastH3 checkpoints](fasth3-distilled.md#ref2va-pdd-students) for the +sampling contract and checkpoint layout. + +This page covers running the NVFP4 stack on one consumer Blackwell device +(RTX 5090, RTX PRO 6000, DGX Spark) and the closest equivalent on an RTX 4090, +which has no FP4 hardware. + +## Composing the model directory + +The NVFP4 checkpoint is +[FastH3 OmniRef QAD](https://huggingface.co/FastVideo/FastH3-OmniRef-NVFP4-QAD). +`basic_fasth3_omniref_pdd.py` accepts a complete checkpoint or a transformer-only +export. For an export containing `transformer_ref/`, the two scheduler configs +and `fastvideo_inference.json`, it composes the missing components from base +MiniMax-H3 into one directory of symlinks and runs the stack: + +```bash +python examples/inference/basic/basic_fasth3_omniref_pdd.py \ + --model-path \ + --image reference.png --audio voice.wav \ + --prompt "a person on a sofa, lit by a warm lamp, speaking to camera" \ + --profile spark +``` + +`--profile` selects the device's offload policy, geometry and environment. +Without it the script keeps its data-center defaults. Each profile also has a +recipe under `examples/inference/basic/basic_fasth3_omniref_*.yaml`; point its +`model_path` at the composed directory and run it with +`fastvideo generate --config `. + +The Blackwell profiles need a packed NVFP4 OmniRef export. The RTX 4090 needs a +dense or FP8 export: a packed NVFP4 export omits the quantized linears from its +dense shards and cannot load as FP8. + +Pass `--adaln-table /path/to/adaln_table.pt` to drop the AdaLN projections (the +24/32 GB cards need it to keep the DiT resident), and +`--decoder /path/to/decoder-nvfp4.safetensors` to swap in the pre-packed NVFP4 +decoder. Write an AdaLN table once with `FASTVIDEO_H3_ADALN_DUMP`; build a +decoder file with `scripts/distill/minimax_h3_nvfp4_decoder/export_deploy.py`. + +## Device profiles + +Every Blackwell profile runs the same numerics, so a clip is the same clip on +every device. Only the memory policy and the geometry differ. + +| Profile | Device | Geometry | Offload | +|---|---|---|---| +| `spark` | DGX Spark, GB10, 128 GB unified | 480p | none: encoder, DiT and decoder resident | +| `rtx-pro-6000` | RTX PRO 6000, sm_120, 96 GB | 768p | none: encoder, DiT and decoder resident | +| `rtx-5090` | RTX 5090, sm_120, 32 GB | 480p | non-resident encoder, sequential denoise-module load, AdaLN table | +| `rtx-4090` | RTX 4090, sm_89, 24 GB | 480p | layerwise DiT stream (16 blocks resident), non-resident encoder, FFN chunk 16384, expandable segments | + +The 96 GB and 128 GB profiles keep the DiT whole. The 32 GB 5090 also keeps it +resident but offloads the text encoder; the 24 GB 4090 streams the FP8 DiT +layerwise with 16 blocks resident. The 24 and 32 GB cards replace the AdaLN +projections (a third of the DiT's bytes on disk) with a precomputed table. + +The `spark` and `rtx-5090` profiles enable the bit-identical lossless set by +default: exact fused kernels, heads-first VSA tiles, parallel VAE paths, and the +content-keyed reference memo. On a repeated reference this removes conditioning +and latent preparation from the warm path, worth 5-16% of warm wall time on +Spark and 3-5 s on the 5090; a first-time reference sees no regression. +`--lossless-accel` enables the same set on the other profiles. + +The Blackwell stack is: + +- **Linears:** NVFP4 from the packed export. Blocks the calibration kept in + bf16 are absent from the packed file and load from the dense shards; the + export records its own skip list in the safetensors metadata. +- **Attention:** VSA sparsity 0.9, 128-token tiles, reference keep 0.1. Video + query tiles run on the FP4 sparse kernel. Dense prefix query tiles - the + text, reference and audio rows - are split off and run at 8-bit Q/K with bf16 + P.V, keeping the conditioning out of the FP4 rounding. The split applies on + the single-rank and sequence-parallel paths alike. +- **Text encoder:** the NVFP4 Qwen3-VL encoder, including reference vision + tokens. +- **Decoder:** the pre-packed NVFP4 decoder (`decoder-nvfp4.safetensors`), read + straight onto the GPU without materializing full-precision weights. + +The RTX 4090 profile uses FP8 linears, the unquantized Triton emulator for the +checkpoint's fixed 128-token VSA tiles and the VAE folder's dense decoder. The +NVFP4 decoder's GEMMs need FP4 hardware, so `--decoder` does not apply there. +Layerwise offload is mandatory: the 47 GB FP8 DiT cannot be resident on 24 GB. +The profile keeps 16 blocks resident, chunks the FFN at 16384 tokens and +disables pinned swap. Repeated requests in one process exceed the 24 GB GPU and +the 57.7 GiB host cgroup once the 16 GB text encoder reloads; run one process +per clip on this card. + +The Blackwell YAML examples also select FP4 attention and the prefix precision +through `pipeline.experimental.h3_vsa_fp4`. Explicit YAML values take precedence +over the corresponding environment flags. A missing native extension raises an +error. To choose BF16 VSA for OmniRef's 128-token tiles, set +`h3_vsa_fp4.enabled: false` and +`FASTVIDEO_H3_VSA_FINE_ATTENTION=emulator-bf16`. + +The Blackwell profiles select the FP8 prefix split to protect the conditioning +queries from FP4 rounding. The split runs prefix rows densely in bf16 after +8-bit Q/K rounding, which costs about 24% of +denoise time on one Spark. The split applies at `sp_size` 1 and above; the +sequence-parallel path keeps the same prefix and video row split. + +`--attention-prefix bf16` is a higher-precision diagnostic mode, not a faster +one. Keeping every row on the FP4 kernel (`FASTVIDEO_H3_VSA_FP4_PREFIX=` empty) +is a low-level escape hatch, not a release mode: on a matched 480p/5 s clip it +sits 4.3 dB farther from the bf16 attention reference on every frame (20.96 vs +25.26 dB PSNR; SSIM 0.785 vs 0.852; structural correlation 0.864 vs 0.945). +Use the 36L decoder for the best-quality clip and the 8L decoder on memory-bound +devices; both run the FP8 split. + +## Environment + +The profiles set these; set them yourself when driving the YAML recipes. + +| Variable | Value | Why | +|---|---|---| +| `FASTVIDEO_H3_ADALN_TABLE` | path, every device | drops the 26 GB of `adaln_proj`; the 24/32 GB profiles need it | +| `FASTVIDEO_ATTENTION_BACKEND` | `VIDEO_SPARSE_ATTN_H3` | the trained backend; the contract fixes the tile size | +| `FASTVIDEO_H3_VSA_FP4` | `1` on Blackwell, `0` on the 4090 | FP4 sparse attention; sm_89 has no FP4 | +| `FASTVIDEO_H3_VSA_FP4_PREFIX` | `fp8` | the prefix split; `bf16` is the higher-precision variant | +| `FASTVIDEO_H3_VSA_FP4_TWO_LEVEL_P` | `1` on QAD profiles | the P quantizer QAD trained with | +| `FASTVIDEO_H3_VSA_FP4_K_SMOOTH` | `1` on QAD profiles | subtract the per-head key mean before FP4 quantization, as QAD trained | +| `FASTVIDEO_H3_GATE_CHUNK_ROWS` | `8192` on QAD profiles | bit-identical gate chunking; caps the 768p10 transient | +| `FASTVIDEO_H3_VSA_FINE_ATTENTION` | `emulator-bf16` on the 4090 | bf16 VSA on sm_89 | +| `FASTVIDEO_H3_VAE_NVFP4_DECODER` | path to `decoder-nvfp4.safetensors` | the pre-packed NVFP4 decoder | +| `FASTVIDEO_MINIMAX_H3_FUSIONS` | `all` | bit-exact fused kernels | +| `FASTVIDEO_VSA_TRITON` | `1` | Triton VSA on the consumer parts | +| `FASTVIDEO_VSA_PTX_BLACKWELL`, `FASTVIDEO_FA4` | `0` | no data-center kernel and no FlashAttention 4 | +| `FASTVIDEO_NVFP4_MM_ROW_SPLIT` | `table` on Spark, globally `off` | opt into the cuDNN row chunks measured on GB10 | +| `FASTVIDEO_NVFP4_MM_BACKEND` | `cudnn` on Spark | faster than CUTLASS on the H3 FFN shapes on GB10, bit-identical | +| `FASTVIDEO_H3_VAE_TILE_BATCH` | `12` Spark, `4` 5090/4090 | decoder tiles per call | +| `FASTVIDEO_H3_VAE_INT8_OVERLAY` | `0` on Spark | the INT8 ConvRot overlay is slower on GB10 | +| `FASTVIDEO_LAYERWISE_OFFLOAD_BUFFERS`, `FASTVIDEO_LAYERWISE_RESIDENT_BLOCKS` | `1`, `16` on the 4090 | stream the packed FP8 weights; keep 16 blocks resident | +| `FASTVIDEO_H3_FFN_CHUNK_TOKENS`, `FASTVIDEO_H3_PINNED_SWAP` | `16384`, `0` on the 4090 | bound the FFN activation and avoid pinned host duplicates | +| `FASTVIDEO_ULYSSES_A2A`, `FASTVIDEO_MINIMAX_H3_EXACT_KERNELS`, `FASTVIDEO_H3_VSA_HEADS_FIRST_TILE`, `FASTVIDEO_H3_VAE_TILE_PARALLEL`, `FASTVIDEO_H3_REF2VA_MEMO_ENTRIES` | lossless set, on for Spark/5090 | bit-identical speed-ups; the memo helps repeated references | +| `PYTORCH_CUDA_ALLOC_CONF` | `expandable_segments:True` on the 4090 | reduces allocator fragmentation | + +`fastvideo-kernel` must be built for the device: `12.0a` on sm_120 and `12.1a` +on GB10 (CUDA 13). Building on the machine resolves the right one. + +### Fallbacks + +The FP4 route needs that kernel build. These are explicit alternatives; a +missing native kernel raises rather than silently selecting another mode. + +1. **FP4 kernel + prefix split** (the profiles above). +2. **`FASTVIDEO_H3_VSA_FINE_ATTENTION=emulator-mix8`**: Triton FP4 Q/K for video + query tiles, per-token FP8 Q/K for prefix tiles, bf16 P.V and K smoothing. + Needs no consumer-Blackwell extension but is slower, and its numerics differ + from the native kernel (which also quantizes P.V). +3. **bf16 VSA**: set `FASTVIDEO_H3_VSA_FP4=0` and + `FASTVIDEO_H3_VSA_FINE_ATTENTION=emulator-bf16`. This runs unquantized + attention on the same tiles with Triton on consumer CUDA GPUs, including + sm_89. + +A nonempty fine-attention override takes priority over +`FASTVIDEO_H3_VSA_FP4=1`. These overrides use the eager path; disable inference +compilation when using them. + +The prefix rows are rounded to FP8 and attended in bf16; this protects their +precision rather than accelerating them. For the FP4 video-query kernel, the +Blackwell profiles enable two-level probability quantization +(`FASTVIDEO_H3_VSA_FP4_TWO_LEVEL_P=1`) and key smoothing +(`FASTVIDEO_H3_VSA_FP4_K_SMOOTH=1`), matching those settings in the OmniRef QAD +v2 training recipe. The kernel defaults to single-level probability +quantization without key smoothing when these switches are unset. Matching +these two settings does not establish exact parity for the entire attention +path, including the protected prefix and sequence-parallel exchanges. + +## Memory + +What each component costs at 480p with one image reference. The AdaLN column is +why the 24 and 32 GB cards pass `FASTVIDEO_H3_ADALN_TABLE`. + +| Component | NVFP4, bf16 AdaLN | with the AdaLN table | +|---|---|---| +| DiT weights on disk | 44.5 GB | 18.5 GB | +| Qwen3-VL text encoder | 24.7 GiB | same | +| 36-block decoder, frozen | ~1.4 GiB | same | +| 8-block `fast-8L` decoder | 0.31 GiB | same | + +Budget about 69 GiB for the DiT and encoder when both stay resident, before +the decoder and activations. The encoder and decoder figures are measured; the +rest are estimates. Both the RTX 5090 and RTX 4090 offload the encoder and load +denoise modules on demand; the 4090 also uses the larger FP8 DiT weights. + +## Multiple Sparks + +Two Sparks run the same stack with `--num-gpus 2`, Ray and sequence +parallelism over a direct link (`basic_fasth3_omniref_spark_pair.yaml`). +Sequence parallelism cannot use three ranks: MiniMax-H3 has 56 attention +heads, so `sp_size` must divide 56. For three Sparks, enable the encoder split +(`FASTVIDEO_H3_ENCODER_SPLIT=1`, `FASTVIDEO_H3_ENCODER_NODES=1`): one node +runs the Qwen3-VL conditioner and the other two run the DiT and VAEs with +`sp_size=2`. The prompt-embedding broadcast still spans all three ranks, so +all three nodes need a common fast link; the split does not fit a triangle of +separate point-to-point subnets without routes between them. + +## Apple Silicon + +`fastvideo/mlx_runtime` has an MLX MiniMax-H3 runtime, but its pipeline is +T2V/FL2VA-class. Running OmniRef on a Mac needs an MLX Ref2VA path and an MLX +build of the NVFP4 decoder; it is not implemented. diff --git a/docs/inference/support_matrix.md b/docs/inference/support_matrix.md index ec3a64409..215e07311 100644 --- a/docs/inference/support_matrix.md +++ b/docs/inference/support_matrix.md @@ -242,6 +242,17 @@ resolve default pipeline and sampling configuration for it. `FastVideo/GEN3C-Cosmos-7B-Diffusers`) or convert locally with `scripts/checkpoint_conversion/convert_gen3c_to_fastvideo.py`. +## FastH3 OmniRef on consumer devices + +OmniRef (Ref2VA) runs the same NVFP4 stack on one consumer Blackwell device - +RTX 5090, RTX PRO 6000 and DGX Spark - with a packed NVFP4 export, FP4 sparse +attention with the dense prefix rows split off, and the pre-packed NVFP4 +decoder. An RTX 4090 runs the closest sm_89 equivalent (FP8 linears, bf16 VSA). +Per-device profiles, recipes and the environment are in +[FastH3 OmniRef inference](fasth3_omniref.md). + +There is no MLX Ref2VA path yet: the Apple Silicon rows above are T2V/FL2VA. + ## Hardware and OS Per the installation guides: diff --git a/examples/inference/basic/README.md b/examples/inference/basic/README.md index a3e063349..7857eb643 100644 --- a/examples/inference/basic/README.md +++ b/examples/inference/basic/README.md @@ -225,7 +225,10 @@ python examples/inference/basic/basic_fasth3_omniref_pdd.py \ ``` See [FastH3 distilled checkpoint schedules](https://github.com/hao-ai-lab/FastVideo/blob/main/docs/inference/fasth3-distilled.md#ref2va-pdd-students) -for the contract and the reference-video sparsity policy. +for the contract and the reference-video sparsity policy. For the consumer NVFP4 +stack (RTX 5090, RTX PRO 6000, DGX Spark; FP8 on the RTX 4090) add +`--profile `, and see +[FastH3 OmniRef inference](https://github.com/hao-ai-lab/FastVideo/blob/main/docs/inference/fasth3_omniref.md). ## Basic Walkthrough diff --git a/examples/inference/basic/basic_fasth3_omniref_pdd.py b/examples/inference/basic/basic_fasth3_omniref_pdd.py index a93d57f81..a3f7aed2c 100644 --- a/examples/inference/basic/basic_fasth3_omniref_pdd.py +++ b/examples/inference/basic/basic_fasth3_omniref_pdd.py @@ -22,6 +22,13 @@ 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``. +``--profile {spark,rtx-pro-6000,rtx-5090,rtx-4090}`` runs the consumer stack on +one of those devices: a packed NVFP4 export, FP4 sparse attention with the dense +prefix rows split off at 8-bit Q/K, and the pre-packed NVFP4 decoder, with the +device's own offload policy and geometry. The 4090, having no FP4 hardware, runs +the closest equivalent from a dense/FP8 export (FP8 linears, bf16 VSA, the dense decoder). See +docs/inference/fasth3_omniref.md. + ``--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 @@ -35,11 +42,13 @@ import argparse import json import os from collections.abc import Sequence +from dataclasses import dataclass from pathlib import Path from typing import Any from fastvideo import VideoGenerator from fastvideo.api import ( + CompileConfig, ComponentConfig, EngineConfig, GenerationRequest, @@ -49,12 +58,14 @@ from fastvideo.api import ( OutputConfig, ParallelismConfig, PipelineSelection, + QuantizationConfig, SamplingConfig, ) from fastvideo.configs.pipelines.minimax_h3 import parse_base_model_revision from fastvideo.pipelines.basic.minimax_h3 import MiniMaxH3Reference CONTRACT = "fastvideo_inference.json" +DEFAULT_HEIGHT, DEFAULT_WIDTH = 480, 832 # Components the distilled export replaces, and the ones it shares with base MiniMax-H3. EXPORT_COMPONENTS = ("transformer_ref", "scheduler", "audio_scheduler") BASE_COMPONENTS = ("text_encoder", "tokenizer", "processor", "vae", "audio_vae") @@ -114,14 +125,35 @@ def build_parser() -> argparse.ArgumentParser: help=f"an ordered {flag} reference (repeatable)") parser.add_argument("--prompt", required=True) parser.add_argument("--output", default="outputs/fasth3_omniref_pdd") - parser.add_argument("--height", type=int, default=480) - parser.add_argument("--width", type=int, default=832) + parser.add_argument("--height", type=int, default=None) + parser.add_argument("--width", type=int, default=None) 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)") + parser.add_argument("--profile", + choices=sorted(DEVICE_PROFILES), + default=None, + help="run the full NVFP4 consumer stack for one device, with its offload policy and " + "environment (see docs/inference/fasth3_omniref.md); the checkpoint must be " + "the packed NVFP4 export on Blackwell, or a dense/FP8 export on the 4090") + parser.add_argument("--attention-prefix", + choices=("default", "fp8", "bf16"), + default="default", + help="numerics for the dense attention prefix rows (text, reference and audio). " + "'fp8' is the validated-quality mode and the device-profile default; 'bf16' is " + "higher precision but slower; 'default' keeps the profile's choice") + parser.add_argument("--decoder", + default=None, + help="pre-packed NVFP4 VAE decoder (decoder-nvfp4.safetensors) used by --profile " + "instead of the VAE folder's own decoder") + parser.add_argument("--adaln-table", + default=None, + help="precomputed AdaLN modulation table, so the 26 GB of adaln_proj never loads " + "(write one once with FASTVIDEO_H3_ADALN_DUMP). Strongly recommended: the 24/32 GB " + "profiles need it to keep the DiT resident") return parser @@ -220,6 +252,9 @@ def resolve_model(args: argparse.Namespace) -> tuple[Path, dict[str, Any]]: export_dir = _snapshot(args.model_path, args.revision, [CONTRACT, *MANIFESTS, *(f"{name}/**" for name in EXPORT_COMPONENTS)]) contract = load_contract(export_dir) + if args.profile == "rtx-4090" and (export_dir / "transformer_ref" / "nvfp4_weights.safetensors").is_file(): + raise ValueError("rtx-4090 requires a dense or FP8 OmniRef export. The packed NVFP4 export omits " + "the weights needed for FP8 conversion; pass the dense/FP8 export with --model-path.") base_source, base_revision = base_model_source(args, contract) base_dir = _snapshot(base_source, base_revision, [*MANIFESTS, *(f"{name}/**" for name in BASE_COMPONENTS)]) composed_dir = Path(args.composed_dir) if args.composed_dir else Path(args.output) / "fasth3_omniref_model" @@ -232,8 +267,174 @@ def apply_lossless_accel_env() -> None: os.environ.setdefault(name, value) -def build_generator_config(model_dir: Path, num_gpus: int, lossless_accel: bool = False) -> GeneratorConfig: - if lossless_accel: +_CONSUMER_ENV = { + # The checkpoint's trained attention backend; the contract fixes the tile size. + "FASTVIDEO_ATTENTION_BACKEND": "VIDEO_SPARSE_ATTN_H3", + # FP4 sparse attention on the consumer-Blackwell kernel, with the prefix split. + "FASTVIDEO_H3_VSA_FP4": "1", + "FASTVIDEO_H3_VSA_FP4_PREFIX": "fp8", + # QAD trained the video rows with two-level P and K smoothing; match it. + "FASTVIDEO_H3_VSA_FP4_TWO_LEVEL_P": "1", + "FASTVIDEO_H3_VSA_FP4_K_SMOOTH": "1", + # Bit-identical gate chunking; caps the 1.2 GiB transient at 768p10. + "FASTVIDEO_H3_GATE_CHUNK_ROWS": "8192", + # Bit-exact kernels, and none of the data-center-only routes. + "FASTVIDEO_MINIMAX_H3_FUSIONS": "all", + "FASTVIDEO_VSA_TRITON": "1", + "FASTVIDEO_VSA_PTX_BLACKWELL": "0", + "FASTVIDEO_FA4": "0", +} + + +@dataclass(frozen=True) +class DeviceProfile: + """Everything that differs between consumer devices for one OmniRef clip.""" + + label: str + width: int + height: int + offload: OffloadConfig + env: dict[str, str] + # Defer the denoise modules until the sampler needs them, for the cards that + # cannot hold the text encoder and the DiT at once (the 24/32 GB recipes use it). + sequential_load: bool = False + # Linear quantization. sm_89 has no FP4 hardware, so the 4090 profile is the + # closest equivalent rather than the FP4 stack, in the DiT and - because the + # NVFP4 decoder's GEMMs need it too - in the VAE. + transformer_quant: str = "NVFP4" + layer_profile: str | None = "h3_dit_vsa" + nvfp4_decoder: bool = True + + +DEVICE_PROFILES: dict[str, DeviceProfile] = { + "spark": DeviceProfile( + label="NVIDIA DGX Spark (GB10, 128 GB unified)", + width=832, + height=480, + # 128 GB unified holds the encoder, the DiT and the decoder at once: + # the same resident policy as the released Spark recipe. + offload=OffloadConfig(dit=False, + dit_layerwise=False, + text_encoder=False, + vae=False, + pin_cpu_memory=False, + lazy_module_load=False), + env={ + **_CONSUMER_ENV, + # Byte-identical at both shapes and both decoders; the reference + # memo saves 5-16% of warm wall time on repeated references. + **LOSSLESS_ACCEL_ENV, + # cuDNN mm_fp4 is 2.7-3x faster than CUTLASS on the H3 FFN shapes on GB10. + "FASTVIDEO_NVFP4_MM_BACKEND": "cudnn", + "FASTVIDEO_NVFP4_MM_ROW_SPLIT": "table", + "FASTVIDEO_H3_VAE_TILE_BATCH": "12", + "FASTVIDEO_H3_VAE_INT8_OVERLAY": "0", + }, + ), + "rtx-pro-6000": DeviceProfile( + label="NVIDIA RTX PRO 6000 (96 GB)", + width=1344, + height=768, + # 96 GB fits the encoder, the DiT and the decoder at once. + offload=OffloadConfig(dit=False, + dit_layerwise=False, + text_encoder=False, + vae=False, + pin_cpu_memory=True, + lazy_module_load=False), + env={**_CONSUMER_ENV, "FASTVIDEO_H3_VAE_TILE_BATCH": "8"}, + ), + "rtx-5090": DeviceProfile( + label="NVIDIA RTX 5090 (32 GB)", + width=832, + height=480, + # The encoder is not resident, and the denoise modules load only when the + # sampler needs them. The DiT itself stays whole: no layerwise streaming. + offload=OffloadConfig(dit=False, dit_layerwise=False, text_encoder=True, vae=False, pin_cpu_memory=True), + sequential_load=True, + env={ + **_CONSUMER_ENV, + # Bit-identical at both shapes and both decoders on the 5090; the + # reference memo and VAE paths save 3-5 s on repeated requests. + **LOSSLESS_ACCEL_ENV, + "FASTVIDEO_H3_VAE_TILE_BATCH": "4", + }, + ), + "rtx-4090": DeviceProfile( + label="NVIDIA RTX 4090 (24 GB, sm_89, no FP4)", + width=832, + height=480, + # The 47 GB FP8 DiT cannot be resident on 24 GB: stream it, keeping the + # validated 16 blocks resident. The encoder stays non-resident. + offload=OffloadConfig(dit=False, dit_layerwise=True, text_encoder=True, vae=False, pin_cpu_memory=True), + sequential_load=True, + env={ + # No FP4 hardware: FP8 linears and bf16 VSA attention instead, and the + # dense light VAE because the NVFP4 decoder's GEMMs need FP4 too. + "FASTVIDEO_ATTENTION_BACKEND": "VIDEO_SPARSE_ATTN_H3", + "FASTVIDEO_H3_VSA_FP4": "0", + "FASTVIDEO_H3_VSA_FINE_ATTENTION": "emulator-bf16", + "FASTVIDEO_MINIMAX_H3_FUSIONS": "all", + "FASTVIDEO_VSA_TRITON": "1", + "FASTVIDEO_VSA_PTX_BLACKWELL": "0", + "FASTVIDEO_FA4": "0", + # Validated 24 GB layerwise configuration: stream the packed weights + # too, keep 16 blocks resident, and chunk the FFN. + "FASTVIDEO_LAYERWISE_OFFLOAD_BUFFERS": "1", + "FASTVIDEO_LAYERWISE_RESIDENT_BLOCKS": "16", + "FASTVIDEO_H3_FFN_CHUNK_TOKENS": "16384", + "FASTVIDEO_H3_PINNED_SWAP": "0", + # A 24 GB card fragments easily once the DiT is streamed. + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + "FASTVIDEO_H3_VAE_TILE_BATCH": "4", + }, + transformer_quant="FP8", + layer_profile=None, + nvfp4_decoder=False, + ), +} + + +def apply_attention_prefix(choice: str) -> None: + """Apply ``--attention-prefix``; ``default`` keeps the profile's choice.""" + if choice != "default": + os.environ["FASTVIDEO_H3_VSA_FP4_PREFIX"] = choice + + +def apply_profile(profile: DeviceProfile, decoder: str | None, adaln_table: str | None = None) -> None: + """Set a profile's environment, keeping any value already in the environment.""" + for name, value in profile.env.items(): + os.environ.setdefault(name, value) + if decoder and profile.nvfp4_decoder: + os.environ.setdefault("FASTVIDEO_H3_VAE_NVFP4_DECODER", decoder) + elif decoder: + print(f"--decoder ignored on {profile.label}: the NVFP4 decoder needs FP4 hardware; " + "this profile decodes with the VAE folder's own (dense) decoder") + if adaln_table: + # The AdaLN projections are a third of the DiT's bytes on disk; the + # precomputed table drops them entirely. Every profile takes it: on the + # 24/32 GB cards it is what makes residency possible, and elsewhere it + # buys headroom for the reference tokens. + os.environ.setdefault("FASTVIDEO_H3_ADALN_TABLE", adaln_table) + + +def build_generator_config(model_dir: Path, + num_gpus: int, + lossless_accel: bool = False, + profile: DeviceProfile | None = None) -> GeneratorConfig: + quantization = None + if profile is not None: + # `h3_dit_vsa` matches the packed export the consumer weights ship; the + # export's metadata names its bf16 blocks, so no skip list is needed. + quantization = QuantizationConfig(transformer_quant=profile.transformer_quant, + layer_profile=profile.layer_profile) + offload = profile.offload + # Explicit: on unified memory the auto rule would defer the DiT and load + # 50+ GB at the first denoise call even when everything fits. + experimental = {"h3_sequential_load": profile.sequential_load} + if num_gpus > 1: + experimental.update({"vae_parallel_decode": True, "vae_parallel_encode": True}) + elif 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 @@ -247,17 +448,27 @@ def build_generator_config(model_dir: Path, num_gpus: int, lossless_accel: bool model_path=str(model_dir), engine=EngineConfig( num_gpus=num_gpus, - # Shard the DiT across the GPUs rather than holding a full copy on each. - use_fsdp_inference=num_gpus > 1, + # Multi-GPU device profiles run on Ray (the Spark pair spans two hosts). + execution_backend="ray" if num_gpus > 1 and profile is not None else "mp", + # A packed NVFP4 export is replicated across GPUs; the tokens are split + # by sequence parallelism. Data-center checkpoints keep mp and FSDP. + use_fsdp_inference=num_gpus > 1 and profile is None, parallelism=ParallelismConfig(tp_size=1, sp_size=num_gpus), offload=offload, + quantization=quantization, + # The NVFP4 decoder installs its own fused CUDA graphs; compiling the VAE + # on top trips an inductor cudagraph assertion at 768p with the 36-block student. + compile=CompileConfig(), ), 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, + vae_tiling=profile is not None, + experimental={ + **experimental, "video_decode_backend": "h3-vae" + } if profile is not None else experimental, ), ) @@ -273,9 +484,19 @@ 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] - if args.lossless_accel: + profile = DEVICE_PROFILES[args.profile] if args.profile else None + if profile is not None: + apply_profile(profile, args.decoder, args.adaln_table) + print(f"Profile: {profile.label}") + elif args.lossless_accel: apply_lossless_accel_env() - generator = VideoGenerator.from_config(build_generator_config(model_dir, args.num_gpus, args.lossless_accel)) + apply_attention_prefix(args.attention_prefix) + if args.width is None: + args.width = profile.width if profile is not None else DEFAULT_WIDTH + if args.height is None: + args.height = profile.height if profile is not None else DEFAULT_HEIGHT + generator = VideoGenerator.from_config( + build_generator_config(model_dir, args.num_gpus, args.lossless_accel, profile)) try: result = generator.generate( GenerationRequest( diff --git a/examples/inference/basic/basic_fasth3_omniref_rtx4090.yaml b/examples/inference/basic/basic_fasth3_omniref_rtx4090.yaml new file mode 100644 index 000000000..f549deedc --- /dev/null +++ b/examples/inference/basic/basic_fasth3_omniref_rtx4090.yaml @@ -0,0 +1,78 @@ +# FastH3 OmniRef (Ref2VA) on one RTX 4090 (24 GB, no FP4): the FP8/bf16 equivalent, one clip per prompt. +# +# Point model_path at a composed OmniRef model directory: the distilled +# dense or FP8 transformer_ref export (without nvfp4_weights.safetensors) symlinked +# together with base MiniMax-H3's text_encoder, tokenizer, processor, vae and +# audio_vae. examples/inference/basic/basic_fasth3_omniref_pdd.py builds that +# directory and runs it with --profile rtx-4090. +# +# The contract in the export's fastvideo_inference.json fixes the attention +# backend, the VSA tile size and the eight-forward PDD partition; the run only +# picks the device's offload policy and its environment: +# +# sm_89 has no FP4 hardware, so this profile is the closest equivalent rather than +# the FP4 stack: FP8 linears, bf16 VSA attention, the dense light VAE decoder (the +# NVFP4 one needs FP4 hardware too), a non-resident encoder and expandable segments +# to reduce memory fragmentation. Layerwise offload is mandatory: the 47 GB FP8 +# DiT cannot be resident on 24 GB. Run one process per clip; repeated requests in +# one process exceed the 24 GB GPU and the host cgroup when the encoder reloads. +# +# FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 \ +# FASTVIDEO_H3_VSA_FP4=0 \ +# FASTVIDEO_H3_VSA_FINE_ATTENTION=emulator-bf16 \ +# FASTVIDEO_LAYERWISE_OFFLOAD_BUFFERS=1 \ +# FASTVIDEO_LAYERWISE_RESIDENT_BLOCKS=16 \ +# FASTVIDEO_H3_FFN_CHUNK_TOKENS=16384 \ +# FASTVIDEO_H3_PINNED_SWAP=0 \ +# PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \ +# FASTVIDEO_H3_ADALN_TABLE=/path/to/adaln_table.pt \ +# fastvideo generate --config examples/inference/basic/basic_fasth3_omniref_rtx4090.yaml +# +# See docs/inference/fasth3_omniref.md. +generator: + model_path: ./FastH3-OmniRef-Dense-Composed + engine: + num_gpus: 1 + use_fsdp_inference: false + quantization: + transformer_quant: FP8 + parallelism: + tp_size: 1 + sp_size: 1 + offload: + dit: false + dit_layerwise: true + text_encoder: true + vae: false + pin_cpu_memory: true + lazy_module_load: false + pipeline: + workload_type: i2v + components: + override_pipeline_cls_name: MiniMaxH3Ref2VAModularPipeline + experimental: + attention_backend: VIDEO_SPARSE_ATTN_H3 + video_decode_backend: h3-vae + h3_sequential_load: true +request: + prompt: >- + A person on a sofa, lit by a warm lamp, speaking to camera. + negative_prompt: "" + inputs: + references: + # Ordered: the clip is conditioned on these. Replace with your own media. + - source: assets/images/couple.jpg + media_type: image + sampling: + seed: 1234 + height: 480 + width: 832 + num_frames: 124 + fps: 24 + num_inference_steps: 8 # eight fused PDD blocks; the contract sets it + guidance_scale: 1.0 + batch_cfg: false + output: + output_path: outputs/basic_fasth3_omniref_rtx4090/ + save_video: true + return_frames: false diff --git a/examples/inference/basic/basic_fasth3_omniref_rtx5090.yaml b/examples/inference/basic/basic_fasth3_omniref_rtx5090.yaml new file mode 100644 index 000000000..88e1c77d6 --- /dev/null +++ b/examples/inference/basic/basic_fasth3_omniref_rtx5090.yaml @@ -0,0 +1,82 @@ +# FastH3 OmniRef (Ref2VA) on one RTX 5090 (32 GB): the full NVFP4 stack, one clip per prompt. +# +# Point model_path at a composed OmniRef model directory: the distilled +# transformer_ref export (with its packed nvfp4_weights.safetensors) symlinked +# together with base MiniMax-H3's text_encoder, tokenizer, processor, vae and +# audio_vae. examples/inference/basic/basic_fasth3_omniref_pdd.py builds that +# directory and runs the same stack with --profile rtx-5090. +# +# The contract in the export's fastvideo_inference.json fixes the attention +# backend, the VSA tile size and the eight-forward PDD partition; the run only +# picks the device's offload policy and its environment: +# +# 32 GB: the text encoder stays off the GPU, the denoise modules load on demand and +# the AdaLN table replaces adaln_proj. The DiT itself stays whole - no layerwise streaming. +# FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 \ +# FASTVIDEO_H3_VSA_FP4=1 FASTVIDEO_H3_VSA_FP4_PREFIX=fp8 \ +# FASTVIDEO_H3_VSA_FP4_TWO_LEVEL_P=1 FASTVIDEO_H3_VSA_FP4_K_SMOOTH=1 \ +# FASTVIDEO_H3_GATE_CHUNK_ROWS=8192 \ +# FASTVIDEO_ULYSSES_A2A=auto FASTVIDEO_MINIMAX_H3_EXACT_KERNELS=all \ +# FASTVIDEO_H3_VSA_HEADS_FIRST_TILE=1 FASTVIDEO_H3_VAE_TILE_PARALLEL=1 \ +# FASTVIDEO_H3_REF2VA_MEMO_ENTRIES=16 \ +# FASTVIDEO_MINIMAX_H3_FUSIONS=all \ +# FASTVIDEO_VSA_TRITON=1 FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0 \ +# FASTVIDEO_H3_ADALN_TABLE=/path/to/adaln_table.pt \ +# FASTVIDEO_H3_VAE_NVFP4_DECODER=/path/to/decoder-nvfp4.safetensors \ +# FASTVIDEO_H3_VAE_TILE_BATCH=4 \ +# fastvideo generate --config examples/inference/basic/basic_fasth3_omniref_rtx5090.yaml +# +# See docs/inference/fasth3_omniref.md. +generator: + model_path: ./FastH3-OmniRef-Composed + engine: + num_gpus: 1 + use_fsdp_inference: false + quantization: + transformer_quant: NVFP4 + layer_profile: h3_dit_vsa + parallelism: + tp_size: 1 + sp_size: 1 + offload: + dit: false + dit_layerwise: false + text_encoder: true + vae: false + pin_cpu_memory: true + lazy_module_load: false + pipeline: + workload_type: i2v + components: + override_pipeline_cls_name: MiniMaxH3Ref2VAModularPipeline + experimental: + attention_backend: VIDEO_SPARSE_ATTN_H3 + h3_vsa_fp4: + enabled: true + prefix_precision: fp8 + two_level_p: true + smooth_k: true + video_decode_backend: h3-vae + h3_sequential_load: true +request: + prompt: >- + A person on a sofa, lit by a warm lamp, speaking to camera. + negative_prompt: "" + inputs: + references: + # Ordered: the clip is conditioned on these. Replace with your own media. + - source: assets/images/couple.jpg + media_type: image + sampling: + seed: 1234 + height: 480 + width: 832 + num_frames: 124 + fps: 24 + num_inference_steps: 8 # eight fused PDD blocks; the contract sets it + guidance_scale: 1.0 + batch_cfg: false + output: + output_path: outputs/basic_fasth3_omniref_rtx5090/ + save_video: true + return_frames: false diff --git a/examples/inference/basic/basic_fasth3_omniref_rtx_pro6000.yaml b/examples/inference/basic/basic_fasth3_omniref_rtx_pro6000.yaml new file mode 100644 index 000000000..8b8a9c496 --- /dev/null +++ b/examples/inference/basic/basic_fasth3_omniref_rtx_pro6000.yaml @@ -0,0 +1,78 @@ +# FastH3 OmniRef (Ref2VA) on one RTX PRO 6000 (96 GB): the full NVFP4 stack, one clip per prompt. +# +# Point model_path at a composed OmniRef model directory: the distilled +# transformer_ref export (with its packed nvfp4_weights.safetensors) symlinked +# together with base MiniMax-H3's text_encoder, tokenizer, processor, vae and +# audio_vae. examples/inference/basic/basic_fasth3_omniref_pdd.py builds that +# directory and runs the same stack with --profile rtx-pro-6000. +# +# The contract in the export's fastvideo_inference.json fixes the attention +# backend, the VSA tile size and the eight-forward PDD partition; the run only +# picks the device's offload policy and its environment. 96 GB fits the encoder, +# the DiT and the decoder at once. +# +# FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 \ +# FASTVIDEO_H3_VSA_FP4=1 FASTVIDEO_H3_VSA_FP4_PREFIX=fp8 \ +# FASTVIDEO_H3_VSA_FP4_TWO_LEVEL_P=1 FASTVIDEO_H3_VSA_FP4_K_SMOOTH=1 \ +# FASTVIDEO_H3_GATE_CHUNK_ROWS=8192 \ +# FASTVIDEO_MINIMAX_H3_FUSIONS=all \ +# FASTVIDEO_VSA_TRITON=1 FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0 \ +# FASTVIDEO_H3_ADALN_TABLE=/path/to/adaln_table.pt \ +# FASTVIDEO_H3_VAE_NVFP4_DECODER=/path/to/decoder-nvfp4.safetensors \ +# FASTVIDEO_H3_VAE_TILE_BATCH=8 \ +# fastvideo generate --config examples/inference/basic/basic_fasth3_omniref_rtx_pro6000.yaml +# +# See docs/inference/fasth3_omniref.md. +generator: + model_path: ./FastH3-OmniRef-Composed + engine: + num_gpus: 1 + use_fsdp_inference: false + quantization: + transformer_quant: NVFP4 + layer_profile: h3_dit_vsa + parallelism: + tp_size: 1 + sp_size: 1 + offload: + dit: false + dit_layerwise: false + text_encoder: false + vae: false + pin_cpu_memory: true + lazy_module_load: false + pipeline: + workload_type: i2v + components: + override_pipeline_cls_name: MiniMaxH3Ref2VAModularPipeline + experimental: + attention_backend: VIDEO_SPARSE_ATTN_H3 + h3_vsa_fp4: + enabled: true + prefix_precision: fp8 + two_level_p: true + smooth_k: true + video_decode_backend: h3-vae + h3_sequential_load: false +request: + prompt: >- + A person on a sofa, lit by a warm lamp, speaking to camera. + negative_prompt: "" + inputs: + references: + # Ordered: the clip is conditioned on these. Replace with your own media. + - source: assets/images/couple.jpg + media_type: image + sampling: + seed: 1234 + height: 768 + width: 1344 + num_frames: 124 + fps: 24 + num_inference_steps: 8 # eight fused PDD blocks; the contract sets it + guidance_scale: 1.0 + batch_cfg: false + output: + output_path: outputs/basic_fasth3_omniref_rtx_pro6000/ + save_video: true + return_frames: false diff --git a/examples/inference/basic/basic_fasth3_omniref_spark.yaml b/examples/inference/basic/basic_fasth3_omniref_spark.yaml new file mode 100644 index 000000000..684c779ed --- /dev/null +++ b/examples/inference/basic/basic_fasth3_omniref_spark.yaml @@ -0,0 +1,83 @@ +# FastH3 OmniRef (Ref2VA) on one DGX Spark (GB10, unified memory): the full NVFP4 stack, one clip per prompt. +# +# Point model_path at a composed OmniRef model directory: the distilled +# transformer_ref export (with its packed nvfp4_weights.safetensors) symlinked +# together with base MiniMax-H3's text_encoder, tokenizer, processor, vae and +# audio_vae. examples/inference/basic/basic_fasth3_omniref_pdd.py builds that +# directory and runs the same stack with --profile spark. +# +# The contract in the export's fastvideo_inference.json fixes the attention +# backend, the VSA tile size and the eight-forward PDD partition; the run only +# picks the device's offload policy and its environment: +# +# GB10: everything stays resident; cuDNN mm_fp4 and the dense batched decoder. +# FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 \ +# FASTVIDEO_H3_VSA_FP4=1 FASTVIDEO_H3_VSA_FP4_PREFIX=fp8 \ +# FASTVIDEO_H3_VSA_FP4_TWO_LEVEL_P=1 FASTVIDEO_H3_VSA_FP4_K_SMOOTH=1 \ +# FASTVIDEO_H3_GATE_CHUNK_ROWS=8192 \ +# FASTVIDEO_ULYSSES_A2A=auto FASTVIDEO_MINIMAX_H3_EXACT_KERNELS=all \ +# FASTVIDEO_H3_VSA_HEADS_FIRST_TILE=1 FASTVIDEO_H3_VAE_TILE_PARALLEL=1 \ +# FASTVIDEO_H3_REF2VA_MEMO_ENTRIES=16 \ +# FASTVIDEO_MINIMAX_H3_FUSIONS=all \ +# FASTVIDEO_VSA_TRITON=1 FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0 \ +# FASTVIDEO_H3_ADALN_TABLE=/path/to/adaln_table.pt \ +# FASTVIDEO_H3_VAE_NVFP4_DECODER=/path/to/decoder-nvfp4.safetensors \ +# FASTVIDEO_NVFP4_MM_BACKEND=cudnn FASTVIDEO_NVFP4_MM_ROW_SPLIT=table \ +# FASTVIDEO_H3_VAE_TILE_BATCH=12 FASTVIDEO_H3_VAE_INT8_OVERLAY=0 \ +# fastvideo generate --config examples/inference/basic/basic_fasth3_omniref_spark.yaml +# +# See docs/inference/fasth3_omniref.md. +generator: + model_path: ./FastH3-OmniRef-Composed + engine: + num_gpus: 1 + use_fsdp_inference: false + quantization: + transformer_quant: NVFP4 + layer_profile: h3_dit_vsa + parallelism: + tp_size: 1 + sp_size: 1 + offload: + dit: false + dit_layerwise: false + text_encoder: false + image_encoder: false + vae: false + pin_cpu_memory: false + lazy_module_load: false + pipeline: + workload_type: i2v + components: + override_pipeline_cls_name: MiniMaxH3Ref2VAModularPipeline + experimental: + attention_backend: VIDEO_SPARSE_ATTN_H3 + h3_vsa_fp4: + enabled: true + prefix_precision: fp8 + two_level_p: true + smooth_k: true + video_decode_backend: h3-vae + h3_sequential_load: false +request: + prompt: >- + A person on a sofa, lit by a warm lamp, speaking to camera. + negative_prompt: "" + inputs: + references: + # Ordered: the clip is conditioned on these. Replace with your own media. + - source: assets/images/couple.jpg + media_type: image + sampling: + seed: 1234 + height: 480 + width: 832 + num_frames: 124 + fps: 24 + num_inference_steps: 8 # eight fused PDD blocks; the contract sets it + guidance_scale: 1.0 + batch_cfg: false + output: + output_path: outputs/basic_fasth3_omniref_spark/ + save_video: true + return_frames: false diff --git a/examples/inference/basic/basic_fasth3_omniref_spark_pair.yaml b/examples/inference/basic/basic_fasth3_omniref_spark_pair.yaml new file mode 100644 index 000000000..85a47ffd8 --- /dev/null +++ b/examples/inference/basic/basic_fasth3_omniref_spark_pair.yaml @@ -0,0 +1,91 @@ +# FastH3 OmniRef (Ref2VA) across two DGX Sparks over QSFP RoCE, eight fused PDD +# forwards per clip. Start Ray on both Sparks and source +# examples/inference/optimizations/spark_pair_env.sh first; see +# docs/getting_started/installation/spark_pair.md. +# +# Point model_path at a composed OmniRef model directory (see +# basic_fasth3_omniref_pdd.py): the distilled transformer_ref export, with its +# packed nvfp4_weights.safetensors, symlinked with base MiniMax-H3's encoder, +# tokenizer, processor and VAEs. +# +# GB10 uses Triton VSA. The two nodes share the sequence (sp_size 2), which is +# the only split this pipeline has: H3 generates a clip at a time, so there is +# no depth to pipeline. The FP8 prefix split runs on the sequence-parallel path +# too: prefix rows stay dense and video rows stay on the FP4 kernel. +# +# FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 \ +# FASTVIDEO_H3_VSA_FP4=1 FASTVIDEO_H3_VSA_FP4_PREFIX=fp8 \ +# FASTVIDEO_H3_VSA_FP4_TWO_LEVEL_P=1 FASTVIDEO_H3_VSA_FP4_K_SMOOTH=1 \ +# FASTVIDEO_H3_GATE_CHUNK_ROWS=8192 \ +# FASTVIDEO_ULYSSES_A2A=auto FASTVIDEO_MINIMAX_H3_EXACT_KERNELS=all \ +# FASTVIDEO_H3_VSA_HEADS_FIRST_TILE=1 FASTVIDEO_H3_VAE_TILE_PARALLEL=1 \ +# FASTVIDEO_H3_REF2VA_MEMO_ENTRIES=16 \ +# FASTVIDEO_NVFP4_MM_BACKEND=cudnn \ +# FASTVIDEO_H3_VAE_TILE_BATCH=12 FASTVIDEO_H3_VAE_INT8_OVERLAY=0 \ +# FASTVIDEO_MINIMAX_H3_FUSIONS=all \ +# FASTVIDEO_VSA_TRITON=1 FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0 \ +# FASTVIDEO_H3_ADALN_TABLE=/path/to/adaln_table.pt \ +# FASTVIDEO_H3_VAE_NVFP4_DECODER=/path/to/decoder-nvfp4.safetensors \ +# fastvideo generate --config examples/inference/basic/basic_fasth3_omniref_spark_pair.yaml +# +# See docs/inference/fasth3_omniref.md. +generator: + model_path: ./FastH3-OmniRef-Composed + engine: + num_gpus: 2 + execution_backend: ray + use_fsdp_inference: false + quantization: + transformer_quant: NVFP4 + layer_profile: h3_dit_vsa + parallelism: + tp_size: 1 + sp_size: 2 + offload: + dit: false + dit_layerwise: false + text_encoder: false + image_encoder: false + vae: false + pin_cpu_memory: false + lazy_module_load: false + compile: + enabled: false + vae_enabled: true + pipeline: + workload_type: i2v + components: + override_pipeline_cls_name: MiniMaxH3Ref2VAModularPipeline + vae_tiling: true + experimental: + attention_backend: VIDEO_SPARSE_ATTN_H3 + h3_vsa_fp4: + enabled: true + prefix_precision: fp8 + two_level_p: true + smooth_k: true + vae_parallel_decode: true + vae_parallel_encode: true + video_decode_backend: h3-vae +request: + prompt: >- + A person on a sofa, lit by a warm lamp, speaking to camera. + negative_prompt: "" + inputs: + references: + # Ordered: the clip is conditioned on these. Replace with your own media. + - source: assets/images/couple.jpg + media_type: image + sampling: + seed: 1234 + height: 480 + width: 832 + num_frames: 124 + fps: 24 + num_inference_steps: 8 # eight fused PDD blocks; the contract sets it + guidance_scale: 1.0 + batch_cfg: false + output: + output_path: outputs/fasth3_omniref_spark_pair/ + save_video: true + return_frames: false diff --git a/examples/inference/basic/basic_fasth3_spark_pair_qad_v2_nvfp4.yaml b/examples/inference/basic/basic_fasth3_spark_pair_qad_v2_nvfp4.yaml new file mode 100644 index 000000000..caec53606 --- /dev/null +++ b/examples/inference/basic/basic_fasth3_spark_pair_qad_v2_nvfp4.yaml @@ -0,0 +1,72 @@ +# Start Ray on both Sparks and source spark_pair_env.sh first; see spark_pair.md. +# FastH3-QAD-v2 eight-forward video+audio on two DGX Sparks over QSFP RoCE. +# Download the complete release checkpoint as described in +# docs/getting_started/installation/spark_performance.md. +# +# FP4 attention needs fastvideo-kernel built for 12.1a and eager DiT inference. +# A missing extension is an error. Set h3_vsa_fp4.enabled to false for BF16 VSA. +# FastH3-QAD-v2 retains its trained 64-token VSA layout. The protected FP4 attention +# settings below are an inference choice, not its linears-only QAD recipe. +# The NVFP4 encoder and lightweight VAE ship in the same repository. +# +# FASTVIDEO_H3_VSA_FP4=1 FASTVIDEO_MINIMAX_H3_FUSIONS=all FASTVIDEO_NVFP4_MM_BACKEND=cutlass \ +# FASTVIDEO_H3_VAE_TILE_BATCH=1 \ +# FASTVIDEO_VSA_TRITON=1 FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0 \ +# FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 FASTVIDEO_STAGE_LOGGING=1 \ +# fastvideo generate --config examples/inference/basic/basic_fasth3_spark_pair_qad_v2_nvfp4.yaml +generator: + model_path: FastVideo/FastH3-QAD-v2 + engine: + num_gpus: 2 + execution_backend: ray + use_fsdp_inference: false + quantization: + transformer_quant: NVFP4 + layer_profile: h3_dit_vsa + parallelism: + tp_size: 1 + sp_size: 2 + offload: + dit: false + dit_layerwise: false + text_encoder: false + image_encoder: false + vae: false + pin_cpu_memory: false + lazy_module_load: false + compile: + enabled: false + vae_enabled: false + pipeline: + workload_type: t2v + vae_tiling: true + experimental: + attention_backend: VIDEO_SPARSE_ATTN_H3 + h3_vsa_fp4: + enabled: true + prefix_precision: fp8 + two_level_p: true + smooth_k: true + VSA_sparsity: 0.8 + VSA_tile_size: 64 + h3_sequential_load: false + inference_torch_compile: false + vae_parallel_decode: true + vae_parallel_decode_strategy: gather + video_decode_backend: h3-vae +request: + prompt: A quiet pottery studio with a potter finishing a bowl at the wheel. + negative_prompt: "" + sampling: + seed: 1234 + height: 480 + width: 832 + num_frames: 124 + fps: 24 + num_inference_steps: 9 # nine sigma points, eight DiT forwards + guidance_scale: 1.0 + batch_cfg: false + output: + output_path: outputs/fasth3_spark_pair_qad_v2_nvfp4/ + save_video: true + return_frames: false diff --git a/examples/inference/basic/basic_fasth3_spark_pair_v2_nvfp4.yaml b/examples/inference/basic/basic_fasth3_spark_pair_v2_nvfp4.yaml index 3e76091b2..efaa1b3e2 100644 --- a/examples/inference/basic/basic_fasth3_spark_pair_v2_nvfp4.yaml +++ b/examples/inference/basic/basic_fasth3_spark_pair_v2_nvfp4.yaml @@ -3,10 +3,11 @@ # Download the complete release checkpoint as described in # docs/getting_started/installation/spark_performance.md. # -# GB10 uses Triton VSA. The release packs attention, FFN and VSA gate weights. +# FP4 attention needs fastvideo-kernel built for 12.1a and eager DiT inference. +# A missing extension is an error. Set h3_vsa_fp4.enabled to false for BF16 VSA. # The NVFP4 encoder and lightweight VAE ship in the same repository. # -# FASTVIDEO_MINIMAX_H3_FUSIONS=all FASTVIDEO_NVFP4_MM_BACKEND=cutlass \ +# FASTVIDEO_H3_VSA_FP4=1 FASTVIDEO_MINIMAX_H3_FUSIONS=all FASTVIDEO_NVFP4_MM_BACKEND=cutlass \ # FASTVIDEO_H3_VAE_TILE_BATCH=1 \ # FASTVIDEO_VSA_TRITON=1 FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0 \ # FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 FASTVIDEO_STAGE_LOGGING=1 \ @@ -39,6 +40,8 @@ generator: vae_tiling: true experimental: attention_backend: VIDEO_SPARSE_ATTN_H3 + h3_vsa_fp4: + enabled: true VSA_sparsity: 0.8 VSA_tile_size: 64 h3_sequential_load: false diff --git a/examples/inference/basic/basic_fasth3_spark_pruned_nvfp4.yaml b/examples/inference/basic/basic_fasth3_spark_pruned_nvfp4.yaml index d9b86d72f..11fb6efb8 100644 --- a/examples/inference/basic/basic_fasth3_spark_pruned_nvfp4.yaml +++ b/examples/inference/basic/basic_fasth3_spark_pruned_nvfp4.yaml @@ -4,9 +4,16 @@ # # GB10 uses Triton VSA. The release packs attention, FFN and VSA gate weights. # The NVFP4 encoder and lightweight VAE ship in the same repository. +# The VAE decodes with its dense weights, 12 tiles per call and a compiled +# decoder: on GB10 the INT8 ConvRot overlay is slower and less accurate. +# NVFP4 GEMMs use FlashInfer's cuDNN backend, 2-3x faster than CUTLASS on the +# H3 FFN shapes on GB10 with identical outputs. +# Optional, off by default: FASTVIDEO_H3_VSA_FP4=1 runs VSA on the FP4 attention kernel +# (fastvideo-kernel built for 12.1a). It is faster but changes the output for now; +# see spark_performance.md. # -# FASTVIDEO_MINIMAX_H3_FUSIONS=all FASTVIDEO_NVFP4_MM_BACKEND=cutlass \ -# FASTVIDEO_H3_VAE_TILE_BATCH=1 \ +# FASTVIDEO_MINIMAX_H3_FUSIONS=all FASTVIDEO_NVFP4_MM_BACKEND=cudnn \ +# FASTVIDEO_H3_VAE_TILE_BATCH=12 FASTVIDEO_H3_VAE_INT8_OVERLAY=0 \ # FASTVIDEO_VSA_TRITON=1 FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0 \ # FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 FASTVIDEO_STAGE_LOGGING=1 \ # fastvideo generate --config examples/inference/basic/basic_fasth3_spark_pruned_nvfp4.yaml @@ -31,7 +38,7 @@ generator: lazy_module_load: false compile: enabled: false - vae_enabled: false + vae_enabled: true pipeline: workload_type: t2v vae_tiling: true diff --git a/examples/inference/basic/basic_fasth3_spark_qad_v2_nvfp4.yaml b/examples/inference/basic/basic_fasth3_spark_qad_v2_nvfp4.yaml new file mode 100644 index 000000000..610f44393 --- /dev/null +++ b/examples/inference/basic/basic_fasth3_spark_qad_v2_nvfp4.yaml @@ -0,0 +1,69 @@ +# FastH3-QAD-v2 eight-forward video+audio on one DGX Spark. +# Download the complete release checkpoint as described in +# docs/getting_started/installation/spark_performance.md. +# +# FP4 attention needs fastvideo-kernel built for 12.1a and eager DiT inference. +# A missing extension is an error. Set h3_vsa_fp4.enabled to false for BF16 VSA. +# FastH3-QAD-v2 retains its trained 64-token VSA layout. The protected FP4 attention +# settings below are an inference choice, not its linears-only QAD recipe. +# The NVFP4 encoder and lightweight VAE ship in the same repository. +# +# FASTVIDEO_H3_VSA_FP4=1 FASTVIDEO_MINIMAX_H3_FUSIONS=all FASTVIDEO_NVFP4_MM_BACKEND=cudnn FASTVIDEO_NVFP4_MM_ROW_SPLIT=table \ +# FASTVIDEO_H3_VAE_TILE_BATCH=12 FASTVIDEO_H3_VAE_INT8_OVERLAY=0 \ +# FASTVIDEO_VSA_TRITON=1 FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0 \ +# FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 FASTVIDEO_STAGE_LOGGING=1 \ +# fastvideo generate --config examples/inference/basic/basic_fasth3_spark_qad_v2_nvfp4.yaml +generator: + model_path: FastVideo/FastH3-QAD-v2 + engine: + num_gpus: 1 + use_fsdp_inference: false + quantization: + transformer_quant: NVFP4 + layer_profile: h3_dit_vsa + parallelism: + tp_size: 1 + sp_size: 1 + offload: + dit: false + dit_layerwise: false + text_encoder: false + image_encoder: false + vae: false + pin_cpu_memory: false + lazy_module_load: false + compile: + enabled: false + vae_enabled: true + pipeline: + workload_type: t2v + vae_tiling: true + experimental: + attention_backend: VIDEO_SPARSE_ATTN_H3 + h3_vsa_fp4: + enabled: true + prefix_precision: fp8 + two_level_p: true + smooth_k: true + VSA_sparsity: 0.8 + VSA_tile_size: 64 + h3_sequential_load: false + inference_torch_compile: false + vae_parallel_decode: false + video_decode_backend: h3-vae +request: + prompt: A quiet pottery studio with a potter finishing a bowl at the wheel. + negative_prompt: "" + sampling: + seed: 1234 + height: 480 + width: 832 + num_frames: 124 + fps: 24 + num_inference_steps: 9 # nine sigma points, eight DiT forwards + guidance_scale: 1.0 + batch_cfg: false + output: + output_path: outputs/fasth3_spark_qad_v2_nvfp4/ + save_video: true + return_frames: false diff --git a/examples/inference/basic/basic_fasth3_spark_v2_nvfp4.yaml b/examples/inference/basic/basic_fasth3_spark_v2_nvfp4.yaml index 023010956..69ad31fb1 100644 --- a/examples/inference/basic/basic_fasth3_spark_v2_nvfp4.yaml +++ b/examples/inference/basic/basic_fasth3_spark_v2_nvfp4.yaml @@ -2,11 +2,12 @@ # Download the complete release checkpoint as described in # docs/getting_started/installation/spark_performance.md. # -# GB10 uses Triton VSA. The release packs attention, FFN and VSA gate weights. +# FP4 attention needs fastvideo-kernel built for 12.1a and eager DiT inference. +# A missing extension is an error. Set h3_vsa_fp4.enabled to false for BF16 VSA. # The NVFP4 encoder and lightweight VAE ship in the same repository. # -# FASTVIDEO_MINIMAX_H3_FUSIONS=all FASTVIDEO_NVFP4_MM_BACKEND=cutlass \ -# FASTVIDEO_H3_VAE_TILE_BATCH=1 \ +# FASTVIDEO_H3_VSA_FP4=1 FASTVIDEO_MINIMAX_H3_FUSIONS=all FASTVIDEO_NVFP4_MM_BACKEND=cudnn FASTVIDEO_NVFP4_MM_ROW_SPLIT=table \ +# FASTVIDEO_H3_VAE_TILE_BATCH=12 FASTVIDEO_H3_VAE_INT8_OVERLAY=0 \ # FASTVIDEO_VSA_TRITON=1 FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0 \ # FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 FASTVIDEO_STAGE_LOGGING=1 \ # fastvideo generate --config examples/inference/basic/basic_fasth3_spark_v2_nvfp4.yaml @@ -31,12 +32,14 @@ generator: lazy_module_load: false compile: enabled: false - vae_enabled: false + vae_enabled: true pipeline: workload_type: t2v vae_tiling: true experimental: attention_backend: VIDEO_SPARSE_ATTN_H3 + h3_vsa_fp4: + enabled: true VSA_sparsity: 0.8 VSA_tile_size: 64 h3_sequential_load: false diff --git a/fastvideo-kernel/attn_qat_infer/api.py b/fastvideo-kernel/attn_qat_infer/api.py index e91196a82..f70d29218 100644 --- a/fastvideo-kernel/attn_qat_infer/api.py +++ b/fastvideo-kernel/attn_qat_infer/api.py @@ -23,6 +23,10 @@ from torch.nn.functional import scaled_dot_product_attention as sdpa import fp4attn_cuda import fp4quant_cuda +# ``sageattn_blackwell_sparse_bshd`` sizes its block lists and delta_s by the key length, +# so its query may cover only part of the keys (the MiniMax-H3 FP4 prefix split). +KEY_LENGTH_SIZING = True + # Centralized block size configuration for sageattn_blackwell kernels # These should match the values in fastvideo/attention/backends/sageattn/blackwell/block_config.h BLOCK_M = 128 # Block size for M dimension (query sequence length) @@ -348,14 +352,15 @@ def sageattn_blackwell_sparse_bshd(q, unsmoothed. Returns ``[B, H, L, D]``. """ batch, seq_len, heads, _ = q.shape - if seq_len % BLOCK_M: - raise ValueError(f"sequence length {seq_len} must be a multiple of {BLOCK_M}") + kv_len = k.shape[1] + if seq_len % BLOCK_M or kv_len % BLOCK_M: + raise ValueError(f"query length {seq_len} and key length {kv_len} must be multiples of {BLOCK_M}") if validate: - check_sparse_block_lists(q2k_idx, q2k_num, seq_len) + check_sparse_block_lists(q2k_idx, q2k_num, kv_len) qh, kh, vh = (x.transpose(1, 2) for x in (q, k, v)) - delta_s = _zero_delta_s(batch, heads, seq_len, q.device) + delta_s = _zero_delta_s(batch, heads, kv_len, q.device) return blockscaled_fp4_attn_sparse(scale_and_quant_fp4(qh), scale_and_quant_fp4_permute(kh), - scale_and_quant_fp4_transpose(vh), delta_s, seq_len, q2k_idx, q2k_num, kv_valid, + scale_and_quant_fp4_transpose(vh), delta_s, kv_len, q2k_idx, q2k_num, kv_valid, q2k_quad, False, q.dtype == torch.bfloat16, single_level_p_quant, sm_scale)[0] diff --git a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_reference.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_reference.py new file mode 100644 index 000000000..d4e7554cc --- /dev/null +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_reference.py @@ -0,0 +1,264 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Slow PyTorch reference for block-sparse FP4 attention as ``attn_qat_infer`` computes it. + +This is the numeric specification that ``attn_qat_vsa_train`` (the Triton +fake-quant training kernel) must reproduce. It follows the deployed sm_120/121 +kernel (``sageattn_blackwell_sparse_bshd`` on lists from +``vsa_tile_mask_to_fp4_blocks``) step by step: + +* Q and K: NVFP4 E2M1 with one E4M3 scale per 16 channels along D, no global + scale (``scaled_fp4_quant(x, ..., 1)``): ``s = e4m3(amax / 6)``, value + ``e2m1(x / s) * s``, round-to-nearest-even, saturating. +* V: the same, but grouped along tokens (``scale_and_quant_fp4_transpose``). +* Each 128-row query block visits its listed 128-key blocks in ascending block + order (the lists are stored descending and read from the last entry). +* P: the online softmax of ``softmax_fused.h``. For each visited block the row + max is updated first, ``P~ = exp2(S * log2e * scale - m_scaled)``; per 16 keys + ``s_P = max(P~) / 6`` is computed in fp32, ``P~ / s_P`` is rounded to E2M1 and + the MMA multiplies by ``e4m3(s_P)``. Single-level P has ``m_scaled = m``; + two-level P adds ``log2(1 / (448 * 6))`` so ``P~`` spans ``[0, 2688]``. +* The row sum uses the unquantized ``P~``; the output is ``acc / l``. +* Masked keys (``kv_valid`` halves) score ``-inf``; a fully masked 16-key group + has ``s_P = 0`` and contributes exactly zero. + +``first_block_max_floor`` models an uninitialized-register hazard in the +deployed kernel: on the first visited block, ``AbsMaxP`` is read before it is +written (``fmaxf(AbsMaxP, s)``). ``None`` assumes the compiler folds that to +the true max; a number (e.g. ``0.0``) floors the first block's group maxima +and hence the row max. The cross-device fidelity check decides which applies. +""" +from __future__ import annotations + +import math + +import torch + +QUANT_GROUP = 16 +BLOCK = 128 +HALF = BLOCK // 2 +E4M3_MAX = 448.0 +LOG2E = 1.4426950408889634 +LOG2_INV_FP4_MAX = -2.584962500721156 # log2(1 / 6) +LOG2_INV_TWO_LEVEL = -11.392317422778762 # log2(1 / (448 * 6)) + + +def e4m3_rn(x: torch.Tensor) -> torch.Tensor: + """Round non-negative fp32 scales to E4M3 (RNE, saturating at 448) and back to fp32.""" + return x.clamp(max=E4M3_MAX).to(torch.float8_e4m3fn).to(torch.float32) + + +def e2m1_rn(y: torch.Tensor) -> torch.Tensor: + """Round fp32 values to the E2M1 grid {0, .5, 1, 1.5, 2, 3, 4, 6} (RNE, saturating).""" + a = y.abs() + q = torch.where( + a <= 0.25, 0.0, + torch.where( + a < 0.75, 0.5, + torch.where( + a <= 1.25, 1.0, + torch.where(a < 1.75, 1.5, + torch.where(a <= 2.5, 2.0, torch.where(a < 3.5, 3.0, torch.where(a <= 5.0, 4.0, + 6.0))))))) + return torch.copysign(q, y) + + +def nvfp4_fake_quant_lastdim(x: torch.Tensor) -> torch.Tensor: + """Fake-quantize along the last dim in groups of 16 (no global scale); returns fp32.""" + if x.shape[-1] % QUANT_GROUP: + raise ValueError(f"last dim {x.shape[-1]} is not a multiple of {QUANT_GROUP}") + xf = x.float().unflatten(-1, (-1, QUANT_GROUP)) + scale = e4m3_rn(xf.abs().amax(dim=-1, keepdim=True) / 6.0) + inv = torch.where(scale == 0, torch.zeros_like(scale), 1.0 / scale) + return (e2m1_rn(xf * inv) * scale).flatten(-2) + + +def fake_quant_qkv(q: torch.Tensor, k: torch.Tensor, + v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """``[B, H, L, D]`` inputs -> dequantized Q (along D), K (along D), V (along L); fp32.""" + v_hat = nvfp4_fake_quant_lastdim(v.transpose(-1, -2)).transpose(-1, -2) + return nvfp4_fake_quant_lastdim(q), nvfp4_fake_quant_lastdim(k), v_hat + + +def smooth_keys(k: torch.Tensor) -> torch.Tensor: + """K smoothing as ``preprocess_qkv`` (``k -= k.mean(dim=-2)``): per-head token mean, in the input dtype. + + The dtype matters: FP4 rounding is so input-sensitive that smoothing in + fp32 instead of bf16 alone moves the attention output by ~8% rel-L2. The + mean runs over every row given (tile-ordered, pad rows included); a + deployed K-smoothing path must use the same rows. + """ + return k - k.mean(dim=-2, keepdim=True) + + +def key_valid_mask(kv_valid: torch.Tensor | None, n_blocks: int, device: torch.device) -> torch.Tensor: + """``[n_blocks, 128]`` bool: which key columns of each block are real tokens.""" + if kv_valid is None: + return torch.ones(n_blocks, BLOCK, dtype=torch.bool, device=device) + col = torch.arange(BLOCK, device=device) + limit = torch.where(col < HALF, kv_valid[0::2, None], kv_valid[1::2, None]) + return (col % HALF)[None, :] < limit + + +def tile128_mask_to_fp4_blocks(tile_mask: torch.Tensor, + tile_valid: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """``attn_qat_infer.api.vsa_tile_mask_to_fp4_blocks(mask, 128, tile_valid)`` without the CUDA extensions. + + ``attn_qat_infer`` imports its sm_12x extensions at module load, so data + center GPUs cannot call it; this is the same computation for 128-token + tiles (one tile per kernel block, so no quadrant masks and no padding + block): descending lists and per-64-key-half valid counts. The fidelity + check script asserts equality with the real function on sm_12x. + """ + n_tiles = tile_mask.shape[-1] + half_valid = (tile_valid.to(torch.int32)[:, None] - + torch.tensor([0, HALF], device=tile_mask.device, dtype=torch.int32)[None, :]).clamp(0, HALF) + block_mask = tile_mask & (tile_valid > 0)[None, None, None, :] + valid = half_valid.reshape(-1) + n_blocks = n_tiles + rev = block_mask.flip(-1) + pos = rev.cumsum(-1, dtype=torch.int32) - 1 + q2k_num = (pos[..., -1] + 1).contiguous() + cols = torch.arange(n_blocks - 1, -1, -1, device=tile_mask.device, dtype=torch.int32).expand_as(pos) + slot = torch.where(rev, pos, torch.full_like(pos, n_blocks)).long() + q2k_idx = torch.zeros((*tile_mask.shape[:2], n_blocks, n_blocks + 1), device=tile_mask.device, dtype=torch.int32) + q2k_idx.scatter_(-1, slot, cols) + return q2k_idx[..., :n_blocks].contiguous(), q2k_num, valid.contiguous() + + +def invert_block_lists(q2k_idx: torch.Tensor, q2k_num: torch.Tensor, + n_kblocks: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Query->key lists to key->query lists, keeping each entry's slot in its query list. + + Returns ``(k2q_idx [B, H, NK, M], k2q_slot [B, H, NK, M], k2q_num [B, H, NK])`` (int32). + """ + batch, heads, n_qblocks, width = q2k_idx.shape + device = q2k_idx.device + slots = torch.arange(width, device=device, dtype=torch.int32).expand_as(q2k_idx) + listed = slots < q2k_num[..., None] + slot_of = torch.full((batch, heads, n_qblocks, n_kblocks + 1), -1, device=device, dtype=torch.int32) + slot_of.scatter_(-1, torch.where(listed, q2k_idx.long(), n_kblocks), torch.where(listed, slots, -1)) + by_key = slot_of[..., :n_kblocks].transpose(-1, -2) # [B, H, NK, M] + present = by_key >= 0 + k2q_num = present.sum(-1, dtype=torch.int32) + position = torch.where(present, present.cumsum(-1) - 1, n_qblocks).long() + queries = torch.arange(n_qblocks, device=device, dtype=torch.int32).expand_as(by_key) + k2q_idx = torch.zeros((batch, heads, n_kblocks, n_qblocks + 1), device=device, dtype=torch.int32) + k2q_slot = torch.zeros_like(k2q_idx) + k2q_idx.scatter_(-1, position, queries) + k2q_slot.scatter_(-1, position, by_key) + return k2q_idx[..., :n_qblocks].contiguous(), k2q_slot[..., :n_qblocks].contiguous(), k2q_num.contiguous() + + +def visit_order(q2k_idx: torch.Tensor, q2k_num: torch.Tensor, b: int, h: int, m: int) -> list[int]: + """Blocks query block ``m`` visits, in kernel order (last list entry first).""" + count = int(q2k_num[b, h, m]) + return [int(block) for block in q2k_idx[b, h, m, :count].flip(0)] + + +@torch.no_grad() +def fp4_vsa_attention_reference(q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, + kv_valid: torch.Tensor | None = None, + *, + two_level_p: bool = False, + smooth_k: bool = False, + quantize: bool = True, + first_block_max_floor: float | None = None, + sm_scale: float | None = None, + return_dv_weights: bool = False): + """Emulate the deployed block-sparse FP4 attention on ``[B, H, L, D]`` inputs (L a multiple of 128). + + ``quantize=False`` runs the same visit order with exact softmax numerics + (BF16-input reference). Returns ``[B, H, L, D]`` fp32, and with + ``return_dv_weights`` also the dense ``[B, H, L, L]`` matrix ``W`` with + ``out = W @ v_hat`` (the forward-consistent dV weights). + """ + batch, heads, seq_len, dim = q.shape + if seq_len % BLOCK: + raise ValueError(f"sequence length {seq_len} must be a multiple of {BLOCK}") + scale_log2 = (sm_scale if sm_scale is not None else 1.0 / math.sqrt(dim)) * LOG2E + if smooth_k: + k = smooth_keys(k) + if quantize: + q_hat, k_hat, v_hat = fake_quant_qkv(q, k, v) + else: + q_hat, k_hat, v_hat = q.float(), k.float(), v.float() + offset = LOG2_INV_TWO_LEVEL if two_level_p else 0.0 + n_blocks = seq_len // BLOCK + valid = key_valid_mask(kv_valid, n_blocks, q.device) + out = torch.zeros(batch, heads, seq_len, dim, dtype=torch.float32, device=q.device) + weights = torch.zeros(batch, heads, seq_len, seq_len, device=q.device) if return_dv_weights else None + for b in range(batch): + for h in range(heads): + for m in range(n_blocks): + rows = slice(m * BLOCK, (m + 1) * BLOCK) + row_max = torch.full((BLOCK, ), -math.inf, device=q.device) + row_sum = torch.zeros(BLOCK, device=q.device) + acc = torch.zeros(BLOCK, dim, device=q.device) + visited: list[tuple[int, torch.Tensor, torch.Tensor]] = [] + for position, n in enumerate(visit_order(q2k_idx, q2k_num, b, h, m)): + cols = slice(n * BLOCK, (n + 1) * BLOCK) + scores = q_hat[b, h, rows] @ k_hat[b, h, cols].T + scores = scores.masked_fill(~valid[n][None, :], -math.inf) + group_max = scores.unflatten(-1, (-1, QUANT_GROUP)).amax(dim=-1) + if position == 0 and first_block_max_floor is not None: + group_max = group_max.clamp(min=first_block_max_floor) + new_max = torch.maximum(row_max, group_max.amax(dim=-1)) + m_scaled = new_max * scale_log2 + offset + p_tilde = torch.exp2(scores * scale_log2 - m_scaled[:, None]) + alpha = torch.exp2((row_max - new_max) * scale_log2) + alpha = torch.where(torch.isinf(row_max), torch.zeros_like(alpha), alpha) + row_sum = row_sum * alpha + p_tilde.sum(dim=-1) + if quantize: + s_p = torch.exp2(group_max * scale_log2 - m_scaled[:, None] + LOG2_INV_FP4_MAX) + s_full = s_p.repeat_interleave(QUANT_GROUP, dim=-1) + ratio = torch.where(s_full == 0, torch.zeros_like(p_tilde), p_tilde / s_full) + p_hat = e2m1_rn(ratio) * e4m3_rn(s_full) + else: + p_hat = p_tilde + acc = acc * alpha[:, None] + p_hat @ v_hat[b, h, cols] + visited.append((n, p_hat, new_max)) + row_max = new_max + inv = torch.where((row_sum == 0) | torch.isnan(row_sum), torch.zeros_like(row_sum), 1.0 / row_sum) + out[b, h, rows] = acc * inv[:, None] + if weights is not None: + for n, p_hat, block_max in visited: + decay = torch.exp2((block_max - row_max) * scale_log2) * inv + weights[b, h, rows, n * BLOCK:(n + 1) * BLOCK] = p_hat * decay[:, None] + if weights is not None: + return out, weights + return out + + +def dense_block_mask(q2k_idx: torch.Tensor, q2k_num: torch.Tensor, seq_len: int, + kv_valid: torch.Tensor | None = None) -> torch.Tensor: + """Token-level ``[B, H, L, L]`` bool mask of the listed blocks (and valid keys).""" + batch, heads, n_blocks, width = q2k_idx.shape + listed = torch.arange(width, device=q2k_idx.device) < q2k_num[..., None] + block_mask = torch.zeros(batch, heads, n_blocks, n_blocks + 1, dtype=torch.bool, device=q2k_idx.device) + block_mask.scatter_(-1, torch.where(listed, q2k_idx.long(), n_blocks), True) + block_mask = block_mask[..., :n_blocks] + token = block_mask.repeat_interleave(BLOCK, dim=-2).repeat_interleave(BLOCK, dim=-1)[..., :seq_len, :seq_len] + valid = key_valid_mask(kv_valid, n_blocks, q2k_idx.device).reshape(-1)[:seq_len] + return token & valid + + +def ste_gradients_reference(q_hat: torch.Tensor, k_hat: torch.Tensor, v_hat: torch.Tensor, token_mask: torch.Tensor, + dv_weights: torch.Tensor, grad_out: torch.Tensor, + sm_scale: float) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """The training kernel's STE gradients, computed densely in fp32. + + dQ and dK are the exact gradients of the high-precision surrogate + ``softmax(q_hat k_hat^T) v_hat`` on the sparse mask; dV uses the forward's + quantized, online-softmax weights (forward-consistent dV). + """ + q_hat, k_hat, v_hat = (t.detach().float().requires_grad_(True) for t in (q_hat, k_hat, v_hat)) + scores = (q_hat @ k_hat.transpose(-1, -2)) * sm_scale + probs = torch.softmax(scores.masked_fill(~token_mask, -math.inf), dim=-1) + surrogate = probs @ v_hat + dq, dk = torch.autograd.grad(surrogate, (q_hat, k_hat), grad_out.float()) + dv = dv_weights.transpose(-1, -2) @ grad_out.float() + return dq, dk, dv diff --git a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_train.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_train.py new file mode 100644 index 000000000..8a36fefcb --- /dev/null +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_train.py @@ -0,0 +1,515 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Block-sparse FP4 Attn-QAT for VSA tile-128 training (forward + backward). + +The forward reproduces the deployed sm_120/121 sparse FP4 kernel +(``attn_qat_infer.api.sageattn_blackwell_sparse_bshd`` on lists from +``vsa_tile_mask_to_fp4_blocks(mask, 128, vbs)``) on any CUDA GPU with Triton, +so a GB200 can train against Spark/5090 numerics. The numeric contract is +spelled out in ``attn_qat_vsa_reference`` (the slow PyTorch specification the +unit tests compare against): + +* Q/K/V are fake-quantized once per call (NVFP4, per-16 E4M3 scales, no global + scale; V grouped along tokens) and stored in BF16, which holds every + dequantized value exactly. +* Every query row visits its block's listed 128-key blocks in ascending order; + the running max is updated per 128-key block before P is formed, P is + quantized per 16 keys (single- or two-level), the row sum is unquantized. + +Backward is a straight-through estimator with one deliberate choice per +operand (the lessons of the GB10 Attn-QAT dV fix): + +* dV is forward-consistent: it is built from the forward's quantized + online-softmax weights, recomputed from the saved per-(row, visited block) + running maxima, never by re-quantizing the normalized probabilities (whose + E4M3 group scales underflow on long sequences and silence whole heads). +* dQ/dK are the exact gradients of the high-precision surrogate + ``softmax(q_hat k_hat^T) v_hat`` on the same sparse mask; delta uses the + high-precision output. +* The Q/K/V quantizers pass gradients through unchanged. + +Rows are independent, so programs cover 64 query rows (the forward keeps whole +128-key blocks, where the running max is defined); the backward splits keys +into 64-column halves, which the per-16 P groups allow. +""" +from __future__ import annotations + +import math + + +import torch +import triton +import triton.language as tl +from triton.language.extra import libdevice + +from .attn_qat_vsa_reference import BLOCK, LOG2_INV_TWO_LEVEL, LOG2E, invert_block_lists + +_ROWS = 64 # query rows per program +_HALF = BLOCK // 2 +_QUANT_ROWS = 64 # fake-quant rows per program (a multiple of 16; L is a multiple of 128) +_WARPS = 8 # 8 beats 4 by 15-20% on GB200 at 35k-73k tokens + + +@triton.jit +def _e2m1_rn(y): + """Round to the E2M1 grid {0, .5, 1, 1.5, 2, 3, 4, 6}, ties to even, saturating at 6. + + The grid is 0.5-spaced below 2, 1-spaced in [2, 4) and 2-spaced above, so + round-half-even on the scaled value is exact (ties land on even codes). + """ + a = tl.minimum(tl.abs(y), 6.0) + q = tl.where(a < 2.0, + libdevice.rint(a * 2.0) * 0.5, + tl.where(a < 4.0, libdevice.rint(a), + libdevice.rint(a * 0.5) * 2.0)) + return tl.where(y < 0, -q, q) + + +@triton.jit +def _e4m3_rn(x): + """Round non-negative fp32 to E4M3 (RN, saturating at 448) and back.""" + return tl.minimum(x, 448.0).to(tl.float8e4nv).to(tl.float32) + + +@triton.jit +def _quantized_p(s, m_scaled, first_visit, scale_log2, FLOOR: tl.constexpr, FLOOR_VALUE: tl.constexpr, + ROWS: tl.constexpr, COLS: tl.constexpr): + """``P~`` and its fake-quantized value for raw scores ``s`` (masked = -inf) under row maxima ``m_scaled``.""" + p_tilde = tl.math.exp2(s * scale_log2 - m_scaled[:, None]) + group_max = tl.max(tl.reshape(s, (ROWS, COLS // 16, 16)), axis=2) + if FLOOR: + group_max = tl.where(first_visit, tl.maximum(group_max, FLOOR_VALUE), group_max) + s_p = tl.math.exp2(group_max * scale_log2 - m_scaled[:, None] - 2.584962500721156) + # Per-group values (fp32 divisor, E4M3 multiplier) broadcast back over each group's 16 keys. + s_full = tl.reshape(tl.broadcast_to(s_p[:, :, None], (ROWS, COLS // 16, 16)), (ROWS, COLS)) + s_e4m3 = tl.reshape(tl.broadcast_to(_e4m3_rn(s_p)[:, :, None], (ROWS, COLS // 16, 16)), (ROWS, COLS)) + ratio = tl.where(s_full == 0.0, 0.0, p_tilde / s_full) + return p_tilde, _e2m1_rn(ratio) * s_e4m3 + + +@triton.jit +def _fake_quant_kernel(X, Y, n_rows, D: tl.constexpr, ROWS: tl.constexpr, ALONG_ROWS: tl.constexpr): + """NVFP4 fake quant of a ``[n_rows, D]`` matrix, groups of 16 along D or (``ALONG_ROWS``) along rows.""" + pid = tl.program_id(0).to(tl.int64) + offs_r = pid * ROWS + tl.arange(0, ROWS) + offs_d = tl.arange(0, D) + x = tl.load(X + offs_r[:, None] * D + offs_d[None, :]).to(tl.float32) + if ALONG_ROWS: + x3 = tl.reshape(x, (ROWS // 16, 16, D)) + amax = tl.max(tl.abs(x3), axis=1) + scale = _e4m3_rn(amax / 6.0) + inv = tl.where(scale == 0.0, 0.0, 1.0 / scale) + y3 = _e2m1_rn(x3 * inv[:, None, :]) * scale[:, None, :] + y = tl.reshape(y3, (ROWS, D)) + else: + x3 = tl.reshape(x, (ROWS, D // 16, 16)) + amax = tl.max(tl.abs(x3), axis=2) + scale = _e4m3_rn(amax / 6.0) + inv = tl.where(scale == 0.0, 0.0, 1.0 / scale) + y3 = _e2m1_rn(x3 * inv[:, :, None]) * scale[:, :, None] + y = tl.reshape(y3, (ROWS, D)) + tl.store(Y + offs_r[:, None] * D + offs_d[None, :], y.to(Y.dtype.element_ty)) + + +def fake_quant_qkv_triton(q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + quantize_v: bool = True) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """``fake_quant_qkv`` (Q/K along D, V along tokens) as one Triton pass each; BF16 ``[B, H, L, D]`` out.""" + outs = [] + for x, along_rows in ((q, False), (k, False), (v, True)): + x = x.contiguous() + if along_rows and not quantize_v: + outs.append(x.to(torch.bfloat16)) + continue + y = torch.empty(x.shape, device=x.device, dtype=torch.bfloat16) + rows = x.numel() // x.shape[-1] + _fake_quant_kernel[(rows // _QUANT_ROWS, )](x, y, rows, D=x.shape[-1], ROWS=_QUANT_ROWS, + ALONG_ROWS=along_rows, num_warps=4) + outs.append(y) + return outs[0], outs[1], outs[2] + + +@triton.jit +def _key_ok(KV_VALID, n_block, cols, HAS_VALID: tl.constexpr, HALF: tl.constexpr): + """Valid key columns (block-local ``cols`` in [0, 128)) of 128-key block ``n_block``.""" + if HAS_VALID: + valid0 = tl.load(KV_VALID + 2 * n_block) + valid1 = tl.load(KV_VALID + 2 * n_block + 1) + limit = tl.where(cols < HALF, valid0, valid1) + return (cols % HALF) < limit + return cols >= 0 + + +@triton.jit +def _fwd_kernel(Q, K, V, O, OHP, LSE, MRUN, Q2K_IDX, Q2K_NUM, KV_VALID, BIAS, QGROUP, n_groups, seq_len, n_qblocks, + list_width, scale_log2, OFFSET: tl.constexpr, QUANT: tl.constexpr, HAS_VALID: tl.constexpr, + HAS_BIAS: tl.constexpr, FLOOR: tl.constexpr, FLOOR_VALUE: tl.constexpr, HIGH_PREC_O: tl.constexpr, + SAVE_MRUN: tl.constexpr, + D: tl.constexpr, ROWS: tl.constexpr, BLOCK_N: tl.constexpr): + pid = tl.program_id(0) + bh = tl.program_id(1).to(tl.int64) + halves: tl.constexpr = BLOCK_N // ROWS + m_block = pid // halves + row0 = (pid % halves) * ROWS + offs_r = tl.arange(0, ROWS) + offs_n = tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, D) + base = bh * seq_len * D + rows = m_block * BLOCK_N + row0 + offs_r + q = tl.load(Q + base + rows[:, None] * D + offs_d[None, :]) + meta = bh * n_qblocks + m_block + count = tl.load(Q2K_NUM + meta) + m_i = tl.full((ROWS, ), float("-inf"), tl.float32) + l_i = tl.zeros((ROWS, ), tl.float32) + acc = tl.zeros((ROWS, D), tl.float32) + acc_hp = tl.zeros((ROWS, D), tl.float32) + if HAS_BIAS: + bias_row = BIAS + (bh * n_groups + tl.load(QGROUP + m_block)) * seq_len + for it in range(0, count): + slot = count - 1 - it + n_block = tl.load(Q2K_IDX + meta * list_width + slot) + keys = n_block * BLOCK_N + offs_n + k = tl.load(K + base + keys[:, None] * D + offs_d[None, :]) + s = tl.dot(q, tl.trans(k)) + if HAS_BIAS: # per-key score correction (Q smoothing: mu_q . k_j, high precision) + s = s + tl.load(bias_row + keys)[None, :] + ok = _key_ok(KV_VALID, n_block, offs_n, HAS_VALID, BLOCK_N // 2) + s = tl.where(ok[None, :], s, float("-inf")) + group_max = tl.max(tl.reshape(s, (ROWS, BLOCK_N // 16, 16)), axis=2) + if FLOOR: + group_max = tl.where(it == 0, tl.maximum(group_max, FLOOR_VALUE), group_max) + m_new = tl.maximum(m_i, tl.max(group_max, axis=1)) + m_scaled = m_new * scale_log2 + OFFSET + alpha = tl.math.exp2((m_i - m_new) * scale_log2) + if QUANT: + p_tilde, p_hat = _quantized_p(s, m_scaled, it == 0, scale_log2, FLOOR, FLOOR_VALUE, ROWS, BLOCK_N) + else: + p_tilde = tl.math.exp2(s * scale_log2 - m_scaled[:, None]) + p_hat = p_tilde + l_i = l_i * alpha + tl.sum(p_tilde, axis=1) + v = tl.load(V + base + keys[:, None] * D + offs_d[None, :]) + acc = acc * alpha[:, None] + tl.dot(p_hat.to(tl.bfloat16), v) + if HIGH_PREC_O: + acc_hp = acc_hp * alpha[:, None] + tl.dot(p_tilde.to(tl.bfloat16), v) + if SAVE_MRUN: + tl.store(MRUN + (meta * list_width + slot) * BLOCK_N + row0 + offs_r, m_new) + m_i = m_new + inv = tl.where((l_i == 0.0) | (l_i != l_i), 0.0, 1.0 / l_i) + tl.store(O + base + rows[:, None] * D + offs_d[None, :], (acc * inv[:, None]).to(O.dtype.element_ty)) + if HIGH_PREC_O: + tl.store(OHP + base + rows[:, None] * D + offs_d[None, :], (acc_hp * inv[:, None]).to(OHP.dtype.element_ty)) + tl.store(LSE + bh * seq_len + rows, m_i * scale_log2 + OFFSET + tl.math.log2(l_i)) + + +@triton.jit +def _bwd_dq_kernel(Q, K, V, DO, LSE, DELTA, DQ, Q2K_IDX, Q2K_NUM, KV_VALID, seq_len, n_qblocks, list_width, + scale_log2, sm_scale, HAS_VALID: tl.constexpr, D: tl.constexpr, ROWS: tl.constexpr, + BLOCK_N: tl.constexpr): + pid = tl.program_id(0) + bh = tl.program_id(1).to(tl.int64) + halves: tl.constexpr = BLOCK_N // ROWS + m_block = pid // halves + offs_r = tl.arange(0, ROWS) + offs_c = tl.arange(0, ROWS) # key columns per step: one 64-key half + offs_d = tl.arange(0, D) + base = bh * seq_len * D + rows = m_block * BLOCK_N + (pid % halves) * ROWS + offs_r + q = tl.load(Q + base + rows[:, None] * D + offs_d[None, :]) + do = tl.load(DO + base + rows[:, None] * D + offs_d[None, :]) + lse = tl.load(LSE + bh * seq_len + rows) + delta = tl.load(DELTA + bh * seq_len + rows) + meta = bh * n_qblocks + m_block + count = tl.load(Q2K_NUM + meta) + dq = tl.zeros((ROWS, D), tl.float32) + for it in range(0, count * halves): + n_block = tl.load(Q2K_IDX + meta * list_width + it // halves) + cols = (it % halves) * ROWS + offs_c + keys = n_block * BLOCK_N + cols + k = tl.load(K + base + keys[:, None] * D + offs_d[None, :]) + v = tl.load(V + base + keys[:, None] * D + offs_d[None, :]) + s = tl.dot(q, tl.trans(k)) + ok = _key_ok(KV_VALID, n_block, cols, HAS_VALID, BLOCK_N // 2) + p = tl.where(ok[None, :], tl.math.exp2(s * scale_log2 - lse[:, None]), 0.0) + dp = tl.dot(do, tl.trans(v)) + ds = p * (dp - delta[:, None]) + dq += tl.dot(ds.to(tl.bfloat16), k) + tl.store(DQ + base + rows[:, None] * D + offs_d[None, :], (dq * sm_scale).to(DQ.dtype.element_ty)) + + +@triton.jit +def _bwd_dkdv_kernel(Q, K, V, DO, LSE, DELTA, MRUN, DK, DV, K2Q_IDX, K2Q_SLOT, K2Q_NUM, Q2K_NUM, KV_VALID, seq_len, + n_qblocks, n_kblocks, list_width, scale_log2, sm_scale, OFFSET: tl.constexpr, + QUANT: tl.constexpr, HAS_VALID: tl.constexpr, FLOOR: tl.constexpr, + FLOOR_VALUE: tl.constexpr, D: tl.constexpr, ROWS: tl.constexpr, BLOCK_N: tl.constexpr): + pid = tl.program_id(0) + bh = tl.program_id(1).to(tl.int64) + halves: tl.constexpr = BLOCK_N // ROWS + n_block = pid // halves + offs_r = tl.arange(0, ROWS) + offs_d = tl.arange(0, D) + cols = (pid % halves) * ROWS + tl.arange(0, ROWS) + base = bh * seq_len * D + keys = n_block * BLOCK_N + cols + k = tl.load(K + base + keys[:, None] * D + offs_d[None, :]) + v = tl.load(V + base + keys[:, None] * D + offs_d[None, :]) + ok = _key_ok(KV_VALID, n_block, cols, HAS_VALID, BLOCK_N // 2) + kmeta = bh * n_kblocks + n_block + count = tl.load(K2Q_NUM + kmeta) + dk = tl.zeros((ROWS, D), tl.float32) + dv = tl.zeros((ROWS, D), tl.float32) + for it in range(0, count * halves): + entry = kmeta * n_qblocks + it // halves + m_block = tl.load(K2Q_IDX + entry) + slot = tl.load(K2Q_SLOT + entry) + row0 = (it % halves) * ROWS + rows = m_block * BLOCK_N + row0 + offs_r + q = tl.load(Q + base + rows[:, None] * D + offs_d[None, :]) + do = tl.load(DO + base + rows[:, None] * D + offs_d[None, :]) + lse = tl.load(LSE + bh * seq_len + rows) + delta = tl.load(DELTA + bh * seq_len + rows) + s = tl.dot(q, tl.trans(k)) + s = tl.where(ok[None, :], s, float("-inf")) + p = tl.math.exp2(s * scale_log2 - lse[:, None]) + if QUANT: + qmeta = bh * n_qblocks + m_block + first_visit = slot == tl.load(Q2K_NUM + qmeta) - 1 + m_run = tl.load(MRUN + (qmeta * list_width + slot) * BLOCK_N + row0 + offs_r) + m_scaled = m_run * scale_log2 + OFFSET + _, p_hat = _quantized_p(s, m_scaled, first_visit, scale_log2, FLOOR, FLOOR_VALUE, ROWS, ROWS) + weight = p_hat * tl.math.exp2(m_scaled - lse)[:, None] + else: + weight = p + dv += tl.dot(tl.trans(weight.to(tl.bfloat16)), do) + dp = tl.dot(do, tl.trans(v)) + ds = p * (dp - delta[:, None]) + dk += tl.dot(tl.trans(ds.to(tl.bfloat16)), q) + tl.store(DK + base + keys[:, None] * D + offs_d[None, :], (dk * sm_scale).to(DK.dtype.element_ty)) + tl.store(DV + base + keys[:, None] * D + offs_d[None, :], dv.to(DV.dtype.element_ty)) + + +class _FP4VSAAttnQAT(torch.autograd.Function): + + @staticmethod + def forward(ctx, q, k, v, q2k_idx, q2k_num, kv_valid, quantize, quantize_pv, two_level_p, first_block_max_floor, + sm_scale, high_prec_o, key_bias=None, q_group=None, need_backward=False): # type: ignore[override] + batch, heads, seq_len, dim = q.shape + if seq_len % BLOCK or dim not in (64, 128): + raise ValueError(f"need L % {BLOCK} == 0 and D in (64, 128), got L={seq_len}, D={dim}") + n_blocks = seq_len // BLOCK + if q2k_idx.shape[:3] != (batch, heads, n_blocks): + raise ValueError(f"q2k_idx {tuple(q2k_idx.shape)} does not match {(batch, heads, n_blocks)}") + if quantize: + q_hat, k_hat, v_hat = fake_quant_qkv_triton(q, k, v, quantize_v=quantize_pv) + else: + q_hat, k_hat, v_hat = (t.to(torch.bfloat16).contiguous() for t in (q, k, v)) + q2k_idx = q2k_idx.to(torch.int32).contiguous() + q2k_num = q2k_num.to(torch.int32).contiguous() + has_valid = kv_valid is not None + kv_valid_arg = kv_valid.to(torch.int32).contiguous() if has_valid else q2k_num + width = q2k_idx.shape[-1] + out = torch.empty_like(q_hat) + out_hp = torch.empty_like(q_hat) if high_prec_o else out + lse = torch.empty((batch, heads, seq_len), device=q.device, dtype=torch.float32) + scale_log2 = sm_scale * LOG2E + offset = LOG2_INV_TWO_LEVEL if two_level_p else 0.0 + floor = first_block_max_floor is not None + # Only the P quantizer reads the kernel's QUANT path; QK-only mode runs it with exact softmax. + quantize = bool(quantize and quantize_pv) + # dV only reads these maxima for quantized P. Inference and QK-only + # backward need no quadratic buffer; LSE is an unused pointer there. + save_mrun = bool(need_backward and quantize) + m_run = (torch.empty((batch, heads, n_blocks, width, BLOCK), device=q.device, dtype=torch.float32) + if save_mrun else lse) + has_bias = key_bias is not None + if has_bias: + key_bias = key_bias.float().contiguous() + q_group = (q_group if q_group is not None else torch.zeros(n_blocks, dtype=torch.int32, + device=q.device)).to(torch.int32).contiguous() + grid = (n_blocks * (BLOCK // _ROWS), batch * heads) + _fwd_kernel[grid](q_hat, k_hat, v_hat, out, out_hp, lse, m_run, q2k_idx, q2k_num, kv_valid_arg, + key_bias if has_bias else lse, q_group if has_bias else q2k_num, + key_bias.shape[2] if has_bias else 1, seq_len, n_blocks, width, scale_log2, OFFSET=offset, + QUANT=bool(quantize), HAS_VALID=has_valid, HAS_BIAS=has_bias, FLOOR=floor, FLOOR_VALUE=float(first_block_max_floor or 0.0), + HIGH_PREC_O=bool(high_prec_o), SAVE_MRUN=save_mrun, D=dim, ROWS=_ROWS, BLOCK_N=BLOCK, num_warps=_WARPS, + num_stages=2) + if need_backward: + ctx.save_for_backward(q_hat, k_hat, v_hat, out_hp, lse, m_run, q2k_idx, q2k_num, kv_valid_arg) + ctx.options = (bool(quantize), has_valid, floor, float(first_block_max_floor or 0.0), offset, sm_scale) + ctx.has_bias = has_bias + return out + + @staticmethod + def backward(ctx, grad_out): # type: ignore[override] + if ctx.has_bias: + raise NotImplementedError("the Q-smoothing score bias is an eval-only emulator option (no backward)") + q_hat, k_hat, v_hat, out_hp, lse, m_run, q2k_idx, q2k_num, kv_valid = ctx.saved_tensors + quantize, has_valid, floor, floor_value, offset, sm_scale = ctx.options + batch, heads, seq_len, dim = q_hat.shape + n_blocks = seq_len // BLOCK + width = q2k_idx.shape[-1] + grad_out = grad_out.to(torch.bfloat16).contiguous() + delta = (grad_out.float() * out_hp.float()).sum(-1).contiguous() + k2q_idx, k2q_slot, k2q_num = invert_block_lists(q2k_idx, q2k_num, n_blocks) + dq = torch.empty_like(q_hat) + dk = torch.empty_like(k_hat) + dv = torch.empty_like(v_hat) + scale_log2 = sm_scale * LOG2E + grid = (n_blocks * (BLOCK // _ROWS), batch * heads) + _bwd_dq_kernel[grid](q_hat, k_hat, v_hat, grad_out, lse, delta, dq, q2k_idx, q2k_num, kv_valid, seq_len, + n_blocks, width, scale_log2, sm_scale, HAS_VALID=has_valid, D=dim, ROWS=_ROWS, + BLOCK_N=BLOCK, num_warps=_WARPS, num_stages=2) + _bwd_dkdv_kernel[grid](q_hat, k_hat, v_hat, grad_out, lse, delta, m_run, dk, dv, k2q_idx, k2q_slot, k2q_num, + q2k_num, kv_valid, seq_len, n_blocks, n_blocks, width, scale_log2, sm_scale, + OFFSET=offset, QUANT=quantize, HAS_VALID=has_valid, FLOOR=floor, + FLOOR_VALUE=floor_value, D=dim, ROWS=_ROWS, BLOCK_N=BLOCK, num_warps=_WARPS, + num_stages=2) + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None + + +def _hadamard(dim: int, device: torch.device) -> torch.Tensor: + """Normalized Sylvester Hadamard matrix ``[dim, dim]`` (orthogonal, symmetric), fp32.""" + h = torch.ones(1, 1, device=device) + while h.shape[0] < dim: + h = torch.cat([torch.cat([h, h], 1), torch.cat([h, -h], 1)], 0) + if h.shape[0] != dim: + raise ValueError(f"Hadamard rotation needs a power-of-two head dim, got {dim}") + return h / math.sqrt(dim) + + +def fp8_per_token_fake_quant(x: torch.Tensor) -> torch.Tensor: + """E4M3 fake quant with one fp32 scale per token (row), ``amax / 448``; fp32 out.""" + xf = x.float() + scale = xf.abs().amax(dim=-1, keepdim=True) / 448.0 + inv = torch.where(scale == 0, torch.zeros_like(scale), 1.0 / scale) + return (xf * inv).clamp(-448.0, 448.0).to(torch.float8_e4m3fn).float() * scale + + +def mxfp8_fake_quant(x: torch.Tensor) -> torch.Tensor: + """MXFP8 fake quant along the last dim: E4M3 values, one E8M0 (power-of-two) scale per 32; fp32 out. + + The scale is ``2**ceil(log2(amax / 448))``, bit-identical to ``flashinfer.mxfp8_quantize`` (the FA4 + MXFP8 Q/K operands) on sm_100. + """ + if x.shape[-1] % 32: + raise ValueError(f"last dim {x.shape[-1]} is not a multiple of 32") + xf = x.float().unflatten(-1, (-1, 32)) + amax = xf.abs().amax(dim=-1, keepdim=True).clamp_min(1e-30) + scale = torch.exp2(torch.ceil(torch.log2(amax / 448.0)).clamp(-127.0, 127.0)) + return ((xf / scale).clamp(-448.0, 448.0).to(torch.float8_e4m3fn).float() * scale).flatten(-2) + + +def _query_groups(n_blocks: int, num_prefix_tiles: int | None, segmented: bool, device: torch.device) -> torch.Tensor: + """Per query tile: 0 for dense-prefix tiles, 1 for video tiles (all 0 unless ``segmented``).""" + groups = torch.zeros(n_blocks, dtype=torch.int32, device=device) + if segmented: + groups[int(num_prefix_tiles or 0):] = 1 + return groups + + +def _smooth_queries(q: torch.Tensor, k: torch.Tensor, groups: torch.Tensor, + n_groups: int) -> tuple[torch.Tensor, torch.Tensor]: + """``q - mu_g`` per query group and the exact correction ``bias[g, j] = mu_g . k_j`` (``[B, H, G, L]``).""" + row_group = groups.repeat_interleave(BLOCK).long() + qf = q.float() + mus = [] + for g in range(n_groups): + rows = row_group == g + mus.append(qf[:, :, rows].mean(dim=-2) if bool(rows.any()) else qf.new_zeros(qf.shape[:2] + qf.shape[-1:])) + mu = torch.stack(mus, dim=2) # [B, H, G, D] + bias = torch.einsum("bhgd,bhld->bhgl", mu, k.float()) + return qf - mu[:, :, row_group], bias + + +def fp4_vsa_attn_qat(q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, + kv_valid: torch.Tensor | None = None, + *, + quantize: bool = True, + quantize_pv: bool = True, + two_level_p: bool = False, + smooth_k: bool = False, + first_block_max_floor: float | None = None, + sm_scale: float | None = None, + high_prec_o: bool = True, + hadamard: bool = False, + smooth_q: str | None = None, + prefix_precision: str | None = None, + num_prefix_tiles: int | None = None) -> torch.Tensor: + """Block-sparse attention with the deployed FP4 numerics in the forward and STE gradients. + + ``q``/``k``/``v`` are ``[B, H, L, D]`` (L a multiple of 128, tile-ordered, + pad rows zero); the lists are ``vsa_tile_mask_to_fp4_blocks(mask, 128, vbs)`` + outputs (``q2k_quad`` must be None: 128-token tiles map 1:1 onto blocks). + ``smooth_k`` subtracts each head's token mean from K first (exact for + softmax, differentiable through autograd). ``quantize_pv=False`` quantizes + Q/K only and keeps P and V in BF16 (the FA4 NVFP4-QK + BF16-PV kernel). + Returns ``[B, H, L, D]`` BF16. + + Q/K-only error-reduction options (eval studies; the last three are forward-only): + + * ``hadamard``: rotate Q and K by a normalized 128x128 Hadamard matrix + (after K smoothing) before NVFP4 quantization. Exact in real arithmetic. + * ``smooth_q``: ``"global"`` subtracts each head's mean query, ``"segment"`` + one mean for the dense-prefix query tiles and one for the video tiles; + the exact correction ``mu_q . k_j`` is added to every score as a + high-precision per-key bias. + * ``prefix_precision``: ``"bf16"``, ``"fp8"`` (E4M3, per-token scales) or + ``"mxfp8"`` (E4M3, power-of-two scale per 32 channels, the FA4 kernel's + MXFP8 operands) Q/K for the first ``num_prefix_tiles`` query tiles (text, + reference and audio rows), as a second call over the same lists; video + rows stay FP4. + """ + qk_only = quantize and not quantize_pv + if qk_only and (two_level_p or first_block_max_floor is not None): + raise ValueError("two_level_p and first_block_max_floor shape the P quantizer; QK-only mode has none") + # ``hadamard`` and ``smooth_q`` rewrite Q/K, which only the QK-only route reads. + # ``prefix_precision`` also works with FP4 P.V, so the FP4 call below takes the + # caller's P quantizer settings rather than assuming the QK-only ones. + if (hadamard or smooth_q) and not qk_only: + raise ValueError("hadamard / smooth_q are Q/K-only (quantize_pv=False) options") + if smooth_q not in (None, "global", "segment"): + raise ValueError(f"smooth_q must be None, 'global' or 'segment', got {smooth_q!r}") + if prefix_precision not in (None, "bf16", "fp8", "mxfp8"): + raise ValueError(f"prefix_precision must be None, 'bf16', 'fp8' or 'mxfp8', got {prefix_precision!r}") + if (smooth_q == "segment" or prefix_precision) and num_prefix_tiles is None: + raise ValueError("smooth_q='segment' and prefix_precision need num_prefix_tiles") + if smooth_k: + k = k - k.mean(dim=-2, keepdim=True) + scale = sm_scale if sm_scale is not None else 1.0 / math.sqrt(q.shape[-1]) + # Autograd.Function.forward runs with grad disabled even during training. + # Capture the caller's grad mode here, including no_grad on grad-enabled inputs. + need_backward = torch.is_grad_enabled() and any(t.requires_grad for t in (q, k, v)) + if not (hadamard or smooth_q or prefix_precision): + return _FP4VSAAttnQAT.apply(q, k, v, q2k_idx, q2k_num, kv_valid, quantize, quantize_pv, two_level_p, + first_block_max_floor, scale, high_prec_o, None, None, need_backward) + n_blocks = q.shape[2] // BLOCK + hp_tiles = min(int(num_prefix_tiles or 0), n_blocks) if prefix_precision else 0 + num_fp4 = q2k_num + if hp_tiles: + tile = torch.arange(n_blocks, device=q2k_num.device) + num_fp4 = q2k_num.masked_fill(tile < hp_tiles, 0) + q_lo, k_lo = q, k + if hadamard: + rot = _hadamard(q.shape[-1], q.device) + q_lo, k_lo = q.float() @ rot, k.float() @ rot + bias = groups = None + if smooth_q: + n_groups = 2 if smooth_q == "segment" else 1 + groups = _query_groups(n_blocks, num_prefix_tiles, smooth_q == "segment", q.device) + q_lo, bias = _smooth_queries(q_lo, k_lo, groups, n_groups) + out = _FP4VSAAttnQAT.apply(q_lo, k_lo, v, q2k_idx, num_fp4, kv_valid, quantize, quantize_pv, two_level_p, + first_block_max_floor, scale, high_prec_o, bias, groups, need_backward) + if not hp_tiles: + return out + num_hp = q2k_num - num_fp4 + hp_quant = {"bf16": lambda t: t, "fp8": fp8_per_token_fake_quant, "mxfp8": mxfp8_fake_quant}[prefix_precision] + q_hp, k_hp = hp_quant(q), hp_quant(k) + out_hp = _FP4VSAAttnQAT.apply(q_hp, k_hp, v, q2k_idx, num_hp, kv_valid, False, True, False, None, scale, + high_prec_o, None, None, need_backward) + return out + out_hp + + +__all__ = ["fake_quant_qkv_triton", "fp4_vsa_attn_qat", "fp8_per_token_fake_quant", "mxfp8_fake_quant"] diff --git a/fastvideo-kernel/tests/test_attn_qat_infer_sparse.py b/fastvideo-kernel/tests/test_attn_qat_infer_sparse.py index 3684b2156..2c4b21917 100644 --- a/fastvideo-kernel/tests/test_attn_qat_infer_sparse.py +++ b/fastvideo-kernel/tests/test_attn_qat_infer_sparse.py @@ -133,6 +133,21 @@ def test_bshd_entry_matches_bhsd() -> None: assert torch.equal(bhsd, bshd) +def test_bshd_entry_accepts_video_only_queries_over_the_full_keys() -> None: + # The FP4 prefix split attends the prefix query tiles densely and sends only the video + # query blocks here, with the keys still spanning the whole sequence. The key length, + # not the query length, bounds the block lists and sizes delta_s. + q, k, v, sizes, mask, _ = _case(128, (256, 100), 6, {2: 50}) + idx, num, kv_valid, quad = vsa_tile_mask_to_fp4_blocks(mask, 128, sizes) + head = 3 # prefix tiles: two full chunks of the 256-token segment and one 100-token chunk + bshd = tuple(x.transpose(1, 2).contiguous() for x in (q, k, v)) + full = sageattn_blackwell_sparse_bshd(*bshd, idx, num, kv_valid, quad) + video = sageattn_blackwell_sparse_bshd(bshd[0][:, head * BLOCK_M:], bshd[1], bshd[2], + idx[:, :, head:].contiguous(), num[:, :, head:].contiguous(), kv_valid, + None if quad is None else quad[:, :, head:].contiguous()) + assert torch.equal(video, full[:, :, head * BLOCK_M:]) + + def test_first_visited_block_check_rejects_unanchored_rows() -> None: # Without the exempt prefix a query block may start from a block that masks # one of its 64-row halves entirely; validate=True must refuse it. diff --git a/fastvideo-kernel/tests/test_attn_qat_vsa_train.py b/fastvideo-kernel/tests/test_attn_qat_vsa_train.py new file mode 100644 index 000000000..dfb7491e6 --- /dev/null +++ b/fastvideo-kernel/tests/test_attn_qat_vsa_train.py @@ -0,0 +1,303 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Block-sparse FP4 Attn-QAT (VSA tile 128): the PyTorch specification and the Triton kernel. + +CPU tests pin the reference emulator (E2M1/E4M3 rounding, list builder, online +softmax, dV weights). GPU tests (any CUDA GPU with Triton; run on GB200) +compare the Triton forward and STE backward with the reference. +""" +from __future__ import annotations + +import importlib +import importlib.util +import math +import sys +from pathlib import Path + +import pytest +import torch + +_DIR = Path(__file__).resolve().parents[1] / "python" / "fastvideo_kernel" / "triton_kernels" +_PKG = "_test_qat_vsa_kernels" + + +def _module(name: str): + if _PKG not in sys.modules: + spec = importlib.util.spec_from_loader(_PKG, loader=None, is_package=True) + spec.submodule_search_locations.append(str(_DIR)) + sys.modules[_PKG] = importlib.util.module_from_spec(spec) + return importlib.import_module(f"{_PKG}.{name}") + + +ref = _module("attn_qat_vsa_reference") +needs_gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton kernel needs a CUDA GPU") + + +def _problem(n_tiles: int, heads: int = 2, keep: float = 0.4, seed: int = 0, device: str = "cpu", + partial_tail: bool = True): + """Random tile-ordered Q/K/V with a VSA-like mask: block 0 always kept, tail tile partially valid.""" + gen = torch.Generator().manual_seed(seed) + seq_len = n_tiles * ref.BLOCK + tile_valid = torch.full((n_tiles, ), ref.BLOCK, dtype=torch.int32) + if partial_tail: + tile_valid[-1] = 37 + tile_valid[n_tiles // 2] = 100 + token_valid = (torch.arange(ref.BLOCK)[None, :] < tile_valid[:, None]).reshape(-1) + q, k, v = (torch.randn(1, heads, seq_len, 128, generator=gen) * token_valid[None, None, :, None] + for _ in range(3)) + # RMS-normed Q/K have per-channel structure; add a channel bias so smoothing matters. + k = k + 0.5 * torch.randn(1, heads, 1, 128, generator=gen) * token_valid[None, None, :, None] + mask = torch.rand(1, heads, n_tiles, n_tiles, generator=gen) < keep + mask[..., 0] = True + q2k_idx, q2k_num, kv_valid = ref.tile128_mask_to_fp4_blocks(mask, tile_valid) + move = lambda t: t.to(device) # noqa: E731 + return (move(q.bfloat16()), move(k.bfloat16()), move(v.bfloat16()), move(q2k_idx), move(q2k_num), + move(kv_valid), move(mask), move(tile_valid)) + + +def test_e2m1_grid_and_ties_round_to_even(): + values = torch.tensor([0.0, 0.25, 0.26, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0, 5.1, 7.0, -0.75, -5.0]) + expected = torch.tensor([0.0, 0.0, 0.5, 1.0, 1.0, 2.0, 2.0, 4.0, 4.0, 6.0, 6.0, -1.0, -4.0]) + assert torch.equal(ref.e2m1_rn(values), expected) + + +def test_e4m3_rounding_saturates_and_flushes(): + values = torch.tensor([448.0, 1000.0, 1.0625, 2.0**-10, 2.0**-9, 0.0]) + out = ref.e4m3_rn(values) + assert out[0] == 448 and out[1] == 448 + assert out[2] == 1.0 # tie between 1.0 and 1.125 rounds to even + assert out[3] == 0.0 and out[4] == 2.0**-9 and out[5] == 0.0 + + +def test_fake_quant_matches_nvfp4_definition(): + x = torch.tensor([[0.0] * 15 + [6.0], [1.0] * 16]).float() + out = ref.nvfp4_fake_quant_lastdim(x) + assert torch.equal(out[0], x[0]) # amax 6 -> scale exactly 1 + # amax 1 -> scale e4m3(1/6) = 0.171875, 1 / 0.171875 = 5.82 -> E2M1 6 -> 1.03125 + assert torch.equal(out[1], torch.full((16, ), 1.03125)) + assert ref.nvfp4_fake_quant_lastdim(torch.zeros(1, 16)).abs().sum() == 0 + + +def test_list_builder_is_descending_and_counts_valid_tiles(): + _, _, _, q2k_idx, q2k_num, kv_valid, mask, tile_valid = _problem(6) + for h in range(mask.shape[1]): + for m in range(6): + count = int(q2k_num[0, h, m]) + listed = q2k_idx[0, h, m, :count].tolist() + assert listed == sorted(listed, reverse=True) + assert listed == sorted(torch.nonzero(mask[0, h, m]).flatten().tolist(), reverse=True) + assert listed[-1] == 0 + assert kv_valid.tolist()[-2:] == [37, 0] + assert kv_valid.numel() == 2 * tile_valid.numel() + + +def test_invert_block_lists_round_trips(): + _, _, _, q2k_idx, q2k_num, _, mask, _ = _problem(5, heads=3, seed=3) + k2q_idx, k2q_slot, k2q_num = ref.invert_block_lists(q2k_idx, q2k_num, 5) + for h in range(3): + for n in range(5): + entries = [(int(k2q_idx[0, h, n, i]), int(k2q_slot[0, h, n, i])) for i in range(int(k2q_num[0, h, n]))] + assert [m for m, _ in entries] == torch.nonzero(mask[0, h, :, n]).flatten().tolist() + for m, slot in entries: + assert int(q2k_idx[0, h, m, slot]) == n + + +def test_reference_without_quantization_is_masked_softmax(): + q, k, v, q2k_idx, q2k_num, kv_valid, _, _ = _problem(4) + out = ref.fp4_vsa_attention_reference(q, k, v, q2k_idx, q2k_num, kv_valid, quantize=False) + token_mask = ref.dense_block_mask(q2k_idx, q2k_num, q.shape[2], kv_valid) + scores = (q.float() @ k.float().transpose(-1, -2)) / math.sqrt(128) + expected = torch.softmax(scores.masked_fill(~token_mask, -math.inf), dim=-1) @ v.float() + torch.testing.assert_close(out, expected, rtol=1e-4, atol=1e-5) + + +@pytest.mark.parametrize("two_level_p", [False, True]) +def test_reference_dv_weights_reproduce_output(two_level_p): + q, k, v, q2k_idx, q2k_num, kv_valid, _, _ = _problem(3, seed=1) + out, weights = ref.fp4_vsa_attention_reference(q, k, v, q2k_idx, q2k_num, kv_valid, two_level_p=two_level_p, + return_dv_weights=True) + _, _, v_hat = ref.fake_quant_qkv(q, k, v) + torch.testing.assert_close(weights @ v_hat, out, rtol=1e-4, atol=1e-5) + + +def test_two_level_p_is_closer_to_exact_than_single_level(): + q, k, v, q2k_idx, q2k_num, kv_valid, _, _ = _problem(6, seed=2, keep=0.8) + exact = ref.fp4_vsa_attention_reference(q, k, v, q2k_idx, q2k_num, kv_valid, quantize=False) + errors = {} + for two_level in (False, True): + out = ref.fp4_vsa_attention_reference(q, k, v, q2k_idx, q2k_num, kv_valid, two_level_p=two_level) + errors[two_level] = float((out - exact).norm() / exact.norm()) + assert errors[True] < errors[False] + + +# ----------------------------------------------------------------------------- Triton (GPU) +def _triton(): + return _module("attn_qat_vsa_train") + + +def _rel(a: torch.Tensor, b: torch.Tensor) -> float: + return float((a.float() - b.float()).norm() / b.float().norm().clamp_min(1e-12)) + + +@needs_gpu +@pytest.mark.parametrize("prefix_precision", [None, "fp8"]) +def test_inference_matches_training_without_running_maxima(monkeypatch, prefix_precision): + train = _triton() + q, k, v, idx, counts, valid, _, _ = _problem(4, heads=2, device="cuda") + q.requires_grad_(True) + options = dict(prefix_precision=prefix_precision, num_prefix_tiles=1 if prefix_precision else None) + expected = train.fp4_vsa_attn_qat(q, k, v, idx, counts, valid, **options).detach() + original = torch.empty + + def allocate(shape, **kwargs): + assert len(shape) != 5, "inference must not allocate per-block running maxima" + return original(shape, **kwargs) + + monkeypatch.setattr(torch, "empty", allocate) + with torch.no_grad(): + actual = train.fp4_vsa_attn_qat(q, k, v, idx, counts, valid, **options) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@needs_gpu +@pytest.mark.parametrize("two_level_p,smooth_k", [(False, False), (True, False), (False, True), (True, True)]) +def test_triton_forward_matches_reference(two_level_p, smooth_k): + q, k, v, q2k_idx, q2k_num, kv_valid, _, _ = _problem(6, heads=2, seed=4, device="cuda") + expected = ref.fp4_vsa_attention_reference(q, k, v, q2k_idx, q2k_num, kv_valid, two_level_p=two_level_p, + smooth_k=smooth_k) + out = _triton().fp4_vsa_attn_qat(q, k, v, q2k_idx, q2k_num, kv_valid, two_level_p=two_level_p, + smooth_k=smooth_k) + exact = ref.fp4_vsa_attention_reference(q, k, v, q2k_idx, q2k_num, kv_valid, quantize=False) + # The emulator must sit far closer to the reference than FP4 sits to exact attention. + assert _rel(out, expected) < 2e-3 < 0.1 * _rel(expected, exact) + + +@needs_gpu +def test_triton_forward_unquantized_matches_softmax(): + q, k, v, q2k_idx, q2k_num, kv_valid, _, _ = _problem(5, seed=5, device="cuda") + expected = ref.fp4_vsa_attention_reference(q, k, v, q2k_idx, q2k_num, kv_valid, quantize=False) + out = _triton().fp4_vsa_attn_qat(q, k, v, q2k_idx, q2k_num, kv_valid, quantize=False) + assert _rel(out, expected) < 5e-3 + + +@needs_gpu +def test_triton_first_block_floor_matches_reference(): + q, k, v, q2k_idx, q2k_num, kv_valid, _, _ = _problem(4, seed=6, device="cuda") + expected = ref.fp4_vsa_attention_reference(q, k, v, q2k_idx, q2k_num, kv_valid, first_block_max_floor=0.0) + out = _triton().fp4_vsa_attn_qat(q, k, v, q2k_idx, q2k_num, kv_valid, first_block_max_floor=0.0) + assert _rel(out, expected) < 2e-3 + + +@needs_gpu +@pytest.mark.parametrize("two_level_p", [False, True]) +def test_triton_backward_matches_ste_reference(two_level_p): + q, k, v, q2k_idx, q2k_num, kv_valid, _, _ = _problem(5, heads=2, seed=7, device="cuda") + q, k, v = (t.clone().requires_grad_(True) for t in (q, k, v)) + out = _triton().fp4_vsa_attn_qat(q, k, v, q2k_idx, q2k_num, kv_valid, two_level_p=two_level_p) + grad_out = torch.randn_like(out) + dq, dk, dv = torch.autograd.grad(out, (q, k, v), grad_out) + _, weights = ref.fp4_vsa_attention_reference(q, k, v, q2k_idx, q2k_num, kv_valid, two_level_p=two_level_p, + return_dv_weights=True) + q_hat, k_hat, v_hat = ref.fake_quant_qkv(q.detach(), k.detach(), v.detach()) + token_mask = ref.dense_block_mask(q2k_idx, q2k_num, q.shape[2], kv_valid) + expected = ref.ste_gradients_reference(q_hat, k_hat, v_hat, token_mask, weights, grad_out, 1 / math.sqrt(128)) + for name, got, want in zip("qkv", (dq, dk, dv), expected, strict=True): + assert _rel(got, want) < 2e-2, name + + +@needs_gpu +def test_triton_qk_only_is_exact_softmax_on_quantized_qk(): + """QK-only mode = exact sparse softmax over NVFP4 Q/K with BF16 V; STE gradients reach Q, K and V.""" + q, k, v, q2k_idx, q2k_num, kv_valid, _, _ = _problem(5, heads=2, seed=10, device="cuda") + q, k, v = (t.clone().requires_grad_(True) for t in (q, k, v)) + out = _triton().fp4_vsa_attn_qat(q, k, v, q2k_idx, q2k_num, kv_valid, quantize_pv=False, smooth_k=True) + k_smooth = (k - k.mean(dim=-2, keepdim=True)).detach() + q_hat, k_hat, _ = ref.fake_quant_qkv(q.detach(), k_smooth, v.detach()) + expected = ref.fp4_vsa_attention_reference(q_hat.bfloat16(), k_hat.bfloat16(), v.detach(), q2k_idx, q2k_num, + kv_valid, quantize=False) + full = _triton().fp4_vsa_attn_qat(q, k, v, q2k_idx, q2k_num, kv_valid, smooth_k=True) + assert _rel(out, expected) < 5e-3 < _rel(full, expected) + grads = torch.autograd.grad(out, (q, k, v), torch.randn_like(out)) + assert all(bool(g.float().norm() > 0) for g in grads) + with pytest.raises(ValueError): + _triton().fp4_vsa_attn_qat(q, k, v, q2k_idx, q2k_num, kv_valid, quantize_pv=False, two_level_p=True) + + +@needs_gpu +def test_triton_dv_reaches_every_head_on_long_sequences(): + """Regression for the GB10 Attn-QAT bug: re-quantized normalized P underflowed and zeroed dV per head.""" + q, k, v, q2k_idx, q2k_num, kv_valid, _, _ = _problem(64, heads=4, keep=0.9, seed=8, device="cuda") + v = v.clone().requires_grad_(True) + out = _triton().fp4_vsa_attn_qat(q, k, v, q2k_idx, q2k_num, kv_valid) + (dv, ) = torch.autograd.grad(out, (v, ), torch.ones_like(out)) + per_head = dv.float().norm(dim=(-1, -2))[0] + assert bool((per_head > 0.1 * per_head.mean()).all()), per_head + + +@needs_gpu +def test_triton_fake_quant_is_bit_identical_to_reference(): + q, k, v, *_ = _problem(4, heads=3, seed=9, device="cuda") + q = q * 37.0 # exercise large and subnormal E4M3 scales + got = _triton().fake_quant_qkv_triton(q, k, v) + want = ref.fake_quant_qkv(q, k, v) + for name, a, b in zip("qkv", got, want, strict=True): + assert torch.equal(a.float(), b.to(torch.bfloat16).float()), name + + +def _bias_reference(q, k, v, token_mask, bias): + """fp32 softmax((q k^T + bias[key]) / sqrt(D)) v over ``token_mask``.""" + s = (q.float() @ k.float().transpose(-1, -2) + bias[..., None, :]) / math.sqrt(q.shape[-1]) + p = torch.softmax(s.masked_fill(~token_mask, float("-inf")), -1).nan_to_num(0.0) + return p @ v.float() + + +@needs_gpu +def test_triton_key_bias_matches_reference(): + """The eval-only per-key score bias (Q-smoothing correction) enters the scores before the softmax.""" + q, k, v, q2k_idx, q2k_num, kv_valid, _, _ = _problem(5, heads=2, seed=11, device="cuda") + bias = 3.0 * torch.randn(1, 2, 1, q.shape[2], device="cuda") + out = _triton()._FP4VSAAttnQAT.apply(q, k, v, q2k_idx, q2k_num, kv_valid, False, True, False, None, + 1 / math.sqrt(128), True, bias, None) + token_mask = ref.dense_block_mask(q2k_idx, q2k_num, q.shape[2], kv_valid) + assert _rel(out, _bias_reference(q, k, v, token_mask, bias[:, :, 0])) < 5e-3 + + +@pytest.mark.parametrize("n_groups", [1, 2]) +def test_query_smoothing_correction_is_exact(n_groups): + q, k, *_ = _problem(4, heads=2, seed=12) + groups = torch.tensor([0, 0, 1, 1], dtype=torch.int32) if n_groups == 2 else torch.zeros(4, dtype=torch.int32) + q_c, bias = _triton()._smooth_queries(q, k, groups, n_groups) + row_group = groups.repeat_interleave(ref.BLOCK).long() + rebuilt = q_c @ k.float().transpose(-1, -2) + bias[:, :, row_group] + assert _rel(rebuilt, q.float() @ k.float().transpose(-1, -2)) < 1e-5 + assert q_c[:, :, row_group == 0].mean(-2).abs().max() < 1e-5 + + +def test_hadamard_is_orthogonal(): + h = _triton()._hadamard(128, torch.device("cpu")) + assert torch.allclose(h @ h.T, torch.eye(128), atol=1e-6) + + +@needs_gpu +def test_triton_qk_variants_reduce_to_known_modes(): + """Hadamard = QK-only on rotated inputs; BF16 prefix rows = exact rows; all-prefix = no quantization.""" + q, k, v, q2k_idx, q2k_num, kv_valid, _, _ = _problem(6, heads=2, seed=13, device="cuda") + fn = _triton().fp4_vsa_attn_qat + qk = dict(quantize_pv=False, smooth_k=True) + rot = _triton()._hadamard(128, q.device) + k_s = k - k.mean(dim=-2, keepdim=True) + q_hat, k_hat, _ = ref.fake_quant_qkv(q.float() @ rot, k_s.float() @ rot, v) + expected = ref.fp4_vsa_attention_reference(q_hat.bfloat16(), k_hat.bfloat16(), v, q2k_idx, q2k_num, kv_valid, + quantize=False) + assert _rel(fn(q, k, v, q2k_idx, q2k_num, kv_valid, hadamard=True, **qk), expected) < 5e-3 + exact = fn(q, k, v, q2k_idx, q2k_num, kv_valid, quantize=False, smooth_k=True) + plain = fn(q, k, v, q2k_idx, q2k_num, kv_valid, **qk) + mixed = fn(q, k, v, q2k_idx, q2k_num, kv_valid, prefix_precision="bf16", num_prefix_tiles=2, **qk) + rows = 2 * ref.BLOCK + assert _rel(mixed[:, :, :rows], exact[:, :, :rows]) < 1e-3 + assert _rel(mixed[:, :, rows:], plain[:, :, rows:]) < 1e-3 + all_hp = fn(q, k, v, q2k_idx, q2k_num, kv_valid, prefix_precision="bf16", num_prefix_tiles=6, **qk) + assert _rel(all_hp, exact) < 1e-3 + smoothed = fn(q, k, v, q2k_idx, q2k_num, kv_valid, smooth_q="global", **qk) + assert _rel(smoothed, exact) < 0.5 + with pytest.raises(ValueError): + fn(q, k, v, q2k_idx, q2k_num, kv_valid, hadamard=True) # P.V quantized: not a Q/K-only mode diff --git a/fastvideo/attention/backends/fp4_vsa_qat.py b/fastvideo/attention/backends/fp4_vsa_qat.py new file mode 100644 index 000000000..db9cb6bb5 --- /dev/null +++ b/fastvideo/attention/backends/fp4_vsa_qat.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +"""VSA-H3 tile-128 fine attention with the deployed sparse-FP4 numerics, for quantization-aware training. + +``FP4VSAFineAttention`` plugs into ``MiniMaxH3VSAImpl.fine_attention_override``: +VSA-H3's own selection (tile pooling, top-k mask, gated compression branch) +runs unchanged and in BF16, and only the block-sparse attention is replaced by +``fastvideo_kernel.triton_kernels.attn_qat_vsa_train.fp4_vsa_attn_qat``, which +emulates ``attn_qat_infer``'s sparse FP4 kernel on any CUDA GPU (forward) and +provides straight-through gradients (backward). + +Training images may carry a prebuilt ``fastvideo_kernel`` that predates these +modules; ``load_qat_vsa_modules`` then loads them from this checkout's +``fastvideo-kernel/python`` tree under a private package name. +""" +from __future__ import annotations + +import importlib +import importlib.util +import sys +from dataclasses import dataclass +from pathlib import Path +from types import ModuleType +from typing import Any + +import torch + +_PRIVATE_PACKAGE = "_fastvideo_qat_vsa_kernels" +_KERNEL_DIR = Path( + __file__).resolve().parents[3] / "fastvideo-kernel" / "python" / "fastvideo_kernel" / "triton_kernels" + + +def _load(name: str) -> ModuleType: + try: + return importlib.import_module(f"fastvideo_kernel.triton_kernels.{name}") + except ImportError: + pass + except KeyError: + # An earlier failed `fastvideo_kernel` import leaves a None entry in + # sys.modules, and the machinery then raises KeyError instead of + # ImportError. Same meaning for us: fall back to this checkout's tree. + pass + if _PRIVATE_PACKAGE not in sys.modules: + if not (_KERNEL_DIR / f"{name}.py").is_file(): + raise ImportError(f"{name}.py is neither in the installed fastvideo_kernel nor in {_KERNEL_DIR}") + spec = importlib.util.spec_from_loader(_PRIVATE_PACKAGE, loader=None, is_package=True) + assert spec is not None and spec.submodule_search_locations is not None + spec.submodule_search_locations.append(str(_KERNEL_DIR)) + sys.modules[_PRIVATE_PACKAGE] = importlib.util.module_from_spec(spec) + return importlib.import_module(f"{_PRIVATE_PACKAGE}.{name}") + + +def load_qat_vsa_modules() -> tuple[ModuleType, ModuleType]: + """``(attn_qat_vsa_reference, attn_qat_vsa_train)``; the second needs Triton.""" + return _load("attn_qat_vsa_reference"), _load("attn_qat_vsa_train") + + +@dataclass(frozen=True) +class FP4AttentionNumerics: + """Numeric mode of the deployed kernel; mirror whatever inference ships. + + ``quantize=False`` keeps the emulator's sparse structure with exact + softmax (a BF16, grad-capable tile-128 route). ``quantize_pv=False`` + quantizes Q/K only, with P and V in BF16 (FA4 NVFP4-QK + BF16 PV). + """ + quantize: bool = True + quantize_pv: bool = True + two_level_p: bool = False + smooth_k: bool = False + first_block_max_floor: float | None = None + # Q/K-only error-reduction studies (forward-only except ``hadamard``); see ``fp4_vsa_attn_qat``. + hadamard: bool = False + smooth_q: str | None = None + prefix_precision: str | None = None + + +class FP4VSAFineAttention: + """``fine_attention_override`` callable: ``(q, k, v, mask, metadata) -> [B, S, H, D]``.""" + + def __init__(self, numerics: FP4AttentionNumerics) -> None: + self.numerics = numerics + self._reference, self._train = load_qat_vsa_modules() + + def __call__(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: torch.Tensor, + metadata: Any) -> torch.Tensor: + if int(metadata.tile_elems) != 128: + raise ValueError(f"the FP4 VSA emulator models 128-token tiles only, got {metadata.tile_elems}") + q2k_idx, q2k_num, kv_valid = self._reference.tile128_mask_to_fp4_blocks(mask, metadata.variable_block_sizes) + q, k, v = (t.transpose(1, 2) for t in (query, key, value)) + num_prefix_tiles = int(getattr(metadata, "num_prefix_tiles", 0) or 0) + out = self._train.fp4_vsa_attn_qat(q, + k, + v, + q2k_idx, + q2k_num, + kv_valid, + quantize=self.numerics.quantize, + quantize_pv=self.numerics.quantize_pv, + two_level_p=self.numerics.two_level_p, + smooth_k=self.numerics.smooth_k, + first_block_max_floor=self.numerics.first_block_max_floor, + hadamard=self.numerics.hadamard, + smooth_q=self.numerics.smooth_q, + prefix_precision=self.numerics.prefix_precision, + num_prefix_tiles=num_prefix_tiles) + return out.transpose(1, 2) + + +__all__ = ["FP4AttentionNumerics", "FP4VSAFineAttention", "load_qat_vsa_modules"] diff --git a/fastvideo/attention/backends/video_sparse_attn_h3.py b/fastvideo/attention/backends/video_sparse_attn_h3.py index e369cf44f..745234eac 100644 --- a/fastvideo/attention/backends/video_sparse_attn_h3.py +++ b/fastvideo/attention/backends/video_sparse_attn_h3.py @@ -52,11 +52,12 @@ the original logical tiles. Grad-tracking forwards and every backward stay on the Triton kernels. If the env is set but a precondition fails, the route logs one warning and falls back. -Tile 128 has exactly one implementation: the same ptx_blackwell -forward, which carries a 128-token block instantiation. It needs no opt-in, -pads an odd tile count with the same zero-valid partner tile, runs no-grad -forwards only, and fails closed (no Triton fallback) when the extension or -device cannot run it. +Tile 128 prefers the same ptx_blackwell forward, which carries a 128-token +block instantiation. It needs no opt-in, pads an odd tile count with the same +zero-valid partner tile, and runs no-grad forwards only. When the extension or +device cannot run it, each 128-token tile splits into two contiguous 64-token +halves on the Triton kernels; if those are not installed either, the call fails +closed. """ import functools @@ -555,6 +556,32 @@ def _ptx_blackwell_unavailable_reason(ptx_blackwell_mod: Any, query_bhsd: torch. return None +def _is_datacenter_blackwell(device: torch.device) -> bool: + """sm_100/sm_103 (B200, GB200, B300), the devices the ptx_blackwell kernel targets.""" + return device.type == "cuda" and torch.cuda.get_device_capability(device)[0] == 10 + + +@functools.cache +def _fine_attention_from_name(name: str) -> Any: + kind, _, mode = name.partition("-") + modes = {"emulator": ("mix8", "mixmx", "qk", "bf16")} + if kind not in modes or mode not in modes[kind]: + accepted = ", ".join(f"{prefix}-{value}" for prefix, values in modes.items() for value in values) + raise ValueError(f"Invalid FASTVIDEO_H3_VSA_FINE_ATTENTION={name!r}; accepted names: {accepted}") + from fastvideo.attention.backends.fp4_vsa_qat import FP4AttentionNumerics, FP4VSAFineAttention + if mode == "bf16": + return FP4VSAFineAttention(FP4AttentionNumerics(quantize=False)) + prefix = {"mix8": "fp8", "mixmx": "mxfp8", "qk": None}[mode] + return FP4VSAFineAttention( + FP4AttentionNumerics(quantize=True, quantize_pv=False, smooth_k=True, prefix_precision=prefix)) + + +def fine_attention_from_env() -> Any: + """The ``FASTVIDEO_H3_VSA_FINE_ATTENTION`` replacement of the tile-128 fine kernel (shared), or None.""" + name = envs.FASTVIDEO_H3_VSA_FINE_ATTENTION.get() + return _fine_attention_from_name(name) if name else None + + class MiniMaxH3VSAImpl(AttentionImpl): def __init__( @@ -589,10 +616,21 @@ class MiniMaxH3VSAImpl(AttentionImpl): # 128-token blocks, an even tile count with the partner tile, integer # tile sizes), which every call meets, so it is evaluated once. self._tile128_route: dict[tuple[torch.device, torch.dtype, int], str | None] = {} + # Optional replacement of the fine block-sparse kernel only (selection, + # mask and the gated compression branch are unchanged). Called as + # ``fn(query, key, value, mask, metadata) -> [B, S, H, D]`` on the + # logical tile-ordered rows. Quantization-aware training installs its + # FP4 emulator here; it is also the only grad-capable tile-128 route. + self.fine_attention_override: Any = fine_attention_from_env() def prepare_for_compile(self, device: torch.device) -> None: """Tensorize per-layer state shared by every torch.compile route.""" self._compile_layer_idx = torch.tensor(self.layer_idx, device=device, dtype=torch.int64) + if self.fine_attention_override is not None: + # forward() cannot run the emulator inside a traced graph, and logging there + # would break fullgraph capture, so say it here. + logger.warning_once("FASTVIDEO_H3_VSA_FINE_ATTENTION is ignored inside torch.compile; compiled " + "VSA-H3 layers run the native fine attention. Keep the DiT eager to use it.") def prepare_for_regional_compile(self, device: torch.device) -> str | None: """Resolve the inference-only ptx_blackwell route before fullgraph capture. @@ -838,7 +876,9 @@ class MiniMaxH3VSAImpl(AttentionImpl): # preserves the eager dense-layer contract without a Python branch. mask = mask | force_dense - if tile_elems in _PTX_BLACKWELL_TILE_ELEMS: + if self.fine_attention_override is not None and not compiling: + out = self.fine_attention_override(logical_query, logical_key, logical_value, mask, attn_metadata) + elif tile_elems in _PTX_BLACKWELL_TILE_ELEMS: # Native 64/128-token path: the block map is already at the # kernels' granularity. These entries take BHSD ([B, H, S_pad, D]); # mirror block_sparse_attn_256_bshd's Triton branch and transpose @@ -877,9 +917,10 @@ class MiniMaxH3VSAImpl(AttentionImpl): # fwd/bwd pairing is possible, but it is not built here. grad_mode = torch.is_grad_enabled() and (query.requires_grad or key.requires_grad or value.requires_grad) use_ptx_blackwell = False + split_tile128 = False if tile_elems == 128: - # The CUDA kernel's 128-token instantiation is the only tile-128 - # implementation, so it needs no opt-in and never falls back. + # The CUDA kernel's 128-token instantiation needs no opt-in; without it, + # the Triton tile-64 kernels run each tile as two halves. if grad_mode: raise NotImplementedError("VSA-H3 tile 128 supports no-grad inference forwards only.") route_key = (q_bhsd.device, q_bhsd.dtype, int(q_bhsd.shape[-1])) @@ -888,12 +929,20 @@ class MiniMaxH3VSAImpl(AttentionImpl): ptx_blackwell_variable_block_sizes, grad_mode) if reason is None and map_to_index is None: reason = "fastvideo_kernel.triton_kernels.index (map_to_index) is not importable" + if reason is not None and not compiling and _is_datacenter_blackwell(q_bhsd.device): + # These GPUs should have the native kernel; a broken build + # would otherwise only show up as a slow run. + logger.warning_once(f"VSA-H3 tile 128 on a data-center Blackwell GPU is falling back to " + f"the slower Triton tile-64 kernels: {reason}.") self._tile128_route[route_key] = reason reason = self._tile128_route[route_key] - if reason is not None: - raise RuntimeError(f"VSA-H3 tile 128 requires the ptx_blackwell block-sparse kernel: " - f"{reason}. Use VSA_tile_size=64 or 256 elsewhere.") - use_ptx_blackwell = True + if reason is None: + use_ptx_blackwell = True + elif block_sparse_attn_64_bhsd is None: + raise RuntimeError(f"VSA-H3 tile 128 requires the ptx_blackwell block-sparse kernel or " + f"the Triton tile-64 kernels: {reason}.") + else: + split_tile128 = True elif regional_compiling: if self._regional_compile_ptx_blackwell_enabled is None: raise RuntimeError( @@ -961,7 +1010,17 @@ class MiniMaxH3VSAImpl(AttentionImpl): v_bhsd = v_bhsd[:, :, :logical_seq_len] if not sm89_strided: q_bhsd, k_bhsd, v_bhsd = (t.contiguous() for t in (q_bhsd, k_bhsd, v_bhsd)) - if sm89_route: + if split_tile128: + # No sm_100a kernel on this device: run each 128-token tile as two + # contiguous 64-token halves on the Triton kernels. Every half inherits + # its tile's mask entries, and a partial tile's valid tokens fill the + # first half before the second, so the selection is unchanged. + logger.info_once("MiniMax-H3 VSA tile-128 forward: Triton tile-64 kernels on split tiles") + mask64 = mask.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1) + vbs = attn_metadata.variable_block_sizes + vbs64 = torch.stack((vbs.clamp(max=64), (vbs - 64).clamp(min=0)), dim=-1).flatten() + out_bhsd, _ = block_sparse_attn_64_bhsd(q_bhsd, k_bhsd, v_bhsd, mask64, vbs64) + elif sm89_route: from fastvideo.attention.backends.minimax_h3_sparse_int8 import sparse_sm89_attention logger.info_once(f"MiniMax-H3 VSA tile-64 forward: sm89 {self._sm89_kernel} QK / BF16 PV") out_bhsd = sparse_sm89_attention(q_bhsd, diff --git a/fastvideo/configs/models/dits/minimax_h3.py b/fastvideo/configs/models/dits/minimax_h3.py index b67f2993f..5a3272b69 100644 --- a/fastvideo/configs/models/dits/minimax_h3.py +++ b/fastvideo/configs/models/dits/minimax_h3.py @@ -4,6 +4,7 @@ from __future__ import annotations from dataclasses import dataclass, field +from typing import Any from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig from fastvideo.layers.pdd import validate_pdd_steps @@ -93,3 +94,4 @@ class MiniMaxH3Config(DiTConfig): # FastVideo's Fully Sharded Data Parallel (FSDP) loading path requires one # parameter dtype, while H3 inference keeps boundary projections in FP32. uniform_parameter_dtype: bool = False + vsa_fp4: dict[str, Any] = field(default_factory=dict) diff --git a/fastvideo/envs.py b/fastvideo/envs.py index b219b26fc..778835292 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -587,7 +587,48 @@ FASTVIDEO_FAD_REF_FEATURES = EnvStr(None, 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).") + "(sm_120 or GB10 sm_121, no-grad, eager inference).") +FASTVIDEO_H3_VSA_FP4_PREFIX = EnvChoice( + "", + choices=("", "fp8", "bf16"), + category="attention", + doc="Prefix split for the FP4 sparse attention kernel: keep the dense prefix query tiles (text, " + "reference and audio rows) off the FP4 kernel and run those rows densely with fp8 or bf16 Q/K and " + "bf16 P.V, while video query tiles stay on the FP4 kernel. The fp8 variant is the validated " + "training-time numerics and recovers most of the audio error that FP4 quantizes away. Use " + "64/128-token tiles. Empty keeps each row on the FP4 kernel.") +FASTVIDEO_H3_VSA_FP4_PREFIX_STEPS = EnvInt( + -1, + category="attention", + doc="Number of leading denoising steps that use the FP4 prefix split (see " + "FASTVIDEO_H3_VSA_FP4_PREFIX). -1 applies it to every step; 0 disables it. Early steps set the " + "global structure, so a small value keeps most of the split's quality benefit for a fraction of " + "its cost.") +FASTVIDEO_H3_VSA_FP4_TWO_LEVEL_P = EnvBool( + False, + category="attention", + doc="Quantize the FP4 attention probabilities with the two-level scheme the QAD student trained " + "with (rowmax-derived s_P1, then per-block s_P2) instead of the kernel's single-level default.") +FASTVIDEO_H3_VSA_FP4_K_SMOOTH = EnvBool( + False, + category="attention", + doc="Subtract the per-head key mean before FP4 key quantization, matching the emulator and QAD " + "training numerics. Exact for the softmax; it only changes quantization error.") +FASTVIDEO_H3_GATE_CHUNK_ROWS = EnvInt( + 0, + category="performance", + doc="Rows per chunk when the VSA FP4 gate branch expands the pooled compression output back " + "to token rows. 0 keeps the single full-size index_select; a positive value caps the transient " + "at chunk_rows*heads*head_dim and is bit-identical.") +FASTVIDEO_H3_VSA_FINE_ATTENTION = EnvChoice( + "", + choices=("", "emulator-mix8", "emulator-mixmx", "emulator-qk", "emulator-bf16"), + category="attention", + doc="Replace tile-128 fine attention, preserving VSA selection and compression, with the Triton " + "emulator of the sparse FP4 kernel: emulator-mix8 = NVFP4 Q/K video, per-token FP8 Q/K prefix; " + "emulator-mixmx = MXFP8 prefix; emulator-qk = NVFP4 Q/K everywhere; emulator-bf16 = unquantized " + "BF16. emulator-bf16 runs on any CUDA GPU. Nonempty overrides take priority over " + "FASTVIDEO_H3_VSA_FP4. Empty keeps the native route.") FASTVIDEO_H3_VSA_HEADS_FIRST_TILE = EnvBool(False, category="performance", doc="VSA-H3 64/128-token tiles: scatter rows straight into the heads-first " @@ -627,6 +668,12 @@ FASTVIDEO_LORA_SIDE_BRANCH = EnvBool(False, FASTVIDEO_NVFP4_MM_BACKEND = EnvStr("auto", category="performance", doc="FlashInfer mm_fp4 backend for NVFP4 linears, e.g. auto or cutlass.") +FASTVIDEO_NVFP4_MM_ROW_SPLIT = EnvChoice( + "off", + choices=("off", "table", "auto"), + category="performance", + doc="Run large cuDNN NVFP4 GEMMs over row chunks where its plan heuristic is slow: off, table (fixed " + "per-shape chunk counts) or auto (time 1/2/4 chunks once per shape; not deterministic across runs).") FASTVIDEO_NVFP4_ACT_AMAX = EnvPath(None, category="performance", doc="JSON of calibrated NVFP4 input amax per linear, keyed b. or " @@ -675,6 +722,23 @@ FASTVIDEO_H3_VAE_TILE_PARALLEL = EnvBool(False, 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_NVFP4_FUSED = EnvBool(True, + category="performance", + doc="Run NVFP4 MiniMax-H3 VAE decoder blocks through the bit-exact fused " + "Triton kernels at inference.") +FASTVIDEO_H3_VAE_NVFP4_CUDA_GRAPH = EnvBool(True, + category="performance", + doc="Replay small fused NVFP4 MiniMax-H3 VAE decoder calls (at most " + "four 1797-token tiles) as per-shape CUDA graphs.") +FASTVIDEO_H3_VAE_NVFP4_DECODER = EnvPath(None, + category="performance", + doc="Exported NVFP4 MiniMax-H3 video decoder (.pt, or .safetensors with " + "bf16 masters or pre-packed NVFP4 weights) that replaces the VAE decoder " + "at load; skips the INT8 ConvRot overlay.") +FASTVIDEO_H3_VAE_INT8_OVERLAY = EnvBool(True, + category="performance", + doc="Apply the INT8 ConvRot decoder overlay when the MiniMax-H3 VAE folder " + "ships one; 0 decodes with the dense weights.") FASTVIDEO_H3_VAE_INT8_SHARED_QKV = EnvBool(False, category="performance", doc="Share the INT8 activation rotation and quantization across the " diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 082e39abb..c52276cb8 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -292,6 +292,8 @@ class FastVideoArgs: # which keeps the pipeline-level compile semantics. inference_torch_compile: bool = False + h3_vsa_fp4: dict[str, Any] = field(default_factory=dict) + disable_autocast: bool = False # VSA parameters. None means unset: __post_init__ fills the checkpoint's @@ -438,6 +440,16 @@ class FastVideoArgs: env_backend = envs.FASTVIDEO_ATTENTION_BACKEND.get() if env_backend is not None and backend_name_to_enum(env_backend) is not None: self.attention_backend = env_backend + allowed_fp4 = {"enabled", "prefix_precision", "two_level_p", "smooth_k"} + if not isinstance(self.h3_vsa_fp4, dict) or self.h3_vsa_fp4.keys() - allowed_fp4: + raise ValueError(f"h3_vsa_fp4 must be a mapping with keys from {sorted(allowed_fp4)}") + for key in ("enabled", "two_level_p", "smooth_k"): + if key in self.h3_vsa_fp4 and not isinstance(self.h3_vsa_fp4[key], bool): + raise ValueError(f"h3_vsa_fp4.{key} must be a boolean") + if self.h3_vsa_fp4.get("prefix_precision", "") not in ("", "fp8", "bf16"): + raise ValueError("h3_vsa_fp4.prefix_precision must be empty, fp8 or bf16") + if self.h3_vsa_fp4.get("enabled") and (self.inference_torch_compile or self.enable_torch_compile): + raise ValueError("h3_vsa_fp4 requires eager DiT inference; disable DiT torch.compile") self._fold_vae_parallel_env() self._fold_h3_encoder_split_env() # Runs after FASTVIDEO_ATTENTION_BACKEND is copied into attention_backend, diff --git a/fastvideo/layers/fp4linear.py b/fastvideo/layers/fp4linear.py index 82a2c7573..04ab6bf3f 100644 --- a/fastvideo/layers/fp4linear.py +++ b/fastvideo/layers/fp4linear.py @@ -17,11 +17,17 @@ def _require_flashinfer() -> Any: return flashinfer -@torch.compile def _global_sf(t: torch.Tensor) -> torch.Tensor: - maxabs = t.float().abs().nan_to_num().max() - maxabs = maxabs.clamp(min=1e-12) - return (448.0 * 6.0) / maxabs + # Eager amax: the previous torch.compile'd reduction lowered to a single-block + # Triton kernel (~38 ms for a 21k x 2048 activation on GB200). Max is exact + # in any order, so finite input gives the same scale bit for bit; NaN and inf + # keep the elementwise reduction's behavior. No host sync. + maxabs = torch.nan_to_num(t.abs()).amax().float() + # The reduction above runs in the input dtype, so an infinite input becomes the + # fp16/bf16 maximum. The reference reduction replaces it in float32. + # new_full fills on the device; torch.tensor would copy from the host and sync. + maxabs = torch.where(torch.isinf(t).any(), maxabs.new_full((), torch.finfo(torch.float32).max), maxabs) + return (448.0 * 6.0) / maxabs.clamp(min=1e-12) class _LinearFWD4BWD16Fn(torch.autograd.Function): diff --git a/fastvideo/layers/quantization/nvfp4_config.py b/fastvideo/layers/quantization/nvfp4_config.py index 17460f46f..acecea385 100644 --- a/fastvideo/layers/quantization/nvfp4_config.py +++ b/fastvideo/layers/quantization/nvfp4_config.py @@ -96,6 +96,33 @@ H3_NVFP4_DIT_BUFFER_NAMES = ( # Optional per-layer static activation global scale (448 * 6 / calibrated input amax). # Exports without it quantize activations with the unit global scale. H3_NVFP4_DIT_INPUT_SF_NAME = "_nvfp4_input_global_sf" +# Safetensors metadata key of a packed export: JSON list of main blocks whose linears stay bf16 +# (they are absent from the export and load from the dense shards). +H3_NVFP4_DIT_SKIP_BLOCKS_KEY = "skip_blocks" +_H3_MAIN_BLOCK = re.compile(r"(?:^|\.)transformer_blocks\.(\d+)\.") + + +def minimax_h3_block_index(prefix: str) -> int | None: + """Main transformer-block index of a MiniMax-H3 module prefix, or None outside the main blocks.""" + match = _H3_MAIN_BLOCK.search(prefix) + return int(match.group(1)) if match is not None else None + + +def read_minimax_h3_nvfp4_dit_skip_blocks(path: str) -> tuple[int, ...]: + """The bf16 main blocks a packed NVFP4H3 export declares (empty for exports that quantize every block).""" + import json + + from safetensors import safe_open + + with safe_open(path, framework="pt", device="cpu") as reader: + raw = (reader.metadata() or {}).get(H3_NVFP4_DIT_SKIP_BLOCKS_KEY) + if raw is None: + return () + blocks = json.loads(raw) + if not isinstance(blocks, list) or not all( + isinstance(b, int) and not isinstance(b, bool) and b >= 0 for b in blocks): + raise ValueError(f"{path}: metadata {H3_NVFP4_DIT_SKIP_BLOCKS_KEY!r} must be a JSON list of block indices") + return tuple(sorted(set(blocks))) def is_ltx2_nvfp4_linear_prefix(prefix: str) -> bool: @@ -203,14 +230,17 @@ def nvfp4_quantize_fenced( sfLayout=SfLayout(sf_layout), do_shuffle=do_shuffle, enable_pdl=False) - # With FlashInfer 0.6.18 on GB10, queued activation quantization - # plus GEMM can diverge. Completing quantization while its padded - # input is alive prevents the observed intermittent corruption. + # With FlashInfer 0.6.18 and 0.7.0 on GB10, queued activation + # quantization plus GEMM can diverge: identical seeded H3 requests + # produce different latents. Completing each quantization on the + # host prevents it, at no measurable denoise cost. torch.cuda.current_stream(x.device).synchronize() return quantized, scales _OPS_REGISTERED = False +# Cleared the first time FlashInfer returns 128x4 scales without row padding. +_FLASHINFER_PADS_128X4_SCALES = True def _register_ops_once() -> None: @@ -241,10 +271,20 @@ def _register_ops_once() -> None: sf_layout: int, do_shuffle: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: - del global_sf, sf_layout, do_shuffle + del global_sf + SfLayout, _, _ = _require_flashinfer() m, k = x.shape quantized = torch.empty((m, (k + 1) // 2), device=x.device, dtype=torch.uint8) - scales = torch.empty((m, (k + 15) // 16), device=x.device, dtype=torch.uint8) + # Mirror FlashInfer's shapes: swizzled scale layouts pad rows to the + # tile (128 or 8) and scale columns to 4; linear (and shuffled 128x4) + # scales are unpadded. ``_nvfp4_quantize`` relies on the 128x4 padding. + if do_shuffle or sf_layout == SfLayout.layout_linear.value: + scales = torch.empty((m, (k + 15) // 16), device=x.device, dtype=torch.uint8) + else: + row_tile = 8 if sf_layout == SfLayout.layout_8x4.value else 128 + scales = torch.empty((m + (-m) % row_tile, ((k + 15) // 16 + 3) // 4 * 4), + device=x.device, + dtype=torch.uint8) return quantized, scales @torch.library.custom_op( @@ -329,21 +369,34 @@ def _nvfp4_quantize( elif global_sf.device != x.device: global_sf = global_sf.to(device=x.device) if sf_layout == SfLayout.layout_linear.value: - x_for_quant = x - logical_rows = x.shape[0] - else: - # Sequence-parallel can feed either logical rows or row-padded - # rows. Normalize to the kernel tile shape for swizzled layouts - # so both paths share a stable quantization contract. - row_tile = 8 if sf_layout == SfLayout.layout_8x4.value else 128 - logical_rows = x.shape[0] - pad_rows = (-logical_rows) % row_tile - x_for_quant = F.pad(x, (0, 0, 0, pad_rows)) - + return torch.ops.fastvideo_fp4.nvfp4_quantize(x, global_sf, sf_layout, do_shuffle) + # Sequence-parallel can feed either logical rows or row-padded + # rows. Normalize to the kernel tile shape for swizzled layouts + # so both paths share a stable quantization contract. + row_tile = 8 if sf_layout == SfLayout.layout_8x4.value else 128 + logical_rows = x.shape[0] + padded_rows = logical_rows + (-logical_rows) % row_tile + global _FLASHINFER_PADS_128X4_SCALES + if sf_layout == SfLayout.layout_128x4.value and not do_shuffle and _FLASHINFER_PADS_128X4_SCALES: + # FlashInfer pads 128x4 scale factors to whole row tiles itself, and + # quantizes the logical rows bit for bit like a zero-padded copy, so + # skip that copy (a full activation write per linear). The fake op + # reports the same padded scales, so a compiled graph takes this + # path unconditionally, without a guard on the row count. + quantized, scales = torch.ops.fastvideo_fp4.nvfp4_quantize(x, global_sf, sf_layout, do_shuffle) + # Columns are padded to whole 4-wide tiles too, so compare against the + # fully padded size; logical rows alone can reach the unpadded column count. + padded_cols = ((x.shape[-1] + 15) // 16 + 3) // 4 * 4 + if torch.compiler.is_compiling() or scales.numel() >= padded_rows * padded_cols: + return quantized, scales + # Only an unexpected FlashInfer build gets here; quantize it padded + # from now on rather than quantizing every activation twice. + logger.warning("FlashInfer returned unpadded 128x4 NVFP4 scales; " + "falling back to quantizing a row-padded activation copy.") + _FLASHINFER_PADS_128X4_SCALES = False + x_for_quant = F.pad(x, (0, 0, 0, padded_rows - logical_rows)) quantized, scales = torch.ops.fastvideo_fp4.nvfp4_quantize(x_for_quant, global_sf, sf_layout, do_shuffle) - if sf_layout != SfLayout.layout_linear.value: - quantized = quantized.narrow(0, 0, logical_rows) - return quantized, scales + return quantized.narrow(0, 0, logical_rows), scales def _mm_fp4( @@ -365,6 +418,10 @@ def _mm_fp4( raise TypeError(f"Unsupported kwargs for _mm_fp4: {sorted(kwargs)}") if alpha is not None and not torch.is_tensor(alpha): alpha = torch.tensor(alpha, device=a.device, dtype=torch.float32) + if (out is None and backend == "cudnn" and a.shape[0] >= _SPLIT_MIN_ROWS + and envs.FASTVIDEO_NVFP4_MM_ROW_SPLIT.get() != "off" and not torch.compiler.is_compiling()): + return _mm_fp4_row_split(a, b, a_scale, b_scale, alpha, out_dtype, block_size, use_8x4_sf_layout, backend, + use_nvfp4) return torch.ops.fastvideo_fp4.mm_fp4( a, b, @@ -380,6 +437,78 @@ def _mm_fp4( ) +_SPLIT_MIN_ROWS = 16384 +_SPLIT_CANDIDATES = (1, 2, 4) +_SplitKey = tuple[int, int, int, str, int, bool, torch.dtype] +_SPLIT_CHOICE: dict[_SplitKey, int] = {} +# Measured on GB10 (cuDNN 9.20 via FlashInfer 0.6): (packed K, N) -> chunks. H3's FF down projection is the only +# shape that falls off the plan heuristic; the others are fastest whole. +_SPLIT_TABLE: dict[tuple[int, int], int] = {(7168, 5376): 2} + + +def _mm_fp4_row_chunks(a: torch.Tensor, a_scale: torch.Tensor, parts: int) -> list[tuple[torch.Tensor, torch.Tensor]]: + """Split activations and their 128x4-swizzled scales into ``parts`` row chunks on 128-row boundaries.""" + rows = a.shape[0] + step = max(rows // parts // 128 * 128, 128) + bounds = [i * step for i in range(parts)] + [rows] + chunks = [] + for lo, hi in zip(bounds[:-1], bounds[1:], strict=True): + if hi <= lo: + continue + # The 128x4 layout interleaves rows within a tile, so the last chunk needs + # the whole final scale tile even though its activation slice ends at `hi`. + scale_hi = min(((hi + 127) // 128) * 128, a_scale.shape[0]) + chunks.append((a[lo:hi], a_scale[lo:scale_hi])) + return chunks + + +def _mm_fp4_chunked(a, b, a_scale, b_scale, alpha, out_dtype, block_size, use_8x4_sf_layout, backend, use_nvfp4, + parts: int) -> torch.Tensor: + chunks = _mm_fp4_row_chunks(a, a_scale, parts) if parts > 1 else [(a, a_scale)] + outs = [ + torch.ops.fastvideo_fp4.mm_fp4(ca, b, cs, b_scale, alpha, out_dtype, None, block_size, use_8x4_sf_layout, + backend, use_nvfp4) for ca, cs in chunks + ] + return outs[0] if len(outs) == 1 else torch.cat(outs, dim=0) + + +def _mm_fp4_row_split(a, b, a_scale, b_scale, alpha, out_dtype, block_size, use_8x4_sf_layout, backend, + use_nvfp4) -> torch.Tensor: + """``mm_fp4`` over row chunks when that is faster. + + The cuDNN FP4 plan heuristic falls off a cliff for some (rows, K, N) shapes: H3's FF + down projection (K=14336, N=5376, 35k rows) runs at ~110 TFLOPs whole but ~360 TFLOPs in + two halves (and N=14336 does the reverse). The default ``table`` mode uses fixed per-shape + chunk counts, so every run takes the same kernels. Chunks accumulate in a different order from + the whole-matrix plan, so a few elements round one bf16 step differently (about 1e-8 of outputs, + at massive-activation rows); ``auto`` times 1/2/4 chunks once per shape instead. + """ + key = (a.shape[0], a.shape[1], b.shape[1], backend, block_size, use_8x4_sf_layout, out_dtype) + args = (a, b, a_scale, b_scale, alpha, out_dtype, block_size, use_8x4_sf_layout, backend, use_nvfp4) + # Splitting slices the 128x4 scale rows; an 8x4 layout packs them differently. + if use_8x4_sf_layout or a_scale.dim() != 2 or a_scale.shape[0] < a.shape[0]: + return _mm_fp4_chunked(*args, 1) + if envs.FASTVIDEO_NVFP4_MM_ROW_SPLIT.get() == "table": + _SPLIT_CHOICE.setdefault(key, _SPLIT_TABLE.get((a.shape[1], b.shape[1]), 1)) + if key not in _SPLIT_CHOICE: + import time + timings = {} + for parts in _SPLIT_CANDIDATES: + _mm_fp4_chunked(*args, parts) + torch.cuda.synchronize() + start = time.perf_counter() + for _ in range(2): + _mm_fp4_chunked(*args, parts) + torch.cuda.synchronize() + timings[parts] = time.perf_counter() - start + _SPLIT_CHOICE[key] = min(timings, key=timings.__getitem__) + logger.info("NVFP4 mm_fp4 row split for %s: %d chunk(s) (%s ms)", key, _SPLIT_CHOICE[key], { + k: round(v * 500, 1) + for k, v in timings.items() + }) + return _mm_fp4_chunked(*args, _SPLIT_CHOICE[key]) + + def _mm_fp4_backend() -> str: """FlashInfer ``mm_fp4`` backend (``FASTVIDEO_NVFP4_MM_BACKEND``, default ``auto``). @@ -602,10 +731,19 @@ class NVFP4Config(QuantizationConfig): checkpoints such as FastH3 V2 NVFP4). ``layer_profile="h3_dit_vsa"`` is ``h3_dit`` plus each block's VSA compression gate ``attn.to_gate_compress`` (VSA-distilled students). + ``skip_blocks`` keeps every linear of those MiniMax-H3 main blocks in + bf16 (mixed exports such as FastH3 OmniRef QAD, blocks 42-49). A packed + export records its own list, which the loader adopts when this is empty. """ - def __init__(self, layer_profile: str = "refine", retain_original_weights: bool | None = None): + def __init__(self, + layer_profile: str = "refine", + retain_original_weights: bool | None = None, + skip_blocks: tuple[int, ...] | list[int] = ()): super().__init__() + if any(isinstance(b, bool) or not isinstance(b, int) or b < 0 for b in skip_blocks): + raise ValueError(f"NVFP4Config.skip_blocks must be non-negative block indices, got {skip_blocks!r}") + self.skip_blocks: tuple[int, ...] = tuple(sorted(set(skip_blocks))) if layer_profile not in ("base", "refine", "h3_dit", "h3_dit_ffn", "h3_dit_vsa"): raise ValueError("NVFP4Config.layer_profile must be one of 'base', 'refine', 'h3_dit', " f"'h3_dit_ffn', or 'h3_dit_vsa', got {layer_profile!r}") @@ -638,13 +776,22 @@ class NVFP4Config(QuantizationConfig): return cls( layer_profile=config.get("layer_profile", "refine"), retain_original_weights=config.get("retain_original_weights"), + skip_blocks=tuple(config.get("skip_blocks", ())), ) + def with_skip_blocks(self, skip_blocks: tuple[int, ...]) -> NVFP4Config: + """A copy of this config with ``skip_blocks`` replaced.""" + return NVFP4Config(layer_profile=self.layer_profile, + retain_original_weights=self.retain_original_weights, + skip_blocks=skip_blocks) + def get_quant_method(self, layer: torch.nn.Module, prefix: str): from fastvideo.layers.linear import LinearBase if not isinstance(layer, LinearBase): return None + if self.skip_blocks and minimax_h3_block_index(prefix) in self.skip_blocks: + return None if self.layer_profile == "h3_dit": tagged = is_minimax_h3_nvfp4_dit_linear_prefix(prefix) elif self.layer_profile == "h3_dit_vsa": @@ -815,6 +962,7 @@ __all__ = [ "H3_NVFP4_DIT_BUFFER_NAMES", "H3_NVFP4_DIT_EXPORT_FILENAME", "H3_NVFP4_DIT_INPUT_SF_NAME", + "H3_NVFP4_DIT_SKIP_BLOCKS_KEY", "NVFP4Config", "NVFP4QuantizeMethod", "convert_model_to_nvfp4", @@ -825,5 +973,7 @@ __all__ = [ "is_minimax_h3_nvfp4_dit_linear_prefix", "is_minimax_h3_nvfp4_linear_prefix", "load_minimax_h3_nvfp4_dit_export", + "minimax_h3_block_index", "nvfp4_linear_weight_param_names", + "read_minimax_h3_nvfp4_dit_skip_blocks", ] diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py index df464c4a7..119b70b73 100644 --- a/fastvideo/models/dits/minimax_h3.py +++ b/fastvideo/models/dits/minimax_h3.py @@ -206,6 +206,7 @@ class MiniMaxH3Attention(nn.Module): fuse_qknorm_rope: bool = False, fa4_packed_varlen: bool = False, exact_rope: bool = False, + vsa_fp4: dict[str, Any] | None = None, ) -> None: super().__init__() self.num_attention_heads = num_attention_heads @@ -264,7 +265,11 @@ class MiniMaxH3Attention(nn.Module): # Opt-in inference route: VSA-H3 selection on the block-sparse FP4 # kernel (see minimax_h3_vsa_fp4); grad and compile keep the generic path. self._layer_idx = layer_idx_from_prefix(prefix, default=-1) - self._vsa_fp4 = use_vsa and vsa_fp4_requested() + self._vsa_fp4_options = dict(vsa_fp4 or {}) + requested = self._vsa_fp4_options.get("enabled") + if requested and not use_vsa: + raise ValueError("h3_vsa_fp4 requires the VIDEO_SPARSE_ATTN_H3 attention backend") + self._vsa_fp4 = use_vsa and (vsa_fp4_requested() if requested is None else requested) self._vsa_tile_first = use_vsa and envs.FASTVIDEO_H3_VSA_TILE_FIRST.get() self.to_gate_compress: ReplicatedLinear | None = None # None = unchecked; the first forward tests the loaded weight once and @@ -350,7 +355,8 @@ class MiniMaxH3Attention(nn.Module): rotary_emb: tuple[torch.Tensor, torch.Tensor] | None, original_seq_len: int, ) -> torch.Tensor: - if (self._vsa_fp4 and rotary_emb is not None and not torch.is_grad_enabled() + fine_override = getattr(getattr(self.distributed_attention, "attn_impl", None), "fine_attention_override", None) + if (self._vsa_fp4 and fine_override is None and rotary_emb is not None and not torch.is_grad_enabled() and not torch.compiler.is_compiling()): meta = get_forward_context().attn_metadata # Exempt mode lists the first prefix tile for every query, which the @@ -632,6 +638,7 @@ class MiniMaxH3TransformerBlock(nn.Module): fuse_swiglu: bool = False, fa4_packed_varlen: bool = False, exact_kernels: frozenset[str] = frozenset(), + vsa_fp4: dict[str, Any] | None = None, ) -> None: super().__init__() self.norm1 = nn.RMSNorm(hidden_size, eps=norm_eps) @@ -646,6 +653,7 @@ class MiniMaxH3TransformerBlock(nn.Module): fuse_qknorm_rope=fuse_qknorm_rope, fa4_packed_varlen=fa4_packed_varlen, exact_rope="rope" in exact_kernels, + vsa_fp4=vsa_fp4, ) self.norm2 = nn.RMSNorm(hidden_size, eps=norm_eps) self.ff = MiniMaxH3FeedForward( @@ -907,6 +915,7 @@ class MiniMaxH3Transformer3DModel(BaseDiT): fuse_swiglu="swiglu" in self.enabled_fusions, fa4_packed_varlen=envs.FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN.get(), exact_kernels=self.enabled_exact_kernels, + vsa_fp4=config.vsa_fp4, ) for index in range(arch.num_layers) ]) self.norm_out = MiniMaxH3AdaLayerNormOut( diff --git a/fastvideo/models/dits/minimax_h3_vsa_fp4.py b/fastvideo/models/dits/minimax_h3_vsa_fp4.py index 6c18e150a..4186cc8bb 100644 --- a/fastvideo/models/dits/minimax_h3_vsa_fp4.py +++ b/fastvideo/models/dits/minimax_h3_vsa_fp4.py @@ -1,8 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 """Inference fast path: VSA-H3 attention on the block-sparse FP4 kernel. -Opt-in with ``FASTVIDEO_H3_VSA_FP4=1`` (no-grad, single sequence-parallel -rank, ``fastvideo-kernel`` built with ``attn_qat_infer``). The selection is +Opt-in with ``h3_vsa_fp4.enabled`` or ``FASTVIDEO_H3_VSA_FP4=1`` (no-grad, +eager inference, ``fastvideo-kernel`` built with ``attn_qat_infer``). The selection is VSA-H3's own: tile pooling, top-k block mask and the gated compression branch are unchanged; only the block-sparse attention itself runs on SageAttention3's FP4 kernel (BF16 Triton otherwise), with 64-token tiles carried by quadrant @@ -28,22 +28,47 @@ import fastvideo.envs as envs from fastvideo.attention.backends.video_sparse_attn_h3 import (MiniMaxH3VSAMetadata, _build_block_mask, _pool_tiles) -def _sparse_fp4_attention(api: Any, *args: Any) -> torch.Tensor: +def _fp4_kernel_kwargs(two_level_p: bool | None = None) -> dict[str, object]: + """Quantization options for the FP4 video kernel. + + The native route defaults to single-level P; the QAD student trained with + two-level P, selected here by an inference-only switch. + """ + kwargs: dict[str, object] = {} + enabled = envs.FASTVIDEO_H3_VSA_FP4_TWO_LEVEL_P.get() if two_level_p is None else two_level_p + if enabled: + kwargs["single_level_p_quant"] = False + return kwargs + + +def _sparse_fp4_attention(api: Any, *args: Any, two_level_p: bool | None = None, + smooth_k: bool | None = None) -> torch.Tensor: """Run the block-sparse FP4 kernel on lists built by vsa_tile_mask_to_fp4_blocks. Those lists are in range by construction, so skip the host-sync validation where the installed fastvideo-kernel supports turning it off; older releases have no ``validate`` argument. + + ``FASTVIDEO_H3_VSA_FP4_K_SMOOTH`` subtracts the per-head key mean before the kernel + quantizes K. The shift is exact under softmax, so it only changes quantization error. """ + smooth = envs.FASTVIDEO_H3_VSA_FP4_K_SMOOTH.get() if smooth_k is None else smooth_k + if smooth: + args = (args[0], args[1] - args[1].mean(dim=1, keepdim=True), *args[2:]) try: - return api.sageattn_blackwell_sparse_bshd(*args, validate=False) + return api.sageattn_blackwell_sparse_bshd(*args, validate=False, **_fp4_kernel_kwargs(two_level_p)) except TypeError as error: if "validate" not in str(error): raise - return api.sageattn_blackwell_sparse_bshd(*args) + return api.sageattn_blackwell_sparse_bshd(*args, **_fp4_kernel_kwargs(two_level_p)) _BLOCK = 128 +def _resolve_prefix_mode(mode: str, steps: int, step: int) -> str: + """The prefix split for one denoising step; ``steps < 0`` applies it to every step.""" + return mode if steps < 0 or step < steps else "" + + class _StageTimer: """Opt-in (``FASTVIDEO_H3_SP_PROFILE=1``) CUDA-event spans summed per stage over one DiT forward. @@ -53,13 +78,17 @@ class _StageTimer: def __init__(self) -> None: self.enabled = envs.FASTVIDEO_H3_SP_PROFILE.get() + self.nvtx = envs.FASTVIDEO_NVTX_PROFILE.get() and torch.cuda.is_available() self._spans: list[tuple[str, torch.cuda.Event, torch.cuda.Event]] = [] def span(self, name: str): import contextlib - if not self.enabled: - return contextlib.nullcontext() - return self._record(name) + if self.enabled: + return self._record(name) + if self.nvtx: + from fastvideo.profiler import nvtx_range + return nvtx_range(f"h3.{name}") + return contextlib.nullcontext() def _record(self, name: str): import contextlib @@ -95,12 +124,33 @@ def vsa_fp4_requested() -> bool: def _api() -> Any: global _fp4_api if _fp4_api is None: - import attn_qat_infer.api as api + try: + import attn_qat_infer.api as api + except (ImportError, OSError) as error: + raise RuntimeError("MiniMax-H3 FP4 attention was requested but attn_qat_infer could not load. " + "Build fastvideo-kernel for this GPU from the same checkout; " + "disable FP4 attention explicitly to use BF16 VSA.") from error _fp4_api = api return _fp4_api -class _TileLayout: +class _KeyIndexCache: + """The prefix split's valid-key positions, computed once per step layout. + + ``kv_valid`` comes from the tile sizes alone, so every block of the forward + gets the same index; only the first block pays the host sync. + """ + tile: int + n_tiles: int + _valid_key_index: torch.Tensor | None = None + + def valid_key_index(self, kv_valid: torch.Tensor, half: int) -> torch.Tensor: + if self._valid_key_index is None: + self._valid_key_index = _valid_key_index(kv_valid, half, self.n_tiles * self.tile) + return self._valid_key_index + + +class _TileLayout(_KeyIndexCache): """Per-step tile-order state shared by every block of one forward.""" def __init__(self, meta: MiniMaxH3VSAMetadata, rotary_emb: tuple[torch.Tensor, torch.Tensor]) -> None: @@ -157,6 +207,93 @@ def _shared_input_projections(linears: tuple[Any, ...], x: torch.Tensor) -> list return [m.apply(linear, x, linear.bias, pre_quantized=pre) for m, linear in zip(methods, linears, strict=True)] +@torch.compile(dynamic=True, fullgraph=True) +def _fp8_per_token(x: torch.Tensor) -> torch.Tensor: + """``[..., D]`` BF16 -> per-(row, head) FP8 E4M3 -> BF16, the prefix Q/K precision of the QK study. + + Element for element the emulator's ``fp8_per_token_fake_quant``. The zero + scale is spelled as a reciprocal rather than a clamp so that a row with a + subnormal amax rounds the same way in both. + """ + xf = x.float() + scale = xf.abs().amax(dim=-1, keepdim=True) / _FP8_MAX + inv = torch.where(scale == 0, torch.zeros_like(scale), 1.0 / scale) + return (xf * inv).clamp(-_FP8_MAX, _FP8_MAX).to(_FP8).float().mul(scale).to(torch.bfloat16) + + +def _half_row_valid(kv_valid: torch.Tensor, half: int) -> torch.Tensor: + """Per-half valid-key counts -> a ``[rows]`` bool over the tile-ordered, zero-padded keys.""" + return (torch.arange(half, device=kv_valid.device).unsqueeze(0) < kv_valid.unsqueeze(1)).reshape(-1) + + +def _valid_key_index(kv_valid: torch.Tensor, half: int, logical: int) -> torch.Tensor: + """Positions of the valid keys among the first ``logical`` tile-ordered rows (a host sync).""" + return _half_row_valid(kv_valid, half)[:logical].nonzero().squeeze(1) + + +def _prefix_dense_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, prefix_rows: int, + logical: int, valid_index: torch.Tensor, precision: str) -> torch.Tensor: + """Dense attention for the prefix query rows, ``[B, H, prefix_rows, D]``. + + VSA gives every prefix (non-video) query tile every key column, so those rows need no + block lists. Running them off the FP4 kernel lets their Q/K stay at 8-bit precision, + which is where most of the FP4 attention error sits (the dense prefix rows carry the + text, reference and audio queries). + + K smoothing is exact - it shifts every score row by the same constant - and runs + before the FP8 rounding, in the order the emulator and the QAD student use. + """ + import torch.nn.functional as F + + q, k, v = query[:, :prefix_rows], key[:, :logical], value[:, :logical] + k = k - k.mean(dim=1, keepdim=True) + if precision == "fp8": + q, k = _fp8_per_token(q), _fp8_per_token(k) + # Attend only the valid keys. Gathering them keeps SDPA on its unmasked (flash) kernel; a boolean + # mask over all rows falls back to the memory-efficient kernel, 3-4x slower at OmniRef prefix sizes. + k, v = k.index_select(1, valid_index), v.index_select(1, valid_index) + return F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), + scale=1.0 / math.sqrt(q.shape[-1])) + + +def _fp4_fine_attention(api: Any, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, + q2k_idx: torch.Tensor, q2k_num: torch.Tensor, kv_valid: torch.Tensor, + q2k_quad: torch.Tensor | None, *, logical: int, prefix_rows: int, precision: str, + options: dict[str, Any], layout: _KeyIndexCache | None = None) -> torch.Tensor: + numerics = {"two_level_p": options.get("two_level_p"), "smooth_k": options.get("smooth_k")} + if not prefix_rows: + return _sparse_fp4_attention(api, query, key, value, q2k_idx, q2k_num, kv_valid, q2k_quad, **numerics) + if not getattr(api, "KEY_LENGTH_SIZING", False): + raise RuntimeError("the MiniMax-H3 FP4 prefix split needs a fastvideo-kernel build whose " + "sageattn_blackwell_sparse_bshd sizes by key length; rebuild fastvideo-kernel " + "from this checkout") + if layout is not None: + valid_index = layout.valid_key_index(kv_valid, api.HALF_N) + else: + valid_index = _valid_key_index(kv_valid, api.HALF_N, logical) + with STAGES.span("prefix_attention"): + prefix_out = _prefix_dense_attention(query, key, value, prefix_rows, logical, valid_index, precision) + # A tile-64 prefix can end halfway through a native 128-row block. Keep that whole query block and its + # original quadrant mask, then discard the overlapping prefix half and trailing padding from its output. + head = prefix_rows // _BLOCK + start = head * _BLOCK + video_out = _sparse_fp4_attention(api, query[:, start:], key, value, + q2k_idx[:, :, head:].contiguous(), q2k_num[:, :, head:].contiguous(), kv_valid, + None if q2k_quad is None else q2k_quad[:, :, head:].contiguous(), **numerics) + return torch.cat((prefix_out, video_out[:, :, prefix_rows - start:logical - start]), dim=2) + + +def _prefix_settings(attn: Any, meta: MiniMaxH3VSAMetadata, n_tiles: int) -> tuple[str, int]: + options = getattr(attn, "_vsa_fp4_options", {}) + mode = options.get("prefix_precision") + if mode is None: + mode = envs.FASTVIDEO_H3_VSA_FP4_PREFIX.get() + mode = _resolve_prefix_mode(mode, envs.FASTVIDEO_H3_VSA_FP4_PREFIX_STEPS.get(), int(meta.current_timestep)) + rows = (meta.num_prefix_tiles * int(meta.tile_elems) + if mode and int(meta.tile_elems) in (64, 128) and 0 < meta.num_prefix_tiles < n_tiles else 0) + return mode, rows + + def vsa_tile_first_attention(attn: Any, hidden_states: torch.Tensor, rotary_emb: tuple[torch.Tensor, torch.Tensor], meta: MiniMaxH3VSAMetadata, use_fused_rope: bool) -> torch.Tensor: @@ -249,9 +386,11 @@ def vsa_fp4_attention(attn: Any, hidden_states: torch.Tensor, rotary_emb: tuple[ mask = _build_block_mask(scores, meta.num_prefix_tiles, sparsity, meta.exempt, meta.video_tile_spans, meta.span_sparsities) q2k_idx, q2k_num, kv_valid, q2k_quad = api.vsa_tile_mask_to_fp4_blocks(mask, layout.tile, vbs) + prefix_mode, prefix_rows = _prefix_settings(attn, meta, layout.n_tiles) with STAGES.span("fp4_attention"): - # Lists come from vsa_tile_mask_to_fp4_blocks and are in range; skip the per-call host sync. - out = _sparse_fp4_attention(api, query, key, value, q2k_idx, q2k_num, kv_valid, q2k_quad) + out = _fp4_fine_attention(api, query, key, value, q2k_idx, q2k_num, kv_valid, q2k_quad, + logical=logical, prefix_rows=prefix_rows, precision=prefix_mode, + options=getattr(attn, "_vsa_fp4_options", {}), layout=layout) with STAGES.span("out_untile"): out = out.transpose(1, 2).index_select(1, layout.untile) # [B, L, H, D], packed order if sim_fp8: @@ -262,7 +401,17 @@ def vsa_fp4_attention(attn: Any, hidden_states: torch.Tensor, rotary_emb: tuple[ gate, _ = attn.to_gate_compress(hidden_states) v_pooled = _pool_tiles(value[:, :logical], vbs, layout.tile) out_c = torch.matmul(torch.softmax(scores, dim=-1), v_pooled).permute(0, 2, 1, 3).to(out.dtype) - out = out.addcmul_(out_c.index_select(1, layout.row_tile), gate.unflatten(-1, (heads, dim))) + gate = gate.unflatten(-1, (heads, dim)) + chunk = envs.FASTVIDEO_H3_GATE_CHUNK_ROWS.get() + if chunk > 0 and out.shape[1] > chunk: + # The full gather materializes [B, L, H, D] (1.2 GiB at 768p10) beside the 36L + # decoder. Chunked gathers and the same elementwise FMA are bit-identical. + for start in range(0, out.shape[1], chunk): + stop = min(start + chunk, out.shape[1]) + out[:, start:stop].addcmul_(out_c.index_select(1, layout.row_tile[start:stop]), + gate[:, start:stop]) + else: + out = out.addcmul_(out_c.index_select(1, layout.row_tile), gate) return out.flatten(2, 3) @@ -350,7 +499,7 @@ def _all_to_all(payload: torch.Tensor, scale: torch.Tensor, group: Any) -> tuple return recv, recv_scale -class _SPTileLayout: +class _SPTileLayout(_KeyIndexCache): """Per-step tile-order state for one rank's head subset.""" def __init__(self, meta: MiniMaxH3VSAMetadata, rank: int, local_rows: int) -> None: @@ -409,6 +558,7 @@ def vsa_fp4_attention_sp(attn: Any, hidden_states: torch.Tensor, rotary_emb: tup vbs = meta.variable_block_sizes logical = layout.n_tiles * layout.tile + prefix_mode, prefix_rows = _prefix_settings(attn, meta, layout.n_tiles) with STAGES.span("select_mask"): scores = torch.matmul(_pool_tiles(q_t[:, :logical], vbs, layout.tile), _pool_tiles(k_t[:, :logical], vbs, layout.tile).transpose(-2, -1)) / (dim**0.5) @@ -417,7 +567,9 @@ def vsa_fp4_attention_sp(attn: Any, hidden_states: torch.Tensor, rotary_emb: tup meta.span_sparsities) q2k_idx, q2k_num, kv_valid, q2k_quad = api.vsa_tile_mask_to_fp4_blocks(mask, layout.tile, vbs) with STAGES.span("fp4_attention"): - out_bhsd = _sparse_fp4_attention(api, q_t, k_t, v_t, q2k_idx, q2k_num, kv_valid, q2k_quad) + out_bhsd = _fp4_fine_attention(api, q_t, k_t, v_t, q2k_idx, q2k_num, kv_valid, q2k_quad, + logical=logical, prefix_rows=prefix_rows, precision=prefix_mode, + options=getattr(attn, "_vsa_fp4_options", {}), layout=layout) with STAGES.span("out_pack"): payload, scale = _pack_seq_fp8(out_bhsd, layout.untile, world, local_rows) diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 6d9f3d7d8..35eacaebf 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -827,6 +827,33 @@ class VAELoader(ComponentLoader): return directory return None + @staticmethod + def _defers_h3_nvfp4_decoder(class_name: str | None) -> bool: + return class_name == "AutoencoderKLMiniMaxH3" and bool(envs.FASTVIDEO_H3_VAE_NVFP4_DECODER.get()) + + @staticmethod + def _apply_h3_nvfp4_decoder(vae: nn.Module, path: str, target_device: torch.device) -> None: + """Swap in an exported NVFP4 MiniMax-H3 decoder (``FASTVIDEO_H3_VAE_NVFP4_DECODER``).""" + from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import ( + apply_nvfp4_decoder_checkpoint, + freeze_nvfp4_linears, + load_nvfp4_decoder_checkpoint, + ) + + if not os.path.isfile(path): + raise FileNotFoundError(f"FASTVIDEO_H3_VAE_NVFP4_DECODER={path} does not exist") + loaded_blocks = len(vae.decoder.transformer_blocks) + # The swap runs where the VAE was built (the CPU, see ``_defers_h3_nvfp4_decoder``); a pre-packed + # file never allocates the quantized linears' full-precision weights, a bf16 file packs on the GPU. + device = next(vae.parameters()).device + metadata = apply_nvfp4_decoder_checkpoint(vae, load_nvfp4_decoder_checkpoint(path, device), freeze=False) + vae.to(target_device) + if target_device.type == "cuda": + # Frozen layers hold only packed FP4 weights; packing needs the CUDA FP4 kernels. + freeze_nvfp4_linears(vae.decoder) + logger.info("MiniMax-H3 VAE decoder: NVFP4 %s (%d of %d blocks, act_scale=%s), bf16 decode autocast", path, + len(vae.decoder.transformer_blocks), loaded_blocks, metadata.get("act_scale")) + def load(self, model_path: str, fastvideo_args: FastVideoArgs): """Load the VAE based on the model path, and inference args.""" config = get_diffusers_config(model=model_path) @@ -948,7 +975,11 @@ class VAELoader(ComponentLoader): vae_config = fastvideo_args.pipeline_config.vae_config vae_config.update_model_arch(config) vae_cls, _ = ModelRegistry.resolve_model_cls(class_name) - vae = vae_cls(vae_config).to(target_device) + vae = vae_cls(vae_config) + # An NVFP4 H3 decoder is swapped in on the CPU first, so the dense decoder it replaces + # never reaches the GPU; the VAE moves to the target device after the swap. + if not self._defers_h3_nvfp4_decoder(class_name): + vae = vae.to(target_device) # Find all safetensors files safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors")) @@ -956,9 +987,13 @@ class VAELoader(ComponentLoader): if class_name == "AutoencoderKLMiniMaxH3": from fastvideo.models.vaes.minimax_h3_int8_convrot import ( dense_vae_safetensors, - find_int8_convrot_vae_path, + int8_convrot_overlay_to_apply, ) - int8_convrot_path = find_int8_convrot_vae_path(model_path) + int8_convrot_path = int8_convrot_overlay_to_apply(model_path) + if int8_convrot_path is not None and envs.FASTVIDEO_H3_VAE_NVFP4_DECODER.get(): + logger.info("Skipping MiniMax-H3 INT8 ConvRot VAE overlay %s: FASTVIDEO_H3_VAE_NVFP4_DECODER " + "replaces the decoder", int8_convrot_path) + int8_convrot_path = None safetensors_list = dense_vae_safetensors(safetensors_list) if not safetensors_list: raise ValueError(f"No safetensors files found in {model_path}") @@ -996,9 +1031,12 @@ class VAELoader(ComponentLoader): # Kandinsky6SRVAE: a partially loaded KVAE decodes plausible-looking garbage, so it is strict as well. strict_load = class_name in {"AutoencoderKL", "AutoencoderKLMiniMaxH3", "Kandinsky6SRVAE"} vae.load_state_dict(loaded, strict=strict_load) + del loaded if class_name == "AutoencoderKLMiniMaxH3" and int8_convrot_path is not None: from fastvideo.models.vaes.minimax_h3_int8_convrot import overlay_minimax_h3_int8_convrot_decoder overlay_minimax_h3_int8_convrot_decoder(vae, int8_convrot_path) + if class_name == "AutoencoderKLMiniMaxH3" and envs.FASTVIDEO_H3_VAE_NVFP4_DECODER.get(): + self._apply_h3_nvfp4_decoder(vae, envs.FASTVIDEO_H3_VAE_NVFP4_DECODER.get(), target_device) if (class_name == "AutoencoderKLWan" and getattr(vae.config, "use_light_vae", False) and target_device.type == "cuda" and hasattr(vae, "optimize_memory_format")): vae.optimize_memory_format() @@ -1194,6 +1232,11 @@ class TransformerLoader(ComponentLoader): # Config from Diffusers supersedes fastvideo's model config dit_config = deepcopy(fastvideo_args.pipeline_config.dit_config) dit_config.update_model_arch(config) + if getattr(fastvideo_args, "h3_vsa_fp4", None): + from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3Config + if not isinstance(dit_config, MiniMaxH3Config): + raise ValueError("h3_vsa_fp4 is only supported by MiniMax-H3") + dit_config.vsa_fp4 = dict(fastvideo_args.h3_vsa_fp4) # Generator-only QAT for DMD distillation: the teacher (real_score) and # critic (fake_score) transformers load with this flag set and must stay diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index f93032ae6..d751f8d33 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -6,6 +6,7 @@ from __future__ import annotations import contextlib +import copy import re from collections.abc import Callable, Generator from itertools import chain @@ -67,9 +68,10 @@ def _prequantized_fp8_prefixes(weight_files: list[str]) -> list[str]: def _load_prequantized_fp8(model: nn.Module, weight_files: list[str], targets: dict[str, str], device: torch.device) -> None: - """Attach FP8 W8A8 checkpoint tensors to their FP8 linears as ``convert_model_to_fp8`` would (per channel).""" + """Load packed FP8 linears and preserve the checkpoint's dense linear weights.""" from safetensors import safe_open + from fastvideo.layers.linear import UnquantizedLinearMethod from fastvideo.layers.quantization.fp8_config import FP8QuantizeMethod modules = dict(model.named_modules()) @@ -92,7 +94,17 @@ def _load_prequantized_fp8(model: nn.Module, weight_files: list[str], targets: d tensors[f"{prefix}.weight_scale"].reshape(-1).to(device=device, dtype=torch.float32), persistent=False) module._parameters.pop("weight", None) - logger.info("Loaded %d pre-quantized FP8 linears", len(targets)) + # The packed entries define the export's FP8 scope. A broad FP8Config can + # also tag intentionally dense layers, such as H3's text refiner. Reset + # their method before post-load conversion also quantizes those layers. + packed_modules = {target[:-len(".weight")] for target in targets.values()} + dense_linears = 0 + for name, module in modules.items(): + if (name not in packed_modules and isinstance(getattr(module, "quant_method", None), FP8QuantizeMethod) + and getattr(module, "weight", None) is not None): + module.quant_method = UnquantizedLinearMethod() + dense_linears += 1 + logger.info("Loaded %d pre-quantized FP8 linears; preserved %d dense linears", len(targets), dense_linears) def _maybe_quantize_model(model: nn.Module, *, defer_weight_conversion_until_lora_merge: bool = False) -> None: @@ -228,6 +240,32 @@ def _prepare_model_for_compile(model: nn.Module, *, regional: bool) -> str | Non return unsupported if isinstance(unsupported, str) and unsupported else None +def _adopt_nvfp4_export_skip_blocks(init_params: dict[str, Any], quant_config: Any, + export_path: str) -> dict[str, Any]: + """Match the NVFP4 config's bf16 blocks to the ones the packed export declares. + + A mixed export (e.g. FastH3 OmniRef QAD keeps blocks 42-49 in bf16) leaves + those blocks' linears out of the packed file; they must not be tagged NVFP4, + or the loader would expect packed buffers for them. An empty + ``skip_blocks`` adopts the export's list; a different explicit list is an + error rather than a silent precision change. + """ + from fastvideo.layers.quantization.nvfp4_config import read_minimax_h3_nvfp4_dit_skip_blocks + + export_blocks = read_minimax_h3_nvfp4_dit_skip_blocks(export_path) + config_blocks = tuple(getattr(quant_config, "skip_blocks", ()) or ()) + if export_blocks == config_blocks: + return init_params + if config_blocks: + raise ValueError(f"NVFP4Config.skip_blocks={list(config_blocks)} but the packed export {export_path} keeps " + f"blocks {list(export_blocks)} in bf16; drop skip_blocks to follow the export.") + logger.info("Packed NVFP4 DiT export keeps blocks %s in bf16; leaving their linears unquantized", + list(export_blocks)) + transformer_config = copy.copy(init_params["config"]) + transformer_config.quant_config = quant_config.with_skip_blocks(export_blocks) + return {**init_params, "config": transformer_config} + + def _validate_fsdp_inference_quantization(init_params: dict[str, Any], fsdp_inference: bool) -> None: """Enforce the transformer-quantization allowlist for FSDP inference.""" transformer_config = init_params.get("config") @@ -311,8 +349,17 @@ def maybe_load_fsdp_model( weight_dir_list = dense_transformer_safetensors(weight_dir_list) quant_config = getattr(init_params.get("config"), "quant_config", None) packed_profiles = ("h3_dit", "h3_dit_ffn", "h3_dit_vsa") - packed_nvfp4_export = (packed_candidate - if getattr(quant_config, "layer_profile", None) in packed_profiles else None) + packed_layer_profile = getattr(quant_config, "layer_profile", None) + packed_nvfp4_export = packed_candidate if packed_layer_profile in packed_profiles else None + if packed_candidate is not None and packed_nvfp4_export is None: + # A packed export leaves its quantized linears out of the dense shards, so + # loading it with any other profile misses those weights. + raise ValueError( + f"{packed_candidate} is a packed MiniMax-H3 NVFP4 export, which keeps the quantized linears out of " + f"the dense shards, but the quantization profile is {packed_layer_profile!r}. Load it with an NVFP4 " + f"config (layer_profile one of {', '.join(packed_profiles)}), or point the run at a dense export.") + if packed_nvfp4_export is not None: + init_params = _adopt_nvfp4_export_skip_blocks(init_params, quant_config, packed_nvfp4_export) # NOTE(will): cast_forward_inputs=True shouldn't be needed as we are # manually casting the inputs to the model @@ -403,11 +450,13 @@ def maybe_load_fsdp_model( # while each rank materializes its shard, or the full checkpoint and shards # peak together on the accelerator. Training keeps GPU-direct loading. # A packed NVFP4 export is always staged on CPU. - load_weights_to_cpu = _should_stage_transformer_weights_on_cpu( + # Layerwise offload passes device=cpu but disables dit_cpu_offload, so the + # device itself must select CPU staging or the iterator would fill the GPU. + load_weights_to_cpu = (device.type == "cpu" or _should_stage_transformer_weights_on_cpu( cpu_offload=cpu_offload, fsdp_inference=fsdp_inference, has_unified_memory=has_unified_memory, - ) or packed_nvfp4_export is not None + ) or packed_nvfp4_export is not None) weight_iterator = safetensors_weights_iterator(weight_dir_list, to_cpu=load_weights_to_cpu) logger.info("Loading transformer weights with to_cpu=%s", load_weights_to_cpu) param_names_mapping_fn = get_param_names_mapping(model.param_names_mapping) @@ -952,8 +1001,8 @@ def load_model_from_full_model_state_dict( named_parameters = {_strip_checkpoint_wrapper_prefix(k): v for k, v in model.named_parameters()} named_buffers = {_strip_checkpoint_wrapper_prefix(k): v for k, v in model.named_buffers()} sharded_sd = {} - custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict(full_sd_iterator, - param_names_mapping) # type: ignore + custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict( + full_sd_iterator, param_names_mapping, skip_names=skip_param_names) # type: ignore checkpoint_key_aliases = getattr(model, "_fastvideo_checkpoint_key_aliases", {}) if checkpoint_key_aliases: for original_name, transformed_name in checkpoint_key_aliases.items(): diff --git a/fastvideo/models/loader/utils.py b/fastvideo/models/loader/utils.py index 4783f699e..ad94d8167 100644 --- a/fastvideo/models/loader/utils.py +++ b/fastvideo/models/loader/utils.py @@ -58,7 +58,8 @@ def get_param_names_mapping(mapping_dict: dict[str, str]) -> Callable[[str], tup def hf_to_custom_state_dict( hf_param_sd: dict[str, torch.Tensor] | Iterator[tuple[str, torch.Tensor]], - param_names_mapping: Callable[[str], tuple[str, Any, Any]] + param_names_mapping: Callable[[str], tuple[str, Any, Any]], + skip_names: set[str] | None = None, ) -> tuple[dict[str, torch.Tensor], dict[str, tuple[str, Any, Any]]]: """ Converts a Hugging Face parameter state dictionary to a custom parameter state dictionary. @@ -78,6 +79,10 @@ def hf_to_custom_state_dict( hf_param_sd = hf_param_sd.items() # type: ignore for source_param_name, full_tensor in hf_param_sd: # type: ignore target_param_name, merge_index, num_params_to_merge = param_names_mapping(source_param_name) + if skip_names and target_param_name in skip_names: + # Release the source tensor immediately instead of retaining every + # skipped AdaLN projection in the drained dict (24 GiB on OmniRef). + continue reverse_param_names_mapping[target_param_name] = (source_param_name, merge_index, num_params_to_merge) if merge_index is not None: to_merge_params[target_param_name][merge_index] = full_tensor diff --git a/fastvideo/models/pinned_offload.py b/fastvideo/models/pinned_offload.py index 08c45802a..66c0e4b9a 100644 --- a/fastvideo/models/pinned_offload.py +++ b/fastvideo/models/pinned_offload.py @@ -77,6 +77,17 @@ def load(module: nn.Module, device: torch.device, pin: bool = True) -> nn.Module return module +def refresh_host_copies(module: nn.Module, pin: bool = True) -> None: + """Rebuild an existing host mirror after frozen weights or buffers change. + + Resident modules have no mirror and do not need one. Release old masters + before copying newly packed buffers so both representations do not persist. + """ + if getattr(module, _HOST_ATTR, None) is not None: + delattr(module, _HOST_ATTR) + _host_copies(module, pin and is_pin_memory_available()) + + def unload(module: nn.Module) -> nn.Module: """Point a frozen ``module`` back at its host copies; no device-to-host copy.""" host = getattr(module, _HOST_ATTR, None) @@ -85,4 +96,10 @@ def unload(module: nn.Module) -> nn.Module: for name, t in _tensors(module): if t.data.device.type != "cpu": t.data = host[name] + # Parameters and buffers are not the whole story: a submodule may derive device + # tensors into a plain attribute, which would otherwise stay on the GPU. + for submodule in module.modules(): + clear = getattr(submodule, "clear_offload_caches", None) + if clear is not None: + clear() return module diff --git a/fastvideo/models/vaes/minimax_h3_int8_convrot.py b/fastvideo/models/vaes/minimax_h3_int8_convrot.py index d9a12a01f..0681a6450 100644 --- a/fastvideo/models/vaes/minimax_h3_int8_convrot.py +++ b/fastvideo/models/vaes/minimax_h3_int8_convrot.py @@ -351,3 +351,13 @@ def dense_vae_safetensors(paths: list[str]) -> list[str]: def find_int8_convrot_vae_path(model_path: str | Path) -> Path | None: candidate = Path(model_path) / INT8_CONVROT_FILENAME return candidate if candidate.is_file() else None + + +def int8_convrot_overlay_to_apply(model_path: str | Path) -> Path | None: + """The ConvRot overlay shipped in ``model_path``, unless ``FASTVIDEO_H3_VAE_INT8_OVERLAY=0`` skips it.""" + path = find_int8_convrot_vae_path(model_path) + if path is not None and not envs.FASTVIDEO_H3_VAE_INT8_OVERLAY.get(): + logger.info("Skipping MiniMax-H3 INT8 ConvRot VAE overlay %s (FASTVIDEO_H3_VAE_INT8_OVERLAY=0); " + "decoding with the dense weights", path) + return None + return path diff --git a/fastvideo/models/vaes/minimax_h3_int8_kernels.py b/fastvideo/models/vaes/minimax_h3_int8_kernels.py index e1e5a9fad..d00e8b722 100644 --- a/fastvideo/models/vaes/minimax_h3_int8_kernels.py +++ b/fastvideo/models/vaes/minimax_h3_int8_kernels.py @@ -5,6 +5,7 @@ from __future__ import annotations import torch import triton import triton.language as tl +from triton.language.extra.cuda import libdevice @triton.jit @@ -16,10 +17,11 @@ def _dequant_bias(Acc, XScale, WScale, Bias, Out, M: tl.constexpr, N: tl.constex acc = tl.load(Acc + offsets, valid, 0).to(tl.float32) x_scale = tl.load(XScale + rows * XS, valid, 0).to(tl.float32) weight_scale = tl.load(WScale + cols * WS).to(tl.float32) - out = acc * x_scale - out = out * weight_scale + # Keep each FP32 rounding boundary explicit through vectorization and PTX lowering. + out = libdevice.mul_rn(acc, x_scale) + out = libdevice.mul_rn(out, weight_scale) if HAS_BIAS: - out = out + tl.load(Bias + cols).to(tl.float32) + out = libdevice.add_rn(out, tl.load(Bias + cols).to(tl.float32)) tl.store(Out + offsets, out, valid) diff --git a/fastvideo/models/vaes/minimax_h3_nvfp4_checkpoint.py b/fastvideo/models/vaes/minimax_h3_nvfp4_checkpoint.py new file mode 100644 index 000000000..48e7f7875 --- /dev/null +++ b/fastvideo/models/vaes/minimax_h3_nvfp4_checkpoint.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Safetensors files for exported NVFP4 MiniMax-H3 video decoders. + +Two formats share one file layout: tensors are prefixed ``decoder.`` or +``post_quant_conv.``; the header holds ``format`` and ``nvfp4_metadata`` (JSON). + +* ``fastvideo_h3_decoder_nvfp4_deploy_v1``: NVFP4 linears as bf16 master + weights, packed at load (``NVFP4DecoderLinear.freeze``). +* ``fastvideo_h3_decoder_nvfp4_packed_v1``: the frozen layers' buffers as + written by ``freeze`` on a Blackwell GPU, loaded as is. Per NVFP4 linear + ````: ``.packed_weight`` (uint8 ``[out, in/2]``, two E2M1 codes per + byte, FlashInfer ``fp4_quantize`` order), ``.weight_inv_scale`` (E4M3 + per-16 block scales, FlashInfer ``SfLayout.layout_128x4`` swizzle), + ``.weight_global_sf`` (fp32 scalar), ``.bias`` (in the compute + dtype, which is exactly how inference consumes it) and ``.input_amax`` + (fp32 static activation range). Every other tensor keeps the dtype it is + stored with in the trained decoder, so outputs stay bit-identical. +""" +from __future__ import annotations + +import json +from typing import Any + +import torch +import torch.nn as nn + +from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import ( + NVFP4_DECODER_DEPLOY_FORMAT, + NVFP4DecoderLinear, + check_num_linears, + enable_fused_blocks, + nvfp4_decoder_linear_names, +) + +NVFP4_DECODER_PACKED_FORMAT = "fastvideo_h3_decoder_nvfp4_packed_v1" +NVFP4_DECODER_FORMATS = (NVFP4_DECODER_DEPLOY_FORMAT, NVFP4_DECODER_PACKED_FORMAT) +PACKED_BUFFERS = ("packed_weight", "weight_inv_scale", "weight_global_sf") +GROUPS = ("decoder", "post_quant_conv") +PACKED_LAYOUT = { + "packed_weight": "uint8 [out_features, in_features/2], two E2M1 codes per byte (FlashInfer fp4_quantize)", + "weight_inv_scale": "E4M3 per-16-element block scales, FlashInfer SfLayout.layout_128x4 (swizzled)", + "weight_global_sf": "fp32 scalar, (448 * 6) / amax(bf16 weight)", + "bias": "compute dtype (bf16)", + "input_amax": "fp32 scalar static activation range (required only when act_scale == 'static')", + "other_tensors": "stored dtype of the trained decoder (fp32 norms, projections, registers, post_quant_conv)", +} + + +def write_nvfp4_decoder_safetensors(path: str, checkpoint: dict[str, Any]) -> None: + """Write ``{"format", "metadata", "decoder", "post_quant_conv"}`` as one safetensors file.""" + from safetensors.torch import save_file + + if checkpoint["format"] not in NVFP4_DECODER_FORMATS: + raise ValueError(f"unknown NVFP4 decoder format {checkpoint['format']!r}") + tensors = { + f"{group}.{key}": value.detach().contiguous().cpu().clone() + for group in GROUPS for key, value in checkpoint[group].items() + } + header = {"format": checkpoint["format"], "nvfp4_metadata": json.dumps(checkpoint["metadata"])} + save_file(tensors, path, metadata=header) + + +def read_nvfp4_decoder_safetensors(path: str, device: torch.device | str = "cpu") -> dict[str, Any]: + """Read either safetensors format; packed tensors go straight to ``device`` (bf16 masters stay on CPU).""" + from safetensors import safe_open + from safetensors.torch import load_file + + with safe_open(path, framework="pt") as handle: + header = handle.metadata() or {} + fmt = header.get("format") + if fmt not in NVFP4_DECODER_FORMATS: + raise ValueError(f"{path} is not an NVFP4 decoder safetensors file (format: {fmt!r}); " + f"expected one of {NVFP4_DECODER_FORMATS}") + if "nvfp4_metadata" not in header: + raise ValueError(f"{path} is missing the nvfp4_metadata header") + tensors = load_file(path, device=str(device) if fmt == NVFP4_DECODER_PACKED_FORMAT else "cpu") + checkpoint: dict[str, Any] = {"format": fmt, "metadata": json.loads(header["nvfp4_metadata"])} + for group in GROUPS: + prefix = f"{group}." + checkpoint[group] = {key[len(prefix):]: value for key, value in tensors.items() if key.startswith(prefix)} + unknown = [key for key in tensors if not key.startswith(tuple(f"{group}." for group in GROUPS))] + if unknown: + raise ValueError(f"{path} has tensors outside {GROUPS}: {unknown[:5]}") + return checkpoint + + +def packed_nvfp4_decoder_checkpoint(vae: nn.Module, metadata: dict[str, Any]) -> dict[str, Any]: + """The packed-format checkpoint of a VAE whose NVFP4 decoder linears are all frozen.""" + decoder = vae.decoder + state = dict(decoder.state_dict()) + layers = {name: module for name, module in decoder.named_modules() if isinstance(module, NVFP4DecoderLinear)} + if not layers: + raise ValueError("the VAE decoder has no NVFP4 linears") + for name, layer in layers.items(): + if not layer.frozen: + raise ValueError(f"NVFP4 linear {name} is not frozen; freeze the decoder before exporting packed weights") + for buffer in PACKED_BUFFERS: + state[f"{name}.{buffer}"] = getattr(layer, buffer) + if layer.bias is not None: + state[f"{name}.bias"] = layer.bias.detach().to(layer.compute_dtype) + packed_metadata = {**metadata, "packed_layout": PACKED_LAYOUT, "num_linears": len(layers)} + return { + "format": NVFP4_DECODER_PACKED_FORMAT, + "metadata": packed_metadata, + "decoder": state, + "post_quant_conv": dict(vae.post_quant_conv.state_dict()), + } + + +def apply_packed_nvfp4_decoder(vae: nn.Module, checkpoint: dict[str, Any]) -> None: + """Replace the (already depth-cut) dense decoder linears with frozen layers built from packed tensors. + + The dense linears are dropped one at a time, so no full-precision copy of a quantized weight is + ever made. Everything else is loaded strictly. + """ + metadata = checkpoint["metadata"] + decoder = vae.decoder + state = dict(checkpoint["decoder"]) + names = nvfp4_decoder_linear_names(decoder, tuple(metadata.get("skip_blocks") or ())) + check_num_linears(metadata, len(names)) + act_scale = metadata.get("act_scale", "dynamic") + for name in names: + if act_scale != "static": + # Only static scaling reads input_amax, so a file may omit it; the strict load still needs the key. + state.setdefault(f"{name}.input_amax", torch.zeros((), dtype=torch.float32)) + required = PACKED_BUFFERS + ("input_amax", ) + missing = [f"{name}.{buffer}" for buffer in required if f"{name}.{buffer}" not in state] + if missing: + raise ValueError(f"the packed NVFP4 decoder does not fit this VAE's architecture: missing {missing}") + parent_name, _, child = name.rpartition(".") + parent = decoder.get_submodule(parent_name) + dense = parent[int(child)] if child.isdigit() else getattr(parent, child) + packed = state.pop(f"{name}.packed_weight") + if packed.shape[0] != dense.out_features: + raise ValueError(f"the packed NVFP4 decoder does not fit this VAE's architecture: {name} has " + f"{packed.shape[0]} outputs, the VAE {dense.out_features}") + layer = NVFP4DecoderLinear.from_packed(packed, + state.pop(f"{name}.weight_inv_scale"), + state.pop(f"{name}.weight_global_sf"), + state.get(f"{name}.bias"), + state[f"{name}.input_amax"], + in_features=dense.in_features, + rotation_group=metadata.get("rotation_group"), + compute_dtype=torch.bfloat16, + act_scale=act_scale) + if child.isdigit(): + parent[int(child)] = layer + else: + setattr(parent, child, layer) + del dense + enable_fused_blocks(decoder) + try: + decoder.load_state_dict(state, strict=True) + vae.post_quant_conv.load_state_dict(checkpoint["post_quant_conv"], strict=True) + except RuntimeError as error: + raise ValueError(f"the packed NVFP4 decoder does not fit this VAE's architecture: {error}") from error + vae.requires_grad_(False) diff --git a/fastvideo/models/vaes/minimax_h3_nvfp4_decoder.py b/fastvideo/models/vaes/minimax_h3_nvfp4_decoder.py new file mode 100644 index 000000000..8d41312e3 --- /dev/null +++ b/fastvideo/models/vaes/minimax_h3_nvfp4_decoder.py @@ -0,0 +1,497 @@ +# SPDX-License-Identifier: Apache-2.0 +"""NVFP4 linears for the MiniMax-H3 ViT video decoder. + +Training and inference share one numeric path, so a decoder distilled with +quantization-aware training runs at deployment exactly as it was trained: + +* activations: per-16 E4M3 block scales under a per-tensor global scale that is + ``dynamic`` (abs-max of every call), ``unit`` (1.0, the H3 DiT scheme), or + ``static`` (a per-layer abs-max calibrated once and stored with the weights); +* weights: the same scheme, quantized once at inference and cached; +* GEMM: ``flashinfer.mm_fp4`` (``mm_fp4_backend``), bias added in the activation dtype. + +Training runs through ``_NVFP4DecoderSTE``, whose forward *is* the inference +function and whose backward is full precision (straight-through). An optional block Hadamard rotation (the regular +Hadamard used by the INT8 ConvRot overlay) is folded into the weight along the +input dimension and applied to the activation at runtime; it is exact in full +precision and spreads activation outliers before FP4 rounding. + +Only the transformer-block linears are replaced. ``proj_in`` (K = 24, not +FP4-blockable), ``proj_out``, norms, register tokens and attention stay as they +are. +""" +from __future__ import annotations + +import re +from typing import Any, NamedTuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from fastvideo.models.vaes.minimax_h3_int8_convrot import rotate_activation + +NVFP4_BLOCK_SIZE = 16 +# Block linears of MiniMaxH3VideoTransformerBlock, relative to ``decoder.``. +DECODER_BLOCK_LINEAR = re.compile(r"^transformer_blocks\.(\d+)\.(attn\.to_q|attn\.to_k|attn\.to_v|attn\.to_out\.0" + r"|ff\.net\.0\.proj|ff\.net\.2)$") + + +def _quantized_weight(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Quantize a weight the way ``_LinearFWD4BWD16Fn`` does on every training step.""" + from fastvideo.layers.fp4linear import _global_sf, _require_flashinfer + from fastvideo.layers.quantization.nvfp4_config import nvfp4_quantize_fenced + + flashinfer_mod = _require_flashinfer() + global_sf = _global_sf(weight) + packed, inv_scale = nvfp4_quantize_fenced(weight, global_sf, flashinfer_mod.SfLayout.layout_128x4.value) + return packed, inv_scale, global_sf + + +def mm_fp4_backend(device: torch.device) -> str: + """FlashInfer ``mm_fp4`` backend for one decoder GEMM; the backends return identical outputs. + + DGX Spark (GB10, sm_121) uses cuDNN: up to 1.4x faster than CUTLASS on the tile-batched + FFN input projection and on par elsewhere. Other GPUs use CUTLASS. + """ + from fastvideo.layers.quantization.nvfp4_config import _is_dgx_spark + + if device.type != "cuda": + return "cutlass" + index = device.index if device.index is not None else torch.cuda.current_device() + return "cudnn" if _is_dgx_spark(index) else "cutlass" + + +ACT_SCALES = ("dynamic", "unit", "static") +FP4_E4M3_RANGE = 448.0 * 6.0 + + +def nvfp4_linear_inference(x: torch.Tensor, packed: torch.Tensor, inv_scale: torch.Tensor, global_sf_w: torch.Tensor, + bias: torch.Tensor | None, out_features: int, + global_sf_x: torch.Tensor | None) -> torch.Tensor: + """FP4 GEMM. ``global_sf_x=None`` computes the activation scale dynamically.""" + from fastvideo.layers.fp4linear import _global_sf, _require_flashinfer + from fastvideo.layers.quantization.nvfp4_config import nvfp4_quantize_fenced + + flashinfer_mod = _require_flashinfer() + orig_shape = x.shape + x2d = x.reshape(-1, orig_shape[-1]).contiguous() + if global_sf_x is None: + global_sf_x = _global_sf(x2d) + x_fp4, x_inv_scale = nvfp4_quantize_fenced(x2d, global_sf_x, flashinfer_mod.SfLayout.layout_128x4.value) + out = torch.empty((x2d.shape[0], out_features), device=x.device, dtype=x.dtype) + flashinfer_mod.mm_fp4( + x_fp4, + packed.T, + x_inv_scale, + inv_scale.T, + 1.0 / (global_sf_x * global_sf_w), + x.dtype, + out, + block_size=NVFP4_BLOCK_SIZE, + use_8x4_sf_layout=False, + backend=mm_fp4_backend(x.device), + ) + if bias is not None: + out.add_(bias.to(x.dtype)) + return out.reshape(*orig_shape[:-1], out_features) + + +class NVFP4FusedState(NamedTuple): + """Everything the fused inference kernels need from one ``NVFP4DecoderLinear``. + + Each tensor is computed with the same expression as ``nvfp4_linear_inference``, so a GEMM + fed from here matches the eager linear bit for bit. + """ + packed: torch.Tensor + inv_scale: torch.Tensor + global_sf_x: torch.Tensor + alpha: torch.Tensor + bias: torch.Tensor + out_features: int + + +class _NVFP4DecoderSTE(torch.autograd.Function): + """Forward = ``nvfp4_linear_inference`` on a freshly quantized weight; backward = full precision.""" + + @staticmethod + def forward(ctx, x, weight, bias, global_sf_x): # type: ignore[override] + packed = _quantized_weight(weight.detach().to(x.dtype)) + out = nvfp4_linear_inference(x, *packed, bias, weight.shape[0], global_sf_x) + ctx.save_for_backward(x, weight) + ctx.has_bias = bias is not None + return out + + @staticmethod + def backward(ctx, grad_out): # type: ignore[override] + x, weight = ctx.saved_tensors + grad_2d = grad_out.reshape(-1, grad_out.shape[-1]) + grad_x = (grad_2d @ weight.to(grad_2d.dtype)).reshape(x.shape) + grad_w = grad_2d.t() @ x.reshape(-1, x.shape[-1]).to(grad_2d.dtype) + grad_b = grad_2d.sum(dim=0) if ctx.has_bias else None + return grad_x, grad_w.to(weight.dtype), grad_b, None + + +class NVFP4DecoderLinear(nn.Module): + """Drop-in for an ``nn.Linear`` in the H3 decoder blocks. + + ``weight`` is the full-precision master (already rotated when + ``rotation_group`` is set). With gradients enabled the forward is the STE; + otherwise the cached packed weight is used, so call ``invalidate()`` after + any weight update (the trainer does after each optimizer step). + + ``act_scale="static"`` uses the persistent ``input_amax`` buffer. Set + ``calibrating = True`` to run dynamically while folding each call's + abs-max into it. + + ``freeze()`` turns the layer inference-only: the packed weight becomes a + buffer and the master ``weight`` is dropped (``weight is None``). + """ + + def __init__(self, weight: torch.Tensor, bias: torch.Tensor | None, rotation_group: int | None, + compute_dtype: torch.dtype, act_scale: str = "dynamic") -> None: + super().__init__() + out_features, in_features = weight.shape + if in_features % NVFP4_BLOCK_SIZE: + raise ValueError(f"NVFP4 needs in_features divisible by {NVFP4_BLOCK_SIZE}, got {in_features}") + if rotation_group is not None and in_features % rotation_group: + raise ValueError(f"rotation group {rotation_group} does not divide in_features {in_features}") + if act_scale not in ACT_SCALES: + raise ValueError(f"act_scale must be one of {ACT_SCALES}, got {act_scale!r}") + self.in_features = in_features + self.out_features = out_features + self.rotation_group = rotation_group + self.compute_dtype = compute_dtype + self.act_scale = act_scale + self.calibrating = False + # Run the master weight densely (bf16) instead of NVFP4; for sensitivity analysis and mixed precision. + self.dense_bypass = False + self.weight = nn.Parameter(weight.detach().clone()) + self.bias = nn.Parameter(bias.detach().clone()) if bias is not None else None + self.register_buffer("input_amax", torch.zeros((), dtype=torch.float32, device=weight.device)) + self._packed: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None + self._fused: tuple[tuple[int, int], NVFP4FusedState] | None = None + + @classmethod + def from_linear(cls, linear: nn.Linear, *, rotation_group: int | None, compute_dtype: torch.dtype, + act_scale: str = "dynamic") -> NVFP4DecoderLinear: + weight = linear.weight.detach().float() + if rotation_group is not None: + # y = x W^T = (x H)(W H)^T for orthonormal symmetric H: rotate W along its input dim. + weight = rotate_activation(weight, rotation_group) + bias = linear.bias.detach().float() if linear.bias is not None else None + return cls(weight, bias, rotation_group, compute_dtype, act_scale) + + @classmethod + def from_packed(cls, packed: torch.Tensor, inv_scale: torch.Tensor, global_sf_w: torch.Tensor, + bias: torch.Tensor | None, input_amax: torch.Tensor, *, in_features: int, + rotation_group: int | None, compute_dtype: torch.dtype, act_scale: str) -> NVFP4DecoderLinear: + """A frozen layer built from already packed tensors (``freeze()``'s buffers), with no master weight. + + ``packed`` holds two FP4 codes per byte, so its shape is ``[out_features, in_features // 2]``. + Nothing full-precision is allocated: the constructor's placeholder weight lives on ``meta``. + """ + out_features = packed.shape[0] + if packed.dtype != torch.uint8 or packed.shape[1] * 2 != in_features: + raise ValueError(f"packed NVFP4 weight must be uint8 [{out_features}, {in_features // 2}], " + f"got {packed.dtype} {list(packed.shape)}") + layer = cls(torch.empty(out_features, in_features, device="meta"), None, rotation_group, compute_dtype, + act_scale) + del layer.weight + layer.register_parameter("weight", None) + layer.bias = nn.Parameter(bias, requires_grad=False) if bias is not None else None + layer.register_buffer("input_amax", input_amax.to(torch.float32)) + layer.register_buffer("packed_weight", packed, persistent=False) + layer.register_buffer("weight_inv_scale", inv_scale, persistent=False) + layer.register_buffer("weight_global_sf", global_sf_w, persistent=False) + return layer + + @property + def frozen(self) -> bool: + return self.weight is None + + def invalidate(self) -> None: + if self.frozen: + raise RuntimeError("a frozen NVFP4DecoderLinear has no master weight to re-quantize") + self._packed = None + self._fused = None + + def _packed_weight(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if self.frozen: + return self.packed_weight, self.weight_inv_scale, self.weight_global_sf + if self._packed is None: + self._packed = _quantized_weight(self.weight.detach().to(self.compute_dtype)) + return self._packed + + @torch.no_grad() + def freeze(self) -> None: + """Pack the weight once and drop the full-precision master (inference only, irreversible). + + Resident memory falls to the packed FP4 weight and its scales. Outputs are unchanged: the + same packed tensors feed the eager GEMM and ``fused_state()``. Buffers are non-persistent, + so a frozen layer's ``state_dict()`` holds only ``bias`` and ``input_amax``. + """ + if self.frozen: + return + if self.dense_bypass or self.calibrating: + raise RuntimeError("cannot freeze an NVFP4DecoderLinear in dense-bypass or calibration mode") + packed, inv_scale, global_sf_w = self._packed_weight() + del self.weight + self.register_parameter("weight", None) + self.register_buffer("packed_weight", packed, persistent=False) + self.register_buffer("weight_inv_scale", inv_scale, persistent=False) + self.register_buffer("weight_global_sf", global_sf_w, persistent=False) + self._packed = None + self._fused = None + + def fused_state(self) -> NVFP4FusedState | None: + """Cached inference state for the fused decoder kernels, or None when they cannot reproduce this layer. + + The fused kernels quantize with a precomputed global scale, so a dynamic (per-call abs-max) + scale, calibration, and the activation rotation keep the eager path. + """ + if (self.dense_bypass or self.rotation_group is not None or self.calibrating or self.act_scale == "dynamic" + or self.compute_dtype != torch.bfloat16 or self.bias is None): + return None + packed, inv_scale, global_sf_w = self._packed_weight() + # The data pointer re-keys the cache when a frozen layer's buffers move (``module.to``). + version = (self.input_amax._version, packed.data_ptr()) + if self._fused is None or self._fused[0] != version: + global_sf_x = self._activation_global_sf(self.input_amax) + alpha = 1.0 / (global_sf_x * global_sf_w) + state = NVFP4FusedState(packed, inv_scale, global_sf_x, alpha, self.bias.detach().to(self.compute_dtype), + self.out_features) + self._fused = (version, state) + return self._fused[1] + + def clear_offload_caches(self) -> None: + """Drop the packed and fused caches, which are plain attributes and so invisible to offload. + + ``pinned_offload`` moves parameters and buffers only; these caches hold device + tensors, so without this they stay on the GPU after an unload and accumulate + across clips. The data-pointer keys re-create them on the next forward. + """ + self._packed = None + self._fused = None + + def extra_repr(self) -> str: + return (f"in_features={self.in_features}, out_features={self.out_features}, " + f"rotation_group={self.rotation_group}, act_scale={self.act_scale}, frozen={self.frozen}") + + def _activation_global_sf(self, x: torch.Tensor) -> torch.Tensor | None: + if self.calibrating: + self.input_amax.copy_(torch.maximum(self.input_amax, x.detach().abs().amax().float())) + return None + if self.act_scale == "dynamic": + return None + if self.act_scale == "unit": + return torch.ones((), dtype=torch.float32, device=x.device) + return FP4_E4M3_RANGE / self.input_amax.clamp(min=1e-12) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x.to(self.compute_dtype) + if self.rotation_group is not None: + x = rotate_activation(x, self.rotation_group) + if self.dense_bypass: + if self.frozen: + raise RuntimeError("a frozen NVFP4DecoderLinear has no master weight for the dense bypass") + bias = self.bias.to(x.dtype) if self.bias is not None else None + return F.linear(x, self.weight.to(x.dtype), bias) + global_sf_x = self._activation_global_sf(x) + if torch.is_grad_enabled() and not self.frozen and self.weight.requires_grad: + self.invalidate() + return _NVFP4DecoderSTE.apply(x, self.weight, self.bias, global_sf_x) + return nvfp4_linear_inference(x, *self._packed_weight(), self.bias, self.out_features, global_sf_x) + + +def nvfp4_decoder_linear_names(decoder: nn.Module, skip_blocks: tuple[int, ...] = ()) -> list[str]: + """Names (relative to the decoder) of the block linears to replace.""" + names = [] + for name, module in decoder.named_modules(): + match = DECODER_BLOCK_LINEAR.match(name) + if match and type(module) is nn.Linear and int(match.group(1)) not in skip_blocks: + names.append(name) + return names + + +def convert_decoder_to_nvfp4(decoder: nn.Module, + *, + rotation_group: int | None = None, + skip_blocks: tuple[int, ...] = (), + compute_dtype: torch.dtype = torch.bfloat16, + act_scale: str = "dynamic") -> list[str]: + """Replace the decoder's block linears with ``NVFP4DecoderLinear``; return the replaced names.""" + names = nvfp4_decoder_linear_names(decoder, skip_blocks) + if not names: + raise ValueError("No MiniMax-H3 decoder block linears found to convert") + for name in names: + parent_name, _, child = name.rpartition(".") + parent = decoder.get_submodule(parent_name) + original = getattr(parent, child) if not child.isdigit() else parent[int(child)] + replacement = NVFP4DecoderLinear.from_linear(original, + rotation_group=rotation_group, + compute_dtype=compute_dtype, + act_scale=act_scale) + if child.isdigit(): + parent[int(child)] = replacement + else: + setattr(parent, child, replacement) + enable_fused_blocks(decoder) + return names + + +def enable_fused_blocks(decoder: nn.Module) -> None: + """Run eligible NVFP4 block stacks through the bit-exact fused kernels (see that module).""" + try: + from fastvideo.models.vaes.minimax_h3_nvfp4_fused import fused_nvfp4_blocks_forward + except ImportError: + # Triton-less hosts (CPU, macOS) keep the eager NVFP4 path; the fused kernels are an exact speedup only. + return + decoder.fused_blocks_forward = fused_nvfp4_blocks_forward + + +def nvfp4_linears(module: nn.Module) -> list[NVFP4DecoderLinear]: + """Every NVFP4 decoder linear under ``module``; empty for a stub that is not an ``nn.Module``.""" + if not isinstance(module, nn.Module): + return [] + return [sub for sub in module.modules() if isinstance(sub, NVFP4DecoderLinear)] + + +@torch.no_grad() +def calibrate_static_scales(module: nn.Module, run_batches, margin: float = 1.0) -> int: + """Record each NVFP4 linear's activation abs-max over ``run_batches()`` (dynamic numerics while recording). + + ``margin`` > 1 widens the stored range. Returns the number of calibrated linears. + """ + layers = nvfp4_linears(module) + for layer in layers: + layer.input_amax.zero_() + layer.calibrating = True + try: + run_batches() + finally: + for layer in layers: + layer.calibrating = False + for layer in layers: + if layer.input_amax.item() <= 0: + raise RuntimeError("an NVFP4 decoder linear saw no activations during calibration") + layer.input_amax.mul_(margin) + return len(layers) + + +def freeze_nvfp4_linears(module: nn.Module) -> int: + """``freeze()`` every NVFP4 linear under ``module`` (inference only); return how many.""" + layers = nvfp4_linears(module) + for layer in layers: + layer.freeze() + return len(layers) + + +def freeze_nvfp4_linears_on_device(module: nn.Module) -> int: + """``freeze_nvfp4_linears`` once ``module`` is on an accelerator; 0 when it is not. + + ``_apply_h3_nvfp4_decoder`` cannot freeze a VAE that was built on the CPU, + because packing needs the CUDA FP4 kernels. With VAE offload that is every + load, so the layers keep their fp32 masters and a plain-attribute packed + cache for the whole decode. Freeze at the first CUDA load instead: the packed + weight is then a buffer, which is both smaller and movable by ``pinned_offload``. + """ + if not isinstance(module, nn.Module): + return 0 + layers = [layer for layer in nvfp4_linears(module) + if not layer.frozen and not layer.dense_bypass and not layer.calibrating] + parameters = next(module.parameters(), None) + if not layers or parameters is None or parameters.device.type != "cuda": + return 0 + for layer in layers: + layer.freeze() + return len(layers) + + +def keep_evenly_spaced_blocks(decoder: nn.Module, count: int) -> list[int]: + """Depth-cut the decoder to ``count`` evenly spaced blocks (first and last kept); return the kept indices.""" + total = len(decoder.transformer_blocks) + if not 1 <= count <= total: + raise ValueError(f"cannot keep {count} of {total} decoder blocks") + indices = sorted({round(i * (total - 1) / max(count - 1, 1)) for i in range(count)}) + decoder.transformer_blocks = nn.ModuleList(decoder.transformer_blocks[i] for i in indices) + return indices + + +# ``scripts/distill/minimax_h3_nvfp4_decoder/export_deploy.py`` output. +NVFP4_DECODER_DEPLOY_FORMAT = "fastvideo_h3_decoder_nvfp4_deploy_v1" + + +def load_nvfp4_decoder_checkpoint(path: str, device: torch.device | str = "cpu") -> dict[str, Any]: + """Read an exported NVFP4 decoder (tensors, plain metadata) and check its format. + + ``.safetensors`` files may hold bf16 masters (deploy format) or pre-packed NVFP4 weights; + packed tensors are read straight onto ``device``. ``.pt`` files are deploy format only. + """ + if str(path).endswith(".safetensors"): + from fastvideo.models.vaes.minimax_h3_nvfp4_checkpoint import read_nvfp4_decoder_safetensors + + return read_nvfp4_decoder_safetensors(path, device) + checkpoint = torch.load(path, map_location="cpu", mmap=True, weights_only=True) + if not isinstance(checkpoint, dict) or checkpoint.get("format") != NVFP4_DECODER_DEPLOY_FORMAT: + found = checkpoint.get("format") if isinstance(checkpoint, dict) else type(checkpoint).__name__ + raise ValueError(f"{path} is not a {NVFP4_DECODER_DEPLOY_FORMAT} checkpoint (format: {found!r})") + missing = {"decoder", "post_quant_conv", "metadata"} - checkpoint.keys() + if missing: + raise ValueError(f"{path} is missing {sorted(missing)}") + return checkpoint + + +def apply_nvfp4_decoder_checkpoint(vae: nn.Module, checkpoint: dict[str, Any], *, freeze: bool) -> dict[str, Any]: + """Turn a dense ``AutoencoderKLMiniMaxH3``'s decoder into an exported NVFP4 decoder, in place. + + The VAE is depth-cut to the checkpoint's ``student_layers`` blocks when it has more (a light + or full VAE can host a shallower student), converted with the checkpoint's NVFP4 settings, + and loaded strictly, so an architecture mismatch fails here instead of decoding garbage. + ``freeze`` packs the weights and drops the masters (needs the VAE on a CUDA device). A + pre-packed checkpoint is always loaded frozen, without materializing any master weight. + Sets ``vae.decode_autocast_dtype`` to bf16, the dtype these decoders are trained and + validated under. Returns the checkpoint metadata. + """ + metadata = checkpoint["metadata"] + decoder = vae.decoder + _depth_cut_for(decoder, metadata) + if checkpoint["format"] != NVFP4_DECODER_DEPLOY_FORMAT: + from fastvideo.models.vaes.minimax_h3_nvfp4_checkpoint import apply_packed_nvfp4_decoder + + apply_packed_nvfp4_decoder(vae, checkpoint) + vae.decode_autocast_dtype = torch.bfloat16 + return metadata + act_scale = metadata.get("act_scale", "dynamic") + names = convert_decoder_to_nvfp4(decoder, + rotation_group=metadata.get("rotation_group"), + skip_blocks=tuple(metadata.get("skip_blocks") or ()), + compute_dtype=torch.bfloat16, + act_scale=act_scale) + check_num_linears(metadata, len(names)) + try: + decoder.load_state_dict(checkpoint["decoder"], strict=True) + vae.post_quant_conv.load_state_dict(checkpoint["post_quant_conv"], strict=True) + except RuntimeError as error: + raise ValueError(f"the NVFP4 decoder does not fit this VAE's architecture: {error}") from error + vae.requires_grad_(False) + if freeze: + freeze_nvfp4_linears(decoder) + vae.decode_autocast_dtype = torch.bfloat16 + return metadata + + +def check_num_linears(metadata: dict[str, Any], found: int) -> None: + expected = metadata.get("num_linears") + if expected is not None and int(expected) != found: + raise ValueError(f"the NVFP4 decoder has {expected} NVFP4 linears, the converted VAE {found}") + + +def _depth_cut_for(decoder: nn.Module, metadata: dict[str, Any]) -> None: + available = len(decoder.transformer_blocks) + student_layers = int(metadata.get("student_layers", available)) + if student_layers > available: + raise ValueError(f"the NVFP4 decoder needs {student_layers} decoder blocks but the loaded VAE has " + f"{available}; load a VAE with at least that many (e.g. the full MiniMax-H3 vae/)") + if student_layers < available: + keep_evenly_spaced_blocks(decoder, student_layers) diff --git a/fastvideo/models/vaes/minimax_h3_nvfp4_fused.py b/fastvideo/models/vaes/minimax_h3_nvfp4_fused.py new file mode 100644 index 000000000..d88be6e0d --- /dev/null +++ b/fastvideo/models/vaes/minimax_h3_nvfp4_fused.py @@ -0,0 +1,683 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Bit-exact fused inference kernels for the NVFP4 MiniMax-H3 decoder blocks. + +The eager NVFP4 block (``NVFP4DecoderLinear`` inside ``MiniMaxH3VideoTransformerBlock``) +spends most of its time in small elementwise kernels around the FP4 GEMMs: bias adds, +RMSNorms, casts, RoPE, SiLU, residuals and three separate quantizations of the shared +q/k/v input. These Triton kernels fuse them while reproducing every rounding step of +the eager path, so the decoder output is bit-identical: + +* ``rms_norm`` follows ATen's ``vectorized_layer_norm_kernel`` reduction order + (128 threads, 4-wide vectors, sequential FMA per thread, warp shuffle-down tree, + then a 4-warp tree), so the fp32 statistics match bit for bit; +* NVFP4 quantization follows FlashInfer's ``cvt_warp_fp16_to_fp4`` (approximate + reciprocals, E4M3 block scales, ``layout_128x4`` swizzle, zeroed padding rows); +* every bf16 rounding of the op-by-op eager path is kept, and products that eager + rounds before an add use ``mul.rn`` so they are never contracted into an FMA. + +``fastvideo/tests/vaes/test_minimax_h3_nvfp4_decoder.py`` checks each kernel and whole +decodes against the eager path with ``torch.equal``. +""" +from __future__ import annotations + +import functools +from typing import Any, NamedTuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +import triton +import triton.language as tl + +import fastvideo.envs as envs +from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import NVFP4DecoderLinear, NVFP4FusedState + +NVFP4_BLOCK = 16 +SF_ROW_TILE = 128 +# ATen's vectorized layer/RMS norm: 4 warps x 32 lanes, 4 elements per vector load. +_NORM_THREADS = 128 +_NORM_VEC = 4 + + +@triton.jit +def _rcp_approx_ftz(x): + return tl.inline_asm_elementwise("rcp.approx.ftz.f32 $0, $1;", + "=r,r", [x], + dtype=tl.float32, + is_pure=True, + pack=1) + + +@triton.jit +def _rsqrt_approx(x): + return tl.inline_asm_elementwise("rsqrt.approx.f32 $0, $1;", "=r,r", [x], dtype=tl.float32, is_pure=True, pack=1) + + +@triton.jit +def _mul_rn(a, b): + # A rounded product that the compiler may not contract into an FMA with a following add. + 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 _e4m3_bits(x): + return tl.inline_asm_elementwise( + "{ .reg .b16 t; cvt.rn.satfinite.e4m3x2.f32 t, $1, $1; cvt.u32.u16 $0, t; }", + "=r,r", [x], + dtype=tl.int32, + is_pure=True, + pack=1) & 0xFF + + +@triton.jit +def _e4m3_to_f32(bits): + exponent = (bits >> 3) & 0xF + mantissa = bits & 0x7 + normal = (((exponent + 120) << 23) | (mantissa << 20)).to(tl.float32, bitcast=True) + return tl.where(exponent == 0, mantissa.to(tl.float32) * 0.001953125, normal) + + +@triton.jit +def _e2m1x2(lo, hi): + return tl.inline_asm_elementwise( + "{ .reg .b8 t; cvt.rn.satfinite.e2m1x2.f32 t, $2, $1; cvt.u32.u8 $0, t; }", + "=r,r,r", [lo, hi], + dtype=tl.int32, + is_pure=True, + pack=1) + + +@triton.jit +def _bf16(x): + """Round fp32 to the nearest-even bf16 value (kept in fp32). + + Inline PTX, because the compiler may fold a float ``truncf``/``extf`` round trip into a + following add and skip the rounding that eager applies. + """ + return tl.inline_asm_elementwise("{ .reg .b16 t; cvt.rn.bf16.f32 t, $1; cvt.f32.bf16 $0, t; }", + "=r,r", [x], + dtype=tl.float32, + is_pure=True, + pack=1) + + +@triton.jit +def _quantize_tile(x, gsf, R: tl.constexpr, N: tl.constexpr): + """NVFP4-quantize ``x`` [R, N] (fp32 holding bf16 values) like FlashInfer; return (packed, sf bits).""" + blocks = tl.reshape(x, (R, N // 16, 16)) + vec_max = tl.max(tl.abs(blocks), axis=2) + sf_value = gsf * (vec_max * _rcp_approx_ftz(tl.full(vec_max.shape, 6.0, tl.float32))) + sf_bits = _e4m3_bits(sf_value) + out_scale = _rcp_approx_ftz(_e4m3_to_f32(sf_bits) * _rcp_approx_ftz(gsf + tl.zeros(vec_max.shape, tl.float32))) + out_scale = tl.where(vec_max != 0, out_scale, 0.0) + scaled = tl.reshape(blocks * out_scale[:, :, None], (R, N // 2, 2)) + lo, hi = tl.split(scaled) + return _e2m1x2(lo, hi).to(tl.uint8), sf_bits.to(tl.uint8) + + +@triton.jit +def _sf_offsets(rows, sf_cols, NUM_SF_COLS: tl.constexpr): + """``layout_128x4`` swizzled offsets of block scales (rows [R, 1], sf_cols [1, C]).""" + k_tiles: tl.constexpr = NUM_SF_COLS // 4 + return ((rows // 128) * (k_tiles * 512) + (sf_cols // 4) * 512 + (rows % 32) * 16 + ((rows % 128) // 32) * 4 + + (sf_cols % 4)) + + +@triton.jit +def _store_quantized(q_ptr, sf_ptr, packed, sf_bits, rows, col0, row_ok, sf_row_ok, R: tl.constexpr, + N: tl.constexpr, K: tl.constexpr): + """Store a quantized [R, N] tile that starts at input column ``col0`` of a K-wide activation.""" + rows64 = rows.to(tl.int64) + byte_cols = col0 // 2 + tl.arange(0, N // 2) + tl.store(q_ptr + rows64[:, None] * (K // 2) + byte_cols[None, :], packed, mask=row_ok[:, None]) + sf_cols = col0 // 16 + tl.arange(0, N // 16) + tl.store(sf_ptr + _sf_offsets(rows64[:, None], sf_cols[None, :], K // 16), sf_bits, mask=sf_row_ok[:, None]) + + +@triton.jit +def _lane_tree_sum16(s, ROWS: tl.constexpr): + """Sum [ROWS, 16] lanes in warp shuffle-down order: lane t pairs with t + 8, then + 4, + 2, + 1.""" + s = tl.sum(tl.reshape(s, (ROWS, 2, 8)), axis=1) + s = tl.sum(tl.reshape(s, (ROWS, 2, 4)), axis=1) + s = tl.sum(tl.reshape(s, (ROWS, 2, 2)), axis=1) + return tl.sum(s, axis=1) + + +@triton.jit +def _lane_tree_sum32(s, ROWS: tl.constexpr): + """Sum [ROWS, 32] lanes in warp shuffle-down order (lane t pairs with t + 16 first).""" + return _lane_tree_sum16(tl.sum(tl.reshape(s, (ROWS, 2, 16)), axis=1), ROWS) + + +@triton.jit +def _residual_norm_quant_kernel(h_ptr, o_ptr, bias_ptr, scale_ptr, gamma_ptr, gsf_ptr, q_ptr, sf_ptr, M, M_PAD, eps, + D: tl.constexpr, R: tl.constexpr, HAS_RESIDUAL: tl.constexpr, + HAS_NORM: tl.constexpr): + """``h += bf16(o + bias) * scale`` (optional), then NVFP4(bf16(rms_norm(h) * gamma)) (optional).""" + rows = tl.program_id(0) * R + tl.arange(0, R) + row_ok = rows < M + rows64 = rows.to(tl.int64) + acc = tl.zeros((R, 128), tl.float32) + for k in tl.static_range(D // 512): + cols = k * 512 + tl.arange(0, 512) + offs = rows64[:, None] * D + cols[None, :] + h = tl.load(h_ptr + offs, mask=row_ok[:, None], other=0.0) + if HAS_RESIDUAL: + o = tl.load(o_ptr + offs, mask=row_ok[:, None], other=0.0).to(tl.float32) + biased = _bf16(o + tl.load(bias_ptr + cols).to(tl.float32)[None, :]) + h = h + _mul_rn(biased, tl.load(scale_ptr + cols)[None, :]) + tl.store(h_ptr + offs, h, mask=row_ok[:, None]) + if HAS_NORM: + # Column 512k + 4t + i belongs to ATen thread t; each thread sums its 4-vectors in order. + evens, odds = tl.split(tl.reshape(h, (R, 128, 2, 2))) + x0, x2 = tl.split(evens) + x1, x3 = tl.split(odds) + acc = tl.fma(x0, x0, acc) + acc = tl.fma(x1, x1, acc) + acc = tl.fma(x2, x2, acc) + acc = tl.fma(x3, x3, acc) + if HAS_NORM: + warp_sums = _lane_tree_sum32(tl.reshape(acc, (R * 4, 32)), R * 4) + sigma2 = tl.sum(tl.sum(tl.reshape(warp_sums, (R, 2, 2)), axis=1), axis=1) + rstd = _rsqrt_approx(sigma2 * (1.0 / D) + eps) + gsf = tl.load(gsf_ptr) + if HAS_RESIDUAL: + tl.debug_barrier() + sf_row_ok = rows < M_PAD + for k in tl.static_range(D // 512): + cols = k * 512 + tl.arange(0, 512) + h = tl.load(h_ptr + rows64[:, None] * D + cols[None, :], mask=row_ok[:, None], other=0.0) + gamma = tl.load(gamma_ptr + cols) + normed = _bf16(gamma[None, :] * (rstd[:, None] * h)) + packed, sf_bits = _quantize_tile(normed, gsf, R, 512) + _store_quantized(q_ptr, sf_ptr, packed, sf_bits, rows, k * 512, row_ok, sf_row_ok, R, 512, D) + + +@triton.jit +def _quantize_kernel(x_ptr, gsf_ptr, q_ptr, sf_ptr, M, M_PAD, K: tl.constexpr, R: tl.constexpr, BN: tl.constexpr): + rows = tl.program_id(0) * R + tl.arange(0, R) + cols = tl.program_id(1) * BN + tl.arange(0, BN) + row_ok = rows < M + x = tl.load(x_ptr + rows.to(tl.int64)[:, None] * K + cols[None, :], mask=row_ok[:, None], other=0.0) + packed, sf_bits = _quantize_tile(x.to(tl.float32), tl.load(gsf_ptr), R, BN) + _store_quantized(q_ptr, sf_ptr, packed, sf_bits, rows, tl.program_id(1) * BN, row_ok, rows < M_PAD, R, BN, K) + + +@triton.jit +def _swiglu_quantize_kernel(o_ptr, bias_ptr, silu_ptr, gsf_ptr, q_ptr, sf_ptr, M, M_PAD, F: tl.constexpr, + R: tl.constexpr, BN: tl.constexpr): + """NVFP4(bf16(bf16(value + b) * silu(bf16(gate + b)))) on a value-first packed projection. + + ``silu_ptr`` maps every bf16 bit pattern to eager ``F.silu`` of it, so SiLU costs one cached load. + """ + rows = tl.program_id(0) * R + tl.arange(0, R) + cols = tl.program_id(1) * BN + tl.arange(0, BN) + row_ok = rows < M + offs = rows.to(tl.int64)[:, None] * (2 * F) + cols[None, :] + value = _bf16( + tl.load(o_ptr + offs, mask=row_ok[:, None], other=0.0).to(tl.float32) + + tl.load(bias_ptr + cols).to(tl.float32)[None, :]) + gate = (tl.load(o_ptr + offs + F, mask=row_ok[:, None], other=0.0).to(tl.float32) + + tl.load(bias_ptr + F + cols).to(tl.float32)[None, :]).to(tl.bfloat16) + silu = tl.load(silu_ptr + gate.to(tl.uint16, bitcast=True).to(tl.int32)).to(tl.float32) + # Padding rows (only their zero block scales are stored) must not pick up the bias. + hidden = tl.where(row_ok[:, None], _bf16(value * silu), 0.0) + packed, sf_bits = _quantize_tile(hidden, tl.load(gsf_ptr), R, BN) + _store_quantized(q_ptr, sf_ptr, packed, sf_bits, rows, tl.program_id(1) * BN, row_ok, rows < M_PAD, R, BN, F) + + +@triton.jit +def _qkv_epilogue_kernel(q_ptr, k_ptr, v_ptr, bq_ptr, bk_ptr, bv_ptr, cos_ptr, sin_ptr, q_out, k_out, v_out, M, eps, + H: tl.constexpr, HD: tl.constexpr, ROT: tl.constexpr, R: tl.constexpr): + """Bias + per-head RMSNorm + partial RoPE for q and k, bias for v (eager rounding throughout).""" + D: tl.constexpr = H * HD + which = tl.program_id(1) + if which == 2: + rows = tl.program_id(0) * R + tl.arange(0, R) + cols = tl.arange(0, D) + offs = rows.to(tl.int64)[:, None] * D + cols[None, :] + mask = (rows < M)[:, None] + v = tl.load(v_ptr + offs, mask=mask, other=0.0).to(tl.float32) + tl.load(bv_ptr + cols).to(tl.float32)[None, :] + tl.store(v_out + offs, v.to(tl.bfloat16), mask=mask) + else: + src = q_ptr + bias_ptr = bq_ptr + dst = q_out + if which == 1: + src = k_ptr + bias_ptr = bk_ptr + dst = k_out + RH: tl.constexpr = R * H + head_rows = tl.arange(0, RH) + tokens = tl.program_id(0) * R + head_rows // H + token_ok = (tokens < M)[:, None] + base = tokens.to(tl.int64)[:, None] * D + ((head_rows % H) * HD)[:, None] + bias_base = ((head_rows % H) * HD)[:, None] + d = tl.arange(0, HD)[None, :] + x = _bf16( + tl.load(src + base + d, mask=token_ok, other=0.0).to(tl.float32) + + tl.load(bias_ptr + bias_base + d).to(tl.float32)) + # ATen's per-row RMSNorm over HD = 64: thread t < 16 sums elements 4t..4t+3 in order. + evens, odds = tl.split(tl.reshape(x, (RH, HD // 4, 2, 2))) + x0, x2 = tl.split(evens) + x1, x3 = tl.split(odds) + acc = tl.fma(x0, x0, tl.zeros((RH, HD // 4), tl.float32)) + acc = tl.fma(x1, x1, acc) + acc = tl.fma(x2, x2, acc) + acc = tl.fma(x3, x3, acc) + sigma2 = _lane_tree_sum16(acc, RH) + rstd = _rsqrt_approx(sigma2 * (1.0 / HD) + eps)[:, None] + normed = _bf16(rstd * x) + half: tl.constexpr = ROT // 2 + rotary = d < ROT + partner = tl.where(d < half, d + half, d - half) + partner_x = _bf16( + tl.load(src + base + partner, mask=token_ok & rotary, other=0.0).to(tl.float32) + + tl.load(bias_ptr + bias_base + partner, mask=rotary, other=0.0).to(tl.float32)) + partner_normed = _bf16(rstd * partner_x) + rotated = tl.where(d < half, -partner_normed, partner_normed) + trig = tokens.to(tl.int64)[:, None] * ROT + d + cos = tl.load(cos_ptr + trig, mask=token_ok & rotary, other=0.0).to(tl.float32) + sin = tl.load(sin_ptr + trig, mask=token_ok & rotary, other=0.0).to(tl.float32) + roped = _bf16(_bf16(normed * cos) + _bf16(rotated * sin)) + tl.store(dst + base + d, tl.where(rotary, roped, normed).to(tl.bfloat16), mask=token_ok) + + +def _padded_rows(m: int) -> int: + return (m + SF_ROW_TILE - 1) // SF_ROW_TILE * SF_ROW_TILE + + +def new_quantized(m: int, k: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: + """Buffers shaped like FlashInfer's ``nvfp4_quantize`` output for an [m, k] activation.""" + packed = torch.empty((m, k // 2), dtype=torch.uint8, device=device) + scales = torch.empty((_padded_rows(m), k // NVFP4_BLOCK), dtype=torch.uint8, device=device) + return packed, scales + + +def residual_norm_quantize(h: torch.Tensor, residual: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None, + norm: tuple[torch.Tensor, float, torch.Tensor] | None, + rows_per_program: int = 4) -> tuple[torch.Tensor, torch.Tensor] | None: + """In place ``h += bf16(o + bias) * scale`` and/or NVFP4 of ``bf16(rms_norm(h, gamma, eps))``. + + ``h`` is a contiguous fp32 [M, D] tensor; ``residual`` is ``(o, bias_bf16, scale_fp32)``; + ``norm`` is ``(gamma_fp32, eps, global_sf)``. Returns the quantized activation when ``norm`` is set. + """ + m, d = h.shape + if d % 512: + raise ValueError(f"hidden size must be a multiple of 512, got {d}") + quantized = new_quantized(m, d, h.device) if norm is not None else (h, h) + o, bias, scale = residual if residual is not None else (h, h, h) + gamma, eps, gsf = norm if norm is not None else (h, 0.0, h) + rows = _padded_rows(m) if norm is not None else m + grid = (triton.cdiv(rows, rows_per_program), ) + _residual_norm_quant_kernel[grid](h, + o, + bias, + scale, + gamma, + gsf, + quantized[0], + quantized[1], + m, + rows, + eps, + D=d, + R=rows_per_program, + HAS_RESIDUAL=residual is not None, + HAS_NORM=norm is not None, + num_warps=4) + return quantized if norm is not None else None + + +def quantize(x: torch.Tensor, gsf: torch.Tensor, rows_per_program: int = 8, + block_n: int = 512) -> tuple[torch.Tensor, torch.Tensor]: + """FlashInfer-exact NVFP4 quantization (``layout_128x4``) of a contiguous bf16 [M, K] tensor.""" + m, k = x.shape + quantized = new_quantized(m, k, x.device) + grid = (triton.cdiv(_padded_rows(m), rows_per_program), triton.cdiv(k, block_n)) + _quantize_kernel[grid](x, gsf, quantized[0], quantized[1], m, _padded_rows(m), K=k, R=rows_per_program, BN=block_n) + return quantized + + +@functools.cache +def _silu_table(device_index: int) -> torch.Tensor: + """Eager ``F.silu`` of every bf16 value, indexed by its 16-bit pattern.""" + patterns = torch.arange(1 << 16, dtype=torch.int32, device=torch.device("cuda", device_index)) + return F.silu(patterns.to(torch.int16).view(torch.bfloat16)) # int32 -> int16 wraps to the bit pattern + + +def swiglu_quantize(o: torch.Tensor, bias: torch.Tensor, gsf: torch.Tensor, rows_per_program: int = 2, + block_n: int = 512) -> tuple[torch.Tensor, torch.Tensor]: + """NVFP4 of the H3 SwiGLU of a value-first packed projection ``o`` [M, 2F] (bias not yet added).""" + m, two_f = o.shape + f = two_f // 2 + quantized = new_quantized(m, f, o.device) + grid = (triton.cdiv(_padded_rows(m), rows_per_program), triton.cdiv(f, block_n)) + _swiglu_quantize_kernel[grid](o, + bias, + _silu_table(o.device.index), + gsf, + quantized[0], + quantized[1], + m, + _padded_rows(m), + F=f, + R=rows_per_program, + BN=block_n, + num_warps=2) + return quantized + + +def qkv_epilogue(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, biases: tuple[torch.Tensor, torch.Tensor, + torch.Tensor], cos: torch.Tensor, + sin: torch.Tensor, heads: int, eps: float, + rows_per_program: int = 2) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Projected q/k/v [M, H*HD] bf16 (bias not yet added) -> roped q, k and biased v, each [M, H, HD]. + + ``cos``/``sin`` are bf16 [M, ROT] (the eager path's ``cos.to(query.dtype)``). + """ + m, d = q.shape + head_dim = d // heads + outs = [torch.empty((m, heads, head_dim), dtype=torch.bfloat16, device=q.device) for _ in range(3)] + grid = (triton.cdiv(m, rows_per_program), 3) + _qkv_epilogue_kernel[grid](q, + k, + v, + *biases, + cos, + sin, + *outs, + m, + eps, + H=heads, + HD=head_dim, + ROT=cos.shape[-1], + R=rows_per_program, + num_warps=4) + return outs[0], outs[1], outs[2] + + +class _BlockPlan(NamedTuple): + qkv: tuple[NVFP4FusedState, NVFP4FusedState, NVFP4FusedState] + out: NVFP4FusedState + ff_in: NVFP4FusedState + ff_out: NVFP4FusedState + norm1: tuple[torch.Tensor, float] + norm2: tuple[torch.Tensor, float] + qk_eps: float + scale1: torch.Tensor + scale2: torch.Tensor + heads: int + attention: Any + + +def _rms_norm_args(norm: nn.Module, affine: bool) -> tuple[torch.Tensor | None, float] | None: + if not isinstance(norm, nn.RMSNorm) or (norm.weight is not None) != affine: + return None + if affine and norm.weight.dtype != torch.float32: + return None + return norm.weight, norm.eps if norm.eps is not None else torch.finfo(torch.float32).eps + + +def _clear_block_plan(module: nn.Module) -> None: + """Drop a block's cached fused plan, whose tensors point at the previous weight storage.""" + module.__dict__.pop("_nvfp4_fused_plan", None) + + +def _block_plan(block: nn.Module) -> _BlockPlan | None: + """The fused plan of one ``MiniMaxH3VideoTransformerBlock``, or None if it must run eagerly.""" + attn, ff = block.attn, block.ff + linears = (attn.to_q, attn.to_k, attn.to_v, attn.to_out[0], ff.net[0].proj, ff.net[2]) + if not all(isinstance(layer, NVFP4DecoderLinear) for layer in linears): + return None + # The plan holds views of the parameter storage, so the offload walk has to + # drop it when the module moves to the host (parameters/buffers alone miss it). + if "clear_offload_caches" not in block.__dict__: + block.__dict__["clear_offload_caches"] = lambda module=block: _clear_block_plan(module) + states = [layer.fused_state() for layer in linears] + if any(state is None for state in states): + return None + key = tuple(id(state) for state in states) + cached = getattr(block, "_nvfp4_fused_plan", None) + if cached is not None and cached[0] == key: + return cached[1] + norm1, norm2 = _rms_norm_args(block.norm1, True), _rms_norm_args(block.norm2, True) + norm_q, norm_k = _rms_norm_args(attn.norm_q, False), _rms_norm_args(attn.norm_k, False) + dim = attn.heads * attn.dim_head + # The quantization kernels tile the inner dimension by 512 columns without a column mask. + ff_inner = ff.net[0].proj.out_features + plan = None + if (None not in (norm1, norm2, norm_q, norm_k) and norm_q[1] == norm_k[1] and attn.attn_impl is not None + and attn.dim_head == 64 and attn.heads & (attn.heads - 1) == 0 and dim % 512 == 0 + and ff_inner % 512 == 0 + and block.scale1.dtype == block.scale2.dtype == torch.float32 + and isinstance(ff.net[0].activation, nn.SiLU) + and all(state.global_sf_x is states[0].global_sf_x or torch.equal(state.global_sf_x, states[0].global_sf_x) + for state in states[1:3])): + plan = _BlockPlan(tuple(states[:3]), states[3], states[4], states[5], norm1, norm2, norm_q[1], block.scale1, + block.scale2, attn.heads, attn.attn_impl) + block._nvfp4_fused_plan = (key, plan) + return plan + + +_GEMM_WORKSPACE_BYTES = 40 << 20 +_GEMM_TUNING_REPEATS = 3 + + +@functools.cache +def _sm100_cutlass_fp4_module(device_index: int) -> Any | None: + """FlashInfer's CUTLASS FP4 GEMM extension on sm_100, called without ``mm_fp4``'s per-call checks. + + ``mm_fp4`` costs ~100 us of host time per call (a single 1797-token tile GEMM runs in ~12 us). + Other architectures keep ``mm_fp4``. + """ + if torch.cuda.get_device_capability(device_index) != (10, 0): + return None + try: + from flashinfer.gemm.gemm_base import gen_gemm_sm100_module_cutlass_fp4 + + return gen_gemm_sm100_module_cutlass_fp4().build_and_load() + except (ImportError, AttributeError, RuntimeError): + return None + + +@functools.cache +def _gemm_workspace(device_index: int) -> torch.Tensor: + return torch.empty(_GEMM_WORKSPACE_BYTES, dtype=torch.uint8, device=torch.device("cuda", device_index)) + + +_GEMM_TACTICS: dict[tuple[int, int, int, int], int] = {} + + +def _gemm_tactic(module: Any, args: tuple, out: torch.Tensor) -> int: + """Fastest CUTLASS tactic for this GEMM shape, measured once. + + Every tactic computes the same K-ordered fp32 accumulation, so the choice changes speed only + (``test_cutlass_fp4_tactics_are_bit_identical`` checks this). + """ + key = (out.device.index, out.shape[0], out.shape[1], args[0].shape[1]) + if key not in _GEMM_TACTICS: + timings = [] + for tactic in range(-1, module.fp4_gemm_tactic_num()): + try: + module.fp4_gemm(*args, out, _gemm_workspace(out.device.index), tactic) + except RuntimeError: + continue + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(_GEMM_TUNING_REPEATS): + module.fp4_gemm(*args, out, _gemm_workspace(out.device.index), tactic) + end.record() + end.synchronize() + timings.append((start.elapsed_time(end), tactic)) + _GEMM_TACTICS[key] = min(timings)[1] if timings else -1 + return _GEMM_TACTICS[key] + + +def _gemm(quantized: tuple[torch.Tensor, torch.Tensor], state: NVFP4FusedState) -> torch.Tensor: + """``nvfp4_linear_inference``'s GEMM (bias excluded) on an already quantized activation.""" + x_fp4, x_inv_scale = quantized + out = torch.empty((x_fp4.shape[0], state.out_features), device=x_fp4.device, dtype=torch.bfloat16) + module = _sm100_cutlass_fp4_module(out.device.index) + if module is not None: + # mm_fp4(backend="cutlass") hands the module the same operands (it transposes b and b's scales back). + args = (x_fp4, state.packed, x_inv_scale, state.inv_scale, state.alpha) + key = (out.device.index, out.shape[0], out.shape[1], x_fp4.shape[1]) + # Tuning synchronizes, so a CUDA graph capture uses an already tuned (or the default) tactic. + tactic = (_GEMM_TACTICS.get(key, -1) + if torch.cuda.is_current_stream_capturing() else _gemm_tactic(module, args, out)) + module.fp4_gemm(*args, out, _gemm_workspace(out.device.index), tactic) + return out + import flashinfer + + from fastvideo.models.vaes import minimax_h3_nvfp4_decoder + + backend = minimax_h3_nvfp4_decoder.mm_fp4_backend(out.device) + flashinfer.mm_fp4(x_fp4, + state.packed.T, + x_inv_scale, + state.inv_scale.T, + state.alpha, + torch.bfloat16, + out, + block_size=NVFP4_BLOCK, + use_8x4_sf_layout=False, + backend=backend) + return out + + +def _run_blocks(plans: list[_BlockPlan], hidden: torch.Tensor, rotary_emb: tuple[torch.Tensor, + torch.Tensor]) -> torch.Tensor: + batch, seq, dim = hidden.shape + rows = batch * seq + cos, sin = (t.to(torch.bfloat16).expand(batch, seq, 1, -1).reshape(rows, -1).contiguous() for t in rotary_emb) + h = hidden.reshape(rows, dim).clone() + first = plans[0] + quantized = residual_norm_quantize(h, None, (*first.norm1, first.qkv[0].global_sf_x)) + for index, plan in enumerate(plans): + projections = [_gemm(quantized, state) for state in plan.qkv] + query, key, value = qkv_epilogue(*projections, tuple(state.bias for state in plan.qkv), cos, sin, plan.heads, + plan.qk_eps) + del projections + shape = (batch, seq, plan.heads, dim // plan.heads) + attended = plan.attention.forward(query.view(shape), key.view(shape), value.view(shape), None) + del query, key, value + attended = attended.reshape(rows, dim).contiguous() + out = _gemm(quantize(attended, plan.out.global_sf_x), plan.out) + del attended + quantized = residual_norm_quantize(h, (out, plan.out.bias, plan.scale1), + (*plan.norm2, plan.ff_in.global_sf_x)) + out = _gemm(quantized, plan.ff_in) + out = _gemm(swiglu_quantize(out, plan.ff_in.bias, plan.ff_out.global_sf_x), plan.ff_out) + following = plans[index + 1] if index + 1 < len(plans) else None + quantized = residual_norm_quantize(h, (out, plan.ff_out.bias, plan.scale2), + None if following is None else + (*following.norm1, following.qkv[0].global_sf_x)) + del out + return h.view(batch, seq, dim) + + +# Calls up to this many tokens are launch-bound (a 1797-token tile runs ~400 kernels in ~5 ms of GPU +# time), so their block stack is captured once per shape and replayed as a CUDA graph. Larger calls +# are GPU-bound and run directly, without holding a graph memory pool. +GRAPH_MAX_ROWS = 4 * 1797 + + +# Replayed geometries kept at once. Each entry holds static tensors and a graph's memory pool, so a +# long-lived generator that walks through tile batching choices must not keep every shape it has seen. +GRAPH_CACHE_LIMIT = 2 + + +class _GraphedStack(NamedTuple): + plan_ids: tuple[int, ...] + graph: torch.cuda.CUDAGraph + hidden: torch.Tensor + cos: torch.Tensor + sin: torch.Tensor + out: torch.Tensor + + +def _weight_pointers(blocks: nn.ModuleList) -> tuple[int, ...]: + """A signature of where the blocks' weights live. + + A captured graph bakes the weight addresses into its kernels, so a replay is + only valid while they stay put. ``pinned_offload`` hands out fresh device + storage on every load, which changes them; including them in the cache key + re-captures instead of replaying against freed addresses. + """ + return tuple(t.data_ptr() for t in (*blocks.parameters(), *blocks.buffers())) + + +def _trim_graphs(graphs: dict) -> None: + """Keep the most recent ``GRAPH_CACHE_LIMIT`` geometries (the ``warm`` markers are tiny).""" + geometries = [key for key in graphs if not (isinstance(key, tuple) and key[:1] == ("warm", ))] + for key in geometries[:-GRAPH_CACHE_LIMIT]: + graphs.pop(key, None) + graphs.pop(("warm", key), None) + + +def _clear_graph_entries(graphs: dict) -> None: + """Release every captured graph (and its memory pool) in ``graphs``.""" + for entry in graphs.values(): + if isinstance(entry, _GraphedStack): + entry.graph.reset() + graphs.clear() + + +def _graphed_blocks(blocks: nn.ModuleList, plans: list[_BlockPlan], hidden: torch.Tensor, + rotary_emb: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor: + """``_run_blocks`` through a CUDA graph captured per input geometry (same kernels, same bits).""" + key = (hidden.device, tuple(hidden.shape), tuple(rotary_emb[0].shape), rotary_emb[0].dtype, + _weight_pointers(blocks)) + graphs: dict = blocks.__dict__.setdefault("_nvfp4_cuda_graphs", {}) + if "clear_offload_caches" not in blocks.__dict__: + # The graph cache lives on this ModuleList, not on a parameter or buffer of + # it, so the generic offload walk needs a hook to find it. An unload hands + # out fresh device storage, and a replay against the old addresses is wrong. + blocks.__dict__["clear_offload_caches"] = lambda: _clear_graph_entries(graphs) + plan_ids = tuple(id(plan) for plan in plans) + entry = graphs.get(key) + if entry is None or entry.plan_ids != plan_ids: + if graphs.get(("warm", key)) != plan_ids: + # First call: run eagerly so Triton compiles, GEMM tactics are tuned and lazy tables exist. + graphs[("warm", key)] = plan_ids + return _run_blocks(plans, hidden, rotary_emb) + static_hidden = hidden.clone() + static_cos, static_sin = (t.clone() for t in rotary_emb) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + static_out = _run_blocks(plans, static_hidden, (static_cos, static_sin)) + entry = _GraphedStack(plan_ids, graph, static_hidden, static_cos, static_sin, static_out) + graphs[key] = entry + _trim_graphs(graphs) + entry.hidden.copy_(hidden) + entry.cos.copy_(rotary_emb[0]) + entry.sin.copy_(rotary_emb[1]) + entry.graph.replay() + return entry.out.clone() + + +@torch.compiler.disable +def fused_nvfp4_blocks_forward(blocks: nn.ModuleList, hidden: torch.Tensor, + rotary_emb: tuple[torch.Tensor, torch.Tensor] | None) -> torch.Tensor | None: + """Run every decoder block through the fused kernels; None (caller runs the blocks) when not applicable. + + Applies at inference (no grad) to fp32 residual streams on CUDA when every block's linears are NVFP4 + with a precomputed activation scale. Bit-identical to the eager block loop. ``torch.compile`` treats it + as a boundary: a compiled decoder compiles the code around the block stack and runs the stack here. + Calls of at most ``GRAPH_MAX_ROWS`` tokens replay a per-shape CUDA graph + (``FASTVIDEO_H3_VAE_NVFP4_CUDA_GRAPH``). + """ + if (torch.compiler.is_compiling() or torch.is_grad_enabled() or not envs.FASTVIDEO_H3_VAE_NVFP4_FUSED.get() + or rotary_emb is None or len(blocks) == 0 or hidden.dtype != torch.float32 or not hidden.is_cuda): + return None + plans = [_block_plan(block) for block in blocks] + if any(plan is None for plan in plans): + return None + if (hidden.shape[0] * hidden.shape[1] <= GRAPH_MAX_ROWS and envs.FASTVIDEO_H3_VAE_NVFP4_CUDA_GRAPH.get() + and not torch.cuda.is_current_stream_capturing()): + return _graphed_blocks(blocks, plans, hidden, rotary_emb) + return _run_blocks(plans, hidden, rotary_emb) diff --git a/fastvideo/models/vaes/minimax_h3_video.py b/fastvideo/models/vaes/minimax_h3_video.py index fd9ea6cc6..fa1ecbdbf 100644 --- a/fastvideo/models/vaes/minimax_h3_video.py +++ b/fastvideo/models/vaes/minimax_h3_video.py @@ -7,7 +7,7 @@ This module intentionally uses only PyTorch and FastVideo configuration types. """ import math -from collections.abc import Iterator +from collections.abc import Callable, Iterator from dataclasses import dataclass import torch @@ -477,6 +477,8 @@ class MiniMaxH3VideoViTDecoder3d(nn.Module): self.norm_out = nn.LayerNorm(dim, elementwise_affine=True, eps=norm_eps) self.proj_out = nn.Linear(dim, out_channels * patch_size_t * patch_size * patch_size) self.gradient_checkpointing = False + self.fused_blocks_forward: Callable[[nn.ModuleList, torch.Tensor, tuple[torch.Tensor, torch.Tensor]], + torch.Tensor | None] | None = None def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: """Decode one latent spatial input through the H3 video transformer.""" @@ -501,7 +503,12 @@ class MiniMaxH3VideoViTDecoder3d(nn.Module): suffix_ids = position_ids.new_zeros((batch_size, self.num_register_tokens + 1, 3)) rotary_emb = self.rope(torch.cat([position_ids, suffix_ids], dim=1)) - for block in self.transformer_blocks: + # Quantized decoders (e.g. NVFP4) may install a fused inference path; it returns None when inapplicable. + fused = self.fused_blocks_forward(self.transformer_blocks, hidden_states, + rotary_emb) if self.fused_blocks_forward is not None else None + if fused is not None: + hidden_states = fused + for block in self.transformer_blocks if fused is None else (): if torch.is_grad_enabled() and self.gradient_checkpointing: hidden_states = checkpoint(block, hidden_states, rotary_emb, use_reentrant=False) else: @@ -550,6 +557,8 @@ class AutoencoderKLMiniMaxH3(nn.Module): # ``prepare_for_compile`` flips this per instance when the pipeline opts # into ``enable_torch_compile_vae``; default instances stay fully eager. _tile_helpers_compiled = False + # Autocast dtype of the pipeline's decode stage; NVFP4 decoders (``apply_nvfp4_decoder_checkpoint``) use bf16. + decode_autocast_dtype: torch.dtype = torch.float16 def __init__(self, config: MiniMaxH3VideoVAEConfig) -> None: super().__init__() diff --git a/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py index 50e81879a..5d4260bed 100644 --- a/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py +++ b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py @@ -637,9 +637,11 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): logger.info("Restored MiniMax-H3 denoise modules to %s", device) def _run_condition_then_denoise(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: - for name in ("input_preparation_stage", "conditioning_stage"): - batch = self._stage_name_mapping[name](batch, fastvideo_args) - self._release_text_encoder() + try: + for name in ("input_preparation_stage", "conditioning_stage"): + batch = self._stage_name_mapping[name](batch, fastvideo_args) + finally: + self._release_text_encoder() self._load_denoise_modules(fastvideo_args) if not self._unified_memory_host(): self._restore_denoise_modules(fastvideo_args) @@ -694,6 +696,10 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): tokenizer=self.get_module("tokenizer"), processor=self.get_module("processor"), ref2va=ref2va, + # Sequential placement reuses one pinned encoder copy. A stage-level + # .to("cpu") would bypass it and allocate new host weights every request. + manage_cpu_offload=(h3_encoder_split_enabled(fastvideo_args) + or not self._defer_denoise_modules(fastvideo_args)), ), ) diff --git a/fastvideo/pipelines/basic/minimax_h3/reference.py b/fastvideo/pipelines/basic/minimax_h3/reference.py index 8f8188bf8..712cdd6bf 100644 --- a/fastvideo/pipelines/basic/minimax_h3/reference.py +++ b/fastvideo/pipelines/basic/minimax_h3/reference.py @@ -3,6 +3,7 @@ from __future__ import annotations +from collections.abc import Mapping from dataclasses import dataclass, field import math from numbers import Integral, Real @@ -74,14 +75,28 @@ class MiniMaxH3PreparedReference: num_audio_latents: int = 0 +def coerce_reference(reference: Any) -> MiniMaxH3Reference: + """One Ref2VA entry: a typed reference, a config mapping, or a bare media path. + + A config file deserializes its request as plain data, so ``inputs.references`` + reach here as mappings (or, from a hand-written list, as paths). Both mean the + same thing as a typed reference. + """ + if isinstance(reference, MiniMaxH3Reference): + return reference + if isinstance(reference, Mapping): + return MiniMaxH3Reference(**reference) + if isinstance(reference, str | os.PathLike): + return MiniMaxH3Reference(source=str(reference)) + raise TypeError(f"Ref2VA entries must be a MiniMaxH3Reference, a mapping or a path, got {reference!r}") + + def validate_references(references: list[Any]) -> list[MiniMaxH3Reference]: """Validate the ordered Ref2VA media list and its per-modality limits.""" if not references: raise ValueError("MiniMax-H3 Ref2VA requires at least one reference.") - if not all(isinstance(reference, MiniMaxH3Reference) for reference in references): - raise TypeError("Every Ref2VA entry must be a MiniMaxH3Reference.") - typed = list(references) + typed = [coerce_reference(reference) for reference in references] counts = { media_type: sum(reference.media_type == media_type for reference in typed) for media_type in ("image", "video", "audio") @@ -370,6 +385,7 @@ __all__ = [ "MINIMAX_H3_QWEN_TEMPORAL_PATCH", "MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS", "MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE", + "coerce_reference", "MiniMaxH3PreparedReference", "MiniMaxH3Reference", "decode_reference_audio", 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 02145e9b8..667e9f88f 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_conditioning.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_conditioning.py @@ -108,12 +108,14 @@ class MiniMaxH3ConditioningStage(PipelineStage): processor: Any, *, ref2va: bool = False, + manage_cpu_offload: bool = True, ) -> None: super().__init__() self.conditioner = conditioner self.tokenizer = tokenizer self.processor = processor self.ref2va = ref2va + self.manage_cpu_offload = manage_cpu_offload self._ref2va_memo = ContentMemo() def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: @@ -276,8 +278,8 @@ class MiniMaxH3ConditioningStage(PipelineStage): """Encode one H3 prompt presentation and attach its packed text features.""" device = get_local_torch_device() first_param = next(self.conditioner.parameters(), None) - moved_for_forward = (fastvideo_args.text_encoder_cpu_offload and first_param is not None - and not isinstance(first_param, DTensor) + moved_for_forward = (getattr(self, "manage_cpu_offload", True) and fastvideo_args.text_encoder_cpu_offload + and first_param is not None and not isinstance(first_param, DTensor) and getattr(self.conditioner, "_h3_encoder_layerwise_device", None) is None) if moved_for_forward: self.conditioner.to(device) diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py index 8c55dd323..fdfe8dd18 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py @@ -27,6 +27,7 @@ from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.pipelines.stages.base import PipelineStage from fastvideo.pipelines.stages.validators import StageValidators as V from fastvideo.pipelines.stages.validators import VerificationResult +from fastvideo.platforms import current_platform from fastvideo.utils import allocate_cpu_tensor_with_pin_fallback logger = init_logger(__name__) @@ -67,6 +68,31 @@ class MiniMaxH3VideoDecodingStage(PipelineStage): def __init__(self, vae: AutoencoderKLMiniMaxH3 | None) -> None: super().__init__() self.vae = vae + self._output_buffer: tuple[tuple[tuple[int, ...], bool], torch.Tensor] | None = None + self._unified_memory: bool | None = None + + def _pixel_buffer(self, shape: tuple[int, ...], pin_memory: bool, device: torch.device) -> torch.Tensor: + """The CPU FP32 pixel buffer, reused across requests of the same geometry on unified memory. + + The stage returns a uint8 copy, so the buffer never escapes. Reuse matters on + unified-memory GPUs (DGX Spark): device-to-host copies into never-touched pageable + pages fault through the driver at ~150 MB/s (+3-7 s per 480p-768p clip whenever the + allocator hands back fresh pages), while populated pages copy at full speed. + ``zero_`` populates a new buffer from the CPU once. Discrete GPUs gain nothing from + holding up to 4.27 GB of host memory between requests, so they allocate per request. + """ + if self._unified_memory is None: + self._unified_memory = current_platform.has_unified_memory(device.index or 0) + if not self._unified_memory: + return allocate_cpu_tensor_with_pin_fallback(shape, dtype=torch.float32, pin_memory=pin_memory) + key = (tuple(shape), bool(pin_memory)) + if self._output_buffer is None or self._output_buffer[0] != key: + self._output_buffer = None + buffer = allocate_cpu_tensor_with_pin_fallback(shape, dtype=torch.float32, pin_memory=pin_memory) + if not buffer.is_pinned(): + buffer.zero_() + self._output_buffer = (key, buffer) + return self._output_buffer[1] def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: result = VerificationResult() @@ -130,6 +156,11 @@ class MiniMaxH3VideoDecodingStage(PipelineStage): raise RuntimeError("MiniMax-H3 full VAE decode requires a loaded video VAE.") pinned_offload.load(self.vae, device, pin=fastvideo_args.pin_cpu_memory) try: + # An NVFP4 decoder swapped in on the CPU could not pack its linears there; + # freeze them now that the VAE is on the device, on the first decode. + from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import freeze_nvfp4_linears_on_device + if freeze_nvfp4_linears_on_device(self.vae): + pinned_offload.refresh_host_copies(self.vae, pin=fastvideo_args.pin_cpu_memory) latents = self.vae.denormalize_latents(latents.to(device=device, dtype=torch.float32)) if fastvideo_args.output_type == "latent": # No collectives on this path, so uniform participation is @@ -139,16 +170,15 @@ class MiniMaxH3VideoDecodingStage(PipelineStage): output = None if is_output_rank: - output = allocate_cpu_tensor_with_pin_fallback( - self.vae.decoded_pixel_shape(latents.shape), - dtype=torch.float32, - pin_memory=fastvideo_args.pin_cpu_memory, - ) + output = self._pixel_buffer(tuple(self.vae.decoded_pixel_shape(latents.shape)), + fastvideo_args.pin_cpu_memory, device) # Attribute the streamed decoder computation while retaining # per-chunk device-to-host transfer and pinned-buffer reuse. + # fp16 for the release decoders; NVFP4 decoders were trained under bf16. + autocast_dtype = getattr(self.vae, "decode_autocast_dtype", torch.float16) with ( nvtx_range("minimax_h3.vae"), - torch.autocast(device_type=device.type, dtype=torch.float16, enabled=device.type == "cuda"), + torch.autocast(device_type=device.type, dtype=autocast_dtype, enabled=device.type == "cuda"), ): if parallel: strategy = fastvideo_args.vae_parallel_decode_strategy or DEFAULT_DECODE_GATHER_STRATEGY diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py index e04ac8d33..2cb35f7d5 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py @@ -324,11 +324,13 @@ class MiniMaxH3DenoisingStage(PipelineStage): # CUDAGraphs that has been overwritten" (surfaces at sp=1; # sp>1 is masked by collective-induced graph breaks). torch.compiler.cudagraph_mark_step_begin() - with self._pdd_head_fusion(pdd_plan, index), trace_step(index), set_forward_context( - current_timestep=index, - attn_metadata=attn_metadata, - forward_batch=batch, - ): + with self._pdd_head_fusion( + pdd_plan, + index), trace_step(index), nvtx_range(f"minimax_h3.step{index}"), set_forward_context( + current_timestep=index, + attn_metadata=attn_metadata, + forward_batch=batch, + ): video_velocity, audio_velocity = self.transformer( hidden_states=batch.latents[None], audio_hidden_states=batch.audio_latents[None], diff --git a/fastvideo/tests/attention/test_minimax_h3_fine_override.py b/fastvideo/tests/attention/test_minimax_h3_fine_override.py new file mode 100644 index 000000000..23cfacbf3 --- /dev/null +++ b/fastvideo/tests/attention/test_minimax_h3_fine_override.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Explicit fine attention must bypass the native consumer FP4 shortcut.""" +from types import SimpleNamespace + +import pytest +import torch + +from fastvideo import envs +from fastvideo.attention.backends import fp4_vsa_qat, video_sparse_attn_h3 +from fastvideo.models.dits import minimax_h3, minimax_h3_vsa_fp4 + + +def test_fine_override_bypasses_native_fp4(monkeypatch): + class PairIdentity(torch.nn.Module): + def forward(self, value): + return value, None + + class FakeAttention(torch.nn.Module): + def __init__(self): + super().__init__() + self.attn_impl = SimpleNamespace(fine_attention_override=object()) + self.calls = 0 + + def forward(self, q, k, v, **kwargs): + self.calls += 1 + return v, None + + attn = minimax_h3.MiniMaxH3Attention.__new__(minimax_h3.MiniMaxH3Attention) + torch.nn.Module.__init__(attn) + attn._vsa_fp4 = True + attn._vsa_tile_first = False + attn.distributed_attention = FakeAttention() + attn.to_q = attn.to_k = attn.to_v = attn.to_out = PairIdentity() + attn.num_attention_heads, attn.attention_head_dim = 1, 8 + attn.norm_q = attn.norm_k = torch.nn.Identity() + attn.fuse_qknorm_rope = False + attn.to_gate_compress = None + monkeypatch.setattr(attn, "_rotate", lambda value, rotary: value) + monkeypatch.setattr(minimax_h3_vsa_fp4, "_shared_input_projections", lambda layers, value: (value, value, value)) + + def unexpected_native(*args): + raise AssertionError("the native route must not run when a fine override exists") + + # The native branch reads metadata before projecting. Neither may be reached. + monkeypatch.setattr(minimax_h3, "get_forward_context", unexpected_native) + monkeypatch.setattr(minimax_h3, "vsa_fp4_attention", unexpected_native) + value = torch.randn(1, 3, 8) + with torch.no_grad(): + result = attn(value, (torch.ones(3, 8), torch.zeros(3, 8)), 3) + torch.testing.assert_close(result, value) + assert attn.distributed_attention.calls == 1 + + +def test_bf16_override_keeps_tile128_without_quantization(monkeypatch, env_overrides): + numerics = [] + monkeypatch.setattr(fp4_vsa_qat, "FP4VSAFineAttention", lambda config: numerics.append(config) or config) + video_sparse_attn_h3._fine_attention_from_name.cache_clear() + try: + env_overrides.enter_context(envs.FASTVIDEO_H3_VSA_FINE_ATTENTION.override("emulator-bf16")) + result = video_sparse_attn_h3.fine_attention_from_env() + assert result.quantize is False + assert result.prefix_precision is None + assert len(numerics) == 1 + finally: + video_sparse_attn_h3._fine_attention_from_name.cache_clear() + + +@pytest.mark.parametrize("override", [None, object()]) +def test_compile_warns_that_the_fine_override_is_ignored(monkeypatch, override): + warnings = [] + monkeypatch.setattr(video_sparse_attn_h3.logger, "warning_once", warnings.append) + impl = SimpleNamespace(fine_attention_override=override, layer_idx=3) + video_sparse_attn_h3.MiniMaxH3VSAImpl.prepare_for_compile(impl, torch.device("cpu")) + assert impl._compile_layer_idx.item() == 3 + assert len(warnings) == (override is not None) + if override is not None: + assert "FASTVIDEO_H3_VSA_FINE_ATTENTION is ignored inside torch.compile" in warnings[0] + + +@pytest.mark.parametrize("name", ["", "emulator", "mix8", "emulator-", "other-mix8", "other-bf16", + "emulator-unknown"]) +def test_invalid_fine_attention_names_raise_actionable_value_error(name): + with pytest.raises(ValueError, match="FASTVIDEO_H3_VSA_FINE_ATTENTION") as error: + video_sparse_attn_h3._fine_attention_from_name(name) + for accepted in ("emulator-mix8", "emulator-mixmx", "emulator-qk", "emulator-bf16"): + assert accepted in str(error.value) + + +@pytest.mark.parametrize("mode", ["mix8", "mixmx", "qk"]) +def test_valid_fine_attention_names_keep_implementation_paths(monkeypatch, mode): + monkeypatch.setattr(fp4_vsa_qat, "FP4VSAFineAttention", lambda config: ("emulator", config)) + video_sparse_attn_h3._fine_attention_from_name.cache_clear() + try: + kind, config = video_sparse_attn_h3._fine_attention_from_name(f"emulator-{mode}") + assert kind == "emulator" + assert config.prefix_precision == {"mix8": "fp8", "mixmx": "mxfp8", "qk": None}[mode] + assert config.quantize and not config.quantize_pv and config.smooth_k + finally: + video_sparse_attn_h3._fine_attention_from_name.cache_clear() diff --git a/fastvideo/tests/attention/test_qat_vsa_inference_memory.py b/fastvideo/tests/attention/test_qat_vsa_inference_memory.py new file mode 100644 index 000000000..d00f5d949 --- /dev/null +++ b/fastvideo/tests/attention/test_qat_vsa_inference_memory.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU checks of the emulator's host allocation decisions; kernel launches are replaced. + +Load the two host definitions from the kernel file so these tests do not need +Triton or a CUDA driver. Numerical forward/backward coverage lives in +fastvideo-kernel/tests/test_attn_qat_vsa_train.py and requires CUDA. +""" +import ast +import math +from pathlib import Path + +import pytest +import torch + + +def _host_api(calls): + path = Path(__file__).resolve().parents[3] / "fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_train.py" + tree = ast.parse(path.read_text()) + tree.body = [node for node in tree.body if isinstance(node, (ast.ClassDef, ast.FunctionDef)) + and node.name in {"_FP4VSAAttnQAT", "fp4_vsa_attn_qat"}] + + class Launch: + def __getitem__(self, grid): + def run(*args, **kwargs): + calls.append(kwargs["SAVE_MRUN"]) + args[3].copy_(args[0]) + args[4].copy_(args[0]) + args[5].zero_() + return run + + class BackwardLaunch: + def __init__(self, outputs): + self.outputs = outputs + + def __getitem__(self, grid): + def run(*args, **kwargs): + for index in self.outputs: + args[index].fill_(1) + return run + + namespace = dict(torch=torch, math=math, BLOCK=128, LOG2E=math.log2(math.e), LOG2_INV_TWO_LEVEL=0.0, + _ROWS=64, _WARPS=8, _fwd_kernel=Launch(), + fp8_per_token_fake_quant=lambda value: value, mxfp8_fake_quant=lambda value: value, + _bwd_dq_kernel=BackwardLaunch((6,)), _bwd_dkdv_kernel=BackwardLaunch((7, 8)), + invert_block_lists=lambda idx, count, n: (idx, idx, count), + fake_quant_qkv_triton=lambda q, k, v, **kw: (q, k, v)) + exec(compile(tree, str(path), "exec"), namespace) + return namespace["fp4_vsa_attn_qat"] + + +@pytest.mark.parametrize("grad_enabled,requires_grad,quantize_pv,prefix,expected", [ + (False, False, True, None, [False]), + (False, True, True, None, [False]), + (True, False, True, None, [False]), + (True, True, True, None, [True]), + (True, True, False, None, [False]), + (False, True, True, "bf16", [False, False]), + (True, True, True, "bf16", [True, False]), +]) +def test_running_maxima_only_allocated_for_quantized_p_backward( + monkeypatch, grad_enabled, requires_grad, quantize_pv, prefix, expected): + launches = [] + allocations = [] + attention = _host_api(launches) + q, k, v = [torch.ones(1, 2, 384, 128, dtype=torch.bfloat16, requires_grad=requires_grad) for _ in range(3)] + indices = torch.zeros(1, 2, 3, 3, dtype=torch.int32) + counts = torch.ones(1, 2, 3, dtype=torch.int32) + original = torch.empty + + def allocate(shape, **kwargs): + allocations.append(tuple(shape)) + return original(shape, **kwargs) + + monkeypatch.setattr(torch, "empty", allocate) + with torch.set_grad_enabled(grad_enabled): + out = attention(q, k, v, indices, counts, quantize_pv=quantize_pv, + prefix_precision=prefix, num_prefix_tiles=1 if prefix else None) + assert launches == expected + assert allocations.count((1, 2, 3, 3, 128)) == sum(expected) + assert out.requires_grad is (grad_enabled and requires_grad) + if out.requires_grad: + out.sum().backward() + assert all(t.grad is not None and torch.isfinite(t.grad).all() for t in (q, k, v)) diff --git a/fastvideo/tests/attention/test_vsa_fp4_prefix_split.py b/fastvideo/tests/attention/test_vsa_fp4_prefix_split.py new file mode 100644 index 000000000..8b555ddec --- /dev/null +++ b/fastvideo/tests/attention/test_vsa_fp4_prefix_split.py @@ -0,0 +1,297 @@ +# SPDX-License-Identifier: Apache-2.0 +"""The FP4 prefix split's row bookkeeping and dense prefix attention. + +The split takes the dense prefix query tiles (text, reference, audio rows) off +the FP4 sparse kernel and attends them densely at 8-bit Q/K. These tests cover +the parts that do not need a Blackwell GPU: the FP8 rounding, the valid-key +vector the kernel's half counts expand to, and the dense attention itself - +checked against the training emulator, which is the numeric spec. +""" + +from __future__ import annotations + +import pytest +import torch + +from fastvideo.models.dits.minimax_h3_vsa_fp4 import (_fp8_per_token, _half_row_valid, _prefix_dense_attention, + _resolve_prefix_mode) + +HALF = 64 + + +def _index(valid: torch.Tensor) -> torch.Tensor: + return valid.nonzero().squeeze(1) + + +def _emulator_fp8(x: torch.Tensor) -> torch.Tensor: + """The emulator's per-token E4M3 fake quant; it needs Triton to import.""" + try: + from fastvideo_kernel.triton_kernels.attn_qat_vsa_train import fp8_per_token_fake_quant + except (ImportError, KeyError): # KeyError: a failed fastvideo_kernel import left a None entry + pytest.skip("the Triton emulator is unavailable") + return fp8_per_token_fake_quant(x) + + +def _reference_dense_attention(q, k, v, valid, scale): + """fp32 softmax attention over ``[B, L, H, D]`` with the invalid keys masked out.""" + scores = torch.einsum("blhd,bkhd->bhlk", q.float(), k.float()) * scale + scores = scores.masked_fill(~valid[None, None, None, :], float("-inf")) + return torch.einsum("bhlk,bkhd->blhd", torch.softmax(scores, dim=-1), v.float()).to(v.dtype) + + +def _case(rows=1, prefix=8, kv_len=16, heads=2, dim=128, tail_invalid=4): + torch.manual_seed(0) + q = torch.randn(rows, prefix, heads, dim, dtype=torch.bfloat16) + k = torch.randn(rows, kv_len, heads, dim, dtype=torch.bfloat16) + v = torch.randn(rows, kv_len, heads, dim, dtype=torch.bfloat16) + valid = torch.ones(kv_len, dtype=torch.bool) + valid[kv_len - tail_invalid:] = False + return q, k, v, valid, dim + + +def test_fp8_per_token_rounds_within_one_e4m3_step() -> None: + """FP8 E4M3 has 3 mantissa bits, so a per-row scale bounds the error at ~6%.""" + torch.manual_seed(0) + x = torch.randn(4, 3, 5, 128, dtype=torch.bfloat16) + y = _fp8_per_token(x) + assert y.shape == x.shape and y.dtype == torch.bfloat16 + assert torch.allclose(y, x, rtol=0.1, atol=1e-2) + # Per (row, head) scale: a row whose magnitude is 1000x the others keeps its own precision. + x[0, 0, 0] *= 1000.0 + assert ((_fp8_per_token(x)[0, 0, 0] - x[0, 0, 0]).norm() / x[0, 0, 0].norm()).item() < 0.1 + + +def test_fp8_per_token_keeps_zero_rows_zero() -> None: + """Pad rows stay exactly zero, which the kernel's zero-padded k/v relies on.""" + x = torch.zeros(2, 2, 128, dtype=torch.bfloat16) + assert torch.equal(_fp8_per_token(x), x) + + +def test_fp8_per_token_matches_the_emulator() -> None: + torch.manual_seed(0) + x = torch.randn(1, 16, 2, 128, dtype=torch.bfloat16) + assert torch.equal(_fp8_per_token(x), _emulator_fp8(x).to(torch.bfloat16)) + + +def test_half_row_valid_expands_counts_to_rows() -> None: + kv_valid = torch.tensor([HALF, 3, 0, HALF], dtype=torch.int32) + valid = _half_row_valid(kv_valid, HALF) + assert valid.shape == (4 * HALF, ) + assert valid.sum().item() == 2 * HALF + 3 + assert valid[:HALF].all() and not valid[HALF + 3:HALF + HALF].any() + assert not valid[2 * HALF:3 * HALF].any() + + +@pytest.mark.parametrize(("mode", "steps", "step", "expected"), [ + ("fp8", -1, 0, "fp8"), + ("fp8", -1, 7, "fp8"), + ("bf16", 3, 2, "bf16"), + ("fp8", 3, 3, ""), + ("fp8", 0, 0, ""), + ("", 4, 0, ""), +]) +def test_prefix_mode_step_cutoff(mode: str, steps: int, step: int, expected: str) -> None: + """A non-negative step limit applies the split only to the leading steps.""" + assert _resolve_prefix_mode(mode, steps, step) == expected + + +@pytest.mark.parametrize("precision", ["fp8", "bf16"]) +def test_prefix_dense_attention_matches_the_emulator_numerics(precision: str) -> None: + """Smooth K, then round it - the order the emulator and the student use.""" + q, k, v, valid, dim = _case() + + out = _prefix_dense_attention(q, k, v, q.shape[1], k.shape[1], _index(valid), precision) + + assert out.shape == (1, 2, 8, dim) and out.dtype == torch.bfloat16 + k_ref = k - k.mean(dim=1, keepdim=True) + q_ref = q + if precision == "fp8": + q_ref, k_ref = _emulator_fp8(q), _emulator_fp8(k_ref) + expected = _reference_dense_attention(q_ref, k_ref, v, valid, dim**-0.5).transpose(1, 2) + + assert (out.float() - expected.float()).norm() / expected.float().norm() < 2e-2 + + +def test_smoothing_happens_before_the_fp8_rounding() -> None: + """Rounding before smoothing gives different keys, so the order is load-bearing.""" + q, k, v, valid, _ = _case() + smooth_then_round = _emulator_fp8(k - k.mean(dim=1, keepdim=True)) + round_then_smooth = _emulator_fp8(k) - _emulator_fp8(k).mean(dim=1, keepdim=True) + assert not torch.equal(smooth_then_round, round_then_smooth) + + out = _prefix_dense_attention(q, k, v, q.shape[1], k.shape[1], _index(valid), "fp8") + expected = _reference_dense_attention(q, smooth_then_round, v, valid, k.shape[-1]**-0.5).transpose(1, 2) + wrong = _reference_dense_attention(q, round_then_smooth, v, valid, k.shape[-1]**-0.5).transpose(1, 2) + + error = (out.float() - expected.float()).norm() + assert error < (out.float() - wrong.float()).norm() + + +def test_prefix_dense_attention_ignores_invalid_keys() -> None: + """Masked-out keys must not reach the output, however extreme they are.""" + torch.manual_seed(0) + rows, prefix, valid_len, heads, dim = 1, 4, 6, 2, 128 + q = torch.randn(rows, prefix, heads, dim, dtype=torch.bfloat16) + k = torch.randn(rows, valid_len + 2, heads, dim, dtype=torch.bfloat16) + v = torch.randn(rows, valid_len + 2, heads, dim, dtype=torch.bfloat16) + valid = torch.ones(valid_len + 2, dtype=torch.bfloat16).bool() + valid[valid_len:] = False + + masked = _prefix_dense_attention(q, k, v, prefix, valid_len + 2, _index(valid), "bf16") + + # Truncating to the valid rows changes only the smoothing mean, an exact shift. + k[:, valid_len:] = 0.0 + v[:, valid_len:] = 0.0 + k = k - (k[:, :valid_len].mean(dim=1, keepdim=True) - k.mean(dim=1, keepdim=True)) + truncated = _prefix_dense_attention(q, k[:, :valid_len], v[:, :valid_len], prefix, valid_len, + torch.arange(valid_len), "bf16") + + assert (masked.float() - truncated.float()).norm() / truncated.float().norm() < 5e-3 + + +@pytest.mark.parametrize("tile,prefix_tiles,n_tiles", [(64, 1, 5), (64, 2, 5), (64, 3, 6), (128, 1, 3)]) +@pytest.mark.parametrize("heads", [1, 4]) +def test_prefix_split_keeps_original_video_rows_and_quadrants(monkeypatch, tile, prefix_tiles, n_tiles, heads): + from types import SimpleNamespace + from fastvideo.models.dits import minimax_h3_vsa_fp4 as impl + + logical = tile * n_tiles + padded = (logical + 127) // 128 * 128 + prefix_rows = prefix_tiles * tile + head = prefix_rows // 128 + query = torch.arange(padded).view(1, padded, 1, 1).expand(1, padded, heads, 8).float() + key, value = torch.randn_like(query), torch.randn_like(query) + blocks = padded // 128 + indices = torch.arange(heads * blocks * blocks, dtype=torch.int32).view(1, heads, blocks, blocks) + counts = torch.ones(1, heads, blocks, dtype=torch.int32) + quadrants = (indices % 15 + 1).to(torch.uint8) if tile == 64 else None + valid = torch.full((padded // 64,), 64, dtype=torch.int32) + valid[logical // 64 - 1] = 17 + valid[logical // 64:] = 0 + calls = [] + + def native(q, k, v, idx, num, kv, quad, **kwargs): + calls.append(q.shape[1]) + assert q.shape[1] % 128 == 0 + assert k is key and v is value and kv is valid + assert torch.equal(q, query[:, head * 128:]) + assert torch.equal(idx, indices[:, :, head:]) and idx.is_contiguous() + assert torch.equal(num, counts[:, :, head:]) and num.is_contiguous() + if quadrants is None: + assert quad is None + else: + assert torch.equal(quad, quadrants[:, :, head:]) and quad.is_contiguous() + return q.transpose(1, 2) + + monkeypatch.setattr(impl, "_sparse_fp4_attention", lambda api, *args, **kwargs: native(*args, **kwargs)) + api = SimpleNamespace(KEY_LENGTH_SIZING=True, HALF_N=64) + result = impl._fp4_fine_attention(api, query, key, value, indices, counts, valid, quadrants, + logical=logical, prefix_rows=prefix_rows, precision="bf16", options={}) + expected_prefix = impl._prefix_dense_attention(query, key, value, prefix_rows, logical, + _index(impl._half_row_valid(valid, 64)[:logical]), "bf16") + torch.testing.assert_close(result[:, :, :prefix_rows], expected_prefix) + torch.testing.assert_close(result[:, :, prefix_rows:], query[:, prefix_rows:logical].transpose(1, 2)) + assert result.shape == (1, heads, logical, 8) + assert calls == [padded - head * 128] + layout = impl._KeyIndexCache() + layout.tile, layout.n_tiles = tile, n_tiles + cached = impl._fp4_fine_attention(api, query, key, value, indices, counts, valid, quadrants, logical=logical, + prefix_rows=prefix_rows, precision="bf16", options={}, layout=layout) + assert torch.equal(cached, result) + + +def test_layout_computes_the_valid_key_index_once(monkeypatch) -> None: + """kv_valid depends only on the tile sizes, so every block of a step shares one index (one host sync).""" + from fastvideo.models.dits import minimax_h3_vsa_fp4 as impl + + calls = [] + real = impl._valid_key_index + monkeypatch.setattr(impl, "_valid_key_index", lambda *args: calls.append(args) or real(*args)) + layout = impl._KeyIndexCache() + layout.tile, layout.n_tiles = HALF, 3 + kv_valid = torch.tensor([HALF, 17, 0, 0], dtype=torch.int32) + first = layout.valid_key_index(kv_valid, HALF) + assert layout.valid_key_index(kv_valid.clone(), HALF) is first + assert len(calls) == 1 + assert torch.equal(first, torch.arange(HALF + 17)) + + +@pytest.mark.parametrize("tile", [64, 128]) +def test_explicit_prefix_precision_overrides_environment(tile): + from types import SimpleNamespace + from fastvideo import envs + from fastvideo.models.dits.minimax_h3_vsa_fp4 import _prefix_settings + + meta = SimpleNamespace(tile_elems=tile, num_prefix_tiles=3, current_timestep=0) + with envs.FASTVIDEO_H3_VSA_FP4_PREFIX.override(""), envs.FASTVIDEO_H3_VSA_FP4_PREFIX_STEPS.override(-1): + assert _prefix_settings(SimpleNamespace(_vsa_fp4_options={"prefix_precision": "fp8"}), meta, 5) == ( + "fp8", 3 * tile) + assert _prefix_settings(SimpleNamespace(), meta, 5) == ("", 0) + with envs.FASTVIDEO_H3_VSA_FP4_PREFIX.override("fp8"), envs.FASTVIDEO_H3_VSA_FP4_PREFIX_STEPS.override(0): + assert _prefix_settings(SimpleNamespace(), meta, 5) == ("", 0) + + +def test_explicit_video_numerics_override_environment(): + from types import SimpleNamespace + from fastvideo import envs + from fastvideo.models.dits.minimax_h3_vsa_fp4 import _sparse_fp4_attention + + key = torch.randn(1, 128, 2, 8) + captured = [] + api = SimpleNamespace(sageattn_blackwell_sparse_bshd=lambda *args, **kwargs: captured.append((args, kwargs))) + with envs.FASTVIDEO_H3_VSA_FP4_K_SMOOTH.override(False), envs.FASTVIDEO_H3_VSA_FP4_TWO_LEVEL_P.override(False): + _sparse_fp4_attention(api, key, key, key, two_level_p=True, smooth_k=True) + torch.testing.assert_close(captured[0][0][1], key - key.mean(dim=1, keepdim=True)) + assert captured[0][1] == {"validate": False, "single_level_p_quant": False} + captured.clear() + with envs.FASTVIDEO_H3_VSA_FP4_K_SMOOTH.override(True), envs.FASTVIDEO_H3_VSA_FP4_TWO_LEVEL_P.override(True): + _sparse_fp4_attention(api, key, key, key, two_level_p=False, smooth_k=False) + assert captured[0][0][1] is key + assert captured[0][1] == {"validate": False} + + +def test_missing_native_extension_reports_required_build(monkeypatch): + import builtins + from fastvideo.models.dits import minimax_h3_vsa_fp4 as impl + + original = builtins.__import__ + + def missing(name, *args, **kwargs): + if name == "attn_qat_infer.api": + raise ImportError("test missing extension") + return original(name, *args, **kwargs) + + monkeypatch.setattr(impl, "_fp4_api", None) + monkeypatch.setattr(builtins, "__import__", missing) + with pytest.raises(RuntimeError, match="Build fastvideo-kernel"): + impl._api() + + +@pytest.mark.parametrize("tile,prefix_tiles,n_tiles", [(64, 1, 5), (64, 3, 6), (128, 1, 3)]) +def test_native_prefix_split_preserves_video_attention(tile, prefix_tiles, n_tiles): + """Requires consumer Blackwell and the native sparse FP4 extension from this checkout.""" + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 12: + pytest.skip("requires sm_120 or sm_121 native FP4 attention") + api = pytest.importorskip("attn_qat_infer.api") + from fastvideo.models.dits import minimax_h3_vsa_fp4 as impl + + torch.manual_seed(1234) + logical, prefix_rows = n_tiles * tile, prefix_tiles * tile + padded = (logical + 127) // 128 * 128 + query, key, value = [torch.randn(1, padded, 4, 128, device="cuda", dtype=torch.bfloat16) for _ in range(3)] + valid = torch.full((n_tiles,), tile, device="cuda", dtype=torch.int32) + valid[-1] = tile // 2 + mask = torch.eye(n_tiles, device="cuda", dtype=torch.bool).expand(1, 4, -1, -1).clone() + mask[:, :, :prefix_tiles] = True + mask[:, :, :, :prefix_tiles] = True + lists = api.vsa_tile_mask_to_fp4_blocks(mask, tile, valid) + options = {"two_level_p": True, "smooth_k": True} + baseline = impl._fp4_fine_attention(api, query, key, value, *lists, logical=logical, prefix_rows=0, + precision="", options=options) + split = impl._fp4_fine_attention(api, query, key, value, *lists, logical=logical, prefix_rows=prefix_rows, + precision="fp8", options=options) + assert split.isfinite().all() + torch.testing.assert_close(split[:, :, prefix_rows:], baseline[:, :, prefix_rows:logical], rtol=0, atol=0) + expected = impl._prefix_dense_attention(query, key, value, prefix_rows, logical, + impl._valid_key_index(lists[2], 64, logical), "fp8") + torch.testing.assert_close(split[:, :, :prefix_rows], expected, rtol=0, atol=0) diff --git a/fastvideo/tests/attention/test_vsa_h3_ref2va_regions.py b/fastvideo/tests/attention/test_vsa_h3_ref2va_regions.py index cfe7728dd..cebde2ab9 100644 --- a/fastvideo/tests/attention/test_vsa_h3_ref2va_regions.py +++ b/fastvideo/tests/attention/test_vsa_h3_ref2va_regions.py @@ -479,9 +479,46 @@ def test_forward_tile128_even_tile_count_no_partner_tile(tile128): assert len(tile128.calls) == 1 -def test_forward_tile128_unavailable_cuda_kernel(tile128, monkeypatch): +def test_forward_tile128_falls_back_to_triton_when_cuda_kernel_unavailable(tile128, monkeypatch): + meta = _build_r2v(sparsity=0.5, tile_size=128) + n_tiles = meta.variable_block_sizes.numel() + tile128.supported = False + calls = [] + + def triton(q, k, v, mask, variable_block_sizes): + calls.append((q, mask, variable_block_sizes)) + return torch.zeros_like(q), None + + monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", triton) + _, out = _run_tile128(meta) + assert len(calls) == 1 + q, mask, vbs = calls[0] + assert q.shape[2] == n_tiles * 128 + assert mask.shape[-2:] == (2 * n_tiles, 2 * n_tiles) + expected_vbs = torch.stack((meta.variable_block_sizes.clamp(max=64), + (meta.variable_block_sizes - 64).clamp(min=0)), dim=-1).flatten() + torch.testing.assert_close(vbs, expected_vbs) + assert out.shape == (1, n_tiles * 128, _HEADS, _DIM) + + +@pytest.mark.parametrize("datacenter", [False, True]) +def test_forward_tile128_fallback_warns_only_on_datacenter_blackwell(tile128, monkeypatch, datacenter): meta = _build_r2v(sparsity=0.5, tile_size=128) tile128.supported = False + warnings = [] + monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", lambda q, k, v, mask, vbs: (torch.zeros_like(q), None)) + monkeypatch.setattr(vsa_h3, "_is_datacenter_blackwell", lambda device: datacenter) + monkeypatch.setattr(vsa_h3.logger, "warning_once", warnings.append) + _run_tile128(meta) + assert len(warnings) == int(datacenter) + if datacenter: + assert "falling back to the slower Triton tile-64 kernels" in warnings[0] + + +def test_forward_tile128_unavailable_cuda_kernel_and_triton(tile128, monkeypatch): + meta = _build_r2v(sparsity=0.5, tile_size=128) + tile128.supported = False + monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", None) with pytest.raises(RuntimeError, match="tile 128 requires the ptx_blackwell block-sparse kernel"): _run_tile128(meta) monkeypatch.setattr(vsa_h3, "_ptx_blackwell", None) @@ -489,7 +526,7 @@ def test_forward_tile128_unavailable_cuda_kernel(tile128, monkeypatch): _run_tile128(meta) -def test_forward_tile128_route_resolved_once_per_device(tile128): +def test_forward_tile128_route_resolved_once_per_device(tile128, monkeypatch): meta = _build_r2v(sparsity=0.5, tile_size=128) probes = [] tile128.is_supported = lambda q, variable_block_sizes: probes.append(q.device) or True @@ -499,14 +536,17 @@ def test_forward_tile128_route_resolved_once_per_device(tile128): query, key, value = impl.preprocess_qkv(raw, meta).chunk(3, dim=0) impl.forward(query, key, value, None, meta) assert len(tile128.calls) == 3 and probes == [_CPU] - # A failed probe is cached too, and every later call still fails closed. + # A failed probe is cached too; later calls take the Triton fallback without probing again. + calls = [] + monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", + lambda q, k, v, mask, vbs: (calls.append(q.shape), (torch.zeros_like(q), None))[1]) tile128.is_supported = lambda q, variable_block_sizes: probes.append(q.device) and False impl = _impl(_DIM) raw = torch.randn(3, meta.total_seq_length, _HEADS, _DIM, dtype=torch.bfloat16) query, key, value = impl.preprocess_qkv(raw, meta).chunk(3, dim=0) for _ in range(2): - with pytest.raises(RuntimeError, match="tile 128 requires the ptx_blackwell block-sparse kernel"): - impl.forward(query, key, value, None, meta) + impl.forward(query, key, value, None, meta) + assert len(calls) == 2 assert probes == [_CPU, _CPU] diff --git a/fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py b/fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py index 0c83eceda..468e01de7 100644 --- a/fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py +++ b/fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py @@ -4,8 +4,10 @@ from __future__ import annotations import contextlib import importlib.util +import sys import json from pathlib import Path +from types import SimpleNamespace import pytest @@ -33,7 +35,14 @@ def _load_example(): spec = importlib.util.spec_from_file_location("basic_fasth3_omniref_pdd", EXAMPLE_PATH) assert spec is not None and spec.loader is not None module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) + # `@dataclass` resolves the example's postponed annotations through the + # defining module's __dict__, so it has to be registered before it runs. + sys.modules[spec.name] = module + try: + spec.loader.exec_module(module) + except BaseException: + sys.modules.pop(spec.name, None) + raise return module @@ -119,10 +128,23 @@ def test_build_generator_config_leaves_trained_settings_to_checkpoint(tmp_path): def test_build_generator_config_multi_gpu_shards_dit(tmp_path, num_gpus, sharded): config = example.build_generator_config(tmp_path, num_gpus) assert config.engine.use_fsdp_inference is sharded + assert config.engine.execution_backend == "mp" assert config.engine.parallelism.sp_size == num_gpus assert config.pipeline.components.override_pipeline_cls_name == "MiniMaxH3Ref2VAModularPipeline" +@pytest.mark.parametrize("num_gpus,backend", [(1, "mp"), (2, "ray")]) +def test_multi_gpu_profile_runs_on_ray_without_fsdp(tmp_path, num_gpus, backend): + config = example.build_generator_config(tmp_path, num_gpus, profile=example.DEVICE_PROFILES["spark"]) + assert config.engine.execution_backend == backend + assert config.engine.use_fsdp_inference is False + + +@pytest.mark.parametrize("num_gpus", [1, 4]) +def test_lossless_accel_config_keeps_mp(tmp_path, num_gpus): + assert example.build_generator_config(tmp_path, num_gpus, lossless_accel=True).engine.execution_backend == "mp" + + 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) @@ -159,3 +181,171 @@ 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 + + +@pytest.mark.parametrize("choice", ["default", "fp8", "bf16"]) +def test_attention_prefix_flag_parses(choice): + argv = ["--model-path", "m", "--prompt", "p", "--image", "a.png", "--attention-prefix", choice] + assert example.parse_args(argv).attention_prefix == choice + + +@pytest.mark.parametrize(("choice", "expected"), [ + ("default", "fp8"), + ("fp8", "fp8"), + ("bf16", "bf16"), +]) +def test_attention_prefix_keeps_the_profile_default_or_overrides_it(choice, expected): + from fastvideo import envs + + with envs.FASTVIDEO_H3_VSA_FP4_PREFIX.override("fp8"): + example.apply_attention_prefix(choice) + assert envs.FASTVIDEO_H3_VSA_FP4_PREFIX.get() == expected + + +def test_4090_rejects_packed_export_before_downloading_base(tmp_path, monkeypatch): + export = tmp_path / "packed" + (export / "transformer_ref").mkdir(parents=True) + (export / "transformer_ref" / "nvfp4_weights.safetensors").touch() + (export / example.CONTRACT).write_text(json.dumps({**CONTRACT, "model_type": "ref2va"})) + calls = [] + + def snapshot(*args): + calls.append(args) + return export + + monkeypatch.setattr(example, "_snapshot", snapshot) + with pytest.raises(ValueError, match="requires a dense or FP8"): + example.resolve_model(_args("--profile", "rtx-4090")) + assert len(calls) == 1 + + +def test_4090_profile_uses_fp8_and_tile128_bf16_fallback(tmp_path): + import yaml + + profile = example.DEVICE_PROFILES["rtx-4090"] + config = example.build_generator_config(tmp_path, 1, profile=profile) + assert config.engine.quantization.transformer_quant == "FP8" + assert profile.env["FASTVIDEO_H3_VSA_FP4"] == "0" + assert profile.env["FASTVIDEO_H3_VSA_FINE_ATTENTION"] == "emulator-bf16" + recipe = yaml.safe_load(EXAMPLE_PATH.with_name("basic_fasth3_omniref_rtx4090.yaml").read_text()) + assert "NVFP4" not in recipe["generator"]["model_path"] + assert recipe["generator"]["engine"]["quantization"]["transformer_quant"] == "FP8" + + +def test_4090_accepts_dense_export(tmp_path, monkeypatch): + export = _snapshot_dir(tmp_path / "export", example.EXPORT_COMPONENTS, _EXPORT_ONLY_MANIFEST) + (export / example.CONTRACT).write_text(json.dumps({**CONTRACT, "model_type": "ref2va"})) + base = _snapshot_dir(tmp_path / "base", example.BASE_COMPONENTS, _FULL_MANIFEST) + snapshots = iter((export, base)) + monkeypatch.setattr(example, "_snapshot", lambda *args: next(snapshots)) + args = _args("--profile", "rtx-4090", "--composed-dir", str(tmp_path / "composed")) + composed, _ = example.resolve_model(args) + assert (composed / "transformer_ref").resolve() == (export / "transformer_ref").resolve() + + +@pytest.mark.parametrize("profile,dimensions,expected", [ + (None, (), (832, 480)), + ("rtx-pro-6000", (), (1344, 768)), + ("rtx-pro-6000", ("--width", "832", "--height", "480"), (832, 480)), + ("rtx-pro-6000", ("--width", "832"), (832, 768)), + ("rtx-pro-6000", ("--height", "480"), (1344, 480)), + ("rtx-pro-6000", ("--width", "1024", "--height", "640"), (1024, 640)), + (None, ("--width", "1024"), (1024, 480)), + (None, ("--height", "640"), (832, 640)), +]) +def test_main_preserves_explicit_dimensions_before_profile_defaults(tmp_path, monkeypatch, profile, dimensions, expected): + options = ["--output", str(tmp_path), *dimensions] + if profile: + options += ["--profile", profile] + args = _args(*options) + if "--width" not in dimensions: + assert args.width is None + if "--height" not in dimensions: + assert args.height is None + requests = [] + shutdown = [] + + def generate(request): + requests.append(request) + return SimpleNamespace(video_path="out.mp4", generation_time=None, peak_memory_mb=None) + + generator = SimpleNamespace(generate=generate, shutdown=lambda: shutdown.append(True)) + monkeypatch.setattr(example, "parse_args", lambda: args) + monkeypatch.setattr(example, "resolve_model", lambda args: (tmp_path, CONTRACT)) + monkeypatch.setattr(example, "apply_profile", lambda *args: None) + monkeypatch.setattr(example.VideoGenerator, "from_config", lambda config: generator) + example.main() + assert (requests[0].sampling.width, requests[0].sampling.height) == expected + assert shutdown == [True] + + +@pytest.mark.parametrize("profile", ["spark", "spark_pair"]) +@pytest.mark.parametrize("version", ["v2", "qad_v2"]) +def test_spark_recipe_selects_its_checkpoint_and_fp4(profile, version, monkeypatch, tmp_path): + import yaml + from fastvideo.api.compat import generator_config_to_fastvideo_args + from fastvideo.api.parser import parse_config + from fastvideo.api.schema import GeneratorConfig + from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig + from fastvideo.fastvideo_args import FastVideoArgs + + path = EXAMPLE_PATH.with_name(f"basic_fasth3_{profile}_{version}_nvfp4.yaml") + raw = yaml.safe_load(path.read_text())["generator"] + expected_model = { + "v2": "FastVideo/FastVideo-FastH3-8-Step-V2-NVFP4-Consumer", + "qad_v2": "FastVideo/FastH3-QAD-v2", + }[version] + assert raw["model_path"] == expected_model + monkeypatch.setattr("fastvideo.registry.get_pipeline_config_cls_from_name", lambda path: MiniMaxH3PipelineConfig) + original_from_kwargs = FastVideoArgs.from_kwargs + + def from_kwargs(**kwargs): + kwargs["model_path"] = str(tmp_path) + kwargs["pipeline_config"] = MiniMaxH3PipelineConfig() + return original_from_kwargs(**kwargs) + + # Supply the pipeline locally so this configuration test never downloads a checkpoint. + monkeypatch.setattr(FastVideoArgs, "from_kwargs", from_kwargs) + args = generator_config_to_fastvideo_args(parse_config(GeneratorConfig, raw)) + assert args.h3_vsa_fp4["enabled"] is True + assert args.VSA_tile_size == 64 and args.VSA_sparsity == 0.8 + assert args.sp_size == (2 if profile == "spark_pair" else 1) + if version == "qad_v2": + assert args.h3_vsa_fp4 == {"enabled": True, "prefix_precision": "fp8", "two_level_p": True, "smooth_k": True} + + +@pytest.mark.parametrize("device", ["spark", "spark_pair", "rtx5090", "rtx_pro6000"]) +def test_blackwell_omniref_yaml_enables_protected_fp4(device): + import yaml + + raw = yaml.safe_load(EXAMPLE_PATH.with_name(f"basic_fasth3_omniref_{device}.yaml").read_text()) + assert raw["generator"]["pipeline"]["experimental"]["h3_vsa_fp4"] == { + "enabled": True, "prefix_precision": "fp8", "two_level_p": True, "smooth_k": True} + + +@pytest.mark.parametrize("options", [None, {"unknown": True}, {"enabled": "true"}, {"smooth_k": 1}, + {"prefix_precision": "fp4"}]) +def test_invalid_fp4_config_is_rejected_before_loading(options): + from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig + from fastvideo.fastvideo_args import FastVideoArgs + + with pytest.raises(ValueError, match="h3_vsa_fp4"): + FastVideoArgs(model_path="unused", pipeline_config=MiniMaxH3PipelineConfig(), h3_vsa_fp4=options) + + +@pytest.mark.parametrize("compile_flag", ["inference_torch_compile", "enable_torch_compile"]) +def test_explicit_fp4_cannot_silently_compile_to_bf16(compile_flag): + from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig + from fastvideo.fastvideo_args import FastVideoArgs + + with pytest.raises(ValueError, match="eager DiT inference"): + FastVideoArgs(model_path="unused", pipeline_config=MiniMaxH3PipelineConfig(), h3_vsa_fp4={"enabled": True}, + **{compile_flag: True}) + + +def test_row_split_is_an_explicit_spark_opt_in(): + from fastvideo import envs + + assert envs.environment_variables["FASTVIDEO_NVFP4_MM_ROW_SPLIT"].default == "off" + assert example.DEVICE_PROFILES["spark"].env["FASTVIDEO_NVFP4_MM_ROW_SPLIT"] == "table" + assert "FASTVIDEO_NVFP4_MM_ROW_SPLIT" not in example.DEVICE_PROFILES["rtx-pro-6000"].env diff --git a/fastvideo/tests/loader/test_lora_patch.py b/fastvideo/tests/loader/test_lora_patch.py index d350822e4..e4b10ed06 100644 --- a/fastvideo/tests/loader/test_lora_patch.py +++ b/fastvideo/tests/loader/test_lora_patch.py @@ -371,11 +371,13 @@ def test_a_fully_applied_payload_warns_about_nothing(tmp_path, caplog): assert [r for r in caplog.records if r.levelno >= logging.WARNING] == [] -def _stub_packed_nvfp4_export(monkeypatch): +def _stub_packed_nvfp4_export(monkeypatch, tmp_path): from fastvideo.layers.quantization import nvfp4_config + packed_path = tmp_path / "packed.safetensors" + save_file({"placeholder": torch.zeros(1)}, str(packed_path)) loaded = [] - monkeypatch.setattr(nvfp4_config, "find_minimax_h3_nvfp4_dit_export", lambda _paths: "packed.safetensors") + monkeypatch.setattr(nvfp4_config, "find_minimax_h3_nvfp4_dit_export", lambda _paths: str(packed_path)) monkeypatch.setattr(nvfp4_config, "nvfp4_linear_weight_param_names", lambda _model: {"linear.weight"}) monkeypatch.setattr(nvfp4_config, "load_minimax_h3_nvfp4_dit_export", lambda _model, path, device: loaded.append(path)) @@ -395,17 +397,17 @@ def _packed_init_params(): return {"config": SimpleNamespace(quant_config=SimpleNamespace(layer_profile="h3_dit"))} -def test_packed_nvfp4_export_loads_with_a_low_rank_only_adapter(monkeypatch): - loaded = _stub_packed_nvfp4_export(monkeypatch) +def test_packed_nvfp4_export_loads_with_a_low_rank_only_adapter(monkeypatch, tmp_path): + loaded = _stub_packed_nvfp4_export(monkeypatch, tmp_path) monkeypatch.setattr(DenseLoRAPatch, "from_adapter", classmethod(lambda *_args, **_kwargs: None)) _load_tiny_model(_PackedTinyModel, lora_path="adapter.safetensors", init_params=_packed_init_params()) - assert loaded == ["packed.safetensors"] + assert loaded == [str(tmp_path / "packed.safetensors")] -def test_packed_nvfp4_export_rejects_an_adapter_with_dense_payload(monkeypatch): - loaded = _stub_packed_nvfp4_export(monkeypatch) +def test_packed_nvfp4_export_rejects_an_adapter_with_dense_payload(monkeypatch, tmp_path): + loaded = _stub_packed_nvfp4_export(monkeypatch, tmp_path) patch = DenseLoRAPatch([], {"linear.weight": ("adapter.safetensors", "linear.diff")}, {}, 1.0) monkeypatch.setattr(DenseLoRAPatch, "from_adapter", classmethod(lambda *_args, **_kwargs: patch)) diff --git a/fastvideo/tests/loader/test_minimax_h3_pdd_loading.py b/fastvideo/tests/loader/test_minimax_h3_pdd_loading.py index 808a0a2ed..12fce7478 100644 --- a/fastvideo/tests/loader/test_minimax_h3_pdd_loading.py +++ b/fastvideo/tests/loader/test_minimax_h3_pdd_loading.py @@ -130,7 +130,7 @@ def _write_transformer_ref(root, *, config_pdd_steps, weight_pdd_steps=PDD_STEPS return path, tensors -def _load(path): +def _load(path, **overrides): args = SimpleNamespace( pipeline_config=MiniMaxH3PipelineConfig(), override_transformer_cls_name=None, @@ -151,6 +151,8 @@ def _load(path): inference_mode=True, dit_layerwise_offload=False, ) + for name, value in overrides.items(): + setattr(args, name, value) return component_loader.TransformerLoader().load(str(path), args) @@ -229,3 +231,14 @@ def test_fuse_pdd_block_weighted_mean_of_widened_heads(tmp_path, cpu_loader, sin expected = torch.einsum("n,brnc->brc", alpha.float(), all_heads[..., 1:3, :].float()) assert block.shape == expected.shape torch.testing.assert_close(block.float(), expected, rtol=1e-4, atol=1e-5) + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_load_explicit_fp4_settings_reaches_main_blocks_only(tmp_path, vsa_gates, enabled): + path, _ = _write_transformer_ref(tmp_path, config_pdd_steps=PDD_STEPS) + options = {"enabled": enabled, "prefix_precision": "fp8", "two_level_p": True, "smooth_k": True} + with envs.FASTVIDEO_H3_VSA_FP4.override(not enabled): + model = _load(path, h3_vsa_fp4=options) + assert all(block.attn._vsa_fp4 is enabled for block in model.transformer_blocks) + assert all(block.attn._vsa_fp4_options == options for block in model.transformer_blocks) + assert not model.token_refiner.refiner_blocks[0].attn._vsa_fp4 diff --git a/fastvideo/tests/loader/test_prequantized_fp8.py b/fastvideo/tests/loader/test_prequantized_fp8.py new file mode 100644 index 000000000..efc3d91b1 --- /dev/null +++ b/fastvideo/tests/loader/test_prequantized_fp8.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Packed H3 FP8 exports preserve dense linears through post-load conversion. + +Uses a tiny synthetic checkpoint on CPU, plus CUDA when available. No model +downloads or distributed process groups are required. +""" +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import save_file + +from fastvideo.layers.linear import ReplicatedLinear, UnquantizedLinearMethod +from fastvideo.layers.quantization.fp8_config import FP8Config, FP8QuantizeMethod +from fastvideo.models.loader.fsdp_load import maybe_load_fsdp_model + + +class MiniMaxH3Tiny(torch.nn.Module): + # Different checkpoint names exercise source-to-runtime name mapping. + param_names_mapping = {r"^export_main\.": "main.", r"^export_refiner\.": "refiner."} + + def __init__(self, config): + super().__init__() + self.main = ReplicatedLinear(128, 64, bias=False, params_dtype=torch.bfloat16, + quant_config=config.quant_config, prefix="transformer_blocks.0.attn.to_q") + self.refiner = ReplicatedLinear(128, 64, bias=True, params_dtype=torch.bfloat16, + quant_config=config.quant_config, prefix="token_refiner.0.attn.to_q") + + +@pytest.fixture(params=["cpu", "cuda"]) +def device(request): + if request.param == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA companion check requires a GPU") + if request.param == "cuda": + return torch.device("cuda", torch.cuda.current_device()) + return torch.device(request.param) + + +def _load(paths, device, quant_config): + return maybe_load_fsdp_model( + model_cls=MiniMaxH3Tiny, + init_params={"config": SimpleNamespace(quant_config=quant_config)}, + weight_dir_list=[str(p) for p in paths], + device=device, + hsdp_replicate_dim=1, + hsdp_shard_dim=1, + default_dtype=torch.bfloat16, + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + cpu_offload=True, # Stage source tensors on CPU on both test devices. + training_mode=False, + ) + + +def _dense_weights(): + generator = torch.Generator().manual_seed(42) + return { + "export_main.weight": torch.randn(64, 128, generator=generator).to(torch.bfloat16), + "export_refiner.weight": torch.randn(64, 128, generator=generator).to(torch.bfloat16), + "export_refiner.bias": torch.randn(64, generator=generator).to(torch.bfloat16), + } + + +def test_packed_fp8_preserves_dense_refiner(tmp_path, device): + weights = _dense_weights() + packed_weight = weights.pop("export_main.weight").to(torch.float8_e4m3fn) + scale = torch.linspace(0.01, 0.02, 64, dtype=torch.float32).reshape(64, 1) + # Weight and scale may live in separate shards. + first, second = tmp_path / "first.safetensors", tmp_path / "second.safetensors" + save_file({"export_main.weight": packed_weight}, str(first)) + save_file({**weights, "export_main.weight_scale": scale}, str(second)) + + model = _load([first, second], device, FP8Config()) + + assert isinstance(model.main.quant_method, FP8QuantizeMethod) + assert model.main.quant_method.granularity == "channel" + assert model.main._fp8_weight.dtype == torch.float8_e4m3fn + assert model.main._fp8_weight.device == device + assert torch.equal(model.main._fp8_weight.cpu().float(), packed_weight.float()) + assert torch.equal(model.main._fp8_weight_scale.cpu(), scale.flatten()) + assert getattr(model.main, "weight", None) is None + + assert isinstance(model.refiner.quant_method, UnquantizedLinearMethod) + assert not hasattr(model.refiner, "_fp8_weight") + assert model.refiner.weight.dtype == torch.bfloat16 + assert torch.equal(model.refiner.weight.cpu(), weights["export_refiner.weight"]) + assert torch.equal(model.refiner.bias.cpu(), weights["export_refiner.bias"]) + x = torch.ones(16, 128, dtype=torch.bfloat16, device=device) + expected = torch.nn.functional.linear(x, weights["export_refiner.weight"].to(device), + weights["export_refiner.bias"].to(device)) + assert torch.equal(model.refiner(x)[0], expected) + if device.type == "cuda": + main_output = model.main(x)[0] + assert main_output.shape == (16, 64) + assert torch.isfinite(main_output).all() + + +def test_dense_checkpoint_still_uses_runtime_fp8(tmp_path, device): + path = tmp_path / "dense.safetensors" + save_file(_dense_weights(), str(path)) + + model = _load([path], device, FP8Config()) + + for layer in (model.main, model.refiner): + assert isinstance(layer.quant_method, FP8QuantizeMethod) + assert layer._fp8_weight.dtype == torch.float8_e4m3fn + assert getattr(layer, "weight", None) is None + + +def test_packed_fp8_requires_matching_quant_config(tmp_path): + path = tmp_path / "packed.safetensors" + weights = _dense_weights() + weights["export_main.weight"] = weights["export_main.weight"].to(torch.float8_e4m3fn) + weights["export_main.weight_scale"] = torch.ones(64, dtype=torch.float32) + save_file(weights, str(path)) + + with pytest.raises(RuntimeError, match="not an FP8 linear"): + _load([path], torch.device("cpu"), None) diff --git a/fastvideo/tests/ops/quantization/test_nvfp4_row_split.py b/fastvideo/tests/ops/quantization/test_nvfp4_row_split.py new file mode 100644 index 000000000..da186fc18 --- /dev/null +++ b/fastvideo/tests/ops/quantization/test_nvfp4_row_split.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Row chunking for the NVFP4 mm_fp4 wrapper: chunks must tile the rows exactly, on 128-row boundaries.""" +import math + +import pytest +import torch + +from fastvideo.layers.quantization import nvfp4_config +from fastvideo.layers.quantization.nvfp4_config import _mm_fp4_row_chunks + + +@pytest.mark.parametrize("rows", [35328, 35000, 16384, 20481]) +@pytest.mark.parametrize("parts", [1, 2, 4]) +def test_row_chunks_cover_rows_in_order(rows: int, parts: int) -> None: + a = torch.arange(rows * 2, dtype=torch.float32).view(rows, 2) + scale = torch.arange(rows + (-rows) % 128, dtype=torch.float32).view(-1, 1) + chunks = _mm_fp4_row_chunks(a, scale, parts) + assert torch.equal(torch.cat([c for c, _ in chunks]), a) + # every chunk start is a 128-row scale tile boundary + offset = 0 + for chunk, chunk_scale in chunks: + assert offset % 128 == 0 + assert chunk_scale[0, 0].item() == offset + scale_hi = min(math.ceil((offset + chunk.shape[0]) / 128) * 128, scale.shape[0]) + assert chunk_scale.shape[0] == scale_hi - offset + offset += chunk.shape[0] + assert offset == rows + + +@pytest.mark.parametrize(("use_8x4", "a_scale_rows"), [(True, 20000), (False, None)]) +def test_row_split_stays_whole_without_128x4_scale_rows(monkeypatch, use_8x4, a_scale_rows): + calls = [] + + def fake_chunked(*args): + calls.append(args[-1]) + return None + + monkeypatch.setattr(nvfp4_config, "_mm_fp4_chunked", fake_chunked) + a = torch.zeros(256, 64) + a_scale = torch.zeros(256, 1) if a_scale_rows else torch.zeros(256) + assert nvfp4_config._mm_fp4_row_split(a, torch.zeros(32, 64), a_scale, torch.zeros(32, 1), None, + torch.bfloat16, 16, use_8x4, "cudnn", True) is None + assert calls == [1] + + +@pytest.mark.parametrize("policy,expected", [("off", "whole"), ("table", "split")]) +def test_large_cudnn_call_splits_only_when_opted_in(monkeypatch, policy, expected): + from fastvideo import envs + + calls = [] + monkeypatch.setattr(nvfp4_config, "_register_ops_once", lambda: None) + monkeypatch.setattr(torch.ops.fastvideo_fp4, "mm_fp4", lambda *args: calls.append("whole"), raising=False) + monkeypatch.setattr(nvfp4_config, "_mm_fp4_row_split", lambda *args: calls.append("split")) + with envs.FASTVIDEO_NVFP4_MM_ROW_SPLIT.override(policy): + nvfp4_config._mm_fp4(torch.empty(16384, 32), torch.empty(16, 64), torch.empty(16384, 1), + torch.empty(16, 1), None, torch.bfloat16, None, backend="cudnn") + assert calls == [expected] diff --git a/fastvideo/tests/ops/quantization/test_nvfp4_skip_blocks.py b/fastvideo/tests/ops/quantization/test_nvfp4_skip_blocks.py new file mode 100644 index 000000000..ea030219b --- /dev/null +++ b/fastvideo/tests/ops/quantization/test_nvfp4_skip_blocks.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +"""NVFP4 DiT exports that keep some MiniMax-H3 main blocks in bf16 (FastH3 OmniRef QAD: blocks 42-49).""" +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import save_file + +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.quantization.nvfp4_config import ( + NVFP4Config, + NVFP4QuantizeMethod, + minimax_h3_block_index, + read_minimax_h3_nvfp4_dit_skip_blocks, +) +from fastvideo.models.loader.fsdp_load import _adopt_nvfp4_export_skip_blocks + + +def _linear() -> ReplicatedLinear: + return ReplicatedLinear(32, 32, bias=False) + + +def test_block_index_parses_main_blocks_only(): + assert minimax_h3_block_index("transformer_blocks.42.attn.to_q") == 42 + assert minimax_h3_block_index("minimax_h3.transformer_blocks.7.ff.fc_in") == 7 + assert minimax_h3_block_index("token_refiner.refiner_blocks.1.attn.to_q") is None + assert minimax_h3_block_index("proj_out") is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="NVFP4QuantizeMethod allocates its unit scale on CUDA") +def test_skip_blocks_leave_listed_blocks_unquantized(): + config = NVFP4Config(layer_profile="h3_dit_vsa", skip_blocks=[49, 42, 42]) + assert config.skip_blocks == (42, 49) + layer = _linear() + assert isinstance(config.get_quant_method(layer, "transformer_blocks.41.attn.to_gate_compress"), + NVFP4QuantizeMethod) + assert isinstance(config.get_quant_method(layer, "transformer_blocks.41.ff.fc_out"), NVFP4QuantizeMethod) + for prefix in ("transformer_blocks.42.attn.to_q", "transformer_blocks.49.ff.fc_in", + "transformer_blocks.42.attn.to_gate_compress"): + assert config.get_quant_method(layer, prefix) is None + + +def test_skip_blocks_round_trip_and_copy(): + config = NVFP4Config.from_config({"layer_profile": "h3_dit_vsa", "skip_blocks": [3, 1]}) + assert config.skip_blocks == (1, 3) + copy = config.with_skip_blocks((5, )) + assert copy.skip_blocks == (5, ) and config.skip_blocks == (1, 3) + assert copy.layer_profile == "h3_dit_vsa" + + +@pytest.mark.parametrize("bad", [(-1, ), (True, ), ("3", )]) +def test_skip_blocks_rejects_bad_indices(bad): + with pytest.raises(ValueError): + NVFP4Config(layer_profile="h3_dit", skip_blocks=bad) + + +def _export(tmp_path, metadata): + path = tmp_path / "nvfp4_weights.safetensors" + save_file({"transformer_blocks.0.ff.fc_in::_nvfp4_alpha": torch.ones(())}, str(path), metadata=metadata) + return str(path) + + +def test_export_metadata_skip_blocks(tmp_path): + assert read_minimax_h3_nvfp4_dit_skip_blocks(_export(tmp_path, None)) == () + assert read_minimax_h3_nvfp4_dit_skip_blocks(_export(tmp_path, {"skip_blocks": "[49, 42]"})) == (42, 49) + with pytest.raises(ValueError): + read_minimax_h3_nvfp4_dit_skip_blocks(_export(tmp_path, {"skip_blocks": "{\"a\": 1}"})) + + +def test_loader_adopts_export_skip_blocks_without_mutating_config(tmp_path): + path = _export(tmp_path, {"skip_blocks": "[42, 43]"}) + quant_config = NVFP4Config(layer_profile="h3_dit_vsa") + dit_config = SimpleNamespace(quant_config=quant_config, name="dit") + params = {"config": dit_config, "hf_config": {}} + adopted = _adopt_nvfp4_export_skip_blocks(params, quant_config, path) + assert adopted["config"].quant_config.skip_blocks == (42, 43) + assert dit_config.quant_config is quant_config and quant_config.skip_blocks == () + assert adopted["hf_config"] is params["hf_config"] + + +def test_loader_keeps_matching_and_rejects_conflicting_skip_blocks(tmp_path): + path = _export(tmp_path, {"skip_blocks": "[42]"}) + matching = NVFP4Config(layer_profile="h3_dit_vsa", skip_blocks=(42, )) + params = {"config": SimpleNamespace(quant_config=matching)} + assert _adopt_nvfp4_export_skip_blocks(params, matching, path) is params + conflicting = NVFP4Config(layer_profile="h3_dit_vsa", skip_blocks=(41, )) + with pytest.raises(ValueError, match="skip_blocks"): + _adopt_nvfp4_export_skip_blocks({"config": SimpleNamespace(quant_config=conflicting)}, conflicting, path) + + +def _converter(): + import importlib.util + from pathlib import Path + + script = Path(__file__).resolve().parents[4] / "scripts/checkpoint_conversion/convert_minimax_h3_modelopt_nvfp4_dit.py" + spec = importlib.util.spec_from_file_location("convert_minimax_h3_modelopt_nvfp4_dit", script) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_converter_parses_block_ranges(): + converter = _converter() + assert converter.parse_blocks("42-49") == tuple(range(42, 50)) + assert converter.parse_blocks("3, 7,42-43") == (3, 7, 42, 43) + assert converter.parse_blocks("3,3,7-7") == (3, 7) + assert converter.block_of("transformer_blocks.12.ff.net.2") == 12 + assert converter.block_of("proj_out") is None + + +@pytest.mark.parametrize("text", ["49-42", "3,7,49-42", "7-3,42-49"]) +def test_converter_rejects_inverted_block_ranges(text): + with pytest.raises(ValueError, match=r"end \(\d+\) is below start \(\d+\)"): + _converter().parse_blocks(text) + + +def test_converter_reads_full_name_amax_keys(): + converter = _converter() + table = {"transformer_blocks.0.ff.fc_in": 3.0, "b1.ff.fc_out": 5.0} + assert converter.act_amax_for("transformer_blocks.0.ff.fc_in", table, required=True) == 3.0 + assert converter.act_amax_for("transformer_blocks.1.ff.fc_out", table, required=True) == 5.0 + assert converter.act_amax_for("transformer_blocks.2.attn.to_q", table, required=False) is None + with pytest.raises(SystemExit): + converter.act_amax_for("transformer_blocks.2.ff.fc_in", table, required=True) diff --git a/fastvideo/tests/ops/quantization/test_nvfp4_spark_fence.py b/fastvideo/tests/ops/quantization/test_nvfp4_spark_fence.py index fda1aedc8..95ea73669 100644 --- a/fastvideo/tests/ops/quantization/test_nvfp4_spark_fence.py +++ b/fastvideo/tests/ops/quantization/test_nvfp4_spark_fence.py @@ -48,6 +48,8 @@ def _spy_sync(monkeypatch): return syncs +@pytest.mark.skipif(torch.cuda.is_available() and torch.cuda.get_device_capability() == (12, 1), + reason="needs a non-GB10 Blackwell GPU") def test_fenced_quantize_matches_flashinfer_without_a_fence_off_spark(monkeypatch) -> None: real, pdl = _spy_quantize(monkeypatch) syncs = _spy_sync(monkeypatch) @@ -94,3 +96,100 @@ def test_qat_linear_quantizes_through_the_fence_with_unchanged_numerics(monkeypa reference = torch.nn.functional.linear(x.float(), weight.float(), bias.float()) relative = (on_spark.float() - reference).norm() / reference.norm() assert relative < 0.15 + + +@pytest.mark.parametrize("rows", [200, 256]) +def test_swizzled_quantize_skips_the_pad_copy_bit_exactly(monkeypatch, rows: int) -> None: + torch.manual_seed(2) + x = torch.randn(rows, 256, device="cuda", dtype=torch.bfloat16) + global_sf = (448.0 * 6.0) / x.float().abs().amax() + expected = flashinfer.nvfp4_quantize(_pad_rows(x), global_sf, sfLayout=LAYOUT_128X4, do_shuffle=False) + # nv.F is torch.nn.functional, so the spy sees every pad; install it after the reference. + pads = [] + real_pad = torch.nn.functional.pad + monkeypatch.setattr(nv.F, "pad", lambda *args, **kwargs: pads.append(args) or real_pad(*args, **kwargs)) + quantized, scales = nv._nvfp4_quantize(x, global_sf, sfLayout=LAYOUT_128X4) + assert pads == [] + assert quantized.shape[0] == rows + assert torch.equal(quantized, expected[0][:rows]) and torch.equal(scales, expected[1]) + + +@pytest.mark.parametrize("rows", [200, 256]) +@pytest.mark.parametrize("layout", ["layout_128x4", "layout_8x4", "layout_linear"]) +@pytest.mark.parametrize("do_shuffle", [False, True]) +def test_fake_quantize_reports_flashinfer_shapes(rows: int, layout: str, do_shuffle: bool) -> None: + from torch._subclasses.fake_tensor import FakeTensorMode + + if do_shuffle and (layout != "layout_128x4" or rows % 128): + pytest.skip("FlashInfer only shuffles 128x4, and _nvfp4_quantize row-pads shuffled inputs") + sf_layout = getattr(flashinfer.SfLayout, layout).value + nv._register_ops_once() + torch.manual_seed(3) + x = torch.randn(rows, 256, device="cuda", dtype=torch.bfloat16) + global_sf = (448.0 * 6.0) / x.float().abs().amax() + real = torch.ops.fastvideo_fp4.nvfp4_quantize(x, global_sf, sf_layout, do_shuffle) + with FakeTensorMode() as mode: + fake = torch.ops.fastvideo_fp4.nvfp4_quantize(mode.from_tensor(x), mode.from_tensor(global_sf), sf_layout, + do_shuffle) + assert [tuple(t.shape) for t in fake] == [tuple(t.shape) for t in real] + + +def test_compiled_swizzled_quantize_quantizes_once_without_row_guards(monkeypatch) -> None: + from torch._dynamo.testing import CompileCounterWithBackend + + calls = [] + real_fenced = nv.nvfp4_quantize_fenced + monkeypatch.setattr(nv, "nvfp4_quantize_fenced", lambda x, *args: calls.append(x.shape[0]) or real_fenced(x, *args)) + counter = CompileCounterWithBackend("inductor") + torch._dynamo.reset() + compiled = torch.compile(lambda x, sf: nv._nvfp4_quantize(x, sf, sfLayout=LAYOUT_128X4), + backend=counter, + dynamic=True, + fullgraph=True) + for rows in (200, 256, 300): + torch.manual_seed(rows) + x = torch.randn(rows, 256, device="cuda", dtype=torch.bfloat16) + global_sf = (448.0 * 6.0) / x.float().abs().amax() + calls.clear() + quantized, scales = compiled(x, global_sf) + expected = flashinfer.nvfp4_quantize(_pad_rows(x), global_sf, sfLayout=LAYOUT_128X4, do_shuffle=False) + assert calls == [rows] + assert torch.equal(quantized, expected[0][:rows]) and torch.equal(scales, expected[1]) + assert counter.frame_count == 1 + torch._dynamo.reset() + + +# (255, 2064): K/16 = 129 pads to 132 scale columns, so 255 unpadded rows already +# hold more scales than 256 rows x 129 logical columns. +@pytest.mark.parametrize("rows,k", [(200, 256), (255, 2064)]) +def test_unpadded_flashinfer_scales_fall_back_to_the_pad_copy_once(monkeypatch, rows: int, k: int) -> None: + monkeypatch.setattr(nv, "_FLASHINFER_PADS_128X4_SCALES", True) + real_fenced = nv.nvfp4_quantize_fenced + calls = [] + + def unpadded(x, global_sf, sf_layout, do_shuffle=False): + calls.append(x.shape[0]) + quantized, scales = real_fenced(x, global_sf, sf_layout, do_shuffle) + return quantized, scales[:x.shape[0]] + + monkeypatch.setattr(nv, "nvfp4_quantize_fenced", unpadded) + torch.manual_seed(0) + x = torch.randn(rows, k, device="cuda", dtype=torch.bfloat16) + global_sf = (448.0 * 6.0) / x.float().abs().amax() + nv._nvfp4_quantize(x, global_sf, sfLayout=LAYOUT_128X4) + assert calls == [rows, 256] and nv._FLASHINFER_PADS_128X4_SCALES is False + calls.clear() + nv._nvfp4_quantize(x, global_sf, sfLayout=LAYOUT_128X4) + assert calls == [256] + + +def test_global_sf_does_not_sync() -> None: + x = torch.randn(64, 256, device="cuda", dtype=torch.bfloat16) + x[3, 5] = float("inf") + torch.cuda.synchronize() + torch.cuda.set_sync_debug_mode("error") + try: + scale = fp4linear._global_sf(x) + finally: + torch.cuda.set_sync_debug_mode("default") + assert scale.item() == pytest.approx((448.0 * 6.0) / torch.finfo(torch.float32).max) diff --git a/fastvideo/tests/stages/test_lazy_module_load.py b/fastvideo/tests/stages/test_lazy_module_load.py index 57bf8658e..b65cbcfde 100644 --- a/fastvideo/tests/stages/test_lazy_module_load.py +++ b/fastvideo/tests/stages/test_lazy_module_load.py @@ -12,6 +12,7 @@ from types import SimpleNamespace import pytest +from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase from fastvideo.pipelines.lazy_module import LazyModule, is_lazy_module from fastvideo.pipelines.stages.base import PipelineStage @@ -722,7 +723,7 @@ def test_real_load_modules_defers_nothing_when_the_pipeline_opts_out(monkeypatch assert sorted(loaded) == ["scheduler", "transformer", "vae"] -def test_building_the_real_h3_stages_materializes_nothing(): +def test_building_the_real_h3_stages_materializes_nothing(tmp_path): # `DenoisingStage.__init__` in the shared stage set reads # `transformer.hidden_size` to pick an attention backend, which would pull # the DiT in during post_init. H3's stages must not acquire that habit. @@ -748,7 +749,7 @@ def test_building_the_real_h3_stages_materializes_nothing(): "scheduler": object(), "audio_scheduler": object(), } - args = SimpleNamespace(pipeline_config=MiniMaxH3PipelineConfig()) + args = FastVideoArgs(model_path=str(tmp_path), pipeline_config=MiniMaxH3PipelineConfig(), lazy_module_load=True) pipeline._add_stages(args, ref2va=False) @@ -766,7 +767,7 @@ def test_building_the_real_h3_stages_materializes_nothing(): assert id(pipeline.modules["vae"]) in decode_held -def test_h3_lazy_release_drops_dit_before_vae_decode(): +def test_h3_lazy_release_drops_dit_before_vae_decode(tmp_path): from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import MiniMaxH3Pipeline @@ -783,7 +784,7 @@ def test_h3_lazy_release_drops_dit_before_vae_decode(): "scheduler": object(), "audio_scheduler": object(), } - args = SimpleNamespace(pipeline_config=MiniMaxH3PipelineConfig()) + args = FastVideoArgs(model_path=str(tmp_path), pipeline_config=MiniMaxH3PipelineConfig(), lazy_module_load=True) pipeline._add_stages(args, ref2va=False) schedule = pipeline._build_lazy_release_schedule() names = {pipeline._stages[index]._pipeline_stage_name: modules for index, modules in schedule.items()} diff --git a/fastvideo/tests/stages/test_minimax_h3_consumer_regressions.py b/fastvideo/tests/stages/test_minimax_h3_consumer_regressions.py new file mode 100644 index 000000000..d3ba2b94d --- /dev/null +++ b/fastvideo/tests/stages/test_minimax_h3_consumer_regressions.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Host-copy layout changes for frozen offload, without model weights.""" +from types import SimpleNamespace + +import pytest +import torch + +from fastvideo.models import pinned_offload +from fastvideo.pipelines.basic.minimax_h3.stages import minimax_h3_decoding as decoding +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import MINIMAX_H3_LAYOUT_KEY + + +def _replace_master_with_packed(module): + del module.weight + module.register_buffer("packed_weight", torch.tensor([12, 34], dtype=torch.uint8)) + + +def test_refresh_offload_host_copies_replaces_deleted_master(): + module = torch.nn.Linear(2, 2) + pinned_offload._host_copies(module, pin=False) + _replace_master_with_packed(module) + pinned_offload.refresh_host_copies(module, pin=False) + host = getattr(module, pinned_offload._HOST_ATTR) + assert set(host) == {"bias", "buffer:packed_weight"} + torch.testing.assert_close(host["buffer:packed_weight"], module.packed_weight) + + +def test_resident_module_does_not_gain_host_mirror(): + module = torch.nn.Linear(2, 2) + _replace_master_with_packed(module) + pinned_offload.refresh_host_copies(module, pin=False) + assert not hasattr(module, pinned_offload._HOST_ATTR) + + +def test_unload_runs_submodule_device_cache_hooks(): + """Plain-attribute device caches (fused plans, CUDA graphs) must be released with the weights.""" + module = torch.nn.Module() + child = torch.nn.Module() + module.add_module("child", child) + cleared = [] + child.__dict__["clear_offload_caches"] = lambda: cleared.append(True) + setattr(module, pinned_offload._HOST_ATTR, {}) + pinned_offload.unload(module) + assert cleared == [True] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="offload reload needs CUDA") +def test_packed_offload_reloads_twice_without_old_weight_names(): + module = torch.nn.Linear(2, 2).requires_grad_(False) + device = torch.device("cuda", torch.cuda.current_device()) + pinned_offload.load(module, device, pin=False) + _replace_master_with_packed(module) + module.packed_weight = module.packed_weight.to(device) + pinned_offload.refresh_host_copies(module, pin=False) + for _ in range(2): + pinned_offload.unload(module) + assert module.packed_weight.device.type == "cpu" + pinned_offload.load(module, device, pin=False) + assert module.packed_weight.device == device + assert module.packed_weight.tolist() == [12, 34] + + +def test_decode_refreshes_mirror_before_first_unload(monkeypatch): + import fastvideo.models.vaes.minimax_h3_nvfp4_decoder as nvfp4 + + vae = torch.nn.Linear(2, 2).requires_grad_(False) + vae.denormalize_latents = lambda z: z + pinned_offload._host_copies(vae, pin=False) + layout = SimpleNamespace(num_condition_video_rows=0, num_condition_audio_rows=0) + batch = SimpleNamespace(extra={MINIMAX_H3_LAYOUT_KEY: layout}, latents=torch.zeros(1, 2), + raw_latent_shape=(1, 2, 1, 1, 1), output=None) + args = SimpleNamespace(vae_parallel_decode=False, pin_cpu_memory=False, vae_cpu_offload=True, + output_type="latent") + monkeypatch.setattr(decoding, "_decode_participation", lambda *a: (None, True, False)) + monkeypatch.setattr(decoding, "_layout", lambda batch: layout) + monkeypatch.setattr(decoding, "get_local_torch_device", lambda: torch.device("cpu")) + monkeypatch.setattr(decoding, "h3_dit_patch_size", lambda *a: (1, 1, 1)) + monkeypatch.setattr(decoding, "unpatchify_video_tokens", lambda tokens, *a: tokens) + + def freeze(module): + _replace_master_with_packed(module) + return 1 + + def unload(module): + # The real CUDA unload indexes every current tensor name in this mirror. + host = getattr(module, pinned_offload._HOST_ATTR) + for name, tensor in pinned_offload._tensors(module): + torch.testing.assert_close(host[name], tensor) + + monkeypatch.setattr(nvfp4, "freeze_nvfp4_linears_on_device", freeze) + monkeypatch.setattr(pinned_offload, "unload", unload) + decoding.MiniMaxH3VideoDecodingStage(vae).forward(batch, args) diff --git a/fastvideo/tests/stages/test_minimax_h3_sequential_start.py b/fastvideo/tests/stages/test_minimax_h3_sequential_start.py index 6094be15a..f549a57c5 100644 --- a/fastvideo/tests/stages/test_minimax_h3_sequential_start.py +++ b/fastvideo/tests/stages/test_minimax_h3_sequential_start.py @@ -105,6 +105,7 @@ def test_inference_defers_dit_and_vae_until_after_conditioning(monkeypatch) -> N assert list(pipeline._stage_name_mapping) == ["input_preparation_stage", "conditioning_stage"] condition_stage = pipeline._stage_name_mapping["conditioning_stage"] + assert condition_stage.manage_cpu_offload is False passthrough = lambda batch, _args: batch monkeypatch.setattr(pipeline._stage_name_mapping["input_preparation_stage"], "forward", passthrough) monkeypatch.setattr(condition_stage, "forward", passthrough) @@ -148,6 +149,51 @@ def test_inference_defers_dit_and_vae_until_after_conditioning(monkeypatch) -> N assert encoder.moved_to[moves_before_second] == torch.device("cpu") +@pytest.mark.parametrize("manage_cpu_offload", [True, False]) +@pytest.mark.parametrize("fail_encoding", [True, False]) +def test_conditioning_respects_pipeline_offload_ownership(monkeypatch, manage_cpu_offload, fail_encoding): + import fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_conditioning as conditioning + + moves = [] + encoder = SimpleNamespace(parameters=lambda: iter([torch.ones(1)]), to=lambda device: moves.append(device)) + stage = conditioning.MiniMaxH3ConditioningStage(encoder, None, None, + manage_cpu_offload=manage_cpu_offload) + + def encode(*_args): + if fail_encoding: + raise RuntimeError("encoding failed") + return torch.ones(1, 2, 4), torch.zeros(2, dtype=torch.int32) + + monkeypatch.setattr(conditioning, "get_local_torch_device", lambda: torch.device("cpu")) + monkeypatch.setattr(stage, "_encode_fl2va", encode) + batch = ForwardBatch(data_type="video", prompt="placement") + args = SimpleNamespace(text_encoder_cpu_offload=True) + if fail_encoding: + with pytest.raises(RuntimeError, match="encoding failed"): + stage.forward(batch, args) + else: + assert stage.forward(batch, args).prompt_embeds[0].shape == (1, 2, 4) + assert moves == ([torch.device("cpu"), "cpu"] if manage_cpu_offload else []) + + +def test_sequential_releases_encoder_when_conditioning_fails(monkeypatch): + pipeline = MiniMaxH3Pipeline.__new__(MiniMaxH3Pipeline) + events = [] + + def fail(*_args): + raise RuntimeError("encoding failed") + + pipeline._stage_name_mapping = { + "input_preparation_stage": lambda batch, _args: batch, + "conditioning_stage": fail, + } + monkeypatch.setattr(pipeline, "_release_text_encoder", lambda: events.append("release")) + monkeypatch.setattr(pipeline, "_load_denoise_modules", lambda *_: pytest.fail("Do not load after encode failure")) + with pytest.raises(RuntimeError, match="encoding failed"): + pipeline._run_condition_then_denoise(ForwardBatch(data_type="video"), SimpleNamespace()) + assert events == ["release"] + + def test_sequential_skips_host_offload_for_dtensor_params(monkeypatch) -> None: class _FakeDTensor: pass diff --git a/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py b/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py index fb602bf60..4ff3e986d 100644 --- a/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py +++ b/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py @@ -3,6 +3,7 @@ from types import SimpleNamespace import numpy as np +import pytest import torch from fastvideo.pipelines.basic.minimax_h3 import encoder_split @@ -175,6 +176,66 @@ def test_decode_stage_requests_pin_fallback_output_buffer(monkeypatch) -> None: } +@pytest.mark.parametrize("unified_memory", [True, False]) +def test_decode_stage_reuses_its_pixel_buffer_only_on_unified_memory(monkeypatch, unified_memory) -> None: + """On unified memory, same-geometry requests reuse one pre-touched CPU buffer; returned + frames never alias it. Discrete GPUs allocate per request and hold nothing between them.""" + latent_shape = (1, 4, 2, 4, 4) + rows = patchify_video_latents(torch.randn(latent_shape), (1, 1, 1)) + allocations = [] + buffers = [] + + class VAE: + + fill = 0.25 + + def to(self, device): + return self + + def denormalize_latents(self, decoded_latents): + return decoded_latents + + def decoded_pixel_shape(self, shape): + return (1, 3, 5, 16, 16) + + def decode_to_pixels(self, decoded_latents, output): + buffers.append((output, bool(torch.isnan(output).any()))) + output.fill_(VAE.fill) + + def fake_allocate(size, *, dtype=None, pin_memory=False): + allocations.append(size) + return torch.full(size, float("nan"), dtype=dtype) + + monkeypatch.setattr(minimax_h3_decoding, "get_local_torch_device", lambda: torch.device("cpu")) + monkeypatch.setattr(minimax_h3_decoding, "allocate_cpu_tensor_with_pin_fallback", fake_allocate) + monkeypatch.setattr(minimax_h3_decoding, "current_platform", + SimpleNamespace(has_unified_memory=lambda device_id: unified_memory)) + stage = MiniMaxH3VideoDecodingStage(VAE()) + args = SimpleNamespace(output_type="pil", + pin_cpu_memory=False, + vae_cpu_offload=False, + vae_parallel_decode=False, + pipeline_config=SimpleNamespace(dit_config=SimpleNamespace(patch_size=(1, 1, 1)))) + + def decode(): + batch = ForwardBatch(data_type="video", latents=rows.clone(), raw_latent_shape=latent_shape) + batch.extra[MINIMAX_H3_LAYOUT_KEY] = _layout(rows.shape[0], latent_shape) + return stage.forward(batch, args).output + + first = decode() + VAE.fill = 0.5 + second = decode() + if unified_memory: + assert allocations == [(1, 3, 5, 16, 16)] + assert buffers[0][0] is buffers[1][0] + assert not buffers[0][1] # zeroed (pages touched) before the decoder writes it + else: + assert allocations == [(1, 3, 5, 16, 16)] * 2 + assert buffers[0][0] is not buffers[1][0] + assert stage._output_buffer is None + assert torch.equal(first, torch.full_like(first, 63)) and torch.equal(second, torch.full_like(second, 127)) + + def test_decode_stages_skip_vae_on_non_output_rank(monkeypatch) -> None: class VAE: diff --git a/fastvideo/tests/vaes/test_minimax_h3_int8_convrot.py b/fastvideo/tests/vaes/test_minimax_h3_int8_convrot.py index 165ccc41e..99ecc2935 100644 --- a/fastvideo/tests/vaes/test_minimax_h3_int8_convrot.py +++ b/fastvideo/tests/vaes/test_minimax_h3_int8_convrot.py @@ -5,13 +5,17 @@ from __future__ import annotations import json import torch +import pytest import torch.nn as nn from safetensors.torch import save_file +import fastvideo.envs as envs from fastvideo.models.vaes.minimax_h3_int8_convrot import ( + INT8_CONVROT_FILENAME, Int8ConvRotLinear, _int8_linear_from_tensors, dense_vae_safetensors, + int8_convrot_overlay_to_apply, overlay_minimax_h3_int8_convrot_decoder, parse_comfy_quant_marker, regular_hadamard, @@ -82,6 +86,20 @@ def test_dense_vae_safetensors_drops_convrot_overlay() -> None: ] +@pytest.mark.parametrize(("enabled", "ships_overlay", "applied"), [ + (True, True, True), + (False, True, False), + (True, False, False), +]) +def test_int8_convrot_overlay_follows_the_overlay_switch(tmp_path, env_overrides, enabled: bool, ships_overlay: bool, + applied: bool) -> None: + if ships_overlay: + (tmp_path / INT8_CONVROT_FILENAME).write_bytes(b"") + env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_INT8_OVERLAY.override(enabled)) + expected = tmp_path / INT8_CONVROT_FILENAME if applied else None + assert int8_convrot_overlay_to_apply(tmp_path) == expected + + def test_parse_comfy_quant_marker_reads_padded_uint8() -> None: payload = json.dumps({"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": 256}).encode() blob = torch.zeros(72, dtype=torch.uint8) diff --git a/fastvideo/tests/vaes/test_minimax_h3_int8_shared.py b/fastvideo/tests/vaes/test_minimax_h3_int8_shared.py index 87caab3ee..5528e1e23 100644 --- a/fastvideo/tests/vaes/test_minimax_h3_int8_shared.py +++ b/fastvideo/tests/vaes/test_minimax_h3_int8_shared.py @@ -93,3 +93,25 @@ def test_fused_int8_epilogue_large_accumulators_and_small_scales(dtype, has_bias expected = expected + bias.float() actual = fused_int8_dequant_bias(acc, x_scale, w_scale, bias, dtype) torch.testing.assert_close(actual, expected.to(dtype), rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for fused INT8 epilogue") +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize(("acc_value", "x_scale_value", "weight_scale_value", "bias_value"), [ + (41104, 0.026947021484375, 0.007559554651379585, 0.3650893568992615), + (69320, 0.0205841064453125, 0.007878062315285206, -0.43253710865974426), +]) +def test_fused_int8_epilogue_preserves_fp16_rounding_boundaries( + dtype, acc_value, x_scale_value, weight_scale_value, bias_value, +): + from fastvideo.models.vaes.minimax_h3_int8_kernels import fused_int8_dequant_bias + + # Captured CI failures: contracting the final multiply and bias add changes + # which side of an FP16 midpoint the FP32 result lands on. + acc = torch.full((129, 128), acc_value, device="cuda", dtype=torch.int32) + x_scale = torch.full((129, 1), x_scale_value, device="cuda", dtype=torch.float16) + weight_scale = torch.full((128, 1), weight_scale_value, device="cuda", dtype=torch.float32) + bias = torch.full((128,), bias_value, device="cuda", dtype=torch.float32) + expected = acc.float() * x_scale.float() * weight_scale.t().float() + bias + actual = fused_int8_dequant_bias(acc, x_scale, weight_scale, bias, dtype) + torch.testing.assert_close(actual, expected.to(dtype), rtol=0, atol=0) diff --git a/fastvideo/tests/vaes/test_minimax_h3_nvfp4_decoder.py b/fastvideo/tests/vaes/test_minimax_h3_nvfp4_decoder.py new file mode 100644 index 000000000..bcfa97e2e --- /dev/null +++ b/fastvideo/tests/vaes/test_minimax_h3_nvfp4_decoder.py @@ -0,0 +1,783 @@ +# SPDX-License-Identifier: Apache-2.0 +"""NVFP4 MiniMax-H3 decoder linears: conversion contract, rotation exactness, train/deploy parity.""" +import pytest +import torch +import torch.nn as nn + +from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import ( + NVFP4DecoderLinear, + convert_decoder_to_nvfp4, + nvfp4_decoder_linear_names, +) +from fastvideo.models.vaes.minimax_h3_video import MiniMaxH3VideoViTDecoder3d +from fastvideo.models.vaes.minimax_h3_int8_convrot import rotate_activation + + +def _tiny_decoder(num_layers: int = 3) -> MiniMaxH3VideoViTDecoder3d: + torch.manual_seed(0) + return MiniMaxH3VideoViTDecoder3d(in_channels=24, + out_channels=3, + patch_size=2, + patch_size_t=1, + num_layers=num_layers, + num_attention_heads=4, + attention_head_dim=64, + num_register_tokens=4, + ffn_mult=4, + rope_theta=100.0, + rope_dim_ratio=0.75, + norm_eps=1e-5) + + +def test_selects_only_block_linears(): + decoder = _tiny_decoder() + names = nvfp4_decoder_linear_names(decoder) + assert len(names) == 3 * 6 + assert all(name.startswith("transformer_blocks.") for name in names) + assert "proj_in" not in names and "proj_out" not in names + + +def test_skip_blocks_keeps_them_dense(): + decoder = _tiny_decoder() + names = convert_decoder_to_nvfp4(decoder, skip_blocks=(0, 2)) + assert {name.split(".")[1] for name in names} == {"1"} + assert type(decoder.transformer_blocks[0].attn.to_q) is nn.Linear + assert isinstance(decoder.transformer_blocks[1].attn.to_q, NVFP4DecoderLinear) + assert isinstance(decoder.transformer_blocks[1].ff.net[2], NVFP4DecoderLinear) + assert isinstance(decoder.transformer_blocks[1].attn.to_out[0], NVFP4DecoderLinear) + + +@pytest.mark.parametrize("group", [16, 64, 256]) +def test_rotation_is_exact_in_full_precision(group): + torch.manual_seed(0) + linear = nn.Linear(256, 96) + rotated = NVFP4DecoderLinear.from_linear(linear, rotation_group=group, compute_dtype=torch.float32) + x = torch.randn(5, 7, 256, dtype=torch.float64) + expected = x @ linear.weight.double().T + linear.bias.double() + actual = rotate_activation(x, group) @ rotated.weight.double().T + rotated.bias.double() + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) + + +def test_rejects_non_blockable_input(): + with pytest.raises(ValueError, match="divisible"): + NVFP4DecoderLinear.from_linear(nn.Linear(24, 32), rotation_group=None, compute_dtype=torch.bfloat16) + + +def _fp4_available() -> bool: + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10: + return False + try: + import flashinfer # noqa: F401 + except ImportError: + return False + return True + + +def test_rejects_unknown_activation_scale(): + with pytest.raises(ValueError, match="act_scale"): + NVFP4DecoderLinear.from_linear(nn.Linear(32, 32), rotation_group=None, compute_dtype=torch.bfloat16, + act_scale="per-token") + + +@pytest.mark.skipif(not _fp4_available(), reason="NVFP4 GEMM needs a Blackwell GPU and flashinfer") +@pytest.mark.parametrize("act_scale", ["dynamic", "unit", "static"]) +@pytest.mark.parametrize("group", [None, 256]) +def test_inference_matches_training_forward_bitwise(group, act_scale): + from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import calibrate_static_scales + + torch.manual_seed(0) + linear = nn.Linear(2048, 768).cuda() + layer = NVFP4DecoderLinear.from_linear(linear, rotation_group=group, compute_dtype=torch.bfloat16, + act_scale=act_scale).cuda() + x = torch.randn(3, 300, 2048, device="cuda", dtype=torch.bfloat16) + if act_scale == "static": + assert calibrate_static_scales(layer, lambda: layer(x)) == 1 + assert layer.input_amax.item() > 0 + with torch.enable_grad(): + train_out = layer(x) + with torch.no_grad(): + deploy_out = layer(x) + assert torch.equal(train_out.detach(), deploy_out) + relative = (deploy_out.float() - linear(x.float())).norm() / linear(x.float()).norm() + assert relative < 0.15 + + +@pytest.mark.skipif(not _fp4_available(), reason="NVFP4 GEMM needs a Blackwell GPU and flashinfer") +def test_training_forward_has_gradients(): + torch.manual_seed(0) + layer = NVFP4DecoderLinear.from_linear(nn.Linear(256, 128), rotation_group=64, + compute_dtype=torch.bfloat16).cuda() + x = torch.randn(64, 256, device="cuda", dtype=torch.bfloat16, requires_grad=True) + layer(x).float().square().mean().backward() + assert layer.weight.grad is not None and torch.isfinite(layer.weight.grad).all() + assert x.grad is not None and torch.isfinite(x.grad).all() + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +def test_global_scale_matches_reference_reduction(dtype): + from fastvideo.layers.fp4linear import _global_sf + + torch.manual_seed(0) + x = (torch.randn(1797, 256) * 7).to(dtype) + expected = (448.0 * 6.0) / x.float().abs().nan_to_num().max().clamp(min=1e-12) + assert torch.equal(_global_sf(x), expected) + with_inf = x.clone() + with_inf[3, 5] = float("inf") + expected_inf = (448.0 * 6.0) / with_inf.float().abs().nan_to_num().max().clamp(min=1e-12) + assert torch.equal(_global_sf(with_inf), expected_inf) + + +# Fused inference kernels: every one must match the eager op sequence bit for bit. +_needs_fp4 = pytest.mark.skipif(not _fp4_available(), reason="NVFP4 GEMM needs a Blackwell GPU and flashinfer") + + +def _flashinfer_quantize(x, global_sf): + import flashinfer + + from fastvideo.layers.quantization.nvfp4_config import nvfp4_quantize_fenced + + return nvfp4_quantize_fenced(x, global_sf, flashinfer.SfLayout.layout_128x4.value) + + +def _assert_quantized_equal(actual, expected): + assert actual[0].shape == expected[0].shape and actual[1].shape == expected[1].shape + assert torch.equal(actual[0].view(torch.uint8), expected[0].view(torch.uint8)) + assert torch.equal(actual[1].view(torch.uint8), expected[1].view(torch.uint8)) + + +@_needs_fp4 +@pytest.mark.parametrize("global_sf", [1.0, 2688.0 / 40.0]) +@pytest.mark.parametrize("magnitude", [3.0, 3000.0]) +def test_fused_quantize_matches_flashinfer(global_sf, magnitude): + from fastvideo.models.vaes.minimax_h3_nvfp4_fused import quantize + + torch.manual_seed(0) + x = (torch.randn(1797 * 2 + 5, 2048, device="cuda") * magnitude).bfloat16() + x[7, :16] = 0 # an all-zero block + sf = torch.tensor(global_sf, device="cuda") + _assert_quantized_equal(quantize(x, sf), _flashinfer_quantize(x, sf)) + + +@_needs_fp4 +@pytest.mark.parametrize("residual", [False, True]) +@pytest.mark.parametrize("norm", [False, True]) +def test_fused_residual_norm_quantize_matches_eager(residual, norm): + import torch.nn.functional as F + + from fastvideo.models.vaes.minimax_h3_nvfp4_fused import residual_norm_quantize + + if not (residual or norm): + pytest.skip("nothing to fuse") + torch.manual_seed(0) + rows, dim, eps = 1797 * 2 + 5, 2048, 1e-5 + h = torch.randn(rows, dim, device="cuda") * 5 + o = (torch.randn(rows, dim, device="cuda") * 2).bfloat16() + bias = torch.randn(dim, device="cuda") * 0.1 + scale = torch.randn(dim, device="cuda") * 0.3 + gamma = torch.rand(dim, device="cuda") + 0.5 + sf = torch.ones((), device="cuda") + expected_h = h + if residual: + biased = o.clone() + biased.add_(bias.to(torch.bfloat16)) + expected_h = h + biased * scale + actual_h = h.clone() + quantized = residual_norm_quantize(actual_h, (o, bias.to(torch.bfloat16), scale) if residual else None, + (gamma, eps, sf) if norm else None) + assert torch.equal(actual_h, expected_h) + if norm: + normed = F.rms_norm(expected_h, (dim, ), gamma, eps).to(torch.bfloat16) + _assert_quantized_equal(quantized, _flashinfer_quantize(normed, sf)) + + +@_needs_fp4 +def test_fused_swiglu_quantize_matches_eager(): + import torch.nn.functional as F + + from fastvideo.models.vaes.minimax_h3_nvfp4_fused import swiglu_quantize + + torch.manual_seed(0) + rows, ffn = 1797 + 5, 8192 + o = (torch.randn(rows, 2 * ffn, device="cuda") * 2).bfloat16() + bias = torch.randn(2 * ffn, device="cuda").bfloat16() + sf = torch.ones((), device="cuda") + biased = o.clone() + biased.add_(bias) + value, gate = biased.chunk(2, dim=-1) + expected = _flashinfer_quantize((value * F.silu(gate)).contiguous(), sf) + _assert_quantized_equal(swiglu_quantize(o, bias, sf), expected) + + +@_needs_fp4 +def test_fused_qkv_epilogue_matches_eager_attention_prologue(): + import torch.nn.functional as F + + from fastvideo.models.vaes.minimax_h3_nvfp4_fused import qkv_epilogue + + torch.manual_seed(0) + batch, seq, heads, head_dim, rotary, eps = 2, 1797, 32, 64, 48, 1e-5 + rows = batch * seq + projections = [(torch.randn(rows, heads * head_dim, device="cuda") * 2).bfloat16() for _ in range(3)] + biases = tuple((torch.randn(heads * head_dim, device="cuda") * 0.2).bfloat16() for _ in range(3)) + cos = torch.randn(batch, seq, 1, rotary, device="cuda").cos().to(torch.bfloat16) + sin = torch.randn(batch, seq, 1, rotary, device="cuda").sin().to(torch.bfloat16) + expected = [] + for projection, bias in zip(projections, biases, strict=True): + biased = projection.clone() + biased.add_(bias) + expected.append(biased.view(batch, seq, heads, head_dim)) + for index in range(2): # the attention module's q/k RMSNorm and partial RoPE, op by op + x = F.rms_norm(expected[index].float(), (head_dim, ), None, eps).to(torch.bfloat16) + x_rotary, x_pass = x[..., :rotary], x[..., rotary:] + first, second = x_rotary.chunk(2, dim=-1) + rotated = torch.cat([-second, first], dim=-1) + expected[index] = torch.cat([x_rotary * cos + rotated * sin, x_pass], dim=-1) + actual = qkv_epilogue(*projections, biases, cos.reshape(rows, rotary), sin.reshape(rows, rotary), heads, eps) + for got, want in zip(actual, expected, strict=True): + assert torch.equal(got.view(batch, seq, heads, head_dim), want) + + +def _fused_test_decoder(act_scale: str, num_layers: int = 2, convert: bool = True) -> MiniMaxH3VideoViTDecoder3d: + torch.manual_seed(0) + decoder = MiniMaxH3VideoViTDecoder3d(in_channels=24, + out_channels=3, + patch_size=2, + patch_size_t=1, + num_layers=num_layers, + num_attention_heads=8, + attention_head_dim=64, + num_register_tokens=4, + ffn_mult=4, + rope_theta=100.0, + rope_dim_ratio=0.75, + norm_eps=1e-5).cuda() + with torch.no_grad(): + for block in decoder.transformer_blocks: + block.scale1.normal_(0, 0.5) + block.scale2.normal_(0, 0.5) + decoder.register_tokens.normal_() + if convert: + convert_decoder_to_nvfp4(decoder, compute_dtype=torch.bfloat16, act_scale=act_scale) + return decoder.requires_grad_(False).eval() + + +@_needs_fp4 +@pytest.mark.parametrize("act_scale", ["unit", "static"]) +@pytest.mark.parametrize("autocast", [True, False]) +def test_fused_decoder_matches_eager_decoder_bitwise(act_scale, autocast): + from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import calibrate_static_scales + + decoder = _fused_test_decoder(act_scale) + latents = torch.randn(3, 24, 3, 8, 8, device="cuda") + fused = decoder.fused_blocks_forward + assert fused is not None + + def run(): + context = torch.autocast("cuda", dtype=torch.bfloat16) if autocast else torch.autocast("cuda", enabled=False) + with torch.no_grad(), context: + return decoder(latents) + + if act_scale == "static": + decoder.fused_blocks_forward = None + calibrate_static_scales(decoder, run) + decoder.fused_blocks_forward = fused + fused_out = run() + assert all(block._nvfp4_fused_plan[1] is not None for block in decoder.transformer_blocks) + decoder.fused_blocks_forward = None + eager_out = run() + assert torch.equal(fused_out, eager_out) + + +@_needs_fp4 +def test_fused_decoder_falls_back_when_inapplicable(): + from fastvideo.models.vaes.minimax_h3_nvfp4_fused import fused_nvfp4_blocks_forward + + hidden = torch.randn(1, 64, 512, device="cuda") + rotary = (torch.ones(1, 64, 1, 48, device="cuda"), torch.zeros(1, 64, 1, 48, device="cuda")) + dynamic = _fused_test_decoder("dynamic", num_layers=1) + with torch.no_grad(): + assert fused_nvfp4_blocks_forward(dynamic.transformer_blocks, hidden, rotary) is None + unit = _fused_test_decoder("unit", num_layers=1) + with torch.no_grad(): + assert fused_nvfp4_blocks_forward(unit.transformer_blocks, hidden, rotary) is not None + assert fused_nvfp4_blocks_forward(unit.transformer_blocks, hidden.bfloat16(), rotary) is None + with torch.enable_grad(): + assert fused_nvfp4_blocks_forward(unit.transformer_blocks, hidden, rotary) is None + + +@_needs_fp4 +@pytest.mark.parametrize("rows,out_features,in_features", [(1797, 2048, 2048), (3 * 1797, 4096, 2048), + (1797, 2048, 8192)]) +def test_cutlass_fp4_tactics_are_bit_identical(rows, out_features, in_features): + import flashinfer + + from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import NVFP4FusedState, _quantized_weight + from fastvideo.models.vaes.minimax_h3_nvfp4_fused import (_gemm, _gemm_workspace, _sm100_cutlass_fp4_module, + quantize) + + torch.manual_seed(0) + sf = torch.ones((), device="cuda") + quantized = quantize(torch.randn(rows, in_features, device="cuda").bfloat16(), sf) + packed, inv_scale, global_sf_w = _quantized_weight( + (torch.randn(out_features, in_features, device="cuda") * 0.02).bfloat16()) + alpha = 1.0 / (sf * global_sf_w) + expected = torch.empty(rows, out_features, device="cuda", dtype=torch.bfloat16) + flashinfer.mm_fp4(quantized[0], packed.T, quantized[1], inv_scale.T, alpha, torch.bfloat16, expected, + block_size=16, use_8x4_sf_layout=False, backend="cutlass") + state = NVFP4FusedState(packed, inv_scale, sf, alpha, torch.zeros(out_features, device="cuda"), out_features) + assert torch.equal(_gemm(quantized, state), expected) + module = _sm100_cutlass_fp4_module(torch.cuda.current_device()) + if module is None: + pytest.skip("the direct CUTLASS path is sm_100 only") + for tactic in range(module.fp4_gemm_tactic_num()): + out = torch.empty_like(expected) + module.fp4_gemm(quantized[0], packed, quantized[1], inv_scale, alpha, out, _gemm_workspace(out.device.index), + tactic) + assert torch.equal(out, expected), f"tactic {tactic}" + + +@_needs_fp4 +def test_fused_decoder_cuda_graph_replay_matches_eager(): + decoder = _fused_test_decoder("unit") + fused = decoder.fused_blocks_forward + inputs = [torch.randn(2, 24, 3, 8, 8, device="cuda") for _ in range(4)] + + def run(latents): + with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): + return decoder(latents) + + # Call 1 warms up eagerly, call 2 captures the graph, calls 3-4 replay it on new inputs. + fused_outs = [run(latents) for latents in inputs] + assert len(decoder.transformer_blocks.__dict__["_nvfp4_cuda_graphs"]) == 2 + decoder.fused_blocks_forward = None + for latents, fused_out in zip(inputs, fused_outs, strict=True): + assert torch.equal(fused_out, run(latents)) + decoder.fused_blocks_forward = fused + + +def test_dense_bypass_runs_the_master_weight(): + torch.manual_seed(0) + linear = nn.Linear(64, 32) + layer = NVFP4DecoderLinear.from_linear(linear, rotation_group=None, compute_dtype=torch.float32, act_scale="unit") + layer.dense_bypass = True + x = torch.randn(5, 64) + torch.testing.assert_close(layer(x), linear(x)) + assert layer.fused_state() is None + + +# Inference-only freeze: the packed weight replaces the master, outputs stay bit-identical. +@_needs_fp4 +@pytest.mark.parametrize("fused", [True, False]) +def test_frozen_decoder_matches_unfrozen_bitwise(fused): + from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import freeze_nvfp4_linears, nvfp4_linears + + decoder = _fused_test_decoder("unit") + if not fused: + decoder.fused_blocks_forward = None + latents = torch.randn(2, 24, 3, 8, 8, device="cuda") + + def run(): + with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): + return decoder(latents) + + before = run() + master_bytes = sum(layer.weight.numel() * layer.weight.element_size() for layer in nvfp4_linears(decoder)) + assert freeze_nvfp4_linears(decoder) == 2 * 6 + layers = nvfp4_linears(decoder) + assert all(layer.frozen and layer.weight is None for layer in layers) + names = {name for name, module in decoder.named_modules() if isinstance(module, NVFP4DecoderLinear)} + assert not names & {key.rpartition(".")[0] for key in decoder.state_dict() if key.endswith(".weight")} + packed_bytes = sum(buf.numel() * buf.element_size() for layer in layers for buf in layer.buffers()) + assert packed_bytes < master_bytes / 6 + assert torch.equal(run(), before) + if fused: + assert all(block._nvfp4_fused_plan[1] is not None for block in decoder.transformer_blocks) + with pytest.raises(RuntimeError, match="frozen"): + layers[0].invalidate() + + +@_needs_fp4 +def test_frozen_linear_survives_module_to(): + torch.manual_seed(0) + layer = NVFP4DecoderLinear.from_linear(nn.Linear(256, 128), rotation_group=None, compute_dtype=torch.bfloat16, + act_scale="unit").cuda().requires_grad_(False) + x = torch.randn(64, 256, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + expected = layer(x) + state = layer.fused_state() + layer.freeze() + assert torch.equal(layer(x), expected) + assert torch.equal(layer.fused_state().packed, state.packed) + layer.cpu().cuda() + assert torch.equal(layer(x), expected) + with torch.enable_grad(): + assert torch.equal(layer(x), expected) + + +# Loading an exported decoder (export_deploy.py format) into a dense H3 VAE. +_TINY_HEAD_DIM = 16 + + +def _tiny_vae_arch(decoder_layers: int, heads: int = 2) -> dict: + return { + "_class_name": "AutoencoderKLMiniMaxH3", + "latent_channels": 4, + "block_out_channels": [32, 32], + "layers_per_block": 1, + "spatial_downsample_factors": [2, 2], + "temporal_downsample_factors": [2, 2], + "decoder_num_layers": decoder_layers, + "decoder_num_attention_heads": heads, + "decoder_attention_head_dim": _TINY_HEAD_DIM, + "decoder_num_register_tokens": 2, + "decoder_ffn_mult": 2, + "latents_mean": [0.0] * 4, + "latents_std": [1.0] * 4, + } + + +def _tiny_h3_vae(arch: dict): + from fastvideo.configs.models.vaes.minimax_h3_video import MiniMaxH3VideoVAEArchConfig, MiniMaxH3VideoVAEConfig + from fastvideo.models.vaes.minimax_h3_video import AutoencoderKLMiniMaxH3 + + fields = {key: tuple(value) if isinstance(value, list) else value for key, value in arch.items() if key[0] != "_"} + return AutoencoderKLMiniMaxH3(MiniMaxH3VideoVAEConfig(arch_config=MiniMaxH3VideoVAEArchConfig(**fields))).eval() + + +def _write_vae_dir(path, decoder_layers: int): + from safetensors.torch import save_file + + torch.manual_seed(1) + path.mkdir() + arch = _tiny_vae_arch(decoder_layers) + (path / "config.json").write_text(__import__("json").dumps(arch)) + save_file({key: value.contiguous() for key, value in _tiny_h3_vae(arch).state_dict().items()}, + str(path / "diffusion_pytorch_model.safetensors")) + + +def _write_deploy_checkpoint(path, student_layers: int, base_layers: int, heads: int = 2, **metadata): + from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import (NVFP4_DECODER_DEPLOY_FORMAT, + keep_evenly_spaced_blocks) + + torch.manual_seed(2) + vae = _tiny_h3_vae(_tiny_vae_arch(base_layers, heads)) + keep_evenly_spaced_blocks(vae.decoder, student_layers) + names = convert_decoder_to_nvfp4(vae.decoder, compute_dtype=torch.bfloat16, act_scale="unit") + with torch.no_grad(): + for parameter in vae.decoder.parameters(): + parameter.normal_() + decoder = {} + for key, value in vae.decoder.state_dict().items(): + module, _, param = key.rpartition(".") + decoder[key] = value.bfloat16() if module in names and param in ("weight", "bias") else value + torch.save( + { + "format": NVFP4_DECODER_DEPLOY_FORMAT, + "decoder": decoder, + "post_quant_conv": vae.post_quant_conv.state_dict(), + "metadata": { + "act_scale": "unit", + "rotation_group": None, + "skip_blocks": [], + "num_linears": len(names), + "student_layers": student_layers, + **metadata, + }, + }, path) + return decoder + + +def _load_h3_vae(vae_dir, cpu: bool = True): + from fastvideo.configs.models.vaes.minimax_h3_video import MiniMaxH3VideoVAEConfig + from fastvideo.configs.pipelines import PipelineConfig + from fastvideo.fastvideo_args import FastVideoArgs + from fastvideo.models.loader.component_loader import VAELoader + + args = FastVideoArgs(model_path=str(vae_dir), + pipeline_config=PipelineConfig(vae_config=MiniMaxH3VideoVAEConfig(), vae_precision="fp32")) + args.vae_cpu_offload = cpu # load on CPU + return VAELoader().load(str(vae_dir), args) + + +def test_loader_applies_exported_nvfp4_decoder(tmp_path): + import fastvideo.envs as envs + + _write_vae_dir(tmp_path / "vae", decoder_layers=5) + expected = _write_deploy_checkpoint(tmp_path / "d3.pt", student_layers=3, base_layers=5) + dense = _load_h3_vae(tmp_path / "vae") + assert len(dense.decoder.transformer_blocks) == 5 and dense.decode_autocast_dtype == torch.float16 + with envs.FASTVIDEO_H3_VAE_NVFP4_DECODER.override(str(tmp_path / "d3.pt")): + vae = _load_h3_vae(tmp_path / "vae") + assert len(vae.decoder.transformer_blocks) == 3 + assert vae.decode_autocast_dtype == torch.bfloat16 + layer = vae.decoder.transformer_blocks[1].ff.net[2] + assert isinstance(layer, NVFP4DecoderLinear) and layer.act_scale == "unit" and not layer.frozen # CPU + state = vae.decoder.state_dict() + for key, value in expected.items(): + assert torch.equal(state[key].to(value.dtype).cpu(), value), key + assert not any(parameter.requires_grad for parameter in vae.parameters()) + + +@pytest.mark.parametrize("student_layers,base_layers,heads,match", [(6, 6, 2, "needs 6 decoder blocks"), + (2, 2, 4, "does not fit")]) +def test_loader_refuses_incompatible_nvfp4_decoder(tmp_path, student_layers, base_layers, heads, match): + import fastvideo.envs as envs + + _write_vae_dir(tmp_path / "vae", decoder_layers=4) + _write_deploy_checkpoint(tmp_path / "bad.pt", student_layers, base_layers, heads) + with envs.FASTVIDEO_H3_VAE_NVFP4_DECODER.override(str(tmp_path / "bad.pt")), pytest.raises(ValueError, + match=match): + _load_h3_vae(tmp_path / "vae") + + +def test_loader_rejects_non_deploy_checkpoint(tmp_path): + import fastvideo.envs as envs + + _write_vae_dir(tmp_path / "vae", decoder_layers=2) + torch.save({"decoder": {}, "step": 3}, tmp_path / "train.pt") + with envs.FASTVIDEO_H3_VAE_NVFP4_DECODER.override(str(tmp_path / "train.pt")), pytest.raises(ValueError, + match="format"): + _load_h3_vae(tmp_path / "vae") + with envs.FASTVIDEO_H3_VAE_NVFP4_DECODER.override(str(tmp_path / "missing.pt")), pytest.raises( + FileNotFoundError): + _load_h3_vae(tmp_path / "vae") + + +def test_keep_evenly_spaced_blocks_matches_the_f8_cut(): + from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import keep_evenly_spaced_blocks + + decoder = _tiny_decoder(num_layers=26) + blocks = list(decoder.transformer_blocks) + assert keep_evenly_spaced_blocks(decoder, 8) == [0, 4, 7, 11, 14, 18, 21, 25] + assert list(decoder.transformer_blocks) == [blocks[i] for i in (0, 4, 7, 11, 14, 18, 21, 25)] + + +@_needs_fp4 +@pytest.mark.parametrize("rows,out_features,in_features", [(1797, 2048, 2048), (12 * 1797, 16384, 2048), + (12 * 1797, 2048, 8192)]) +def test_decoder_gemm_backend_matches_cutlass(rows, out_features, in_features): + import flashinfer + + from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import _quantized_weight, mm_fp4_backend + from fastvideo.models.vaes.minimax_h3_nvfp4_fused import quantize + + backend = mm_fp4_backend(torch.device("cuda")) + if backend == "cutlass": + pytest.skip("this GPU decodes with the CUTLASS backend") + torch.manual_seed(0) + sf = torch.ones((), device="cuda") + x_fp4, x_scale = quantize(torch.randn(rows, in_features, device="cuda").bfloat16(), sf) + packed, inv_scale, global_sf_w = _quantized_weight( + (torch.randn(out_features, in_features, device="cuda") * 0.02).bfloat16()) + outs = [] + for name in ("cutlass", backend): + out = torch.empty(rows, out_features, device="cuda", dtype=torch.bfloat16) + flashinfer.mm_fp4(x_fp4, packed.T, x_scale, inv_scale.T, 1.0 / (sf * global_sf_w), torch.bfloat16, out, + block_size=16, use_8x4_sf_layout=False, backend=name) + outs.append(out) + assert torch.equal(outs[0], outs[1]) + + +# Safetensors exports: bf16 masters read directly, and pre-packed NVFP4 weights. +def _deploy_to_safetensors(pt_path, out_path): + from fastvideo.models.vaes.minimax_h3_nvfp4_checkpoint import write_nvfp4_decoder_safetensors + from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import load_nvfp4_decoder_checkpoint + + write_nvfp4_decoder_safetensors(str(out_path), load_nvfp4_decoder_checkpoint(str(pt_path))) + + +def test_loader_reads_bf16_safetensors_like_the_pt(tmp_path): + import fastvideo.envs as envs + + _write_vae_dir(tmp_path / "vae", decoder_layers=5) + _write_deploy_checkpoint(tmp_path / "d3.pt", student_layers=3, base_layers=5) + _deploy_to_safetensors(tmp_path / "d3.pt", tmp_path / "d3.safetensors") + states = [] + for name in ("d3.pt", "d3.safetensors"): + with envs.FASTVIDEO_H3_VAE_NVFP4_DECODER.override(str(tmp_path / name)): + vae = _load_h3_vae(tmp_path / "vae") + assert len(vae.decoder.transformer_blocks) == 3 and vae.decode_autocast_dtype == torch.bfloat16 + states.append({**vae.decoder.state_dict(), **vae.post_quant_conv.state_dict()}) + assert states[0].keys() == states[1].keys() + assert all(torch.equal(states[0][key], states[1][key]) for key in states[0]) + + +def test_safetensors_reader_rejects_foreign_files(tmp_path): + from safetensors.torch import save_file + + from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import load_nvfp4_decoder_checkpoint + + save_file({"decoder.x": torch.zeros(2)}, str(tmp_path / "plain.safetensors")) + with pytest.raises(ValueError, match="format"): + load_nvfp4_decoder_checkpoint(str(tmp_path / "plain.safetensors")) + save_file({"decoder.x": torch.zeros(2)}, str(tmp_path / "nometa.safetensors"), + metadata={"format": "fastvideo_h3_decoder_nvfp4_packed_v1"}) + with pytest.raises(ValueError, match="nvfp4_metadata"): + load_nvfp4_decoder_checkpoint(str(tmp_path / "nometa.safetensors")) + + +def _fake_packed_checkpoint(vae, act_scale="unit"): + """A packed-format checkpoint with arbitrary (not GEMM-meaningful) codes, for CPU-side contract tests.""" + from fastvideo.models.vaes.minimax_h3_nvfp4_checkpoint import NVFP4_DECODER_PACKED_FORMAT + + names = nvfp4_decoder_linear_names(vae.decoder) + state = {} + generator = torch.Generator().manual_seed(3) + for key, value in vae.decoder.state_dict().items(): + module, _, param = key.rpartition(".") + if module in names and param == "weight": + out_features, in_features = value.shape + state[f"{module}.packed_weight"] = torch.randint(0, 256, (out_features, in_features // 2), + dtype=torch.uint8, generator=generator) + state[f"{module}.weight_inv_scale"] = torch.randint(0, 256, (out_features * in_features // 16, ), + dtype=torch.uint8, generator=generator) + state[f"{module}.weight_global_sf"] = torch.rand((), generator=generator) + state[f"{module}.input_amax"] = torch.zeros(()) + elif module in names: + state[key] = value.bfloat16() + else: + state[key] = torch.randn(value.shape, generator=generator).to(value.dtype) + return { + "format": NVFP4_DECODER_PACKED_FORMAT, + "metadata": { + "act_scale": act_scale, + "rotation_group": None, + "skip_blocks": [], + "num_linears": len(names), + "student_layers": len(vae.decoder.transformer_blocks), + }, + "decoder": state, + "post_quant_conv": vae.post_quant_conv.state_dict(), + } + + +def test_loader_applies_packed_safetensors_without_master_weights(tmp_path): + import fastvideo.envs as envs + from fastvideo.models.vaes.minimax_h3_nvfp4_checkpoint import write_nvfp4_decoder_safetensors + from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import keep_evenly_spaced_blocks, nvfp4_linears + + _write_vae_dir(tmp_path / "vae", decoder_layers=5) + student = _tiny_h3_vae(_tiny_vae_arch(5)) + keep_evenly_spaced_blocks(student.decoder, 3) + checkpoint = _fake_packed_checkpoint(student) + write_nvfp4_decoder_safetensors(str(tmp_path / "p3.safetensors"), checkpoint) + with envs.FASTVIDEO_H3_VAE_NVFP4_DECODER.override(str(tmp_path / "p3.safetensors")): + vae = _load_h3_vae(tmp_path / "vae") + layers = nvfp4_linears(vae.decoder) + assert len(vae.decoder.transformer_blocks) == 3 and len(layers) == checkpoint["metadata"]["num_linears"] + assert all(layer.frozen and layer.weight is None and layer.act_scale == "unit" for layer in layers) + assert vae.decode_autocast_dtype == torch.bfloat16 + named = dict(vae.decoder.named_modules()) + for key, value in checkpoint["decoder"].items(): + module, _, param = key.rpartition(".") + actual = getattr(named[module], param) + assert actual.dtype == value.dtype and torch.equal(actual.cpu(), value), key + assert not any(parameter.requires_grad for parameter in vae.parameters()) + + +def test_packed_loader_refuses_mismatched_architecture(tmp_path): + from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import apply_nvfp4_decoder_checkpoint + + checkpoint = _fake_packed_checkpoint(_tiny_h3_vae(_tiny_vae_arch(2, heads=4))) + with pytest.raises(ValueError, match="does not fit"): + apply_nvfp4_decoder_checkpoint(_tiny_h3_vae(_tiny_vae_arch(2, heads=2)), checkpoint, freeze=True) + checkpoint = _fake_packed_checkpoint(_tiny_h3_vae(_tiny_vae_arch(2))) + del checkpoint["decoder"]["transformer_blocks.0.ff.net.2.packed_weight"] + with pytest.raises(ValueError, match="missing"): + apply_nvfp4_decoder_checkpoint(_tiny_h3_vae(_tiny_vae_arch(2)), checkpoint, freeze=True) + + +@pytest.mark.parametrize("act_scale", ["unit", "dynamic", "static"]) +def test_packed_loader_needs_input_amax_only_for_static_scales(act_scale): + from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import apply_nvfp4_decoder_checkpoint, nvfp4_linears + + checkpoint = _fake_packed_checkpoint(_tiny_h3_vae(_tiny_vae_arch(2)), act_scale=act_scale) + for key in [key for key in checkpoint["decoder"] if key.endswith(".input_amax")]: + del checkpoint["decoder"][key] + vae = _tiny_h3_vae(_tiny_vae_arch(2)) + if act_scale == "static": + with pytest.raises(ValueError, match="missing.*input_amax"): + apply_nvfp4_decoder_checkpoint(vae, checkpoint, freeze=True) + return + apply_nvfp4_decoder_checkpoint(vae, checkpoint, freeze=True) + layers = nvfp4_linears(vae.decoder) + assert layers and all(layer.input_amax.item() == 0 for layer in layers) + + +def test_export_deploy_strips_base_layer_from_wrapped_training_keys(tmp_path, monkeypatch): + import importlib.util + import json + import sys + from pathlib import Path + + expected = _write_deploy_checkpoint(tmp_path / "reference.pt", student_layers=2, base_layers=2) + reference = torch.load(tmp_path / "reference.pt", weights_only=False) + names = nvfp4_decoder_linear_names(_tiny_h3_vae(_tiny_vae_arch(2)).decoder) + wrapped = {} + for key, value in expected.items(): + module, _, param = key.rpartition(".") + # Store the training master in fp32 under a wrapper's base_layer, as a QAD run may. + wrapped[f"{module}.base_layer.{param}" if module in names else key] = value.float() + torch.save({"decoder": wrapped, "post_quant_conv": reference["post_quant_conv"], "step": 7}, tmp_path / "last.pt") + (tmp_path / "metadata.json").write_text(json.dumps(reference["metadata"])) + script = Path(__file__).parents[3] / "scripts/distill/minimax_h3_nvfp4_decoder/export_deploy.py" + spec = importlib.util.spec_from_file_location("export_deploy", script) + export_deploy = importlib.util.module_from_spec(spec) + spec.loader.exec_module(export_deploy) + monkeypatch.setattr(sys, "argv", ["export_deploy", str(tmp_path / "last.pt"), str(tmp_path / "deploy.pt")]) + export_deploy.main() + + exported = torch.load(tmp_path / "deploy.pt", weights_only=False)["decoder"] + assert exported.keys() == expected.keys() + for key, value in expected.items(): + module, _, param = key.rpartition(".") + dtype = torch.bfloat16 if module in names and param in ("weight", "bias") else torch.float32 + assert exported[key].dtype == dtype and torch.equal(exported[key], value.float().to(dtype)), key + + +class _DecoderHost(nn.Module): + """The two VAE submodules an NVFP4 decoder export covers.""" + + def __init__(self, decoder: nn.Module) -> None: + super().__init__() + self.decoder = decoder + self.post_quant_conv = nn.Conv3d(24, 24, 1) + + +@_needs_fp4 +@pytest.mark.parametrize("fused", [True, False]) +@pytest.mark.parametrize("act_scale", ["unit", "static"]) +def test_packed_export_round_trip_is_bit_identical(tmp_path, fused, act_scale): + from fastvideo.models.vaes.minimax_h3_nvfp4_checkpoint import (packed_nvfp4_decoder_checkpoint, + read_nvfp4_decoder_safetensors, + write_nvfp4_decoder_safetensors) + from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import (apply_nvfp4_decoder_checkpoint, + freeze_nvfp4_linears, nvfp4_linears) + + reference = _fused_test_decoder(act_scale) + with torch.no_grad(): + for layer in nvfp4_linears(reference): + layer.input_amax.fill_(4.0) + freeze_nvfp4_linears(reference) + host = _DecoderHost(reference).cuda() + metadata = {"act_scale": act_scale, "rotation_group": None, "skip_blocks": [], "student_layers": 2} + write_nvfp4_decoder_safetensors(str(tmp_path / "p.safetensors"), packed_nvfp4_decoder_checkpoint(host, metadata)) + + loaded = _DecoderHost(_fused_test_decoder(act_scale, convert=False)).cuda() + with torch.no_grad(): + for parameter in loaded.parameters(): + parameter.zero_() # every value must come from the file + checkpoint = read_nvfp4_decoder_safetensors(str(tmp_path / "p.safetensors"), "cuda") + torch.cuda.reset_peak_memory_stats() + before = torch.cuda.memory_allocated() + apply_nvfp4_decoder_checkpoint(loaded, checkpoint, freeze=True) + # No master weight: the dense linears are only released; at most a few scalars are allocated, + # far below one bf16 copy of the smallest quantized weight (512 x 512). + assert torch.cuda.max_memory_allocated() - before < 512 * 512 * 2 // 16 + assert all(layer.frozen for layer in nvfp4_linears(loaded.decoder)) + if not fused: + reference.fused_blocks_forward = None + loaded.decoder.fused_blocks_forward = None + latents = torch.randn(2, 24, 3, 8, 8, device="cuda") + with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): + assert torch.equal(loaded.decoder(latents), reference(latents)) diff --git a/pyproject.toml b/pyproject.toml index 11bb254d4..10c3edae9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -119,10 +119,11 @@ members = ["apps/dreamverse", "apps/infinite_livestream"] # PyPI's aarch64 fastvideo-kernel wheels do not cover the NVIDIA DGX Spark / # GB10 (sm_121), so on ARM build from the in-tree source instead. The FastH3 # extra also selects the in-tree source on x86_64 so a PR checkout can build -# its matching ptx_blackwell route before that version reaches PyPI. Other x86_64 -# installs keep the published wheel. This table is uv-only: plain pip ignores -# it. Pass --no-sources to bypass it (e.g. when the kernel is already -# installed, as in the Spark install guide). +# its matching ptx_blackwell route and key-length-sized sparse FP4 entry point +# before that version reaches PyPI. Other x86_64 installs keep the published +# wheel. This table is uv-only: plain pip ignores it. Pass --no-sources to +# bypass it (e.g. when the kernel is already installed, as in the Spark +# install guide). fastvideo-kernel = [ { path = "fastvideo-kernel", marker = "platform_machine == 'aarch64'" }, { path = "fastvideo-kernel", marker = "platform_machine == 'x86_64'", extra = "fasth3" }, diff --git a/scripts/checkpoint_conversion/convert_minimax_h3_modelopt_nvfp4_dit.py b/scripts/checkpoint_conversion/convert_minimax_h3_modelopt_nvfp4_dit.py index 463a8371c..ba03be7d4 100644 --- a/scripts/checkpoint_conversion/convert_minimax_h3_modelopt_nvfp4_dit.py +++ b/scripts/checkpoint_conversion/convert_minimax_h3_modelopt_nvfp4_dit.py @@ -34,6 +34,16 @@ adds a calibrated static activation scale per linear (``_nvfp4_input_global_sf`` = 448 * 6 / amax) from a JSON of input amax keyed ``b.`` (e.g. ``b3.ff.fc_in``); without it activations use the unit global scale, which saturates inputs above 2688 (H3's ``ff.fc_out``). +Every FFN linear needs an entry; attention projections and gates without one +keep the unit scale, as in the V2 Consumer layout (OmniRef: +``--quantize-ffn --act-amax amax.json --quantize-attention --quantize-gate`` +on the bf16 ``transformer_ref``). + +``--skip-blocks 42-49`` keeps every linear of those main blocks in bf16: they +stay in the dense shards, are absent from the packed file, and the export's +safetensors metadata lists them (``skip_blocks``) so the loader leaves them +unquantized. FastH3 OmniRef QAD ships this way: this converter packs its +student's bf16 ``transformer_ref``. Every exported linear is probed on random BF16 rows through the same ``mm_fp4`` path the loader runs; the relative error against a BF16 matmul with @@ -78,6 +88,40 @@ def fastvideo_module_name(diffusers_prefix: str) -> str: return name +def act_amax_for(module: str, amax_table: dict[str, float], required: bool) -> float | None: + """The calibrated input amax of ``module``, or None for the unit scale. + + Keys are ``b.`` or the full module name (the QAD exports use the latter). + FFN linears must have an entry; attention projections and gates without one keep the unit + activation scale (the V2 Consumer layout calibrates the FFN only). + """ + block = re.match(r"transformer_blocks\.(\d+)\.(.+)$", module) + keys = (f"b{block.group(1)}.{block.group(2)}", module) if block else (module, ) + for key in keys: + if key in amax_table: + return amax_table[key] + if required: + raise SystemExit(f"--act-amax has no entry {keys[0]!r} for {module}; nothing written") + return None + + +def parse_blocks(text: str) -> tuple[int, ...]: + """``"42-49"`` or ``"3,7,42-49"`` -> sorted block indices.""" + blocks: set[int] = set() + for part in filter(None, (p.strip() for p in text.split(","))): + low, _, high = part.partition("-") + start, end = int(low), int(high or low) + if end < start: + raise ValueError(f"Invalid block range {part!r}: end ({end}) is below start ({start})") + blocks.update(range(start, end + 1)) + return tuple(sorted(blocks)) + + +def block_of(prefix: str) -> int | None: + match = re.match(r"transformer_blocks\.(\d+)\.", prefix) + return int(match.group(1)) if match else None + + def dequantize_modelopt(packed: torch.Tensor, scale: torch.Tensor, scale_2: torch.Tensor) -> torch.Tensor: """E2M1 low nibble = even column, high nibble = odd column (ModelOpt / CUTLASS order).""" lut = _E2M1.to(packed.device) @@ -162,6 +206,8 @@ def main() -> None: parser.add_argument("--quantize-gate", action="store_true", help="also quantize attn.to_gate_compress") parser.add_argument("--quantize-ffn", action="store_true", help="quantize the FFN linears from a bf16 source") parser.add_argument("--act-amax", type=Path, help="JSON of calibrated input amax per linear") + parser.add_argument("--skip-blocks", type=parse_blocks, default=(), + help="main blocks kept in bf16, e.g. 42-49 (recorded in the export metadata)") parser.add_argument("--max-probe-error", type=float, default=0.3) parser.add_argument("--dense-shard-gb", type=float, default=5.0) args = parser.parse_args() @@ -182,7 +228,10 @@ def main() -> None: modelopt_set = set(modelopt) attention = sorted(k[:-len(".weight")] for k in weight_map if k.endswith(".weight") and dense_pattern.match(k[:-len(".weight")]) - and k[:-len(".weight")] not in modelopt_set) if dense_pattern else [] + and k[:-len(".weight")] not in modelopt_set + and block_of(k) not in args.skip_blocks) if dense_pattern else [] + if any(block_of(p) in args.skip_blocks for p in modelopt): + parser.error("--skip-blocks names blocks whose linears the ModelOpt source already quantized") if args.quantize_ffn and any(_BLOCK_FFN.match(p) for p in modelopt): parser.error("--quantize-ffn needs a bf16 source; this one already holds ModelOpt FFN linears") amax_table = None @@ -214,15 +263,14 @@ def main() -> None: raise SystemExit(f"probe error {error:.3f} on {prefix} exceeds {args.max_probe_error}; nothing written") module = fastvideo_module_name(prefix) if amax_table is not None: - block = re.match(r"transformer_blocks\.(\d+)\.(.+)$", module) - key = f"b{block.group(1)}.{block.group(2)}" - if key not in amax_table: - raise SystemExit(f"--act-amax has no entry {key!r} for {module}; nothing written") - buffers["_nvfp4_input_global_sf"] = torch.tensor((448.0 * 6.0) / max(amax_table[key], 1e-12), - dtype=torch.float32) + amax = act_amax_for(module, amax_table, required=bool(_BLOCK_FFN.match(prefix))) + if amax is not None: + buffers["_nvfp4_input_global_sf"] = torch.tensor((448.0 * 6.0) / max(amax, 1e-12), + dtype=torch.float32) for name, value in buffers.items(): export[f"{module}::{name}"] = value.cpu() - save_file(export, str(args.dst / EXPORT_FILENAME)) + metadata = {"skip_blocks": json.dumps(list(args.skip_blocks))} if args.skip_blocks else None + save_file(export, str(args.dst / EXPORT_FILENAME), metadata=metadata) dense_keys = [k for k in weight_map if k not in modelopt_keys and k not in attention_keys] shard, shard_bytes, shards = {}, 0, [] @@ -254,6 +302,7 @@ def main() -> None: "gate_linears": sum(1 for p in attention if _BLOCK_GATE.match(p)), "min_scale_byte_agreement": round(min(agreements), 4) if agreements else None, "mean_scale_byte_agreement": round(sum(agreements) / len(agreements), 4) if agreements else None, + "skip_blocks": list(args.skip_blocks), "dense_tensors": len(dense_keys), "dense_shards": len(shards)})) diff --git a/scripts/distill/minimax_h3_nvfp4_decoder/export_deploy.py b/scripts/distill/minimax_h3_nvfp4_decoder/export_deploy.py new file mode 100644 index 000000000..7c3b7f5d9 --- /dev/null +++ b/scripts/distill/minimax_h3_nvfp4_decoder/export_deploy.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Export a QAD training checkpoint to an inference-only file. + +Drops the optimizer state and stores each ``NVFP4DecoderLinear`` master weight +and bias in bf16: both the training forward and inference quantize from the +bf16 cast (``weight.to(compute_dtype)``), so the packed NVFP4 weights, and +therefore every decoded pixel, are unchanged. All other tensors (norms, +scales, register tokens, ``proj_in``/``proj_out``, ``input_amax``) keep their +dtype. The run's ``metadata.json`` (student layers, act scale, rotation) is +embedded so the file is self-describing. +""" +from __future__ import annotations + +import argparse +import json +import re +from pathlib import Path + +import torch + +# Same pattern as minimax_h3_nvfp4_decoder.DECODER_BLOCK_LINEAR; inlined so the export runs on GPU-less hosts +# (importing fastvideo initializes Triton, which needs a driver). +DECODER_BLOCK_LINEAR = re.compile(r"^transformer_blocks\.(\d+)\.(attn\.to_q|attn\.to_k|attn\.to_v|attn\.to_out\.0" + r"|ff\.net\.0\.proj|ff\.net\.2)$") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("checkpoint", help="QAD training checkpoint (best.pt / last.pt)") + parser.add_argument("output", help="destination .pt") + parser.add_argument("--metadata", help="run metadata.json (default: next to the checkpoint)") + args = parser.parse_args() + + state = torch.load(args.checkpoint, map_location="cpu", mmap=True, weights_only=False) + metadata_path = Path(args.metadata or Path(args.checkpoint).with_name("metadata.json")) + metadata = {**state.get("metadata", {}), **json.loads(metadata_path.read_text())} + decoder = {} + for key, value in state["decoder"].items(): + # A wrapped training linear stores its tensors under ``.base_layer``; the + # inference decoder names them without it and loads strictly. + key = re.sub(r"\.base_layer(\.[^.]+)$", r"\1", key) + module, _, param = key.rpartition(".") + nvfp4 = DECODER_BLOCK_LINEAR.match(module) is not None + decoder[key] = value.to(torch.bfloat16) if nvfp4 and param in ("weight", "bias") else value.clone() + exported = { + "format": "fastvideo_h3_decoder_nvfp4_deploy_v1", + "decoder": decoder, + "post_quant_conv": {k: v.clone() for k, v in state["post_quant_conv"].items()}, + "metadata": { + **metadata, "source_step": state.get("step") + }, + } + torch.save(exported, args.output) + size = Path(args.output).stat().st_size / 2**30 + print(json.dumps({"output": args.output, "gib": round(size, 2), "metadata": exported["metadata"]})) + + +if __name__ == "__main__": + main() diff --git a/scripts/distill/minimax_h3_nvfp4_decoder/export_packed.py b/scripts/distill/minimax_h3_nvfp4_decoder/export_packed.py new file mode 100644 index 000000000..abeebe0a1 --- /dev/null +++ b/scripts/distill/minimax_h3_nvfp4_decoder/export_packed.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Export a pre-packed NVFP4 decoder (``fastvideo_h3_decoder_nvfp4_packed_v1`` safetensors). + +Loads an exported bf16-master decoder (``export_deploy.py`` .pt or its safetensors) into a dense +H3 VAE exactly as the pipeline loader does, freezes it on a Blackwell GPU, and writes the frozen +layers' packed tensors. The loader then skips the bf16 masters entirely. + + python export_packed.py decoder.safetensors /path/to/MiniMax-H3/vae decoder-nvfp4.safetensors +""" +from __future__ import annotations + +import argparse +import glob +import json +import os + +import torch + +from fastvideo.models.vaes.minimax_h3_nvfp4_checkpoint import ( + packed_nvfp4_decoder_checkpoint, write_nvfp4_decoder_safetensors) +from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import ( + apply_nvfp4_decoder_checkpoint, load_nvfp4_decoder_checkpoint) + + +def _load_h3_vae(vae_dir: str, device: torch.device) -> torch.nn.Module: + """The dense H3 VAE the decoder checkpoint was distilled against.""" + from safetensors.torch import load_file + + from fastvideo.configs.models.vaes.minimax_h3_video import MiniMaxH3VideoVAEConfig + from fastvideo.models.vaes.minimax_h3_int8_convrot import dense_vae_safetensors + from fastvideo.models.vaes.minimax_h3_video import AutoencoderKLMiniMaxH3 + + config = MiniMaxH3VideoVAEConfig() + with open(os.path.join(vae_dir, "config.json")) as handle: + config.update_model_arch(json.load(handle)) + vae = AutoencoderKLMiniMaxH3(config).to(device) + state: dict[str, torch.Tensor] = {} + for path in dense_vae_safetensors(sorted(glob.glob(os.path.join(vae_dir, "*.safetensors")))): + state.update(load_file(path)) + vae.load_state_dict(state, strict=True) + return vae.eval().requires_grad_(False) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("decoder", help="bf16-master decoder (.pt from export_deploy.py, or .safetensors)") + parser.add_argument("host_vae_dir", help="MiniMax-H3 vae/ folder with at least the decoder's block count") + parser.add_argument("output", help="destination .safetensors") + args = parser.parse_args() + if not torch.cuda.is_available(): + raise SystemExit("packing needs a CUDA Blackwell GPU (FlashInfer FP4 quantization)") + + device = torch.device("cuda", torch.cuda.current_device()) + vae = _load_h3_vae(args.host_vae_dir, device) + metadata = apply_nvfp4_decoder_checkpoint(vae, load_nvfp4_decoder_checkpoint(args.decoder), freeze=True) + import flashinfer + + major, minor = torch.cuda.get_device_capability(device) + packed_on = f"{torch.cuda.get_device_name(device)} sm_{major}{minor}" + packed = packed_nvfp4_decoder_checkpoint(vae, { + **metadata, "packed_on": packed_on, + "flashinfer_version": flashinfer.__version__ + }) + write_nvfp4_decoder_safetensors(args.output, packed) + print(json.dumps({"output": args.output, "bytes": os.path.getsize(args.output), "metadata": packed["metadata"]})) + + +if __name__ == "__main__": + main()